Low-Memory Neural Network Training: A Technical Report

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.

Enabling On-Device Smartphone GPU based Training: Lessons Learned

“…However, since convolutional layers are computationally more expensive than fully-connected layers (i.e., our target to improve in this work) as analyzed in [17] and the real bottleneck of ondevice training is memory bound as in Section 4, analyzing and improving computationally expensive convolutional layers can be a potential future direction. Our finding, memory bottleneck issue, also suggests investigating how to enable ondevice training with memory optimization in terms of model, optimizer, and the activation [18].…”

Enabling Binary Neural Network Training on the Edge

“…In order to quantify the potential gains from approximation, we conducted a variable representation and lifetime analysis of Algorithm 1 following the approach taken by Sohoni et al (2019). Table 2 lists the properties of all variables in Algorithm 1, with each variable's contributions to the total footprint shown for a representative example.…”

“…Despite featuring binary forward propagation, existing BNN training approaches perform backward propagation using high-precision floating-point data types-typically float32-often making training infeasible on resourceconstrained devices. The high-precision activations used between forward and backward propagation commonly constitute the largest proportion of the total memory footprint of a training run (Sohoni et al, 2019;Cai et al, 2020). Moreover, backward propagation with high-precision gradients is costly, challenging the energy limitations of edge platforms.…”

Feature Alignment As A Generative Process

“…By constraining the trainable parameters, such as the weights, to be updated only by local variables (the information contained in the neurons that share the same parameter), we can reduce the memory requirements to load a model in hardware such as CPUs and GPUs. This constraint can save memory resources and has many potential applications, from low-memory devices 7,8 to train large batch sizes 9,10 , and, even further, to train very large neural networks 11 .…”