E2E: 2-Layer MLP with Backprop - End-to-End Implementations

Implement E2E: 2-Layer MLP with Backprop in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

End-to-End 2-Layer MLP with Backprop

What you'll build

A complete 2-layer classifier in plain NumPy: forward pass, softmax cross-entropy, hand-derived backward pass, an SGD step and a training loop. This is the "implement backprop from scratch" interview question.

Architecture

Step	Operation	Output shape
input	`X`	`(batch, input_dim)`
layer 1	`z1 = X @ W1 + b1`	`(batch, hidden_dim)`
activation	`h = relu(z1)`	`(batch, hidden_dim)`
layer 2	`logits = h @ W2 + b2`	`(batch, output_dim)`
loss	`mean_i −log softmax(logits)[i, y_i]`	scalar

The constructor (given) uses He initialization, W ~ N(0, 2 / fan_in), with np.random.seed(42), and zero biases.

Backward pass

With the loss averaged over the batch N, softmax and cross-entropy combine into a simple gradient:

d_logits = (probs - y_onehot) / N         # (N, output_dim)
dW2 = h.T @ d_logits                      # (hidden_dim, output_dim)
db2 = d_logits.sum(axis=0)                # (output_dim,)
d_h = d_logits @ W2.T                     # (N, hidden_dim)
d_z1 = d_h * (z1 > 0)                     # ReLU passes gradient only where z1 > 0
dW1 = X.T @ d_z1                          # (input_dim, hidden_dim)
db1 = d_z1.sum(axis=0)                    # (hidden_dim,)

Method contract

Method	Does	Returns
`forward(X)`	computes logits and stores `self.cache = {'X', 'z1', 'h', 'logits'}`	`logits (batch, output_dim)`
`compute_loss(logits, y)`	softmax (max-subtracted) + mean cross-entropy; stores `self.cache['probs']`	float
`backward(y)`	uses the cache from `forward` and `compute_loss`	dict `{'dW1', 'db1', 'dW2', 'db2'}`, each shaped like its parameter
`train_step(X, y, lr=0.01)`	forward, loss, backward, then `param -= lr * grad` for all four	the loss before the update (float)
`train_mlp(X_train, y_train, epochs=100, lr=0.01)`	builds `MLP(input_dim, 64, num_classes)` (classes = `len(np.unique(y_train))`) and calls `train_step` on the full batch each epoch	`(model, losses)`, with `len(losses) == epochs`

Inputs

X: float array (batch, input_dim)
y: int array (batch,) with class ids in [0, output_dim)

Edge cases and conventions

Clip the picked probabilities to [1e-15, 1] before the log (or use log-softmax) so a confident wrong prediction does not give log(0).
The ReLU derivative at exactly 0 is taken as 0 (z1 > 0).

Constraints

NumPy only; vectorise over the batch (no per-example loops).

Follow-up

How would you verify the gradients? (Finite differences: (L(w + ε) − L(w − ε)) / 2ε.)
Add L2 weight decay or momentum; what changes in backward and train_step?

Examples

Input: m = MLP(3, 2, 2); X = np.array([[1.0, -1.0, 0.5], [0.0, 2.0, -1.0]]); m.forward(X)
Output: array([[0.0, 0.0], [0.7148, 2.4115]])

Both hidden units are negative for row 0, so ReLU zeroes them and the logits equal b2 = 0

Input: m.compute_loss(m.cache["logits"], np.array([1, 0]))
Output: 1.2791

Row 0: −log 0.5 = 0.6931; row 1: −log softmax([0.7148, 2.4115])[0] = 1.8650; the mean is 1.2791

Input: [m.train_step(X, np.array([1, 0]), lr=0.5) for _ in range(3)]
Output: [1.2791, 0.4081, 0.3437]

Each value is the loss before that step, so the first equals the loss above

Practice in Python

import numpy as np

class MLP:
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        """Initialize a 2-layer MLP (given; keep as is)."""
        np.random.seed(42)
        # He initialization: N(0, 2 / fan_in), suited to ReLU
        self.W1 = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim)
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(hidden_dim, output_dim) * np.sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros(output_dim)

        # Cache for backward pass
        self.cache = {}

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

    def softmax(self, 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 forward(self, X: np.ndarray) -> np.ndarray:
        """
        Forward pass. Store X, z1, h and logits in self.cache.

        Args:
            X: Input (batch, input_dim)

        Returns:
            logits: (batch, output_dim)
        """
        # Your code here
        pass

    def compute_loss(self, logits: np.ndarray, y: np.ndarray) -> float:
        """
        Mean softmax cross-entropy. Store the softmax probabilities in self.cache['probs'].

        Args:
            logits: Model output (batch, output_dim)
            y: True labels (batch,) as integers

        Returns:
            loss: Python float
        """
        # Your code here
        pass

    def backward(self, y: np.ndarray) -> dict:
        """
        Backward pass for the mean loss, using self.cache.

        Args:
            y: True labels (batch,) as integers

        Returns:
            gradients: {'dW1', 'db1', 'dW2', 'db2'}, each shaped like its parameter
        """
        # Your code here
        pass

    def train_step(self, X: np.ndarray, y: np.ndarray, lr: float = 0.01) -> float:
        """
        Forward, loss, backward, SGD update of W1, b1, W2, b2.

        Returns:
            loss: Loss value before the update
        """
        # Your code here
        pass


def train_mlp(X_train, y_train, epochs=100, lr=0.01):
    """
    Build MLP(input_dim, 64, num_classes) and run `epochs` full-batch train steps.

    Returns:
        (model, losses): the trained MLP and the list of per-epoch losses
    """
    # Your code here
    pass

The interactive editor loads with Python and NumPy. Run your implementation against 9 tests, or use the progressive hints.

More exercises in this section →