A Fix for Softmax and Cross-Entropy

概要

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously-observed "Slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's $ε$ hyperparameter.

参考リンク

• Paper ↗

リンク先のライセンス・再現性・データ条件をご確認ください。自動分類や実装の目安は保証ではありません。

実装を始めるプロンプト

参考リンクを確認したうえで、次の内容をコピーして利用できます。

# 論文・記事
A Fix for Softmax and Cross-Entropy

# 概要
Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much…

# 原典
https://arxiv.org/abs/2610.31642

上記の手法を実装してください。実行前に依存関係とライセンスを確認してください。