Activation Checkpointing

Backpropagation needs the activations the forward pass produced, so the obvious implementation stores all of them — which for a deep model is far more memory than the weights themselves. Checkpointing keeps only a few and recomputes the rest on the way back. It buys memory with time, at a rate good enough that almost every large training run takes the deal.

Store the input to every k-th layer and rerun the segment between checkpoints during the backward pass. At k near the square root of the depth, activation memory falls to about the square root of what it was, for one extra forward pass per step. The alternative is not training without recomputation: it is a smaller model or a smaller batch.

For n layers, storing everything costs O(n) memory and no recomputation; storing every k-th costs O(n/k + k), minimised at k = √n for O(√n). The extra compute is one forward pass per step, so a constant factor of time buys an asymptotic factor of memory — which is why the trade is nearly always worth taking.