Pretraining Transformers with Quantized Softmax in Attention
K-interval quantized softmax nearly matches full-softmax pretraining loss when calibration and the gradient surrogate are chosen carefully.
The paper studies Transformer pretraining when softmax is replaced by K-interval attention, which approximates the exponential with K+1 grid values and therefore changes both the forward pass and the gradients. Detaching row extrema leaves the forward computation unchanged but later increases validation loss. At 124M parameters and 2.5B tokens, fixed-window calibration with a post-normalization surrogate trails full softmax by +0.019 nats at K=4, while a pre-normalization surrogate yields a +0.004 nat gap at K=16.
- K-interval attention approximates softmax exponentials with K+1 calibrated grid values.
- Hard rounding at K=4 with min-max calibration causes a large loss gap.
- Best K=4 setup trails softmax by +0.019 nats on a 124M model.
- At K=16, a pre-normalization surrogate narrows the gap to +0.004 nats.
From the original article
1 min read at huggingface.co
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.
By Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.