Neural Network Forward Pass in NumPy: MLP From Scratch

Given a batch X (N × D) and parameters W1 (D × H), b1 (H), W2 (H × C), b2 (C), return the N × C matrix of class probabilities P = softmax(max(0, XW1 + b1)W2 + b2), with the softmax taken over each row.

Examples

Input X = [[1, 2]], W1 = [[1, 0], [0, 1]], b1 = [0, 0], W2 = [[1, 0], [0, 1]], b2 = [0, 0]

Output [[0.2689, 0.7311]]

Identity weights: the logits are the inputs, and softmax([1, 2]) follows.

Input X = [[-1, -1]], W1 = [[1, 0], [0, 1]], b1 = [0, 0], W2 = [[1, 0], [0, 1]], b2 = [0, 0]

Output [[0.5, 0.5]]

Every hidden unit is off. The logits are equal, so the output is uniform.

Input X = [[1, 0], [0, 1]], W1 = [[2, -1], [-1, 2]], b1 = [0.5, 0.5], W2 = [[1, -1], [-1, 1]], b2 = [0, 0.1]

Output [[0.9926, 0.0074], [0.0061, 0.9939]]

Two examples at once. Each row sums to 1.

Constraints

1 ≤ N ≤ 256, 1 ≤ D, H ≤ 128, 2 ≤ C ≤ 32
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Three lines: a linear layer, a ReLU, a linear layer. Then a softmax over axis 1.

Hint 2

b1 has shape (H,). Adding it to an (N, H) matrix broadcasts it to every row, which is what you want.

Hint 3

Subtract each row's maximum before exponentiating: keepdims=True keeps the shapes compatible.

Where candidates go wrong

Taking the softmax over the whole matrix, or over axis 0, so that columns sum to 1 instead of rows.
Using np.max(x, 0) for ReLU. That is the maximum along axis 0, not an elementwise maximum with zero. Use np.maximum.
Shifting by the maximum of the whole batch instead of each row.

The idea

A linear layer is one matrix product plus a broadcast bias. ReLU is np.maximum(0, ·). The softmax is taken per row, shifted by each row's own maximum so that the largest exponent is 0.

Everything is a whole-batch operation, so the cost is the two matrix products, O(N(DH + HC)), with no Python loop at all.

Time O(N(DH + HC)) · Space O(N(H + C))

Reference solution

import numpy as np

def mlp_forward(X, W1, b1, W2, b2):
    X, W1, b1, W2, b2 = (np.asarray(a, dtype=float) for a in (X, W1, b1, W2, b2))
    H = np.maximum(0.0, X @ W1 + b1)          # (N, H); the bias broadcasts over rows
    Z = H @ W2 + b2                            # (N, C) logits
    Z = Z - Z.max(axis=1, keepdims=True)       # per-row shift
    E = np.exp(Z)
    return E / E.sum(axis=1, keepdims=True)

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

A warm-up that interviewers use to see how you handle shapes and broadcasting. It is also the first half of the backpropagation problem, so getting it clean pays twice.

The theory behind it is in the lesson Backpropagation from scratch.

Keep going

Easy10 min

Forward pass of a two-layer network