论文

KV缓存压缩的风险

The risk of KV cache compression

模型推理KV Cache

摘要

Transformer 对长序列的推理非常昂贵,因为 softmax 注意力会重复从大型 KV 缓存中读取。解决这个瓶颈的流行方法是 KV 缓存压缩,它用紧凑摘要代替完整缓存。尽管具有实际重要性,但此类摘要的设计很大程度上是由经验实验驱动的。在理论方面,现有的结果表明,KV 缓存压缩在最坏的情况下是不可能的,但对于在可能进行精确压缩的情况下设计算法几乎没有提供系统指导。我们通过根据缓存的固有可压缩性来描述 KV 缓存压缩的极小极大风险来弥补这一差距,揭示何时以及如何可以进行精确压缩。这些结果产生了因果屏蔽下 KV 缓存压缩的新颖设计原理,可有效映射到预填充和自回归解码,同时实现极小极大最优风险。我们在实用算法中实例化了这些原理,并在有针对性的实验中在 LongBench 上报告了有希望的性能。总的来说,我们的结果为实际的 KV 缓存压缩提供了一条原则性的途径,并提供了理论保证。