论文
随机注意力:重新思考 KV 缓存驱逐以实现高效推理
Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning
摘要
大语言模型在需要扩展推理的任务上实现了优越的性能,但长的思想链使KV缓存成为严重的内存瓶颈。现有的 KV 缓存压缩方法共享一个范例:通过对每个缓存词元稍后的重要性的估计来对每个缓存词元进行评分,并保留得分最高的词元。我们表明选择信号几乎没有任何贡献。随机注意力在每个注意力头内保持提示和驱逐均匀随机,根本不计算分数;在四种模型和六种推理任务中,它与最强的先验驱逐器相匹配,同时比 vLLM 部署中的吞吐量高出 32-43%。对照实验通过以下方式解释了这一点:1)提示符是缓存的脆弱部分,选择器之间的大部分差距只是它们的选择信号是否碰巧保留了它; 2)推理轨迹通过两个层面的冗余来保护自己免受驱逐,即在文本中(模型重申其工作时仍然需要的内容)和跨注意力头(每个都保留自己的轨迹副本),因此一旦提示是安全的,随机抽取将保留模型仍然需要的足够副本,并且不需要分数来挑选它们。我们的代码可在 https://github.com/SalesforceAIResearch/Random-Attention 公开获取。