现代人工智能系统,即大型语言模型,在理解和生成人类语言方面已变得异常出色。它们能够撰写故事、解决数学问题并回答复杂的问题。然而,这些系统面临着一个显著的物理限制:它们一次只能在“工作记忆”中持有一定量的信息。当用户要求模型阅读一篇非常长的文档时,例如一整部小说或一系列法律记录,系统往往会表现得力不从心。随着输入内容的增长,它往往会遗忘文本开头的细节,或者在处理过程中失去对主旨的把握。这是因为模型用于连接词语的数学过程,随着文本规模的扩大,会变得极其昂贵且缓慢。为了解决这个问题,研究人员尝试了两种主要方法。一种是通过昂贵且耗时的重新训练,强迫模型学习如何处理更长的文本;另一种是使用一个辅助系统,在模型阅读之前先搜索文本中的相关部分,但这种辅助系统通常依赖于简单的词汇匹配,从而忽略了深层含义。
来自中国文华学院的一组研究人员提出了另一种解决方案,该方案既不需要重新训练,也不需要外部搜索工具。他们发现,模型本身已经具备了在长文中寻找正确信息的能力,只是需要被引导去使用这种能力。通过观察模型内部注意力机制的工作方式,他们发现,即使那些句子被埋藏在海量文档的深处,模型也会自然而然地将注意力集中在包含问题答案的特定句子之上。这种模式在模型处理网络的最后一层最为明显。研究人员构建了一种名为 InfiniRetri 的方法来利用这种自然行为。该方法并非试图将文本压缩成较小的、难以阅读的片段,也不依赖于独立的搜索引擎,而是迭代地阅读长文档的各个块(chunks),询问模型根据其自身的内部焦点识别出最相关的句子,并将仅有的这些完整句子保存到一个小型、高效的记忆缓存中。
这个过程就像一位细心的读者逐页阅读厚书。研究人员将长文本拆分为易于处理的部分。对于每个部分,他们将文本与用户的提问一起输入到模型中。随后,模型会观察其内部机制对该部分不同部分的关注程度。研究人员提取了这种注意力的图谱,特别观察了问题词汇是如何与文本中的词汇相连的。他们发现,模型的注意力评分充当了一个精确的信号,标示出了哪些句子持有答案。系统随后选取评分最高的句子,并将它们保存到一个轻量级的缓存中。这个循环在处理文档的每一个块时都会重复进行。当整个文本处理完毕时,缓存中仅包含了回答问题所需的关键句子,有效地过滤掉了其余的噪声。最后,模型阅读这个经过精选的小型句子集合,并生成最终答案。
这种方法的测试结果令人瞩目。在测试模型是否能在海量文档中找到隐藏事实的任务中,一个使用该方法的极小版本模型,即使在文本包含超过一百万个单词的情况下,也实现了完美的准确率。这是一个显著的进步,因为以往的方法在文本超过 32,000 个单词时往往就会失效。该方法也被证明在处理需要跨多份文档寻找信息的问答任务时非常有效,其表现优于其他依赖复杂训练或外部数据库的高级技术。值得注意的是,这种高性能是在仅使用两 GB 计算机内存的情况下实现的,这一容量足以在标准的消费级硬件上运行。研究人员指出,当存在一个明确的问题时,该方法的效果最好,因为系统依靠问题来引导其注意力。虽然在处理像总结整本书这样(目标是压缩所有内容而非寻找特定事实)的任务时,提升效果较小,但能够在无需昂贵训练的情况下,从几乎无限长度的文本中检索精确信息,代表了向前迈出的重要一步。这项研究表明,通过信任模型自身的内部信号,我们可以解锁其在任何规模文本上进行阅读与推理的能力。
技术摘要:InfiniRetri
问题陈述
大型语言模型(LLMs)在输入长度增加时会面临持续的性能下降,这影响了对长文档的直接检索和复杂推理。现有方案面临显著局限性:
- 上下文窗口扩展: 通过持续训练来扩展上下文窗口的成本极高,因为注意力机制具有二次方复杂度。
- 提示词压缩: 当前的方法(如 SelectiveContext, LLMLingua)通常在标记(token)层面进行操作,丢弃语义单元并产生不可读的文本,这在不进行训练的情况下会损害性能。
- KV 缓存压缩: H2O 或 SnapKV 等方法仅保留部分键值(Key-Value)状态,无法以无需训练的方式压缩所有状态,且往往忽略了细粒度的注意力模式。
- 检索增强生成 (RAG): 标准 RAG 依赖外部嵌入模型,这些模型仅能捕捉表层的词汇相关性,难以弥合语义鸿沟或建立多个检索段落之间的关联。
方法论:InfiniRetri
作者提出了 InfiniRetri,一种无需训练的方法,它利用预训练 LLM 内在的注意力机制,从任意长的输入中检索并缓存相关信息。该方法构建在对 注意力对齐检索 (Attention-Aligned Retrieval, AAR) 模式的发现之上。
核心洞察:注意力对齐检索 (AAR)
通过逐层可视化和分析,作者观察到:
- 检索能力: 预训练 LLM(特别是最后一个 Transformer 层)中的“问题-上下文”注意力子矩阵提供了一个精确且可解释的信号,用于识别与查询相关的内容。模型天生“知道”哪些上下文跨度是相关的。
- 层深度: 这种检索模式随深度增加而增强,在最后一个 Transformer 层达到顶峰,在该层中,相关上下文与无关上下文之间的注意力最具辨别力。
工作流
InfiniRetri 通过一个包含五个步骤的迭代循环来处理长上下文:
- 分块 (Chunking): 将长上下文在句子边界处分割成片段,以保持语义完整性。
- 合并 (Merging): 将之前检索到的句子的轻量级外部缓存前置到当前块中。
- LLM 推理: 将合并后的上下文和用户问题输入 LLM。
- 注意力中检索:
- 提取最后一层的注意力矩阵。
- 切片出 问题-上下文子矩阵(排除注意力汇点/attention sinks)。
- 使用一维求和池化(1-D sum-pooling)在短语层面聚合注意力分数。
- 根据聚合的分数识别出前 K 个最相关的候选位置。
- 缓存更新: InfiniRetri 不存储单个标记,而是检索包含前 K 个标记的 完整句子,并将其添加到缓存中。
处理完所有块后,使用累积的缓存和原始问题生成最终答案。
核心贡献
- 识别 AAR 模式: 本文指出,预训练 LLM 最后一层的“问题-上下文”注意力子矩阵可以作为准确、可靠且可解释的检索信号,从而消除了对外部嵌入模型的需求。
- InfiniRetri 框架: 这是一种无需训练的方法,它通过迭代方式将完整的、相关的句子提炼到轻量级缓存中。它符合提示词压缩的目标,但在句子层面进行操作,具有更好的可解释性和语义连贯性。
- 全面评估: 该方法在三类基准测试(提示词压缩、KV 缓存压缩和 RAG)中进行了评估,展示了卓越的性能。
实验结果
InfiniRetri 在 LongBench、LongBenchV2、∞Bench 和 大海捞针 (Needle-In-a-Haystack, NIH) 测试中进行了评估。
- LongBench 性能:
- 提示词压缩: 在平均得分上超过了之前的无需训练的最优方法 (DAC) 11.4% (42.05 vs. 37.76),且压缩率显著降低 (0.26 vs. 0.5)。
- KV 缓存压缩: 在 Qwen2-7B 上整体超越了最强的先前方法 CompilerKV(后者需要离线校准)15.9%。值得注意的是,在多文档问答任务中实现了 238.7% 的提升 (13.40 → 45.38)。
- RAG: 优于领先的基于嵌入的检索器 (BGE-M3, Stella-v5) 和先进的 RAG 系统 (GraphRAG, RQ-RAG)。仅次于受益于全阶段训练的 MemoRAG。
- 超长上下文能力:
- NIH 测试: 在 100 万个标记的 NIH 测试中,使用 0.5B 参数的模型实现了 100% 的准确率,且仅需 2GB GPU 显存。
- ∞Bench: 一个 1.5B 模型在超长检索任务(214K tokens)中达到了 100% 的准确率,超越了配备 KV 检索方法的更大模型。
- LongBenchV2: 使用 InfiniRetri 的 7B 模型达到了与 72B 模型相当的性能,特别是在长文本任务上。
- 复杂度: 通过将上下文剪枝为少量的相关句子缓存,该方法将推理复杂度从二次方 O(N2) 降低到近似线性 O(mN),其中 m 是块的大小。
重要性与主张
论文声称 InfiniRetri 建立了一种基于 LLM 内在能力 的信息检索新范式。通过利用模型自身的注意力模块,它在无需外部检索组件或训练的情况下,弥合了查询与文档之间的语义鸿沟。
作者强调:
- 可解释性: 与标记级压缩不同,InfiniRetri 保留了完整的句子,使检索过程透明且人类可读。
- 高效性: 它大幅降低了内存开销(例如,100 万个标记仅需 2GB)和计算成本,同时保持或提高了准确性。
- 泛化性: 该方法适用于任何基于注意力的 LLM,并且理论上支持无限长的上下文。
承认的局限性:
论文指出,InfiniRetri 需要预定义的提问(限制了没有问题生成环节的开放式场景),并且在摘要任务上的提升有限,因为这类任务可能需要聚合多个窗口之外的信息,超出了当前的单次检索机制。此外,性能对超参数调优较为敏感,特别是短语大小相对于目标答案长度的关系。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。