技術要約：LLMのFP4事前学習における収縮バイアス（Shrinkage Bias）の再考

問題提起

大規模言語モデル（LLM）のスケール拡大に伴い、計算コストとメモリコストの削減のために、より低精度なフォーマットへの移行が必要となっている。FP8は広く採用されているが、次世代のハードウェア（NVIDIA Blackwell/Rubin、AMD MI350など）はネイティブなFP4演算をサポートしている。現在のNVFP4のようなE2M1形式（指数部2ビット、仮数部1ビット）を主体としたFP4学習レシピは、アウトライヤー（外れ値）への感度を緩和するためにランダム・ハダマール変換（RHT）を組み合わせて使用している。

しかし、これらのE2M1ベースのレシピは、学習の不安定性とBF16ベースラインに対する大幅な損失（loss）の劣化という問題を抱えている。本論文は、これまで見落とされていた根本的な制限である「収縮バイアス（Shrinkage Bias）」を特定している。これは、E2M1のような非一様量子化フォーマットに固有の、系統的な負の丸め誤差である。E2M1の表現可能なビン（bin）の幾何学的な非対称性により、Round-to-Nearest-Even（RTNE：最近接偶数への丸め）において一貫した「収縮」が発生する。重要な点は、RHTはアウトライヤーを分散させコードブックの利用率を高めるよう設計されているものの、E2M1においては、テンソルの質量を最も非対称な丸めビンへと移動させてしまうことで、結果的にこのバイアスを悪化させてしまうという主張である。これにより、深いネットワークを通じて乗法的な信号減衰が引き起こされる。

手法

1. 収縮バイアスの理論的分析

著者らは、RTNEビンの幾何学的非対称性に起因する、正規化されたマグニチュード空間における負の期待誤差として、収縮バイアスを定式化した。

• 幾何学的起源: E2M1のような非一様グリッドでは、表現可能な値の間隔が変化する（例：0.5から1.0、1.5へ）。レベル $q_i$ において、左側のビン幅 $\ell_i$ と右側のビン幅 $r_i$ はしばしば不一致となる。もし $r_i > \ell_i$ であれば、一様な局所密度下での条件付き期待誤差は負となる（$\frac{\ell_i - r_i}{2} < 0$）。
• 系統的影響: 平均ゼロのノイズとは異なり、このバイアスは行列積（GEMM）における系統的な減衰因子（$\alpha < 1$）として作用する。$K$ 個のレイヤーにわたって、これらの因子は乗法的に蓄積し（$\prod \alpha_k \approx \exp(-\sum \delta_k)$）、系統的な信号減衰を引き起こす。
• RHTとの相互作用: RHTは、テンソルを「ダイナミックレンジ制限型」から「局所解像度制限型」へと変換する。これはバケット利用率を向上させる一方で、データをE2M1の非対称性が最も顕著になる中間マグニチュード領域へと強制的に押し込み、収縮バイアスを増幅させる。

2. UFP4 レシピ

一様グリッドがこの幾何学的誤差を回避できるという知見に基づき、著者らはE1M2/INT4スタイルの一様グリッドを用いた**UFP4（Uniform FP4）**という学習レシピを提案している。

• 一様グリッド: E2M1とは異なり、E1M2グリッドはレベル間の間隔が等しい（$\ell_i = \ell_i$）ため、幾何学的な誤差を排除している。
• フルRHTカバレッジ: 一様グリッドはバイアスを持たないため、UFP4では、フォワード（FPROP）、データ勾配（DGRAD）、ウェイト勾配（WGRAD）の3つすべての学習用GEMMパスに対して安全にRHTを適用できる。従来のE2M1レシピでは、非リーフパスでの誤差蓄積を避けるために、RHTをウェイト勾配パスに限定していた。
• 確率的丸め（Stochastic Rounding: SR）: 勾配の期待値を保持するために、SRはアップストリームの勾配（$dY$）にのみ適用され、それ以外ではRTNEが使用される。

主な貢献

1. 収縮バイアスの特定: 非一様フォーマット（E2M1）が、分布に起因する量子化誤差とは異なる、幾何学的非対称性に由来する系統的な負の丸ing誤差を本質的に抱えていることを形式的に定義し、実証した。
2. 不安定性の統一的説明: 既存のE2M1ベースのレシピにおける学習の不安定性に対し、収縮バイアスがいかに乗法的に蓄積し、RHTによって増幅されるかという理論的および経験的な説明を提供した。
3. UFP4の提案: 一様グリッド（E1M2/INT4）を利用した新しい4ビット学習レシピを提案。これにより、すべての学習用GEMMに対して広範なRHTの適用が可能となり、改善されたバケット利用率をより高い量子化品質へと転換できることを示した。
4. 産業規模での検証: Dense 1.5B、MoE 7.9B、MoE 124Bモデルを用いて検証を行い、一様グリッドが実用的な第一級のFP4学習フォーマットであることを示した。

実験結果

量子化およびGEMM診断

• RHTの効果: アウトライヤーの多いテンソルにおいて、RHTは両方のフォーマットで有効なバケット比を大幅に改善する。しかし、その結果としての信号対量子化ノイズ比（SQNR）は分岐する。E1M2はRHT後にSQNRが向上するのに対し、E2M1は（データが非対称なビンに移動するため）SQNRが低下（負の $\Delta$SQNR）することが多い。
• GEMM出力: 単一GEMMテストにおいて、E1M2はRHT適用後のフォワードおよびバックワードパスの両方で出力SQNRを維持または向上させるが、E2M1は特にバックワードパスにおいて劣化が生じる。

エンドツーエンド事前学習

• 損失の劣化: 長期事前学習タスクにおいて、UFP4（E1M2）は、強力なE2M1ベースラインと比較して、一貫してBF16に対する相対的な損失劣化が低い。
  • Dense 1.5B: 相対誤差は1.2570%（E2M1）から0.9673%（UFP4）へと減少。
  • MoE 7.9B: 2.3596%から1.8469%へと減少。
  • MoE 124B: 1.7308%から1.3863%へと減少。
• スケーリング則: UFP4の優位性は、モデルスケール（10Mから324Mパラメータ）を横断して持続する。適合された損失曲線は、E1M2とBF16のギャップがE2M1よりも小さく、FP4によるペナルティが計算スケールとともに増大しないことを示している。
• アブレーション研究:
  • RHTの範囲: 全てのGEMMに対してフルRHTを適用することでUFP4は最良の結果を得たが、E2M1レシピのようにRHTを制限すると性能が制限される。
  • 範囲制限: E2M1のダイナミックレンジを制限（最大値をキャップするなど）することで一様グリッドを模倣しようとする試みは、ダイナミックレンジを犠牲にするだけであり、バイアスを完全には排除できず、E1M2の性能には及ばなかった。

効率性

• カーネル融合: RHTをFP4量子化と効率的に融合できることを示している。テストされたハードウェア（SM90/SM100）において、融合されたRHT+量子化のレイテンシは、スタンドアロンの量子化に対してわずか1.06倍〜1.07倍であり、非融合の実装（約1.4倍〜1.6倍）と比較して非常に高速である。

意義と主張

本論文は、E2M1を唯一の実行可能なFP4学習フォーマットとして扱う現在の業界標準は、その幾何学的制限により最適ではないと主張している。本研究の意義は以下の点にある：

1. 現状への挑戦: 将来のMLアクセラレータは、E2M1だけに頼るのではなく、E1M2/INT4スタイルの一様4ビットグリッドを第一級の学習プリミティブとしてサポートすべきであることを示唆している。
2. フォーマットとレジームの不一致: RHTのようなアウトライヤー緩和技術の有効性は、基礎となるグリッドの幾何学に依存している。RHTは一様グリッドには有益だが、深い学習レジームにおける非一様グリッドには有害である。
3. 実践的なレシピ: UFP4は、複雑な推定器の修正を必要とせず、単にグリッドの幾何学をRHT後のテンソル分布に合わせるだけで、FP4とBF16の間のギャップを縮小できる、具体的かつスケーラブルなレシピを提供する。

著者らは、E2M1は範囲制限のあるワークロードや推論には適しているものの、学習用ハードウェアは、FP4事前学習のポテンシャルを最大限に引き出すために、一様グリッドをサポートするように進化すべきであると結論付けている。