REVIEW 2 major objections 4 minor 69 references
SoftWater: Class-Aware Rate Allocation for Softmax Quantization
T0 review · 2 major / 4 minor · reviewed 2026-08-16 · deepseek-v4-flash
Pith's one-line read SoftWater shows that treating an LLM's softmax head as a KL rate-distortion problem lets a 2-bit head remove 45–60% of stored bytes at a few percent perplexity cost.
desk verdict A solid new idea for quantizing the softmax head under KL, with broad empirical support and a fixable theoretical gap around the dither assumption. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the softmax KL Hessian E_X[(diag(p)−pp^T)⊗XX^T], a Kn×Kn matrix coupling classes and features. A dither argument drops the off-diagonal −pp^T cross-class terms, and the separability assumption factors the remainder as λ̄⊗Σ_X, replacing the Kn×Kn Cholesky factorization with one n×n factorization rescaled per class. The per-class, per-column grid spacing α_{(k−1)n+i}=c_SW/(λ̄$_k^{{1/2}}$|ℓ_ii|) is the identity that carries the rate allocation: it gives fine grids to frequent, low-variance classes and coarse grids to rare ones, and it reduces to WaterSIC when the class distribution is uniform. A smoothing prior p̃=(1−ε)p+ε/K keeps uncovered classes inside the regime where the second-order Taylor expansion of KL is valid and interpolates between SoftWater and WaterSIC as ε approaches one. The same Cholesky factor powers successive interference cancellation at O($n^{3}$) cost, with both statistics accumulated in a single calibration forward pass.
What would settle it
Quantize a head on a model class where frequent tokens have systematically different hidden-state covariance from rare tokens, visible as a large Cov_X(λ_k(X), (Δ_k X)^2) in Eq. (19). Compute the relative proxy error ρ of Eq. (22); the paper reports ρ ≤ 0.10 on its five models, and the claim stands only if ρ stays near that level. A finding of ρ ≈ 0.25 or worse on a realistic model would show that the separability assumption, not the KL objective, is carrying the reported gains.
Extended reading notes
Core claim
The central claim is that the error metric induced by KL divergence on a softmax head is not the weighted mean-squared error weight I⊗Σ_X but E_X[(diag(p)−pp^T)⊗XX^T], so rate should be allocated along both a feature axis and a class axis. With a separability assumption that λ and XX^T are uncorrelated, the error weight becomes λ̄⊗Σ_X, and the resulting lattice grid spacing is α_{(k−1)n+i}=c_SW/(λ̄$_k^{{1/2}}$|ℓ_ii|), inversely proportional to class curvature and to the feature-side Cholesky diagonal. SoftWater measures λ̄_k=E_X[p_k(1−p_k)] from the same calibration forward pass as Σ_X, smooths p with a uniform prior to cap rare-class spacing, and encodes with successive interference cancellation. The paper reports that this allocation reduces head-induced KL by 6.5–8.3 times versus the released WaterSIC head at 2 bits and makes a 2-bit tied head practical on fully quantized models.
Load-bearing premise
The whole method rests on assuming that a class's average softmax curvature is uncorrelated with the second-order geometry of the hidden states feeding the head, so the full error matrix factors into one class-side vector and one feature-side matrix; the paper measures this proxy error at no more than 10% on the five tested models but does not claim a guarantee in general.
Editorial extensions
If this is right
- At matched head rates, SoftWater cuts head-induced KL by 6.5–8.3 times on WikiText-2 at 2 bits and beats the released WaterSIC head on 59 of 60 head-rate/model test points.
- A 2-bit SoftWater head on Llama-3.2-1B-Instruct with quantized bodies removes 45–60% of stored bytes for a 2.9–3.7% perplexity increase, and a 4-bit head costs 0.2–0.3% while still cutting file size by at least 38%.
- SoftWater nearly halves the zero-shot accuracy lost to a 2-bit head on every tested quantized body, with the largest gains on LAMBADA, the task scored by next-token prediction.
- Calibrating the class-side statistic on the deployment domain gives the lowest KL on that domain at every rate, so the head can be targeted at a specific corpus.
- The scheme contains WaterSIC as the uniform-class limit, so the standard WMSE treatment of the head is recovered when all classes are assumed equally likely.
Reading between the lines
- Editorial inference: the same per-class rate allocation should transfer to any linear-softmax layer with a skewed output distribution, such as MoE routers with unbalanced expert utilization or speech and vision classifiers with Zipfian class frequencies, since the method needs only a class-frequency statistic from calibration data.
- Editorial inference: on tied-embedding models the head weights double as the input embedding, so the class-aware grids affect both directions of the model; the reported gains suggest an embedding-aware allocation could improve input-side quantization as well, although that is not tested in the paper.
- Editorial inference: because the class-side statistic is a context frequency rather than a token count, increasing the number of distinct calibration contexts helps more than increasing total tokens; a testable corollary is that chunking calibration data into shorter sequences should yield head quantizers that transfer further off domain.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes SoftWater, a post-training quantization method for the final softmax layer (the head) of LLMs. It treats head quantization as a rate-distortion problem under the KL divergence between the original and quantized output distributions, derives a second-order expansion, and approximates the resulting Kn x Kn Hessian by a separable, class-diagonal surrogate. The method allocates lattice spacing per class and per feature according to alpha = c_SW / (tilde_lambda_k^{1/2} |ell_ii|), with class scales estimated from a single calibration forward pass and a smoothing prior that caps rare-class spacings. Experiments on five models (1B to 32B) compare SoftWater with the released WaterSIC quantizer at matched rates, in head-only and whole-model settings, and also include domain-targeted calibration, zero-shot accuracy, statistic ablations, and a separability diagnostic. The paper claims 59 of 60 head-only wins, 6.5x to 8.3x KL reduction at 2 bits on WikiText2, and 45% to 60% stored-byte reduction on a 2-bit tied head for about 3% to 4% perplexity increase.
Significance. Softmax heads occupy 15% to 30% of parameters in small LLMs and are normally left in high precision, so a method that quantizes them under the correct output metric is practically significant. If the claims hold, SoftWater is a useful contribution: a 2-bit head becomes close to lossless, the gains are attributed to class-aware allocation rather than to pipeline refinements, and the comparison baseline is the released WaterSIC quantizer. The paper has notable strengths: rates are matched to within 0.005 bits, the baseline is an existing released method, the evaluation spans five models and multiple datasets, and the authors explicitly acknowledge that the separability surrogate is not guaranteed in general. The main gaps are the unmeasured cross-class correlation terms in the dither argument and the fact that the separability validation is shown for only one model despite the broader claim in the text.
major comments (2)
- [§4.3, Eq. (8)] The reduction from the full KL Hessian (diag(p)-pp^T) tensor XX^T to the diagonal lambda tensor XX^T is justified by a dither argument in which quantization errors are uncorrelated across rows, so the cross-class terms -p_k p_l drop out. The dither is then dropped from the deployed scheme with the assertion that undithered errors already behave as if uncorrelated across rows. This step is load-bearing because the per-class allocation in Eq. (14) is derived entirely from tilde_lambda_k, and no measurement in the paper supports the undithered assertion. The diagnostic in Table 6a defines Q_k and S_k using only the diagonal class-side curvatures, so it cannot detect cross-class correlations. Please provide a direct measurement on the actual quantized Delta, for example the normalized contribution sum_{k != l} E_X[-p_k p_l (Delta_k X)(Delta_l X)] / sum_k S_k, or otherwise justify the step. Without this, the claim that the gains follow from a principled KL-optimal allocation is not fully supported, although the empirical wins over WaterSIC stand independently.
- [§5.6, Table 6a; §4.3] The text in §4.3 states that the separability proxy underestimates the true distortion by at most 10% across models, rates, and datasets, but Table 6a is titled 'Separability diagnostics on Llama-3.2-1B' and reports only that one model, with rates 2, 3, and 4 on WT2 and C4. Because Assumption 1 is what reduces the Kn x Kn factorization into one n x n Cholesky factorization with class rescaling, the breadth of this validation matters for the generality of the theoretical claim. Please add the same diagnostic for the remaining four models, or restrict the claim in the text to the model actually measured and state explicitly that the other models are currently untested on this diagnostic.
minor comments (4)
- [Abstract and §5.1, Table 2] The abstract's statement that SoftWater cuts head-induced KL by 6.5x to 8.3x at 2 bits is not qualified by dataset; the C4 ratios in Table 2 are 3.5x to 5.1x, so the unqualified range should specify WikiText2.
- [§5.3, Table 10] The claim that matching calibration to the deployment domain gives the lowest KL on that domain throughout is partly by construction, because KL is measured on the same domain used for calibration; the calibration and evaluation slices are disjoint, but a held-out slice of the same domain or a mismatched-domain analysis would make the claim more informative.
- [§5, baseline description] The paper describes the baseline as 'the full released WaterSIC pipeline minus finetuning' while SoftWater uses none of that pipeline's refinements; spelling out that the baseline retains the orthogonal refinements while SoftWater does not would make the comparison easier to interpret.
- [§5.6, Eq. (22)] The 'Bits' row in Table 6a is defined by the formula (1/2) log_2(1+rho) but the derivation of this high-rate bit penalty is stated only briefly; a one-sentence derivation or a pointer to WaterSIC's rate analysis would help the reader verify the 0.07-bit penalty.
Circularity Check
No significant circularity: the central comparison is against the external released WaterSIC baseline at matched rates, and SoftWater's derivation is self-contained.
full rationale
The headline result is an empirical comparison against an externally released quantizer (WaterSIC) at matched bit rates, so it is not self-referential. The derivation chain is self-contained: Eqs. (5)-(7) obtain the softmax KL Hessian from a second-order expansion; Eq. (8) uses a subtractive-dither argument to drop cross-class terms; Assumption 1 (Eq. 11) is an explicit separability assumption whose induced error is bounded in Eqs. (18)-(19) and measured in §5.6. The grid spacing (Eq. 14) follows from applying the same high-rate SIC analysis as WaterSIC to the surrogate lambda_bar ⊗ Sigma_X, and c_SW is a rate knob tuned by secant search, not fit to the KL or perplexity metrics. The reported KL and perplexity numbers are measured on the quantized models, not computed from the proxy. No load-bearing result is justified by a self-citation: the authors cite no prior work of their own, and the WaterSIC/SIC machinery is used as an external baseline and encoder. The domain-targeting result (§5.3) is not circular: it is an empirical crossover table, and while calibrating on a domain naturally biases the class statistic toward that domain, the claim is supported by the reported mismatched-arm comparisons. The main caveat is that Eq. (8)'s assertion that undithered errors behave as uncorrelated across rows is not directly validated, since Table 6a only tests the separability Assumption 1 via the diagonal lambda_k, but that is a correctness or evidence gap, not a circular reduction. There is no step where a fitted parameter is renamed a prediction or where an input is equivalent to the output by construction.
Assumptions & free parameters
free parameters (3)
- epsilon (smoothing prior weight) =
0.1 (default)
- c_SW (lattice density) =
set by secant search to match target bit rate
- delta (Cholesky damping) =
1e-6 * mean(diag(Sigma_X))
assumptions (4)
- domain assumption Dither ensemble makes quantization errors zero-mean and uncorrelated across rows, so cross-class terms -p_k p_l contribute nothing to expected distortion
- domain assumption Assumption 1 (Separability): E_X[lambda tensor X X^T] = E_X[lambda] tensor E_X[X X^T]
- domain assumption Second-order Taylor expansion of D_KL at delta=0 is accurate over the perturbations produced by the quantizer
- standard math Softmax KL Hessian at delta=0 is diag(p)-pp^T and log-sum-exp Taylor expansion
Cite this review
Pith. "Pith review of SoftWater: Class-Aware Rate Allocation for Softmax Quantization." pith.science (2026). https://pith.science/paper/AO4RZQ25
@misc{pith2026260812026,
author = {Pith},
title = {Pith review of: SoftWater: Class-Aware Rate Allocation for Softmax Quantization},
year = {2026},
howpublished = {\url{https://pith.science/paper/AO4RZQ25}},
note = {Machine review of arXiv:2608.12026}
}
abstract
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\times Kn$ Cholesky with one $n\times 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\times$--$8.3\times$ 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.
Reference graph
Works this paper leans on
-
[1]
doi:10.48550/arXiv.2604.02638 , abstract =
Nishida, Yasushi , month = apr, year =. doi:10.48550/arXiv.2604.02638 , abstract =
-
[2]
doi:10.48550/arXiv.2504.19874 , abstract =
Zandieh, Amir and Daliri, Majid and Hadian, Majid and Mirrokni, Vahab , month = apr, year =. doi:10.48550/arXiv.2504.19874 , abstract =
-
[3]
doi:10.48550/arXiv.2406.03482 , abstract =
Zandieh, Amir and Daliri, Majid and Han, Insu , month = jul, year =. doi:10.48550/arXiv.2406.03482 , abstract =
-
[4]
Katharopoulos, Angelos and Vyas, Apoorv and Pappas, Nikolaos and Fleuret, François , month = aug, year =. Transformers are. doi:10.48550/arXiv.2006.16236 , abstract =
-
[5]
Dao, Tri and Gu, Albert , month = may, year =. Transformers are. doi:10.48550/arXiv.2405.21060 , abstract =
-
[6]
Gu, Albert and Dao, Tri , month = may, year =. Mamba:. doi:10.48550/arXiv.2312.00752 , abstract =
-
[7]
Birnick, Johann , month = mar, year =. The. doi:10.48550/arXiv.2508.01077 , abstract =
-
[8]
Federici, Marco and Breugel, Boris van and Whatmough, Paul and Nagel, Markus , month = mar, year =. Dissecting. doi:10.48550/arXiv.2603.04359 , abstract =
Show all 69 references
-
[9]
doi:10.48550/arXiv.2511.04063 , abstract =
Shao, Yuantian and Chen, Yuanteng and Wang, Peisong and Yu, Jianlin and Lin, Jing and Yao, Yiwu and Wei, Zhihui and Cheng, Jian , month = nov, year =. doi:10.48550/arXiv.2511.04063 , abstract =
-
[10]
doi:10.48550/arXiv.2511.10645 , abstract =
Liang, Yesheng and Chen, Haisheng and Zhang, Zihan and Han, Song and Liu, Zhijian , month = feb, year =. doi:10.48550/arXiv.2511.10645 , abstract =
-
[11]
doi:10.48550/arXiv.2405.16406 , abstract =
Liu, Zechun and Zhao, Changsheng and Fedorov, Igor and Soran, Bilge and Choudhary, Dhruv and Krishnamoorthi, Raghuraman and Chandra, Vikas and Tian, Yuandong and Blankevoort, Tijmen , month = feb, year =. doi:10.48550/arXiv.2405.16406 , abstract =
-
[12]
doi:10.13140/RG.2.2.28167.37282 , abstract =
Liu, Zirui and Yuan, Jiayi and Jin, Hongye and Zhong, Shaochen and Xu, Zhaozhuo and Braverman, Vladimir and Chen, Beidi and Hu, Xia , year =. doi:10.13140/RG.2.2.28167.37282 , abstract =
-
[13]
and Massaroli, Stefano and Moro, Alessandro and Smith, Jimmy T
Parnichkun, Rom N. and Massaroli, Stefano and Moro, Alessandro and Smith, Jimmy T. H. and Hasani, Ramin and Lechner, Mathias and An, Qi and Ré, Christopher and Asama, Hajime and Ermon, Stefano and Suzuki, Taiji and Yamashita, Atsushi and Poli, Michael , month = jun, year =. St...
-
[14]
and Li, Bo and Cameron, Pashmina and Jaggi, Martin and Alistarh, Dan and Hoefler, Torsten and Hensman, James , month = oct, year =
Ashkboos, Saleh and Mohtashami, Amirkeivan and Croci, Maximilian L. and Li, Bo and Cameron, Pashmina and Jaggi, Martin and Alistarh, Dan and Hoefler, Torsten and Hensman, James , month = oct, year =. doi:10.48550/arXiv.2404.00456 , abstract =
- [15]
- [16]
- [17]
- [18]
- [19]
-
[20]
Alignment-
Wee, Sunghyun and Kim, Suyoung and Kim, Hyeonjin and Hwang, Kyomin and Kwak, Nojun , month = jan, year =. Alignment-. doi:10.48550/arXiv.2511.07842 , abstract =
- [21]
-
[22]
and Colombo, Maurizio and Damiani, Ernesto and Asal, Rasool and Almemari, Al Anoud and Alhammadi, Yousof , month = dec, year =
Al Qassem, Lamees M. and Colombo, Maurizio and Damiani, Ernesto and Asal, Rasool and Almemari, Al Anoud and Alhammadi, Yousof , month = dec, year =. Quantization in. 2024. doi:10.1109/CloudCom62794.2024.00020 , urldate =
2024
- [23]
- [24]
- [25]
- [26]
- [27]
- [28]
-
[29]
Pretraining
NVIDIA and Abecassis, Felix and Agrusa, Anjulie and Ahn, Dong and Alben, Jonah and Alborghetti, Stefania and Andersch, Michael and Arayandi, Sivakumar and Bjorlin, Alexis and Blakeman, Aaron and Briones, Evan and Buck, Ian and Catanzaro, Bryan and Choi, Jinhang and Chrzanowski...
- [30]
-
[31]
Guo, Hang and Li, Yawei and Benini, Luca , month = sep, year =. Optimal. doi:10.48550/arXiv.2509.11177 , abstract =
- [32]
- [33]
- [34]
- [35]
- [36]
- [37]
- [38]
- [39]
- [40]
-
[41]
and Bich, Philippe and Zhuang, Jiawei and Çelik, Ahmet and Benfenati, Luca and Cavigelli, Lukas , month = jan, year =
Müller, Lorenz K. and Bich, Philippe and Zhuang, Jiawei and Çelik, Ahmet and Benfenati, Luca and Cavigelli, Lukas , month = jan, year =. doi:10.48550/arXiv.2509.22944 , abstract =
- [42]
- [43]
- [44]
- [45]
- [46]
- [47]
- [48]
- [49]
- [50]
- [51]
- [52]
-
[53]
doi:10.48550/arXiv.2603.14591 , abstract =
Tranheden, Wilhelm and Ahmed, Shahnawaz and Dubhashi, Devdatt and Matthiesen, Jonna and Essen, Hannes von , month = mar, year =. doi:10.48550/arXiv.2603.14591 , abstract =
- [54]
- [55]
- [56]
- [57]
-
[58]
Jean, Sébastien and Cho, Kyunghyun and Memisevic, Roland and Bengio, Yoshua , year =. On. Proceedings of the 53rd. doi:10.3115/v1/P15-1001 , language =
- [59]
- [60]
-
[61]
and Song, Hyun Oh , month = sep, year =
Kim, Jinuk and Halabi, Marwa El and Park, Wonpyo and Schaefer, Clemens JS and Lee, Deokjae and Park, Yeonhong and Lee, Jae W. and Song, Hyun Oh , month = sep, year =. doi:10.48550/arXiv.2505.07004 , abstract =
-
[62]
Human behavior and the principle of least effort
Zipf, George Kingsley , year =. Human behavior and the principle of least effort. , abstract =
-
[63]
Psychonomic Bulletin & Review , author =
Zipf’s word frequency law in natural language:. Psychonomic Bulletin & Review , author =. 2014 , pages =. doi:10.3758/s13423-014-0585-6 , abstract =
2014 doi
- [64]
- [65]
- [66]
- [67]
-
[68]
Zhang, Shihao and Zhang, Haoyu and Colbert, Ian and Saab, Rayan , month = feb, year =. Qronos:. doi:10.48550/arXiv.2505.11695 , abstract =
-
[69]
Journal of Computational and Applied Mathematics , author =
The ubiquitous. Journal of Computational and Applied Mathematics , author =. 2000 , pages =. doi:10.1016/S0377-0427(00)00393-9 , language =
2000 doi
Reviewed August 16, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.