Rematerialization¶
Recompute a needed intermediate at its later use instead of retaining or retrieving it, trading extra work for less storage or data movement.
Core Idea¶
Rematerialization recomputes a needed intermediate at its later use instead of retaining its earlier value or fetching a stored copy. The reconstruction must be possible from valid dependencies and produce a result acceptable under the program's semantics. Extra work buys reduced storage occupancy or data movement. Briggs, Cooper and Torczon describe the compiler case as choosing recomputation when it is cheaper than storing and reloading a value that cannot remain in a register.[^ref-f50b18aab535]
Scope of Application¶
In compiler register allocation, a long-lived value can occupy a scarce register or be spilled to memory. If a correct and inexpensive recipe is available at a later instruction, the allocator can regenerate it instead. In JAX reverse-mode differentiation, jax.checkpoint/jax.remat can avoid retaining selected forward residuals until the backward pass, recomputing them from saved inputs when needed. The storage and cost units differ, but the later-use reconstruction pattern is the same.[ref-f50b18aab535][ref-e444ebc0588d]
Clarity¶
The method requires a future use, an available reconstruction recipe, and an avoided retained or retrieved copy. Dropping a value never used again is dead-value elimination. Reading from an offloaded CPU-memory copy is retrieval, not recomputation. Repeating an expression after mutable state changes may fail semantic correctness even if the expression's text is unchanged.[^ref-e444ebc0588d]
Manages Complexity¶
Rematerialization replaces a long retention interval with a local reconstruction obligation. That can ease register pressure or accelerator-memory demand. It is selective, not a policy to recompute everything: JAX's official worked example saves some checkpoint inputs and outputs while recomputing internal residuals; policies can preserve expensive operations and offload others. The correct choice depends on compute, storage and transfer costs.[^ref-e444ebc0588d]
Abstract Reasoning¶
Identify the intermediate and later consumer. Compare keeping it live, saving and reloading it, offloading it, and recreating it near use. Verify that all needed inputs and effects still yield an acceptable value. Then compare costs under the actual hardware and workload. A safe reconstruction may be unprofitable; a cheap one may be invalid. These are separate decisions.[ref-f50b18aab535][ref-e444ebc0588d]
Knowledge Transfer¶
The compiler and JAX cases instantiate the same role mapping: later-needed value → costly retention → dependency recipe → later reconstruction → equivalent value → storage-versus-work comparison. In a compiler, registers and stack traffic dominate; in gradient checkpointing, forward residuals, device memory and backward FLOPs dominate. The workspace DAG leaves this provisionally unparented: Register Allocation is not universal, and the live Trade-offs prime requires a nontrivial frontier and substitution structure not necessary for every rematerialization.[ref-f50b18aab535][ref-e444ebc0588d]
[^ref-f50b18aab535]: Preston Briggs, Keith D. Cooper and Linda Torczon, “Rematerialization,” Proceedings of PLDI 1992, pp. 311–321, original ACM publisher abstract. [^ref-e444ebc0588d]: JAX project, “Gradient checkpointing with jax.checkpoint (jax.remat)”, TL;DR, saved-residual example, policies and offloading, accessed 2026-09-30.
Neighborhood in Abstraction Space¶
Rematerialization sits in a moderately populated region (53rd percentile for distinctiveness): it has near-neighbors but no dense thicket of look-alikes.
Family — Program Execution & Runtime Concepts (27 abstractions)
Nearest neighbors
- Thrashing — 0.87
- Dynamic Problem — 0.86
- Elapsed-Time Memory Decay — 0.86
- Data Migration — 0.85
- Incremental Backup — 0.85
Computed from structural-signature embeddings · 2026-10-08