Pretraining Transformers with Quantized Softmax in Attention
Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski
cs.LG
2026-09-27
伊利诺伊大学用 K-interval attention 系统比较量化 softmax 的行校准、硬取整与 STE 放置。124M 模型训 2.5B token，固定窗口加后归一化 STE 在 K=4 只比 softmax 差 +0.019 nats；切断行极值的梯度、前向不变也会晚发散 0.74 nats。
低精度 Transformer 训练已经把线性投影、越来越多的 attention 矩阵乘推进窄格式，指数和行归约仍故意留在 FP16/FP32。EXAQ、IndexSoftmax、NVIDIA 的 MXFP8 配方都把 softmax 路径降精度标成未测或未来工作。推理时近似算子只扰动一个固定函数。预训练时它定义学习规则：模型看见近似前向，更新用的是实现挂上去的那条梯度。
精确 softmax 里对行最大值做 stop-gradient 无害，平移不变让这条梯度精确抵消。同一习惯搬到按行校准的网格上就不成立，行极值还在定格子。量化感知训练做 fake quantization 时，observer 统计同样常不进 autograd。伊利诺伊大学把问题收成一个每项前向、反向都有闭式的算子族，在模型、数据、优化器对齐的预训练里逐项对照。实验在 fp32 attention 算术里做，训练开 TF32 matmul，不测 kernel，也不声称加速。
K-interval attention 用 K+1 个表格值，在 K 个区间上重建 e^z，再归一化成注意力权重。三个设计轴决定一个算子：
完整校准梯度会穿过行极值。平移不变要求一行分数的梯度求和为零。只保留插值斜率、把极值 detach 掉，这条恒等式就被破坏。类比很直接：精确 softmax 减最大值只是数值稳定；校准网格上最大值还在决定格宽，摘掉它等于改了学习规则。
最干净的对照是可微的 LERP K=32。完整校准梯度相对 softmax 的验证损失差 ΔNLL=−0.0008（95% CI [−0.004, +0.002]），检测不到差别。同一前向、把行极值 detach，终点高 +0.737 nats。轨迹前 25M token 还贴着 softmax，约 30M 处分开，行跨度随后涨一个数量级以上。把 detach 梯度投影回零和子空间，终点比 plain detach 更差。只留最大值相关的校准梯度，几乎收回全部缺口；只留最小值通道几乎收不回。
硬取整在 K=4 上，校准和 STE 位置强交互：
| 条件 | 250M ΔNLL | 2.5B ΔNLL |
| MinMax–Weight K=4 | +0.389 | +0.890 |
| MinMax–Prob K=4 | +0.074 | +0.061 |
| FWM–Weight K=4 | +0.039 | +0.043 |
| FWM–Prob K=4 | +0.007 | +0.019 |
| LERP K=4 | +0.052 | +0.005 |
改 STE 到概率侧，或把 MinMax 换成 FWM，都能消掉大部分缺口；两个一起改，再多出来的很少。这组符号在 1B 参数、五颗种子上都保持。MinMax–Weight K=4 是唯一优化异常：250M 时 45% 的 step 被 clip；2.5B 中途 375M 处塌到 +3.66 nats，单步梯度范数 10^4，最后 1.5B token 里 98% 的 step 被 clip，再部分回到 +0.89 平台。没有任何 Prob-STE 或 FWM run 出现这种 clip。
LERP 在 K∈{4,16,32} 上 2.5B 终点都落在 softmax 的 0.005 nats 以内；K=4 从 125M 的 +0.086 收到 +0.005。FWM–Weight 在 K=16 是 +0.004。0.01 nats 大约对应 1% perplexity 比，+0.89 nats 大约 2.4 倍。五种子上条件排序可复现（Kendall's W=0.91），按 K 平均，FWM 相对 MinMax 低 0.067 nats，Prob-STE 相对 Weight-STE 低 0.021。同一种子、124M、100M token 上，BF16 指数相对 softmax 为 −0.0034 nats，带 STE 的 FP8 取整指数为 +0.0028，远小于同设置 MinMax K=4 的缺口。失败不是「指数不够多 bit」这么简单。
下游按 2.5B ΔNLL 预先分成大缺口和近基线。大缺口组 5 个条件 × 7 个评测配置共 35 次对照，点估计全低于 softmax，27 次 q<0.05；MinMax–Weight K=4 掉 3.9–19.4 个准确率点。近基线组 35 次里 8 次 q<0.05，正负都有，单种子无法归因到算子。评测时换成精确 softmax，42 格里 37 格缺口变大；LERP K=2 在 WikiText-103 测试集从 +0.005 拉到 +1.814。权重已经绑在训练时的那个算子上。
给想把 softmax 也量化进预训练的人看：前向近似和反向规则要一起评。QAT 里把 observer detach 掉、softmax 实现里对行最大值 stop-gradient，在精确指数下是数值习惯，在校准网格上会改学习规则。
这篇不提供更快的 kernel，也不证明能上 7B。它把格子怎么定、硬取整还是插值、STE 放归一化前还是后拆开，并给出匹配实验。只抄推理侧查表 softmax 进预训练，MinMax 配 Weight-STE 在粗网格上是明确的失败组合；FWM 或 Prob-STE 把同一硬前向拉回接近基线。
对现有 FP8/FP4 attention 配方，这解释了指数路径为何仍留在高精度，以及哪几处 Jacobian 不能省。机制清楚的渐进工作。
作者列得很清楚。A–C 套件单种子；五种子套件的方差不能外推到 2.5B。没有学习率扫描，2.5B 上的持续性只对测试过的 cosine 日程成立。τ=6 来自冻结权重的推理筛选，最后一轮还用了后来 C/D 套件的 WikiText-103 验证集，没有证明对训练最优。1B 套件是另一语料上的早期探针（0.1 token/参数），没有超过 1B 参数或 2.5B token 的 run。
两个 STE 的历史前向并不 bitwise 相同，是 surrogate 对照的混淆项；冻结 checkpoint 上均值差最多 3.4×10^{-4} nats，训练期累积没在 A–C 上复测。套件 D 的训练栈和算子族部分绑定，Prob-STE 多在云端 RTX 5090，其余多在本地 3090。范围-分辨率反馈只是假设，文中没有任何干预能识别机制。小 NLL 缺口不能当成下游等价。