IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference
作者: Xintong Yang, Hao Gu, Binxing Xu, Lujun Li, Bei Liu, Jiacheng Liu, Qiyuan Zhu, Sirui Han, Yike Guo
分类: cs.CL, cs.AI
发布日期: 2026-05-25
💡 一句话要点
IndexMem：利用潜在记忆学习KV缓存淘汰策略，提升长文本LLM推理性能
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长文本推理 KV缓存 淘汰策略 可学习索引 潜在记忆 大语言模型 信息压缩
📋 核心要点
- 现有KV缓存淘汰策略依赖启发式方法，无法准确捕捉token的重要性分布，导致长文本推理性能受限。
- 提出IndexMem，通过可学习的索引器预测KV重要性，并使用潜在记忆模块补偿淘汰token带来的信息损失。
- 实验表明，IndexMem在多个模型和数据集上显著提升了长文本推理性能，尤其在激进淘汰策略下效果更佳。
📝 摘要（中文）
大型语言模型(LLMs)越来越多地需要在长上下文中运行，但标准softmax注意力机制产生的KV缓存随着序列长度线性增长，迅速成为长上下文推理的瓶颈。一个实用的补救方法是淘汰不太重要的KV条目；然而，现有的淘汰策略大多是启发式的，难以捕捉token重要性的丰富、输入相关的分布。本文提出了一种可学习的索引器，用于预测KV重要性，从而更准确地保留关键token。同时，简单地永久淘汰token会丢弃它们的信息，导致不可逆的遗忘和长距离检索性能下降。为了解决这个问题，我们提出了一个轻量级的潜在记忆模块，将淘汰的token压缩成一个紧凑的、在线更新的状态，并提供残差读出，以补偿因KV淘汰而损失的注意力贡献。总的来说，我们的方法能够在有限的KV预算下实现准确的长上下文推理，在Qwen、Mistral和Llama模型上，RULER（4K/16K）上实现了持续的改进（在激进的淘汰策略下高达25个点），显著提高了Needle-in-a-Haystack检索的稳定性，并且与现有的淘汰策略相比，具有更优越的LongBench分数和压缩曲线。
🔬 方法详解
问题定义：长文本LLM推理中，KV缓存的大小随序列长度线性增长，成为性能瓶颈。现有的KV缓存淘汰策略，如LRU、FIFO等，主要基于启发式规则，无法有效区分token的重要性，导致关键信息被淘汰，影响长文本推理的准确性。
核心思路：IndexMem的核心思路是学习一个索引器来预测每个token的KV值的重要性，并结合一个潜在记忆模块来存储被淘汰的token的信息，从而在有限的KV缓存预算下，尽可能保留重要的信息，并减少信息损失。通过这种方式，模型可以更好地处理长上下文，提高推理的准确性和效率。
技术框架：IndexMem主要包含两个模块：1) 可学习的索引器：该模块负责预测每个token的KV值的重要性得分。索引器可以是任何可学习的模型，例如小型神经网络。2) 潜在记忆模块：该模块用于存储被淘汰的token的信息。该模块将淘汰的token压缩成一个紧凑的向量表示，并在线更新该向量。在推理时，该模块提供残差读出，以补偿因KV淘汰而损失的注意力贡献。整体流程是，对于每个token，首先使用索引器预测其重要性得分，然后根据得分决定是否保留该token的KV值。如果token的KV值被淘汰，则将其信息存储到潜在记忆模块中。
关键创新：IndexMem的关键创新在于：1) 提出了一种可学习的索引器，用于预测KV值的重要性，从而可以更准确地保留关键信息。2) 引入了潜在记忆模块，用于存储被淘汰的token的信息，从而减少了信息损失。这两个创新点共同作用，使得IndexMem能够在有限的KV缓存预算下，实现更准确的长文本推理。
关键设计：索引器可以使用多层感知机（MLP）或Transformer等结构，输入为token的embedding，输出为重要性得分。潜在记忆模块可以使用循环神经网络（RNN）或Transformer等结构，将淘汰的token的embedding作为输入，输出为压缩后的向量表示。损失函数可以包括预测重要性得分的损失和重构token embedding的损失。具体参数设置需要根据具体的模型和数据集进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，IndexMem在RULER数据集上，针对Qwen、Mistral和Llama等模型，在4K/16K上下文长度下，相比现有淘汰策略，性能提升高达25个点（在激进淘汰策略下）。同时，IndexMem显著提高了Needle-in-a-Haystack检索的稳定性，并在LongBench数据集上取得了更优越的分数和压缩曲线。
🎯 应用场景
IndexMem可应用于各种需要处理长文本的LLM应用场景，例如长文档摘要、问答系统、代码生成、对话系统等。通过更有效地利用有限的KV缓存资源，IndexMem可以显著提升这些应用的性能和效率，降低计算成本，并支持更大规模的上下文处理。
📄 摘要（原文）
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.