[論文レビュー] Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction

Georgii Novikov, Daniel Bershatsky+4 more|arXiv (Cornell University)|Feb 1, 2022

Advanced Neural Network Applications被引用数 5

ひとこと要約

この論文では、勾配の量子化を用いて、ディープラーニングの学習中にメモリ使用量を削減するための、点関数の活性化関数の勾配を区分的定数近似で量子化する手法を提案している。バックプロパゲーションの際、フル精度の入力を保持する代わりに3ビットのインデックスを保存することで、ベンチマーク上で精度の低下なしに最大40％のメモリ削減を達成している。動的計画法を用いて最適な量子化スキームを計算し、ReLU、GELU、Swishなどの一般的な活性化関数の即時置き換えが可能である。

ABSTRACT

Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, one needs to store the input to each operation in the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which -- as we show -- can be significantly reduced by quantization of the gradients. We propose a systematic approach to compute optimal quantization of the retained gradients of the pointwise nonlinear functions with only a few bits per each element. We show that such approximation can be achieved by computing optimal piecewise-constant approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.

研究の動機と目的

• バックプロパゲーション中に保存される活性化テンソルのメモリ使用量を削減すること。
• 現代のアーキテクチャにおいて、点関数非線形関数のフル精度入力を保存する際の見過ごされがちなメモリコストに対処すること。
• データに依存しない、即時置き換え可能な活性化関数のための方法を開発し、トレーニングパイプラインを変更せずに低ビットの勾配保存を可能にすること。
• 多様なアーキテクチャとタスクにおいて、モデルの収束性と精度を維持したまま顕著なメモリ削減を達成すること。
• 微調整やアーキテクチャ変更を必要とせず、Transformerを含むあらゆるモデルに適用可能な汎用的なフレームワークを提供すること。

提案手法

• 勾配関数の導関数の最適な区分的定数近似を計算するために動的計画法を用い、量子化誤差を最小化する。
• バックワードパスにおいて、フル精度の活性化入力を、バケットインデックス（例：3ビット）のみに置き換えることで、1テンソル要素あたりのメモリ使用量を著しく削減する。
• 元の活性化入力ではなく、量子化済みの導関数値を保存することで、劣化した圧縮を実現し、勾配への影響は最小限に抑える。
• PyTorchパイプラインにシームレスに統合可能な、一般的な活性化関数（ReLU、GELU、Swish、Sigmoid）の即時置き換え実装を提供する。
• データ分布やモデル構造に依存せず、Transformer や EfficientNet を含むあらゆるモデルに普遍的に適用可能である。
• 最小限のパフォーマンスオーバーヘッドで、効率的な推論と学習を実現するCUDA最適化カーネルを採用する。

実験結果

リサーチクエスチョン

• RQ1 点関数活性化関数の勾配を少数ビットに量子化することで、バックプロパゲーション中のメモリ使用量を削減できるか、かつモデルの精度が低下しないか？
• RQ2 低ビット表現に最適な活性化関数の導関数の区分的定数近似は何か？ これは量子化誤差を最小化する。
• RQ3 3ビットまたは4ビットの活性化勾配量子化を用いることで、Transformer や EfficientNet、GPT-2 などの多様なモデルで実際にどれほどのメモリ削減が達成できるか？
• RQ4 提案手法は、異なるモデルやトレーニング段階（例：事前学習モデルの微調整）においても即時置き換えとして適用可能か？
• RQ5 データに依存しない設計であるにもかかわらず、異なるデータ分布やモデルアーキテクチャにおいても本手法は有効であるか？

主な発見

• GPT2 や RoBERTa-large などのモデルでは、1活性化要素あたり3ビットインデックスを保存することで、最大40％のメモリ削減を達成している。
• 3ビット量子化を用いることで、ResNet、EfficientNet、BERTベースのモデルを含む複数のベンチマークで平均して20％のメモリ削減が達成された。
• 本手法は、ResNet-50、EfficientNet-B7、RoBERTa などのすべての評価対象モデルにおいて、フル精度学習と同等の訓練収束性とテスト精度を維持している。
• 勾配の導関数の最適な区分的定数近似は動的計画法により計算され、ReLU、GELU、Swish、Sigmoid などのすべてのテスト対象活性化関数で有効である。
• 既存のトレーニングパイプラインと互換性があり、アーキテクチャ変更なしに、微調整段階を含むあらゆる段階で適用可能である。
• CUDA実装により低オーバーヘッドが達成されており、実際のトレーニングワークフローへの実用的導入が可能で、パフォーマンスへの影響は最小限に抑えられている。