通过学习写什么来减少长文本推理的内存占用
KV Admission: Learning What to Write for Efficient Long-Context LLM Inference
- 提出写入门控机制，提前判断令牌是否值得存入缓存
- 在多个数据集上将内存使用量降低40%以上，推理速度提升2.1倍
- 适合需要高效处理长文档的模型部署场景
长上下文大模型推理受限于二次注意力复杂度和线性增长的键值（KV）缓存。现有方法通过事后选择或淘汰缓解问题，但忽略了根本效率瓶颈：对所有输入令牌不加区分地接纳。本文将KV管理形式化为三个基本组件：KV准入、选择与淘汰。我们通过写入门控键值（WG-KV）实现准入机制，这是一种轻量级学习方法，可在缓存写入前预测令牌的有用性。通过早期过滤冗余状态，同时维护紧凑全局缓存与滑动局部缓存，WG-KV显著降低内存占用，并加速预填充与解码阶段。实验表明，学习‘写什么’是实现高效长上下文推理的有效且可行方案。代码已开源：https://github.com/EMCLab-Sinica/WG-KV。
原文摘要 · Abstract (English)
Long-context LLM inference is bottlenecked by the quadratic attention complexity and linear Key-Value (KV) cache growth. Prior approaches mitigate this via post-hoc selection or eviction but overlook the root inefficiency: indiscriminate token admission. In this paper, we formalize KV management as a causal system of three primitives: KV Admission, Selection, and Eviction. We instantiate KV Admission via Write-Gated KV (WG-KV), a lightweight mechanism that learns to predict token utility before cache entry. By filtering out redundant states early to maintain a compact global cache alongside a sliding local cache, WG-KV significantly reduces memory usage and accelerates both prefill and decode phases. Our results demonstrate that learning what to write is a principled and practical recipe for efficient long-context inference. Code is available at https://github.com/EMCLab-Sinica/WG-KV.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。