Dynamic Tensor Rematerialization

• Dynamic Tensor Rematerialization (DTR) is an approach that reduces memory consumption in deep learning by dynamically trading memory usage against recomputation cost.
• It extends static checkpointing to arbitrary computation graphs using both MILP-based optimization and online greedy heuristics for adaptive tensor scheduling.
• DTR integrates with system architectures like paging and heterogeneous device scheduling to enable efficient training on memory-constrained hardware.

Dynamic Tensor Rematerialization (DTR) is an algorithmic and systems approach for reducing memory consumption during the training and inference of deep neural networks, particularly in resource-constrained or dynamic environments. DTR trades off memory usage against recomputation cost by dynamically deciding which intermediate tensor activations to retain, evict, and rematerialize on demand—in contrast to prior static checkpointing methods, which fix these decisions offline. Its scope now spans from optimal planning (e.g., Checkmate), to greedy runtime heuristics (DTR proper), to integration with paging (POET), heterogeneous scheduling (XEngine), advanced constraint programming (Moccasin), memory system awareness (Coop), and extension to dynamic shape graphs (BladeDISC++).

1. Foundations: Rematerialization and Generalization Beyond Checkpointing

Traditional checkpointing strategies—for example, the $\mathcal{O}(\sqrt{n})$ rule and Griewank’s log-$n$ method—were designed for linear computation graphs with uniform memory and compute cost, achieving tractable tradeoffs in backward passes (Jain et al., 2019). DTR extends these ideas to arbitrary computation graphs (DAGs), accommodating nonuniform tensor sizes and computational profiles (Kirisame et al., 2020). Rematerialization here refers to intentionally discarding intermediate activations once they are not immediately required, and regenerating them later (via recomputation) when dependencies dictate. This generalization encompasses complex architectures with residuals, skips, and branches, far beyond the reach of earlier analytic formulas.

2. Algorithmic Strategies: From MILP-Optimized Schedules to Online Heuristics

Static approaches such as Checkmate frame rematerialization scheduling as an integer linear program (ILP) or mixed-integer linear program (MILP). In this context, binary matrices $S$ and $R$ encode at each timestep whether a tensor is checkpointed (in memory) or recomputed, respectively. The objective is typically:

$\min \sum_{t, i} C_i \cdot R_{t,i}$

where $C_i$ is the profiled compute cost of operator $i$, subject to constraints that preserve dependency correctness and enforce hardware-specific memory budgets. MILP solvers (e.g., Gurobi) can obtain optimal or near-optimal schedules, with solve times ranging from seconds to under an hour for mid-sized networks, and two-phase rounding approximations for scalability (Jain et al., 2019).

DTR, in contrast, is an online, greedy algorithm. It dynamically interposes on tensor allocations at runtime, ranks tensors for eviction using heuristics that combine staleness, memory size, and recomputation cost (including downstream evicted dependencies). The runtime adapts to actual execution traces, enabling support for dynamic control flows not possible in static checkpointing frameworks (Kirisame et al., 2020). Specifically, eviction decisions frequently use a score such as:

$h(t) = \frac{c(t)}{m(t) \cdot s(t)}$

where $c(t)$ is the rematerialization cost (including recursively evicted “neighborhood”), $m(t)$ the tensor’s memory size, and $s(t)$ its staleness.

Constraint programming, as in Moccasin, further innovates by modeling retention intervals for tensors and solving minimization subject to memory and precedence via $O(n)$ integer variables rather than $O(n^2)$ Boolean variables, delivering an order of magnitude speedup on large computation graphs (Bartan et al., 2023).

3. Integration with System Architectures and Advanced Memory Models

POET demonstrates the integration of rematerialization with paging, jointly optimizing schedules that weigh (i) recomputation cost, (ii) page-in and page-out energy, and (iii) throughput constraints, using a MILP framework. Rematerialization is applied for lightweight operations, paging for compute-intensive activations, with auxiliary storage (flash, SD card) used for cheap swaps. The joint optimization enables fine-tuning large architectures (ResNet-18, BERT) on devices with 32KB RAM, without violating backpropagation correctness (Patil et al., 2022).

XEngine generalizes the MILP approach to mixed-integer quadratic programming (MIQP), handling operator placement, memory usage, and transfer costs across heterogeneous devices (CPU, GPU). Its objective function minimizes end-to-end compute and communication cost, substantially outperforming static single-device schedulers like Checkmate—achieving up to 22.5% runtime speedup in VGG19 inference schedules under tight memory budgets (Schuler et al., 2022).

BladeDISC++ introduces symbolic shape rematerialization for dynamic graphs, analyzing tensor sizes algebraically (e.g., @S0 $= 12 \cdot$ @S1) at compile time. Regeneration branches are inserted during compilation, and runtime eviction is triggered adaptively according to actual memory pressure. This compilation–runtime hybrid approach achieves memory efficiency on par with static shape optimization, facilitating dynamic shape graph training (e.g., Llama-2-1b) (Yuan et al., 22 Dec 2024).

4. Challenges, Limitations, and Recent Remedies

Memory fragmentation—where evictions from disjoint regions yield unusable free blocks—undermines many heuristic-based rematerialization systems. Coop addresses this by enforcing contiguous evictions via a sliding window algorithm, partitioning tensors by cost density, and exploiting recomputable in-place mutations to minimize actual memory fragmentation and compute overhead (Zhang et al., 2023). Given a required allocation $M_R$, Coop finds the least-cost set $S$ contiguous in memory such that $M(S, L) \geq M_R$:

$$\arg\min_{S, L} \sum_{t \in S} h(t), \quad \text{subject to} \quad M(S, L) \ge M_R$$

Empirically, Coop achieves up to $2\times$ memory savings, reduces search latency up to $10\times$, and brings fragmentation below $5\%$ in representative large-scale models.

Adversarial graph constructions, as analyzed in (Kirisame et al., 2020), prove that any deterministic online eviction strategy can be forced to incur $\Omega(N/B)$ times more tensor operations than an offline optimal, setting lower bounds on worst-case overhead. Trade-offs persist between richness of runtime metadata, search space size, and schedule sub-optimality in dynamic settings.

5. Performance Metrics and Experimental Evaluations

The evaluation of DTR methods encompasses metrics such as:

• Peak memory usage reduction (e.g., Checkmate allows a VGG19 batch size increase from 167 to 289 on a V100 GPU (Jain et al., 2019))
• Compute overhead, measured as additional tensor operations and runtime slowdown relative to unconstrained execution (Kirisame et al., 2020)
• Energy savings, as reported by POET (e.g., up to 35% lower energy overhead relative to prior baselines) (Patil et al., 2022)
• Scalability of optimization solvers (Moccasin is up to $10\times$ faster than prior MILP solvers, scaling to graphs with $>1000$ nodes (Bartan et al., 2023))
• Search latency and fragmentation rate (Coop: $10\times$ lower search latency, $<5\%$ fragmentation) (Zhang et al., 2023)
• Throughput (BladeDISC++ achieves comparable tokens/sec to static optimization in dynamic shape training) (Yuan et al., 22 Dec 2024)

These empirical results substantiate that DTR strategies consistently enable larger batch sizes, input resolutions, and broader architecture exploration within fixed device memory budgets.

6. Applications and Extension to Specialized Settings

Dynamic tensor rematerialization is critical for:

• Training and deploying large networks (ResNets, Transformers, U-Nets, GPT-3 variants) on GPUs, edge devices, FPGAs, and microcontrollers with severely limited RAM (Patil et al., 2022, Zhang et al., 2023)
• Models with dynamic control flow, e.g., TreeLSTMs, variable-length sequence models, where static checkpointing is infeasible (Kirisame et al., 2020)
• Heterogeneous system environments requiring distributed scheduling and load balancing across CPUs/GPUs (Schuler et al., 2022)
• Sparse tensor decomposition applications such as multi-threaded spMTTKRP, where dynamic remapping and high locality layouts (FLYCOO) further reduce data movement (Wijeratne et al., 2023)
• Dynamic shape compiler scenarios, where symbolic reasoning over unknown dimensions allows runtime-efficient scheduling and rematerialization (Yuan et al., 22 Dec 2024)
• Integration with paging and external memory hierarchies for energy-aware training with hard throughput deadlines (Patil et al., 2022)

7. Broader Implications and Future Directions

The evolution of DTR enables training regimes decoupled from memory limitations, facilitating privacy-preserving edge personalization and enabling larger, deeper networks on constrained hardware. Advances in symbolic shape analysis and runtime adaptation foreshadow growing adoption in dynamic graph compilers and variable-shape workloads (Yuan et al., 22 Dec 2024). Further integration of co-optimization for memory management (allocation/rematerialization), concurrent processing techniques, and energy/latency-aware scheduling will likely characterize future research. The persistent challenge remains bridging the gap between static optimality and dynamic system demands, especially under adversarial topologies and highly nonuniform cost profiles. Efforts to reduce memory fragmentation, latency, and runtime overhead are central to advancing both theoretical understanding and practical deployment of dynamic tensor rematerialization across machine learning, signal processing, and scientific computing domains.