MLP Forward & Backward Pass - Neural Networks

MLP Forward & Backward Pass

What you'll build

Two functions for a 2-layer MLP classifier: mlp_forward computes class probabilities and a cache, and mlp_backward uses the cache to compute the gradients of the mean cross-entropy loss with respect to every weight and bias, by hand (backpropagation).

Background

Backprop is the chain rule applied layer by layer, from the loss back to the inputs, reusing values saved during the forward pass. "Write the backward pass of a small MLP in NumPy" is one of the most common ML coding interview questions, and getting the shapes right is most of the work.

Architecture

X (m, n_features) -> [W1, b1] -> ReLU -> (m, n_hidden) -> [W2, b2] -> softmax -> (m, n_classes)

Forward pass

Z1 = X @ W1 + b1          # (m, n_hidden)
A1 = relu(Z1)             # (m, n_hidden)
Z2 = A1 @ W2 + b2         # (m, n_classes), the logits
output = softmax(Z2)      # row-wise probabilities
cache = (X, Z1, A1, Z2, W1, W2)

Backward pass

The loss is the mean cross-entropy L = -(1/m) * sum(Y * log(output)), with Y one-hot. Softmax followed by cross-entropy has a famously simple gradient with respect to the logits, output - Y (per example; the 1/m comes from the mean):

dZ2 = output - Y                      # (m, n_classes)
dW2 = (1/m) * A1.T @ dZ2              # (n_hidden, n_classes)
db2 = (1/m) * sum(dZ2, axis=0)        # (n_classes,)
dA1 = dZ2 @ W2.T                      # (m, n_hidden)
dZ1 = dA1 * (Z1 > 0)                  # ReLU passes gradient only where Z1 > 0
dW1 = (1/m) * X.T @ dZ1               # (n_features, n_hidden)
db1 = (1/m) * sum(dZ1, axis=0)        # (n_hidden,)

Every gradient has the same shape as its parameter: a quick sanity check.

Inputs

• X: (m, n_features); W1: (n_features, n_hidden); b1: (n_hidden,); W2: (n_hidden, n_classes); b2: (n_classes,)
• Y: one-hot labels (m, n_classes); output and cache come from mlp_forward

Expected return format

• mlp_forward returns the tuple (output, cache): output has shape (m, n_classes) with rows summing to 1, and cache is the tuple (X, Z1, A1, Z2, W1, W2) in that order.
• mlp_backward returns a dictionary with these keys:
  • 'dW1': (n_features, n_hidden)
  • 'db1': (n_hidden,)
  • 'dW2': (n_hidden, n_classes)
  • 'db2': (n_classes,)

Edge cases and conventions

• The ReLU derivative at exactly Z1 = 0 is taken as 0 (Z1 > 0).
• Gradients are averaged over the batch (divide by m), matching a mean loss.
• The relu and softmax helpers (stable, row-wise) are provided in the starter code.

Constraints

NumPy only, fully vectorised over the batch (no loops over examples). Don't modify the inputs.

Follow-up

• Check your gradients numerically: perturb one weight by ±1e-5 and compare (L+ - L-) / 2e-5 with your analytic value.
• Add L2 weight decay: what changes in dW1 and dW2? (Add λ * W.)

Input: mlp_forward(np.array([[1.0, 0.0], [0.0, 1.0]]), np.eye(2), np.zeros(2), np.eye(2), np.zeros(2))[0]
Output: [[0.7311, 0.2689], [0.2689, 0.7311]] (rounded)

Input: X = [[1, 2]], W1 = I, b1 = 0, W2 = zeros((2, 2)), b2 = 0, Y = [[1, 0]]: mlp_backward(Y, *mlp_forward(X, W1, b1, W2, b2))
Output: {'dW1': [[0, 0], [0, 0]], 'db1': [0, 0], 'dW2': [[-0.5, 0.5], [-1.0, 1.0]], 'db2': [-0.5, 0.5]}

import numpy as np

def relu(x):
    return np.maximum(0, x)

def softmax(x):
    exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
    return exp_x / np.sum(exp_x, axis=1, keepdims=True)

def mlp_forward(X, W1, b1, W2, b2):
    """
    Forward pass for a 2-layer MLP (ReLU hidden layer, softmax output).

    Args:
        X: Input data (m, n_features)
        W1: First layer weights (n_features, n_hidden)
        b1: First layer bias (n_hidden,)
        W2: Second layer weights (n_hidden, n_classes)
        b2: Second layer bias (n_classes,)

    Returns:
        output: Class probabilities (m, n_classes), rows sum to 1
        cache: Tuple (X, Z1, A1, Z2, W1, W2) for backprop
    """
    # Your code here
    pass

def mlp_backward(Y, output, cache):
    """
    Backward pass for the 2-layer MLP with mean cross-entropy loss.

    Args:
        Y: One-hot labels (m, n_classes)
        output: Probabilities from mlp_forward (m, n_classes)
        cache: Tuple (X, Z1, A1, Z2, W1, W2) from mlp_forward

    Returns:
        grads: dict with 'dW1' (n_features, n_hidden), 'db1' (n_hidden,),
               'dW2' (n_hidden, n_classes), 'db2' (n_classes,)
    """
    X, Z1, A1, Z2, W1, W2 = cache
    m = X.shape[0]  # batch size

    # Your code here
    pass