针对大模型软分类层量化难题，提出按类别自适应分配比特率的新方法。
SoftWater: Class-Aware Rate Allocation for Softmax Quantization
- 将输出层量化建模为基于KL散度的率失真问题，结合特征协方差与类别曲率加权误差
- 在1B至32B模型上，2比特头可降6.5至8.3倍头层引入的KL损失，仅增2.9%困惑度
- 无需复杂优化流程，适配低资源部署，尤其适合词汇分布不均的大模型
后训练量化常保留软分类层为高精度。然而在现代小规模语言模型中，该层占总参数15%–30%，即使标称2比特模型，其fp16头仍存数倍比特/权重。本文将软分类层量化视为在原始与量化输出分布间最小化KL散度的率失真问题。二阶分析揭示类相关几何结构：量化误差由特征协方差与类别特异性软分类曲率共同加权。通过可分离近似，将原$Kn imes Kn$ Cholesky分解简化为单个$n imes n$分解并按类别缩放，使格点可通过逐次干扰消除编码，统计量仅需一次前向传播获取。所得方法SoftWater对高频、低方差类别分配细网格，稀有类别用粗网格，在齐普夫分布下表现显著。在五个1B至32B模型上，于相同头层比特率下，SoftWater优于已发布WaterSIC量化器（线性层均方误差最优但非输出层KL最优）在59/60测试点上表现，且无需其任何优化改进，2比特时头层引入的KL下降6.5×–8.3×。对Llama-3.2-1B-Instruct模型，量化主体后，2比特头可减少45%–60%存储字节，困惑度仅上升2.9%–3.7%。因类别侧统计来自校准数据，校准域匹配部署域可实现该域最低KL。同一模型下，4比特头近乎无损，2比特头困惑度损失低于4%，表明此类模型的头层量化具有可行性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。