Softmax と Cross-Entropy に潜む

論文が狙うのは、ほとんどの実務者が目を向けない数値計算の配管部分、すなわちSoftmaxとCross-Entropyの受け渡しであり、そこでは勾配が「確率マイナス正解クラスの指示関数」に簡略化される。著者らは、指数化の前に最大ロジットを引くという標準的な手法が、それ自体は数値的に安全であっても、確信度が精度依存のしきい値を超えると丸め誤差が正解クラスの勾配を消す経路を生むことを示した。修正版は、log-sum-expを正解以外のロジットのみで計算し、損失をsoftplusで包むことで、「1足す微小な数」をそもそも形成しない。

Macでの再現は、バグが理論上のものではないという独立した証拠を加える。CPUとApple GPU（MPS）上のPyTorchで、標準の交差エントロピーは確信度の高い例に対して勾配の合計がゼロではなく+3.0e-08となった一方、修正版は正確にゼロを返した。生徒と教師が同一の重みで開始する蒸留の設定では、標準実装はfp64でも3ステップ以内に乖離したが、修正版はそのまま留まった。

実務上の結論は損失をどこで計算するかで決まる。Hugging Face transformersは損失の前にロジットをfp32へキャストするため、再現では事前学習済みLLMで崖を超えるトークンの割合は小さかった。しかしbf16のロジットをそのまま損失に渡すと、コードデータ上の8Bモデルでは最大47%のトークンが崖の向こう側に位置し、正解クラスの勾配がゼロに丸められていた。パイプラインが影響を受けるかどうかを決めるのは、バグの存在自体よりもこの違いだろう。