PaperlayerView exact revision on arXiv

Source-linked AI summary

Low-Memory Neural Network Training: A Technical Report

Nimit S. Sohoni, Christopher R. Aberger, Megan Leszczynski, Jian Zhang, Christopher Ré

arXiv:1904.10631v2cs.LGstat.ML

TL;DR

Training memory is an important bottleneck, yet training-memory reduction has received less comprehensive study than inference-memory reduction. The paper profiles two deep-learning benchmarks and evaluates four memory-reduction techniques individually and in combination. Appropriate combinations reduce training memory substantially on both tasks, with small reported accuracy or BLEU drops.

  • Problem

    Training is often far more memory-intensive than inference, yet techniques for reducing training memory and their interactions have not been comprehensively studied.

    Source: §1 · pp. 1, 4

  • Method

    The paper profiles training memory and evaluates sparsity, low precision, microbatching, and gradient checkpointing individually and in combinations.

    Source: §1 · p. 4

  • Results

    Appropriate combinations cut WideResNet training memory by up to 60.7x with a 0.4% accuracy drop, and DC-Transformer memory by up to 8.7x with a 0.15 BLEU drop.

    Source: §1 · p. 4

  • Takeaways & Limitations

    The techniques can be combined to achieve near-state-of-the-art results with comparable computation and an order of magnitude less memory than standard training approaches.

    Source: §6 · p. 26

    Microbatching changes batch-normalization statistics, which can affect final model accuracy.

    Source: §3.3 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

Training can require far more memory than inference, while many inference-memory techniques do not apply directly to training. This paper profiles training memory and evaluates four reduction techniques, individually and in combination, on image-classification and translation benchmarks.

1 Introduction

  • Training is often far more memory-intensive than inference, and many inference-memory techniques require first training a full-size model, limiting their applicability to training-memory reduction.The paper notes that neural networks commonly contain millions or billions of parameters, making memory a hardware bottleneck.Source: §1 · p. 1
  • Sparsity can reduce model and optimizer memory, while low precision reduces all three memory components; the reported accuracy costs vary by model and technique.WideResNet tolerates 70% sparsity with a 0.3% accuracy loss; DC-Transformer tolerates 60% sparsity with a 0.8 BLEU drop. Half precision causes no WideResNet accuracy drop and a 0.15 BLEU drop for DC-Transformer.Source: §1 · p. 3
  • Microbatching reduces activation memory but can alter batch-normalization statistics, while checkpointing saves activation memory by recomputing unstored activations during backpropagation.Microbatch sizes as small as 10 cause no WideResNet accuracy loss in the reported experiments; checkpointing increases computation.Source: §1 · p. 3
  • Combining techniques matters because their interactions can make the best settings in isolation suboptimal for accuracy under a fixed training-memory budget.The paper also frames training-memory reduction as relevant to constraints on model size and input dimensionality.Source: §1 · pp. 2, 4
  • Up to 60.7x less training memory with a 0.4% accuracy drop on WideResNet, and up to 8.7x less with a 0.15 BLEU drop on DC-Transformer, using appropriate technique combinations.For WideResNet, sparse models can also achieve higher accuracy than smaller models at the same memory budget; for DC-Transformer, smaller models outperform sparse ones.Source: §1 · p. 4
  • The study profiles WideResNet on CIFAR-10 and DynamicConv Transformer on IWSLT’14 German-to-English, then evaluates sparsity, low precision, microbatching, and gradient checkpointing.It examines their individual memory, accuracy, and computation tradeoffs and how to combine them.Source: §1 · pp. 3–4

2 Preliminaries

The section categorizes training memory into model, optimizer, and activation components, then profiles their contributions and describes how the profiler estimates memory. Across the two studied architectures, activations are the largest component and model memory the smallest.

2 Preliminaries

  • The section profiles training memory in WideResNet and DC-Transformer to inform the memory-reduction techniques studied later.It focuses on model, optimizer, and activation memory.Source: §2 · p. 5

2.1 Types of Memory

  • Model memory stores parameters, optimizer memory stores gradients and momentum buffers, and activation memory stores layer outputs needed for backpropagation.Activation memory includes forward activations and gradients with respect to those activations during backpropagation.Source: §2.1 · p. 5

2.2.1 Types of Memory

  • In both studied architectures, model memory is the smallest of the three categories, while activations consume most memory in standard training.Model memory is 2–3x less than optimizer memory in these settings and can be orders of magnitude less than activation memory.Source: §2.2.1 · pp. 5–6
  • Figure 1 presents training-memory consumption as pie charts for WideResNet on CIFAR-10 and DC-Transformer on IWSLT’14 German-to-English.Optimizer memory is 2x model memory for WideResNet with Nesterov momentum and 3x for DC-Transformer with Adam.Source: Figure 1 · p. 6, §2.2.1 · p. 6
  • The study does not examine updating a momentum buffer in place without a separate gradient buffer, which would prevent standard gradient clipping and may cause numerical issues, especially with low precision.The authors note that this alternative can save the separate gradient buffer.Source: §2.2.1 · p. 6

2.2.2 Profiler Implementation

  • The profiler estimates model and optimizer memory from tensor sizes and activation memory by traversing PyTorch’s autograd graph and summing required activation-input storage.The calculated total matches PyTorch’s maximum allocated-memory profiling within 5% on WideResNet and 11% on DC-Transformer.Source: §2.2.2 · p. 7

3 Techniques and Tradeoffs

This section evaluates four techniques for reducing training memory and examines their individual memory–accuracy or memory–computation tradeoffs. It describes how sparsity, low precision, microbatching, and checkpointing change storage or recomputation.

3 Techniques and Tradeoffs

  • The study compares sparsity, low precision, microbatching, and checkpointing; all but checkpointing are assessed mainly for memory–accuracy tradeoffs, and checkpointing for memory–computation tradeoffs.For the first three techniques, the section reports no adverse tradeoff in total FLOPs; checkpointing increases computation through recomputation.Source: §3 · p. 7

3.1 Sparsity

  • Dynamic sparse reparameterization prunes small-magnitude weights and introduces new nonzeros, keeping the model’s nonzero count constant during training.Gradients and momentum share the model’s sparsity pattern, reducing both model and optimizer memory.Source: §3.1 · pp. 7–8
  • At 80% sparsity in convolutions, test accuracy drops less than 0.5% relative to the dense baseline.Figure 3 plots test accuracy against the percentage of nonzero parameters for WRN-28-2 on CIFAR-10.Source: Figure 3 · p. 9, §3.1 · p. 9

3.2 Low Precision

  • FP16 training uses 16-bit quantities instead of standard 32-bit single precision, reducing memory across model, optimizer, and activation storage.Low precision can also reduce computation, though numerical issues are possible.Source: §3.2 · p. 10
  • On CIFAR-10, FP16 yields 94.43 ± 0.17% test accuracy for WideResNet, versus 94.37 ± 0.07% with FP32.The study uses dynamic loss scaling and keeps batch-normalization parameters in FP32.Source: §3.2 · p. 10

3.3 Microbatching

  • Microbatch sizes as small as 10 preserve WideResNet accuracy versus the size-100 baseline, while accuracy degrades at sizes 4 or 2.Microbatching accumulates gradients before updating parameters; with batch normalization, its statistics differ from those of the full minibatch.Source: §3.3 · pp. 10–11, Figure 4 · p. 12

3.4 Checkpointing

  • Checkpoint-residual-2* reduces WideResNet activation storage by 5.8x with a 30% increase in total FLOPs.Checkpointing stores selected activations and recomputes discarded ones during backpropagation; it trades computation for memory without a memory–accuracy tradeoff.Source: §3.4 · pp. 12, 14, Figure 5 · p. 13

4 Experiments

The experiments compare memory–accuracy tradeoffs for WideResNet on CIFAR-10 and DC-Transformer on IWSLT’14 German-to-English, including combinations of memory-reduction techniques and smaller-network baselines. Combining techniques substantially reduces training memory on both benchmarks, while the relative performance of sparsity and smaller architectures differs between them.

4.2.2 Tradeoff Studies

  • 6.7 MB (60.7x below the 404.8 MB baseline) yields 93.95% CIFAR-10 accuracy versus 94.37% for baseline WideResNet training.The setting uses 30% nonzeros, microbatch size 10, FP16, and checkpoint-residual-2*.Source: §4.2.2 · pp. 16–17, Table 3 · p. 19
  • FP16 and FP32 show no significant WideResNet accuracy difference across tested sparsity levels, supporting FP16 use in subsequent experiments.Figure 6 plots test accuracy against sparsity for both precisions.Source: §4.2.2 · p. 16, Figure 6 · p. 16
  • Combining techniques gives better accuracy at a fixed memory budget than increasing sparsity or reducing microbatch size alone; checkpoint-residual-2* cuts activation memory about 5.8x at roughly 30% more FLOPs.Across the WideResNet settings, low precision and checkpointing alone reduce memory 9.5x; adding microbatching reaches 33.2x with 0.09% accuracy loss, and sparsity reaches 60.7x with 0.42% loss.Source: §4.2.2 · p. 16, §4.2.5 · p. 19

4.2.3 Microbatching with Low Precision

  • At microbatch sizes 10 and 4, simulated and true microbatching differ by only 0.06%, within random variation.The comparison uses FP16 and dense WideResNet training; true microbatching slightly outperforms simulation at size 4.Source: §4.2.3 · p. 17, Table 2 · p. 18

4.2.4 Baselines

  • For WideResNet, sparsity preserves accuracy better than reducing depth or width at a given memory budget until the crossover below 5 MB, where accuracy has already degraded substantially.The study notes that sparsity and architecture changes could be combined, but does not test that combination.Source: §4.2.4 · p. 18

4.3.5 Takeaways

  • 8.7x lower training memory gives DC-Transformer a BLEU score of 34.99 versus the baseline 35.15, a drop of 0.15.The reduction combines FP16, microbatching, and checkpointing.Source: §4.3 · p. 19, §4.3.5 · p. 23

4.3.2 Tradeoff Studies

  • For DC-Transformer, FP16 has a small BLEU cost that grows with sparsity, while 250-token microbatching is mathematically equivalent to standard training and checkpoint-residual-1 cuts activation memory about 5.7x at roughly 30% more FLOPs.The model uses layer normalization rather than batch normalization, so examples are not interdependent across the minibatch.Source: §4.3.2 · p. 21

4.3.4 Baselines

  • At a fixed memory budget, narrower DC-Transformer models generally achieve higher BLEU than sparse models, while shallower models perform slightly worse than sparse ones.Figure 10 plots test BLEU against training memory for shallow, narrow, and sparse variants; all use FP16, 250-token microbatches, and checkpoint-residual-1.Source: Figure 10 · p. 22, §4.3.4 · p. 22

5 Related Work

The section surveys four areas of training-memory research and other compression approaches, distinguishing inference-focused methods from approaches that reduce memory during training.

5.1 Sparse Neural Networks

  • Prior sparse-training work often targets inference memory, while sparsity alone may not yield low-memory training when activations exceed even dense-model memory.Earlier approaches pruned during training or adaptively evolved a fixed parameter budget; the section distinguishes these from reducing the overall training-memory footprint.Source: §5.1 · p. 23

5.2 Low Precision Training

  • Half-precision training can reduce storage, but some approaches retain 32-bit or master-weight copies, and all-half storage saves at most 2x.The most aggressive quantization approaches discussed can also lose accuracy relative to state of the art.Source: §5.2 · p. 24

5.3 Small-batch Training

  • Small minibatches have been studied for generalization, whereas microbatching splits minibatches into independently computed gradients that are accumulated.Small-minibatch effects may depend on hyperparameters and setting; alternatives to batch normalization can still lose accuracy at very small microbatch sizes.Source: §5.3 · pp. 24–25

5.4 Gradient Checkpointing

  • Checkpointing reduces activation storage without affecting accuracy, while reversible networks avoid storing most activations but require architecture changes.The section notes practical checkpointing strategies and a fused normalization-and-activation approach that can avoid storing up to 50% of activations.Source: §5.4 · p. 25

5.5 Other Techniques

  • Post-training compression and distillation reduce inference-model size but require first training a full-size or larger network; parameter-efficient architectures are another route.Pruning and quantization can preserve final accuracy after fine-tuning, while distillation trains a smaller model to imitate a larger one.Source: §5.5 · p. 25
  • Other approaches reduce memory through structured weight matrices, grouped optimizer state, or activation quantization, sparsification, allocation, and offloading.The section identifies activation-compression methods as orthogonal to the techniques studied and a possible direction for future work.Source: §5.5 · pp. 25–26

6 Conclusion

The conclusion summarizes the paper’s study of training-memory reduction and its reported memory, accuracy, and computation tradeoffs.

6 Conclusion

  • Combining techniques appropriately yields near-state-of-the-art results with comparable computation and an order of magnitude less training memory than standard approaches.The study examines model, optimizer, and activation memory alongside final accuracy.Source: §6 · p. 26

A.1 Experimental Statistics

The appendix reports trial counts and the method used to compute confidence-interval half-widths.

A.1 Experimental Statistics

  • WideResNet results average 5 independent random-seed trials per setting, DC-Transformer results average 3, and 95% confidence-interval half-width is standard error × 1.96.The paper uses fewer DC-Transformer trials because run-to-run variability is lower for that model.Source: §A.1 · p. 34

A.2 Hyperparameters

The appendix identifies the sources of the default hyperparameters used in the two benchmark experiments.

A.2 Hyperparameters

A.3 Hyperparameter Tuning Procedure

The appendix describes how learning-rate and sparse-rewiring hyperparameters were selected separately for each architecture and sparsity setting, using validation performance before held-out testing.

A.3 Hyperparameter Tuning Procedure

  • Hyperparameters were tuned separately for each architecture and sparsity setting, with learning rate and, for sparse networks, rewiring frequency and fraction among the search parameters.The selected parameters were then used across experiments sharing that architectural setting.Source: §A.3 · p. 34
  • WideResNet searches used 45,000 training examples and a randomly sampled validation set, with five seeds; final tests trained on all 50,000 examples and used the held-out 10,000-example test set.The selected setting maximized mean validation accuracy.Source: §A.3 · p. 34
  • Learning-rate multipliers were searched relative to each model’s default schedule, with WideResNet results in Table 6 and DC-Transformer results in Table 7.Table 5 lists hyperparameters for all experiments.Source: §A.3 · pp. 34–35, Table 5 · p. 35
  • DC-Transformer searches used the predefined training and validation sets and three seeds, selecting the setting with the best mean validation BLEU before evaluation on the predefined held-out test set.The authors report three trials were sufficient to identify the best hyperparameter within 95% confidence in this setting.Source: §A.3 · pp. 34–35

A.3.1 Dynamic Sparse Reparameterization Details

The appendix specifies model-specific dynamic sparse reparameterization settings, including rewiring searches and validation-based selection for each sparsity level.

A.3.1 Dynamic Sparse Reparameterization Details

  • WideResNet’s rewiring fraction was set to 0.01377866 by default, matching 20,000 parameters on the CIFAR-10 architecture, and frequency and fraction were grid-searched for each sparsity level.The selected combination was based on validation accuracy and used for test experiments at that sparsity level.Source: §A.3.1 · pp. 35–36
  • For WideResNet, rewiring frequency and fraction were searched using separate multiplier grids around the defaults for each sparsity level.The frequency multipliers were {0.125, 0.25, 0.5, 1, 2, 4}; fraction multipliers were {0.25, 0.5, 1, 2, 4, 8}.Source: §A.3.1 · p. 36
  • For DC-Transformer, rewiring periods were halved to match its 50,000-update training schedule, and grid searches included a no-rewiring option that produced suboptimal BLEU at every sparsity level.The selected frequency and fraction were based on validation BLEU and used for test experiments at that sparsity level.Source: §A.3.1 · p. 36

B Activation Memory Profiler Implementation

The profiler estimates peak activation memory across forward and backward computation, adjusting its accounting for checkpointed execution and recomputation.

B Activation Memory Profiler Implementation

  • Without checkpointing, the profiler traverses the autograd graph in reverse topological order and finds the maximum stored activation memory needed while backpropagating.It counts forward activations needed for gradients and backward activations, while filtering operations such as reshapes and transposes that do not require their inputs for gradients.Source: B Activation Memory Profiler… · p. 36
  • The profiler excludes temporary memory for intermediate computations because it is generally a small fraction of total usage and can depend heavily on implementation.Activation memory includes gradients with respect to forward activations, while weight gradients are separately counted in the model gradient buffer.Source: B Activation Memory Profiler… · pp. 36, 38
  • With checkpointing, the profiler counts stored checkpoint-node outputs and the activations needed within each segment, then computes the maximum memory required at any one time.It also accounts for small cached quantities used to recompute some layers; that additional memory is less than 1% of total usage.Source: B Activation Memory Profiler… · p. 38
  • The recomputation-based checkpointing strategy is standard for GPUs, while future work would consider paging to off-device memory for other computing substrates.The profiler’s checkpointing approach recomputes activations as needed during backpropagation.Source: B Activation Memory Profiler… · p. 38
Loading 1904.10631v2…