大型语言模型(LLM)在文本生成过程中,需要缓存所有先前计算过的键值(Key-Value, KV)对,这个KV缓存的大小与序列长度呈线性增长,成为了服务部署时的主要内存瓶颈。通过将KV缓存量化为更低的比特数,可以显著降低成本。然而,当前的量化方法对所有注意力头分配相同的位宽,忽视了不同注意力头在模型中重要性的巨大差异。直观地看,一个自然的想法是给重要的头分配更多的比特,而给不那么重要的头分配更少的比特,即混合精度量化。但最新研究揭示了一个隐藏的陷阱:每个量化器遵循不同的失真曲线 D(b)=α·β^{-b},其中衰减率β在不同量化器设计间的变化范围从3.6到5.3。如果错误地将一种量化器的失真模型应用到另一种量化器上,会导致比特分配顺序完全颠倒,从而使性能甚至比均匀量化还要差。这种失败模式被称为“失真模型不匹配”。
为了解决这一问题,来自中国的研究团队提出了RateQuant方法。该方法首先使用一个小的校准集来拟合每个量化器的失真模型,然后利用率失真理论中的反水填充(reverse waterfilling)方法,以闭合形式求解最优的比特分配问题。实验结果表明,在Qwen3-8B模型上,以平均2.5比特为目标时,经过校准的RateQuant将KIVI量化方法的困惑度(perplexity)从49.3大幅降低到14.9,降幅达到70%,并在QuaRot方法的基础上进一步提升了6.6的困惑度。整个校准过程在单个GPU上仅需1.6秒,并且在推理阶段没有任何额外开销。
这项成果为大语言模型的高效部署提供了全新的思路,表明基于率失真理论的混合精度量化可以在不增加推理成本的前提下显著提升模型性能。论文共18页,包含7张图和5张表格,为后续研究提供了丰富的参考。