Pretraining Transformers with Quantized Softmax in Attention
Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski
University of Illinois Urbana-Champaign
Published September 27, 2026 · on the feed September 30, 2026 · ▲ 5 · arXiv 2609.33591
The 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.
The Brief
2:00A short podcast summary of the paper. Written and voiced by AI.
Transcript
Pretraining Transformers with Quantized Softmax in Attention. By Shangzhen Zhu and two others, from the University of Illinois Urbana-Champaign. Published on the twenty-seventh of September. No code repository is linked to it.
Low-precision Transformer systems increasingly quantize the matrix multiplications inside attention, while softmax often stays at higher precision. During pretraining, approximating softmax is more than a change to the forward pass. It also changes the gradients that train the model. The study's claim is that the way the approximation is built, and especially the way its backward pass is defined, decides how close pretraining stays to ordinary softmax.
The vehicle is K-interval attention, which approximates the exponential with K plus one grid values. Three choices are varied: per-row calibration of the grid, interpolation versus hard rounding, and where a straight-through surrogate sits relative to normalization. The authors derive the matching backward rules, including derivatives through the calibration, and compare the options in pretraining runs matched on model, data, and optimizer.
One result stands apart. Detaching the row extrema leaves the forward computation exactly as it was, yet validation loss rises after a delay. With hard rounding at K equals four, min-max calibration combined with a surrogate placed before normalization produces a large loss gap. Changing either one of those two choices shrinks the gap substantially.
The numbers come from a single scale: one hundred twenty-four million parameters trained on two point five billion tokens. There, fixed-window calibration with a surrogate after normalization lands zero point zero one nine nats of validation loss above softmax at K equals four. With the surrogate before normalization, the gap is zero point zero zero four nats at K equals sixteen.
A quantized softmax in pretraining turns out to be two decisions, not one. The forward approximation is the visible one. The gradient it implies is hidden, and a change that leaves the forward pass untouched can still surface later in the loss curve.
Paper Quest · 7 chapters · about 39 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
Cheaper softmax looks like a forward-pass problem. In pretraining, it is a gradient problem too.
You are pretraining a Transformer and want the attention softmax cheaper too. The plan: read the exponential off a small grid of values, calibrated for each row of scores.
The road ahead
- 1 · The problem ~6 min When you train with an approximate softmax, the gradient you attach to it decides what the model learns, not just the forward numbers.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min For an approximate softmax in pretraining, the backward rule is the learning rule, so missing gradient terms or poor design choices change what the model learns.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Fit a grid to each row, rebuild the exponential from it, add a surrogate gradient where needed, derive exact backward rules, then compare in matched runs.
- 6 · The experiment ~6 min You’ll guess: Detaching the row extrema in the approximate softmax leaves the forward computation exactly the same. What happened to validation loss during pretraining?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.