# Pretraining Transformers with Quantized Softmax in Attention

arXiv: https://arxiv.org/abs/2609.33591 · HTML: https://mlsanity.com/paper/2609.33591/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 36 references that are in the map.
- First posted on arXiv 27 Sep 2026.

## Builds on (most cited first, 29 of 29)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Learned Step Size Quantization](https://mlsanity.com/paper/1902.08153/) (2019)
- [PACT: Parameterized Clipping Activation for Quantized Neural Networks](https://mlsanity.com/paper/1805.06085/) (2018)
- [I-BERT: Integer-only BERT Quantization](https://mlsanity.com/paper/2101.01321/) (2021)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration](https://mlsanity.com/paper/2410.02367/) (2024)
- [Trained Quantization Thresholds for Accurate and Efficient Fixed-Point Inference of Deep Neural Networks](https://mlsanity.com/paper/1903.08066/) (2019)
- [FP8-LM: Training FP8 Large Language Models](https://mlsanity.com/paper/2310.18313/) (2023)
- [Theory, Analysis, and Best Practices for Sigmoid Self-Attention](https://mlsanity.com/paper/2409.04431/) (2024)
- [Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformers](https://mlsanity.com/paper/2103.09301/) (2021)
- [A Study on ReLU and Softmax in Transformer](https://mlsanity.com/paper/2302.06461/) (2023)
- [Replacing softmax with ReLU in Vision Transformers](https://mlsanity.com/paper/2309.08586/) (2023)
- [Sparse Attention with Linear Units](https://mlsanity.com/paper/2104.07012/) (2021)
- [Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems](https://mlsanity.com/paper/2310.12956/) (2023)
- [ConSmax: Hardware-Friendly Alternative Softmax with Learnable Parameters](https://mlsanity.com/paper/2402.10930/) (2024)
- [ITA: An Energy-Efficient Attention and Softmax Accelerator for Quantized Transformers](https://mlsanity.com/paper/2307.03493/) (2023)
- [Towards Fully FP8 GEMM LLM Training at Scale](https://mlsanity.com/paper/2505.20524/) (2025)
- [SageBwd: A Trainable Low-bit Attention](https://mlsanity.com/paper/2603.02170/) (2026)
- [Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention](https://mlsanity.com/paper/2607.04422/) (2026)
- [EXAQ: Exponent Aware Quantization For LLMs Acceleration](https://mlsanity.com/paper/2410.03185/) (2024)
- [IntAttention: A Fully Integer Attention Pipeline for Efficient Edge Inference](https://mlsanity.com/paper/2511.21513/) (2025)
- [BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScaling](https://mlsanity.com/paper/2602.02071/) (2026)

## Abstract (arXiv)

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.