sumio

Softmaxとクロスエントロピーの標準的な実装に2種類の数値誤差が存在し、従来のGrokkingの研究で観測されていたSlingshot現象（突然の損失急増）が、この数値誤差によって説明できることが示された。さらに、蒸留で教師と生徒の出力が近い場合にも問題が発生しうることが明らかになった。 “A Fix for Softmax and Cross-Entropy” * Softmax Collapse 通常のクロスエントロピー損失では、正解クラスのLogitに対する勾配は gₜ = qₜ - 1 である。ここでqₜは、Softmaxが正解クラスに割り当てる確率である。 正解クラスに割り当てる確率が大きくなるにつれて勾配は小さくなるが、数学的にはゼロにはならない。しかし、浮動小数点演算ではqₜがちょうど1に丸められ、勾配がゼロになってしまう。 一方、不正解クラスに対する勾配は、より小さな値まで保持される。 本来、すべてのクラスに対するLogitの勾配の総和はゼロ、つまり ∑ₖ gₖ = 0 を満たす必要がある。これは、すべてのクラスのLogitに同じ値を加えてもSoftmaxが変わらないという性質に対応する。しかし、この数値誤差によって勾配の総和がゼロでなくなる。 正解クラスのLogitと、それ以外のLogitのLogSumExpとの差をマージンmと定義すると、FP32ではmが約16.6を超えた時点でこの問題が発生する。 例えば、Logitが [9, -9, -9] であり、最初のクラスが正解の場合、通常実装の勾配は、マージンが18と16.6を超えてしまうので0となり、 [0, 1.523 × 10⁻⁸, 1.523 × 10⁻⁸] となるが、正しい勾配は [-3.046 × 10⁻⁸, 1.523 × 10⁻⁸, 1.523 × 10⁻⁸] である。 この違いは絶対値としては小さいものの、Adamのように勾配の大きさを使って更新量を正規化するOptimizerでは、重大な影響を与えうる。 特にGrokkingでは、ネットワークが訓練データを記憶した後、かなり時間が経ってから汎化性能が急上昇する現象が知られている。その過程では、Lossがゼロ近傍から急激に増大し、再びゼロに戻るSlingshot現象も観測されていた。このSlingshot現象が次の過程で発生することを示している。 1. モデルが訓練データを記憶すると、正解Logitのマージンが増大する。 2. マージンがFP32の約16.6を超えると、正解クラスの勾配がゼロに丸められる。 3. 勾配が極めて小さくなると、Adamの二次モーメント推定値が減衰する。 4. Weight Decayなどによってマージンが再び16.6未満になると、正解クラスの勾配が復活する。 5. 小さくなった二次モーメントによって勾配が正規化されるため、非常に大きなパラメータ更新が発生し、Lossが急増する。 さらに、AdamWのεとWeight Decayが、学習後のLogitのマージンを決めることも分析している。 Weight Decayが働く定常状態では、マージンは近似的に m* ≈ log(1/ε) となる。例えばε = 10⁻⁸の場合、m* ≈ 18.4となり、FP32で数値誤差が発生する16.6を超えてしまう。これもSlingshotが発生する原因となる。 この問題の解決方法は非常に単純であり、クロスエントロピーの計算に必要な L = log(1 + e⁻ᵐ) を、数値的に安定したSoftplus関数 L = softplus(-m) で計算する。 さらに、正解Logitとその他のLogitとの差を使って損失を表現することで、自動微分によって正解クラスの勾配が、その他のクラスの勾配の負の総和として計算されるようにする。 これによって、非常に小さな勾配も保持され、勾配の総和がゼロになる性質も維持できる。 * Softmax Cancellation もう一つの問題は、LLMの蒸留やKL正則化などで発生しうる。 TeacherとStudentの、確率分布をp, q としたとき、Soft-target Cross-EntropyのStudent側のLogitに対する勾配は g = q - p となる。 このときTeacherとStudentが完全に一致すれば、勾配はゼロになるはずである。 しかし、従来の実装ではTeacherとStudentのSoftmaxを別々に計算してから引き算するため、近い値同士の減算による桁落ちが発生する。さらに、前向き計算と後ろ向き計算で演算経路が異なることによって、両者が完全に一致していてもゼロでない勾配が発生する場合がある。 これは、Teacherの重みをコピーしてStudentを初期化する場合などに問題となる。 実際に論文では、Teacherの重みからStudentを初期化して蒸留を続ける実験において、従来実装ではBF16の数値誤差がAdamによって増幅され、StudentがTeacherから乖離して精度が低下することを確認している。 この問題も計算方法の工夫によって解決できる。具体的には、TeacherとStudentの確率を別々に計算して引き算する代わりに、Logitの差から出発し、expm1やlog1pなどの数値的に安定した関数を用いて勾配を計算する。 これによって、TeacherとStudentが一致する場合には勾配が正確にゼロとなる。 コメント ==== 今回指摘された数値誤差の原因となる計算は、PyTorchやJAXを含む現在の標準実装にも残っている。今後、これらの標準実装でも対応が進むと思われるが、大規模LLMの学習で使われるFused Kernelなどについても、同様の問題が存在するかどうかの確認が必要である。 このような一見わずかな数値誤差が、学習性能や学習中の不可解な挙動に大きな影響を与える、数値計算の方法によって勾配が変わり、Optimizerを通じて学習全体の挙動が変化するのは今後注意が必要である。 将来的には、コーディングエージェントが、勾配の総和がゼロになるといった数学的に保たれるべき不変性を理解し、実装の中でそれらが破られる可能性のある箇所を自動的に検出・指摘できるようになることを期待する。