II · THE IDEA · ARTIFICIAL INTELLIGENCE
Memory-Efficient Backpropagation with Reversible Layers
▶ Listen · narrated
Training deep networks usually means keeping every activation until the backward pass finishes. Reversible layers remove that obligation by making each block mathematically invertible.
At a glance
- What it is
- Backpropagation that recalculates activations instead of storing them
- Core idea
- Design residual blocks so each step is invertible
- Trade-off
- Extra forward-style compute for lower activation memory
- Source paper
- The Reversible Residual Network (arXiv 1707.04585)
Think of a long chain of locked boxes, each built so that knowing what came out tells you exactly what went in. While you walk forward you can throw away the notes about each box’s insides. On the way back, when you need those insides to work out how to adjust the box, you unlock it in reverse and recreate the notes, then throw them away again.
That is the plain idea behind reversible residual layers. Ordinary training keeps a full set of notes — the activations — because the backward pass needs them. If each block is invertible, the notes are optional: you recompute them when gradients arrive. You spend extra arithmetic to avoid carrying a growing pile of tensors through the whole depth of the network.
In standard reverse-mode differentiation through a residual stack, the forward pass materialises activations that the backward pass consumes when applying the vector–Jacobian products for each block. Peak memory therefore scales with depth times the size of those tensors (modulo any overlapping or streaming tricks).
A reversible residual block is parameterised so an inverse exists in closed form from the block outputs back to its inputs, typically via additive coupling across feature partitions: one half is updated by a residual function of the other, then roles swap, yielding a composition that unrolls by subtraction in reverse order. During training the implementation may discard the block’s intermediate activations after the forward evaluation, retain the block outputs (or a minimal boundary state), and on the backward pass invoke the inverse to reconstruct the forward activations before running the usual local backward kernels.
The gradient computation itself remains ordinary backpropagation on the reconstructed tensors; what changes is the storage policy. The dominant extra cost is recomputation of the residual functions during the backward pass, so wall-clock time and energy per step rise while activation memory across the reversible span falls. The technique is complementary to, rather than identical with, generic gradient checkpointing: checkpointing rematerialises from saved segment boundaries in arbitrary graphs, whereas reversibility builds rematerialisation into the block’s algebra and can avoid storing even those boundaries inside the reversible region. Limitations follow directly: the architecture must actually be invertible as implemented (including numerics), non-reversible stems and heads still need ordinary storage or checkpoints, and the compute–memory trade only pays when memory is the binding constraint.
Look closer
Storage is optional when the block can be undone
In ordinary residual training, the backward pass needs the activations produced on the way forward, so those tensors are retained until gradients finish. A reversible residual block is built so its inputs can be recovered from its outputs. Once that holds, the forward pass need not keep the intermediates: the backward pass reconstructs them when required, then discards them again.
The residual path makes invertibility practical
A plain stack of irreversible transforms cannot be run backwards without stored state. The reversible residual construction couples additive updates so each stage can be peeled off in reverse order. Invertibility is therefore an architectural property of the block, not an approximation applied after training.
Memory moves; arithmetic does not vanish
What is saved is activation storage across depth. What is spent is recomputation while gradients flow. The method shifts the bottleneck from memory capacity toward compute budget. Whether that shift helps depends on hardware and network shape, and cannot be read off from the idea alone.
The story
Standard backpropagation is exact about gradients and greedy about memory. During the forward pass each layer writes activations that the matching backward step will later combine with local derivatives. In a deep residual stack those tensors accumulate quickly. Depth, batch size and feature width all push the same way: a training step can become limited by how many activations fit, not only by how fast the hardware multiplies.
The reversible residual network attacks that pressure at the level of block design. If a block is invertible, its input is a deterministic function of its output. The training loop can therefore release the block’s internal activations after the forward evaluation, retain only what is needed to continue the chain, and, when the backward pass reaches that block, rebuild the missing activations by applying the inverse. Gradients then proceed with respect to the reconstructed values. The aim remains exact backpropagation for that architecture; the change is where the activations come from — recalculation rather than a retained stash.
Residual structure is what makes the scheme workable in practice. Additive coupling lets one partition of the features be updated from the other in an order that reverses by subtraction in the opposite order. The block can stay expressive enough to serve as a residual unit, yet remain algebraically undoable without storing the pre-update tensors. Nothing in that description fixes a particular width, depth or dataset; it is a constraint on how the unit is wired.
The cost appears on the clock, not in the gradient formula. Each reversed block pays for another evaluation of its internal functions during the backward pass. Training therefore trades activation memory for extra compute. On machines where memory is the binding constraint, that trade can unlock larger batches or deeper stacks. On machines where compute is already saturated and memory is plentiful, the same trade may simply slow the step.
The editorial point is architectural rather than mystical. One does not discover invertibility after the fact in an arbitrary network. One designs blocks that are invertible, accepts the recomputation, and stops carrying a full activation tape for those blocks. Backpropagation without storing activations, in the sense of the reversible residual network, is that design choice carried through the training loop.
Why it mattered then
As residual networks grew deeper, activation storage became a first-order training constraint rather than a bookkeeping detail. Methods that reduced that footprint without abandoning exact gradients mattered because they attacked the limit that batch size and depth hit first on a given accelerator. Reversible residual blocks offered a structural answer: change the block so the tape is optional, and spend compute to rebuild what used to be retained. In its own moment that was a way to keep scaling depth and training throughput inside fixed memory budgets.
Why it matters now
Memory pressure has not left training. Large models, long sequences and wide activations still force choices between batch size, model size and checkpointing strategy. Reversible and invertible designs remain one branch of that toolkit, alongside gradient checkpointing and other recomputation schemes. The underlying lesson still applies directly: if a unit can be inverted, activation memory becomes a negotiable cost rather than a fixed tax on depth, and engineers can decide where on the compute–memory curve a given run should sit.
The surprising detail
The method does not compress activations or approximate them. It refuses to keep them at all for the reversible spans, then regenerates the exact tensors needed for the backward step by running the block backwards. The surprising part is how ordinary the backward arithmetic stays once reconstruction is done: the novelty is the permission to forget, granted by invertibility, not a new gradient estimator.
What is disputed
The mechanism follows from designing invertible residual blocks so activations can be recalculated rather than stored. Quantitative gains depend on depth, width, batch size and hardware and are not fixed by the idea alone. Unrelated work on excitons and negative-mass electrons does not bear on this training method.
Remember this
If a residual block is invertible, its activations need not be stored: the backward pass can rebuild them, paying compute to buy memory.
Test yourself
A training step is running out of device memory on activation storage, yet the GPU is only partly busy. Why might replacing ordinary residual blocks with reversible ones help in this situation, and what new cost should you expect to see?
Reversible blocks allow the forward pass to discard intermediate activations and reconstruct them during the backward pass, cutting the activation memory footprint across depth. The new cost is extra compute: each reversed block performs additional function evaluations to rebuild the tensors that used to be read from storage. In a memory-bound, compute-rich setting that trade can raise feasible batch size or depth; in a compute-bound setting it may simply lengthen the step.
Go deeper
- [1707.04585] The Reversible Residual Network: Backpropagation Without Storing Activations · arxiv.org
- [2006.14705] Bright excitons with negative-mass electrons · arxiv.org
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.