通过注意力门控动态淘汰冗余缓存，提升大模型推理效率与性能。
In-context KV-Cache Eviction for LLMs via Attention-Gate
- 引入轻量级注意力门控模块，按全局上下文决定缓存哪些键值对。
- 在多个场景下验证，淘汰冗余信息后推理速度提升且准确率不降反升。
- 可微调适配预训练模型，计算开销极小，适合部署在资源受限设备上。
KV-Cache 技术已成为大语言模型推理的标准方法，但常被视为系统瓶颈。本文提出一种新型动态 KV-Cache 淘汰策略，通过在模型中注入轻量级 Attention-Gate 模块实现。该模块接收全局上下文输入，为每个标记生成淘汰标志，使自注意力机制根据标志仅缓存部分键值状态用于下一词预测。Attention-Gate 可针对不同注意力头和层生成差异化标志，并通过持续预训练或监督微调轻松适配预训练模型。所引入的计算与内存开销极小。我们在多种场景下实证评估该方法，结果表明有效淘汰冗余标记不仅能提升效率，还能增强模型性能。
原文摘要 · Abstract (English)
The KV-Cache technique has become the standard for the inference of large language models (LLMs). Yet, it is widely criticized that KV-Cache can become a bottleneck of the LLM inference system. This paper enables a novel dynamic KV-Cache eviction policy by injecting a lightweight module called Attention-Gate to the model. It accepts the global context as input and yields eviction flags for each token. The self-attention modules in the model proceed according to the flags and cache only a subset of the KV states for next token prediction. The Attention-Gates can yield various flags for different heads and layers and be easily tuned on top of a pre-trained LLM via continual pre-training or supervised fine-tuning. The computational and memory overhead introduced by Attention-Gates can be minimal. We empirically evaluate the proposed approach across multiple scenarios, showing that effective eviction of redundant tokens can not only improve efficiency but also enhance performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。