Checkpointing and Reverse-Mode AD

Here is the surprise ending of this course: the most abstract idea in it — Bennett's pebble game, the space–time bargain for undoing a computation — turns out to run the modern world. Every time a neural network learns, it plays Bennett's pebble game. The machinery is called reverse-mode automatic differentiation, better known as backpropagation, and its central problem is exactly the reversible-computing problem: to go backward through a computation, you need the states you passed through on the way forward.

Why going backward needs the forward states

Reverse-mode AD computes the gradient of an output with respect to every input by traversing the computation graph backwards, from output to input, applying the chain rule at each step. But the derivative of each step depends on the values that flowed through it on the forward pass. So you have two choices, and they are precisely Bennett's two choices:

Store everything. Keep every intermediate state of the forward pass. Memory grows in proportion to the depth n of the computation: O(n) memory, no recomputation. This is how naive backprop works, and it is why training deep networks is memory-hungry.
Recompute from checkpoints. Store only a few states (checkpoints) and, when the backward pass needs a state you threw away, recompute it forward from the nearest checkpoint. Less memory, more compute. This is exactly the pebble game's trade.

In deep learning this second option is called gradient checkpointing, and it is standard practice for fitting large models into limited GPU memory.

The √n sweet spot

How many checkpoints should you keep? Suppose the forward pass has n steps and you place c evenly spaced checkpoints. Then:

memory is O(c) — one stored state per checkpoint;
each segment between checkpoints has length n/c, and on the backward pass you recompute one such segment, so the extra compute is O(n/c).

The total cost c + n/c is minimised when the two terms are equal, c = n/c, i.e.

c = \sqrt{n}, \qquad \text{giving memory } O(\sqrt{n}) \text{ at about } 2\times \text{ the compute}.

This is the celebrated square-root rule: with one level of checkpointing you cut memory from O(n) to O(\sqrt{n}) for roughly a doubling of forward work. Nest the trick recursively and you reach Bennett's logarithmic regime — the adjoint-solver community calls the optimal recursive schedule REVOLVE (Griewank and Walther), and it long predates deep learning: PDE, climate and seismic adjoint methods used exactly these binomial checkpoint schedules decades earlier to differentiate through enormous time-stepped simulations.

Invertible networks: skip the storage entirely

There is a third option that is pure reversible computing: make the network invertible by design, so no forward states need storing at all — you recover them by running the layer backward. The trick is a coupling layer. Split the input into two halves x_1, x_2 and compute

y_1 = x_1 + f(x_2), \qquad y_2 = x_2 + g(y_1).

This is algebraically invertible no matter what the functions f and g are, because you can peel it apart in reverse order:

x_2 = y_2 - g(y_1), \qquad x_1 = y_1 - f(x_2).

Given the outputs you reconstruct the inputs exactly — a bijection built from arbitrary neural networks. RevNets and normalizing-flow models are built from stacks of these. Because every layer is reversible, the backward pass reconstructs each activation on the fly instead of reading it from memory: reversible computing has quietly become an architecture choice that trains huge models on modest hardware. Compare this with a standard network's training loop, which must cache every activation.

// A reversible coupling layer. f and g are arbitrary neural nets. function forward(x1: number[], x2: number[]): [number[], number[]] { const y1 = add(x1, f(x2)); const y2 = add(x2, g(y1)); return [y1, y2]; } // Its exact inverse — no stored activations needed. function inverse(y1: number[], y2: number[]): [number[], number[]] { const x2 = sub(y2, g(y1)); // undo the second line first const x1 = sub(y1, f(x2)); // then the first return [x1, x2]; }

The through-line

Everything in this lesson is the pebble game cashed out in engineering. The abstract question "how much scratch memory must I keep to be able to run this computation backward?" is answered every day by machine-learning frameworks, climate models and seismic imaging codes. Reverse-mode AD is the pebble game monetised.

reverse-mode AD (backprop) must traverse a computation backward, needing its forward intermediate states;
store them all → O(n) memory; or checkpoint and recompute → less memory, more compute — Bennett's trade;
one-level checkpointing is optimal at c=\sqrt{n} checkpoints: O(\sqrt n) memory, ~2× compute; recursive REVOLVE schedules reach O(\log n);
invertible coupling layers (y_1=x_1+f(x_2),\,y_2=x_2+g(y_1)) store nothing — the network is reversible by construction.

Long before anyone trained a neural network, computational scientists faced the identical problem in adjoint methods: to compute how a climate or seismic-wave simulation depends on its inputs, you must run the time-stepping backward, which needs the states from the forward run — over thousands of time steps that will not fit in memory. Griewank's REVOLVE algorithm (1992) gave the provably optimal checkpoint schedule, a binomial spacing that achieves O(\log n) memory. Deep learning rediscovered a corner of this decades later and called it gradient checkpointing. The pebble game is older, and wider, than the field that made it famous.

A reversible coupling layer stores no activations, which sounds like something for nothing. It is not. You have traded memory for compute: on the backward pass you must re-evaluate f and g to reconstruct the activations you didn't store. That is the pebble game all over again — you never escape the space–time trade, you only move along its curve. Invertible networks pick the extreme "store almost nothing, recompute as needed" end; naive backprop picks the "store everything" end; gradient checkpointing sits in between. Same curve, different operating point.