[Paper Review] Low-Memory Neural Network Training: A Technical Report

Nimit S. Sohoni, Christopher R. Aberger+3 more|arXiv (Cornell University)|Apr 24, 2019

Advanced Neural Network Applications83 references38 citations

TL;DR

The paper quantifies memory requirements for training and evaluates four techniques—sparsity, low precision, microbatching, and gradient checkpointing—to reduce training memory with controlled accuracy loss on WideResNet and DynamicConv Transformer.

ABSTRACT

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.

Motivation & Objective

• Quantify the memory components involved in neural network training (model, optimizer, activation) and identify their relative contributions.
• Evaluate how four memory-reduction techniques affect training memory and model accuracy in isolation.
• Explore how combining these techniques impacts memory, accuracy, and computation.
• Demonstrate potential memory savings on representative benchmarks (WideResNet on CIFAR-10 and DynamicConv on IWSLT'14 De→En).

Proposed method

• Profile memory usage into model, optimizer, and activation components during training.
• Evaluate sparsity via dynamic sparse reparameterization and measure accuracy under varying nonzero percentages.
• Assess low-precision training with FP16 and dynamic loss scaling on accuracy and FLOPs.
• Test microbatching by simulating smaller microbatches and analyze accuracy with batch-normalization considerations.
• Analyze gradient checkpointing strategies to reduce activation memory and quantify the FLOPs tradeoffs when memory is reduced.

Experimental results

Research questions

• RQ1 What are the dominant memory components during training across representative models?
• RQ2 How do sparsity, low precision, microbatching, and checkpointing individually affect memory and accuracy?
• RQ3 What are the tradeoffs in computation when applying these techniques?
• RQ4 How can these techniques be combined to maximize memory reduction under a given accuracy constraint?
• RQ5 Do the combined techniques enable training with significantly reduced memory for WideResNet and DynamicConv Transformer?

Key findings

• Activation memory dominates total training memory in the studied setups.
• Sparsity up to 70% for WideResNet and 60% for DC-Transformer yields small accuracy/BLEU losses (0.3% and 0.8 BLEU respectively).
• FP16 training does not significantly affect WideResNet accuracy and causes only a 0.15 BLEU drop for DC-Transformer.
• Microbatch sizes as small as 10 preserve WideResNet accuracy; smaller sizes degrade accuracy when extremely small.
• Gradient checkpointing can reduce activation memory by up to about 5.8x while increasing FLOPs, with no loss in accuracy.
• Combined techniques can reduce WideResNet training memory by up to 60.7x with 0.4% accuracy loss, and DC-Transformer memory by up to 8.7x with 0.15 BLEU loss.