[논문 리뷰] Training Deep Nets with Sublinear Memory Cost

Tianqi Chen, Bing Xu외 2명|arXiv (Cornell University)|2016. 04. 21.

Advanced Neural Network Applications참고 문헌 19인용 수 538

한 줄 요약

이 논문은 계산 그래프 기반 방법을 제안하여 서브선형 메모리로 심층 네트워크를 학습시키고, n-layer 네트워크에서 메모리를 추가 순방향 패스와의 트레이드로 O(sqrt(n))의 메모리 사용량을 달성하며, 극단적인 경우에는 O(log n)까지 가능하다고 보여주고, 매우 깊은 ResNets 및 긴 시퀀스 LSTM에서 시연한다.

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.

연구 동기 및 목표

• 학습 중에 중간 특성 맵과 기울기를 저장하여 메모리 사용을 감소시킨다.
• inplace 연산과 메모리 공유를 통한 계산 그래프 기반의 메모리 최적화를 개발한다.
• 깊은 네트워크 학습을 가능하게 하는 제어된 계산-메모리 트레이드오프를 도입한다.
• 메모리 최적화 기법을 딥 러닝 프레임워크에 통합하기 위한 실용 가이드라인을 제공한다.

제안 방법

• 중간 결과 간의 inplace 연산과 메모리 공유를 가능하게 하기 위해 계산 그래프를 분석한다.
• 역전파 중에 중단된 중간값을 재계산하기 위한 그래디언트-그래프 구성 방법(미러 카운트)을 개발한다.
• 네트워크를 k세그먼트로 나누고 세그먼트 출력을 저장하여 추가 순방향 패스를 발생시키는 O(sqrt(n)) 메모리 계획을 제시한다.
• 메모리 최적화된 그래디언트 그래프(Alg. 2)와 예산 기반 계획(Alg. 3)으로 임의 그래프에 일반화한다.
• 재귀적 분할(k=1의 경우)로 메모리를 O(log n)까지 줄일 수 있는 재귀 기반 관점을 제시한다.
• 현존 DL 프레임워크에 이 기술들을 구현하기 위한 프레임워크 가이드라인과 오픈 소스 계획을 제공한다.

실험 결과

연구 질문

• RQ1 중간 특성 맵과 그래디언트를 서브선형 메모리로 저장하면서 학습의 정확성을 유지할 수 있는가?
• RQ2 주어진 네트워크 깊이와 메모리 예산에 대해 최적의 메모리-계산 트레이드오프는 무엇인가?
• RQ3 계산 그래프를 분석하여 학습 중 inplace 연산 및 메모리 공유를 가능하게 할 수 있는가?
• RQ4 이 접근법이 매우 깊은 구조(예: 1000개 이상 레이어)와 긴 시퀀스 모델(LSTMs)에도 과도한 오버헤드 없이 확장되는가?
• RQ5 현재 DL 프레임워크에 메모리 최적화를 통합하기 위한 실용 가이드라인은 무엇인가?

주요 결과

• 네트워크 깊이에 대해 메모리 비용을 선형에서 서브선형으로 줄일 수 있으며, 미니배치당 추가 하나의 순방향 패스만으로 n-layer 네트워크에서 O(sqrt(n)) 메모리를 달성한다.
• 극단적 경우, 추가 순방향 계산이 O(n log n) 정도로만 필요해져 메모리 사용을 O(log n)으로 줄일 수 있다.
• 제안된 방법으로 ImageNet에서 1000-layer ResNet의 메모리 사용이 48G에서 7G로 감소한다.
• 긴 시퀀스에서 복잡한 RNN(특히 LSTM)을 학습할 때도 상당한 메모리 절감이 관찰된다.
• 서브선형 메모리 계획은 선형 메모리 할당에 비해 약 30%의 느려짐을 야기하지만, 큰 메모리 절감에 비해 완만한 오버헤드이다.
• 이 접근법은 기존 프레임워크와 호환되며 다른 메모리 최적화와 결합될 수 있다.