技術概要：ニアロスレス HiF8 W8A8 量子化認識トレーニングのための Max-Window スケール推定

問題定義

リソース制約環境における大規模言語モデル（LLM）の展開には、HiF8（High-precision Float 8）のような 8 ビット浮動小数点（FP8）形式を用いた積極的な量子化が不可欠である。量子化認識トレーニング（QAT）は、事前学習済みモデルを低ビット表現に適応させることを目指すが、交差エントロピー損失のような標準的なトレーニング指標では見えない微妙な失敗モードを導入する。

具体的には、事前学習済みモデルへの QAT の適用は、事前学習された知識を保持しつつ内部表現を量子化ノイズに適応させるという複雑なトレードオフを伴う。著者らは、標準的なトレーニング損失が劣化を警告しないことが多いことを特定し、HiF8 W8A8 QAT において 2 つの明確で直交する失敗モードが生じることを明らかにした：

1. $A_{max}$ 飽和: 遅延テンソルスケーリング（DTS）における時間的遅れに起因する。推定されたスケール因子は過去のデータに基づいているため、現在のステップの真の活性化最大値が予測されたスケールを超えると、値がクリップされ、順伝播活性化が破損する。しかし、ストレートスルー推定器（STE）はクリップが発生しなかったかのように誤って勾配を伝達する。
2. 破滅的忘却: 継続的な事前学習における過度に積極的な学習率に起因し、量子化方式とは無関係に事前学習された常識知識を上書きする。

手法

本研究は、10 億パラメータの OpenPangu-Embedded-1B モデルを HiF8 W8A8 形式で用いて行われた。手法としては、パフォーマンス劣化の原因を解きほぐすための 8 つの制御実験による体系的な調査が含まれる。

量子化設定

• 形式: HiF8 は値の大きさに基づく階層的な仮数精度を使用する（$|x| \le 8$ で 3 ビット、$|x| \le 128$ で 2 ビット、より大きな値で 1 ビット）。
• モジュール選択: 効率性と安定性のバランスを取るため、中間の 21 個のトランスフォーマーブロック内の MLP レイヤー（ゲート、アップ、ダウンの投影）のみを量子化する。アテンション投影、埋め込み、レイヤーノルム、および境界ブロック（最初の 3 ブロックと最後の 2 ブロック）は、外れ値の処理と出力精度を保持するため BF16 のまま残す。
• トレーニングプロトコル: 実験では FineWeb データセットを使用し、グローバルバッチサイズ 1024、シーケンス長 1024、10,000 ステップで行われた。

実験変数

調査の核心は、異なる $A_{max}$ 推定アルゴリズム と 学習率 を比較することにある：

1. $A_{max}$ アルゴリズム:
   • Most-Recent (DTS): 直前のステップの最大値を使用。
   • Exponential Smoothing (DTS): 移動平均を使用。
   • Max-Window (DTS): 64 ステップの履歴ウィンドウ内の最大値を使用。
   • Current Tensor Scaling (CTS): 現在のステップから最大値を計算（2 回目の順伝播が必要）。
2. 学習率: 実験では、BF16 ベースラインに対して当初は安全とみなされた $10^{-4}$ と $10^{-5}$ をテストした。
3. ウォームアップ: 後続の実験では、重みを安定させ、量子化を有効化する前に $A_{max}$ 履歴を代表性のある値で満たすため、500 ステップの BF16 ウォームアップ段階を導入した。

主要な貢献

本論文は、失敗モードの解きほぐしに基づき、4 つの主要な貢献をなしている：

1. 直交する失敗モードの特定: 著者らは、HiF8 QAT におけるベンチマークの劣化が、$A_{max}$ 飽和（量子化固有の問題）と破滅的忘却（トレーニングダイナミクス上の問題）という 2 つの独立した源に起因することを示した。重要なのは、トレーニング損失だけではいずれも検出できず、両方ともダウンストリームの評価でのみ顕在化することである。
2. 保守的な Max-Window 戦略: $A_{max}$ 飽和に対処するため、標準的なアルゴリズムに代わり、64 ステップの履歴ウィンドウにわたる max-アルゴリズム を採用することを提案する。これは、観測ウィンドウ内の任意のステップに対して、スケールに起因する飽和に対するハードな保証を提供する。わずかな過剰な保守性は、有益な正則化として機能する。
3. 2 フェーズトレーニングプロトコル: 破滅的忘却を緩和するため、500 ステップの BF16 ウォームアップ に続き、$10^{-5}$ に低減した学習率で QAT を行うプロトコルを導入する。ウォームアップにより、モデルが安定した重み領域に到達し、$A_{max}$ 履歴に代表性のある値が蓄積され、低い学習率により事前学習された知識の上書きが防止される。
4. 包括的な失敗分析: 本研究は、以前の手法が失敗した理由を記録している：
   • Most-Recent は 1 ステップ遅れの飽和に悩まされる。
   • Exponential Smoothing は非単調な活性化ダイナミクス下では失敗する。
   • CTS は安定した表現を破損させるステップごとの変動を導入する。
   • 高い学習率 ($10^{-4}$) は、量子化の有無にかかわらず、顕著な忘却を引き起こす。

結果

最終的な設定（Max-Window アルゴリズム + 500 ステップのウォームアップ + $lr=10^{-5}$）は、同じ学習率でトレーニングされた対応する BF16 ベースラインと比較して、ニアロスレスな量子化を達成する：

• MMLU: 0.43% の低下（ベースライン 43.36% に対し 43.17%）。
• HellaSwag: 0.58% の低下。
• ARC-Challenge: 0.22% の低下。
• トレーニング損失: 10,000 ステップにわたる HiF8 QAT と BF16 ベースライン間の平均パーセント誤差（APE）はわずか 0.11% である。

注目すべきは、実験によりいずれかの修正単独では不十分であったことが示されたことである。$lr=10^{-4}$ で Max アルゴリズムを使用しても、忘却により ARC タスクで約 3% の低下が生じた。逆に、標準的な most-recent アルゴリズムで $lr=10^{-5}$ を使用しても、飽和に悩まされることになる。

意義と主張

本論文は、事前学習済みモデルに対する成功した HiF8 QAT には、スケール安定性 と 学習ダイナミクス の相互作用を同時に解決する必要があると主張している。この研究の意義は以下の点にある：

• 「不可視」な失敗の診断: 標準的なトレーニング損失が、ダウンストリームベンチマークでのみ現れる飽和と忘却を隠蔽するため、QAT の品質を評価する指標として不十分であることを強調している。
• 堅牢な解決策の提供: 保守的な $A_{max}$ 推定戦略と、低い学習率およびウォームアップを組み合わせることで、10 億パラメータモデルにおいて主要ベンチマークで 1% 未満の低下（ニアロスレス性能）を達成できることを実証している。
• 設計軸の明確化: $A_{max}$ アルゴリズムの選択と学習率は、孤立したハイパーパラメータではなく、相互に依存する重要な設計軸であることを確立している。

著者らは、現在の設定が 10 億パラメータモデルに対して有効である一方で、将来的な研究ではレイヤーごとの学習率、適応的ウォームアップ、高精度レイヤー数の増加などを検討することで、残存するギャップをさらに埋める可能性があると結論づけている。