Backprop through a two-layer ReLU network
Hard25 min10 hidden tests
A pricing team trains a small network that predicts k numbers per listing (for example, nightly price and cleaning fee). The network is Y_hat = relu(X @ W1 + b1) @ W2 + b2, and the loss is the mean squared error over every entry of Y_hat - Y. Write mlp_grads(X, Y, W1, b1, W2, b2) that returns the loss and the gradient of the loss with respect to each of the four parameters, computed by backpropagation. In the first example the loss is 1.0 and dW2 is [[4.0], [0.0]].
Run it forward, then backward
loss = mean((ŷ − y)²)dL/dŷ = 2(ŷ − y) = 2. dW2 = h × 2 = [4, 0], db2 = 2. Back through W2: dh = [2, 4]; the ReLU passes it only where z > 0, so dz = [2, 0] (unit 2 sits at z ≤ 0 and passes nothing). dW1 = x × dz = [[2, 0], [4, 0]], db1 = [2, 0].
Loss 1.
Examples
The second hidden unit sits exactly at 0, so it passes no gradient
mlp_grads([[1.0, 2.0]], [[1.0]], [[1.0, -1.0], [0.5, 0.5]], [0.0, 0.0], [[1.0], [2.0]], [0.0])
returns {'loss': 1, 'dW1': [[2, 0], [4, 0]], 'db1': [2, 0], 'dW2': [[4], [0]], 'db2': [2]}
Two listings, two outputs: the mean divides by 4
mlp_grads([[1.0, 0.0], [0.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.5], [-0.5, 1.0]], [0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]], [0.0, 0.0])
returns {'loss': 0.8125, 'dW1': [[0.5, -0.25], [0, 0.5]], 'db1': [0.5, 0.25], 'dW2': [[0.5, -0.25], [-0.25, 0.375]], 'db2': [0, 0.25]}
Requirements
- Shapes:
X(n, d),Y(n, k),W1(d, h),b1(h,),W2(h, k),b2(k,). Lists or arrays. - The loss is
mean((Y_hat - Y) ** 2)over all n * k entries, as PyTorch'sMSELossdoes by default. - Take the ReLU's derivative as 1 where its input is above 0 and 0 where it is 0 or below.
- Return a dict with
"loss"(a float) and"dW1","db1","dW2","db2"(float arrays shaped like their parameters).
The idea behind it
Backpropagation is the chain rule applied layer by layer, from the loss back to the inputs. Each layer receives the gradient of the loss with respect to its output and returns the gradients for its parameters and its input. For a dense layer out = A @ W + b, the rules are dW = A.T @ dout, db = dout.sum(axis=0) and dA = dout @ W.T. Each gradient has the same shape as the thing it belongs to, which is a quick check. ReLU passes gradient where its input was positive and blocks it elsewhere. Its derivative at exactly 0 is undefined, and frameworks pick 0, as here. The loss sets the starting gradient. For a mean of squares over N entries, it is 2 * (Y_hat - Y) / N. Writing this by hand once makes autograd far easier to debug, because you know what each gradient should look like.