Two-Layer MLP Forward Pass

Implement the forward pass of a two-layer multilayer perceptron (MLP), the smallest network that is more than a single linear layer. The hidden layer applies an affine transform followed by a ReLU nonlinearity, and the output layer applies a second affine transform to produce the logits. Interviewers use this to check that you can chain dense layers correctly, place the nonlinearity where it belongs (after the hidden layer, not after the output), and optionally turn logits into a valid probability distribution with a numerically stable softmax.

Definition

Given a batch X of shape (n, d_in), the network computes two stages:

h      = relu(X @ W1 + b1)     # (n, d_in) @ (d_in, d_hid) + (d_hid,) -> (n, d_hid)
logits = h @ W2 + b2           # (n, d_hid) @ (d_hid, d_out) + (d_out,) -> (n, d_out)

relu(z) = max(0, z) applied elementwise. The logits are the raw, unbounded output scores. If probabilities are requested, apply a row-wise softmax:

probs[i, j] = exp(logits[i, j] - max_k logits[i, k]) / sum_j exp(logits[i, j] - max_k logits[i, k])

Subtracting the per-row max before exp is a stability trick that avoids overflow without changing the result.

Task

Complete mlp_forward(X, W1, b1, W2, b2, return_probs=False). Compute the hidden activations h = relu(X @ W1 + b1), then the logits = h @ W2 + b2. Return logits when return_probs is False. When return_probs is True, also compute a numerically stable row-wise softmax of the logits and return the tuple (logits, probs). The nonlinearity is applied only after the hidden layer, never after the output layer. Everything must be vectorized with NO Python loops over samples. Use plain NumPy only (no PyTorch, TensorFlow, or Keras).

Example

X = np.array([[1.0, 2.0],
              [-1.0, 0.0]])
W1 = np.array([[1.0, -1.0, 0.0],
               [0.0,  1.0, 1.0]])
b1 = np.array([0.0, 0.0, 0.0])
W2 = np.array([[1.0, 0.0],
               [0.0, 1.0],
               [0.0, 1.0]])
b2 = np.array([0.0, 0.0])

mlp_forward(X, W1, b1, W2, b2)
# -> array([[1., 3.],
#           [0., 1.]])

logits, probs = mlp_forward(X, W1, b1, W2, b2, return_probs=True)
# logits -> array([[1., 3.],
#                  [0., 1.]])
# probs  -> array([[0.11920292, 0.88079708],
#                  [0.26894142, 0.73105858]])

For row 0, X @ W1 + b1 = [1, 1, 2]; ReLU leaves it unchanged since every entry is positive, so h = [1, 1, 2], and h @ W2 = [1, 3]. For row 1, X @ W1 + b1 = [-1, 1, 0]; ReLU clamps the -1 to 0, giving h = [0, 1, 0] and logits [0, 1].

Constraints

• 1 <= n <= 10^4; 1 <= d_in, d_hid, d_out <= 10^3; use vectorized NumPy, no Python loop over samples.
• Inner dimensions must agree: X.shape[1] == W1.shape[0], W1.shape[1] == b1.shape[0] == W2.shape[0], W2.shape[1] == b2.shape[0].
• Softmax must be numerically stable (subtract the per-row max before exponentiating); each returned probability row sums to 1.
• Values fit in float64; outputs are dense (n, d_out) float arrays.

Approach

An MLP is just dense layers chained with a nonlinearity between them. Compute the hidden pre-activation X @ W1 + b1, pass it through elementwise ReLU (np.maximum(0, z)) to get h, then run a second affine h @ W2 + b2 to get the logits. The nonlinearity goes only between the two layers; the output layer is left linear so its raw scores are the logits. When probabilities are requested, apply a row-wise softmax, subtracting each row’s max before exponentiating so the exponentials never overflow.

Solution

import numpy as np
from typing import Tuple, Union

def mlp_forward(
    X: np.ndarray,
    W1: np.ndarray,
    b1: np.ndarray,
    W2: np.ndarray,
    b2: np.ndarray,
    return_probs: bool = False,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    X = np.asarray(X, dtype=np.float64)
    W1, b1 = np.asarray(W1, dtype=np.float64), np.asarray(b1, dtype=np.float64)
    W2, b2 = np.asarray(W2, dtype=np.float64), np.asarray(b2, dtype=np.float64)

    h = np.maximum(0.0, X @ W1 + b1)     # (n, d_in) @ (d_in, d_hid) + b1 -> relu -> (n, d_hid)
    logits = h @ W2 + b2                 # (n, d_hid) @ (d_hid, d_out) + b2 -> (n, d_out)

    if not return_probs:
        return logits

    z = logits - logits.max(axis=1, keepdims=True)   # stable softmax: shift each row by its max
    exp = np.exp(z)
    probs = exp / exp.sum(axis=1, keepdims=True)      # normalize each row to sum to 1
    return logits, probs

Walkthrough

On the example, X = [[1, 2], [-1, 0]], W1 = [[1, -1, 0], [0, 1, 1]], b1 = 0, W2 = [[1, 0], [0, 1], [0, 1]], b2 = 0:

• Hidden pre-activation X @ W1 + b1. Row 0: [1*1 + 2*0, 1*(-1) + 2*1, 1*0 + 2*1] = [1, 1, 2]. Row 1: [-1*1 + 0*0, -1*(-1) + 0*1, -1*0 + 0*1] = [-1, 1, 0].
• ReLU. Row 0 has no negatives, so h = [1, 1, 2]. Row 1’s -1 is clamped to 0, so h = [0, 1, 0].
• Output affine h @ W2 + b2. Row 0: [1*1 + 1*0 + 2*0, 1*0 + 1*1 + 2*1] = [1, 3]. Row 1: [0, 1]. That is the returned logits = [[1, 3], [0, 1]].
• With return_probs=True, softmax on row 0 subtracts the max 3: exp([-2, 0]) = [0.1353, 1], normalized to [0.1192, 0.8808]. Row 0’s larger logit 3 dominates. Row 1 gives [0.2689, 0.7311]. Each row sums to 1.

Complexity & notes

• Time O(n * (d_in * d_hid + d_hid * d_out)), dominated by the two matmuls that BLAS runs in optimized C; ReLU and softmax are O(n * d_hid) and O(n * d_out), negligible. Space O(n * d_hid) for the hidden activations plus O(n * d_out) for the output.
• The two subtleties interviewers watch for: the nonlinearity belongs only between layers (ReLU after the hidden layer, nothing after the output, so the logits stay unbounded), and softmax must be numerically stable. Skipping the per-row max subtraction lets exp overflow to inf on large logits, producing nan after the division; subtracting the max shifts the largest exponent to exp(0) = 1 and leaves the normalized result unchanged.
• ReLU as np.maximum(0.0, z) is fully vectorized; a Python loop over entries would be correct but orders of magnitude slower and a red flag in an interview.
• Both keepdims=True reductions are load-bearing: max and sum collapse the d_out axis, and keeping the dimension lets the result broadcast back against the (n, d_out) logits row-wise. Without it the shapes would misalign or broadcast along the wrong axis.
• Returning the logits (not just probabilities) matches how real classifiers are trained: loss functions like cross-entropy operate on logits directly for stability, and softmax is applied only when you need an explicit probability distribution.