💬 NLP

IndexMem: Learned KV-Cache Eviction with Latent Memory for Long-Context LLM Inference

IndexMem 通过结合用于准确令牌淘汰的可学习索引器与保留被丢弃信息的轻量级潜在记忆模块，解决了长上下文大语言模型推理中的 KV 缓存瓶颈，从而在各种模型和基准测试中显著提升了性能与稳定性。

原作者： Xintong Yang, Hao Gu, Binxing Xu, Lujun Li, Bei Liu, Jiacheng Liu, Qiyuan Zhu, Sirui Han, Yike Guo

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Xintong Yang, Hao Gu, Binxing Xu, Lujun Li, Bei Liu, Jiacheng Liu, Qiyuan Zhu, Sirui Han, Yike Guo

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图阅读一本长达 1000 页的巨著，只为回答关于第一页的一个问题。在阅读过程中，你的大脑会自然地试图记住你看到的每一个字。但问题在于：你的“工作记忆”容量是有限的。如果你试图在脑海中同时容纳从第 1 页到第 1000 页的每一个字，你会感到不堪重负、速度变慢，甚至彻底耗尽空间。

这正是大型语言模型（LLM）在处理长文档时面临的难题。它们使用一个名为KV 缓存的“记忆库”来记住已读内容。随着文本变长，这个记忆库会不断膨胀，直到占满计算机内存，导致系统崩溃或停滞。

论文《IndexMem》针对这一问题提出了一种巧妙的两部分解决方案，其作用如同一名智能图书管理员和一个备份金库。

问题所在：“全盘保留”的陷阱

目前，大多数 AI 模型试图记录它们处理的每一个词（token）。这就像试图把你去过的每家商店的收据都塞进钱包里。最终，钱包会变得太重而无法携带。

为了解决这个问题，以往的方法尝试基于简单规则丢弃“不重要”的收据，例如“保留最新的”或“保留数值最大的”。但这些规则往往笨拙。它们可能仅仅因为某个细节来自故事开头（较旧）就将其丢弃，或者仅仅因为某句话是最近出现的就将其保留。一旦被丢弃，该信息便永远消失。

解决方案：IndexMem

作者提出了一种包含两个主要部分的系统：智能索引器和潜在记忆金库。

1. 智能索引器（“图书管理员”）

IndexMem 不使用僵硬的规则来决定保留什么，而是使用一个可学习的索引器。这就像一位受过高度训练的图书管理员，他读过数百万本书，确切知道哪些细节通常对后续回答问题至关重要。

工作原理：随着 AI 阅读，这位“图书管理员”会审视当前的问题（或 AI 即将猜测的下一个词），并预测文本中哪些部分实际上至关重要。
优势：它不仅仅基于“时间远近”进行猜测，而是理解上下文。它可以判断：“尽管这个词来自第 50 页，但它对第 100 页的问题至关重要，因此我们必须保留它。”它学会了更准确地决定哪些内容应留在主内存中，哪些应被剔除。

2. 潜在记忆金库（“备份金库”）

这是论文中最具创新性的想法。过去，如果一个词被踢出主内存，它就永远消失了。如果 AI 后来需要这个词，那将是一场灾难。

IndexMem 引入了一个潜在记忆模块。这就像一个高科技的压缩备份金库。

处理过程：当智能索引器决定将某个词踢出主内存时，它不会直接删除它。相反，它将该信息压缩成一个微小的“摘要”，并存入这个特殊金库中。
检索机制：如果 AI 后来需要回忆该信息，它不会回到已清空的主内存，而是打开金库，取出压缩摘要，并利用它来填补空白。
类比：想象你在为旅行打包。你的行李箱（主内存）只能装下几件衣服。你不得不把最喜欢的毛衣留在家里。与其扔掉它，不如拍一张高分辨率的照片并保存在手机里（潜在记忆）。如果你后来想念那件毛衣，你可以看着照片回忆它的样子和触感，而不必随身携带整件毛衣。

为何这很重要

该论文在各种模型（如 Qwen、Mistral 和 Llama）上针对超长上下文测试了该系统。

更高的准确性：即使为了节省空间而激进地丢弃了 75% 到 90% 的内存，AI 的表现依然出色。它没有忘记“大海捞针”（即回答问题所需的特定细节）。
稳定性：与旧方法不同（如果重要细节位于文本的“困难”位置，旧方法会突然失效），IndexMem 保持了稳定。
效率：它使 AI 能够在标准计算机芯片上运行，而无需依赖庞大且昂贵的超级计算机，因为它不再试图囤积每一个数据片段。

总结

简而言之，IndexMem教会 AI 成为一名更聪明的编辑。它利用一个学习系统来决定将什么保留在即时内存中，并利用一个特殊的“压缩金库”来存储其余部分。这样一来，AI 可以阅读整个图书馆而不会耗尽“脑力”，并且永远不会真正忘记解决问题所需的细节。

技术摘要：IndexMem - 基于潜在记忆的 learned KV 缓存淘汰

1. 问题陈述

大型语言模型（LLM）在处理长上下文时面临一个关键瓶颈：自回归生成所需的键值（KV）缓存随序列长度线性增长，迅速超出 GPU 显存容量并耗尽内存带宽。尽管标准 KV 缓存避免了重新计算注意力，但它造成了一个“内存墙”，限制了可部署的上下文长度（例如，一个 700 亿参数模型处理 100 万 token 的上下文需要约 320 GB 显存）。

现有解决方案主要分为两类，但均存在局限性：

表示压缩：量化和低秩近似等技术减少了每个 token 的占用空间，但并未解决缓存随上下文长度固有的线性增长问题。
Token 淘汰：通过丢弃较不重要的 token 来限制内存占用。然而，当前的淘汰策略大多是启发式的（依赖键的幅度或局部注意力窗口等静态统计信息）且无需训练。它们往往无法捕捉复杂、依赖输入 token 重要性的模式。此外，标准淘汰会永久丢弃信息，导致不可逆的遗忘，并降低需要检索遥远 token 的任务（如“大海捞针”）的性能。

2. 方法论

作者提出了 IndexMem，这是一个结合了用于准确 token 选择的可学习索引器（indexer）和用于从被淘汰 token 中恢复信息的潜在记忆模块的框架。

A. 用于 token 重要性的可学习索引器

IndexMem 不使用启发式分数，而是采用轻量级、可学习的索引器来预测 KV token 的重要性。

架构：该索引器是一个多查询注意力（MQA）风格的模块。它将多头查询投影到较低维空间，并使用源自隐藏状态的单个共享键。它包含 QK 归一化以确保稳定性，并引入可学习的门控机制来调节各头的贡献。
功能：给定隐藏状态 $X$ 和查询 $Q$ ，索引器输出一个稠密的查询到键的分数矩阵 $A$ ，其中 $A_{s,t}$ 衡量键 token $t$ 对查询 $s$ 的重要性。
训练：骨干 LLM 被冻结。索引器通过知识蒸馏进行训练，以匹配骨干模型的注意力行为。目标是最小化教师（骨干注意力 logits）和学生（索引器分数）的聚合每键重要性分布之间的 KL 散度。为了避免 $O(L^2)$ 矩阵带来的内存瓶颈，训练采用流式、FlashAttention 风格的方法来计算运行最大值聚合。
预淘汰：与在淘汰前实例化所有 KV 对的方法不同，索引器支持预淘汰，即在完整预填充计算之前预测保留哪些条目，从而节省内存和计算资源。

B. 用于被淘汰 token 的潜在记忆

为了解决因丢弃 token 而导致的不可逆信息损失，IndexMem 引入了一个潜在记忆模块。

设计：记忆模块并非将被淘汰的 token 作为额外的 KV 对注入（这可能导致注意力下沉问题），而是作为注意力流的显式残差补偿。
机制：
- 快速权重（在线更新）：被淘汰的 token 在推理过程中在线更新固定大小的状态矩阵 $M$ 和稳定器向量 $b$ 。这使用了带衰减的外积累积，使记忆能够适应特定上下文。
- 慢速权重（学习得到）：一个投影层和一个门控函数 $g(q)$ 被训练用于从记忆状态 $M$ 中读取，并决定在给定查询下对它的依赖程度。
输出：最终的注意力输出被增强： $o = o_{attn} + g(q) \cdot m(q)$ ，其中 $o_{attn}$ 是保留 token 上的标准注意力输出， $m(q)$ 是补偿丢失信息的记忆读取结果。

C. 训练协议

该系统采用两阶段训练方案：

索引器训练：单独训练索引器，以学习可靠的 token 重要性估计。
联合训练：联合训练索引器和记忆模块。记忆模块被优化以最小化完整注意力输出与“压缩输出 + 记忆残差”之间的均方误差（MSE）。

3. 主要贡献

可学习索引器：一种基于训练的方法，通过根据上下文依赖模式准确预测 token 重要性，优于启发式淘汰策略。
潜在记忆模块：一种将被淘汰 token 压缩为紧凑、在线更新状态的机制，提供残差读取以恢复淘汰过程中丢失的信息，并缓解不可逆遗忘。
全面评估：在多个模型系列（Qwen、Mistral、Llama）和基准测试（RULER、LongBench、大海捞针）中展示了持续的性能提升。

4. 实验结果

实验在 Qwen3-8B、Mistral-7B 和 Llama-3.1-8B 模型上使用 RULER 和 LongBench 套件进行。

RULER 性能：IndexMem 显著优于基线方法（SnapKV、PyramidKV、TOVA、Expected Attention）。在激进淘汰（90% 压缩）下，IndexMem 保持了 substantially 更高的分数。例如，在 Llama-3.1-8B 上进行 90% 淘汰时，IndexMem 在 RULER-16K 上的得分为 56.0，而 PyramidKV 降至 24.1。
大海捞针（NIAH）：IndexMem 在不同上下文长度（1K–12.5K）和针深度下表现出卓越的稳定性。虽然启发式方法显示出位置依赖性的失败（偏向最近的 token），但 IndexMem 无论针的位置如何都能可靠地检索。潜在记忆特别减少了仅使用索引器变体中出现的“灾难性遗漏”。
LongBench：IndexMem 在多样化任务（HotpotQA、MultiFieldQA、TriviaQA）中实现了最稳健的性能。与启发式方法相比，随着压缩率的增加，其性能下降更为平缓。值得注意的是，在某些任务中，适度的压缩甚至通过移除干扰 token 提高了性能。
效率：该方法引入了极小的参数开销（对于 8B 模型，索引器约 2000 万参数，记忆模块约 50 万参数）。虽然记忆模块增加了少量的更新开销，但整体解码时间与全缓存基线相比仍具有竞争力，且内存使用量随压缩比线性扩展。

5. 意义与主张

该论文声称，IndexMem 为长上下文推理提供了有利的精度 - 内存权衡。通过将淘汰从启发式转向可学习，并引入一种机制来恢复被丢弃的信息，该方法使得在有限的 KV 预算下实现准确的长上下文推理成为可能。

作者将这项工作定位为迈向高效注意力的可学习架构的一步。他们指出，虽然当前方法将淘汰和记忆模块作为冻结骨干的附加组件，但未来的工作可以探索端到端训练，使骨干网络学习原生的淘汰和记忆能力。论文强调，这是一种通用的效率技术，旨在降低长上下文模型的硬件成本并提高其可及性，同时也承认在高风险领域必须谨慎行事，因为遗漏关键细节可能会产生有害后果。