# SoftWater: Class-Aware Rate Allocation for Softmax Quantization

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 24 references that are in the map.
- First posted on arXiv 12 Aug 2026.

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

- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [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)
- [SpinQuant: LLM quantization with learned rotations](https://mlsanity.com/paper/2405.16406/) (2024)
- [OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models](https://mlsanity.com/paper/2308.13137/) (2023)
- [QuIP: 2-Bit Quantization of Large Language Models With Guarantees](https://mlsanity.com/paper/2307.13304/) (2023)
- [QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks](https://mlsanity.com/paper/2402.04396/) (2024)
- [Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning](https://mlsanity.com/paper/2208.11580/) (2022)
- [QTIP: Quantization with Trellises and Incoherence Processing](https://mlsanity.com/paper/2406.11235/) (2024)
- [On Using Very Large Target Vocabulary for Neural Machine Translation](https://mlsanity.com/paper/1412.2007/) (2014)
- [GuidedQuant: Large Language Model Quantization via Exploiting End Loss Guidance](https://mlsanity.com/paper/2505.07004/) (2025)
- [Qronos: Correcting the Past by Shaping the Future... in Post-Training Quantization](https://mlsanity.com/paper/2505.11695/) (2025)
- [NestQuant: Nested Lattice Quantization for Matrix Products and LLMs](https://mlsanity.com/paper/2502.09720/) (2025)
- [The Geometry of LLM Quantization: GPTQ as Babai's Nearest Plane Algorithm](https://mlsanity.com/paper/2507.18553/) (2025)
- [DartQuant: Efficient Rotational Distribution Calibration for LLM Quantization](https://mlsanity.com/paper/2511.04063/) (2025)
- [CSV-Decode: Certifiable Sub-Vocabulary Decoding for Efficient Large Language Model Inference](https://mlsanity.com/paper/2511.21702/) (2025)
- [WaterSIC: Information-Theoretically (Near) Optimal Linear Layer Quantization](https://mlsanity.com/paper/2603.04956/) (2026)
- [VQ-Logits: Compressing the Output Bottleneck of Large Language Models via Vector Quantized Logits](https://mlsanity.com/paper/2505.10202/) (2025)
- [FlashHead: Efficient Drop-In Replacement for the Classification Head in Language Model Inference](https://mlsanity.com/paper/2603.14591/) (2026)
- [LFQ: Logit-aware Final-block Quantization for Boosting the Generation Quality of Low-Bit Quantized LLMs](https://mlsanity.com/paper/2605.29756/) (2026)

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

- [Softmax Reparameterization for Output-Head Quantization](https://mlsanity.com/paper/2609.31291/) (2026)
- [WaterKron and FlipFlop Hessian: Information-Theoretically Grounded Quantization with Kronecker-factored Hessians](https://mlsanity.com/paper/2609.14706/) (2026)

## Abstract (arXiv)

Post-training quantization pipelines routinely leave the softmax output layer in high precision. Yet in small LLMs with modern vocabularies, the head holds 15--30% of all parameters, so a nominal ``2-bit'' model with an fp16 head can store several times as many bits per weight. We pose softmax-layer quantization as a rate-distortion problem under the KL divergence between the original and quantized output distributions. A second-order analysis reveals a class-aware geometry: quantization error is weighted jointly by feature covariance and class-specific softmax curvature. A separability approximation replaces the Kn× Kn Cholesky with one n× n factorization rescaled per class, making the lattice encodable by successive interference cancellation, with both statistics from a single forward pass. The resulting method, SoftWater, gives fine grids to frequent, low-variance classes and coarse grids to rare ones, a large gap under Zipfian token distributions. Across five models from 1B to 32B, SoftWater outperforms the released WaterSIC quantizer (near-optimal under linear-layer WMSE but not output KL) at matched head rates on 59 of 60 test points, using none of that pipeline's refinements and cutting head-induced KL by 6.5×--8.3× at 2 bits. On Llama-3.2-1B-Instruct with quantized bodies, a 2-bit head removes 45--60% of stored bytes for a 2.9--3.7% perplexity increase. Because the class-side statistic comes from calibration data, matching calibration to the deployment domain gives the lowest KL on that domain throughout. On a tied model, a 4-bit head is near-lossless and a 2-bit head costs under 4% perplexity, making head quantization of such models practical.