Gradient Checkpointing for Memory

Backpropagation needs the intermediate outputs of every layer from the forward pass in order to compute each layer's local gradient — you can't apply the chain rule at a node if you never kept that node's value around. The straightforward way to run backpropagation is to simply store every one of those intermediate values as the forward pass runs, then walk backward through them. For a network with hundreds of layers, or very large batches, that storage bill can exceed the memory available on a single GPU, becoming the actual bottleneck to training a bigger model at all — not compute, but memory.

The analogy: keeping every draft versus reconstructing one on demand

Imagine writing a long essay in a hundred numbered drafts, each one small edit from the last. If you want to compare draft 100 back to draft 3 to see exactly what changed, the easy way is to have kept every single draft on your desk. But if desk space is the scarce resource, you could instead keep only every tenth draft, and when you need draft 47 specifically, retype it by re-applying the edits from draft 40 forward — a bit of extra work, but far less paper on the desk at any one time. Gradient checkpointing does exactly this trade with a neural network's intermediate activations: keep only a subset (the "checkpoints"), and when backpropagation needs a value that wasn't kept, recompute it on the fly from the nearest checkpoint.

The trade: memory down, compute up

Ordinary training stores every activation, using memory that grows linearly with the number of layers $L$: roughly $O(L)$ memory, and each activation computed exactly once.

Gradient checkpointing instead stores only a handful of checkpoint activations — often placed every $\sqrt{L}$ layers — and during the backward pass, recomputes the activations between two checkpoints by re-running the forward pass over just that segment, right when the backward pass needs them:

$$\text{memory} \sim O(\sqrt{L}), \qquad \text{extra forward compute} \sim O(1) \text{ extra pass}$$

In words: placing checkpoints roughly every $\sqrt{L}$ layers means you never need to reconstruct more than about $\sqrt{L}$ layers' worth of activations at once, cutting peak memory from growing linearly with depth to growing with its square root — at the cost of roughly one extra forward pass's worth of recomputation overall, since every activation between checkpoints gets computed twice (once during the original forward pass, discarded, then again during the backward pass).

Worked example 1: memory without checkpointing

A network has 100 layers, each layer's activation taking 1 unit of memory to store. Without checkpointing: total stored memory is $100 \times 1 = 100$ units, since every layer's output is kept until the backward pass reaches it.

Worked example 2: memory with checkpointing

Same 100-layer network, checkpointing every $\sqrt{100}=10$ layers — so checkpoints are kept at layers 10, 20, 30, ..., 100: that's 10 stored checkpoint activations. When the backward pass reaches, say, layer 47, it needs the activations for layers 41 through 49 — it recomputes just that 9-layer segment forward from the checkpoint at layer 40, uses those values for the backward step, then discards them again. Peak memory at any moment is roughly the 10 checkpoints plus the small segment being recomputed — around $10+10=20$ units, a fifth of the original 100, at the cost of re-running the forward computation for the in-between layers an extra time.

What this means in practice

Gradient checkpointing is a standard feature in modern deep learning frameworks, usually a one-line flag, and is essential for training very deep networks or large batches on limited GPU memory — it's often what turns "out of memory" into a model that fits and trains successfully. The cost is real: training runs measurably slower (commonly 20–30% more wall-clock time) because of the extra recomputation, so it's a deliberate trade of speed for the ability to fit a bigger model or batch at all.

Practice

1. A network has 400 layers; checkpoints are placed every $\sqrt{400}=20$ layers. How many checkpoints are stored, and roughly how many layers get recomputed at once?
2. Why does checkpointing every single layer (checkpoint spacing of 1) provide no memory benefit at all?
3. In one sentence, why is checkpointing a compute-for-memory trade rather than a free lunch?