[논문 리뷰] Low-Memory Neural Network Training: A Technical Report

Nimit S. Sohoni, Christopher R. Aberger외 3명|arXiv (Cornell University)|2019. 04. 24.

Advanced Neural Network Applications참고 문헌 83인용 수 38

한 줄 요약

이 논문은 WideResNet과 DynamicConv Transformer에서 제어된 정확도 손실을 통해 학습 메모리를 줄이기 위해 희소성, 저정밀도, 마이크로배치 및 그래디언트 체크포인팅의 네 가지 기법을 활용하여 학습 메모리 요구량을 정량화하고 평가한다.

ABSTRACT

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.

연구 동기 및 목표

• 신경망 학습에서 관련된 메모리 구성 요소(모델, 옵티마이저, 활성화)를 정량화하고 각 구성의 상대 기여도를 파악한다.
• 네 가지 메모리 저감 기법이 고립된 상태에서 학습 메모리와 모델 정확도에 미치는 영향을 평가한다.
• 이 기법들을 결합했을 때 메모리, 정확도, 계산에 미치는 영향을 탐구한다.
• 대표 벤치마크(WideResNet on CIFAR-10 and DynamicConv on IWSLT'14 De→En)에서 잠재적 메모리 절감을 시연한다.

제안 방법

• 학습 중 모델, 옵티마이저, 활성화 구성 요소로 메모리 사용을 프로파일링한다.
• 동적 희소 재매개변화를 통해 희소성을 평가하고 비영(非영) 부분의 비율에 따라 정확도를 측정한다.
• FP16과 동적 손실 스케일링의 정확도 및 FLOPs에 대한 영향을 통해 저정밀도 학습을 평가한다.
• 마이크로배치 크기를 더 작게 시뮬레이션하고 배치 정규화 고려 사항을 분석하며 정확도를 평가한다.
• 활성화 메모리를 줄이기 위한 그래디언트 체크포인팅 전략을 분석하고 메모리 감소 시 FLOPs 트레이드를 정량화한다.

실험 결과

연구 질문

• RQ1 대표적인 모델 전반에서 지배적인 메모리 구성 요소는 무엇인가?
• RQ2 희소성, 저정밀도, 마이크로배치, 체크포인팅이 각각 메모리와 정확도에 어떤 영향을 미치는가?
• RQ3 이 기법들을 적용할 때 계산상의 트레이드오프는 무엇인가?
• RQ4 주어진 정확도 제약 하에서 메모리 감소를 극대화하기 위해 이 기법들을 어떻게 결합할 수 있는가?
• RQ5 결합된 기법들이 WideResNet 및 DynamicConv Transformer의 훈련에서 상당한 메모리 감소를 가능하게 하는가?

주요 결과

• 활성화 메모리는 연구에서 다룬 설정에서 총 학습 메모리의 지배적 구성요소이다.
• WideResNet의 경우 희소성 70%, DC-Transformer의 경우 60%에서 작은 정확도/ BLEU 손실(각각 0.3% 및 0.8 BLEU)을 보인다.
• FP16 학습은 WideResNet의 정확도에 큰 영향을 미치지 않으며 DC-Transformer에 대해서는 0.15 BLEU의 감소만 유발한다.
• 마이크로배치 크기를 10으로도 WideResNet의 정확도가 유지되지만, 아주 작은 크기일수록 정확도가 감소한다.
• 그래디언트 체크포인팅은 활성화 메모리를 최대 약 5.8배까지 줄일 수 있지만 FLOPs가 증가하며 정확도 손실은 없다.
• 복합 기법은 WideResNet 학습 메모리를 최대 60.7배까지, DC-Transformer은 최대 8.7배까지 각각 0.4% 정확도 손실 및 0.15 BLEU 손실과 함께 줄일 수 있다.