Backpropagation step by step: the forward and backward pass

wpnews.pro

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.

Forward pass and backpropagation on a classic example #

Take a multilayer perceptron with one hidden layer and two output classes: a linear map gives , a ReLU gives , a second linear map gives the logits (one score per class), a softmax turns them into probabilities , and a cross-entropy lossBinary cross-entropy, opens a note compares with the label. Three inputs, four hidden units, two classes, twenty weights, biases aside.

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

Training means lowering , the loss, which measures the gap between the probabilities and the label . In its simplest form, gradient descent does it one update at a timeBatch size and optimiser, opens a panel: compute for every weight , move each weight a small step against its own derivative, and repeatClassic gradient descent, opens a note.

The difficulty is that sits at the beginning of the network and at the far end, with a ReLU, and a softmax in between, so there is no direct expression for 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 , which says how fast the loss changes as each coordinate of the hidden representation moves. Once the vector reaches a linear mapWhy a linear map, opens a note, 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 rather than looks like a shortcut. Append a constant 1 to the representation and the bias becomes one more column ofWhy a bias is a column, opens a note , 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 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 , , and 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, , 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 diagonalA diagonal Jacobian, opens a note: each coordinate is scaled on its own, by 1 where the unit was open and 0 where it was closed. A softmax Jacobian is fullA full Jacobian, opens a note: 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 maps one representation to the next, , carrying its own parameters , and the last one feeds the loss:

Training needs for every layer, and the chain rule gets there through the representations in between. Write for the gradient arriving at representation , and for the Jacobian of layer .

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 carries weights and . 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 transposeWhy a transpose, opens a panel.

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

For a linear map, the second line is an outer productThe gradient of a layer’s weights, opens a panel of the gradient arriving at the layer with the representation that entered it: Assembled, those lines are one step of training:

One step of training

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:

ReLU is the case . 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 , which gives ; attention multiplies the queries by the keys, . Every one of them is a case of the generic Jacobian, and all are written out entry by entryWhat is in a Jacobian, opens a panel.

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

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:

Write and 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, :

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 again. The second passes the gradient down through the transpose, so that the map below can build its own weight gradient: gives , the ReLU turns it into , and needs for its gradient. Below there is only the input. Passing down through it would give , the gradient with respect to the data, and no weight needs it, so the sweep stops there. is plus one pass down, through :

The same count holds for any stack of linear maps. A map of shape costs multiplications forward and backward, except the first, which costs backward because it passes nothing down. With the weights of the first map, and nudging each weight in turn for one extra forward pass:

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 multiplications and nudging is , 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 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: for the gradient of , for the ReLU mask, for the gradient of and for . Calling their count :

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 they cost . 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, . The forward pass never reads the label, so , , and do not change: , and units 1 and 4 are still closed.

Compute , the gradient at the logits.
Carry it back to , through and the ReLU, with:
Write for the gradient of the loss with respect to that section 01 found with the cat label, and for the one the dog label gives. Both are , each with its own , and they differ by a single factor:Find without computing either matrix, and say why one number is enough.

Solution, exercise 1 #

. With the cat label, section 01 had : the model already puts 0.7892 on dog, so with this label it is right, and the gradient is smaller.
gives , and the closed units 1 and 4 delete their coordinates: .
With two classes , so the gradient at the logits, , is:Every later step multiplies by something the forward pass fixed, the same under both labels: , the ReLU mask, then . A factor in front of comes out of each of them unchanged, so it is the factor between the two matrices:Rows 1 and 4 of both matrices, those of the closed units, are zero.

2The softmax Jacobian

The softmax turns the logits into probabilities:

Show that , where is 1 when and 0 otherwise. In matrix form, .
Show that the entries of sum to zero whatever , so whatever loss follows the softmax.
Add the same number to every logit. Show that does not change, and explain why part 2 already said so: what is the derivative of the loss along that shift?

Solution, exercise 2 #

Write for the denominator, so that . The numerator depends on alone, while depends on every logit:When , both the numerator and depend on , and the quotient rule gives:When , only depends on :The factor picks the case, so one formula covers both:Entry of is and that of is , which gives the matrix form. The network of section 01 has two classes, so and every entry is :
Entry of collects, over every probability , how much moves with :Add up all the entries, then swap the two sums so that the inner one runs along row of :Every row adds up to zero, because the probabilities sum to one:So each term of the outer sum is , and the entries of add up to zero. The argument never used the values of , so it holds for any loss. Written with the vector of ones , the sum is .
Adding to every logit multiplies the numerator and the denominator by :So 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 , which part 2 found to be zero:

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: . 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 does it have, and how many multiplications does the backward sweep of one example take?
How many numbers does one example keep for the backward sweep?
From which batch size do the kept values outnumber the weights?

Solution, exercise 3 #

Write , and for the three maps. A map of shape has 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. 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:Layer by layer, forward:and backward, from the top, where the ReLU masks cost nothing:Adding up each column:The backward column is for the weight gradients plus what passing down costs in and :layer shape weights forward backward total Nudging each weight in turn would cost one forward pass per weight, multiplications.
The backward sweep reads for , each for its ReLU mask, each and for the weight gradient of the layer above it:
A batch of examples keeps numbers, and the weights are stored once:From on, one step keeps more values than the network has weights, and a batch of 256 is already past it.

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.

- [Email contact@fremaconsulting.ch](mailto:contact@fremaconsulting.ch)
- [LinkedIn linkedin.com/company/frema-consulting](https://www.linkedin.com/company/frema-consulting/)
- [GitHub github.com/TomSOhm](https://github.com/TomSOhm)

Future articles?

We're planning more deep dives into:

Graph Neural Networks
Transformer Architecture
System Design for ML

Get them in your inbox

source & further reading

fremaconsulting.ch — original article