🤖 machine learning

SoftWater: Class-Aware Rate Allocation for Softmax Quantization

SoftWater is a class-aware rate allocation method for softmax quantization that solves a KL-divergence-based rate-distortion problem to assign finer grids to frequent, low-variance classes and coarser grids to rare ones, significantly reducing storage overhead in small LLMs while minimizing perplexity degradation.

Original authors: Joao V. Cavalcanti, Ashia C. Wilson

Published 2026-08-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Joao V. Cavalcanti, Ashia C. Wilson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to shrink a giant library of books down to fit into a tiny backpack. In the world of Artificial Intelligence, these "books" are massive language models—computer brains trained to write, chat, and solve problems. To make them fast enough to run on your phone or a small laptop, scientists use a trick called quantization. Think of this like translating a high-definition movie into a pixelated, low-resolution version. You lose some detail, but the file size shrinks dramatically, making it portable.

Usually, when scientists shrink these AI models, they treat every part of the brain the same way. They squeeze the "body" (the part that understands context) and the "head" (the part that decides which word to say next) into the same tiny box. But here's the catch: in modern AI, the "head" is surprisingly heavy. It holds a huge chunk of the model's total memory—sometimes nearly a third! If you squeeze the whole model down to a tiny size but leave the head in its original, heavy format, you aren't actually saving much space. It's like packing a suitcase where you compress your clothes but leave the heavy boots in their original, bulky boxes. The paper we are looking at tackles this specific problem: how to shrink the AI's "word-picking head" without making it stutter or sound robotic.

The authors of this paper, Joao V. Cavalcanti and Ashia C. Wilson from MIT, realized that the old way of shrinking the head was like using a sledgehammer to crack a nut. They treated the problem as a rate-distortion issue. Imagine you have a limited budget of "bits" (the digital currency of storage) to spend on describing the AI's choices. The old methods spent these bits evenly, like giving every word in the dictionary the same amount of attention. But the authors noticed that AI doesn't use words evenly. It uses common words like "the" or "and" constantly, and rare words like "xylophone" or "quintessential" only occasionally.

The paper introduces a new method called SoftWater. Instead of giving every word the same storage space, SoftWater acts like a smart water sprinkler system. It pours a fine, high-resolution mist of bits onto the frequent, low-variance words (the ones the AI uses all the time and uses them the same way). For the rare, unpredictable words, it uses a coarser, wider spray. This "class-aware" approach means the AI can remember the common words perfectly while accepting a little fuzziness on the rare ones, all while using the same total amount of space.

The researchers tested this on five different AI models, ranging from small 1-billion-parameter models to massive 32-billion-parameter ones. They found that SoftWater is significantly better than the previous best method, called WaterSIC. In 59 out of 60 tests, SoftWater produced a model that sounded more natural and made fewer mistakes (measured by something called "KL divergence," which is just a fancy way of saying "how different the AI's choices are from the original").

Here is the magic number: when they squeezed the "head" of a 1-billion-parameter model down to just 2 bits per weight, SoftWater cut the model's total size by 45% to 60% while only making the AI slightly less perfect (a 2.9% to 3.7% increase in confusion, known as perplexity). In contrast, the old method made the AI much more confused for the same amount of space. Even more impressively, when they used a 4-bit head, the model was almost indistinguishable from the original, high-quality version.

The paper also discovered that this method is flexible. If you know the AI will be used for a specific job, like writing legal documents or coding, you can "calibrate" the sprinkler system to that specific topic. This ensures the AI stays sharp on the words it needs most for that specific domain. The authors showed that by matching the calibration data to the real-world use case, the AI performs even better.

In short, SoftWater doesn't just shrink the AI; it shrinks it intelligently. It realizes that not all words are created equal and allocates storage space accordingly. By doing so, it makes it possible to run powerful, high-quality AI models on devices that previously couldn't handle them, all without needing to retrain the model from scratch. It's a practical, mathematically sound way to make AI lighter, faster, and ready for the real world.

Technical Summary: SoftWater – Class-Aware Rate Allocation for Softmax Quantization

1. Problem Statement

Post-training quantization (PTQ) pipelines for Large Language Models (LLMs) typically leave the final softmax output layer (the "head") in high precision (e.g., FP16). While standard linear layers are quantized to minimize Weighted Mean Squared Error (WMSE), the output layer produces a probability distribution, making the Kullback-Leibler (KL) divergence the more appropriate distortion metric.

This omission is increasingly costly in modern small LLMs with large vocabularies. As shown in Table 1 of the paper, the output head can constitute 15–30% of a model's total parameters. Consequently, a "2-bit" model with an FP16 head effectively stores significantly more bits per weight than the nominal rate suggests. Prior attempts to address this either replace the head architecture entirely or apply standard WMSE-based quantization, which fails to account for the specific geometry of the softmax function and the non-uniform distribution of tokens (Zipfian distribution).

The core problem is to formulate softmax layer quantization as a rate-distortion problem under the KL divergence metric, determining how to allocate a fixed bit budget across the weight matrix to minimize the divergence between the original and quantized output distributions.

2. Methodology: SoftWater

The authors propose SoftWater, a method that derives a class-aware rate allocation strategy based on a second-order analysis of the KL divergence.

Objective Derivation:
The distortion is defined as the expected KL divergence between the original distribution $p$ and the quantized distribution $q$ . Using a second-order Taylor expansion of the KL divergence around the logit perturbation $\delta$ , the authors derive an error metric:
$D \approx \frac{1}{2} \text{vec}(\Delta)^\top \mathbb{E}[(\nabla^2 f(0) \otimes XX^\top)] \text{vec}(\Delta)$
where $\nabla^2 f(0) = \text{diag}(p) - pp^\top$ is the Hessian of the softmax function, and $XX^\top$ represents the input feature covariance. This results in a $Kn \times Kn$ error weight matrix coupling classes ( $K$ ) and features ( $n$ ).

Tractable Approximation:
Directly factorizing the $Kn \times Kn$ matrix is computationally prohibitive. The authors introduce two key approximations:

Dither Argument: Assuming subtractive dither, the errors become uncorrelated across rows, allowing the off-diagonal terms of the Hessian to be ignored.
Separability Assumption: The authors assume the softmax curvature $\lambda$ and the input covariance $XX^\top$ are uncorrelated. This allows the expectation to factorize: $\mathbb{E}[\lambda \otimes XX^\top] \approx \mathbb{E}[\lambda] \otimes \mathbb{E}[XX^\top]$ .

This factorization reduces the problem to $K$ independent $n \times n$ Cholesky factorizations (or one rescaled $n \times n$ factorization), making the problem tractable. The resulting grid spacing for class $k$ and feature column $i$ is:
$\alpha_{(k-1)n+i} \propto \frac{1}{\bar{\lambda}_k^{1/2} |\ell_{ii}|}$
where $\bar{\lambda}_k = \mathbb{E}[p_k(1-p_k)]$ is the expected curvature (related to class frequency and variance) and $\ell_{ii}$ are the diagonal entries of the Cholesky factor of the input covariance.

Key Algorithmic Features:

Class-Aware Allocation: The method allocates finer grids (more bits) to frequent, low-variance classes and coarser grids to rare, high-variance classes. This exploits the Zipfian nature of token distributions.
Smoothing Prior: To prevent instability for classes absent from calibration data (where $\bar{\lambda}_k \to 0$ ), the method applies a smoothing prior, mixing the calibration distribution with a uniform distribution. This caps the maximum grid spacing, ensuring the Taylor expansion remains valid.
Implementation: The algorithm uses Successive Interference Cancellation (SIC) for encoding, similar to the WaterSIC algorithm for linear layers, but with class-specific rescaling. It requires only a single forward pass over calibration data to compute the necessary statistics ( $\Sigma_X$ and $\bar{\lambda}$ ).

3. Key Contributions

Problem Formulation: The paper frames softmax quantization as a rate-distortion problem under KL divergence, deriving the induced error metric that couples class statistics with feature covariance.
The SoftWater Algorithm: A practical algorithm that utilizes a separability approximation to reduce the complexity from $O(K^3n^3)$ to $O(n^3)$ (plus linear passes). It introduces a class-side scaling factor that adapts the lattice grid based on class frequency and cross-context variability.
Empirical Validation: The method is validated across five models (1B to 32B parameters) from the Llama and Qwen families.

4. Experimental Results

The authors evaluate SoftWater against the released WaterSIC quantizer (which optimizes for WMSE) and full-precision baselines.

Head-Only Quantization: When quantizing only the head (keeping the body in FP16), SoftWater outperforms WaterSIC on 59 of 60 test points (across different bit rates and models). At 2 bits, SoftWater reduces the KL divergence by 6.5× to 8.3× compared to WaterSIC on WikiText2.
Fully Quantized Models: On Llama-3.2-1B-Instruct with a quantized body (using GuidedQuant), a 2-bit SoftWater head removes 45–60% of the total stored bytes while increasing perplexity by only 2.9–3.7%. In contrast, the WaterSIC baseline incurs a 9.7–17.6% perplexity increase for similar compression.
Domain Adaptation: The paper notes that because the class-side statistic ( $\bar{\lambda}$ ) is derived from calibration data, matching the calibration domain to the deployment domain yields the lowest KL divergence on that specific domain.
Separability: Experiments confirm that the separability assumption holds well, with the proxy underestimating true distortion by at most 10%.

5. Significance and Claims

The paper claims that SoftWater makes head quantization practical for models where the head constitutes a significant portion of the model size. By treating the head as a rate-distortion problem under KL rather than WMSE, the method achieves near-lossless performance at 4 bits and significant compression at 2 bits with minimal perplexity degradation.

The authors emphasize that their gains are attributable solely to the rate allocation strategy, as they do not use the additional refinements (e.g., activation drift correction, residual stream correction) present in the WaterSIC pipeline, which are orthogonal to rate allocation. The method is presented as a drop-in replacement for the head quantization step in existing PTQ pipelines, offering a substantial reduction in model size without requiring architectural changes or fine-tuning.