Transformers pretrained using low precision softmax approximations in attention
Pretraining Transformers with Quantized Softmax in Attention
Machine Learning
Summary
Transformers use a softmax function in their attention mechanism, usually computed with high precision, which can slow down training and increase resource use. The authors explore a low-precision method called K-interval attention that approximates this softmax using a limited set of values, impacting both how the model learns and makes predictions. They test different ways to calibrate and compute this approximation during training, finding setups that minimize accuracy loss. Their results show that quantizing softmax at certain levels can nearly match the performance of full softmax while reducing precision.
What this means in practice
- •For machine learning engineers: Train Transformer models using quantized softmax in attention to reduce compute requirements with minor loss in validation accuracy.
- •For hardware accelerator designers: Design low-precision hardware for Transformer attention matrix with effective softmax approximation methods from this study to improve efficiency.
Authors
Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski
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.