[論文レビュー] Low-Memory Neural Network Training: A Technical Report

Nimit S. Sohoni, Christopher R. Aberger+3 more|arXiv (Cornell University)|Apr 24, 2019

Advanced Neural Network Applications参考文献 83被引用数 38

ひとこと要約

この論文はトレーニングのメモリ要件を定量化し、広汎な WideResNet と DynamicConv Transformer において、訓練メモリを制御可能な精度損失で削減するための 4 つの手法—スパース性、低精度、マイクロバッチ、勾配チェックポイント—を評価する。

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.

研究の動機と目的

• ニューラルネットワーク訓練に関与するメモリ成分（モデル、最適化器、活性化）を定量化し、それぞれの寄与を特定する。
• 4 つのメモリ削減技法が訓練メモリとモデル精度に与える影響を個別に評価する。
• これらの技法を組み合わせた場合のメモリ、精度、計算量への影響を探る。
• 代表的なベンチマーク（CIFAR-10 上の WideResNet と IWSLT'14 De→En の DynamicConv）で潜在的なメモリ節約を実証する。

提案手法

• 訓練中のモデル、最適化器、活性化成分へのメモリ使用をプロファイリングする。
• 動的スパース再パラメータ化を介してスパース性を評価し、非零要素のパーセンテージを変化させた場合の精度を測定する。
• 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 倍、精度損失 0.4% で削減し、DC-Transformer のメモリを最大 8.7 倍、BLEU 損失 0.15 に抑えることができる。