Pretraining Transformers with Quantized Softmax in Attention

AI summary

Overview

Research area: Low-precision / quantization-aware training of Transformer language models, specifically the attention softmax operator and its gradients during pretraining.

Technical level: Advanced. The paper assumes familiarity with attention, softmax normalization, straight-through estimators (STEs), Jacobians, and quantization-aware training. It is written for researchers who work on numerics and training dynamics rather than for general readers.

Scope (one sentence): The paper studies a family of "K-interval" approximate softmax operators in attention and measures, through matched pretraining experiments up to 1B parameters and 2.5B tokens, how calibration, reconstruction, and the placement of the backward surrogate change validation loss and downstream performance.

What This Paper Is About

Low-precision Transformer systems increasingly quantize attention matrix multiplications, but the softmax exponential and its row normalization are usually kept in higher precision (FP16/FP32). Whether an approximate reconstruction of the exponential can be used during pretraining is left open by prior work. The problem is different from inference: during training the operator defines the learning rule, because the model is updated by whatever gradient the implementation attaches to the approximate forward. This paper asks which choices — how the exponential grid is calibrated, how scores are reconstructed onto it, and where the backward surrogate sits relative to normalization — actually matter for pretraining.

Key Contributions

1. A closed-form operator family with matching gradients. The authors define "K-interval attention," which approximates the exponential using K+1 grid values, and derive the corresponding backward rules, including the derivatives through per-row calibration statistics (the row max and min). They call the complete rule the "full calibration gradient" and show it satisfies a zero-sum identity on the score gradients that shift invariance implies.

2. Isolating a gradient defect with an unchanged forward. By detaching the row extrema, they show that the same forward computation with an incomplete backward produces delayed training failure, and that projecting the resulting gradients onto the zero-sum subspace does not repair it.

3. A calibration × surrogate factorial study. Holding Nearest (hard rounding) reconstruction at K=4, they measure how calibration policy (per-row min–max vs. a fixed window anchored at the row max) and surrogate placement (weight-level vs. probability-level) affect loss, stability, and their interaction — replicated on scales, horizons, and five seeds.

4. Evidence that coarse deterministic reconstruction need not prevent near-baseline performance. Piecewise-linear interpolation (LERP) at K ∈ {4, 16, 32} ends within 0.005 nats of softmax at 2.5B tokens, and selected hard configurations also approach softmax.

Main Findings

• Detaching the row extrema causes delayed failure. LERP K=32 with full calibration gradients shows no detectable difference from softmax (ΔNLL = −0.0008, 95% CI [−0.004, +0.002]), while the identical forward with the detached backward ends +0.737 nats higher. The failure is delayed: the run tracks softmax through 25M tokens and separates near 30M tokens. Across the tested runs, detached variants end 0.65–3.07 nats above the matched full-calibration run.

• Restoring zero-sum is not sufficient. Projecting each detached row onto the zero-sum subspace removes the common-mode component exactly, yet the projected run ends worse than plain detach. The maximum-only backward is not strictly zero-sum, yet ends within 0.015 nats of full on both seeds it was tested on (seeds 7 and 42).

• The gradient channel that matters is the row maximum. Retaining only the maximum-dependent calibration gradient recovers almost the whole full-minus-detach gap on both seeds; retaining only the minimum-dependent gradient recovers almost none of it.

• Calibration and surrogate placement interact strongly under hard rounding at K=4. MinMax–Weight ends +0.39 nats above softmax at 250M tokens and +0.89 nats at 2.5B. Moving the surrogate to the probability level (Prob−Weight under MinMax: −0.316 at 250M, −0.829 at 2.5B) or replacing MinMax with FWM (FWM−MinMax under Weight-STE: −0.350 at 250M, −0.848 at 2.5B) each removes most of the deficit; changing both adds little more. The K=4 interaction is +0.284 at 250M and +0.805 at 2.5B.

• These directions hold at larger scale and across seeds. The calibration effect, surrogate effect, and interaction keep their sign at 1B parameters (100M tokens) and on all five seeds of suite D, where the mean interaction is +0.052 with a 95% t-interval of [0.030, 0.075]. The ordering of conditions is reproducible (Kendall's W 0.91).

• The surrogate effect shrinks with K, faster under FWM. In the 250M factorial, the surrogate effect under FWM is about 30× smaller at K=16 than at K=4 and only just detectable, whereas under MinMax it stays large; five seeds show the same decay.

• Selected configurations approach softmax. At 2.5B tokens, fixed-window calibration with a post-normalization surrogate yields a gap of +0.019 nats at K=4, and with a pre-normalization surrogate +0.004 nats at K=16 (abstract's terms; the body calls these Prob-STE and Weight-STE respectively). At 2.5B, MinMax–Weight falls from +0.89 at K=4 to +0.04 at K=64.

• LERP improves with training while hard-rounding gaps plateau. LERP K=4 goes from +0.086 nats at 125M tokens to +0.005 at 2.5B. The hard-vs-LERP ordering consequently changes with training budget: FWM–Weight K=4 minus LERP K=4 is negative at 125M, negative at 250M, and positive at 2.5B.

• A geometry observation: MinMax–Weight K=4 grows its score span. At 2.5B its row span is about seven times softmax's, its native interval width reaches h = 53 nats, and the total variation between native and softmax probabilities is 0.86. Across the five-seed grid, condition-mean ΔNLL has Spearman correlation +0.95 with native width and −0.31 with score span. The authors frame the resulting "range–resolution feedback" idea explicitly as a hypothesis that no intervention in the paper identifies.

• Counter-examples limit that hypothesis. MinMax–Weight K=64's span more than doubles over training while its gap stays between +0.02 and +0.04 nats (h = 1.54, ΔNLL = 0.038). At 250M the MinMax–Weight/Prob gap is already open at 50M tokens while the two geometries nearly coincide. MinMax–Weight K=16 has a gap without any gradient abnormality; only MinMax–Weight K=4 shows an optimization abnormality (45% of steps clipped at 250M).

• Downstream transfer is graded, not binary. On WikiText-103, PTB, and C4 the ordering is preserved (Spearman ≥ 0.95). All 35 large-gap comparisons across five conditions and seven configurations have negative point estimates and 27 reach q<0.05; MinMax–Weight K=4 loses 3.9–19.4 accuracy points, less at larger K. Of the near-baseline group's 35 comparisons, 8 reach q<0.05 (six negative, two positive).

• The trained weights depend on the training operator. Substituting exact softmax at evaluation widens the gap in 37 of 42 cells; for LERP K=2 on the WikiText-103 test set it goes from +0.005 to +1.814 nats.

• Precision alone does not explain the failures. At 124M / 100M tokens (seed 7), BF16-exp ends −0.0034 nats from softmax and FP8-rounded exp +0.0028, far below the same-seed LERP and MinMax K=4 gaps.

Methodology in Plain English

The authors take the standard attention softmax and replace the exponential with a lookup-style approximation. They define a small design space with three knobs:

1. Calibration — where the grid of values sits. MinMax spreads K+1 grid points across the range actually observed in each row, so a wider learned range means a coarser grid. Fixed Window to Max (FWM) instead anchors a fixed 6-nat window below the row maximum and gives zero weight to scores further below it, so the grid width is τ/K regardless of the tail.

2. Reconstruction — how a score is turned into a weight. LERP draws straight lines between adjacent grid points (continuous); Nearest rounds each score to the nearest grid center (hard, so several scores share one weight).

3. Surrogate placement — where the straight-through replacement sits. Weight-STE applies it to the unnormalized weights before normalization; Prob-STE applies it to the normalized probabilities. Both produce the same mathematical hard forward but evaluate the normalization Jacobian at different points.

They derive the exact backward rules on paper, including derivatives through the row max and min, and check those closed forms numerically against automatic differentiation of the frozen training code. Then they run controlled pretraining experiments in which model, data, and optimizer protocol are matched and only the operator changes: four suites covering a 2×2×2 calibration × surrogate × K factorial at 124M parameters / 250M tokens, a long 2.5B-token run at 124M, an ordering probe at 1B parameters / 100M tokens on another corpus, and a full 5×5 grid varied across five seeds. Every checkpoint is evaluated on the full validation split of its corpus through its own frozen training code, with paired per-block differences reported in nats per token and bootstrap or t-intervals as appropriate. A separate geometry probe records row span, interval width, query/key norms, softmax entropy, and the total variation between the approximate and exact probability distributions at each checkpoint. A pre-declared caution margin of +0.01 nats is used to label single-seed contrasts as "separated"; the authors state this is not a significance or equivalence threshold.

Why This Matters

Impact on research. Most low-precision Transformer work quantizes linear layers and attention matmuls while deliberately leaving the softmax exponential in high precision. This paper argues that whether an approximate exponential is trainable depends jointly on the forward approximation and its backward rule, and that a backward which is wrong only in its calibration terms can look fine for tens of millions of tokens before failing. That reframes a number of "training is untested" statements in prior inference-side softmax work as an open empirical question, and it suggests that calibration statistics should not be silently kept off the autograd tape — a practice the paper notes is common in fake-quantization frameworks.

Real-world applications:

• Efficient training hardware and kernels — the operators studied here are the kind a custom attention datapath or lookup-table unit would implement; knowing which backward rules are safe informs what a training-capable kernel must support.
• Quantization-aware training pipelines — the finding that per-row calibration statistics need gradients (and that the max-dependent channel dominates) is directly relevant to how observers and step sizes are handled in QAT tooling.
• Low-precision model deployment and training stacks — teams shipping FP8/FP4 training recipes can use the calibration × surrogate interaction as a checklist item when deciding whether to push precision reduction into the normalization operator.
• Small-model research and reproducibility practice — the paper's protocol (matched operator-only comparisons, per-block paired statistics, frozen per-run evaluation code, and public manifests with code hashes) is a template for evaluating numerical operators rather than just end-to-end models.

Industry relevance. If attention softmax reduction can be moved to a narrow format without loss, the payoff is cheaper and faster pretraining. This paper claims no speed-up and measures no kernels — it deliberately studies training properties in fp32 attention arithmetic with TF32 matmuls — but it identifies which design choices would have to survive a bit-exact hardware implementation, and it shows a plausible failure mode (MinMax–Weight at K=4) that would be costly to discover late in a large run.

Future Directions

1. Identify the mechanism behind the range–resolution feedback. The authors state explicitly that no intervention in this paper identifies the mechanism linking wider intervals to degraded surrogate gradients to further span growth; an intervention that manipulates the geometry directly is the obvious next step.

2. Extend the strict-forward and seed evidence beyond one setting. The bit-exact forward retraining and the channel ablation each cover one setting (seed 7; two seeds), and suites A–C use a single seed. Repeating these across the full factorial, scales, and horizons would test whether the reported directions generalize.

3. Test whether the results hold under learning-rate tuning and longer schedules. No learning-rate sweep was run, so the 2.5B-token persistence of the gaps is specific to the tested cosine schedule; nothing exceeds 1B parameters or 2.5B tokens.

4. Re-examine τ and the tail policy. τ = 6 nats was fixed from an inference-time screen on frozen weights (whose final stage compared τ ∈ {5, 6, 7} on the WikiText-103 validation split later used for suites C and D) and is not shown to be optimal for training. The tail-policy comparison is single-seed and not cost-matched, and the FWM control does not isolate resolution from tail truncation.

5. Turn the graded downstream results into a general claim. Small NLL gaps do not establish equivalence on every benchmark, and the evaluation-time softmax substitution was run on corpora and probes but not on the benchmark battery.

Target Audience

Researchers and engineers working on low-precision or quantization-aware training of Transformers, on attention kernel and hardware design, and on numerical operator design more broadly. It is also relevant to practitioners who maintain QAT or fake-quantization frameworks and need to decide which statistics belong on the autograd tape, and to anyone evaluating whether a training-free operator replacement can be carried into pretraining. It requires comfort with attention mechanics, gradient derivations, and statistical reporting conventions; readers looking for an introductory treatment of quantization will find the paper's closed-form Jacobians and factorial designs demanding.

Authors’ abstract

Low-precision Transformer systems increasingly quantize attention matrix multiplications, while softmax often remains at higher precision. During pretraining, an approximate softmax changes the gradients that train the model as well as its forward computation. We study this interaction with K-interval attention, which approximates the exponential using K+1 grid values. We vary per-row grid calibration, interpolation versus hard rounding, and the placement of a straight-through surrogate relative to normalization. We derive the corresponding backward rules, including calibration derivatives, and compare these choices in pretraining experiments matched on model, data, and optimizer. Detaching the row extrema leaves the forward computation unchanged but produces a delayed increase in validation loss. With hard rounding at K=4, min-max calibration and a pre-normalization surrogate incur a large loss gap; changing either choice substantially reduces it. At 124M parameters and 2.5B training tokens, fixed-window calibration with a post-normalization surrogate yields a validation loss gap of +0.019 nats relative to softmax at K=4, and with a pre-normalization surrogate yields +0.004 nats at K=16.