# GPU Memory Usage Optimization for Backward Propagation in Deep Network Training

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 29 references that are in the map.
- First posted on arXiv 18 Feb 2025.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Learning Transferable Architectures for Scalable Image Recognition](https://mlsanity.com/paper/1707.07012/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Deep Learning with Limited Numerical Precision](https://mlsanity.com/paper/1502.02551/) (2015)
- [The Reversible Residual Network: Backpropagation Without Storing Activations](https://mlsanity.com/paper/1707.04585/) (2017)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [Low-Memory Neural Network Training: A Technical Report](https://mlsanity.com/paper/1904.10631/) (2019)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [Training Large Neural Networks with Constant Memory using a New Execution Algorithm](https://mlsanity.com/paper/2002.05645/) (2020)
- [Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory](https://mlsanity.com/paper/1911.13214/) (2019)
- [TFLMS: Large Model Support in TensorFlow by Graph Rewriting](https://mlsanity.com/paper/1807.02037/) (2018)
- [High-performance Semantic Segmentation Using Very Deep Fully Convolutional Networks](https://mlsanity.com/paper/1604.04339/) (2016)

## Abstract (arXiv)

In modern Deep Learning, it has been a trend to design larger Deep Neural Networks (DNNs) for the execution of more complex tasks and better accuracy. On the other hand, Convolutional Neural Networks (CNNs) have become the standard method for most of computer vision tasks. However, the memory allocation for the intermediate data in convolution layers can cause severe memory pressure during model training. Many solutions have been proposed to resolve the problem. Besides hardware-dependent solutions, a general methodology rematerialization can reduce GPU memory usage by trading computation for memory efficiently. The idea is to select a set of intermediate results during the forward phase as checkpoints, and only save them in memory to reduce memory usage. The backward phase recomputes the intermediate data from the closest checkpoints in memory as needed. This recomputation increases execution time but saves memory by not storing all intermediate results in memory during the forward phase. In this paper, we will focus on efficiently finding the optimal checkpoint subset to achieve the least peak memory usage during the model training. We first describe the theoretical background of the training of a neural network using mathematical equations. We use these equations to identify all essential data required during both forward and backward phases to compute the gradient of weights of the model. We first identify the checkpoint selection problem and propose a dynamic programming algorithm with time complexity O(n3) to solve the problem of finding the optimal checkpoint subset. With extensive experiments, we formulate a more accurate description of the problem using our theoretical analysis and revise the objective function based on the tracing, and propose an O(n)-time algorithm for finding the optimal checkpoint subset.