arXiv:2605.06675cs.LGcs.CL2026-05被引 2

提出新方法,让不同注意力头用不同精度存缓存,大幅降内存且不丢性能。

RateQuant: Optimal Mixed-Precision KV Cache Quantization via Rate-Distortion Theory

论文配图:RateQuant: Optimal Mixed-Precision KV Cache Quantization via Rate-Distortion Theory
图 1 · 摘自论文原文
  • 根据每个注意力头的特性,动态分配不同精度比特数。
  • 在Qwen3-8B上实现2.5比特平均精度,困惑度降低70%。
  • 仅需1.6秒校准,推理零开销,适合部署优化场景。

大型语言模型在生成时会缓存所有先前计算出的键值(KV)对,该缓存随序列长度线性增长,成为服务阶段的主要内存瓶颈。将KV缓存量化到更少比特可降低内存开销,但现有方法对每个注意力头统一分配比特数,忽略了头重要性的差异。一个自然思路是为重要头分配更多比特,其余则减少。然而我们发现,这种混合精度分配存在隐性陷阱:每个量化器遵循不同的失真曲线 D(b)=alpha*beta^{-b},其衰减速率 beta 在3.6至5.3之间变化。若错误使用某一量化器的失真模型去指导另一量化器的比特分配,会导致分配顺序颠倒,性能反而劣于均匀量化。我们称此为失真模型错配。为此提出 RateQuant:从少量校准集拟合每种量化器的失真模型,再通过率失真理论中的反向水填法闭式求解比特分配问题。在 Qwen3-8B 上以 2.5 比特平均精度运行,校准后的 RateQuant 将 KIVI 的困惑度从 49.3 降至 14.9(下降 70%),并使 QuaRot 提升 6.6 点 PPL。整个校准过程仅需单张 GPU 上 1.6 秒,推理阶段无额外开销。

原文摘要 · Abstract (English)

Large language models cache all previously computed key-value (KV) pairs during generation, and this KV cache grows linearly with sequence length, making it a primary memory bottleneck for serving. Quantizing the KV cache to fewer bits reduces this cost, yet all current quantizers assign the same bit-width to every attention head, ignoring the large variation in head importance. A natural idea is to allocate more bits to important heads and fewer to the rest. We show, however, that such mixed-precision allocation has a hidden pitfall: each quantizer follows a different distortion curve D(b)=alpha*beta^{-b}, and the decay rate beta varies from 3.6 to 5.3 across quantizer designs. Applying one quantizer's distortion model to another inverts the allocation order and makes performance worse than uniform quantization. We call this failure mode distortion model mismatch and propose RateQuant to resolve it. RateQuant fits a per-quantizer distortion model from a small calibration set, then solves the resulting bit-allocation problem in closed form via reverse waterfilling from rate-distortion theory. On Qwen3-8B at 2.5 average bits, calibrated RateQuant reduces KIVI's perplexity from 49.3 to 14.9 (70% reduction) and improves QuaRot by 6.6 PPL. The entire calibration takes 1.6 s on a single GPU and adds zero overhead at inference time.

量化注意力缓存混合精度推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。