Pretraining Transformers with Quantized Softmax in Attention
Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski
cs.LG
2026-09-27
K-interval attention for quantized-softmax pretraining: FWM+Prob-STE is +0.019 nats vs softmax at K=4 (124M/2.5B); detaching row extrema later costs +0.74 nats, same forward.
Low-precision Transformer training already pushes linear projections, and increasingly the attention matmuls, into narrow formats. The exponential and its row reduction stay in FP16/FP32 by design. EXAQ, IndexSoftmax, and NVIDIA's MXFP8 recipe all leave a lower-precision softmax path untested or as future work. At inference an approximate operator perturbs a fixed function. In pretraining it defines the learning rule: the model sees the approximate forward and is updated by whatever gradient the implementation attaches.
Detaching the row maximum is harmless for exact softmax. Shift invariance cancels that gradient exactly. The same habit is not harmless on a calibrated grid, because the extrema also set the grid. Fake-quantization observers in QAT frameworks likewise keep their statistics off the autograd tape. This Illinois paper studies a family simple enough that every forward and backward term is closed-form, in pretraining runs matched on model, data, and optimizer. The experiments use fp32 attention arithmetic with TF32 matmuls in training. No kernels are measured and no speedup is claimed.
K-interval attention rebuilds e^z from K+1 tabulated grid values over K intervals, then normalizes. Three axes define an operator:
The full calibration gradient routes through the row extrema. Shift invariance requires the score gradient of a row to sum to zero. Keeping only the interpolation slope and detaching the extrema breaks that identity. Subtracting the max is a numeric convenience for exact softmax; on a calibrated grid the max also sets bin width, so dropping it changes the learning rule.
The cleanest case is differentiable LERP at K=32. With full calibration gradients the validation gap vs softmax is ΔNLL=−0.0008 (95% CI [−0.004, +0.002]), undetectable. The same forward with detached extrema ends +0.737 nats higher. The run tracks softmax for 25M tokens, separates near 30M, and the row span then grows by more than an order of magnitude. Projecting the detached row onto the zero-sum subspace ends worse than plain detach. Keeping only the max-dependent calibration term recovers almost the whole gap; keeping only the min term recovers almost none.
For hard rounding at K=4, calibration and surrogate placement interact strongly:
| condition | 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 |
Moving the surrogate to the probability level, or replacing MinMax with FWM, removes most of the deficit. Doing both adds little more. The signs hold at 1B parameters and on all five seeds. MinMax–Weight K=4 is the only optimization abnormality: 45% of steps clipped at 250M; at 2.5B it collapses to +3.66 nats at 375M tokens, with single-step gradient norms of 10^4 and 98% of steps clipped over the last 1.5B tokens, then partially recovers to a +0.89 plateau. No Prob-STE or FWM run clips that way.
LERP at K in {4, 16, 32} ends within 0.005 nats of softmax at 2.5B; K=4 closes from +0.086 at 125M to +0.005. FWM–Weight at K=16 is +0.004. A 0.01 nat gap is about a 1% perplexity ratio; +0.89 nats is about 2.4×. On five seeds the condition ranking is reproducible (Kendall's W=0.91). Averaged over K, FWM cuts 0.067 nats vs MinMax and Prob-STE cuts 0.021 vs Weight-STE. On the same seed at 124M / 100M tokens, BF16-exp is −0.0034 nats from softmax and an FP8-rounded exp with a weight-level STE is +0.0028, far below the MinMax K=4 gaps. Bit width of the exponential is not what fails here.
Downstream groups were fixed from 2.5B ΔNLL before any benchmark run. All 35 large-gap comparisons (five conditions × seven configs) have negative point estimates vs softmax, 27 at q<0.05; MinMax–Weight K=4 loses 3.9–19.4 accuracy points. Of 35 near-baseline comparisons, 8 reach q<0.05, six negative and two positive; with one seed those small gaps cannot be attributed to the operator. Swapping in exact softmax at evaluation widens the gap in 37 of 42 cells; LERP K=2 on WikiText-103 test goes from +0.005 to +1.814. The trained weights depend on the training operator.
Anyone who wants to quantize softmax during pretraining has to evaluate the forward approximation and the backward rule together. Detaching observers in QAT, or stop-gradient on the row max in a softmax kernel, is a numeric habit under an exact exponential. On a calibrated grid it changes the learning rule.
This paper does not ship a faster kernel and does not show the method at 7B. It separates how the grid is set, interpolation vs hard rounding, and STE placement relative to normalization, under matched training. Copying an inference lookup-table softmax into pretraining with MinMax plus Weight-STE is a clear failure at coarse K. FWM or Prob-STE pull the same hard forward back near baseline.
For current FP8/FP4 attention recipes, that is why the exponential path is still left in high precision, and which Jacobian terms cannot be dropped. Incremental work with a clear mechanism.
The authors are explicit. Suites A–C use one seed; suite D's seed spread does not transfer to 2.5B. There is no learning-rate sweep, so persistence at 2.5B is specific to the tested cosine schedule. τ=6 comes from an inference-time screen on frozen weights, and the last stage used the WikiText-103 validation split later used by suites C and D; it is not shown optimal for training. The 1B suite is an early-regime probe on another corpus (0.1 tokens/parameter). No run exceeds 1B parameters or 2.5B tokens.
The two historical STE forwards are not bitwise identical, a confound of the surrogate contrast. On frozen checkpoints the mean NLL difference is at most 3.4×10^{-4} nats; training-time accumulation was not remeasured on suites A–C. In suite D the training stack is partly tied to operator family: Prob-STE mostly on cloud RTX 5090s, other families mostly on a local 3090. The range–resolution feedback is a hypothesis; no intervention in the paper identifies the mechanism. A small NLL gap does not establish downstream equality.