mlsanityData to 14 Oct 2026

Pretraining Transformers with Quantized Softmax in Attention

arXiv 2609.33591 ↗First posted 27 Sep 2026cs.LGcited by 0

Not on the map. Why?

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.

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

29 of 36 references.

References · 29

  1. 2016Categorical Reparameterization with Gumbel-Softmaxcited by 2,739
  2. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,079
  3. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 2,022
  4. 2022GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformerscited by 1,919
  5. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,386
  6. 2016The Concrete Distribution: A Continuous Relaxation of Discrete Random Variablescited by 1,148
  7. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 694
  8. 2019Learned Step Size Quantizationcited by 628
  9. 2018PACT: Parameterized Clipping Activation for Quantized Neural Networkscited by 467
  10. 2021I-BERT: Integer-only BERT Quantizationcited by 255
  11. 2019Understanding Straight-Through Estimator in Training Activation Quantized Neural Netscited by 231
  12. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 183
17 more
  1. 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 160
  2. 2019Trained Quantization Thresholds for Accurate and Efficient Fixed-Point Inference of Deep Neural Networkscited by 90
  3. 2023FP8-LM: Training FP8 Large Language Modelscited by 77
  4. 2024Theory, Analysis, and Best Practices for Sigmoid Self-Attentioncited by 64
  5. 2021Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformerscited by 54
  6. 2023A Study on ReLU and Softmax in Transformercited by 53
  7. 2023Replacing softmax with ReLU in Vision Transformerscited by 44
  8. 2021Sparse Attention with Linear Unitscited by 41
  9. 2023Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problemscited by 14
  10. 2024ConSmax: Hardware-Friendly Alternative Softmax with Learnable Parameterscited by 14
  11. 2023ITA: An Energy-Efficient Attention and Softmax Accelerator for Quantized Transformerscited by 12
  12. 2025Towards Fully FP8 GEMM LLM Training at Scalecited by 11
  13. 2026SageBwd: A Trainable Low-bit Attentioncited by 8
  14. 2026Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attentioncited by 8
  15. 2024EXAQ: Exponent Aware Quantization For LLMs Accelerationcited by 3
  16. 2025IntAttention: A Fully Integer Attention Pipeline for Efficient Edge Inferencecited by 3
  17. 2026BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScalingcited by 1

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers
  2. 2025Softpick: No Attention Sink, No Massive Activations with Rectified Softmax
  3. 2026Linearizing Vision Transformer with Test-Time Training
  4. 2026EFQ-Softmax: Exp-Free Quantization for Softmax
  5. 2023Towards Accurate Post-Training Quantization for Vision Transformer
  6. 2024From Attention to Activation: Unravelling the Enigmas of Large Language Models
  7. 2023Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing
  8. 2026The Quantization Benefits of Residual-Free Transformers