[论文] A Fix for Softmax and Cross-Entropy

论文概要

研究领域: ML 作者: Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen 发布时间: 2026-10-22 arXiv: 2610.31642

中文摘要

标准的 softmax、log-softmax 和交叉熵实现存在两个数值缺陷。第一个缺陷「Softmax 坍缩」：当目标 logit 远大于其他 logit 时，其梯度会被舍入为零，导致梯度之和不再为零。第二个缺陷「Softmax 抵消」：当学生 logit 接近教师 logit 时，即使两者完全相等也会产生非零梯度。我们提出了修复这两个缺陷的简单实现。修复消除了此前观察到的「弹弓」（Slingshot）现象，加速了带权重衰减的 grokking 过程，防止了无权重衰减时的精度崩溃，阻止了蒸馏过程中学生网络偏离教师网络，并揭示了记忆化之后 logit 增长受权重衰减与 Adam 的 ε 超参数之间平衡的约束。

原文摘要

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. 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. Softmax Cancellation occurs when student logits are close to teacher logits, producing non-zero gradients even when they are equal. We present simple implementations 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.

#论文 #arXiv #ML #小凯