MindRef: 利用分层检索教 LLM 像人类一样记忆

想象一下，你试图回忆几年前读过的一本书中的某个具体细节——比如《哈利·波特》中某个小角色的名字。你的大脑不会线性地扫描你读过的每一本书的每一句话。相反，你可能会进行分层搜索: 首先，你会回忆起具体的书名 (《哈利·波特与火焰杯》) ，然后在脑海中将范围聚焦到相关场景以检索名字。

这种“从文档到细节”的过程高效、灵活且自然。然而，目前人工智能使用的传统信息检索 (IR) 系统并非如此运作。相反，它们依赖于将海量数据库预分割 (切块) 成固定大小 (例如 100 个单词的块) 并对其进行索引。当你提出问题时，系统会搜索这数百万个互不相连的块以寻找匹配项。这往往会破坏上下文并限制灵活性。

但是，如果大型语言模型 (LLM) 能够模仿人类的记忆呢？如果它们能首先识别正确的文档，然后从任意位置开始“阅读”特定段落，而不依赖于任意的预切块呢？

这就是 MindRef 的前提，这是由研究人员 Wang、Xu 和 Ding 提出的一个突破性框架。在这篇文章中，我们将剖析他们的研究论文，了解他们如何将经典数据结构 (前缀树和 FM 索引) 与现代 LLM 相结合，创建一个既在认知上合理又在计算上高效的检索系统。

“切块”带来的问题

在深入探讨解决方案之前，我们需要了解当前方法的瓶颈。

在标准的检索增强生成 (RAG) 流程中，知识源 (如维基百科) 被切分为块。如果相关答案跨越了两个块的边界，或者切块大小太小无法提供上下文，检索模型就会失效。

最近的进展提出了生成式检索 , 即训练 LLM 记住文档标识符 (如标题) 。当被问及问题时，LLM 会生成相关文档的标题。虽然这种方法很有前景，但它通常指向整个文档 (太长而无法完全阅读) 或预定义的段落 ID (这又回到了切块问题) 。

MindRef 背后的研究人员提出了一个简单的问题: “LLM 能否绕过切块，从任意位置回忆参考资料？”

图 1: 人类回忆遗忘信息通常涉及两个步骤: 首先回忆起印象深刻的文档，然后在其中定位具体段落。

如上图所示，目标是模拟人类的两步过程:

  1. 第一阶段: 回忆更广泛的上下文 (文档标题) 。
  2. 第二阶段: 聚焦回答查询所需的具体细节 (段落) 。

MindRef 框架

MindRef 是一个旨在实现“无切块”参考定位的两阶段检索框架。它在推理过程中不依赖 Elasticsearch 或 Faiss 等外部搜索引擎。相反，它利用 LLM 在预训练期间已经存储的参数化知识，并由巧妙的数据结构引导以确保准确性。

该架构在关注点分离方面非常优雅:

图 2: MindRef 框架。首先，所有维基百科标题都存储在前缀树中，然后提示 LLM 在此限制下回忆标题标识符。随后，根据获得的前 k 个文档构建 FM 索引，LLM 在新限制下回忆参考段落。

让我们详细分解这两个阶段。

第一阶段: 粗粒度文档召回

第一步是将搜索范围从“整个互联网”缩小到“少数相关文档”。模型接收一个查询 (例如，“佐助在他的氏族灭亡时几岁？”) ，并必须生成可能包含答案的维基百科页面标题。

然而，LLM 容易产生幻觉。如果我们只是让 LLM “列出一个关于佐助的维基百科页面”，它可能会编造一个不存在的标题。为了防止这种情况，MindRef 使用了通过 Trie (前缀树) 进行的 受限解码 (Constrained Decoding) 。

Trie 的作用

Trie 是一种树形数据结构，用于存储动态字符串集，其中键通常是字符串。在 MindRef 中，研究人员构建了一个包含所有有效维基百科标题的 Trie。

当 LLM 生成文档标题时，它不允许从词汇表中随意选择 token。它受到 Trie 的“引导”。如果 LLM 生成了 “Sasuke Uch”，Trie 会查看存储在其分支中的有效下一个字母 (例如 “i”) ，并强制 LLM 从有效的后续字符中进行选择 (从而生成 “Sasuke Uchiha”) 。这保证了生成的每个标题在数据库中都是实际存在的。

图 3: 受限解码方法: (a) 显示了 LLM 使用前缀树生成标题标识符的过程。(b) 显示了 LLM 通过 FM 索引在文档集中生成段落前缀的过程。

如上图面板 (a) 所示，Trie 限制了搜索空间，确保标题生成的有效性。模型使用束搜索 (beam search) 生成多个候选标题，并对每个标题进行评分。生成的标题 \(t\) 的分数是基于其 token 的平均对数概率计算的:

公式 1: 第一阶段的分数计算

这里，\(y_t\) 代表标题中的 token，分数反映了 LLM 对该标题与提示相关性的置信度。

第二阶段: 细粒度段落召回

一旦系统获得了前 \(k\) 个相关文档 (例如，“宇智波佐助”和“宇智波鼬”的维基百科页面) ，下一个挑战就是找到回答问题的确切句子或段落。

我们不能简单地将整个文档输入 LLM 的上下文窗口——它们可能太长了。我们也不想使用预切的块。相反，我们要让 LLM 从文档中生成相关的逐字文本。

但这再次面临幻觉风险。LLM 可能会生成一个听起来合理但实际上并不存在于源文本中的句子。为了解决这个问题，MindRef 采用了一种更复杂的数据结构: FM 索引 (FM-Index) 。

FM 索引

虽然 Trie 适用于像标题这样的短字符串，但索引成千上万篇维基百科文章的全文会占用过多内存。FM 索引 (基于 Burrows-Wheeler 变换) 是一种压缩数据结构，支持快速子串查询。

对于第一阶段检索到的前 \(k\) 个文档，MindRef 构建了一个针对性的 FM 索引。当 LLM 开始生成段落时，FM 索引充当类似于 Trie 的约束机制。它动态检查目前为止生成的 token，并将 LLM 的下一个 token 概率分布限制为仅包含源文档中实际出现的词。

其底层机制依赖于 Burrows-Wheeler 变换矩阵，该矩阵对文本进行排列以使其可压缩和可搜索:

FM 索引中使用的 Burrows-Wheeler 变换矩阵表示

通过使用 FM 索引，LLM 可以从文档中的任意位置开始生成文本。它实际上是从记忆中“读取”文档，但 FM 索引确保它不会错引任何一个词。

此阶段的评分与第一阶段类似，计算生成段落的置信度:

公式 2: 第二阶段的分数计算

结合信号

检索段落的最终相关性不仅取决于段落与查询的匹配程度，还取决于父文档的相关性。MindRef 使用加权和结合两个阶段的分数:

公式 3: 第一阶段和第二阶段分数的加权和

这里，\(\alpha\) 是一个超参数，用于平衡文档标题相关性与特定段落相关性的重要程度。

优化: 短前缀召回与定位 (SPRL)

逐个 token 地生成完整的 200 个 token 的段落是很慢的。为了使 MindRef 具有实用性，研究人员引入了 短前缀召回与定位 (SPRL) 。

LLM 不需要生成整个段落，而是被提示仅生成一个 短前缀 (例如，相关句子的前 16 个 token) 。由于 FM 索引的约束，这个短序列通常足以在文档中精确定位一个位置。

该过程如下:

  1. 召回: LLM 生成一个短前缀 (例如，“胰凝乳蛋白酶的主要底物是…”) 。
  2. 定位: 系统使用 KMP (Knuth-Morris-Pratt) 字符串匹配算法在文档中找到该前缀的确切偏移量。
  3. 扩展: 一旦找到位置，系统只需直接从原始文件中提取周围的文本 (例如，接下来的 150 个单词) 。

这项技术将推理速度提高了 4 倍 , 同时保持了超过 95% 的准确率。

实验与结果

研究人员在 KILT 基准 上评估了 MindRef，这是一套知识密集型任务，包括开放域问答 (Natural Questions, TriviaQA) 、事实验证 (FEVER) 和对话系统 (Wizard of Wikipedia) 。

他们将 MindRef 与以下模型进行了比较:

  • BM25: 标准的基于关键词的稀疏检索方法。
  • Contriever: 一种流行的无监督密集检索模型。
  • DPR: 密集段落检索 (监督/微调) 。

1. 它能找到正确的文档吗？ (页面级结果)

第一个测试是 MindRef 是否能识别正确的维基百科页面。

表 1: 粗粒度页面级结果 (R-Precision) 。MindRef 显著优于 BM25 和 Contriever。

如表 1 所示，MindRef (使用 LLaMA-2-13b) 在无监督基准测试中占据主导地位。在 FEVER 数据集 (事实验证) 上，它达到了 83.69 的 R-Precision，大大超过了 BM25 (52.09) 。它甚至在某些任务上超越了完全监督的 DPR 模型。这证实了 LLM 拥有强大的内部知识“索引”，可以通过标题生成来解锁。

2. 它能找到正确的答案吗？ (段落级结果)

找到文档只是成功了一半。表 2 显示了 MindRef 检索包含答案的具体段落的能力。

表 2: 细粒度段落级结果。MindRef 表现出色，特别是在 TriviaQA 和 FEVER 上。

MindRef 在 TriviaQA、HotpotQA 和 FEVER 上始终获得最高分。虽然 DPR 在 Natural Questions (NQ) 上表现更好，但 MindRef 在无需检索数据预训练的情况下定位段落的能力令人印象深刻。这表明“先召回再定位”的方法是“先索引再搜索”的可行替代方案。

3. 它对下游任务有帮助吗？

最后，研究人员将检索到的段落输入到阅读器模型中，以回答实际的用户问题。

表 3: 下游任务结果。

表 3 中的结果证实，更好的检索能带来更好的回答。经 MindRef 增强的 LLaMA-2 在 FEVER 上取得了最先进的结果 (78.79% 的准确率) ，并在 TriviaQA 上表现强劲。

消融实验: 什么最重要？

研究人员去除模型的某些部分，以观察是什么驱动了性能。

表 4: 消融研究结果显示了去除权重、SPRL 和第一阶段的影响。

  • w/o weight (无权重) : 去除加权分数 (公式 3) 会损害性能，证明文档相关性 (第一阶段) 是段落选择的关键上下文。
  • w/o SPRL (无 SPRL) : 令人惊讶的是，去除短前缀优化 (转而生成全文) 有时反而会降低性能。作者认为，生成更长的文本会引入更多束搜索偏离或引入噪声的机会。
  • w/o first stage (无第一阶段) : 跳过文档召回并尝试立即搜索整个语料库会导致准确率大幅下降。分层方法至关重要。

分析: 前缀长度的影响

论文中最有趣的分析之一是关于 SPRL 中使用的前缀长度。直觉上，你可能会认为前缀越长，准确率越高。

图 5: 不同前缀 token 长度的实验结果。

图 5 揭示了一个反直觉的发现: 更长并不总是更好。 在 HotpotQA 和 NQ 数据集上，性能在相对较短的前缀长度 (约 8-16 个 token) 处达到峰值，然后下降或趋于平稳。这验证了 SPRL 策略: LLM 只需要记住足以“指向”位置的内容即可。一旦确定了位置，硬盘驱动器就可以完成剩下的工作。

结论与未来启示

MindRef 代表了我们对检索思考方式的重大转变。MindRef 不再将文本视为由外部搜索引擎索引的静态块集合，而是将文本视为 LLM 可以利用其内部记忆进行导航的连续流。

核心要点:

  1. 模仿人类: 两阶段的“文档标题 -> 具体段落”工作流程比传统的密集检索更接近认知过程。
  2. 无需预切块: 通过使用 FM 索引，检索可以从任意字符开始，提供了基于块的方法所缺乏的细粒度灵活性。
  3. 效率: SPRL 方法证明，LLM 在检索阶段最好充当“指针”而非“阅读器”——回忆一个短片段足以找到完整数据。

局限性: 虽然 MindRef 功能强大，但它严重依赖 LLM 的预训练。如果文档标题晦涩难懂或不在训练数据中，模型在第一阶段就难以回忆起它。此外，用新知识更新模型需要重新训练 (或复杂的提示) ，而传统的搜索引擎只需索引新文件即可。

尽管如此，MindRef 让我们窥见了一个未来，在这个未来中，“记忆” (神经网络) 和“存储” (数据库) 之间的界限变得越来越模糊，从而带来更流畅、更强大的人工智能系统。