Softmax Reparameterization for Output-Head Quantization
cs.LG, cs.AI
Submitted: 2026-09-25
Updated: 2026-09-28
Comments: 32 pages, including appendix, ICLR 2027 submission
Code: https://github.com/ggml-org/llama.cpp
Project page: https://neurips2023-enlsp.github.io/papers/paper_47.pdf
License: http://creativecommons.org/licenses/by/4.0/
The gist: Large vocabularies make output heads a substantial inference cost in small language models.
Terminology
Abstract
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.
Sources
- QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs
- GaugeQuant: Online Learning of Quantization-Optimal Bases from LLM Symmetries
- SoftWater: Class-Aware Rate Allocation for Softmax Quantization
- Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning
- GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
- The Llama 3 Herd of Models
- Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling
- Efficient Memory Management for Large Language Model Serving with PagedAttention
- SpinQuant: LLM quantization with learned rotations
- Pointer Sentinel Mixture Models
- Qwen2.5 Technical Report
- Output Embedding Centering for Stable LLM Pretraining
- SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
- SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks