Backpropagation Step by Step: The Forward and Backward Pass

Before we start

Gradient descent is an optimisation method used to minimise a loss function, and minimising that loss is mainly what training a model means.

Many good sources cover the method itself. Foundations of Learning is the closest neighbour to this article and the source of the chain notation in section 02. Gradient Descent: Downhill to a Minimum takes the linear algebra route, which is the one that matters here. For the method as a tool you call, the scikit-learn guide to SGD is the reference, and Google's Machine Learning Crash Course is a good place to practise it. For the intuition, I strongly recommend the two videos below. They are the clearest visual account there is, and if you do not know 3Blue1Brown yet (shame on you ;)), the rest of the channel is just as good.

3Blue1Brown, Neural networks, chapter 2Gradient descent, how neural networks learn

3Blue1Brown, Neural networks, chapter 4Backpropagation calculus

Forward pass and backpropagation on a classic example

Take a multilayer perceptron with one hidden layer and two output classes: a linear map $W_{1}$ gives $z$ , a ReLU gives $h$ , a second linear map $W_{2}$ gives the logits $u$ (one score per class), a softmax turns them into probabilities $y$ , and a cross-entropy loss compares $y$ with the label. Three inputs, four hidden units, two classes, twenty weights, biases aside.

The label $t$ is one-hot: 1 for the true class, 0 for the other. To tell cats from dogs, a cat is $t = (1, 0)$ and a dog $t = (0, 1)$ .

Training means lowering $L$ , the loss, which measures the gap between the probabilities $y$ and the label $t$ . In its simplest form, gradient descent does it one update at a time: compute $\frac{\partial L}{\partial w}$ for every weight $w$ , move each weight a small step against its own derivative, and repeat.

The difficulty is that $W_{1}$ sits at the beginning of the network and $L$ at the far end, with a ReLU, $W_{2}$ and a softmax in between, so there is no direct expression for $\frac{\partial L}{\partial W _{1}}$ to read off.

Backpropagation gets around this with the chain rule: it carries one vector back through the network, the derivative of the loss with respect to whatever the pass has reached. At the hidden layer we write it $g_{h} = \frac{\partial L}{\partial h}$ , which says how fast the loss changes as each coordinate of the hidden representation $h$ moves. Once the vector reaches a linear map, the gradient of the loss with respect to its weights needs only two things:

The gradient that arrived at the layer’s output.
The input the layer received on the forward pass.

Writing $W x$ rather than $W x + b$ looks like a shortcut. Append a constant 1 to the representation and the bias becomes one more column of $W$ , the one multiplying that 1, so nothing that follows changes. Biases do not come up again, and large language models mostly drop them anyway: PaLM uses none and reports steadier training at scale, Llama has none in its linear layers, and Cramming found that removing them makes each step cheaper without slowing the loss. In the numbers here the bias columns are zero and every weight has one decimal place, so every sum has at most four terms and you can check the whole pass on paper.

ReLU set the two negative entries of $z$ to zero, and on the way back it zeroes the same two coordinates of the gradient: the weights into and out of those units get no update from this example.

The loss depends on the input through $z$ , $h$ , $u$ and $y$ in turn, so its derivative with respect to anything early is a sum over every path by which that quantity reached the end. The vector carried backward, $g$ , is that sum. Written as matrix products, each step of it multiplies by the transpose of one Jacobian, the matrix of derivatives of a step’s outputs with respect to its inputs.

A ReLU Jacobian is diagonal: each coordinate is scaled on its own, by 1 where the unit was open and 0 where it was closed. A softmax Jacobian is full: every coordinate of the outgoing gradient mixes all the incoming ones, so both logits receive a gradient even though only one class is true. Through a deep network these factors multiply, one per layer, and whether their product shrinks or grows decides whether the early layers learn at all: the vanishing and exploding gradient problem, which a following article takes up.

Whatever sits inside the box (the two-layer network just folded, a convolution, an attention head or an entire transformer layer) takes a representation and a set of parameters and returns a representation.

If you would rather watch a network like this train than read about it, A Neural Network Playground puts a network of the same kind in a browser tab, with a live decision boundary and a loss curve that either falls or does not.

Going deeper: the same chain with more layers

Nothing in that walk depended on there being two layers. A deep network is the same chain of functions with more links in it: layer $l$ maps one representation to the next, $x_{l} = f_{l} (x_{l - 1}, θ_{l})$ , carrying its own parameters $θ_{l}$ , and the last one feeds the loss:

x_{0} f_{1} θ_{1} x_{1} f_{2} θ_{2} \dots f_{L} θ_{L} x_{L} ⟶ L

Training needs $\frac{\partial L}{\partial θ _{l}}$ for every layer, and the chain rule gets there through the representations in between. Write $g_{l} = \frac{\partial L}{\partial x _{l}}$ for the gradient arriving at representation $l$ , and $J_{l} = \frac{\partial x _{l}}{\partial x _{l - 1}}$ for the Jacobian of layer $l$ .

On the network of section 01, count each linear map with the activation after it as one layer, the usual way of counting, so that every $f_{l}$ carries weights and $L = 2$ . Any other grouping gives the same gradients, since a product of Jacobians does not depend on where the brackets go.

One convention, kept from here on: gradients are column vectors. With it, the backward step is a transpose.

The gradient that reaches a layer does two jobs. Through $J_{l}$ , it carries on to the layer below. Through the derivative of the layer with respect to its own weights $θ_{l}$ , it says how the loss responds when each of those weights moves, which is what training needs to update them:

g_{l - 1} \nabla_{θ_{l}} L = J_{l}^{⊤} g_{l} = (\frac{\partial f _{l}}{\partial θ _{l}})^{⊤} g_{l}

For a linear map, the second line is an outer product of the gradient arriving at the layer with the representation that entered it:

\frac{\partial L}{\partial W _{l}} = \frac{\partial L}{\partial x _{l}} x_{l - 1}^{⊤} = g_{l} x_{l - 1}^{⊤}

Assembled, those lines are one step of training:

One step of training

forward →

x_{l} = f_{l} (x_{l - 1}, θ_{l}) l = 1, \dots, L

← seed

g_{L} = \frac{\partial L}{\partial x _{L}}

← backward

g_{l - 1} \nabla_{θ_{l}} L = J_{l}^{⊤} g_{l} l = L, \dots, 1 = (\frac{\partial f _{l}}{\partial θ _{l}})^{⊤} g_{l}

update

θ_{l} \leftarrow θ_{l} - η \nabla_{θ_{l}} L

The learning rate $η$ sets how far one step goes.

Most layers are built from Jacobians of three shapes, all three met in section 01: a dense matrix for a linear map, a diagonal one for an elementwise activation $φ$ , and a dense symmetric one for a softmax:

J_{W} J_{φ} J_{soft} = W = diag (φ^{'} (z)) = diag (y) - y y^{⊤}

ReLU is the case $φ^{'} (z) = 1 {z > 0}$ . Normalisation and attention are built from these plus one more piece, a product of two values that both depend on the input. Normalisation multiplies the input by the inverse of its own length $∥ x ∥ = x_{1}^{2} + \dots + x_{n}^{2}$ , which gives $\frac{x}{∥ x ∥}$ ; attention multiplies the queries by the keys, $Q K^{⊤}$ . Every one of them is a case of the generic Jacobian, and all are written out entry by entry.

What one backward sweep costs

One backward sweep returns the gradient of every weight. On the network of section 01 its cost can be counted by hand, as the multiplications in its matrix products, since everything else is a handful of comparisons and exponentials:

The multiplications of one step

operationmultiplicationsmemory

forward →

z = W_{1} x

4 \times 3 = 12

keeps

x

h = ReLU (z)

0

keeps

z

u = W_{2} h

2 \times 4 = 8

keeps

h

y = softmax (u)

0

keeps

y

← backward

g_{u} = y - t

0

reads

y

\frac{\partial L}{\partial W _{2}} = g_{u} h^{⊤}

2 \times 4 = 8

reads

h

g_{h} = W_{2}^{⊤} g_{u}

4 \times 2 = 8

reads

W_{2}

g_{z} = 1 {z > 0} ⊙ g_{h}

0

reads

z

\frac{\partial L}{\partial W _{1}} = g_{z} x^{⊤}

4 \times 3 = 12

reads

x

g_{x} = W_{1}^{⊤} g_{z}

0

skipped, no weights below

nudge each weight and run forward again

20 \times 20 = 400

keeps nothing

forward20

backward28

nudging400

Laid out along the network, the top row is the forward pass, the middle row passes the gradient down and the bottom row builds the weight gradients, each product with its cost under it:

x W_{1} 12 12 g_{z} x^{⊤} h g_{z} W_{2} 8 W_{2}^{⊤} 8 8 g_{u} h^{⊤} u g_{u}

Write $C_{forward}$ and $C_{backward}$ for the multiplications of the forward pass and of the backward sweep, the totals of the two halves of the table. Forward, each weight is used in exactly one multiplication, so the forward pass costs as many multiplications as there are weights, $P$ :

C_{forward} = P = 4 \cdot 3 + 2 \cdot 4 = 12 + 8 = 20

Backward, each linear map does the two jobs of section 02, one product each. The first builds its weight gradient, the outer product of the gradient at its output with its input. That uses every weight once, so over both maps it costs $P$ again. The second passes the gradient down through the transpose, so that the map below can build its own weight gradient: $W_{2}^{⊤} g_{u}$ gives $g_{h}$ , the ReLU turns it into $g_{z}$ , and $W_{1}$ needs $g_{z}$ for its gradient. Below $W_{1}$ there is only the input. Passing down through it would give $g_{x} = W_{1}^{⊤} g_{z}$ , the gradient with respect to the data, and no weight needs it, so the sweep stops there. $C_{backward}$ is $P$ plus one pass down, through $W_{2}$ :

C_{backward} = P + W_{2}^{⊤} g_{u} 8 = 20 + 8 = 28 \approx 1.40 P

The same count holds for any stack of linear maps. A map of shape $m \times n$ costs $mn$ multiplications forward and $2 mn$ backward, except the first, which costs $mn$ backward because it passes nothing down. With $P_{1}$ the weights of the first map, and nudging each weight in turn for one extra forward pass:

C_{forward} C_{backward} C_{nudging} = P = 2 P - P_{1} \leq 2 P = P \cdot C_{forward} = P^{2}

The network of section 01 has 20 weights, and the three totals are 20, 28 and 400. With a billion weights the backward sweep is about $2 \times 1 0^{9}$ multiplications and nudging is $1 0^{18}$ , for a finite-difference approximation of the same numbers. The factor 2 is particular to stacked linear maps, but a bound of the same kind holds for any function built from elementary operations: its gradient costs at most a small constant times the function itself, whatever $P$ is (Griewank and Walther, chapter 4). It comes from running backward: the loss is one number and the weights are many, so one sweep from the loss reaches all of them, while carrying derivatives forward with the computation takes one sweep per weight again (Baydin et al.).

The price is memory. The backward rows of the table read values the forward pass produced, so those values stay alive until the sweep comes back for them: $x$ for the gradient of $W_{1}$ , $z$ for the ReLU mask, $h$ for the gradient of $W_{2}$ and $y$ for $g_{u} = y - t$ . Calling their count $K$ :

K = x 3 + z, h 4 + 4 + y 2 = 13

That is 13 numbers against 20 weights. The weights are stored once, but the kept values are stored again for every example of a batch, so with a batch of $B$ they cost $13 B$ . That is why the kept values (the “activations” of training-memory accounts), not the weights, usually set the memory bill of training, and why checkpointing, which keeps only some of them and recomputes the rest during the backward sweep, buys memory back with time.

Three exercises

1The same input, the other label

Keep the input of section 01 but label it a dog, $t = (0, 1)$ . The forward pass never reads the label, so $z$ , $h$ , $u$ and $y$ do not change: $y = (0.2108, 0.7892)$ , and units 1 and 4 are still closed.

Compute $g_{u}$ , the gradient at the logits.
Carry it back to $g_{z}$ , through $g_{h} = W_{2}^{⊤} g_{u}$ and the ReLU, with:
$W_{2} = (1 - 0.3 - 0.7 0.5 0.4 1.2 0.9 - 0.6)$
Write $G_{cat}$ for the gradient of the loss with respect to $W_{1}$ that section 01 found with the cat label, and $G_{dog}$ for the one the dog label gives. Both are $g_{z} x^{⊤}$ , each with its own $g_{z}$ , and they differ by a single factor:
$G_{dog} = c G_{cat}$
Find $c$ without computing either matrix, and say why one number is enough.

Solution

$g_{u} = y - t = (0.2108, - 0.2108)$ . With the cat label, section 01 had $(- 0.7892, 0.7892)$ : the model already puts 0.7892 on dog, so with this label it is right, and the gradient is smaller.
$g_{h} = W_{2}^{⊤} g_{u}$ gives $(0.2741, - 0.253, - 0.1687, 0.3162)$ , and the closed units 1 and 4 delete their coordinates: $g_{z} = (0, - 0.253, - 0.1687, 0)$ .
With two classes $y_{1} + y_{2} = 1$ , so the gradient at the logits, $g_{u} = y - t$ , is:
$t = (1, 0) : g_{u} t = (0, 1) : g_{u} = (- y_{2}, y_{2}) = (y_{1}, - y_{1}) = - \frac{y _{1}}{y _{2}} (- y_{2}, y_{2})$
Every later step multiplies by something the forward pass fixed, the same under both labels: $W_{2}^{⊤}$ , the ReLU mask, then $x^{⊤}$ . A factor in front of $g_{u}$ comes out of each of them unchanged, so it is the factor between the two matrices:
$c = - \frac{y _{1}}{y _{2}} = - 0.2671$
Rows 1 and 4 of both matrices, those of the closed units, are zero.

2The softmax Jacobian

The softmax turns the logits $u$ into probabilities:

y_{k} = \frac{e ^{u_{k}}}{\sum _{j} e ^{u_{j}}}

Show that $\frac{\partial y _{k}}{\partial u _{j}} = y_{k} (δ_{k j} - y_{j})$ , where $δ_{k j}$ is 1 when $k = j$ and 0 otherwise. In matrix form, $J_{soft} = diag (y) - y y^{⊤}$ .
Show that the entries of $g_{u} = J_{soft}^{⊤} g_{y}$ sum to zero whatever $g_{y}$ , so whatever loss follows the softmax.
Add the same number $a$ to every logit. Show that $y$ does not change, and explain why part 2 already said so: what is the derivative of the loss along that shift?

Solution

Write $S$ for the denominator, so that $y_{k} = \frac{e ^{u_{k}}}{S}$ . The numerator depends on $u_{k}$ alone, while $S$ depends on every logit:
$S = j \sum e^{u_{j}} = j \neq = k \sum e^{u_{j}} + e^{u_{k}}, \frac{\partial S}{\partial u _{j}} = e^{u_{j}}$
When $k = j$ , both the numerator and $S$ depend on $u_{k}$ , and the quotient rule gives:
$\frac{\partial y _{k}}{\partial u _{k}} = \frac{e ^{u_{k}} S - e ^{u_{k}} e ^{u_{k}}}{S ^{2}} = \frac{e ^{u_{k}}}{S} - (\frac{e ^{u_{k}}}{S})^{2} = y_{k} (1 - y_{k})$
When $k \neq = j$ , only $S$ depends on $u_{j}$ :
$\frac{\partial y _{k}}{\partial u _{j}} = - \frac{e ^{u_{k}} e ^{u_{j}}}{S ^{2}} = - \frac{e ^{u_{k}}}{S} \cdot \frac{e ^{u_{j}}}{S} = - y_{k} y_{j}$
The factor $δ_{k j}$ picks the case, so one formula covers both:
$y_{k} (δ_{k j} - y_{j}) = {y_{k} (1 - y_{k}) - y_{k} y_{j} k = j, δ_{k j} = 1 k \neq = j, δ_{k j} = 0$
Entry $(k, j)$ of $diag (y)$ is $y_{k} δ_{k j}$ and that of $y y^{⊤}$ is $y_{k} y_{j}$ , which gives the matrix form. The network of section 01 has two classes, so $1 - y_{1} = y_{2}$ and every entry is $\pm y_{1} y_{2}$ :
$y = (0.2108, 0.7892) J_{soft} = (y_{1} y_{2} - y_{1} y_{2} - y_{1} y_{2} y_{1} y_{2}) = (0.1664 - 0.1664 - 0.1664 0.1664)$
Entry $j$ of $g_{u} = J_{soft}^{⊤} g_{y}$ collects, over every probability $y_{k}$ , how much $y_{k}$ moves with $u_{j}$ :
$g_{u, j} = k \sum (J_{soft})_{k j} g_{y, k}$
Add up all the entries, then swap the two sums so that the inner one runs along row $k$ of $J_{soft}$ :
$j \sum g_{u, j} = j \sum k \sum (J_{soft})_{k j} g_{y, k} = k \sum g_{y, k} (j \sum (J_{soft})_{k j})$
Every row adds up to zero, because the probabilities sum to one:
$j \sum y_{k} (δ_{k j} - y_{j}) = y_{k} - y_{k} j \sum y_{j} = y_{k} - y_{k} = 0$
So each term of the outer sum is $g_{y, k} \cdot 0$ , and the entries of $g_{u}$ add up to zero. The argument never used the values of $g_{y}$ , so it holds for any loss. Written with the vector of ones $1$ , the sum is $1^{⊤} g_{u} = 0$ .
Adding $a$ to every logit multiplies the numerator and the denominator by $e^{a}$ :
$\frac{e ^{u_{k} + a}}{\sum _{j} e ^{u_{j} + a}} = \frac{e ^{a} e ^{u_{k}}}{e ^{a} \sum _{j} e ^{u_{j}}} = y_{k}$
So $y$ does not change, and neither does the loss. The gradient says the same thing: the rate at which the loss changes as every logit moves together is its component along $1$ , which part 2 found to be zero:
$\frac{d}{d a} L (u + a 1) = 1^{⊤} g_{u} = 0$

3The cost of a real network

A classifier of handwritten digits takes 784 pixels, has two hidden layers of 512 ReLU units and ends on 10 classes, biases aside: $784 \to 512 \to 512 \to 10$ . Count as the table of section 03 does: multiplications in the matrix products, and for memory the values the backward sweep reads.

How many weights $P$ does it have, and how many multiplications $C_{backward}$ does the backward sweep of one example take?
How many numbers $K$ does one example keep for the backward sweep?
From which batch size $B$ do the kept values outnumber the weights?

Solution

Write $W_{1}$ , $W_{2}$ and $W_{3}$ for the three maps. A map of shape $m \times n$ has $mn$ weights. Forward, each weight is used in one multiplication, so a layer costs as many multiplications as it has weights. Backward, it costs them twice: once for its weight gradient, once to pass the gradient down through its transpose, so that the layer below can build its own. $W_{1}$ pays only once: passing down through it would give the gradient with respect to the pixels, and no weight needs it. Along the network, the top row is the forward pass, the middle row passes the gradient down and the bottom row builds the weight gradients, each product with its cost under it:
$x W_{1} 401408 401408 g_{z_{1}} x^{⊤} h_{1} g_{z_{1}} W_{2} 262144 W_{2}^{⊤} 262144 262144 g_{z_{2}} h_{1}^{⊤} h_{2} g_{z_{2}} W_{3} 5120 W_{3}^{⊤} 5120 5120 g_{u} h_{2}^{⊤} u g_{u}$
Layer by layer, forward:
$z_{1} z_{2} u = W_{1} x = W_{2} h_{1} = W_{3} h_{2} 512 \cdot 784 512 \cdot 512 10 \cdot 512 = 401408 = 262144 = 5120$
and backward, from the top, where the ReLU masks cost nothing:
$\frac{\partial L}{\partial W _{3}} g_{h_{2}} g_{z_{2}} \frac{\partial L}{\partial W _{2}} g_{h_{1}} g_{z_{1}} \frac{\partial L}{\partial W _{1}} = g_{u} h_{2}^{⊤} = W_{3}^{⊤} g_{u} = 1 {z_{2} > 0} ⊙ g_{h_{2}} = g_{z_{2}} h_{1}^{⊤} = W_{2}^{⊤} g_{z_{2}} = 1 {z_{1} > 0} ⊙ g_{h_{1}} = g_{z_{1}} x^{⊤} 10 \cdot 512 10 \cdot 512 512 \cdot 512 512 \cdot 512 512 \cdot 784 = 5120 = 5120 = 0 = 262144 = 262144 = 0 = 401408$
Adding up each column:
layer shape weights forward backward
$W_{1}$ $512 \times 784$ $401408$ $401408$ $401408$
$W_{2}$ $512 \times 512$ $262144$ $262144$ $262144 + 262144$
$W_{3}$ $10 \times 512$ $5120$ $5120$ $5120 + 5120$
total $P = 668672$ $668672$ $935936$
The backward column is $P$ for the weight gradients plus what passing down costs in $W_{2}$ and $W_{3}$ :
$C_{backward} = P + W_{2}^{⊤} g_{z_{2}} 262144 + W_{3}^{⊤} g_{u} 5120 = 668672 + 267264 = 935936 \approx 1.40 P$
Nudging each weight in turn would cost one forward pass per weight, $P \cdot P \approx 4.47 \times 1 0^{11}$ multiplications.
The backward sweep reads $y$ for $g_{u} = y - t$ , each $z$ for its ReLU mask, each $h$ and $x$ for the weight gradient of the layer above it:
$K = x 784 + z_{1}, h_{1} 512 + 512 + z_{2}, h_{2} 512 + 512 + y 10 = 2842$
A batch of $B$ examples keeps $2842 B$ numbers, and the weights are stored once:
$2842 B > 668672 ⟺ B > \frac{668 672}{2 842} \approx 235.3$
From $B = 236$ on, one step keeps more values than the network has weights, and a batch of 256 is already past it.

layer	shape	weights	forward	backward
$W_{1}$	$512 \times 784$	$401408$	$401408$	$401408$
$W_{2}$	$512 \times 512$	$262144$	$262144$	$262144 + 262144$
$W_{3}$	$10 \times 512$	$5120$	$5120$	$5120 + 5120$
total		$P = 668672$	$668672$	$935936$

References

The reading list is what this article takes as known: gradient descent itself, and a network to train in a browser. The references are where each claim of the prose and of the proof panels is established, with a note on what each one was cited for.

Get in touch

We hope you enjoyed this article. If you have questions, found a bug, or just want to chat about AI and engineering, we'd love to hear from you.

Future articles?

We're planning more deep dives into:

Graph Neural Networks
Transformer Architecture
System Design for ML

Get them in your inbox

Backpropagation Step by Step: The Forward and Backward Pass

Forward pass and backpropagation on a classic example

Going deeper: the same chain with more layers

What one backward sweep costs

Three exercises

1The same input, the other label

2The softmax Jacobian

3The cost of a real network

References

Get in touch

Future articles?

Send feedback