# Softmax Reparameterization for Output-Head Quantization

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 32 references that are in the map.
- First posted on arXiv 25 Sep 2026.

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

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [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)
- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs](https://mlsanity.com/paper/2404.00456/) (2024)
- [Up or Down? Adaptive Rounding for Post-Training Quantization](https://mlsanity.com/paper/2004.10568/) (2020)
- [BRECQ: Pushing the Limit of Post-Training Quantization by Block Reconstruction](https://mlsanity.com/paper/2102.05426/) (2021)
- [SpinQuant: LLM quantization with learned rotations](https://mlsanity.com/paper/2405.16406/) (2024)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning](https://mlsanity.com/paper/2208.11580/) (2022)
- [SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration](https://mlsanity.com/paper/2410.02367/) (2024)
- [OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text](https://mlsanity.com/paper/2310.06786/) (2023)
- [Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies](https://mlsanity.com/paper/2407.13623/) (2024)
- [OstQuant: Refining Large Language Model Quantization with Orthogonal and Scaling Transformations for Better Distribution Fitting](https://mlsanity.com/paper/2501.13987/) (2025)
- [Efficient softmax approximation for GPUs](https://mlsanity.com/paper/1609.04309/) (2016)
- [Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling](https://mlsanity.com/paper/2501.16975/) (2025)
- [LeanQuant: Accurate and Scalable Large Language Model Quantization with Loss-error-aware Grid](https://mlsanity.com/paper/2407.10032/) (2024)
- [Sparse Fine-tuning for Inference Acceleration of Large Language Models](https://mlsanity.com/paper/2310.06927/) (2023)
- [Output Embedding Centering for Stable LLM Pretraining](https://mlsanity.com/paper/2601.02031/) (2026)
- [SoftWater: Class-Aware Rate Allocation for Softmax Quantization](https://mlsanity.com/paper/2608.12026/) (2026)
- [GaugeQuant: Online Learning of Quantization-Optimal Bases from LLM Symmetries](https://mlsanity.com/paper/2607.20757/) (2026)
- [ARCHead: Activation-Metric Residual Correction for Large Language Model Output Heads](https://mlsanity.com/paper/2608.02703/) (2026)

## Built on by (most cited first, 1 of 1)

- [Logit Leakage through Floating-Point Bias Addition](https://mlsanity.com/paper/2610.16790/) (2026)

## Abstract (arXiv)

Large vocabularies make output heads a substantial inference cost in small language models. We introduce softmax reparameterization, a post-training method that searches over functionally equivalent output heads before quantization. The method subtracts a scalar multiple of the vocabulary-row mean from every output row and selects the coefficient by validation KL. For linear-softmax heads, these shifts preserve full-precision predictions exactly and require no decoder retraining; a rank-one correction extends the construction to nonlinear logit paths. Across seven output heads and three quantizers, W4 gains are largest where baseline quantization substantially distorts predictions: test KL falls by 93% on XGLM under RTN and by 73--77% on Phi, BLOOM, and BLOOMZ under activation-weighted MSE. Heads with low baseline error change little; at W2, used as a compression stress test, benefits extend more broadly. On Phi, the gains persist under stronger GPTQ calibration; a separate untouched holdout reproduces the improvements on Phi and BLOOM. Frozen WikiText-selected coefficients also transfer without retuning to C4 and OpenWebMath. Residual analysis on Phi shows how fidelity can improve despite greater total logit error: the selected representative reduces error on likely outputs and lowers its Fisher-weighted cost. For shift-compatible heads, the shift adds no inference operation. With the decoder held in BF16, a packed W4 Phi output head reduces batch-one generation latency by 10.8%, and reparameterization preserves this speedup.