mlsanityData to 14 Oct 2026

Softmax Reparameterization for Output-Head Quantization

arXiv 2609.31291 ↗First posted 25 Sep 2026cs.LG cs.AIcited by 1

Not on the map. Why?

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.

Key facts. Cited by 1 paper 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

27 of 32 references.

References · 27

  1. 2024The Llama 3 Herd of Modelscited by 14,715
  2. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
  3. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
  4. 2017On Calibration of Modern Neural Networkscited by 4,960
  5. 2024Qwen2.5 Technical Reportcited by 4,942
  6. 2016Pointer Sentinel Mixture Modelscited by 3,096
  7. 2022GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformerscited by 1,919
  8. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,386
  9. 2023AWQ: Activation-aware Weight Quantization for LLM Compression and Accelerationcited by 1,298
  10. 2024QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMscited by 634
  11. 2020Up or Down? Adaptive Rounding for Post-Training Quantizationcited by 621
  12. 2021BRECQ: Pushing the Limit of Post-Training Quantization by Block Reconstructioncited by 479
15 more
  1. 2024SpinQuant: LLM quantization with learned rotationscited by 471
  2. 2014New insights and perspectives on the natural gradient methodcited by 443
  3. 2022Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruningcited by 322
  4. 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 160
  5. 2023OpenWebMath: An Open Dataset of High-Quality Mathematical Web Textcited by 155
  6. 2024Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabulariescited by 118
  7. 2025OstQuant: Refining Large Language Model Quantization with Orthogonal and Scaling Transformations for Better Distribution Fittingcited by 97
  8. 2016Efficient softmax approximation for GPUscited by 67
  9. 2025Over-Tokenized Transformer: Vocabulary is Generally Worth Scalingcited by 44
  10. 2024LeanQuant: Accurate and Scalable Large Language Model Quantization with Loss-error-aware Gridcited by 22
  11. 2023Sparse Fine-tuning for Inference Acceleration of Large Language Modelscited by 14
  12. 2026Output Embedding Centering for Stable LLM Pretrainingcited by 4
  13. 2026SoftWater: Class-Aware Rate Allocation for Softmax Quantizationcited by 2
  14. 2026GaugeQuant: Online Learning of Quantization-Optimal Bases from LLM Symmetriescited by 1
  15. 2026ARCHead: Activation-Metric Residual Correction for Large Language Model Output Headscited by 1

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026Logit Leakage through Floating-Point Bias Additioncited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025SiLQ: Simple Large Language Model Quantization-Aware Training
  2. 2025Enhancing Ultra-Low-Bit Quantization of Large Language Models Through Saliency-Aware Partial Retraining
  3. 2025VQ-Logits: Compressing the Output Bottleneck of Large Language Models via Vector Quantized Logits
  4. 2025An Extra RMSNorm is All You Need for Fine Tuning to 1.58 Bits
  5. 2024Exploring the Trade-Offs: Quantization Methods, Task Difficulty, and Model Size in Large Language Models From Edge to Giant
  6. 2023SqueezeLLM: Dense-and-Sparse Quantization
  7. 2026Variable Bit-width Quantization: Learning Per-Group Precision for "Bigger-but-Smaller" Language Models
  8. 2023Softmax Bias Correction for Quantized Generative Models