用可学习的索引与潜在记忆，提升长文本推理时的关键词保留能力。
IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference
- 通过学习预测注意力关键信息重要性，优化缓存淘汰策略。
- 在4K/16K上下文测试中，性能最高提升25分，比传统方法更稳定。
- 适合需要高精度长文本处理的场景，如问答、文档理解。
大型语言模型在长上下文推理中面临挑战，因标准softmax注意力导致键值（KV）缓存随序列长度线性增长，成为瓶颈。现有淘汰策略多为启发式方法，难以捕捉输入相关的词重要性分布。本文提出一种可学习的索引器，用于预测KV条目重要性，实现更精准的关键词保留。同时，直接永久淘汰词会丢失信息，造成不可逆遗忘。为此，我们设计了一个轻量级潜在记忆模块，将被淘汰的词压缩为紧凑的在线更新状态，并提供残差读出，以补偿因缓存淘汰而损失的注意力贡献。该方法在有限缓存预算下实现了准确的长上下文推理，在Qwen、Mistral和Llama模型上对RULER（4K/16K）任务均有持续提升（最激进淘汰下最高达25点），显著改善了长距离检索表现，并在LongBench得分和压缩曲线方面优于现有策略。
原文摘要 · Abstract (English)
Large Language Models (LLMs) are increasingly expected to operate over long contexts, yet standard softmax attention incurs a KV cache that grows linearly with sequence length, quickly becoming the bottleneck for long context inference. A practical remedy is to evict less important KV entries; however, existing eviction policies are largely heuristic and struggle to capture the rich, input-dependent distribution of token importance. In this work, we introduce a learnable indexer that predicts KV importance, enabling more accurate retention of critical tokens. Meanwhile, naively evicting tokens permanently discards their information, leading to irreversible forgetting and degraded retrieval over long ranges. To address this, we propose a lightweight latent memory module that compresses evicted tokens into a compact, online-updated state and provides residual readouts to compensate for the attention contributions lost through KV eviction. Collectively, our method enables accurate long-context inference under a bounded KV budget, delivering consistent improvements on RULER (4K/16K) across Qwen, Mistral, and Llama models (up to 25 points under aggressive eviction), markedly more stable Needle-in-a-Haystack retrieval, and superior LongBench scores and compression curves compared to existing eviction policies.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。