現在のLLMの低精度学習でFlashAttention-3を使った場合、backward計算で数値的な問題が生じる可能性があることが報告されている。 実験では、最初の25Bトークン程度までは正常に学習できていたが、その後、勾配ノルムが1000倍近くまで増大し、最終的な損失も悪化した。 問題の大きかった2層についてattentionのbackward計算だけをFP32に置き換えると、モデル全体の勾配ノルムが2桁以上減少した。つまり、巨大な勾配のかなりの部分が、モデルそのものではなくattention kernelの数値誤差に由来していたことがわかった。 この問題は二つからなる。 1つ目は、FlashAttentionにおけるsoftmaxの計算方法である。FlashAttentionではsoftmax計算時に、scaleとrow maxのsubtractをFMAでまとめて計算する。そのため、本来、最大logitはsoftmax内部で厳密に0になるべきところが、丸め誤差によって0にならない場合がある。このわずかな誤差がbackward時に増幅される。 この問題は、先にrow maxを引いてからscaleするように演算順序を変更することで、その大部分を取り除くことができる。 2つ目の問題は、softmaxのscore gradientが持つ対称性が、丸め誤差によって崩れてしまうことである。 softmaxには、すべてのscoreに同じ定数を加えても出力が変わらない、という対称性がある。 例えばsoftmaxに(s1, s2, s3)というscoreを入力するとき、すべてに同じ定数cを加えて(s1+c, s2+c, s3+c)としても、 p_j = exp(s_j+c) / Σ_l exp(s_l+c) であり、exp(c)は分子と分母で打ち消し合うため、softmaxの出力は変わらない。 したがって、この「全scoreを一様に動かす方向」に対する勾配は0でなければならず、softmaxのscore gradient g_j は Σ_j g_j = 0 という性質を厳密に満たす。 しかし実際のFlashAttentionでは、このgradientをBF16にキャストした後に計算を続けるため、丸め誤差によってこの総和がわずかに0からずれてしまう。 すると、本来は消えるはずのkey全体の平均成分がquery gradientに混入する。 特に学習が進みkey normが大きくなると、このわずかな保存則の破れが大きく増幅され、本来のgradientよりも偽のgradientが支配的になることがある。 さらに、この誤ったgradientによってQ/K normが大きくなり、それによって誤差がさらに増幅されるというフィードバックループが生じる。 この問題への対処は比較的単純で、BF16にキャストした後にscore gradientの総和が再び0になるように射影すればよい。論文ではこれをGProjと呼んでいる。 GProjを適用すると、query/key gradientの誤差をほぼFP32 attentionと同程度まで低減できる。 また、QK normalizationを使っているモデルではQ/K normの増大そのものが抑えられるため、今回のような誤差が大きく増幅される状況には入りにくいと考えられる。ただし、QK normalizationが数値誤差そのものを修正するわけではなく、あくまでその増幅を抑える役割であると考えるのが適切である。 コメント === 今回報告されたような問題は非常に厄介である。学習初期にはほとんど問題が見えず、学習が進んで特定の内部状態に入ったときに徐々に顕在化するためである。 この研究が示しているように、モデル内部にはいくつかの対称性があり、それに対応する保存則がbackward計算にも存在する。 低精度計算によってこうした保存則がわずかに破られると、最初は小さな誤差であっても、学習ダイナミクスとのフィードバックループによって徐々に増幅され、最終的には学習全体を不安定化させる場合がある。 この論文が直接扱っているのはBF16だが、最近ではFP8やFP4など、さらに低精度な数値表現を用いた学習も広がっている。 その際には、どの層やどの演算にどの精度を使うべきかが、現在はかなり実験的に決められている。 今回の研究のように、モデル内部に存在する対称性や、それに対応する保存則を明示的に理解することができれば、単に量子化誤差を小さくするだけでなく、どの誤差は許容でき、どの誤差は構造的に危険なのかを判断できるようになる。 その結果として、より低精度での学習を安全に実現したり、低精度化しながら学習安定性や最終性能を維持・改善したりできる可能性がある。