从未来回来：反因果惊喜的键值缓存管理

通过压缩和驱逐策略进行键值（KV）缓存管理已成为近年来的一个重要研究方向。通过缓存后续缩放点积注意力操作所需的先前键和值计算，可以部分缓解 大语言模型（LLM）及其多模态变体在输出生成过程中的计算需求。然而，这会导致另一个问题：生成的 KV 缓存的大小随着上下文长度线性增长，并且当提示或生成的输出很长时，会快速消耗所有可用的 GPU 内存。 KV 缓存管理定期从缓存中修剪条目，从而减少其内存占用，同时尝试保留足够的信息以进行准确生成。副产品是更快的推理速度。我们提出了一种简单而有效的 KV 驱逐方案，其动机是认识到可以从最近的词元中很好地预测到的过去的词元是多余的，并且可以从缓存中删除它们相关的键和值。为了对驱逐条目进行评分，我们按照词元的原始顺序运行模型，重用已存储在 KV 缓存中的键和值表示，并应用反因果注意掩码，以便每个位置仅关注其未来的上下文。这是分布式的，直接与实际的缓存内容相关，并且不需要额外的训练。为了进一步降低成本，我们还提出了一种快速单层近似，将反因果传递限制到最后一个 Transformer 层，以边际精度成本实现每个刷新周期的显着加速。我们在各种开源 LLM 和基准数据集上评估我们的策略，显示出与其他最先进方法相比具有竞争力或改进的性能。参考代码可在 https://github.com/metacognitionai/counter_causal 获取。