SoftWater: Class-Aware Rate Allocation for Softmax Quantization

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.

• The paper introduces SoftWater, a KL-based post-training quantization method that combines softmax curvature, token-class statistics, and feature covariance through a class-rescaled WaterSIC lattice.
• SoftWater outperforms WaterSIC on 59 of 60 matched-rate tests, reduces 2-bit head KL by up to 8.3×, and delivers roughly one extra effective bit of quality.
• The method can remove 45–60% of stored bytes in quantized Llama models with a 2.9–3.7% WikiText-2 perplexity increase, while domain-matched calibration improves transfer-sensitive performance.

SoftWater addresses a gap in post-training quantization (PTQ): the softmax output layer is almost always left in fp16, even though in small LLMs with modern vocabularies the head holds 15–30% of all parameters. A nominal "2-bit" Llama-3.2-1B with an fp16 head stores nearly 5 bits per weight. The paper poses softmax-layer quantization as a rate-distortion problem under the KL divergence between original and quantized output distributions, derives the induced second-order error metric, and shows it factors into a class-side statistic that standard WMSE pipelines ignore entirely. The resulting algorithm, SoftWater, reuses the successive interference cancellation (SIC) encoder of WaterSIC (Lifar et al., 5 Mar 2026) with a class-rescaled lattice, and requires only one extra $K$-vector accumulated during the same calibration forward pass.

Motivation and problem formulation

The standard PTQ objective for linear layers is weighted mean-squared error, $E\Vert(\hat{W}-W)X\Vert^2$, targeted by GPTQ-style sequential rounding, QuIP-style lattice codebooks, and rotation methods. WaterSIC brought this WMSE problem within 0.255 bits of its information-theoretic limit by allocating per-column rates from the Cholesky factor of the input covariance. But a layer feeding a softmax outputs a probability distribution, and the paper argues the correct distortion is the KL divergence $D_{\mathrm{KL}}(p \,\Vert\, q)$ between original and quantized outputs. Prior work on the output layer either replaces the head entirely (adaptive softmax, VQ-logits, FlashHead) or optimizes cross-entropy on the final block (LFQ); none quantize the existing head under KL.

The head-share motivation is concrete: Gemma-3-1B's head is 30.2% of parameters, Qwen3-0.6B's is 26.1%, Llama-3.2-1B's is 21.3%. Because the head is often tied to the input embedding, quantizing it also shrinks the embedding, so the savings compound.

The KL geometry and a tractable proxy

A second-order Taylor expansion of $D_{\mathrm{KL}}(p\,\Vert\,q)$ around zero logit perturbation yields an error weight $E_X[(\mathrm{diag}(p)-pp^\top)\otimes XX^\top]$ — a $Kn\times Kn$ matrix coupling softmax curvature with feature covariance. Direct encoding against it would require a Cholesky factorization costing $O(K^3n^3)$.

The paper reduces this in two steps. First, under a dither argument (which empirically turns out unnecessary, since undithered SIC errors already behave as uncorrelated across rows), the cross-class terms $-p_kp_l$ vanish and only the diagonal curvature $\lambda_k(X)=p_k(1-p_k)$ matters. Second, a separability assumption — that $\lambda$ and $E\Vert(\hat{W}-W)X\Vert^2$0 are uncorrelated under the input distribution — lets the Kronecker factors be estimated independently. The resulting proxy uses $E\Vert(\hat{W}-W)X\Vert^2$1, i.e., class frequency discounted by cross-context variance, paired with the usual covariance $E\Vert(\hat{W}-W)X\Vert^2$2. One $E\Vert(\hat{W}-W)X\Vert^2$3 Cholesky, rescaled per class, gives a lattice encodable by SIC at $E\Vert(\hat{W}-W)X\Vert^2$4 cost, with both statistics from a single forward pass and no backward pass.

The induced grid spacing for class $E\Vert(\hat{W}-W)X\Vert^2$5 and column $E\Vert(\hat{W}-W)X\Vert^2$6 is proportional to $E\Vert(\hat{W}-W)X\Vert^2$7. This is the paper's central claim: SoftWater corrects for in-feature geometry, class frequency, and class variance simultaneously, giving fine grids to frequent, low-variance classes and coarse grids to rare ones. Under Zipfian token distributions this gap spans orders of magnitude, and the paper argues this class side — not the feature side — is where the rate budget of a softmax layer is won or lost. Solving the head via WMSE is equivalent to assuming identically distributed classes.

The raw allocation is least trustworthy exactly where the Taylor expansion breaks: classes absent from calibration data get vanishing $E\Vert(\hat{W}-W)X\Vert^2$8 and hence absurdly coarse grids. A uniform smoothing prior $E\Vert(\hat{W}-W)X\Vert^2$9 caps the class-side spacing ratio at $D_{\mathrm{KL}}(p \,\Vert\, q)$0, bounding logit perturbations in the valid regime. The default $D_{\mathrm{KL}}(p \,\Vert\, q)$1 interpolates between calibration-driven allocation ($D_{\mathrm{KL}}(p \,\Vert\, q)$2) and uniform weighting, which exactly recovers WaterSIC ($D_{\mathrm{KL}}(p \,\Vert\, q)$3) — so SoftWater strictly contains the WMSE treatment as the special case of a uniform deployment assumption.

Head-only results

Across five models (1B to 32B, Llama and Qwen families), with bodies in fp16 and heads quantized at 2, 3, and 4 bits, SoftWater beats the full released WaterSIC pipeline (minus finetuning, and using none of its quality refinements) at matched head rates on 59 of 60 test points. At 2 bits on WikiText2, the KL advantage is $D_{\mathrm{KL}}(p \,\Vert\, q)$4 to $D_{\mathrm{KL}}(p \,\Vert\, q)$5, and in every model a 2-bit SoftWater head outperforms a 3-bit WaterSIC head on WT2 KL — roughly one free bit. The KL gap is nearly flat across head rates but shrinks out of domain ($D_{\mathrm{KL}}(p \,\Vert\, q)$6–$D_{\mathrm{KL}}(p \,\Vert\, q)$7 on C4). Perplexity gains close almost entirely at 4 bits, where SoftWater heads are near-lossless ($D_{\mathrm{KL}}(p \,\Vert\, q)$8–$D_{\mathrm{KL}}(p \,\Vert\, q)$9 PPL).

Quantized bodies and whole-model compression

With released GuidedQuant (Kim et al., 11 May 2025) quantized bodies for Llama-3.2-1B-Instruct (tied head), the head holds 52–68% of stored bytes once the body is compressed, so head quantization is the only remaining lever. A 2-bit SoftWater head removes 45–60% of stored bytes for a 2.9–3.7% WT2 perplexity increase, versus 9.7–17.6% for WaterSIC; SoftWater wins 35 of 36 cells. At 4 bits the head costs 0.2–0.3% perplexity while still removing at least 38% of model size. The paper also notes the head dominates total KL as bodies improve: on a 4-bit body, a 2-bit head already doubles total KL under SoftWater and quintuples it under the released quantizer.

Domain targeting and zero-shot transfer

Because $D_{\mathrm{KL}}(p\,\Vert\,q)$0 is a token-frequency-like statistic, it carries domain information that $D_{\mathrm{KL}}(p\,\Vert\,q)$1 does not. Calibrating on five domains (English and German Wikipedia, Python code, math web, case law), the paper finds that for every evaluation domain and head rate, the lowest KL among all ten arms is the SoftWater head calibrated on that domain, with matching calibration lowering 2-bit KL by $D_{\mathrm{KL}}(p\,\Vert\,q)$2 to $D_{\mathrm{KL}}(p\,\Vert\,q)$3 over the best mismatched arm. This is a double-edged result: seven of the nine losses in the crossover are the code column under non-code calibration, so sharpened allocation toward calibration-frequent tokens costs out-of-domain accuracy.

On zero-shot tasks (lm-evaluation-harness, seven benchmarks) with 2-bit heads on quantized bodies, SoftWater roughly halves the mean accuracy lost to head quantization on every body rate (e.g., from 1.91 to 1.07 points on the 4-bit body). The gain concentrates in LAMBADA — the only task scored by next-token prediction, which KL directly controls — where SoftWater gains 5.2 to 7.1 points at 2 bits. Ranking-based tasks show smaller but consistent gains.

A calibration-diversity finding is notable: because $D_{\mathrm{KL}}(p\,\Vert\,q)$4 is a per-class frequency, context count, not token count, is the binding resource. Re-chunking a fixed 131k-token budget into 128-token contexts lifts the 2-bit-head mean accuracy by 0.35–0.86 points while moving WaterSIC by at most 0.23.

Ablations and validity checks

Three experiments probe the method's assumptions:

• Class-side statistic. Replacing $D_{\mathrm{KL}}(p\,\Vert\,q)$5 with the plain marginal $D_{\mathrm{KL}}(p\,\Vert\,q)$6 performs nearly identically in domain, but SoftWater wins every out-of-domain KL cell and 12 of 15 in-domain ones — the second-moment (variance) correction is small but systematic and is what makes the allocation transfer off the calibration domain. Both beat the corpus unigram frequency $D_{\mathrm{KL}}(p\,\Vert\,q)$7 by roughly $D_{\mathrm{KL}}(p\,\Vert\,q)$8 in KL, showing the model's own marginal carries information corpus counts do not.
• Separability. The measured relative proxy error $D_{\mathrm{KL}}(p\,\Vert\,q)$9 is at most $E_X[(\mathrm{diag}(p)-pp^\top)\otimes XX^\top]$0 across models, rates, and datasets — a rate penalty of at most 0.07 bits, i.e., the proxy underestimates true distortion by at most 10%. An adapted YAQA-style power iteration to the Frobenius-optimal diagonal-Kronecker factorization shows the analytical SoftWater factors start at cosine similarity 0.98–0.99 to the optimum and converge in one round, while identity initializations need two.

Limitations and open questions

The paper is explicit about two bounds on its claims. First, the separable surrogate is validated empirically at ≤10% error on the models measured but carries no general guarantee; the theoretical justification for the cross-row decorrelation relies on a dither argument the final scheme does not deploy, and the separability assumption is plausible rather than proven for arbitrary deployments. Second, the deployed-model experiments (quantized bodies) rest on a single model family and a single released body checkpoint pipeline, so the whole-model compression figures may not transfer. The out-of-domain KL degradation under domain-targeted calibration — and the corresponding accuracy losses on code under non-code calibration — remains an unresolved trade-off between targeting and robustness. Finally, the class-side statistic depends on calibration data matching the deployment domain; the paper does not address what happens under distribution shift between calibration and deployment beyond the crossover experiments.

Conclusion

SoftWater reframes softmax-layer quantization as KL rate-distortion, derives a class-aware second-order geometry, and reduces it to a class-rescaled SIC lattice whose two statistics come from one calibration forward pass. Empirically it dominates the near-optimal WMSE quantizer on 59 of 60 matched-rate test points, cuts 2-bit head KL by up to $E_X[(\mathrm{diag}(p)-pp^\top)\otimes XX^\top]$1, halves the zero-shot accuracy cost of a 2-bit head, and enables up to 60% whole-model size reduction with under 4% perplexity cost. Beyond LLM heads, the analysis applies to any linear-softmax layer with fixed class dimension, including MoE routers, whose skewed expert utilization mirrors the Zipfian structure the method exploits.