Low-Memory Neural Network Training: A Technical Report
By Nimit Sharad Sohoni, Christopher R. Aberger, Megan Leszczynski, Jian Zhang, Christopher R'e. 2019
114 citations recorded by the source corpus.
Memory is increasingly often the bottleneck when training neural network models. Despite this, techniques to lower the overall memory requirements of training have been less widely studied compared to the extensive literature on reducing the memory requirements of inference. In this paper we study a fundamental question: How much memory is actually needed to train a neural network? To answer this question, we profile the overall memory usage of training on two representative deep learning benchmarks -- the WideResNet model for image classification and the DynamicConv Transformer model for machine translation -- and comprehensively evaluate four standard techniques for reducing the training memory requirements: (1) imposing sparsity on the model, (2) using low precision, (3) microbatching, and (4) gradient checkpointing. We explore how each of these techniques in isolation affects both the peak memory usage of training and the quality of the end model, and explore the memory, accuracy, and computation tradeoffs incurred when combining these techniques. Using appropriate combinations of these techniques, we show that it is possible to the reduce the memory required to train a WideResNet-28-2 on CIFAR-10 by up to 60.7x with a 0.4% loss in accuracy, and reduce the memory required to train a DynamicConv model on IWSLT'14 German to English translation by up to 8.7x with a BLEU score drop of 0.15.
Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Efficiency & Compression
How researchers see this paper
The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.
In Artificial Intelligence
In Computer Vision and Pattern Recognition
In Efficiency & Compression
Citations in the corpus
- Reformer: The Efficient Transformer (2020, methodology): “In addition to standard methods such as precision reduction and gradient checkpointing (Sohoni et al., 2019), more efficient versions of the Transformer model’s self-attention mechanism (Sukhbaatar et al., 2019a;b) have also recently been explored.”
- Blockwise Self-Attention for Long Document Understanding (2019, methodology): “In order to identify the bottleneck, we follow the memory model proposed by Sohoni et al. (2019), where the memory usage throughout neural network training is categorized into three main types: (1) Model Memory is used to store model parameters; (2) Optimizer Memory is the additional memory used by…”
- Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning (2020, methodology): “BERT-Large can only be used with access to a Google TPU, and BERT-Base requires some optimization tricks such as gradient checkpointing or gradient accumulation to be trained effectively on consumer hardware (Sohoni et al., 2019).”
- A Survey on Federated Learning for Resource-Constrained IoT Devices (2021, unknown): “Training on edge devices can bene-fit by adapting techniques that have served well in managing this problem in the server context, such as efficient gradient checkpointing [108], tensor rematerialization [109], and recompute [110].”
- Artificial Neural Networks and Machine Learning – ICANN 2021: 30th International Conference on Artificial Neural Networks, Bratislava, Slovakia, September 14–17, 2021, Proceedings, Part III (2021, background): “This also constitutes a realistic scenario on the production level, where a constrained and low model memory footprint is crucial [37].”
- Drawing early-bird tickets: Towards more efficient training of deep networks (2019, unknown): “Note that FP8 and FP32 are counted as the same FLOPs count units (Sohoni et al., 2019).”