Softmax Reparameterization for Output-Head Quantization
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 K...
| Authors | Asim Kadav, Christian Flores, Chirag Arora et al. |
|---|---|
| Primary category | cs.LG |
| Published | 25 September 2026 |
| Last updated at source | 28 September 2026 |
Silytics shows public metadata from arXiv and links to the original. Information may be incomplete or out of date. Last refreshed 17 October 2026. Licenses and terms of the original apply to the work itself. Not financial, legal or professional advice.