[論文レビュー] Training Transformers with 4-bit Integers

Haocheng Xi, Changhao Li+2 more|arXiv (Cornell University)|Jun 21, 2023

Advanced Neural Network Applications被引用数 8

ひとこと要約

この論文は、Hadamard量子化を用いた前向き伝播とビット分割と活用度スコアサンプリングを用いた逆伝播を組み合わせたINT4ベースのトランフォーマー訓練法を提案し、多様なタスクで競争力のある精度とハードウェアに優しいスピードアップを実現する。

ABSTRACT

Quantizing the activation, weight, and gradient to 4-bit is promising to accelerate neural network training. However, existing 4-bit training methods require custom numerical formats which are not supported by contemporary hardware. In this work, we propose a training method for transformers with all matrix multiplications implemented with the INT4 arithmetic. Training with an ultra-low INT4 precision is challenging. To achieve this, we carefully analyze the specific structures of activation and gradients in transformers to propose dedicated quantizers for them. For forward propagation, we identify the challenge of outliers and propose a Hadamard quantizer to suppress the outliers. For backpropagation, we leverage the structural sparsity of gradients by proposing bit splitting and leverage score sampling techniques to quantize gradients accurately. Our algorithm achieves competitive accuracy on a wide range of tasks including natural language understanding, machine translation, and image classification. Unlike previous 4-bit training methods, our algorithm can be implemented on the current generation of GPUs. Our prototypical linear operator implementation is up to 2.2 times faster than the FP16 counterparts and speeds up the training by up to 35.1%.

研究の動機と目的

• 超低精度(INT4)でトランスフォーマーをエンドツーエンド訓練し、計算を加速・メモリ使用量を削減する動機づけ。
• Hadamardベースの変換で活性化の外れ値を抑制する前向き量子化を開発。
• 勾配の構造的スパース性を活用して、ビット分割(BS)と活用度スコアサンプリング(LSS)によりバックプロパゲーションを効率的に量子化。
• 現在のGPUと互換性のある実用的なINT4 MMベース訓練パイプラインを提供。
• NLP、翻訳、ビジョンのベンチマークで競争力のある精度と速度アップを実証。

提案手法

• トランスフォーマーの全ての線形演算を行列乗算(MM)として定式化し、LSQベースの量子化でINT4演算を適用。
• HQ量子化(HQ)を導入し、ブロック対角Hadamard行列で活性化を量子化前に変換して前向きのアクティベーション外れ値を抑制。
• バックプロパゲーションで勾配の構造的スパース性を活用し、勾配を4ビットの上位/下位部に分割するビット分割(BS)と、情報量の多い行を選択するLSSを用いてMM計算を実行。
• activations/weightsの表現をINT4のままに保ちつつ、各MMごとに勾配を動的に量子化し、最終結果を必要に応じてFP16蓄積のためにデ量子化。
• 2つのMM手順を提供：前向き用HQ-MMと後向き用LSS-MMを組み合わせ、INT4 MMベースの訓練をエンドツーエンドで可能に。
• NLPベンチマーク(GLUE、SQuAD、CoNLL、MT)と画像分類(ViT、ImageNet)でFP、INT8、FP4、LSQ+LUQのベースラインと比較評価。

実験結果

リサーチクエスチョン

• RQ1 標準的なハードウェアのみを用いて特殊な数値形式を必要とせず、4ビット整数演算でトランスフォーマーをエンドツーエンド訓練できるか。
• RQ2 超低精度訓練で精度を維持するために、前方の活性化外れ値と後方の勾配のスパース性をどのように対処すべきか。
• RQ3 Hadamardベースの前方量子化とLSSベースの逆伝播の精度・効率トレードオフはどうなるか。

主な発見

• HQ+LSSは、NLP、MT、ビジョンタスクにおいてFPおよび既存の4ビット訓練法と比較して競争力のある精度を達成。
• GLUEおよびSQuADスタイルの評価で、HQ+LSSはLSQ+LUQを上回り、特にBert-baseおよびBert-largeの変種で優位。
• 機械翻訳(WMT)はHQ+LSSで約1.0% BLEUの低下だが、超低端の設定よりは良く、他と同等程度でハードウェアに優しい。
• ViTモデルを用いた画像分類では、完全精度のベースラインに比べて小さな精度低下(多くは1-2%未満)で、HQ+LSSは複数のケースでLSQ+LUQを上回る。
• 本手法は現在のGPUでINT4 MM実装を可能にし、FP16 MMより最大2.2x速く、訓練速度向上は最大35.1%を実現。
• CoLAのアブレーション研究は、前方Hadamard量子化の精度維持の重要性を裏付け、バックプロパゲーションはBS+LSSの恩恵を受けることを確認。