Backpropagation, and where the memory goes

The problem it solves

You have a network with 7 billion parameters and you need the derivative of the loss with respect to each one. The naive approach is to nudge each parameter, rerun the network, and see how the loss changed. That is 7 billion forward passes per training step. At a tenth of a second each, one step would take twenty-two years.

Backpropagation gets all 7 billion in roughly the cost of two forward passes. The trick is that most of the arithmetic is shared, and doing it backwards is what lets you share it.

The idea in one paragraph

Going forward, you compute and keep the intermediate values. Going backward, you carry a single quantity — the derivative of the loss with respect to the current layer's output — and at each layer you do two things with it: use it to compute this layer's weight gradients, and convert it into the same quantity for the layer below. One sweep, one pass, every gradient.

Worked, on a two-layer network

Small enough to check with a pen. One input, two layers, squared loss.

Forward:
  z1 = w1 * x           (cache x)
  a1 = relu(z1)         (cache z1)
  z2 = w2 * a1          (cache a1)
  L  = (z2 - y)^2

Numbers: x = 2, w1 = 0.5, w2 = 3, y = 1.

z1 = 1.0
a1 = relu(1.0) = 1.0
z2 = 3.0
L  = (3 - 1)^2 = 4

Now backwards. Start with the derivative of the loss with respect to its own input.

dL/dz2 = 2(z2 - y) = 2(3 - 1) = 4

layer 2: dL/dw2 = dL/dz2 * a1 = 4 * 1.0 = 4
         dL/da1 = dL/dz2 * w2 = 4 * 3   = 12       <- pass this down

layer 1: dL/dz1 = dL/da1 * relu'(z1) = 12 * 1 = 12   (relu' is 1 since z1 > 0)
         dL/dw1 = dL/dz1 * x = 12 * 2 = 24

Two weight gradients, 4 and 24, obtained in one backward sweep. Notice that computing dL/dw2 needed a1, and computing dL/dw1 needed x and the sign of z1. Every one of those is a forward value that had to be kept. That is where the memory goes.

The memory arithmetic

Here is the practical consequence, and it is the reason training a model needs so much more memory than running it.

For inference, you need the weights and one layer's activations at a time; earlier activations can be discarded as soon as they have been used. For training, every activation needed by the backward pass must survive from the forward pass until the backward pass reaches it. So peak memory scales with depth × batch size × activation size, on top of the weights.

Work a rough example. A 7-billion-parameter model, 32 layers, batch of 8 sequences of 2,048 tokens, hidden width 4,096, activations in 16-bit:

one layer's activations = 8 * 2048 * 4096 * 2 bytes = 134 MB
32 layers, several tensors kept per layer, say 4:
   32 * 4 * 134 MB = 17 GB of activations

against 14 GB for the weights themselves in 16-bit. Activations are the larger half, and they scale with batch size while the weights do not. That is why "reduce the batch size" is the first answer to an out-of-memory error, and why it works when nothing about the model has changed.

Gradient checkpointing, and the trade you are making

The standard remedy is gradient checkpointing: keep only a few activations — say every fourth layer — and recompute the rest during the backward pass. Memory falls by roughly the square root of the depth; compute rises by about 30 per cent, since you do part of the forward pass twice.

It is worth understanding as a pure trade. You are choosing to spend time to save memory, and it is the right trade whenever the alternative is not being able to run at all. One flag in most frameworks.

Two things worth knowing about the real implementation

The graph is built as you go. In PyTorch, every operation on a tensor with requires_grad=True records itself. Calling .backward() walks that record in reverse. This is why a Python if inside a model is fine — the graph records whichever branch actually ran.

Gradients accumulate rather than replace. Calling .backward() twice adds the second set to the first, which is a feature — it is how gradient accumulation lets you simulate a large batch on a small card — and a trap, because forgetting optimizer.zero_grad() silently sums every batch you have ever seen and your model diverges for no visible reason.

The asymmetry worth knowing

Backpropagation is reverse-mode automatic differentiation, and it is cheap when there are many inputs and one output. That is exactly the shape of a loss function: billions of parameters in, one scalar out.

Forward-mode differentiation is the reverse trade and is cheap when there is one input and many outputs. If you ever need the derivative of many outputs with respect to a single parameter, forward mode is the efficient choice, and every serious framework provides it. The reason you rarely hear about it is simply that loss functions return one number.

The rule to keep

Backpropagation is bookkeeping: cache on the way forward, reuse on the way back. The cache is why training memory is dominated by activations, and knowing that turns most out-of-memory errors into a choice between batch size and checkpointing rather than a mystery.