Making AI Training Replay Correct—and Knowing When It Saves Resources

Recomputing intermediate results can reduce the memory needed to train a neural network. But replaying the same code is not enough if the underlying state has changed. This study examines how to preserve the right history—and when doing so actually improves training efficiency.

Why replay needs a memory of its own

Neural-network training uses intermediate results, called activations, to calculate gradients: the quantities that guide parameter updates. Activation recomputation saves accelerator memory by discarding some activations and rebuilding them during the backward pass.

The complication is mutable state. A buffer may have been overwritten since the original calculation. Different tensor views may share the same physical bytes. A random-number generator may have advanced. Replaying unchanged instructions can therefore produce a different computation and different gradients.

In their October 2026 preprint, Annan Zhang and colleagues propose making these dependencies explicit. Their approach combines a replay contract describing the state that must be recovered with a planner deciding how to preserve it economically.

Preserving the right historical state

The contract identifies values by the writes that produced them, rather than by storage addresses alone. It requires replay to reproduce both the relevant storage bytes and random-generator state. Matching visible numerical outputs is not sufficient if shared storage or later random draws behave differently.

For supported framework execution, replay temporarily restores declared historical buffers and random state, reruns the segment, and then reinstates the state that should remain visible outside replay.

The planner chooses among retaining activations, saving whole-buffer snapshots, and preserving selected regions. Selective preservation can be cheaper when only a small fraction of a large buffer changes. However, it must cover later writes that could destroy needed history, not just writes made by the segment being replayed.

Each option is charged for retained memory, temporary workspace, execution time, and preservation and recovery work. The planner finds the best policy within its stated static model when its search frontier is untruncated; that is not a claim of optimality for every possible training schedule or allocator.

Checking correctness separately from speed

The authors report separate numerical and byte-level checks. In a 1,024-program derivative suite, maximum gradient disagreement with an independently organized forward-Jacobian reference was 8.88 × 10⁻¹⁶. A separate 1,024-program suite checked exact storage bytes and generator state.

Their PyTorch and JAX investigation retained all 5,900 execution records, including matches, framework rejections, and divergences. These records represent paired conditions, not 5,900 independent experiments.

Native PyTorch checkpoint variants matched the tested local-view conditions but diverged on hidden-buffer conditions; the declared-state transaction matched both families. JAX functional checkpointing matched all 500 functional conditions, while unsupported in-place and checkpointed impure-holder cases were rejected. The authors distinguish rejection from silent divergence and do not present stateful functions outside ordinary pure-function assumptions as evidence that frameworks violate their contracts.

Where selective preservation helped

The strongest planner comparison used a frozen, large-activation, sparse-write workload on H100 NVL hardware, with fifty fresh processes per cohort.

At a 44 GiB logical budget, joint planning reduced paired step time by 18.16% against a correctly costed planner restricted to activation retention and whole snapshots. The approximate 97.5% interval for their time ratio was 0.8172–0.8195. Measured peak allocation was 42.352 GiB for joint planning and 42.813 GiB for the restricted planner.

At 32 GiB, the restricted planner had no feasible policy within the static model. Joint planning was 4.98% faster than fixed region preservation, but used more memory: 29.936 rather than 26.200 GiB.

These budgets were logical constraints, not enforced hardware limits. Exceeding one was not an observed device out-of-memory event. In other tested workloads, eager activation retention was already faster and smaller, leaving no demonstrated planner advantage.

A complete application revealed the boundary

The CIFAR-100 dense-queue application ran seven methods for 100 encoder epochs and 25 probe epochs at six training seeds. Test-score files were byte-identical across methods within each seed, although this does not establish population-level quality equivalence.

Complete encoder-process measurements covered five seeds. Joint planning used 22.41% more allocated memory than whole snapshots. Its geometric-mean time ratio against whole snapshots was 1.0080: slower at four seeds and faster at one, without an application significance claim.

The conclusion is conditional. Selective preservation helped when writes were sparse and activations large; whole snapshots used less allocated memory in the dense application. Automatic discovery of arbitrary hidden state and aliases, concurrent external effects, and allocator-level optimality remain outside the demonstrated scope.