A trainable trajectory still fails if its history cannot fit
Paper 0044 showed why deep optimization depends on memory in the algorithmic sense. Momentum carries directional history so that persistent gradients accelerate and oscillating directions damp. Paper 0045 addresses a different memory: the physical accelerator storage occupied by intermediate activations during forward and backward computation. The network may have an excellent initialization and optimizer, yet still be impossible to train because the values required by reverse-mode differentiation exceed GPU capacity.
Backpropagation first evaluates a chain or graph of operations. Later, in reverse order, each gradient operator needs some of the values produced earlier. A straightforward system retains every such activation from the forward pass until its final backward use. For an n-layer chain with similarly sized activations, that is O(n) feature-map memory. Deeper residual networks and longer recurrent unrollings increase the stored history even when parameter count changes little.
Chen and colleagues challenge an implicit assumption: a backward operator needs the correct activation, but that activation need not have survived continuously in memory. If a retained checkpoint can regenerate it through deterministic forward operations, storage may be released early and computation repeated later. Training computes the same mathematical gradient while following a different physical schedule.
The paper combines three levels. Computation-graph liveness identifies in-place and reusable buffers. A symbolic gradient-graph transformation mirrors operations whose outputs will be recomputed. A planner selects boundaries under a memory budget, balancing checkpoint storage against the size and cost of replayed segments. The contribution is not merely the slogan “trade compute for memory”; it is a system for expressing, scheduling and measuring that trade.
Causal probe 1 · Why can recomputation preserve the gradient?
Reverse-mode differentiation depends on the numerical forward values supplied to each local derivative. If replay produces those values exactly and in the same dependency context, the local derivatives and their accumulated parameter gradients are unchanged. Only the time at which the values occupy memory differs.
A byte may be reused only after every dependent operation has finished
A computation graph represents operations as nodes and data dependencies as directed edges. The gradient calculation can itself be made explicit as another graph created by reverse topological traversal. Once forward and backward operations share one dependency structure, memory allocation becomes a lifetime problem. Each node output is born when its operation runs and dies after its final consumer finishes.
An in-place operation overwrites an input buffer with its output. This is safe only when no other pending operation needs the original input. Memory sharing assigns the same buffer to different values whose live intervals do not overlap. These rules can remove a large amount of unnecessary allocation without changing arithmetic.
Finding a globally minimal allocation through a conflict graph and graph colouring can require O(n²) work. The paper instead uses a linear-time heuristic. It walks the graph in topological order, counts unfinished dependants and attaches reusable tags to buffers when their values die. A recycled tag tells a later node that it may occupy the same physical storage.
Dependency declarations control what the allocator can prove. If a backward operator is declared to depend on a broad forward output when it really needs only a smaller input, the output’s lifetime is extended and sharing opportunities disappear. The paper therefore advises frameworks to expose minimum dependency requirements. Memory efficiency begins with semantic precision.
Worked example 1 · One buffer, two non-overlapping values
Suppose activation B is last read by operation D at time 7, while activation E is created at time 8. B and E can share a 40 MB buffer. If another branch reads B at time 10, the intervals overlap and reuse would corrupt that branch. A scheduler needs dependency evidence, not visual proximity in source code.
Causal probe 2 · Why does liveness still leave O(n) memory?
In a long training chain, many forward activations remain simultaneously live because their backward consumers have not begun. Sharing removes buffers after death, but it cannot declare a value dead while a future gradient still depends on it. Recomputing the value is what changes that dependency on continuous storage.
Keep the boundaries, replay the interior
Consider an n-layer feed-forward chain divided into k contiguous segments. During the original forward pass, the algorithm stores only each segment’s output checkpoint and drops intermediate activations inside the segment. Backward processing starts from the final segment. Its input checkpoint is loaded, its internal forward values are regenerated into temporary memory, and local backward operators consume them in reverse. The temporary values are released before the preceding segment is processed.
Two forms of storage coexist. The k checkpoints must remain available across much of the run. The currently replayed segment needs memory proportional to its length. If layers are uniform and segments are balanced, peak activation memory is O(k)+O(n/k). Setting k=√n makes the two terms equal, giving O(2√n), conventionally written O(√n).
Each discarded internal activation is recomputed once in the one-level scheme. Across all segments this is almost one additional forward pass. Backward computation is already more expensive than forward computation for many neural operations, so doubling forward work need not double total step time. The paper later measures about 30% additional runtime on its workloads.
The checkpoint plan changes storage and arithmetic order, not model semantics. Parameter gradients are still accumulated from the same activation values. Yet “same” requires care: random operations, mutable state, nondeterministic kernels or side effects must be replayed consistently. The primary’s graphs assume reproducible operators and do not offer a general random-state protocol.
Worked example 2 · Sixty-four layers with eight segments
Full retention stores roughly 65 boundary-and-activation values including the input. Eight equal segments store nine boundaries and need at most seven additional interior values during replay: about 16 live values, versus 65. Fifty-six discarded interiors are recomputed once. Constants differ by convention, but the square-root balance is visible.
Causal probe 3 · Why not use one giant segment?
One checkpoint minimizes persistent storage but replaying the whole network requires all of its internal values simultaneously for that segment’s backward pass. The temporary term becomes O(n). Too many segments make checkpoint storage O(n). The minimum lies between the extremes.
Mirror an operation when its output is cheaper to reproduce than retain
A chain-specific training loop is easy to describe but difficult to integrate into a general deep-learning framework. Real graphs branch, merge and mix large convolutions with tiny elementwise operators. The paper’s Algorithm 2 embeds recomputation in the symbolic gradient graph itself.
Each forward node v receives a mirror count m(v). Zero means its output remains available for ordinary backpropagation. A positive value tells the constructor to create one or more duplicate nodes with the same operator and mirrored predecessors. When the reverse traversal constructs the gradient for v, it references the appropriate mirrored output and inputs. The enlarged graph therefore represents recomputation explicitly.
This representation carries two advantages. First, existing dependency analysis and allocation can see the replayed work rather than receiving it as opaque control flow. Second, the topological order of the transformed graph becomes a logical execution plan. Frameworks that allocate at runtime can add control dependencies so memory is released and recreated in the intended sequence.
Mirror counts may exceed one. A node can be regenerated inside a replay that is itself regenerated, allowing recursive checkpointing. That generality also expands the search space. The paper does not claim that any mirror pattern is good. It separates correctness: construct a valid gradient graph: from planning: choose a pattern with an acceptable peak-memory and compute cost.
Worked example 3 · Mirroring a cheap activation
A convolution produces C, batch normalization produces B and ReLU produces R. Backward for ReLU and batch normalization needs their forward context, but recomputing B and R is much cheaper than repeating the convolution. Retain C, set positive mirror counts for batchnorm and ReLU, and regenerate them from C just before their backward nodes.
Causal probe 4 · Why mirror the graph rather than hide replay in an operator?
Visible nodes expose dependencies, lifetimes and costs to the allocator and scheduler. Hidden replay may be correct locally but prevents global memory planning, higher-order graph transformations and explicit execution ordering.
Equal layer counts are a derivation; real plans balance bytes and time
The square-root result assumes roughly equal activation sizes and operation costs. A convolutional stage can produce a large feature map and be expensive to rerun, while a nearby activation produces the same shape cheaply. Recurrent graphs contain many fine-grained operations and repeated gradient additions. Dividing only by node count can place a checkpoint on the wrong side of both memory and compute.
One quick policy is to drop outputs of low-cost operations. In a convolution–batchnorm–activation pipeline, keep the convolution result but reproduce batch normalization, activation and pooling. The paper reports that this “drop bn, relu” strategy improves the ResNet allocation before the full sublinear planner is applied.
Algorithm 3 accepts candidate stage boundaries and an approximate memory budget B. As it walks the graph, it accumulates output sizes. A candidate crossing the budget becomes a retained stage boundary; nodes inside become mirrors. The algorithm estimates inter-stage checkpoint cost x and the largest sub-stage cost y. A static allocator then evaluates the exact transformed graph, and the appendix grid-searches B around √(xy) to find a good plan.
This is heuristic optimisation over a precise execution model. The budget does not magically predict runtime or fragmentation. It proposes a mirror plan; exact allocation and measurement decide whether that plan is good. The separation between approximation and exact validation is an important systems discipline.
Worked example 4 · Same depth, different plan
Two eight-operation segments each contain one 200 MB activation. In the first, reproducing that activation costs a tiny elementwise operation; in the second, it costs a large convolution. A count-balanced plan treats them equally. A cost-aware plan retains the second activation and spends replay on the first.
Causal probe 5 · Why is a memory-only objective incomplete?
Two plans with the same peak bytes can differ greatly in repeated FLOPs, kernel launches, communication and critical-path delay. A feasible plan must fit; a useful plan must also price the work needed to recreate what it discards.
Replay can itself contain checkpoints
A segment can be treated as one bulk operator whose internal graph is hidden from the outer plan. When its backward pass runs, the segment’s internal graph may use the same checkpointing strategy again. This subroutine view creates a hierarchy: outer checkpoints delimit large regions; inner checkpoints delimit the replay inside one region.
The paper writes the memory recurrence g(n)=k+g(n/(k+1)). Solving it gives k logk+1(n). At k=1, memory becomes O(log n). The price is repeated regeneration across levels, leading to O(n log n) forward computation. Nothing about the recursion is free; it moves the plan along a memory–compute frontier.
The logarithmic result is conceptually important because it shows that linear activation storage is not inherent to exact backpropagation through a chain. It is not the paper’s recommended operating point for every workload. The square-root schedule offers a practical one-extra-forward-pass compromise, while deeper recursion is an extreme for severe capacity constraints.
Checkpointing has a longer automatic-differentiation lineage. The paper cites Griewank and Walther’s REVOLVE and thanks David Warde-Farley for pointing out the connection. Later neural systems formulate optimal rematerialization, heterogeneous-chain dynamic programming, online eviction and reversible architectures. Paper 0045 is a bridge between that lineage and graph-level deep-learning frameworks.
Worked example 5 · Thirty-two layers as nested halves
Store the input and midpoint of a 32-layer chain. To backpropagate the second half, split its 16 layers again, then split the active eight, and so on. Only a logarithmic stack of boundaries is retained, but early subpaths are replayed at multiple recursion levels. Capacity falls because work rises.
Causal probe 6 · What does recursion prove?
It proves a possible asymptotic trade-off for chain-like feature maps. It does not prove that logarithmic memory minimizes wall-clock time, total bytes or energy on a heterogeneous neural system.
A thousand layers fit where two hundred exhausted ordinary planning
The convolutional experiment uses residual networks for ImageNet-scale inputs in MXNet. Batch size is 32 and input shape is 3×224×224. The authors vary depth by expanding residual stages and count one convolution–batchnorm–ReLU group as one layer. A static planner estimates feature-map memory, while nvidia-smi records total runtime GPU memory on a Titan X.
In-place and sharing optimizations reduce memory by a factor of roughly two to three, but the trend remains linear. Dropping batchnorm and ReLU outputs helps further. Even with those constant-factor improvements, the paper says the available GPU could train only about a 200-layer ResNet. The sublinear plan instead stays on a visibly shallower curve and trains a 1,000-layer configuration in less than 7 GB. The abstract describes a reduction from about 48 GB to 7 GB.
Figure 5 separates estimated feature-map cost from total GPU cost. That distinction prevents the asymptotic claim from swallowing parameters, libraries and workspace allocations. Missing ordinary-plan points correspond to out-of-memory failures. The sublinear curve is not just an estimate: the runtime panel confirms that the resulting configurations execute.
The experiment establishes capacity expansion, not better classification. The graph transformations preserve the weight gradient, so the paper focuses on memory and time rather than accuracy. A larger trainable architecture might enable new modeling research, but whether a 1,000-layer model predicts better is a separate question.
Worked example 6 · Constant-factor savings do not alter slope
If optimized sharing turns 48 GB into 20 GB, the improvement is valuable but a tenfold increase in depth still pushes memory towards 200 GB under linear growth. A square-root trend grows by about √10 instead. As models scale, changing the exponent can dominate a large constant improvement.
Unrolled time is another deep chain
The recurrent experiment unrolls a four-layer LSTM with 1,024 hidden states. Batch size is 64; each timestamp receives a continuous 50-dimensional vector and produces a 5,000-class softmax. Longer unrolling can expose longer dependencies, but it also retains more time-indexed activations for backpropagation through time.
Fine-grained graph optimization already helps substantially. In-place accumulation can add each timestamp’s weight gradient into one shared cell instead of allocating a separate full gradient at every time. The sublinear plan then reduces memory by more than four times relative to the optimized plan that does not trade computation for storage.
Runtime is measured over 20 batches on one Titan X. The sublinear plan adds roughly 30% relative to the linear-memory strategy for both ResNet and LSTM scaling. The authors explain that backward work is more expensive than forward work, so one extra forward pass is a smaller fraction of total training than “double forward” might suggest. Both strategies retain a roughly linear time trend as workload grows.
The 30% is not an algorithmic constant. Kernel balance, memory bandwidth, compiler fusion, graph capture, recomputation choice and hardware have changed. The stable claim is structural: recomputation creates a tunable capacity–time exchange, and measured end-to-end overhead can be materially smaller than the raw amount of repeated forward arithmetic.
Worked example 7 · Depth can mean time
An LSTM unrolled for 1,024 steps may reuse the same parameters at every step, so parameter memory is unchanged. Activation memory nevertheless grows with the unrolled computation graph. Checkpointing acts on that temporal depth even though the model’s stored weights are constant.
Causal probe 7 · Why may replay cost less than offload?
Replay consumes arithmetic on the accelerator; offload moves bytes across a slower link and later brings them back. Which wins depends on operator cost and bandwidth. The paper emphasizes that its method adds no PCIe communication and can be combined with offload or parallelism.
Checkpoint laboratory: retain, replay and prove the gradient
The laboratory simulates an exact differentiable chain. Each profile chooses activation sizes and checkpoint boundaries. Full retention and scheduled rematerialization compute the same parameter gradients independently. The chart is the measured live-byte trace, not a stylized complexity curve.
Count what was saved, what was repeated and what was left outside
The primary’s strongest claim combines an asymptotic construction with measured systems evidence. The chain derivation establishes O(√n) feature-map memory under uniform assumptions. The general graph transformation makes recomputation representable in a framework. Static allocation predicts exact feature-map bytes for a chosen plan. Runtime measurements show that those plans execute and remain under physical capacity.
Each layer of evidence answers a different question. The equation says how a balanced chain scales. Algorithm 2 says how to construct a correct mirrored graph. Algorithm 3 says how to search a heterogeneous plan approximately. Figures 5 and 6 say how memory grew on the tested ResNet and LSTM graphs. Figure 7 says what replay cost on one Titan X implementation.
None can replace the others. A clean asymptotic can hide large workspaces. A valid graph can be badly scheduled. A static allocation can miss runtime fragmentation. A single hardware result can obscure general structure. The paper is persuasive because it moves across those levels while distinguishing feature-map estimates from total device memory.
The companion therefore avoids the loose statement “training now needs O(√n) memory”. The bounded statement is: one-level activation checkpointing can train an n-layer uniform chain using O(√n) intermediate feature-map storage, with one additional forward pass, while the paper’s general planner achieved sublinear trends and particular measured savings on its 2016 workloads.
Worked example 8 · Total memory can stop following the activation bound
If checkpoints reduce activations from 40 GB to 5 GB but parameters, optimizer states and workspaces occupy 9 GB, total memory falls from 49 GB to 14 GB, not to 5 GB. Further activation savings eventually meet a floor set by other state.
Evidence boundary · Exact gradients require exact replay conditions
For deterministic pure operators, rematerialization changes schedule rather than mathematics. With dropout, stateful normalization, random sampling, nondeterministic kernels or external side effects, a system must preserve or reconstruct the relevant state. The primary does not supply a universal solution for those cases.
Rematerialization is one axis, not the entire plan
The paper explicitly places CPU/GPU swapping and model parallelism beside its method. Offload retains values elsewhere and pays transfer bandwidth. Model parallelism partitions parameters and computation across devices and pays communication. Checkpointing frees values and pays repeated arithmetic. Because the resources differ, the techniques can be combined.
Later reversible networks reconstruct earlier activations by inverting architectural transformations. That can avoid both storage and generic forward replay for much of the network, but it constrains model design and numerical behavior. Activation compression stores a lossy or compact representation and may change gradient statistics. Tensor rematerialization planners such as Checkmate formulate broader graph schedules with profiled costs and memory constraints.
A responsible system asks which resource is scarce and which substitute is abundant. Compute-rich but memory-limited accelerators favour replay. Fast host links may favour offload. Multiple devices may favour sharding. Pipeline parallelism introduces in-flight activations whose lifetimes interact with checkpoint placement. There is no memory optimization independent of execution context.
Causal probe 8 · Why can two memory techniques compose?
They transform different state or move it through different resources. Checkpoint activations inside each partition, shard parameters across devices and offload the remaining long-lived values: one decision need not invalidate the others if dependencies and schedules are coordinated.
Four images to carry away
1 · The waiting room
A forward activation occupies memory because a future backward node still depends on it. Lifetime, not creation alone, determines cost.
2 · Red boundary stones
Checkpoints remain visible across the journey. Interior values disappear and are walked again only when their segment’s reverse work arrives.
3 · The U-shaped budget
Too few segments create a large temporary replay; too many create a large checkpoint set. Square-root balance sits between them.
4 · The mirrored workshop
Recomputation is explicit graph work. A duplicate operator rebuilds a value, feeds its gradient consumer and then releases the temporary storage.
Optimization becomes executable resource choreography
Earlier canon papers progressively expanded what a learning rule must manage. Robbins–Monro and Kiefer–Wolfowitz organized noisy observations over time. Momentum carried direction across updates. Adam and AdamW stored coordinate-wise optimizer state. XGBoost aligned sufficient statistics with caches and disks. Paper 0045 adds a second schedule around the learning schedule: the physical production, retention, destruction and reproduction of intermediate values.
This distinction matters for agency. An algorithm may specify what mathematical action should occur next, yet the system must make that action possible under bounded resources. Checkpointing is a formal compromise between memory of the past and capacity for further computation. It deletes information only because a procedure and an anchor make the information recoverable.
The analogy to cognitive memory should remain limited. Neural activations are exact numerical intermediates, and rematerialization is deterministic computation. Human forgetting and reconstruction involve uncertainty, transformation and meaning. The useful abstraction is narrower: stored state, reconstructible state and irrecoverable state impose different future options.
As the canon moves from optimization into neural representation, this paper supplies a systems prerequisite. Deeper models, longer sequences and richer internal computation depend not only on gradients and architecture but on whether the execution engine can schedule their histories. Representation power is constrained by resource choreography.
Balance persistent checkpoints against temporary replay
Step 1 · Divide the chain
Split n equal layers into k segments. Store every segment boundary during the first forward pass.
Step 2 · Count persistent storage
There are O(k) retained boundaries. Endpoint conventions change k to k+1 but not the asymptotic.
Step 3 · Count temporary storage
Only one segment is reconstructed for backward at a time. Balanced segments contain O(n/k) internal activations.
Step 4 · Minimize the sum
M(k)=k+n/k. Treating k as continuous, M′(k)=1−n/k². The stationary point is k=√n, and M=2√n.
Step 5 · Price computation
Every discarded interior is replayed once, so one-level checkpointing adds at most one forward traversal. Recursive checkpointing reduces storage further by replaying some operations more than once.
How to audit a training-memory claim
- Name the memory class. Separate activations, gradients, parameters, optimizer state, workspaces, communication buffers and allocator reserve.
- Expose lifetimes. Record which operation produces each tensor, its final consumer and any aliasing or in-place constraints.
- Protect replay semantics. Preserve random state, mutable buffers, side effects and kernel determinism where gradients require exact values.
- Measure both plan and runtime. Compare static live-byte estimates with actual peak device allocation and fragmentation.
- Price the critical path. Count recomputed operators, not only layers, and profile overlap with communication or idle periods.
- Match the workload. Include batch, sequence length, tensor shapes, precision, framework, compiler and hardware.
- Verify mathematical equivalence. Compare losses and gradients against full retention under controlled seeds before using scale results.
What this companion adds
This page rebuilds activation checkpointing from dependency semantics rather than beginning with a big-O slogan. Ten paper-specific animated figures connect forward lifetimes, liveness reuse, segment replay, the square-root balance, mirrored graphs, heterogeneous planning, recursion, ResNet depth, LSTM unrolling and the wider memory toolbox.
The executable laboratory uses real activation values and independent backward computations. It reports checkpoint bytes, recomputed nodes and the exact live-memory trace for uniform, square-root, heterogeneous and full-retention profiles. Four self-tests require gradient equivalence, correct boundary replay, exact peak accounting and deterministic finite execution.
Two independent passes
Scholarly and editorial review: pass. The 12-page arXiv v2 primary was rendered and inspected in full. Figures 1–7, Algorithms 1–3, equations 1–3 and Appendix A were checked against extracted text and original-resolution pages. Feature-map bounds remain separate from total memory, the one-extra-pass statement stays attached to one-level checkpointing, and 2016 hardware evidence is not generalized into a modern performance ranking.
Executable and publication review: pass. All four laboratory profiles reproduce full-retention gradients and pass four exact audits. Responsive rendering is checked at 1440, 768 and 390 pixels; SVG titles and descriptions, reduced-motion behavior, labelled controls and live output are present. Placeholder, AI-telltale and forbidden-term scans are clean.
References and source trail
- Tianqi Chen, Bing Xu, Chiyuan Zhang and Carlos Guestrin, “Training Deep Nets with Sublinear Memory Cost,” arXiv:1604.06174v2 (2016). Canonical record and PDF.
- Andreas Griewank and Andrea Walther, “Algorithm 799: Revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation,” ACM TOMS 26(1), 19–45 (2000). DOI record.
- Minsoo Rhu et al., “vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design,” arXiv:1602.08124 (2016). arXiv record.
- Kaiming He et al., “Identity Mappings in Deep Residual Networks,” arXiv:1603.05027 (2016). arXiv record.
- Sepp Hochreiter and Jürgen Schmidhuber, “Long Short-Term Memory,” Neural Computation 9(8), 1735–1780 (1997). DOI record.
- Paras Jain et al., “Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization,” arXiv:1910.02653. Later planner.
- Olivier Beaumont et al., “Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory,” arXiv:1911.13214. Later heterogeneous analysis.
- Marisa Kirisame et al., “Dynamic Tensor Rematerialization,” arXiv:2006.09616. Later online planner.
- Aidan N. Gomez et al., “The Reversible Residual Network: Backpropagation Without Storing Activations,” arXiv:1707.04585. Reversible architecture context.