{"id":"1f3f52bf-5a82-4638-b94f-731c654404e0","arxiv_id":"2608.12026","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"SoftWater, a KL-divergence-based quantizer for LLM softmax heads, allocates bit rate by class frequency and variance and beats WaterSIC at matched head rates on 59 of 60 test points.","lead":"This paper shows how to quantize the softmax output layer of small language models to 2-4 bits by allocating more bits to frequent, low-variance token classes, cutting head-induced KL error several-fold versus the prior WMSE-based quantizer. A generalist should read it because the head can be 15-30% of parameters, so quantizing it can remove up to 60% of stored bytes with only a few percent perplexity increase.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The dither-based removal of cross-class KL terms (Eq. 8) is never empirically validated; Table 6a's rho tests only the diagonal separable proxy, leaving the core class-aware allocation without direct support.","rationale":"The paper's central empirical claim — that SoftWater outperforms the released WaterSIC quantizer at matched head rates on 59/60 test points — is well supported by the reported tables across five models and multiple rates. The comparison is fair in design: both arms are layerwise without finetuning, rates are matched to within 0.005 bits, and KL is measured on identical token streams. The reader's conditional verdict is appropriate, and the lack of code/data and the WT2-versus-C4 gap in the abstract are legitimate but secondary concerns. Among the theoretical assumptions, the reader identifies separability (Assumption 1) as the weakest, and the paper does provide a partial validation via rho <= 0.10 on one model. However, the dither argument in Eq. (8) is a more load-bearing gap: it is the step that removes the off-diagonal cross-class terms of the KL Hessian, yet it is supported only by an unquantified empirical assertion. The rho diagnostic in Table 6a does not test this step, because both Q_k and S_k use only the diagonal lambda_k. If the undithered SIC errors are correlated across classes, the proxy that justifies the per-class allocation could be substantially wrong, and the method's success would rest on an unexamined property of the errors rather than on the derived KL geometry. The proposed concrete test directly measures the off-diagonal contribution to the second-order KL proxy for the actual deployed quantizer, settling whether the dither assumption lands. Until that check is run, the conditional verdict stands, but the theoretical justification should be read as heuristic rather than established.","tokens_in":24942,"tokens_out":13120,"duration_ms":120165,"concrete_test":"On Llama-3.2-1B with the SoftWater 2-bit head from §5.1, compute the full second-order KL proxy T_full = E_X[Sigma_k p_k(1-p_k) delta_k^2 - 2 Sigma_{k<l} p_k p_l delta_k delta_l] and the diagonal proxy T_diag = E_X[Sigma_k tilde_lambda_k delta_k^2] using the actual undithered logit perturbations delta = Delta X on the WT2 calibration stream. If |T_full - T_diag|/T_full <= 0.10, the dither assumption lands; if the off-diagonal contribution exceeds 10%, the class-aware allocation's theoretical basis is unsupported and the paper's derivation needs revision.","verdict_should_be":"UNCHANGED","load_bearing_attack":"In §4.3, Eq. (8) replaces the full softmax KL Hessian (diag(p)-pp^T) ⊗ XX^T with the diagonal lambda ⊗ XX^T under a dither argument, and the text asserts that 'empirically, undithered errors already behave as if uncorrelated across rows, so we drop it from the final scheme.' No measurement supports this assertion. The validation in §5.6 (Table 6a) defines rho via Q_k = 1/2 Delta_k E[lambda_k(X) XX^T] Delta_k^T and S_k = 1/2 tilde_lambda_k Delta_k Sigma_X Delta_k^T; both use only the diagonal lambda_k, so neither includes the off-diagonal terms -p_k p_l of the true Hessian. The deployed SIC quantizer produces undithered errors whose cross-class correlations are unmeasured. If these correlations are non-negligible, the true KL is not captured by the proxy, and the per-class rate allocation of Eq. (14) — derived entirely from lambda_k — may be far from optimal for the actual KL, weakening the paper's claim that the gains follow from a principled KL-optimal allocation.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":25105,"tokens_out":12952,"duration_ms":128454,"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":[{"comment":"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.","section":"§4.3, Eq. (8)"},{"comment":"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.","section":"§5.6, Table 6a; §4.3"}],"minor_comments":[{"comment":"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.","section":"Abstract and §5.1, Table 2"},{"comment":"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.","section":"§5.3, Table 10"},{"comment":"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.","section":"§5, baseline description"},{"comment":"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.","section":"§5.6, Eq. (22)"}],"recommendation":"major_revision","confidential_remarks":"The empirical comparison is solid and the method is practically promising, but the theoretical reduction needs one additional measurement (cross-class correlation terms) and broader separability diagnostics before I would recommend acceptance. The domain-targeting claim, while partly by construction, is not a reason to reject."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nQuick take on arXiv:2608.12026 (SoftWater). It's a solid paper with a real new idea. The senior author (Wilson) and the student identify that the softmax output head is a significant fraction of parameter storage in small LLMs (15-30%) and that post-training quantization pipelines routinely leave it in fp16. They pose head quantization as a rate-distortion problem under KL divergence between original and quantized output distributions. That's new relative to the WMSE-framed WaterSIC, and the per-class rescaling B=diag(lambda_k^{-1/2}) applied to the WaterSIC lattice is a genuine innovation. It reduces to WaterSIC when the class distribution is uniform.\n\nThe experiments are broad and mostly convincing: five models from 1B to 32B, head-only and fully quantized, multiple rates, zero-shot tasks, domain-targeted calibration. At matched head rates, SoftWater beats the released WaterSIC on 59 of 60 test points. The gains are large: at 2 bits, the KL reduction is 6.5-8.3x on WT2 and 3.5-5.1x on C4. The zero-shot results are consistent with the KL objective: most of the improvement is in LAMBADA, the only task scored by next-token prediction.\n\nThe main soft spots are theoretical. The dither argument in Eq. (8) is used to drop cross-class terms, but the dither is then removed with only an empirical assertion that undithered errors behave the same. The stress-test note is right that the validation in §5.6 (rho) only tests the diagonal separable proxy, not the off-diagonal -p_k p_l terms. So the paper does not directly support the claim that the class-aware allocation is KL-optimal for the deployed scheme. However, this is not fatal: the empirical results stand on their own, and if cross-class correlations were large, one would expect the gains to be less consistent than they are. Still, the authors should either provide the dither variant as a baseline or measure the undithered correlations directly.\n\nThe other issues are minor: no code/data release, so the headline numbers cannot be independently checked; and the abstract's '6.5-8.3x' is WT2-only, with C4 roughly half. The paper is honest about the separability assumption's limitations.\n\nOverall, this deserves serious peer review and likely publication at a strong venue after the theoretical gap is addressed. I'd bring it to a reading group and cite it if I worked on LLM quantization.","headline":"A solid new idea for quantizing the softmax head under KL, with broad empirical support and a fixable theoretical gap around the dither assumption.","tokens_in":25796,"tokens_out":6277,"would_cite":true,"duration_ms":47916,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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.","keywords":["post-training quantization","softmax output layer","KL divergence","rate-distortion","lattice quantization","successive interference cancellation","class-aware rate allocation","LLM compression"],"falsifier":"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.","tokens_in":24640,"feed_emoji":"🧠","tokens_out":7443,"duration_ms":70265,"temperature":0.7,"pith_summary":"SoftWater is a post-training quantization scheme for the softmax output layer (the 'head') of a language model, treating it as a rate-distortion problem under the KL divergence between the original and quantized output distributions. It claims that bits should be allocated per output class as well as per input feature: frequent, low-variance classes get fine grids, rare classes get coarse grids, and under Zipfian token statistics this gap is large. The paper backs the claim with a separable second-order surrogate that makes the lattice encodable by successive interference cancellation, and with experiments across five models from 1B to 32B parameters where SoftWater beats the released WaterSIC quantizer at matched head rates on 59 of 60 test points, reducing head-induced KL by 6.5–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 at a 2.9–3.7% WikiText-2 perplexity cost.","feed_headline":"2-bit LLM heads can shed half the stored bytes for ~3% perplexity","feed_subtitle":"SoftWater allocates head bits by token frequency, cutting head-induced KL up to 8x versus uniform quantization.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the successive-interference-cancellation encoder, the per-column rate allocation formula α_i∝|ℓ_ii|^{-1}, and the released WaterSIC baseline that SoftWater is compared against.","marker":"Lifar et al. (2026)"},{"why":"Establishes the Kronecker-factored curvature precedent that motivates treating the softmax KL Hessian as separable.","marker":"Martens & Grosse (2015)"},{"why":"Supplies the released quantized-body checkpoints used in the full-model experiments with tied heads.","marker":"Kim et al. (2025)"},{"why":"Supplies the power-iteration Hessian sketching used to measure how close the SoftWater factors are to the Frobenius-optimal Kronecker factorization.","marker":"Tseng et al. (2026)"},{"why":"Defines the standard WMSE linear-layer quantization objective that SoftWater replaces for the softmax head.","marker":"Frantar et al. (2023a)"},{"why":"Gives the Kronecker-product rearrangement facts behind the alternating-least-squares diagnostic used to test separability.","marker":"Van Loan (2000)"}],"fun_headline_variants":["SoftWater: token-frequency bit allocation cuts head KL up to 8x at 2-bit","Class-aware head bits: SoftWater cuts KL up to 8x, storage half at 2-bit","SoftWater: 2-bit heads with class-aware grids cut KL 6.5-8x","SoftWater: allocate head bits by class, not WMSE; cut KL up to 8x at 2-bit"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["SoftWater: token-frequency bit allocation cuts head KL up to 8x at 2-bit","Class-aware head bits: SoftWater cuts KL up to 8x, storage half at 2-bit","SoftWater: 2-bit heads with class-aware grids cut KL 6.5-8x","SoftWater: allocate head bits by class, not WMSE; cut KL up to 8x at 2-bit"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.001471,"raw_usage":{"total_tokens":6004,"prompt_tokens":1128,"completion_tokens":4876,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":744,"completion_tokens_details":{"reasoning_tokens":4778}},"tokens_in":744,"tokens_out":4876,"duration_ms":37861,"temperature":1.0,"reasoning_tokens":4778,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T00:19:18.605266+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[],"review_version":1}