# Moccasin: Efficient Tensor Rematerialization for Neural Networks

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 30 references that are in the map.
- First posted on arXiv 27 Apr 2023.

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

- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [CVXPY: A Python-Embedded Modeling Language for Convex Optimization](https://mlsanity.com/paper/1603.00943/) (2016)
- [Efficient Processing of Deep Neural Networks: A Tutorial and Survey](https://mlsanity.com/paper/1703.09039/) (2017)
- [A Rewriting System for Convex Optimization Problems](https://mlsanity.com/paper/1709.04494/) (2017)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [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)
- [Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs](https://mlsanity.com/paper/2111.06483/) (2021)
- [Neural Topological Ordering for Computation Graphs](https://mlsanity.com/paper/2207.05899/) (2022)
- [A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation](https://mlsanity.com/paper/1905.11722/) (2019)
- [Combinatorial Register Allocation and Instruction Scheduling](https://mlsanity.com/paper/1804.02452/) (2018)
- [Survey on Combinatorial Register Allocation and Instruction Scheduling](https://mlsanity.com/paper/1409.7628/) (2014)
- [XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments](https://mlsanity.com/paper/2212.09290/) (2022)

## Abstract (arXiv)

The deployment and training of neural networks on edge computing devices pose many challenges. The low memory nature of edge devices is often one of the biggest limiting factors encountered in the deployment of large neural network models. Tensor rematerialization or recompute is a way to address high memory requirements for neural network training and inference. In this paper we consider the problem of execution time minimization of compute graphs subject to a memory budget. In particular, we develop a new constraint programming formulation called textscMoccasin with only O(n) integer variables, where n is the number of nodes in the compute graph. This is a significant improvement over the works in the recent literature that propose formulations with O(n2) Boolean variables. We present numerical studies that show that our approach is up to an order of magnitude faster than recent work especially for large-scale graphs.