Backpropagation, One VJP per Layer

Here is the chapter's network with real numbers: two inputs, a linear layer W one, a sigmoid, a second linear layer W two, and the squared error loss.

The forward pass runs from left to right and writes every value on the way. It keeps x and h, the amber tags, because the backward pass will need them.

Now the backward pass runs from right to left, in purple. The loss hands back y hat minus y, and each box multiplies by its own Jacobian, transposed: one V J P per box.

On the way back, the two linear layers drop off their weight gradients as outer products. Watch the cells of nabla W one light up: a column times a row.

Below, the loss is drawn as a surface over two of the weights, the first entry of W one and the first entry of W two. The dot is our network, and the red arrow is the gradient that backpropagation just computed: the steepest way up.

Does backpropagation tell the truth? The slider nudges the first weight. The dashed line predicts the change in the loss from the gradient, and the surface shows the real change. For small nudges they agree; for large ones they drift apart.

Your turn. Drag the slider, or drag the dot to move both weights, and find a nudge that lowers the loss.

Any small nudge against the gradient lowers the loss, at the rate backpropagation predicted, and the backward pass read only x and h, never z. Keep exploring: every number above follows the dot.