[論文レビュー] ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

Jianfei Chen, Lianmin Zheng+5 more|arXiv (Cornell University)|Apr 29, 2021

Advanced Memory and Neural Computing参考文献 43被引用数 8

ひとこと要約

ActNNは、ランダム量子化を用いて活性化を2ビット精度に圧縮することで、メモリ効率の高い学習フレームワークを提案する。これにより、精度損失が最小限であるにもかかわらず、顕著なメモリ削減が可能となる。理論的収束保証を提供し、異種性に配慮したミックス・プレシジョン量子化を採用することで、活性化メモリを12倍削減し、PyTorchにおいて最小限のコード変更で最大14倍のバッチサイズを可能にする。

ABSTRACT

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12x, and it enables training with a 6.6x to 14x larger batch size.

研究の動機と目的

• GPUメモリの増加が限定的であるにもかかわらず、大規模ニューラルネットワークの学習における増大するメモリボトルネックに対処すること。
• 多様なアーキテクチャにわたる活性化圧縮学習（ACT）に対して、一般的な理論的収束保証を提供すること。
• 特徴次元、サンプル、層ごとの活性化の非均一性を活用する量子化戦略を設計し、勾配分散を最小限に抑えること。
• PyTorchのような既存のディープラーニングフレームワークへの実用的導入を可能とし、最小限のコード変更で実現すること。
• 画像認識タスクにおいて、精度損失が最小限の高圧縮比（例：2ビット）を達成すること。

提案手法

• ActNNにおける量子化勾配を、完全精度勾配の不偏推定値として扱い、理論的収束を保証する。
• 勾配分散の正確な式を導出し、量子化が学習安定性に与える影響を定量化する。
• 特徴グループごとの活性化統計に応じて量子化レベルを適応的に変更する、グループ別量子化器を導入する。
• 固定メモリ予算下で勾配分散を最小限に抑える、細粒度なミックス・プレシジョン量子化戦略を開発する。
• PyTorchライブラリとしてActNNを実装し、シームレス統合が可能なドロップイン置換レイヤー（例：actnn.Conv2d）を提供する。
• メモリ節約と学習効率のバランスを取るために、動的量子化戦略チューニングを適用する。

実験結果

リサーチクエスチョン

• RQ1 一般のニューラルネットワークアーキテクチャに対して、活性化圧縮学習が理論的に収束することが保証可能か？
• RQ2 量子化が勾配分散に与える影響は何か？ そして、これを正確にモデル化することで、量子化設計を支援できるか？
• RQ3 活性化の非均一性を活用するミックス・プレシジョン量子化戦略は、より高い圧縮比と低い精度損失を達成できるか？
• RQ4 2ビット活性化圧縮は、セグメンテーションや検出を含む多様なビジョンタスクで効果的か？
• RQ5 固定メモリ制約下で、ActNNはモデルの深さ、幅、解像度をどの程度スケーリングできるか？

主な発見

• ActNNは、画像分類、検出、セグメンテーションタスクにおいて、平均で2ビットの活性化圧縮を実現し、精度損失が0.5％未満に抑えられる。
• 同じGPU上で、バッチサイズを6.6倍から14倍まで拡大可能となり、学習スループットが顕著に向上する。
• 同じメモリ予算下で、ResNet-152を6.4倍深く、3.7倍広く、3.1倍高い解像度にスケーリング可能であり、元の学習スループットの64％～155％を維持する。
• ActNNは、活性化メモリフットプリントを12倍削減し、1つのGPUで高解像度セグメンテーションや検出のための大バッチ学習を可能にする。
• 自己教師あり学習では、2ビット活性化を用いたActNNが、1024バッチサイズでImageNetトップ-1精度72.65％を達成し、完全精度ベースラインと同等またはそれを上回る。
• 1.25ビットの活性化でさえ、ActNNは収束し、妥当な結果を生成する。これは、先行研究が4ビット圧縮に限定されていたのに対し、優れた性能を発揮する。