# XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments

arXiv: https://arxiv.org/abs/2212.09290 · HTML: https://mlsanity.com/paper/2212.09290/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 31 references that are in the map.
- First posted on arXiv 19 Dec 2022.

## Builds on (most cited first, 16 of 16)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [The Reversible Residual Network: Backpropagation Without Storing Activations](https://mlsanity.com/paper/1707.04585/) (2017)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks](https://mlsanity.com/paper/1801.04380/) (2018)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging](https://mlsanity.com/paper/2207.07697/) (2022)
- [Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory](https://mlsanity.com/paper/1911.13214/) (2019)
- [A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation](https://mlsanity.com/paper/1905.11722/) (2019)
- [Divide-and-Conquer Checkpointing for Arbitrary Programs with No User Annotation](https://mlsanity.com/paper/1708.06799/) (2017)
- [DELTA: Dynamically Optimizing GPU Memory beyond Tensor Recomputation](https://mlsanity.com/paper/2203.15980/) (2022)
- [Mimose: An Input-Aware Checkpointing Planner for Efficient Training on GPU](https://mlsanity.com/paper/2209.02478/) (2022)

## Built on by (most cited first, 2 of 2)

- [Moccasin: Efficient Tensor Rematerialization for Neural Networks](https://mlsanity.com/paper/2304.14463/) (2023)
- [Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch](https://mlsanity.com/paper/2608.08740/) (2026)

## Abstract (arXiv)

Memory efficiency is crucial in training deep learning networks on resource-restricted devices. During backpropagation, forward tensors are used to calculate gradients. Despite the option of keeping those dependencies in memory until they are reused in backpropagation, some forward tensors can be discarded and recomputed later from saved tensors, so-called checkpoints. This allows, in particular, for resource-constrained heterogeneous environments to make use of all available compute devices. Unfortunately, the definition of these checkpoints is a non-trivial problem and poses a challenge to the programmer - improper or excessive recomputations negate the benefit of checkpointing. In this article, we present XEngine, an approach that schedules network operators to heterogeneous devices in low memory environments by determining checkpoints and recomputations of tensors. Our approach selects suitable resources per timestep and operator and optimizes the end-to-end time for neural networks taking the memory limitation of each device into account. For this, we formulate a mixed-integer quadratic program (MIQP) to schedule operators of deep learning networks on heterogeneous systems. We compare our MIQP solver XEngine against Checkmate, a mixed-integer linear programming (MILP) approach that solves recomputation on a single device. Our solver finds solutions that are up to 22.5 % faster than the fastest Checkmate schedule in which the network is computed exclusively on a single device. We also find valid schedules for networks making use of both central processing units and graphics processing units if memory limitations do not allow scheduling exclusively to the graphics processing unit.