国思软件 - 无需训练，100%完美检索！LLM练出「火眼金睛」，InfiniRetri超长文本一针见血

　　新智元报道

　　编辑：KingHZ

　　LLM 自身有望在无限长 token 下检索信息！无需训练，在检索任务「大海捞针」（Needle-in-a-Haystack）测试中，新方法 InfiniRetri 让有效上下文 token 长度从 32K 扩展至 1000+K，让 7B 模型比肩 72B 模型。

　　全新检索模式：在无限长 token 下，大语言模型自身或能检索信息！

　　受大语言模型（LLM）上下文窗口大小的限制，处理输入 token 数超过上限的各种任务颇具挑战性，无论是简单的直接检索任务，还是复杂的多跳推理任务。

　　尽管新提出的各种方法用来增强大语言模型的长上下文处理能力，但这些方法痛点突出：

　　要么会产生高昂的训练后成本，

　　要么需要额外的工具模块（如检索增强生成 RAG），

　　要么在实际任务中显示出改进，并不明显。

　　研究团队观察了各层注意力分布与生成答案之间的相关性，通过实验证实了注意力分配与检索增强能力是一致的。

　　基于上述见解，研究团队提出了一种全新的方法 InfiniRetri，该方法利用大语言模型自身的注意力信息，实现对任意长度输入的精确检索。

　　实验表明，在 100 万个 token 的「大海捞针」（Needle-In-a-Haystack，NIH）测试中，InfiniRetri 将 5 亿参数的模型从 44.6% 提升到了 100% 的准确率。

　　没有 InfiniRetri 的 NIH 测试只有 44.6% 的准确率

　　要知道 GPT-4 在 NIH 测试中都做不到 100% 准确率。

　　InfiniRetri 一举超过了其他方法或更大的模型，创造了当前最佳（SOTA）结果。

　　值得注意的是，某 7B 模型在 HotpotQA 任务上的得分，超越了其他同等参数规模的模型。

　　类似地，Mistral-7B-Instruct v0.2 作为擅长短文本推理的模型，在长文本任务中的表现也得到了显著提升。

　　此外，新方法在实际基准测试中也取得了显著的性能提升，最大提升幅度达到 288%。

　　另外，无需额外训练，InfiniRetri 就可应用于任何基于 Transformer 的大语言模型，并且能大幅降低长文本推理延迟和计算开销。

　　论文链接：https://arxiv.org/abs/2502.12962

　　项目地址：https://github.com/CapitalCode2020/InfiniRetri2

　　文章主要贡献如下：

创新性提出「注意力分配与检索增强对齐」概念，并成功利用这一特性提升 LLM 处理长文本的能力。
无需额外训练，InfiniRetri 可直接应用于任何基于 Transformer 的 LLM，使其具备处理无限长上下文的能力。
不同于 RAG 依赖外部嵌入模型，提出了「注意力中的检索」这一新颖视角，充分利用 LLM 内在能力增强长文本处理能力。
显著降低推理延迟和计算开销，在大规模数据集上的检索与问答任务中表现优异，展现出在极长文本处理场景中的实际应用价值。
不仅仅是扩展上下文窗口，证明了可以通过多种方式提升 LLM 处理长文本能力。未来的改进可通过在较小上下文窗口内增强模型内部能力，从而获得更优的长文本处理效果。

　　鞭辟入里，灵魂三问

　　近年来，许多主流的 LLM 都在致力于扩展上下文窗口的规模。

　　例如，GPT-4、Llama3 和 DeepSeekV3 支持最长128K的上下文长度，Claude-3 可达200K，而 Gemini-Pro-1.5 和 Qwen2.5-1M 甚至支持1M的上下文长度。

　　然而，实际效果并未完全达到这些模型所宣称的长度。

　　LLM 在处理超长上下文方面仍然面临重大挑战。

　　那自然考虑：是否能不断延长上下文窗口？

　　这是第一个问题，答案是不能，有 3 大原因：

增加计算成本，产生明显的延迟。
输入长度具有长尾效应，简单地扩展上下文窗口带来的提升注定会带来越来越低。
持续和阶段性训练成本极高，绝大多数研究人员无法承受。

　　接下来，进一步考虑第二个问题：要打破处理长上下文时不同窗口之间的信息壁垒，是否存在一种低成本的方法？

　　事实上，业界已经通过检索增强生成（Retrieval-Augmented Generation，RAG）提供了答案。

　　RAG 由两大核心组件组成：检索模块和生成模块。其中，检索模块依赖外部嵌入模型，根据输入查询从长文本中检索相关内容。

　　然而，RAG 系统普遍难以建立检索信息之间的关联。

　　相比之下，在推理过程中，LLM 的注意力机制，能够高效地建立不同信息片段之间的联系。

　　由此引出了第三个问题：为何不直接利用 LLM 自身的检索能力，来处理长文本上下文？

　　为了让 LLM 以低成本的方式压缩并存储过去的键（key）和值（value），研究团队认为，只有打破不同上下文窗口之间的信息壁垒，LLM 才能真正提升其处理长文本的能力。

　　此外，通过观察 LLM 在推理过程中回答问题时的注意力分配模式，研究团队提出：这种注意力分配模式与检索增强能力（RAG）高度契合。

　　研究团队测试了基于所有层的注意力分布来检索答案的准确性。

　　正如下图 3 所示，越接近输出层，这种模式的增强效果就越明显。

　　此外，在第 14 层和第 15 层，检索准确率达到一个局部峰值。

　　因此，InfiniRetri 引入了全新的策略，利用 LLM 自身的注意力信息，而非依赖外部嵌入模型，去提升其长文本处理能力。

　　无需对基于 Transformer 的 LLM 进行额外训练，InfiniRetri 可以开箱即用。得益于此，研究团队在多个模型上进行了全面的对比实验。

　　在跨上下文长度的事实检索任务「大海捞针」（Needle In A Haystack，NIH）中，InfiniRetri 仅使用 5 亿个额外参数，就将模型的上下文长度从原始的 32K 扩展至 100 万个 token（如图 1 所示）。

　　原始方法（上）与 InfiniRetri 方法（下）对比，准确检索的最大文本长度从原始 32K 提升至超过 1M tokens。

　　更值得注意的是，InfiniRetri 在 NIH 任务上实现了无限长度范围内的精准检索，不仅超越了当前主流方法，还有效解决了 NIH 任务的挑战。

　　此外，在 LongBench 提供的 9 个真实数据集上，InfiniRetri 在基于 KV Cache 甚至 Full KV 的主流方法中均取得了超越性的表现，尤其是在多文档问答（Multi-Document QA）任务（如 HotpotQA、2WikiMQA 和 Musique）中，Qwen2-7B-Instruct 采用 InfiniRetri 后，平均性能提升高达 369.6%。

　　InfiniRetri 方法：带着问题阅读

　　那么，如何应用这一模式来处理超出上下文窗口的长文本，并真正提升 LLM 的长文本处理能力呢？

　　在本节中，将 InfiniRetri 方法拆分为三个小节，分别介绍该模式的应用方式，以提升 LLM 的长文本处理能力。

　　如图 4 所示，该方法的完整工作流程包括五个主要步骤，详细介绍这些步骤：

步骤1（切分，Chunk）、步骤2（合并，Merge）和步骤3（推理，Inference）。
步骤4（检索，Retrieval）是方法的核心部分。
步骤5（缓存，Cache）。

　　图4：InfiniRetri 方法在增强 LLM 长上下文处理能力中的完整工作流程

　　文本分割

　　新方法受到人类阅读书籍过程的启发，专门针对 LLM 在处理超出上下文窗口的文本时所面临的挑战。

　　尽管人类的视野有限，一次只能看到一页内容，但仍然可以逐页阅读并理解整本书。在这个过程中，大脑就像一个缓存（cache），通过记忆保留并整合每一页的信息，从而掌握整本书的内容。

　　类似地，InfiniRetri 将整篇文本拆分为连续的文本块（chunk）。

　　这种切分方式虽然与 RAG 相似，但不同之处在于，InfiniRetri 不是并行处理每个文本块，而是按照顺序逐块迭代地处理每个文档。

　　这种方法能够保留文本的顺序信息，更符合人类的阅读习惯。

　　具体而言，如图 4 所示，在步骤1（切分，Chunk）中，研究团队根据句子边界将整个长文本划分为长度大致相等的文档块，其长度由方法参数 ChunkSize³决定。

　　然后，这些文档块依次与缓存（Cache）中保留的 token 进行合并，形成完整的输入序列，称为 MergeToken，并将其输入 LLM 进行处理。

　　InfiniRetri 采用了类似滑动窗口注意力（Slide Window Attention，SWA）的迭代方式，按顺序处理每个文本片段。

　　然而，InfiniRetri 对缓存的处理方式与传统方法有本质区别。

传统缓存通常在每一层存储过去的键值（Key-Value）状态，而新方法则重新定义了缓存概念，改为存储过去的 token ID。
如图4，步骤2（合并，Merge）所示，新方法在输入 LLM 之前，将缓存的 token ID 与当前文本片段的 token 进行合并，从而取代了推理过程中对历史键值状态的合并需求。
因此，在步骤3（推理，Inference）阶段，LLM 仍然使用标准注意力机制，而非 SWA。对于第h层的注意力得分，其计算公式如下：

　　其中，A^h∈R^{n×m}表示查询（query）和键（key）组成的注意力矩阵，n是查询的数量，m是键的数量。

　　在注意力中检索（Retrieval In Attention）

　　在单个上下文窗口内，根据问题 token 准确定位相关的上下文 token，注意力分配模式能够帮助 LLM 找到正确答案。

　　如果在滑动窗口框架内的每次推理过程中持续应用这一模式，理论上，LLM 就可以在保持查询不变的情况下，对整个长文本进行推理。

　　这一过程与人类的阅读方式高度相似，类似于公认的「带着问题阅读」学习策略，即通过问题作为锚点，在 LLM 可处理的范围内逐步整合相关信息。

　　因此，LLM 能否精准检索与问题最相关的文本，是本方法有效性的核心。

　　关键在于设计基于注意力得分分布的 token 检索策略和算法，以确保模型能够在长文本中高效提取关键信息。

　　借鉴实验结果，研究团队选取了多头注意力（Multi-Head Attention）的最后一层，并对所有注意力头的得分进行求和聚合（如公式 2 所示），探索了一种能够准确判断模型关注重点的方法。

　　通过可视化注意力得分，研究团队观察到：与答案相关的信息通常由连续的 token 组成。

　　也就是说，它们以短语级别的粒度存在。

　　这一发现与在图 2c 的实验结果一致，进一步确认了 LLM 在 token 级别上具备较高的注意力精度。

　　因此，希望设计的操作是：计算每个 token 及其相邻 token 在 2D 注意力得分矩阵中的累积注意力得分。

　　计算结果将作为新的特征，在后续的检索过程中用于排序。经过深入分析，发现此操作等效于使用一个填充了 1 的卷积核（kernel）进 1D 卷积。

　　对于查询i和键j，其特征重要性计算如下：

　　1. 聚合所有注意力头的得分（公式2）：

　　2. 基于 1D 卷积计算每个 token 及其相邻 token 的重要性（公式3）：

　　其中，k是 1D 卷积核的大小，对应于方法中的参数 Phrase Token Num。

　　3. 沿矩阵的列方向求和，计算每个上下文 token 的总重要性分数（公式4）：

　　其中，s_i代表第i个上下文 token 的综合重要性分数。

　　4. 最后，选取重要性分数最高的前K个上下文 token，并将其所在句子的所有 token 写入缓存（cache）。这个过程可表示为：

　　即，从所有 token 的重要性分数v中，选择排名前K的 token，并将它们所在的完整句子存入缓存，以便后续推理时使用。

　　缓存句子 Token（Cache Sentence Token）

　　在推理过程中对缓存（cache）的使用方式，InfiniRetri 与传统方法存在本质性区别。

　　相比于直接使用缓存，InfiniRetri 将其用于存储过去的上下文信息。

　　具体而言，主要有以下两点不同：

　　1 新方法在模型外部缓存 token ID，而不是每层的历史键值（Key-Value）状态。具体来说，在推理过程中不使用传统的 Key-Value 缓存，而是在每次推理前，将过去的上下文信息与当前输入合并后再进行推理。

　　2 新方法基于短语级特征进行检索，并在缓存中存储包含 Top-K token 的句子级 token。也就是说，存储的是完整的句子，而不是单独的 token，从而确保检索到的信息更具上下文完整性。

　　事实上，正是这两项创新性改进，使得新方法在无需微调的情况下，就能比传统的 KV 缓存方法更有效地提升 LLM 处理长文本的能力。

　　新方法并不试图压缩缓存中的 token，而是保留句子级别的相关上下文信息。这是因为，句子是最小的完整语义单元，相比于单个 token，更能确保 LLM 对上下文的理解。

　　在 LLM 逐步推理每个文本片段的过程中，缓存中保留的中间结果是动态变化的，它们由先前存储的 token 和当前输入片段的组合决定。因此，在整个过程中，这些中间结果会相对调整和更新，以适应模型的理解需求。

　　AI 界的「大海捞针」

　　大海捞针（Needle-in-a-Haystack，NIH）任务要求模型在一篇「超长文档」（「大海捞针」之海）中，精准检索出一个特定的目标句子（「针」），该句子可以被随机插入到文档的任意位置。

　　通过调整「针」的放置位置（文档深度）和上下文长度，反复测试以衡量模型的表现。

　　为了直观分析模型的检索能力，采用了热力图（heatmap）可视化实验：

绿色代表完美检索（准确找到目标句），
其他颜色表示检索错误。

　　这种可视化方法可以直观展示 LLM 处理长文本的能力上限，因此被广泛用于评估。

　　实验1：Llama3-8B-Instruct 的对比

　　如图 6 所示，在 Llama3-8B-Instruct 模型上测试 NIH 任务，并与以下方法进行了对比：Full KV（完整 KV 缓存）

　　StreamingLLM、H2O、SnapKV、PyramidKV 以及 InfiniRetri（新方法）。

　　在最长 32K token 的输入文本上进行实验，结果表明：

传统 KV 缓存压缩方法虽有所改进，但没有超越 Full KV 的性能。
InfiniRetri 方法的表现优于 FullKV，显著增强了 Llama3-8B-Instruct 处理 NIH 任务的能力，甚至突破了原始 8K token 上下文窗口的限制。

　　实验2：Mistral-7B-Instruct 的对比

　　为了进一步验证 InfiniRetri 的有效性，在 Mistral-7B-Instruct 上扩展了输入长度。

　　Mistral-7B-Instruct 官方支持的上下文窗口为 32K token，对比了 FullKV 和 InfiniRetri 的表现，如图 7 所示：

Mistral-7B-Instruct+FullKV（图 7a）：最多可在 43K token 长度范围内正确完成 NIH 任务。
Mistral-7B-Instruct+InfiniRetri（图 7b）：在相同的参数设置下，新方法不仅超越了 Llama3-8B-Instruct，而且在 NIH 任务上达到了 100% 的检索准确率，并将可处理的输入长度扩展至 1M token，且没有损失准确率。

　　关键发现与额外实验

　　进一步观察到：只要 LLM 在有限上下文窗口内具备足够的检索能力，新方法就可以赋能模型处理超长文本的检索任务，理论上可支持无限长输入。

　　基于这一发现，在更小的开源模型上的额外实验，结果符合预期：

　　新方法将该模型的有效上下文 token 长度从 32K 扩展至 1M+，从而使其在 NIH 任务上具备了近乎无限的长文本处理能力（如图 1 所示）。

　　LongBench 实验

　　从表 1 的整体实验结果来看，新方法是唯一一个在所有模型上全面超越 Full KV 的方法，其中在文档问答（DocumentQA）任务中的提升最为显著。

　　主要实验结果：

LLaMA3-8B-Instruct：相对提升 4.9%（32.92→34.56）
Qwen2-7B-Instruct：相对提升 70.5%（25.11→42.82）
Mistral-7B-Instruct-v0.3：相对提升 55.8%（24.17→37.68）

　　其中，Qwen2-7B-Instruct 在 HotpotQA 任务上的表现提升最为显著，最大增幅达到 288%（14.8→57.52）。

　　关键发现

值得注意的是，Qwen2-7B-Instruct 在 HotpotQA 任务上的得分，超越了其他同等参数规模的模型，表明其在短文本推理方面的优势。这进一步说明，Qwen2-7B-Instruct 通过新方法，能够有效提升其长文本推理能力。
类似地，Mistral-7B-Instruct v0.2 作为擅长短文本推理的模型，在长文本任务中的表现也得到了显著提升。

　　随后，虽然新方法在长文档问答任务中取得了显著改进，但在文档摘要任务上的表现相对较差。

　　这种差异可能源于摘要任务的性质，这些任务通常需要更丰富的上下文信息来生成高质量的输出。

　　新方法无法一次性访问所有相关信息，这在一定程度上限制了它在这些任务中的有效性。

　　与问答和检索任务不同，在这些任务中，答案通常依赖于长上下文的一小部分，摘要任务则高度依赖于对整个上下文的全面理解。

　　因此，新方法可能需要进一步的优化和改进，以更好地应对这些摘要任务。

　　为了进一步评估 InfiniRetri 的有效性，使用最新的 Qwen2.5-7B-Instruct 模型在 LongBenchV2 上进行了额外的实验。

　　正如表 2 所示的结果，在应用了新方法 InfiniRetri 后，在处理 LongBenchV2 上的长文本和中等长度文本方面，Qwen2.5-7B 整体性能与 72B 的对比模型相当，表现出显著的改进。

　　这一结果进一步验证了，只要 LLMs 在短上下文场景中表现出色，新方法就可以有效地提高其处理更长上下文文本的能力。