A Fix for Softmax and Cross-Entropy [1] は、logsoftmax、softmax、交差エントロピーの標準的な実装に 2 つの数値的バグがあると指摘し、それらを修正する実装を提案した論文です。著者は Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen で、2026 年 10 月 22 日に arXiv に投稿されました。

1 つ目は softmax の分母の情報落ち ("Softmax Collapse") です。これにより、学習中に損失が急に跳ね上がったり、テスト精度が急落したりします。
2 つ目は、ソフトラベルの交差エントロピーの勾配を 2 つの確率の引き算で計算する際の丸め誤差のずれ ("Softmax Cancellation") で、予測とソフトラベルが一致していても勾配が 0 にならなくなります。これにより、予測の確率分布がソフトラベルと一致していても、そこからずれていきます。

	バグ1 (Softmax Collapse)	バグ2 (Softmax Cancellation)
原因	情報落ち ($1 + e^{-m}$ が 1 になる)	丸め誤差のずれ ($q$ と $p$ の演算順序の違い)
発生条件	マージン $m$ が大きいとき (モデルの予測確率が正解クラスに偏っているとき)	予測とソフトラベルがほぼ一致するとき
影響範囲	logsoftmax、softmax、交差エントロピー	ソフトラベルの交差エントロピーの勾配

参考文献

以下の内容は arXiv 版 (v1) に基づきます。

文献概要

標準的な実装

JAX、PyTorch、TensorFlow などの深層学習フレームワークでは、logsoftmax、softmax を計算する際、オーバーフローを避けるためにロジットから最大値を引きます。

\tilde z_k = z_k - z_{\max}, \quad
\mathrm{logsoftmax}(z)_k = \tilde z_k - \log \sum_j \exp(\tilde z_j), \quad
\mathrm{softmax}(z)_k = \frac{\exp(\tilde z_k)}{\sum_j \exp(\tilde z_j)}

交差エントロピーは、logsoftmax から $L(z, p) = -\sum_k p_k , \mathrm{logsoftmax}(z)_k$ と出します。

交差エントロピーの勾配は、softmax から $q = \mathrm{softmax}(z)$ として $q - p$ と出します。

ロジットの最大値を引くことでこれらの計算でオーバーフローが発生することは防げます。が、それでも以下の 2 つのバグが残る、というのが著者らの主張です。

【バグ1】 softmax の分母の情報落ち (Softmax Collapse)

※ 「Softmax Collapse」という呼び名は、この論文より前に、Prieto ら [4] が grokking の研究で、softmax の浮動小数点誤差による問題に付けたものです。

標準的な実装では、オーバーフローを避けるためにロジットから最大値を引きます。このとき、モデルの予測確率が正解クラス $t$ に偏っていると、ロジットの最大値は $z_t$ です。すると、最大値を引いた後のロジットは、正解クラスでは $\tilde z_t = z_t - z_t = 0$、ほかのクラスでは $\tilde z_j = z_j - z_t < 0$ となるので、softmax の分母は次のように正解クラスの項 $e^0 = 1$ とそれ以外の項の和に分かれます。

\sum_j \exp(\tilde z_j) = 1 + \sum_{j \neq t} e^{z_j - z_t}

ここで、右辺第 2 項を $e^{-m}$ とおきます。

e^{-m} = \sum_{j \neq t} e^{z_j - z_t}, \quad \text{i.e.} \quad m = -\log \sum_{j \neq t} e^{z_j - z_t}

$m$ は正解クラスのロジットがほかのクラスのロジットをどれだけ上回っているかを表す量 (マージン) で、例えば正解クラス以外のロジットがすべて $z_t - d$ なら、クラス数を $K$ として $m = d - \log(K-1)$ です。これを使うと分母は $1 + e^{-m}$ と書けます。$m$ が大きいと、この足し算は浮動小数点で丸められて 1 になります。丸めが起きるマージンのしきい値は、bf16 (bfloat16) で約 5.5、fp32 (単精度) で約 16.6、fp64 (倍精度) で約 36.7 です。

これが勾配にどう影響するかを見ます。ハードラベルの場合、交差エントロピーの勾配 $q - p$ は、正解クラスでは $q_t - 1$、ほかのクラスでは $q_j$ です。ここで、$q = \mathrm{softmax}(z)$ です。標準的な実装は、$q_t$ と $q_j$ を上の分母を使って次のように計算します。

q_t = \frac{1}{1 + e^{-m}}, \quad q_j = \frac{e^{z_j - z_t}}{1 + e^{-m}} \quad (j \neq t)

分母が 1 に丸められると、$q_t$ は 1 になり、正解クラスの勾配 $q_t - 1$ はちょうど 0 になります。本来の値は $q_t - 1 = -e^{-m}/(1+e^{-m})$ で、小さいながらも 0 ではありません。一方、ほかのクラスの勾配 $q_j$ は、分子 $e^{z_j - z_t}$ が小さい値のまま残るので、指数がアンダーフローするまで正の値を保ちます。その結果、本来 0 になるはずの勾配の総和が 0 になりません。

著者らによれば、この非対称性は、Adam や AdamW のような、勾配の 2 乗の移動平均でパラメータごとに更新幅を調整するオプティマイザで学習するときに、Slingshot [2] (損失の周期的な急上昇) を引き起こします。その仕組みは次の通りです：ミニバッチ内のすべてのサンプルでマージンがしきい値を超えると、パラメータの勾配はほぼ消え、Adam の 2 次モーメント (勾配の 2 乗の移動平均) が 0 に向かって減衰します。その間、weight decay がパラメータを少しずつ縮めます。その後、勾配が戻ってくると、Adam はそれを非常に小さい値で割るため、パラメータが大きく更新されて損失が跳ね上がります。これにより 2 次モーメントがリセットされ、同じことが繰り返されます。
→ 実験では、修正後の実装では Slingshot が起きなくなることが確かめられました。

【バグ1】の修正

【バグ1】の原因は、正解クラスの項 $1$ とそれ以外の項の和 $e^{-m}$ を、$1 + e^{-m}$ という足し算で計算することでした。修正版の実装は、この足し算を行わないようにします。

ハードラベルの交差エントロピーの損失は $L = -\mathrm{logsoftmax}(z)_t$ で、上の分母を使うと次のように書き直せます。

L = \log \sum_j e^{z_j - z_t} = \log\left(1 + e^{-m}\right) = \mathrm{softplus}(-m)

ここで、$\mathrm{softplus}(x) = \log(1 + e^x)$ です。修正版の実装は、この式のとおりに計算します。softplus は $1 + e^x$ を経由せずに計算できるので、$m$ が大きくても情報落ちは起きません。

この修正版の式を自動微分すると、勾配は次のようになります。

\frac{\partial L}{\partial z_t} = -\frac{e^{-m}}{1 + e^{-m}}, \quad \frac{\partial L}{\partial z_j} = \frac{e^{z_j - z_t}}{1 + e^{-m}} \quad (j \neq t)

正解クラスの勾配は、標準的な実装では $q_t - 1$ という引き算で計算していたため 0 になりました。修正版では分子に $e^{-m}$ が残るので、分母 $1 + e^{-m}$ が 1 に丸められても、勾配は本来の値 (約 $-e^{-m}$) になります。その結果、勾配の総和も 0 になります。

この修正は数学的には標準の式と等価です。この修正により、勾配だけでなく元の損失の値も変わる場合があります。

【バグ2】ソフトラベルの交差エントロピーの勾配の丸め誤差のずれ (Softmax Cancellation)

ここでは、ソフトラベル $p$ が、別のロジットのベクトル $z'$ から $p = \mathrm{softmax}(z')$ として与えられる場合を考えます。このとき、交差エントロピーの勾配は $q - p$ です。ここで、$q = \mathrm{softmax}(z)$ です。

$z$ と $z'$ が一致していれば $q$ と $p$ も一致するので、本来、勾配 $q - p$ は 0 です。しかしソフトラベルの交差エントロピーの標準的な実装では、損失の計算では $p$ を softmax の式で求める一方、勾配の計算では $q$ を $\exp(\mathrm{logsoftmax}(z))$ で求めます。$q$ と $p$ はどちらも計算の途中で丸め誤差を含みますが、演算順序が異なるため、その丸め誤差は一致しません。そのため、$z$ と $z'$ が一致していても、計算された $q$ と $p$ は丸め誤差の分だけ食い違います。結果、$q - p$ が 0 になりません。勾配は、fp64 で約 1e-17、fp32 で約 1e-8、bf16 で約 1e-3 と報告されています。よって、bf16、fp32 では Adam の $\epsilon$ を上回ることがあり、パラメータの更新が止まらない恐れがあります。
→ 実際、実験では、標準的な実装で、bf16 で $q$ が $p$ から離れていき、fp32 でも一時的に $q$ と $p$ の KL ダイバージェンスが跳ね上がりました。修正すると、どの精度でも KL ダイバージェンスは 0 のまま保たれました。

【バグ2】の修正

【バグ2】の原因は、ほぼ等しい 2 つの確率 $q$ と $p$ を直接引き算することでした。修正版の実装は、この引き算を行わないようにします。

$q$ と $p$ の比は、ロジットの差 $\delta = z - z'$ だけで書けます。修正版の実装は、$\delta$ から $p$ による加重平均を引いた $d = \delta - \sum_j p_j \delta_j$ を使って、$\log(q_k / p_k)$ と $q_k - p_k$ を次のように計算します。

s_k = d_k - \log\Bigl(1 + \sum_j p_j \bigl(e^{d_j} - 1\bigr)\Bigr) = \log \frac{q_k}{p_k}, \quad
u_k = p_k \bigl(e^{s_k} - 1\bigr) = q_k - p_k

その際、$e^x - 1$ と $\log(1 + x)$ を、$x$ が 0 に近くても精度よく計算できる関数 expm1 と log1p で計算します。$z$ と $z'$ が完全に一致していれば $d = 0$ なので、$s_k = 0$、$u_k = 0$ となり、勾配はちょうど 0 になります。最後に、勾配の総和が 0 になるよう、勾配を $g = u - q \sum_k u_k$ と補正します。

損失の値は標準的な実装のままで、勾配だけを置き換えます。

実験結果

著者らは、grokking [3] (訓練データを記憶しきった後も学習を続けると、ある時点で急にテスト精度が上がる現象) の研究でよく使われる剰余算のタスクで、標準的な実装と修正版の実装を比較しました。このタスクは、【バグ1】が起きやすい状況です。grokking の研究では、モデルが訓練データを記憶した後も長く学習を続けるため、モデルが正解クラスの確率を 1 に近づけようとしてマージン $m$ が大きくなり続けます。また、前述した Slingshot も剰余算での grokking を Adam で調べる中で見つかったもので、Prieto ら [4] も上手く grokking させられないケースを報告しています。著者らは、これらの現象の原因が【バグ1】にあると主張しており、同じ設定で修正版の効果を確かめています。

剰余算のタスクとは、2 つの整数 $a, b$ から $a + b \bmod 31$ などの値を予測する分類問題です。モデルは 1 層の Transformer で、フルバッチ (1 回のパラメータ更新で訓練データ全体を使う方法) の AdamW で学習させています。比較は、数値精度 (bf16、fp32、fp64)、Adam の $\epsilon$、weight decay の有無を変えながら行っています。Adam の $\epsilon$ は、1e-8、1e-12、1e-16 が試されています。

実験で比較対象に使われたのは、jax.nn.log_softmax と jax.nn.softmax、ハードラベルの交差エントロピーとしての optax.softmax_cross_entropy_with_integer_labels、そして jax.nn.log_softmax と jax.nn.softmax を組み合わせて書いたソフトラベルの交差エントロピーです。

論文では、検証したい事柄ごとに実験を行っています。実験ごとの主な結果は以下です。

Slingshot が起きるかの検証: 【バグ1】の修正の効果の検証です。修正後の実装では Slingshot が起きなくなりました。weight decay を用いた設定では、修正後の実装のほうが速く grokking しました。weight decay を用いない設定では、標準実装はマージンが fp32 のしきい値に達した時点でテスト精度が急落しましたが、修正後の実装では急落しませんでした。【バグ1】節で述べた解析及びこの実験結果をもって、著者らは Slingshot は Softmax Collapse によるものだと主張しています。
マージンの理論限界の検証: 【バグ1】に関する理論解析の検証です。著者らは、訓練データを記憶しきった後は、weight decay とパラメータの更新が釣り合うことで、マージンが $m^* \approx \log(1/\epsilon)$ 付近で増えなくなると理論的に導出しました (例えば $\epsilon$ = 1e-8 で約 18.4)。これを確かめるため、マージンがしきい値を超えても正しく計算できる修正版の実装を使い、fp32 で学習させました (この実験のみ、目的上精度が求められるので、オプティマイザに LaProp¹ を使っています)。その結果、マージンは理論どおりの値で増えなくなりました。$\epsilon$ = 1e-8 のときのこの値 (約 18.4) は、fp32 で丸めが起きるしきい値 (約 16.6) をわずかに上回ります。つまり、標準的な実装で学習すると、マージンが【バグ1】の起きる領域に入ってしまいます (この実験で標準的な実装を使って確かめたわけではなく、理論値からの帰結です。著者らは、これが $\epsilon$ = 1e-8 の設定で Slingshot が起きる理由だと説明しています)。
grokking が起きるかの検証: 【バグ1】の修正の効果の検証です。1 層の Transformer を剰余算 ($a + b \bmod 31$) で学習させると、修正後の実装は bf16、fp32、fp64 のどの精度でもテスト精度 99〜100% に達しました。一方、標準実装は bf16 と fp32 で 93〜97% にとどまりました。既存の対策である zero-sum projection (Liu ら) や StableMax (Prieto ら [4]) と比べても、修正後の実装のほうがテスト精度が高くなりました。
予測の確率分布がソフトラベルからずれるかの検証: 【バグ2】の検証です。$a + b \bmod 53$ で学習済みのモデル (教師) と同じ重みで初期化したモデル (生徒) を、教師の出力をソフトラベルとして蒸留すると、修正後の実装では教師との KL ダイバージェンスが 0 のまま保たれました。標準実装では、bf16 で生徒モデルの出力する確率分布が教師モデルの出力する確率分布からずれていきました。

メモ

2 つのバグの原因・発生条件は独立しており、互いに関係していないはずです。
論文の名前「Softmax Cancellation」の cancellation は桁落ちの意味です。論文は、ほぼ等しい確率の引き算で精度が失われることも原因の 1 つとして挙げていますが、この記事では、$z$ と $z'$ が一致していても $q - p$ が 0 にならない直接の原因である丸め誤差のずれとして記述しています。
Prieto ら [4] の Softmax Collapse は、以下のようなものです。
- 【バグ1】と同じ現象を Softmax Collapse と名付け、別の方向から対策した先行研究です。
- Prieto らは、正解クラスのロジットがほかのクラスのロジットより十分大きいときに、浮動小数点の情報落ちで softmax の分母が正解クラスの項だけに等しくなり、損失と正解クラスの勾配が 0 になる現象を Softmax Collapse と定義しています。
- 正則化のない設定で剰余算などのタスクを学習させると、訓練データを記憶した後もロジットが大きくなり続けて Softmax Collapse が起き、学習が止まって grokking が起きなくなると報告しています。
- ロジットが大きくなり続けるのは、訓練データを記憶した後、勾配が予測を変えずにロジットを大きくする方向に揃うためだと分析しています。
- 対策として、softmax を別の関数に置き換える StableMax と、ロジットを大きくするだけの方向の勾配成分を取り除く学習アルゴリズム ⊥Grad を提案しています。
- いずれも関数や学習方法を変えて Softmax Collapse が起きる状況を避けるもので、softmax と交差エントロピーの定義はそのままに計算式の変形だけで情報落ちを解消する今回の論文とは方向性が異なります。
Barron らは、Prieto らに以下のように言及しています。
- Barron らは、Prieto らを、上手く grokking させられないケースを研究した先行研究として挙げています。
- 既存の対策として、代理の損失関数 (Prieto ら)、勾配の射影 (Liu ら)、fp64 への変換を挙げ、自分たちは関数そのものを直す点で異なるとしています。
- StableMax は softmax の $\exp(x)$ を区分的な有理関数に置き換えるもので、指数的ではなく線形に増えるためマージンが 10³〜10⁴ まで膨らみ、どの $\epsilon$ でも修正版よりテスト精度が低いと報告しています。
- 「Softmax Collapse」という呼び名を Prieto らから借りたとは書いておらず、自分たちの Softmax Collapse が Prieto らのものと同じかどうかも述べていません。⊥Grad には触れていません。

Adam のモーメンタムとパラメータごとの更新幅の調整を分離した、Adam の変種のオプティマイザです。Liu Ziyin, Zhikang T. Wang, Masahito Ueda. LaProp: Separating momentum and adaptivity in Adam. arXiv:2002.04839, 2020. https://arxiv.org/abs/2002.04839 ↩