[Paper Review] Training Deep Nets with Sublinear Memory Cost

Tianqi Chen, Bing Xu+2 more|arXiv (Cornell University)|Apr 21, 2016

Advanced Neural Network Applications19 references538 citations

TL;DR

The paper introduces a computation-graph based method to train deep networks with sublinear memory, achieving O(sqrt(n)) memory for an n-layer network by trading memory for extra forward passes, and even O(log n) in a extreme case, demonstrated on very deep ResNets and long-sequence LSTMs.

ABSTRACT

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.

Motivation & Objective

• Reduce memory usage during training by storing fewer intermediate feature maps and gradients.
• Develop a computation-graph based memory optimization with in-place operations and memory sharing.
• Introduce a controlled compute-for-memory trade-off to enable training of deeper networks.
• Provide practical guidelines for integrating memory optimization techniques into deep learning frameworks.

Proposed method

• Analyze the computation graph to enable in-place operations and memory sharing between intermediate results.
• Develop a gradient-graph construction method (mirror counts) to re-compute dropped intermediates during backpropagation.
• Propose an O(sqrt(n)) memory plan by dividing the network into k segments and storing segment outputs, incurring extra forward passes.
• Generalize the approach to arbitrary graphs via a memory-optimized gradient graph (Alg. 2) with a budget-based planning (Alg. 3).
• Provide a recursion-based view showing how memory can be reduced to O(log n) with recursive subdivision (k=1 case).
• Offer framework guidelines and an open-source plan for implementing these techniques in existing DL frameworks.

Experimental results

Research questions

• RQ1 Can intermediate feature maps and gradients be stored with sublinear memory while maintaining training correctness?
• RQ2 What is the optimal memory-to-computation trade-off for given network depth and memory budget?
• RQ3 How can computation graphs be analyzed to enable in-place operations and memory sharing during training?
• RQ4 Can the approach scale to very deep architectures (e.g., 1000+ layers) and long sequence models (e.g., LSTMs) without prohibitive overhead?
• RQ5 What are practical guidelines for integrating memory optimization into current DL frameworks?

Key findings

• Memory cost can be reduced from linear to sublinear in network depth, achieving O(sqrt(n)) memory for an n-layer network with only an extra forward pass per mini-batch.
• In the extreme case, memory usage can be reduced to O(log n) with as little as O(n log n) extra forward computation.
• Experiments show a 1000-layer ResNet memory reduction from 48G to 7G on ImageNet with the proposed method.
• Significant memory reductions are observed for training complex RNNs on long sequences (LSTMs).
• The sublinear memory plan incurs about a 30% slowdown compared to linear memory allocation, a modest overhead for large memory savings.
• The approach is compatible with existing frameworks and can be combined with other memory optimizations.