TL;DR
標準のsoftmax/log-softmax/交差エントロピーには2つの数値バグがある。提案実装は最大値やターゲットを和から除外し、蒸留時はq-pをログ空間で直接計算する。これによりSlingshotを消し、重み減衰ありのgrokkingを速め、重み減衰なしの精度崩壊と蒸留時のドリフトを防ぐ。重み減衰時はロジットマージンがlog(1/ε)付近で平衡する。
既存手法との違い
原論文は、JAX/PyTorch/TensorFlow等で一般的な「最大ロジットを引いてsoftmax/交差エントロピーを計算する」標準実装を比較対象にする。標準実装は、ターゲット確率が高いと1+e^{-m}のe^{-m}が丸められてターゲット勾配が0になるSoftmax Collapseと、蒸留でq=softmax(zS)とp=softmax(zT)を別々に計算してq-pを引くため近接時に精度が落ちるSoftmax Cancellationを起こす。原論文は緩和策として代理損失・勾配射影・fp64キャストにも触れるが、提案はlog-softmax/softmax/交差エントロピー自体を修正する。比較条件はbf16/fp32/fp64、ε∈{10^-8,10^-12,10^-16}、modular addition等のgrokking実験と蒸留の数値例で、標準は精度クリフ（bf16≈5.5、fp32≈16.6、fp64≈36.7）で不正確になる一方、提案はその範囲でも勾配が和ゼロを保つと報告する。
使える条件
対象はsoftmax/log-softmax、hard-target cross-entropy、soft-target cross-entropy（知識蒸留・KL正則化）の実装。入力はKクラスのロジットz、hard target t、またはstudent/teacher logits zS,zT。浮動小数点bf16/fp32/fp64を想定し、勾配は自動微分で得る。原論文の評価はJAX実装、Adam/AdamW（およびRMSProp/LaPropの平衡点）を使うgrokking実験（1層幅64 ReLU Transformer、a+b mod 31等）と蒸留の数値例が中心。大規模LLM、RAG、他フレームワーク、実運用データでの有効性は提示本文からは確認できない。
限界と未検証の範囲
原論文の評価は小規模なmodular addition/grokkingと数値例が中心で、大規模モデル・多様なタスク・RAG/LLM運用への一般化は提示本文では未検証。提示本文はSection 6の途中までで、全ベンチマーク結果や計算コスト・メモリ影響、他言語・他フレームワークでの挙動は確認できない。Soft-target cross-entropyの修正はstop_gradientでカスタム勾配を入れるJAX実装例であり、他フレームワークへの移植性や数値的な限界は本文だけでは判断できない。重み減衰時のマージン平衡m*≈log(1/ε)はAdamW等の正規化勾配とεに依存する解析で、ε=0や異なる条件下の一般化には注意が必要。
解説
ねえ智也、softmaxと交差エントロピーに数値バグがあるって論文、どういうこと？ 普通に使ってる実装が壊れてるの？
壊れてるというか、特定の条件で勾配が変になる。まずSoftmax Collapse。正解クラスに自信がありすぎてマージンが精度の崖を超えると、標準実装は正解ロジットの勾配を0に丸める。でも他のクラスの勾配は残るから、勾配の和が0にならなくなる。
え、和が0じゃなくなるの気持ち悪い。それがSlingshotのlossスパイクにつながるの？
そう。Adam系は勾配の大きさで正規化するから、片方だけ0になると不安定になる。論文ではSlingshotを消すって報告してる。
もう一つのバグは？
Softmax Cancellation。蒸留でstudentとteacherのロジットが近いとき、q=softmax(zS)とp=softmax(zT)を別々に計算してq-pを引くと、ほぼ同じ確率の引き算で精度が落ちる。zS=zTでも非ゼロの勾配が残る。
近いほど正確にしたいのに、そこで壊れるんだ。
直し方はシンプル。log-softmaxとsoftmaxでは最大ロジットを和から除外して、softplusでまとめる。hard-targetの交差エントロピーはターゲットを和から外してmを計算する。
最大値やターゲットを「足さない」って発想なんだ。
soft-targetはq-pをログ空間で直接計算する。δ=zS-zTをpで中心化して、expm1とlog1pで扱う。最後に勾配の和が0になるように補正する。JAXの実装例はstop_gradientで前向きのlossとカスタム勾配を分けてる。
評価は何をやってるの？
modular additionのgrokking。1層幅64のReLU Transformerで、bf16/fp32/fp64、εを1e-8,1e-12,1e-16に変えて、標準と修正版の交差エントロピーを比較してる。重み減衰ありとなしも見てる。
結果は？
修正版は精度の崖を滑らかに超えて、重み減衰ありのgrokkingが速くなる。重み減衰なしでは標準だと精度が61%から23%に落ちる場面で、修正版は85%まで続く。蒸留でもstudentがteacherからドリフトしない。
すごいじゃん。で、なんで重み減衰ありだとマージンがlog(1/ε)付近で止まるの？
それは数値精度の崖じゃなくて、重み減衰とAdamWの正規化勾配の平衡。勾配がe^{-m}で小さくなると、εがあるせいで正規化ステップも小さくなる。だからm*≈log(1/ε)で止まる。ε=1e-8なら約18.4。
なるほど、崖とは別の話なんだ。比較対象は標準実装だけ？
標準はJAX/PyTorch/TensorFlowみたいな最大値シフトの実装。論文は緩和策として代理損失、勾配射影、fp64キャストにも触れるけど、提案はlog-softmax/softmax/交差エントロピー自体を直す。
使える条件は？
Kクラスのロジット、hard targetかstudent/teacherのsoft target、bf16/fp32/fp64で自動微分する想定。評価はJAX実装の小規模なgrokkingと蒸留の数値例が中心。
じゃあ大規模LLMとかRAGでも効くの？
提示された本文からは確認できない。大規模モデルや多様なタスク、他フレームワークへの移植、計算コストやメモリの影響も未検証。Section 6の途中までしか載ってないから、全ベンチマーク結果も判断できない。
わかった。数値の崖を直すだけじゃなくて、平衡点の理解まで話が広がるのが面白いね。じゃあ私のレポートの数値バグもついでに直してほしいな。
それはsoftmaxと関係ないだろ。