II · THE IDEA · ARTIFICIAL INTELLIGENCE
Gradient Checkpointing for Memory-Efficient Training
▶ Listen · narrated
Memory, not arithmetic, often limits how deep a network can be trained on a given device. One response is to store fewer activations and pay for that choice in extra forward computation.
At a glance
- Core idea
- Store sparse checkpoints; recompute the rest in the backward pass
- Typical cost
- About one extra forward pass for O(√n) activation memory
- Deeper saving
- Recursive checkpointing can reduce activation memory to O(log n)
- Gradients
- Unchanged: recomputed activations match the original forward values
- DenseNets
- Concatenated feature maps make naïve activation storage especially heavy
Think of writing an essay and needing to show your working later. The usual habit is to keep every draft page on the desk until the final review. Checkpointing keeps only a few marked pages — say the start of each chapter — and throws the rest away. When the review reaches a chapter whose pages are gone, you recreate that chapter from the marked start page, check it, and discard the extras again before moving on.
In a neural network the “pages” are activations: the intermediate results each layer writes during the forward pass. Backpropagation needs them to compute gradients. Saving all of them uses a lot of memory. Gradient checkpointing saves only a sparse set, frees the rest, and recomputes missing stretches just in time for the backward step. The gradients come out the same; the device holds fewer tensors at once; the cost is extra forward arithmetic. Chen et al. showed that with roughly √n checkpoints along a depth-n chain, activation memory along depth drops to O(√n) for about one extra forward pass. Recursion can push that memory toward O(log n). DenseNets, which concatenate many feature maps, feel the memory pressure especially hard, and memory-efficient implementations use the same rebuild-rather-than-store idea.
Standard reverse-mode autodiff on a depth-n feed-forward chain retains O(n) activations (plus whatever the operator workspace needs) so that each local Jacobian-vector product in the backward pass has its forward inputs. Chen et al. (2016) replace that policy with explicit checkpoints. Partition the chain into segments; persist only segment inputs after the forward pass; free interior activations. On the backward pass, for each segment in reverse order: reload the checkpoint, re-execute the segment forward to materialise interiors, run the local backward, then free the temporary activations.
For balanced segments of length about √n, peak retained activations along depth are O(√n) and each layer is executed twice per step (once in the global forward, once on recompute), i.e. roughly one extra forward of compute. Recursively checkpointing inside segments yields O(log n) activation memory in depth at higher recompute cost. The gradient is algebraically unchanged if recomputation is bit-honest with respect to the original forward.
DenseNets stress a related pattern: layer i receives concat of prior maps in the block, so a naïve tape holds many large tensors simultaneously. Pleiss et al. (2017) give a memory-efficient DenseNet implementation that pre-allocates and shares storage and avoids retaining intermediate concatenations that can be regenerated, so depth and growth rate are less often memory-bound.
Limitations: the classic bounds address depth-varying activation storage on chain-like graphs, not parameter memory, optimiser state, or irregular autograd graphs; checkpoint placement on residual and multi-path nets needs care; recompute increases wall time and energy; mixed precision and fused kernels change constants but not the underlying trade.
Look closer
What is kept, and what is freed
In ordinary backpropagation the forward pass retains every intermediate activation until the matching backward step consumes it. Checkpointing keeps only a sparse set of those tensors — the checkpoints at chosen boundaries — and frees the tensors inside each segment. When the backward pass reaches a segment whose interiors were discarded, it restores the segment input from the checkpoint, reruns that stretch of the forward computation, and only then applies the local backward step.
The √n segment choice
Chen and colleagues analysed a simple policy: divide a depth-n network into roughly √n segments and retain only the inputs to those segments. Peak activation memory along depth then scales as O(√n). Each discarded segment is recomputed once during the backward sweep, so the computational overhead is on the order of one additional forward pass over the network — a deliberate, bounded trade rather than an open-ended slowdown.
Dense connectivity under pressure
A DenseNet layer receives the concatenation of feature maps from preceding layers in its block. A naïve implementation therefore holds many large tensors at once, and memory fails before arithmetic does. Pleiss and colleagues gave a memory-efficient implementation that shares buffers and avoids retaining intermediate concatenations that can be rebuilt, applying the same discard-and-recompute instinct to an architecture whose connectivity makes storage especially expensive.
The story
Deep network training carries two distinct memory burdens. Parameters and their gradients occupy a relatively stable amount of space. Activations — the intermediate tensors written layer by layer during the forward pass — grow with depth, width and batch size, and in the usual algorithm they must all be retained until backpropagation has used them. For a long chain of layers, that second burden dominates.
Gradient checkpointing, set out systematically by Chen and colleagues in 2016, refuses to keep most of those activations. The network is treated as a sequence of segments. Only the tensors at segment boundaries are saved after the forward pass; interiors are freed. When backpropagation arrives at a segment, the implementation restores the segment’s input checkpoint, runs the forward computation again to materialise the missing activations, performs the local backward step, and frees the temporary tensors before moving on. The mathematical gradient does not change. The recomputed values are the same ones the original forward pass produced, so the backward pass sees the same inputs and writes the same parameter updates. What changes is only the schedule of allocation, deallocation and recomputation.
How the segments are chosen sets the memory–compute trade. A simple, analysed policy divides a depth-n chain into about √n equal segments and stores only the segment inputs. Peak activation memory along the depth axis then scales as O(√n) rather than O(n). Because each segment is recomputed once on the way back, the extra arithmetic is roughly one additional forward pass — a predictable cost, not an unbounded one. A recursive version of the same idea nests checkpoints inside checkpoints and can bring activation memory down to O(log n) in depth, at the price of still more recomputation.
The same pressure appears in a sharper form in densely connected architectures. In a DenseNet, each layer is fed the concatenation of earlier feature maps in the block. A straightforward implementation therefore retains a large collection of maps at once, and memory is exhausted long before the floating-point units are. Work by Pleiss and colleagues in 2017 showed how a memory-efficient DenseNet implementation can share storage and drop intermediate concatenations that are cheap to rebuild, so that depth and growth rate are no longer dictated by how many feature maps happen to fit in device memory.
Checkpointing is not free, and it is not a substitute for every other form of thrift. The recomputation costs wall-clock time and energy. The optimal placement of checkpoints depends on the graph: residual branches, multi-scale paths and irregular connectivity all complicate a scheme that is simplest on a pure chain. Still, the central bargain remains clear. Training need not store every activation it will later need; it can store a sparse scaffold and rebuild the rest on demand. Memory is traded for compute in a controlled way, and the gradient at the end is the gradient the full-storage algorithm would have produced.
Why it mattered then
By the mid-2010s, depth had become a principal lever on accuracy, yet device memory had not grown at the same pace as the graphs people wanted to train. Full activation storage meant that many architectures were limited by what a single GPU could hold rather than by optimisation theory or dataset size. Chen and colleagues gave a general algorithm and concrete complexity bounds — O(√n) memory for about one extra forward, O(log n) with recursion — so that memory cost along depth no longer had to be linear. Shortly afterwards, DenseNets exposed a related failure mode: concatenation-heavy connectivity made naïve storage especially wasteful. Memory-efficient implementations showed that the same discard-and-recompute attitude could unlock architectures that would otherwise have been impractical on then-current hardware. The technique mattered because it widened the set of train-able models without waiting for larger chips.
Why it matters now
The same trade still governs large-scale training. Activation memory grows with sequence length, batch size, width and depth; parameters alone do not explain the footprint of a modern training step. Checkpointing — often under names such as activation recomputation — remains a standard tool when fitting a model into a fixed device budget, whether on a single accelerator or across a cluster where memory pressure and communication interact. The original analyses still supply the right mental model: choose what to keep, accept a bounded amount of extra forward work, and leave the gradient mathematically unchanged. Anyone who has reduced batch size only to lose utilisation, or who has watched a training job die with an out-of-memory error deep in the stack, is looking at the problem these papers named and quantified.
The surprising detail
The recursive form of the algorithm does not merely shave a constant off memory use. By nesting checkpoints, it can reduce the activation memory that scales with depth to O(log n), so that a chain twice as deep need not double that part of the footprint. The cost is additional recomputation at each level of the recursion — a pure exchange of arithmetic for storage that the 2016 analysis makes explicit rather than leaving as folklore.
What is disputed
The O(√n) and O(log n) figures describe activation memory that scales with depth under the analysed segment policies for chain-like graphs. They are not a full account of training memory: parameters, optimiser state, workspace buffers and non-chain connectivity all sit outside those bounds. Reported overheads of “about one extra forward” likewise assume the simple segment scheme rather than every production variant.
Remember this
Checkpointing stores a sparse scaffold of activations and recomputes the rest during backpropagation, trading a bounded amount of extra forward work for sublinear activation memory in depth.
Test yourself
You train a plain feed-forward chain of depth n with full activation storage and hit the memory limit. You switch to the simple √n-segment checkpointing scheme from Chen et al. What should you expect to happen to peak activation memory along depth, to the gradient values, and to the amount of forward computation per training step?
Peak activation memory along depth should fall from O(n) to O(√n), because only segment-boundary tensors are retained as checkpoints and interiors are freed. The gradients should match those of the full-storage run, since recomputed activations are the same values the original forward pass produced. Forward computation roughly doubles: each segment is run once on the way up and recomputed once on the way down, on the order of one extra forward pass over the network.
Go deeper
- [1604.06174] Training Deep Nets with Sublinear Memory Cost · arxiv.org
- [1707.06990] Memory-Efficient Implementation of DenseNets · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.