FP4による大規模言語モデルの事前学習について、既存手法を体系的に比較した論文が公開されている。 “What Matters for NVFP4 Training - A Comparative Study of Low-Precision Pretraining Recipes” FP4を使った学習は高速な行列演算や必要なメモリ帯域幅を抑えられる一方、BF16と比較すると精度低下や学習の不安定性が残る。 FP4のE2M1形式では、0、0.5、1、1.5、2、3、4、6と、それらの負の値からなる15種類の数値しか表現できない。ただし、NVFP4では16個の値からなるブロックごとにスケール係数を持ち、それを掛けることで表現範囲を調整する。 今回の論文では、8BのDenseモデルと30B-A3BのMoEモデルを最大1Tトークン学習し、既存のFP4学習手法を比較している。 特に提案するMosaic NVFP4では、以下の工夫を組み合わせている。 ・Forwardでは重みを1次元で量子化し、Backwardでは必要な方向に再量子化する。 ・Backwardの2種類の行列積にRandom Hadamard Transformを適用し、外れ値の影響を軽減する。 ・勾配にはMS-EDEN（後述）を使い、Stochastic Roundingより小さな量子化誤差で不偏性を実現する。 ・その他のテンソルにはFour Over Sixによる適応的なスケール選択を適用し、量子化誤差を削減する。 ・最後の15%の層はBF16のまま維持する。 これらを組み合わせることで、30B-A3Bモデルでは1Tトークン学習後のBF16に対する相対学習損失差を0.81%未満まで縮小し、従来のFP4手法から0.36パーセンテージポイント改善した。 特に厄介なのは、すべての層をFP4にすると、8Bモデルでも300Bトークンを超えたところで学習が発散したことである。また、手法間の優劣は150Bトークンを超えるまで安定せず、短期間の小規模実験だけではFP4学習手法を適切に評価できないことも示された。 ・MS-EDEN 勾配の量子化において、バイアスをゼロに保ちながら量子化誤差の分散を小さくする手法。 個々の値を確率的に丸めるStochastic Roundingとは異なり、まずベクトル全体にRandom Hadamard Transformを適用して値を混合し、FP4に量子化する。その後、補正係数を掛けることで不偏性を回復する。 この補正係数はNVFP4のブロックスケールに組み込まれる。スケールはFP8で格納されるため、補正後のスケールを量子化する際にはStochastic Roundingを使用する。 これにより、各要素を確率的に丸める場合よりも小さな量子化誤差で、不偏な勾配推定を実現できる。 ・Four Over Six FP4のE2M1形式では、表現可能な値の4と6の間に5が存在しないため、この領域で量子化誤差が大きくなりやすい。 通常のスケーリングでは、ブロック内の絶対値の最大値をFP4の6に対応させる。しかし、最大値を4に設定した方が、ブロック全体の量子化誤差が小さくなる場合がある。 Four Over Sixは、ブロックごとに最大値を4に対応させる場合と6に対応させる場合の両方を試し、量子化後の二乗誤差が小さい方を選択する。 コメント ==== FP4は行列積の演算性能を向上させ、必要なメモリ帯域幅も削減できるため有望だが、大規模な事前学習での実績はまだ限られている。 今回の研究でも、かなり学習が進んだ後に発散する現象が確認された。最後の15%の層をBF16に維持することで発散を抑えられ、この条件ではBackwardの量子化手法による性能差もかなり小さくなる。 これは、Backwardの量子化バイアスを抑えるだけでは解決できない問題が残っていることを示している。著者らは量子化による表現精度の問題を示唆しているが、発散の詳しい機構はまだ分かっていない。 学習が進むにつれて小さな数値誤差が蓄積し、学習ダイナミクスが変化する可能性も考えられる。Forward側の表現誤差や、量子化とパラメータ更新の相互作用も重要だと思われる。 これはFP4に限らず、誤差逆伝播法を近似したり、一部を別の計算方式に置き換えたりする手法にも共通する課題である。小規模な実験では問題なく動作していても、モデルを大きくしたり学習を長期間続けたりすると、発散や性能劣化が生じる場合が多い。 深層学習はかなり粗い数値計算でも動作する印象があるが、特に長期間の学習における勾配計算や数値誤差の扱いは、想像以上に繊細である。