Training Deep Nets with Sublinear Memory Cost

TL;DR

This paper introduces a novel algorithm that significantly reduces memory usage during deep neural network training, enabling deeper models with minimal extra computation, thus advancing deep learning capabilities.

Contribution

The authors present a systematic method to train deep networks with sublinear memory cost, using in-place operations and memory sharing to optimize memory efficiency.

Findings

Reduced memory for a 1000-layer residual network from 48G to 7G

Achieved memory reduction with only 30% increase in training time

Enabled training of complex RNNs on long sequences

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.