[論文レビュー] Stable and low-precision training for large-scale vision-language models

Mitchell Wortsman, Tim Dettmers+4 more|arXiv (Cornell University)|Apr 25, 2023

Multimodal Machine Learning Applications被引用数 5

ひとこと要約

本稿では、大規模なビジョン・ランゲージモデルの安定的で効率的な8ビット（int8およびfloat8）学習を実現する新しい線形層「SwitchBack」を提案する。CLIP ViT-Hugeにおいて、bfloat16学習とほぼ同一の精度を達成しつつ、13〜25％の高速化を実現した。また、更新のクリッピングを適用することで損失の急増を解消するAdamW-AdaFactorハイブリッド手法「StableAdamW」を提案し、大規模スケールでの学習安定性を著しく向上させた。

ABSTRACT

We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) For acceleration, we introduce SwitchBack, a linear layer for int8 quantized training which provides a speed-up of 13-25% while matching the performance of bfloat16 training within 0.1 percentage points for the 1B parameter CLIP ViT-Huge -- the largest int8 training to date. Our main focus is int8 as GPU support for float8 is rare, though we also analyze float8 training through simulation. While SwitchBack proves effective for float8, we show that standard techniques are also successful if the network is trained and initialized so that large feature magnitudes are discouraged, which we accomplish via layer-scale initialized with zeros. 2) For stability, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become under-estimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid which avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping at the scales we test.

研究の動機と目的

• 精度を損なわず、低精度（int8およびfloat8）計算を用いて大規模ビジョン・ランゲージモデルの学習を高速化すること。
• 損失の急増がモデル性能を低下させることを特定・是正することで、学習の安定性を向上させること。
• 10億パラメータのCLIP ViT-Hugeモデルにおける、初めての成功したint8学習を実現し、bfloat16性能と0.1％以内の差で一致させること。
• 将来の低精度学習研究のための実用的でオープンソースのツール（例：Tritonカーネル）を開発すること。
• 適切な初期化と適応的最適化が、ネイティブハードウェアサポートがなくても安定したfloat8学習を可能にすることを示すこと。

提案手法

• SwitchBackは、重み勾配の計算に16ビット行列乗算を用い、前方伝搬および入力勾配計算にはint8を適用することで、高次元のアテンション層における量子化ノイズを低減する。
• 本手法はハイブリッド精度戦略を採用：推論および入力勾配には8ビット、勾配蓄積には16ビットを適用し、標準的な量子化手法よりも精度を向上させる。
• float8学習のため、正確なfp8エミュレーションを用いてSwitchBackをシミュレートし、スケールが大きくなると発散するテンソル単位の量子化ベースラインと比較して優れた性能を示した。
• StableAdamWは、AdamWとAdaFactorの更新クリッピング技術を統合し、二乗勾配とその2次モーメント推定値の比を監視することで、動的学習率調整を実現する。
• 損失スケールを変更し、テンソルレベルでInf/NaNをチェックし、固定することで、孤立したレイヤーの不安定性が原因で発生するグローバルな学習失敗を防止する。
• ゼロで初期化されたレイヤー・スケール初期化を用い、大きな特徴量の大きさを抑制することで、標準的な量子化手法を用いた安定なfloat8学習を可能にする。

実験結果

リサーチクエスチョン

• RQ1 SwitchBackのような新しい線形層設計が、大規模ビジョン・ランゲージモデルのint8学習において高い精度を達成しつつ、顕著な高速化を実現できるか？
• RQ2 大規模CLIP学習中に損失の急増が発生する理由は何か？ そして、これを予測または防止できるか？
• RQ3 ネイティブハードウェアサポートがなくても、float8学習を安定化できるか？ どのような初期化や最適化手法がこれを可能にするか？
• RQ4 AdamWとAdaFactorの更新クリッピングの組み合わせ（StableAdamW）は、勾配クリッピングと比較して、安定性および最終的な性能で優れているか？
• RQ5 大きな活性化値と勾配が、低精度学習中に不安定性を引き起こす要因として果たす役割は何か？

主な発見

• SwitchBackは、bfloat16と比較してCLIP ViT-Hugeの学習において13〜25％のエンドツーエンドの高速化を達成し、ゼロショットImageNet精度は0.1パーセンテージポイント以内の差で一致した。
• SwitchBackは、10億パラメータのCLIP ViT-Hugeモデルにおける、これまで報告された中で最大のint8学習を実現し、LLM.int8()を精度面で上回った。
• 標準的なテンソル単位の量子化を用いたfloat8学習は、4億2000万パラメータを超えると発散するが、SwitchBackとゼロで初期化されたレイヤー・スケール初期化を組み合わせることで、安定した学習が可能になった。
• 更新クリッピングを備えたAdamW-AdaFactorハイブリッド手法StableAdamWは、損失の急増を完全に解消し、スケールが大きい状況でも勾配クリッピングを上回るゼロショットImageNet精度を達成した。
• 損失の急増は、常にパッチ埋め込み層における二乗勾配とその移動平均の比の低下に先行しており、2次モーメント推定値による過小評価が原因であると示唆された。
• テンソルレベルでInf/NaNをチェックし、固定された損失スケールを用いることで、グローバルな学習失敗を防止でき、ViT-Hugeスケールでの安定したfp16ミックス精度学習が可能になった。