SoftWater: Class-Aware Rate Allocation for Softmax Quantization

2026-08-12

AI Summary

This work addresses the challenge that the softmax output layer in large language models occupies a significant fraction of model parameters, hindering effective low-bit quantization. The authors formulate this problem as a KL divergence–driven rate-distortion optimization and propose a class-aware quantization grid allocation strategy: fine-grained grids are assigned to frequent, low-variance tokens, while coarse-grained grids are used for rare tokens, aligning with Zipf’s law. Their approach leverages second-order analysis of feature covariances and class-specific softmax curvature, combined with Cholesky approximation and continuous interference cancellation coding to enable efficient quantization. Experiments on models ranging from 1B to 32B parameters show that 2-bit softmax heads reduce output KL divergence by 6.5–8.3× compared to baselines; on Llama-3.2-1B-Instruct, storage overhead is cut by 45–60% with only a 2.9–3.7% increase in perplexity.

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.

Problem

• softmax quantization
• post-training quantization
• rate-distortion
• large language models
• KL divergence

Innovation

• softmax quantization
• rate-distortion optimization
• class-aware allocation
• post-training quantization
• KL divergence