Source-linked AI summary
Training Deep Nets with Sublinear Memory Cost
Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin
TL;DR
Deep neural network training stores feature maps and gradients whose memory cost scales with network depth. The paper introduces computation–memory trade-offs that achieve sublinear memory use, including O(√n) memory with one extra forward pass and a 30% runtime increase in experiments.
Problem
Feature maps and gradients scale linearly with network depth, limiting memory efficiency as deep architectures grow.
Method
The paper combines computation-graph analysis, in-place operations, memory sharing, and recomputation to trade additional computation for lower training memory.
Results
The approach trains n-layer networks with O(√n) feature-map memory and one additional forward pass, while a 1,000-layer ResNet uses less than 7GB with 30% additional runtime.
Takeaways & Limitations
The method enables training deeper convolutional networks and recurrent networks with longer unrolling steps, while also supporting larger batch sizes and data parallelism in some settings.
Abstract
from arXiv · showhide
We propose a systematic approach to reduce the memory consumption of deep neural network training. Specifically, we design an algorithm that costs O(sqrt(n)) memory to train a n layer network, with only the computational cost of an extra forward pass per mini-batch. As many of the state-of-the-art models hit the upper bound of the GPU memory, our algorithm allows deeper and more complex models to be explored, and helps advance the innovations in deep learning research. We focus on reducing the memory cost to store the intermediate feature maps and gradients during training. Computation graph analysis is used for automatic in-place operation and memory sharing optimizations. We show that it is possible to trade computation for memory - giving a more memory efficient training algorithm with a little extra computation cost. In the extreme case, our analysis also shows that the memory consumption can be reduced to O(log n) with as little as O(n log n) extra cost for forward computation. Our experiments show that we can reduce the memory cost of a 1,000-layer deep residual network from 48G to 7G with only 30 percent additional running time cost on ImageNet problems. Similarly, significant memory cost reduction is observed in training complex recurrent neural networks on very long sequences.
1 Introduction
The paper targets the linear memory cost of storing intermediate feature maps and gradients in increasingly deep networks. It introduces computation–memory trade-offs that enable sublinear memory training while supporting deeper models and longer recurrent sequences.
- The proposed method combines automatic in-place and memory-sharing optimizations with a computation-for-memory trade-off.
- O(√n) memory trains an n-layer network with only double the forward-pass computational cost.
- O(log n) feature-map memory is possible in the extreme case, with O(n log n) extra forward computation.
- Intermediate feature maps and gradients scale linearly with network depth, constraining training of deeper architectures.
- The approach supports deeper convolutional networks and recurrent networks with longer unrolling steps.
2 Related Works
The work builds on computation-graph liveness analysis and prior memory–computation trade-offs, extending them into a general methodology for deep neural networks. It is designed to complement existing memory optimizations and orthogonal large-model training approaches.
- Existing frameworks analyze computation graphs for memory or computation optimization but generally do not address the computation–memory trade-off.
- The paper generalizes dropping intermediate results, also known as gradient checkpointing, to convolutional and recurrent neural networks.
- The proposed automatic planning algorithm provides memory plans for general deep networks with sublinear memory cost.
- CPU/GPU swapping and model parallel training are orthogonal approaches that can be combined with the proposed algorithm.
3 Memory Optimization with Computation Graph
The method represents training as a computation graph and uses dependency, liveness, in-place, and memory-sharing analyses to reduce storage for intermediate results. Explicit backward pathways and automated allocation make these optimizations systematic rather than ad hoc.
- A computation graph represents operations as nodes and dependencies as edges, with an explicit backward pathway for gradient calculation.
- Memory allocation targets intermediate outputs and gradients because these nodes consume a large proportion of training memory.
- In-place operation overwrites an input with output, while memory sharing recycles storage from intermediate results that are no longer needed.
- In-place updates are safe only when an input has no other pending use, and shared memory requires non-overlapping node lifetimes.
- The allocation algorithm declares minimal gradient dependencies and applies liveness analysis to enable memory sharing.
- Liveness analysis can reduce prediction memory from O(n) to nearly O(1) through sharing, while training gains are limited to a constant factor.
4 Trade Computation for Memory
The paper reduces training memory by dropping intermediate results and recomputing them during backpropagation, using graph-based planning to balance storage and computation. Segmenting a network yields O(√n) memory with one extra forward pass, while recursive application can reach O(log n) memory at higher computational cost.
- 4.3 An O(√n) Memory Cost Algorithm: O(n) intermediate-result memory is reduced by dividing the network into segments, retaining segment outputs and recomputing dropped results during backpropagation.The approach generalizes from linear chains to computation graphs through automatic gradient-graph construction.
- 4.1 General Methodology: The mirror count function m(v) specifies how many times a node may be recomputed, with m(v)=0 retaining its output and positive values enabling duplication.The construction algorithm also produces a traversal order that supports memory optimization and runtime control-flow dependencies.
- 4.3 An O(√n) Memory Cost Algorithm: Low-cost operations such as batch normalization, activation, and pooling can be dropped while retaining convolution outputs to save memory with little computation overhead.This optimization targets common Conv-BatchNorm-Activation pipelines.
- 4.3 An O(√n) Memory Cost Algorithm: O(2√n) memory results from setting k=√n segments, combining O(n/k) within-segment memory with storage for intermediate segment outputs.The algorithm requires one additional forward pass during training.
- 4.3 An O(√n) Memory Cost Algorithm: Greedy budgeted planning searches allocation plans that balance memory for inter-stage outputs against computation within each segment.Varying the budget produces alternative plans, and static allocation provides their exact memory costs.
- 4.4 More General View: Recursion and Subroutine: O(log2 n) memory is possible when the recursive scheme stores k=1 intermediate result, at the cost of O(log2 n) forward-pass computation.The recursion treats each segment as a bulk operator and applies the optimization within its internal subgraph.
5 Experiments
Experiments show that graph optimizations reduce memory substantially, while sublinear planning enables much deeper networks and long-sequence recurrent models with modest runtime overhead.
- Evaluation setup: The evaluation reports feature-map memory separately from parameter and temporal memory.Static allocation enables exact feature-map accounting, while runtime measurements capture total memory cost.
- Deep convolutional networks: 2–3× lower memory results from system optimizations, but memory still scales linearly with ResNet depth.After optimization, the best available GPU could train only a 200-layer ResNet.
- Deep convolutional networks: <7GB of GPU memory trains a 1,000-layer ResNet using the proposed sublinear plan.The method trades additional computation for a sublinear memory trend.
- LSTM for long sequences: >4× memory reduction is achieved on the long-unrolled LSTM compared with the optimized non-trading plan.In-place optimization helps by adding weight gradients directly to a single memory cell instead of allocating gradients at every timestamp.
- Training speed: 30% additional runtime cost is incurred by sublinear allocation compared with the normal strategy.The overhead is attributed to the double forward cost in gradient calculation.
6 Conclusion
The paper combines graph-based memory optimizations with computation–memory trading to reduce training memory for deep neural networks.
- Conclusion: O(√n) memory trains an n-layer network with no more than one extra forward computation per mini-batch.The approach targets intermediate feature maps and gradients using liveness analysis, in-place operations, and memory sharing.
A Search over Budget B
The budget-search procedure uses approximate plan costs to construct and refine a grid of candidate memory budgets.
- Budget search: A grid search over B selects a good memory plan after approximate plan generation and exact memory-cost calculation.The algorithm first estimates a budget from B = √xy, then refines the solution around that value.
- Budget initialization: B = √xy estimates each stage’s memory cost from inter-stage storage and per-stage execution costs.The search first runs the allocation algorithm with B = 0, then uses the resulting x and y values to set the initial budget.