はじめに
AI の学習で、いちばん基本的な計算といえば Softmax（ソフトマックス） と Cross-Entropy（交差エントロピー） です。画像分類でも、LLM の次トークン予測でも、最後はこの 2 つで「正解をどれだけ当てられたか」を測り、学習の方向（勾配）を決めています。
2026 年 10 月 22 日、Google DeepMind の Jonathan T. Barron さん・Dan Gnanapragasam さん・Michael Tschannen さんが、この基本中の基本に 2 つの数値計算のバグ があるとする論文「A Fix for Softmax and Cross-Entropy」を arXiv に公開しました。PyTorch・JAX・TensorFlow など、どのフレームワークの標準実装にも当てはまる話です。
「標準の関数にバグ」と聞くと大ごとに聞こえますが、数式が間違っているわけではありません。数式は正しいのに、コンピューターの小数の丸め方のせいで、特定の条件で勾配がおかしくなる、という種類の問題です。
この記事では、論文の中身をかみくだいて解説したうえで、手元の Mac（Apple M5 Max）で次の 4 つを実際に確かめます。
- 2 つのバグが PyTorch（CPU・Apple GPU の MPS）と JAX で本当に起きるか
- 論文が示した「損失が周期的に跳ねる現象（スリングショット）」と、修正による解消
- 蒸留で「動かないはずの生徒が教師から離れる」現象
- 実務で気になる「LLM の学習や LoRA にどれくらい関係するか」
先に結論をまとめておきます。
- 2 つのバグは、手元の PyTorch 2.14.1（CPU・MPS）と JAX でそのまま再現しました。修正版は数行の関数で、勾配が正しい値になります。
- スリングショットも再現しました。標準の損失では 3 万ステップの学習中に損失が平均 21.3 回跳ね、修正版では 0 回です。ただし論文の「グロッキング（遅れてやってくる汎化）が速くなる」は、手元の fp32 の実験では再現せず、損失を bf16 で計算した場合にだけ修正版がはっきり上回りました。
- 蒸留では、標準の実装だと fp64 でも 3 ステップ目に生徒が教師から離れ始めました。修正版は勾配がちょうど 0 なので、1 歩も動きません。
- LLM では、損失を fp32 に上げて計算する通常の設定なら、崖を越えるトークンは Qwen3-0.6B で 0〜0.03%、8B モデルで 0.2〜0.5% でした。LoRA で同じデータを 600 ステップ回しても跳ねは出ず、損失もほぼ同じです。
つまり、「入れれば LLM の性能が上がる魔法」ではなく、条件がそろったときに学習を壊さないための安全装置、というのが手元で確かめた結論です。どんな条件なら効くのかを、順番に見ていきます。
この図から分かること：バグ 1 は「確信が強すぎるとき」、バグ 2 は「蒸留で生徒と教師が近いとき」に起きる。どちらも計算の順番を変えるだけで直る
この記事の確認範囲
| 項目 | 内容 |
|---|---|
| 論文 | arXiv:2610.31642 v1（2026-10-22 投稿、CC BY 4.0、10 ページ）。査読・採択の記載はなし |
| 公式コード | 論文中に JAX 実装（約 30 行）が載っている。論文にリポジトリへのリンクはない |
| 確認日 | 2026 年 10 月 26 日 |
| 手元の環境 | Apple M5 Max（メモリ 128GB）、macOS 26.6.2、Python 3.12.15、PyTorch 2.14.1、JAX 0.12.0.1、optax 0.2.8、transformers 5.20.2、peft 0.21.2 |
| 確かめたこと | 2 つのバグの再現、精度の崖の測定、グロッキングと蒸留の小さな再現実験、学習済み LLM のマージン分布、Qwen3-0.6B の LoRA 学習、速度 |
| 確かめていないこと | NVIDIA GPU（CUDA）での挙動、大規模な事前学習、論文と同一のモデル構成（細部が書かれていないため筆者の推定で組んだ） |
想定している読者は、PyTorch でモデルを学習・微調整したことがある方です。ただ、数式は最小限にして図で説明するので、「Softmax は名前だけ知っている」という方でも読み進められるように書いています。
前提：Softmax と Cross-Entropy をおさらい
まず、バグの舞台になる計算をおさらいします。下の図で、①から④の順に追ってみてください。
この図から分かること：勾配は「確率 − 正解の印」で、合計は数式の上ではいつも 0。ロジットマージン m は確信の強さを表す
モデルの最後の層が出す生の点数を ロジット（logit） と呼びます。Softmax はロジットを「合計 1 の確率」に変換し、Cross-Entropy は「正解の確率の −log」を損失にします。正解の確率が 1 に近いほど、損失は 0 に近づきます。
この記事で何度も出てくるのが ロジットマージン m です。論文の定義は「正解のロジット − log Σ（正解以外のロジットの e 乗）」で、ざっくり言えば モデルが正解をどれだけ強く確信しているか を表す数です。m が大きくなるほど、正解の勾配はおよそ −e−m まで小さくなります。m = 18 なら約 −1.5×10−8 で、とても小さいですが 0 ではありません。
もう 1 つ大事なのが、実装上の工夫です。ez をそのまま計算すると、fp32 ではロジットが約 88 を超えたところで桁あふれ（オーバーフロー）します。そこで、どのフレームワークも 最大のロジットを先に引いてから e を計算しています。結果の確率は変わらないので、教科書どおりの安全な書き方です。実は、この「最大値を引く」工夫こそが、バグ 1 の入り口になります。
浮動小数点の「1 の近くは粗い」問題
コンピューターの小数（浮動小数点数）は、有効数字の桁数が決まっています。fp32 ならおよそ 7 桁です。ここから、少し意外な性質が生まれます。
この図から分かること：1 にごく小さい数を足すと、その数は丸めで消える。消える境目を、マージンの単位で表したのが「精度の崖」
0 の近くの数は、fp32 でも 10−38 程度まで有効数字を保って表せます。ところが 1 の近くでは目盛りが粗く、1.0 の次に大きい数は 1.00000012 です。その間の数は表せないので、1.0 + 0.00000003 を計算すると 1.0 に戻ってしまいます。0.00000003 単体なら問題なく表せるのに、1 に足した瞬間に消えるわけです。
論文は、この「1 に足すと消える」境目をマージン m の単位で表し、精度の崖（precision cliff） と呼んでいます。有効ビット数を p とすると崖は ln 2p で、次のようになります。
| 形式 | 有効ビット | 精度の崖 c | 正解の確率でいうと |
|---|---|---|---|
| bf16 | 8 | 約 5.55 | 約 99.6% を超えたあたり |
| fp32 | 24 | 約 16.64 | 約 99.999994% を超えたあたり |
| fp64 | 53 | 約 36.74 | 9 が 15 個以上並ぶあたり |
論文が書いている 5.5・16.6・36.7 と一致します。bf16 は LLM の学習でよく使う 16 ビット形式ですが、有効ビットが少ないぶん、正解の確率が 99.6% を超えたあたりで、もう崖に着いてしまう 点に注目してください。
バグ 1：Softmax Collapse（確信が強すぎると勾配が消える）
何が起きるのか
論文の例で、実際に追いかけてみます。ロジットが [9, −9, −9]、正解が 1 番目だとします。マージンは約 17.3 で、fp32 の崖（16.6）をわずかに越えています。
![Softmax Collapseが起きる流れ。z=[9,−9,−9]から最大値を引くと[0,−18,−18]、分母は1+0.0000000305が1.0に丸まり、正解の確率は1.0、損失は0.0、正解の勾配は1.0−1=0.0。ほかの勾配は1.5e-08のまま残り、合計が+3.05e-08になる](https://qiita-user-contents.imgix.net/https%3A%2F%2Ftakuya-genai.com%2Fwp-content%2Fuploads%2F2026%2F10%2Fsmx-fig04-collapse.gif?ixlib=rb-4.1.1&auto=format&gif-q=60&q=75&s=c7bf35390a3c8615c54789da84de36bc)
この図から分かること：正解の勾配だけが「1 から 1 を引く」計算で消え、ほかは残る。合計が 0 でなくなり、一方向にだけ押す勾配が生まれる
手順を言葉でも確認します。
- 最大値（9）を引くと [0, −18, −18] になります。正解の位置はちょうど 0 です。
- 分母は e0 + e−18 + e−18 = 1 + 0.0000000305。ここで 1 に小さい数を足している ので、fp32 では 1.0 に丸まります。
- 正解の確率は 1 ÷ 1.0 = 1.0、損失は −log 1 = 0.0 です。
- 勾配は「確率 − 正解の印」なので、正解は 1.0 − 1 = 0.0。一方、ほかのクラスは e−18 ÷ 1.0 = 1.5×10−8 で、0 の近くで計算されるので 生き残ります。
本来の勾配は「正解を押し上げる力（−3.05×10−8）」と「ほかを押し下げる力（+1.52×10−8 × 2）」が釣り合って、合計がちょうど 0 になるはずでした。ところが押し上げる力だけが消え、ほかを押し下げる力だけが残ります。これが Softmax Collapse です。
崖を越えたこと自体よりも、「片側だけ消える」非対称さが問題です。この食い違いが、あとで説明する Adam との組み合わせで、学習を周期的に揺さぶります。
直し方：「1 + ごく小さい数」を一度も作らない
論文の修正は、とてもシンプルです。
この図から分かること：1 を足すのは softplus の中だけ。正解のロジットは「ほかとの差」の形でしか出てこないので、勾配の合計は自動的に 0 になる
ポイントは 2 つです。
- 分母の「1（= e0）」を和に入れない。正解以外のロジットだけで log-sum-exp を取り、マージン m を直接作ります。和は 0 の近くの数どうしなので、丸めで消えません。
- 損失は softplus(−m) = log(1 + e−m) でまとめます。softplus は、m が大きいときに e−m とほぼ同じ小さな値を正確に返すように作られた関数なので、損失が 0 に潰れません。
さらに、正解のロジットは「ほかのロジットとの差」の形でしか式に出てきません。そのため自動微分は、正解の勾配を「ほかの勾配の合計のマイナス」として計算します。勾配の合計が、計算の形からして 0 になります（手元の例ではちょうど 0.0）。
論文は cross_entropy に加えて、log_softmax と softmax も同じ考え方で直した JAX 実装を載せています（最大のロジットを「ちょうど 0」に置き、残りでマージンを作る）。この記事では、これを PyTorch に移植して使います。
準備：Mac に検証環境を作る
ここからは手を動かします。Python 3.12 の仮想環境に、PyTorch と JAX を入れます。筆者は uv を使いましたが、python -m venv でも同じです。
# 作業フォルダと仮想環境を作る
mkdir -p softmax-fix-lab/scripts && cd softmax-fix-lab
uv venv --python 3.12 .venv
source .venv/bin/activate
# PyTorch・JAX（CPU 版）・optax、LLM の確認用に transformers と peft
uv pip install torch "jax[cpu]" optax numpy transformers peft pandas pyarrow
次に、論文の修正版を PyTorch に移植したファイルを scripts/fixed_torch.py として保存します。論文の JAX 実装と同じ計算をしています。ignore_index（LLM の学習で使う「損失に含めないトークン」の指定）と reduction は、F.cross_entropy と置き換えやすいように筆者が足しました。
"""論文 "A Fix for Softmax and Cross-Entropy" (arXiv:2610.31642) の修正版を PyTorch に移植したもの。"""
import torch
import torch.nn.functional as F
def _softplus(x: torch.Tensor) -> torch.Tensor:
# log(1 + e^x) を桁落ちなく計算する。MPS の F.softplus は x < -17 付近で 0 を返すため使わない
return torch.logaddexp(x, torch.zeros_like(x))
def log_softmax(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
i_max = z.argmax(dim, keepdim=True)
z_max = z.gather(dim, i_max) # 最大のロジット
is_max = torch.zeros_like(z, dtype=torch.bool).scatter_(dim, i_max, True)
shifted = torch.where(is_max, torch.zeros_like(z), z - z_max) # 最大の位置はちょうど 0
others = shifted.masked_fill(is_max, float("-inf")) # 最大を和から外す
m = -torch.logsumexp(others, dim, keepdim=True) # ロジットマージン
return shifted - _softplus(-m)
def softmax(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
return log_softmax(z, dim).exp()
def cross_entropy(z: torch.Tensor, t: torch.Tensor, ignore_index: int = -100,
reduction: str = "mean") -> torch.Tensor:
"""z: (..., K) のロジット、t: (...) の正解番号。F.cross_entropy の置き換え用。"""
valid = t != ignore_index
t_safe = torch.where(valid, t, torch.zeros_like(t))
z_t = z.gather(-1, t_safe.unsqueeze(-1)) # 正解のロジット
is_t = torch.zeros_like(z, dtype=torch.bool).scatter_(-1, t_safe.unsqueeze(-1), True)
others = (z - z_t).masked_fill(is_t, float("-inf")) # 正解を和から外す
m = -torch.logsumexp(others, -1) # ロジットマージン
loss = _softplus(-m) * valid # -log_softmax(z)_t
if reduction == "none":
return loss
if reduction == "sum":
return loss.sum()
return loss.sum() / valid.sum().clamp(min=1)
def soft_target_cross_entropy(z_S: torch.Tensor, z_T: torch.Tensor) -> torch.Tensor:
"""生徒ロジット z_S と教師ロジット z_T の soft-target CE（1 行ごと）。勾配は q - p を log 空間で計算。"""
z_T = z_T.detach()
log_q = log_softmax(z_S)
q, p = log_q.exp(), softmax(z_T)
with torch.no_grad():
delta = z_S - z_T # ロジットの差
d = delta - (p * delta).sum(-1, keepdim=True) # p で中心化
s = d - torch.log1p((p * torch.expm1(d)).sum(-1, keepdim=True)) # log(q / p)
u = p * torch.expm1(s) # q - p
g = u - q * u.sum(-1, keepdim=True) # 合計 0 に揃える
loss = -(p * log_q.detach()).sum(-1) # 順伝播の損失
return loss + (g * (z_S - z_S.detach())).sum(-1) # 勾配だけ g に差し替え
_softplus だけは論文と違う書き方をしています。論文は jax.nn.softplus を使っていますが、PyTorch の F.softplus を MPS（Apple GPU）で使うと、入力が −17 より小さいあたりで 0 を返してしまいました（CPU では正確です）。修正版の損失がまさにその範囲を使うので、torch.logaddexp(x, 0) に置き換えています。詳しくは最後のトラブルシューティングで触れます。
最後の soft_target_cross_entropy は、後半で扱うバグ 2 の修正版です。
再現 1：2 つのバグと「崖」の位置を測る
論文の例を PyTorch で再現する
まずは、論文の 2 つの数値例を PyTorch で確かめます。次のスクリプトを scripts/demo_bugs.py として保存してください。
"""2 つのバグを、PyTorch の標準関数と修正版で見比べる（CPU と、使えれば MPS）。"""
import torch, torch.nn.functional as F
import fixed_torch as fx
devices = ["cpu"] + (["mps"] if torch.backends.mps.is_available() else [])
for dev in devices:
print(f"===== {dev} =====")
# バグ 1: 確信が強すぎると、正解の勾配が 0 に丸まる
t = torch.tensor([0], device=dev)
for name, ce in [("標準", F.cross_entropy), ("修正版", fx.cross_entropy)]:
z = torch.tensor([[9.0, -9.0, -9.0]], device=dev, requires_grad=True)
loss = ce(z, t)
loss.backward()
g = z.grad[0]
print(f"[Collapse] {name:4s} loss={loss.item():.3e} grad={[f'{v:+.3e}' for v in g.tolist()]} sum={g.sum().item():+.1e}")
# バグ 2: 生徒 = 教師 なのに、勾配が 0 にならない
z_T = torch.tensor([[2.0, 1.0, 0.0]], device=dev)
for name in ["標準", "修正版"]:
z_S = z_T.clone().requires_grad_()
if name == "標準":
loss = -(F.softmax(z_T, -1) * F.log_softmax(z_S, -1)).sum()
else:
loss = fx.soft_target_cross_entropy(z_S, z_T).sum()
loss.backward()
print(f"[Cancel] {name:4s} grad at z_S==z_T = {[f'{v:+.3e}' for v in z_S.grad[0].tolist()]}")
scripts フォルダに移動して実行します。
cd scripts
python demo_bugs.py
手元の M5 Max での出力です。
===== cpu =====
[Collapse] 標準 loss=0.000e+00 grad=['+0.000e+00', '+1.523e-08', '+1.523e-08'] sum=+3.0e-08
[Collapse] 修正版 loss=3.046e-08 grad=['-3.046e-08', '+1.523e-08', '+1.523e-08'] sum=+0.0e+00
[Cancel] 標準 grad at z_S==z_T = ['+5.960e-08', '+1.490e-08', '+7.451e-09']
[Cancel] 修正版 grad at z_S==z_T = ['+0.000e+00', '+0.000e+00', '+0.000e+00']
===== mps =====
[Collapse] 標準 loss=-0.000e+00 grad=['+0.000e+00', '+1.523e-08', '+1.523e-08'] sum=+3.0e-08
[Collapse] 修正版 loss=3.046e-08 grad=['-3.046e-08', '+1.523e-08', '+1.523e-08'] sum=+0.0e+00
[Cancel] 標準 grad at z_S==z_T = ['+0.000e+00', '-1.490e-08', '-1.490e-08']
[Cancel] 修正版 grad at z_S==z_T = ['+0.000e+00', '+0.000e+00', '+0.000e+00']
論文が JAX で示した値（損失 0.0 と 3.046e-8、勾配 1.523e-8 など）と、PyTorch でも同じ値になりました。標準の実装は損失も正解の勾配も 0 で、勾配の合計が +3.0e-08 とずれています。修正版は損失が正しい値になり、合計がちょうど 0 です。
バグ 2 の行も見てください。生徒と教師のロジットがまったく同じなので、本来の勾配は 0 のはずです。標準の実装は 10−8 程度の勾配を返し、しかも CPU と MPS で違う値 になりました。CPU と GPU では計算の実装が違い、丸め誤差の出方も違うためです。修正版はどちらでもちょうど 0 です。
念のため、論文の JAX コードもそのまま実行し、論文に書かれた出力（[−3.046e-08, 1.523e-08, 1.523e-08] など）と一致することを確認しています。
崖の位置を測る
次に、精度の崖が本当に理論値の位置にあるのかを測ります。ロジットを [m, 0] にすると、マージンがちょうど m になります。m を 0 から 100 まで 0.01 刻みで動かし、正解の勾配が 0 に落ちる位置を探しました。中心部分だけ抜き出すと、次のようなコードです。
import torch, torch.nn.functional as F
import fixed_torch as fx
m = torch.arange(0, 100, 0.01, dtype=torch.float64)
for dtype in [torch.bfloat16, torch.float32, torch.float64]:
for name, ce in [("標準", F.cross_entropy), ("修正版", fx.cross_entropy)]:
z = torch.stack([m, torch.zeros_like(m)], 1).to(dtype).requires_grad_()
t = torch.zeros(len(m), dtype=torch.long)
ce(z, t, reduction="sum").backward()
zero = z.grad[:, 0] == 0 # 正解の勾配がちょうど 0 になった行
m_real = (z[:, 0] - z[:, 1]).detach().float() # 丸めた後の実際のマージン
cliff = m_real[zero].min().item() if zero.any() else None
print(f"{str(dtype):15s} {name}: 勾配が 0 になる最小のマージン = {cliff}")
CPU での出力です。
torch.bfloat16 標準: 勾配が 0 になる最小のマージン = 5.5625
torch.bfloat16 修正版: 勾配が 0 になる最小のマージン = 89.0
torch.float32 標準: 勾配が 0 になる最小のマージン = 16.639999389648438
torch.float32 修正版: 勾配が 0 になる最小のマージン = 88.7300033569336
torch.float64 標準: 勾配が 0 になる最小のマージン = 36.7400016784668
torch.float64 修正版: 勾配が 0 になる最小のマージン = None
MPS（device="mps" を足したもの）や、bf16 を fp32 に上げてから渡す場合も含めて測った結果が下の図です。
この図から分かること：崖は理論値どおりの位置にある。ただし「bf16 の崖 5.5」は bf16 のまま計算した場合だけで、環境によって変わる
崖は理論値の 5.55・16.64・36.74 とぴったり一致しました。修正版は、fp32・bf16 ではマージン 87〜89 まで勾配が残ります。そこで 0 になるのは、e−m がその形式で表せる最小の数を下回る（アンダーフロー）ためで、丸めの崖とは別の話です。
意外だったのは bf16 の崖が環境で違う ことです。CPU で bf16 のまま計算すると論文どおり 5.56 ですが、MPS では bf16 を渡しても崖が 16.75 でした。MPS の cross_entropy は、内部で fp32 に上げて計算しているようです。また、bf16 のロジットを .float() で fp32 に上げてから渡すと、どの環境でも崖は 16.64〜16.75 になります。
この点は実務に直結します。Hugging Face transformers の因果 LM（GPT 型）の損失関数は、計算の前にロジットを fp32 に上げています（transformers 5.20.2 の loss/loss_utils.py に logits = logits.float() があります）。つまり、bf16 で学習していても、transformers 経由なら崖は 5.5 ではなく 16.6 です。「bf16 だから崖が 5.5」と決めつけず、自分の環境で測るのが確実です。
スリングショット：Adam と組むと損失が跳ねる理由
ここまでで「崖を越えると、正解の勾配だけが消える」ことが分かりました。では、それが学習にどう響くのでしょうか。鍵は、最適化手法の Adam（アダム） です。
この図から分かること：勾配が消えている間に Adam の分母が縮み、勾配が戻った瞬間に巨大な 1 歩が出る。これが周期的な損失の跳ね
Adam は、各パラメーターの 1 歩を「勾配の平均 ÷（勾配の大きさの平均 + ε）」で決めます。勾配が大きいパラメーターも小さいパラメーターも、だいたい同じ歩幅で進めるようにする仕組みです。ε（イプシロン）は 0 で割らないための小さな数で、既定値は 1e-8 です。
論文の説明を 4 段階にまとめると、次のようになります。
- 学習データを丸暗記するとマージンが伸び、崖を越えた例は正解の勾配が 0 になります。
- 全部の例が崖の向こうへ行くと、勾配がほぼ消えます。このとき Adam の分母（勾配の大きさの平均）がどんどん縮みます。
- weight decay（重みを少しずつ 0 に引き戻す正則化）などで一部の例が崖の手前に戻ると、正解の勾配が急に復活します。
- 復活した勾配を、縮みきった分母で割るので、1 歩が大きくなりすぎて損失が急上昇します。マージンが崩れて崖の手前に戻り、また 1. から繰り返します。
この現象は 2022 年に Thilak らが「スリングショット・メカニズム（Slingshot Mechanism）」として報告していたもので、論文はその正体が Softmax Collapse だと説明しています。修正版では正解の勾配が崖で消えないので、Adam の分子と分母が一緒になめらかに縮み、跳ねが起きません。
なぜ既定の ε = 1e-8 で起きやすいのか
論文には、もう 1 つ面白い分析があります。weight decay がある学習では、マージンは log(1/ε) あたりで伸びが止まる というものです。
この図から分かること：ε = 1e-8 の釣り合い 18.4 は、fp32 の崖 16.6 のすぐ向こう。そこへ向かう途中で、標準の実装は必ず崖を踏む
勾配 |g| ≈ e−m が ε より大きい間、Adam の 1 歩はほぼ 1 です。|g| が ε を下回ると、1 歩は e−m ÷ ε に縮みます。縮み始める位置が m ≈ log(1/ε) で、そこで weight decay と釣り合います。
ε = 1e-8 なら釣り合いは 18.4 で、fp32 の崖 16.6 のすぐ向こうです。だから既定の設定のまま fp32 で丸暗記させると、ちょうど崖を踏みます。逆に ε を 1e-7 以上にすると釣り合いが 16.1 以下になり、崖の手前で止まるので、標準の実装でも跳ねないと論文は述べています。
再現 2：グロッキングでスリングショットを起こす
論文の実験は、グロッキング（grokking） と呼ばれる現象の研究で使われる小さな課題です。グロッキングとは、学習データを丸暗記した後、かなり経ってから急にテストの正解率が上がる（汎化する）現象のことです。課題は「a + b を 31 で割った余り」を当てる足し算で、マージンが大きくなりやすいので Softmax Collapse の観察にぴったりです。
論文 6 章の設定をもとに、PyTorch で次のように組みました。
| 項目 | 論文 | 手元の再現 |
|---|---|---|
| 課題 | a + b mod 31（半分を学習に使う） | 同じ |
| モデル | 1 層・幅 64 の ReLU Transformer | 同じ方針（ヘッド数 4・MLP 幅 256・LayerNorm なしは筆者の推定） |
| 最適化 | 全バッチ AdamW、学習率 1e-3、β 0.9 / 0.98、weight decay 1 | 同じ |
| 長さ | 10,000 エポック（全バッチなので 1 エポック＝1 ステップ） | 30,000 ステップ（10,000 では手元のモデルの多くがまだ汎化しなかったため延長） |
| 比べる条件 | 損失（標準 / 修正版）× 精度 × ε | 同じ考え方。精度は「損失を計算する精度」だけを変えた |
| 乱数の種 | 5 | fp32・bf16 は 15、ほかは 2〜3 |
実験の本体を短くまとめたのが次のスクリプトです（scripts/grok_mini.py）。筆者が実際に回した版は、ログの項目や精度・ε の切り替えが増えているだけで、モデルと学習の中身は同じです。
"""a + b mod 31 のグロッキングで、標準と修正版の cross-entropy を比べる（記事掲載用の短縮版）。
使い方: python grok_mini.py std / python grok_mini.py fix
"""
import sys, math, torch, torch.nn as nn, torch.nn.functional as F
import fixed_torch as fx
LOSS = sys.argv[1] # "std" か "fix"
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 30000
torch.manual_seed(0)
P, D = 31, 64
class Net(nn.Module): # 1 層の Transformer（LayerNorm なし）
def __init__(self):
super().__init__()
self.tok = nn.Embedding(P + 1, D) # 0..30 と "=" の 31 番
self.pos = nn.Parameter(torch.randn(3, D) * 0.02)
self.attn = nn.MultiheadAttention(D, 4, bias=False, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.ReLU(), nn.Linear(4 * D, D))
self.out = nn.Linear(D, P, bias=False)
def forward(self, x):
h = self.tok(x) + self.pos
h = h + self.attn(h, h, h, need_weights=False)[0]
h = h + self.mlp(h)
return self.out(h[:, -1]) # "=" の位置で答えを予測
a, b = torch.meshgrid(torch.arange(P), torch.arange(P), indexing="ij")
X = torch.stack([a.flatten(), b.flatten(), torch.full((P * P,), P)], 1)
Y = (a.flatten() + b.flatten()) % P
perm = torch.randperm(P * P, generator=torch.Generator().manual_seed(1000))
tr, te = perm[: P * P // 2], perm[P * P // 2:] # 半分を学習、残りをテストに
net = Net()
opt = torch.optim.AdamW(net.parameters(), lr=1e-3, betas=(0.9, 0.98), eps=1e-8, weight_decay=1.0)
ce = F.cross_entropy if LOSS == "std" else fx.cross_entropy
for step in range(STEPS + 1):
z = net(X[tr])
loss = ce(z, Y[tr])
if step % 500 == 0:
with torch.no_grad():
z_t = z.gather(1, Y[tr, None])
m = -torch.logsumexp((z - z_t).scatter(1, Y[tr, None], float("-inf")), 1) # マージン
acc = (net(X[te]).argmax(1) == Y[te]).float().mean()
print(f"step {step:5d} loss {F.softplus(-m).mean():.2e} margin {m.mean():5.2f} test_acc {acc:.3f}")
opt.zero_grad()
loss.backward()
opt.step()
ログに出す損失は、標準・修正版のどちらで学習していても、正しい値（softplus(−m)）で計算しています。標準の損失そのものを出すと、崖の向こうで 0.0 と表示されてしまうためです。30,000 ステップの実行は、M5 Max の CPU で 12 本を同時に走らせて 1 本あたり 3〜4 分でした。
python grok_mini.py std 4000 # 標準の損失で 4,000 ステップ
python grok_mini.py fix 4000 # 修正版で 4,000 ステップ
最初の 4,000 ステップだけでも、違いははっきり出ます。左が標準、右が修正版です（表示は筆者が横に並べたもの）。
標準 修正版
step 1500 loss 1.67e-07 margin 15.83 test_acc 0.457 step 1500 loss 9.91e-08 margin 16.43 test_acc 0.455
step 2000 loss 5.21e-04 margin 7.94 test_acc 0.464 step 2000 loss 9.14e-08 margin 16.53 test_acc 0.455
step 2500 loss 3.51e-06 margin 12.98 test_acc 0.449 step 2500 loss 9.12e-08 margin 16.53 test_acc 0.453
step 3000 loss 3.79e-02 margin 12.45 test_acc 0.401 step 3000 loss 8.94e-08 margin 16.54 test_acc 0.453
step 3500 loss 1.16e-04 margin 9.51 test_acc 0.462 step 3500 loss 8.77e-08 margin 16.55 test_acc 0.455
標準は、マージンが 15.83 まで伸びた直後に 7.94 まで崩れ、損失が 1.67e-07 から 5.21e-04 へ約 3,000 倍に跳ねています。修正版は、崖の 16.64 に近づいてもそのまま穏やかに進みます。
損失の跳ねは、はっきり再現した
30,000 ステップ全体を描くと、違いはさらに明らかです。
この図から分かること：標準はマージンが崖で崩れては伸びるのこぎり状になり、跳ねのたびに正解率も揺れる。修正版は崖をなめらかに越える
標準の損失では、マージンが 16 付近まで伸びては 7 前後に崩れる「のこぎり」を、30,000 ステップで 21 回繰り返しました。崩れるたびに損失が数千〜1 万倍に跳ねます（跳ねの回数は 15 本の平均で 21.3 回）。論文 Figure 2 の上段（a + b mod 31、weight decay 1）とそっくりの形です。
修正版は、崖の 16.64 をそのまま越え、17.2 付近でなめらかに止まりました。止まる位置は、さきほどの釣り合いの式 log(1/ε) = 18.4 より少し低い値です。テスト正解率も、修正版は約 1 万ステップで 99% を越えてそのまま 100% を保ったのに対し、標準は約 1.2 万ステップで 95% を越えた後、跳ねのたびに 93〜95% へ揺れ続けました。
ε を変えた実験でも、論文の説明どおりの結果になりました。修正版の最終マージンは、ε = 1e-6・1e-8・1e-10・1e-12 で 12.4・17.0・22.1・26.8 と、理論の log(1/ε)（13.8・18.4・23.0・27.6）に沿って 4.6〜5.0 ずつ上がりました（理論は 4.6 ずつ）。そして ε = 1e-6 にすると、標準の損失でもマージンが 12.4 で止まって崖に届かず、跳ねは 0 回、正解率は 2 本とも 100% でした。「ε を 1e-7 以上にすれば標準でも跳ねない」という論文の主張どおりです。
正解率は、論文ほどはっきりした差が出なかった
では、跳ねが消えると汎化（テストの正解率）はどれだけ良くなるのでしょうか。乱数の種を 15 本に増やして比べました。
この図から分かること：跳ねはすべての種で消えた。正解率は、損失を bf16 で計算すると修正版が明確に良く、fp32 では差が見えなかった
結果は、精度によって分かれました。
| 条件（15 本） | 最終正解率の平均 | 95% 以上で終了 | 95% に届いた | 跳ね（平均） | 最終マージン |
|---|---|---|---|---|---|
| fp32・標準 | 84% | 8/15 | 11/15 | 21.3 回 | 12.6 |
| fp32・修正版 | 80% | 8/15 | 8/15 | 0 回 | 17.0 |
| bf16・標準 | 68% | 4/15 | 5/15 | ―（※） | 4.6 |
| bf16・修正版 | 83% | 10/15 | 10/15 | 0 回 | 17.1 |
※ bf16 の標準はマージンが崖（5.5）の手前で頭打ちになり、損失が十分に下がらないため、跳ねの回数は数えていません。
損失を bf16 で計算した場合は、修正版がはっきり上回りました。 標準はマージンが 5.5 の壁を越えられず平均 4.6 で止まり、95% 以上で終わったのは 15 本中 4 本です。修正版は fp32 と同じく 17.1 まで伸び、10 本が 95% 以上で終わりました。同じ乱数の種どうしで比べると、修正版が先に 95% に届いたのが 8 本、標準が先が 2 本です。
一方、fp32 では差が見えませんでした。 95% 以上で終わったのはどちらも 8 本で、同じ種どうしで「先に 95% に届いた」のも 6 本対 6 本です（残り 3 本はどちらも届かず）。標準が 95% に届いた本数はむしろ多く（11 本対 8 本）、ただしそのうち 3 本は跳ねで正解率が揺れ、最後は 95% を下回って終わりました。修正版は、届いた 8 本とも最後まで 95% 以上を保っています。
論文は同じ課題（10,000 エポック、乱数の種 5 本）で「標準 94 ± 5%、修正版 100 ± 0%」と報告していますが、手元のモデルは 30,000 ステップでも汎化しない種が多く、そもそもグロッキングが遅いようです。モデルの細部（ヘッド数・初期化など）は論文に書かれておらず筆者の推定なので、その違いが効いている可能性があります。手元の実験から言えるのは、次の 3 点です。
- 損失の跳ね（スリングショット）は、修正版で確実に消える
- 損失を bf16 で計算するなら、修正版のほうが明確に学習が進む
- fp32 では「跳ねが消える」「一度届いた正解率が安定する」効果はあるが、汎化が速くなるとまでは確認できなかった
なお、論文の Figure 2 下段にある「weight decay なし・ε = 1e-16 で、跳ねの瞬間に正解率が 61% → 23% に落ちる」実験は、課題とモデル（2 層 MLP で a² + b² mod 97）が違います。手元の Transformer で weight decay なし・ε = 1e-16 を試すと、標準は 3 本平均 22.3 回跳ね、修正版は 0 回でしたが、どちらも 30,000 ステップで汎化しなかったため（最終正解率 35〜46%）、この点は比較できていません。
バグ 2：Softmax Cancellation（蒸留で勾配が 0 にならない）
もう 1 つのバグは、知識蒸留（knowledge distillation） で起きます。蒸留とは、大きな教師モデルの出力する確率分布を「お手本」にして、生徒モデルを学習させる手法です。損失は soft-target cross-entropy（または KL ダイバージェンス）で、その勾配は「生徒の確率 q − 教師の確率 p」になります。
この図から分かること：標準は「ほぼ同じ確率どうしの引き算」をしてしまう。修正版は確率を一度も引かず、差を差のまま計算する
生徒が教師に近づくほど、q と p はほぼ同じ数になります。ほぼ同じ数どうしの引き算は、有効数字がごっそり消える典型的な場面です（桁落ち）。しかも標準の実装では、順伝播で使う softmax と、逆伝播で使う exp(log_softmax) の計算手順が微妙に違うので、生徒と教師がまったく同じでも勾配が 0 になりません。再現 1 の出力で、CPU と MPS がそれぞれ違う「0 でない勾配」を返したのは、このためです。
修正版は、確率どうしを引きません。ロジットの差 δ = z_S − z_T を先に取り、expm1（ex − 1）と log1p（log(1 + x)）という、x が小さくても桁落ちしない専用の関数を使って、log 空間で q − p を直接組み立てます。δ = 0 なら、結果はちょうど 0 です。
論文の例（教師ロジット [2, 1, 0]、生徒に ±10−7 のずれ）では、標準の JAX は 2 番目の勾配を 3 分の 1 に縮め、3 番目の符号を逆にしてしまいます。修正版は fp64 の正しい値と一致しました。手元の JAX・PyTorch の修正版でも、同じ値になっています。
再現 3：教師と同じ重みの生徒は、動かないはずが…
論文 7 章の実験を、手元でも組みました。生徒を 教師とまったく同じ重み で始めて、蒸留を続けます。生徒と教師の出力が同じなので、本来の勾配はちょうど 0 で、生徒は 1 歩も動かないはずです。
| 項目 | 論文 | 手元の再現 |
|---|---|---|
| 課題 | a + b mod 53 | 同じ |
| 教師 | テスト正解率 100% まで学習 | 全ペアで正解率 100% まで学習（筆者の簡略化） |
| 蒸留 | Adam（学習率 2e-3、weight decay 0、ε 1e-8）で 2,000 ステップ、温度 4 | 同じ。半分のペアで蒸留し、残り半分で KL と正解率を測定 |
| 乱数の種 | 3（図は 2 本） | 3 |
損失の部分は、標準が -(F.softmax(z_T / T) * F.log_softmax(z_S / T)).sum(-1).mean()、修正版が fx.soft_target_cross_entropy(z_S / T, z_T / T).mean() です。精度は「損失を計算する精度」だけを fp64・fp32・bf16 で変えました。
この図から分かること：標準の実装は、丸め誤差の残りカスが Adam で増幅されて生徒が教師から離れる。修正版は勾配がちょうど 0 なので、1 歩も動かない
標準の実装では、最初のステップの勾配がちょうど 0 になりませんでした。パラメーターの勾配の大きさ（ノルム）は bf16 で 10−3、fp32 で 10−7、fp64 で 10−16 前後です（論文が挙げる残りカスの目安は約 10−3・10−8・10−17）。そしてこの残りカスから、bf16 は 1 ステップ目、fp32 は 2 ステップ目、fp64 でも 3 ステップ目 に、生徒が教師から大きく離れました（KL が 10−1 前後まで上昇）。正解率は一時 92% 前後まで下がり、2,000 ステップ後には 100% に戻っています。
修正版は、3 つの精度すべてで勾配がちょうど 0 でした。生徒は 1 歩も動かず、2,000 ステップ後も KL はほぼ 0（10−18 程度の計算誤差）のままです。
ここで、論文と違う点が 1 つあります。論文では「fp64 の残りカス（約 10−17）は ε = 10−8 よりはるかに小さいので、生徒はほぼ動かない」とされていました。手元では fp64 でも離れています。fp64 の seed 0 を 1 ステップずつ追うと、勾配の大きさは 1.7×10−16 → 8.7×10−11 → 4.5×10−5 → 1.1 と、1 歩ごとに数万〜数十万倍に育っていました。
理由は Adam の性質にあると筆者は考えています。勾配が ε よりずっと小さいとき、Adam の 1 歩は「勾配 ÷ ε × 学習率」になり、実質の学習率は 2×10−3 ÷ 10−8 = 20 万という巨大な値になります。ずれを元に戻そうとする正しい勾配でさえ、この倍率で行き過ぎてしまい、ずれが雪だるま式に大きくなるわけです。どこまで小さい残りカスで不安定になるかはモデルの形（損失の曲がり具合）で変わるので、論文のモデルでは fp64 が持ちこたえ、手元のモデルでは持ちこたえなかったのだと思います。いずれにしても、勾配がちょうど 0 なら増幅のしようがない という修正版の強みは、手元でもはっきり確認できました。
なお、論文では bf16 の標準で正解率が下がっていましたが、手元では一時的な低下にとどまりました。
実務への影響：LLM の学習ではどうか
ここまでの実験は、どれも小さなモデルでした。読者がいちばん気になるのは「自分の LLM の学習や LoRA に関係あるのか」だと思います。論文自身も結論で「極端なマージンや、教師と生徒のごく小さな差は、通常の大規模学習ではまれ」と書いています。そこで、手元で 3 つの角度から確かめました。
学習済み LLM は、どれくらい「崖の向こう」にいるか
まず、学習済みの LLM が「次のトークン」をどれだけ強く確信しているかを測りました。テキストを読ませて、正解のトークン（実際に次に来たトークン）のマージンを 1 つずつ計算します。使ったのは、手元にあった Qwen3-0.6B と、Qwen3-8B-Base を土台に追加学習された 8B モデル（HybridAgent-8B-Agentic、文字の生成モードで使用）です。データは日本語 Wikipedia（JSQuAD の本文）・英語 Wikipedia（SQuAD v2 の本文）・Python コード（PyTorch 自身のソース）を、それぞれ 16,352 トークンずつ用意しました。
この図から分かること：fp32 の崖を越えるトークンはごくわずか。ただし bf16 の崖なら、8B ではコードの半分近くが越える
| モデル・データ | マージンの中央値 | bf16 の崖（5.55）超え | fp32 の崖（16.64）超え |
|---|---|---|---|
| 0.6B・日本語 | −2.2 | 8.2% | 0.00% |
| 0.6B・英語 | −1.5 | 4.3% | 0.00% |
| 0.6B・コード | 1.9 | 25.8% | 0.03% |
| 8B・日本語 | −0.8 | 13.6% | 0.18% |
| 8B・英語 | 1.0 | 25.2% | 0.31% |
| 8B・コード | 5.1 | 46.9% | 0.53% |
読み取れることは 3 つです。
- fp32 で損失を計算するなら、崖の向こうのトークンはごくわずかです。0.6B で 0〜0.03%、8B でも 0.2〜0.5% でした。transformers の既定の損失はこちらに当たります。
- bf16 のまま計算すると話が変わります。 8B では日本語で 14%、英語で 25%、コードでは 47% のトークンが崖の向こうにいて、正解の勾配が 0 に丸まります。自前の学習ループで bf16 のロジットをそのまま損失に渡している場合は要注意です。
- 大きいモデルほど、そしてコードのような定型的なデータほど、確信が強い（マージンが大きい）傾向があります。
マージンの中央値がマイナスなのは、正解のトークンの確率が 50% に届かない場面が半分以上ある、という意味です。普通の文章を読ませたときの LLM は、意外と「迷いながら」次のトークンを当てています。
LoRA で同じデータを何周もさせてみる
学習済みの状態では崖の向こうは少数でも、少量のデータで何周も微調整すれば、丸暗記が進んでマージンが伸びるはずです。そこで、Qwen3-0.6B を LoRA（モデル本体は固定し、小さな追加の行列だけを学習する手法）で過学習させ、標準と修正版を比べました。
| 項目 | 設定 |
|---|---|
| データ | 日本語 Wikipedia（JSQuAD の学習用の本文）16 本 × 256 トークン、毎ステップ全部を使う |
| LoRA | r = 16、alpha = 32、注意機構と MLP の全線形層、LoRA の重みと Adam の状態は fp32 |
| 最適化 | AdamW、学習率 1e-3、ε = 1e-8、weight decay 0、600 ステップ |
| 損失 | ロジットを fp32 に上げてから、標準 F.cross_entropy / 修正版 |
この図から分かること：LoRA で丸暗記させても跳ねは出ず、標準と修正版はほぼ同じ学習曲線になった
600 ステップで、平均マージンは約 11.8 まで伸び、fp32 の崖を越えたトークンは 2.9%（標準）・3.2%（修正版）まで増えました。bf16 の崖は、最後には 100% のトークンが越えています。それでも、標準で損失の跳ねは出ませんでした。損失が直前より増えた最大の比は 1.88 倍で、グロッキングの実験の数千〜1 万倍とはまったく違います。最終の損失も 1.71×10−5 と 1.64×10−5 で、ほぼ同じです。
グロッキングとの違いは、崖の向こうに行くのが「一部のトークンだけ」だという点にあると考えられます。スリングショットは「全部の例が崖の向こうへ行き、勾配がほぼ消える」ことで起きました（スリングショットの図の ②）。LLM では、600 ステップ丸暗記させても 97% のトークンは崖の手前にいて、勾配はそこから十分に来るので、Adam の分母が縮みきらないのだと思います。
ただし、これは 0.6B・LoRA・600 ステップ・weight decay なしの 1 条件だけの結果です。もっと長く回す、全パラメーターを学習する、bf16 のまま損失を計算する、といった条件では違う結果になる可能性があります。
速さとメモリ
論文は「目立った速度低下なしに置き換えられる」と書いています。ただ、論文の実装は JAX で、JAX はコンパイラー（XLA）が計算をまとめて最適化します。素の PyTorch に移植した版ではどうなるかを、MPS で測りました。語彙は Qwen3 と同じ 151,936、順伝播と逆伝播を合わせた時間の中央値です。
| 条件（MPS） | 標準 F.cross_entropy | 修正版 | 比 |
|---|---|---|---|
| fp32 ロジット・1,024 トークン | 13.3 ms | 39.6 ms | 2.97 倍 |
| fp32 ロジット・4,096 トークン | 51.8 ms | 158.0 ms | 3.05 倍 |
| bf16 ロジットを fp32 に上げて・4,096 トークン | 66.0 ms | 173.3 ms | 2.63 倍 |
| 上の 2 行目を torch.compile した場合 | 27.6 ms | 61.0 ms | 2.21 倍 |
損失の計算だけを見ると、素の PyTorch では 2.6〜3 倍遅くなりました。語彙サイズの中間テンソル（正解を引いた差、マスクなど）を何本も作るためで、逆伝播のために保持するメモリも、4,096 トークンで 0.6GB（fp32）〜3.0GB（bf16 から上げた場合）増えました。torch.compile をかけると差は 2.2 倍まで縮みます。値と勾配は、通常の範囲のロジットでは標準と丸め誤差の範囲で一致しました。
ただし、これは損失の部分だけの話です。モデル全体の 1 ステップで見ると、さきほどの Qwen3-0.6B の LoRA 学習（4,096 トークン、MPS）では、1 ステップが標準 1.64 秒・修正版 1.73 秒（中央値、1 本ずつ測定）で、増えたのは約 5% でした。語彙の大きさは同じでもモデル本体の計算はモデルが大きいほど増えるので、大きなモデルほど損失の部分の割合は小さくなります。
速さが気になる場合は、まず torch.compile を試し、それでも足りなければ fused（1 つのカーネルにまとめた）実装を書くことになります。この記事ではそこまでは踏み込んでいません。
自分のコードに入れるなら
まず「気にするべきか」を判断する
ここまでの結果を、3 つの質問にまとめました。1 つでも「はい」なら、修正版を試す価値があります。
この図から分かること：問題になるのは「低い精度で損失を計算」「ほぼ同じ教師から蒸留」「長い丸暗記」のどれか。通常の LLM の微調整なら優先度は低い
-
Q1：損失を bf16 / fp16 のまま計算している？ 自前の学習ループで、bf16 のロジットをそのまま
F.cross_entropyに渡している場合です。CPU では崖が 5.5 まで下がり、8B モデルでは次トークンの 14〜47% が崖の向こうにいました。修正版にするか、少なくとも損失の前に.float()で fp32 に上げましょう。 -
Q2：生徒が教師とほぼ同じ状態から蒸留する？ 自己蒸留や、同じ重みから始めて全語彙の分布どうしの KL を取る正則化などです。再現 3 のとおり、標準の実装では fp64 でも生徒が離れていきます。修正版の
soft_target_cross_entropyが使えます。 - Q3：同じデータを何周もして丸暗記させ、Adam の ε が 1e-8 以下？ 小さなデータで長く学習する場合です。損失の周期的な跳ねがないか確認し、跳ねるなら修正版にするか、ε を 1e-7 以上にします。
なお、強化学習（RLHF など）で参照モデルとの KL を取る場合は、実装によって事情が違います。サンプルしたトークンの対数確率だけで KL を近似する実装が多く、その場合は「全語彙の分布どうしの引き算」をしないので、論文が扱った形とは異なります。ここは筆者の見立てで、検証はしていません。
Hugging Face の Trainer に差し込む
transformers の Trainer には、損失関数を差し替える compute_loss_func という引数があります。次のファイルを scripts/hf_loss.py として保存します。
"""Hugging Face Trainer の compute_loss_func に渡す、修正版 cross-entropy の因果 LM 損失。"""
from fixed_torch import cross_entropy
def causal_lm_fixed_loss(outputs, labels, num_items_in_batch=None):
logits = outputs["logits"] if isinstance(outputs, dict) else outputs[0]
logits = logits[:, :-1].float() # 位置 i で i+1 番目のトークンを当てる
labels = labels[:, 1:].to(logits.device)
loss_sum = cross_entropy(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1),
ignore_index=-100, reduction="sum")
if num_items_in_batch is None: # 勾配累積なし: 有効トークン数で割る
num_items_in_batch = (labels != -100).sum()
return loss_sum / num_items_in_batch
使うときは、Trainer を作るときに渡すだけです。
from transformers import Trainer
from hf_loss import causal_lm_fixed_loss
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
compute_loss_func=causal_lm_fixed_loss, # ここで損失を差し替える
)
trainer.train()
ラベルを 1 つずらす、ロジットを fp32 に上げる、−100 を無視する、勾配累積のときは num_items_in_batch で割る、という点は transformers 既定の損失と同じにしています。手元では、小さな乱数初期化の Qwen3 で既定の損失と小数第 6 位まで一致すること（6.897326 と 6.897325）と、勾配累積 2 で 5 ステップ学習できることを確認しました。
自分のモデルで「崖を越えるトークン」を数える
自分のモデルとデータで、崖を越えるトークンがどれくらいあるかを数えておくと判断しやすくなります。次のスクリプトは、テキストファイルを読ませて次トークンのマージンを測ります。
"""手元の因果 LM で、次トークンのマージンが精度の崖を越える割合を測る。
使い方: python check_margins.py <モデルのパスか Hugging Face の名前> <テキストファイル>
"""
import sys, math, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id, text_file = sys.argv[1], sys.argv[2]
dev = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16).to(dev).eval()
ids = tok(open(text_file).read(), return_tensors="pt").input_ids[:, :2048].to(dev)
with torch.no_grad():
logits = model(ids).logits[0, :-1].float() # 損失と同じく fp32 に上げる
tgt = ids[0, 1:]
z_t = logits.gather(1, tgt[:, None])
others = (logits - z_t).scatter(1, tgt[:, None], float("-inf"))
m = (-torch.logsumexp(others, 1)).cpu() # 次トークンのマージン
print(f"トークン数 {m.numel()} マージンの中央値 {m.median():.2f} 最大 {m.max():.2f}")
for name, c in [("bf16 の崖", math.log(2 ** 8)), ("fp32 の崖", math.log(2 ** 24))]:
print(f"{name}（{c:.2f}）を越えた割合: {(m > c).float().mean() * 100:.2f}%")
# 例：Qwen3-0.6B に、学習に使う予定のデータの一部を読ませる
python check_margins.py Qwen/Qwen3-0.6B my_train_sample.txt
筆者が短い PyTorch のコード（147 トークン）で試したところ、Qwen3-0.6B では bf16 の崖を越えたトークンが 34.69%、fp32 の崖は 0.00% でした。コードは定型の書き方が多いので、普通の文章より確信が強く出ます。
トラブルシューティング
PyTorch と Mac に移植する中で、筆者が実際に引っかかった点をまとめます。
この図から分かること：つまずきの多くは「Mac の GPU（MPS）の関数の精度」と「16 ビット形式の限界」。どちらも対処は簡単
| 症状 | 原因 | 対処 |
|---|---|---|
| MPS で修正版の損失だけ 0.0 になる（勾配は正しい） | MPS の F.softplus は入力が −17 付近より小さいと 0 を返す（CPU は正確） |
torch.logaddexp(x, torch.zeros_like(x)) で softplus を計算する |
| マージンがちょうど 0 の例だけ勾配が 0 になる | softplus を clamp(x, min=0) + log1p(exp(-abs(x))) で自作すると、x = 0 で微分が 0 になる |
同じく logaddexp を使う（x = 0 での微分は正しく 0.5） |
| bf16 の崖が CPU では 5.56、MPS では 16.75 | MPS の cross_entropy は bf16 を渡しても内部を fp32 で計算している | 崖の位置は環境で違う前提で、自分の環境で測る |
| fp16 だと修正版でもマージン約 11 で勾配が 0 | fp16 は表せる最小の数が大きく、小さな勾配がアンダーフローする | 損失は bf16 か fp32 で計算する |
| 損失の計算が 2.6〜3 倍遅い | 素の PyTorch では、語彙サイズの中間テンソルが何本も作られる |
torch.compile で差が縮む。モデル全体の 1 ステップで見ると差は小さい |
transformers に fixed_cross_entropy という関数がある |
名前が似ているだけで、トークン数で割る処理のための関数 | Trainer の compute_loss_func に修正版を渡す |
MPS で .double() するとエラーになる |
MPS は float64 に対応していない | 正しい値との比較は .cpu().double() の順で CPU に移してから行う |
2 行目は、筆者が最初にやってしまった失敗です。「桁あふれしない softplus」の定番の書き方ですが、x = 0 のところで clamp と abs の微分がどちらも 0 扱いになり、勾配が消えます。マージンがちょうど 0 になるのは「2 クラスでロジットが等しい」ような場合だけですが、2 クラス分類で出力層を 0 で初期化すると、最初のステップで実際に起きます。PyTorch では logaddexp を使うのが安全です。
まとめ
Google DeepMind の論文「A Fix for Softmax and Cross-Entropy」を読み、手元の Mac で再現しました。要点を振り返ります。
- 2 つのバグは本物で、手元の PyTorch・JAX でもそのまま起きる。 Softmax Collapse は「1 + ごく小さい数」が 1 に丸まることで正解の勾配だけが消える問題、Softmax Cancellation は蒸留で「ほぼ同じ確率どうしの引き算」をする問題です。精度の崖は bf16 で 5.55、fp32 で 16.64、fp64 で 36.74 と、理論どおりの位置にありました。
-
修正は数行で、勾配は正しくなる。 1 を足さずにマージンを作って softplus でまとめる、確率を引かずに差を log 空間で作る、の 2 つだけです。PyTorch に移植するときは、MPS の
F.softplusの精度に注意してlogaddexpを使いましょう。 - 学習への効果は条件しだい。 グロッキングでは損失の跳ねが確実に消え、損失を bf16 で計算した場合は修正版が明確に良い結果でした。一方 fp32 では、汎化が速くなるとまでは確認できませんでした。蒸留では、標準の実装は fp64 でも生徒が教師から離れ、修正版は勾配がちょうど 0 で動きませんでした。
- 通常の LLM の微調整では優先度は低い。 損失を fp32 で計算するなら崖の向こうのトークンは 0〜0.5% で、Qwen3-0.6B を LoRA で丸暗記させても跳ねは出ませんでした。気にするべきなのは、損失を bf16 のまま計算している、ほぼ同じ教師から蒸留する、小さなデータで長く丸暗記させる、のどれかに当てはまる場合です。
Softmax と Cross-Entropy が LLM の学習でも使われる基本の計算なのは確かですが、「入れれば LLM が良くなる」とは論文も主張していません。論文の実験はグロッキング用の小さなモデルで、手元の LLM の実験でも差はほとんど出ませんでした。自分の学習が崖の向こうに行く条件かどうかを測ってから入れるのが、いちばん現実的な使い方だと思います。
次の一歩としては、記事中の check_margins.py で自分のモデルとデータのマージンを測ってみることをおすすめします。蒸留や自己蒸留をしている方は、soft_target_cross_entropy に置き換えて、最初の数ステップの勾配が本当に 0 になっているかを確かめてみてください。
参考リソース
- 論文：Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen「A Fix for Softmax and Cross-Entropy」arXiv:2610.31642（2026 年 10 月 22 日、CC BY 4.0）https://arxiv.org/abs/2610.31642
- スリングショット現象の元の報告：Vimal Thilak ほか「The Slingshot Mechanism: An Empirical Study of Adaptive Optimizers and the Grokking Phenomenon」（2022 年、論文の参考文献より）
- 蒸留：Geoffrey Hinton ほか「Distilling the Knowledge in a Neural Network」（2015 年）
- PyTorch のドキュメント：torch.nn.functional.cross_entropy、torch.logaddexp
- Hugging Face transformers：Trainer（
compute_loss_func引数） - この記事で使った主なスクリプト（本文に掲載）：
fixed_torch.py（修正版の PyTorch 移植）、demo_bugs.py（2 つのバグの再現）、grok_mini.py（グロッキング）、hf_loss.py（Trainer 用の損失）、check_margins.py（LLM のマージン）