无限上下文 LLM：超越 RAG 的新范式

当下广为流行的大语言模型，针对的是这样一项任务进行优化：生成看起来仿佛曾经出现在训练语料库中的 token 序列。这与 ChatGPT 或 Perplexity.ai 等用户界面使用 LLM 的方式截然不同——在这些界面中，用户期望模型执行复杂的推理任务，并如实检索事实性强、与当前主题相关的信息。如果我们希望将模型用作通用推理智能体，而不是一只随机鹦鹉，就需要在推理时向它提供所有相关数据，而不能依赖于：（1）关键数据曾经出现在训练语料库中；（2）模型能够回忆起这些数据。此外，要构建真正能够增强人类工作流的应用，展示参考资料或引用、明确指出模型在生成过程中使用了哪些内容也至关重要。

这推动了许多通常被统称为“检索”¹的方法的发展。换言之，就是帮助 LLM 利用相关文档的方法。上下文学习，即在提示词之前将相关文档放入上下文窗口，是最显而易见的第一步。然而，在许多情况下，我们面对的文档长度会超过模型的上下文窗口。RAG² 试图通过选出最合适的文档子集，并将其与用户查询一同提供给模型来绕开这一问题。尽管 RAG 通常很有效，但它从根本上受限于对独立搜索引擎的依赖。例如，我们无法向模型提出需要综合整个文档集合才能回答的问题。此外，由于检索发生在生成之前，在可解释性方面，我们最多只能报告提示词本身包含了哪些文本。这完全无法说明模型在生成过程中实际使用了哪些文本。

微调³试图扩展上下文窗口本身的长度。对于当今的大模型而言，即便拥有专门的机器学习团队，仅运行几个 epoch 的训练也可能是一项不可小觑的工作。此外，这些方法并不能提升模型的可解释性。还有一些方法⁴建议对模型进行结构性修改。其中许多方法都令人振奋，但大多数需要从头开始训练或进行微调，因此很难应用于预训练模型。

在本文中，我们提出并开源了扩展心智 Transformer，它在模型内部对 RAG 进行了泛化。这一简单的数学泛化不仅能带来 RAG 的性能提升，甚至还能实现更多能力，包括全新的生成控制手段和细粒度的因果引用。在易用性方面，我们也兼得两者之长：无缝集成（一切都在模型内部完成），而且无须微调！

扩展心智 Transformer 的美学基础

作为研究动机，我们将介绍一些心灵哲学背景，它们启发了本文的命名方式和方法论。在 Clark 和 Chalmers 于 1998 年发表的《扩展心智》（“The Extended Mind”）中，他们提出：持续且即时可访问、并会被自动认可的外部信息，应当被视为记忆的一部分。更进一步，这种扩展也应被视为心智的一部分。他们将这一思想称为主动外在主义。Otto 的故事则充当了一个直觉泵：

在本文中，我们提出面向 LLM 的主动外在主义：一种在美学上受到扩展心智论题启发、用于增强 Transformer 记忆的机制。我们将实现了主动外在主义的 Transformer 称为扩展心智 Transformer。

扩展心智 Transformer

我们提出的方法与 Wu 等人（2022）⁶的工作非常相似，只需对自注意力机制进行一项简单修改。除了 Transformer 不可或缺的因果自注意力之外，我们还允许每个查询 token 关注固定数量的“外部记忆”。这些记忆存储在不可微分的缓存中。每个解码器层和注意力头都通过余弦相似度来选择需要关注的记忆。更准确地说，我们的注意力计算如下：

其中，(𝐾𝐿,𝑉𝐿) 是来自局部上下文的键值对，(𝐾𝑅,𝑉𝑅) 是来自外部记忆的键值对，⊕ 表示张量拼接。我们对注意力权重进行掩码处理，使每个查询 token 只能关注它自己检索到的键，而不能关注前一个或后一个查询 token 检索到的键。在下文介绍的实验中，我们使用的是通过线性偏置而非位置编码训练的模型。当我们将这些线性偏置应用于注意力权重时，会为所有检索到的记忆分配相同的索引。⁷

重要的是，主动外在主义会精确检索记忆——除了线性偏置的影响之外，它不会对记忆进行摘要，也不会以其他方式削弱记忆。

我们只生成一次外部记忆（键值对），然后以类似于传入先前“缓存”键值对⁸的方式，将这些表示传递给每个解码器层。为了加快 top-k 余弦相似度计算，我们可以使用专门为此目的设计的向量数据库⁹。

我们认为，这种关注外部记忆或信念的方式，是对 RAG 等方法自然而又最优的泛化，并且非常贴近 Otto 与其笔记本之间的关系。信息持续且即时可访问，会被自动认可，并能得到可靠引用。我们设置了一个相似度阈值，以确保始终引用外部记忆（对于生成的每个 token，以及所有解码器层都是如此），但会丢弃未达到某个较低相似度阈值¹⁰的键，以免无关信息干扰模型。

主动外在主义在概念上并不难实现，但确实需要熟悉特定模型的实现方式，因为键值对的存储方式，以及将其读入自注意力计算的方式等细节都需要被接管。

困惑度实验

我们使用困惑度作为衡量模型性能的指标。困惑度用于衡量模型对每个生成 token 的不确定性，与我们的交叉熵损失函数密切相关。若想全面了解困惑度这一指标，我们建议阅读这篇优秀的文章。

下面展示的是使用 Mosaic 的 MPT-7b 模型，在 Wikitext-103 基准¹¹上进行困惑度实验的结果。在困惑度实验中，我们使用 512 个 token 的步长，这意味着，只要确实存在 512 个可供条件化的 token，每个 token 都会以此前至少 512 个 token 为条件。

我们的主动外在主义方法将每个序列分成长度不断增加的块（横轴），并将最后 2048 个 token（最大序列长度）之前的 token 作为外部记忆加以关注。我们展示了不同 k 值下的结果，其中 k 表示为每个查询 token 检索的记忆数量。我们将主动外在主义与两种基线方法进行比较。“truncated”基线在计算困惑度时，直接丢弃最后 2048 个 token 之前的所有 token；“naive”方法则使用输入长度范围内的全部 token，无论序列变得多长。

对于 naive 方法，我们恰好观察到了主动外在主义试图改善的现象：当序列长度超过 2～3k 个 token 后，性能会迅速下降（在本例中，困惑度急剧上升）。

我们可以看到，在 truncated 基准场景下，与仅使用局部注意力相比，主动外在主义带来了明显收益。更令人振奋的是，随着每个查询 token 检索到的记忆数量增加，困惑度还会持续下降。

检索实验

我们还在检索基准上衡量了性能，并将其与 RAG 和简单基线进行了比较。我们的数据集是 Abacus.AI 最近发布的 Long context WikiQA 基准的一个修改版本。

我们的目标是衡量模型在不同文档长度下的检索能力，但同时也希望控制模型在训练期间记忆的事实所产生的影响，因此我们通过将标注答案替换为看似合理但实际错误的答案来修改数据集。例如，在歌曲“These Boots Were Made For Walking”的 Wikipedia 词条中，我们将每一处“Lee Hazlewood”都替换为“Terry Allen”，然后要求模型给出词曲作者的姓名，此时正确答案就变成了“Terry Allen”。

我们的意图是衡量模型优先采用上下文或记忆中的事实，而非训练期间记住的事实的能力。再次强调，如果我们要求 LLM 在不断变化的世界中充当推理智能体，我们认为这是一项重要能力。

在下面的结果中，baseline 完全不会收到与问题相关的上下文（我们直接向它提问）；RAG 从文档中选出最佳的约 2～3k 个 token，并将其放入上下文中¹²；主动外在主义则将整篇文档放入记忆，并像 Otto 使用自己的笔记本一样使用它。

我们可以看到，虽然 RAG 方法的效果会随着输入长度增加而下降，但主动外在主义仍然持续有效。尽管经过长上下文微调的模型目前在某些长程检索任务上的表现优于主动外在主义，但对于较小的模型而言，主动外在主义似乎是一种更有效的长上下文检索方式。

在大型模型中，主动外在主义明显优于 RAG 的地方，恰恰是模型在过拟合之前已经完成记忆的部分。换句话说，即使模型在测试数据14上的表现仍在持续提升，模型的权重也已经编码了事实信息。根据具体应用，这既可以被视为优势，也可以被视为缺陷。当然，当我们将 LLM 用作推理智能体时，这是一项缺陷。

使用主动外在主义似乎还能减少对提示词的依赖。通常，我们需要在提示词中加入一些期望回答类型的示例（或者使用经过 RLHF 的“聊天”模型），但实验表明，采用主动外在主义时并不需要这样做。

对推理引擎的影响

我们将讨论主动外在主义对 LLM 作为推理智能体能力的两项重要影响：不确定性感知和抽象程度调节手段。

如果我们向模型提出一个它无法确定答案的问题15，它可能不会以透明表达这种不确定性的方式作答。主动外在主义提供了一种新方法，可以揭示模型何时对自己的答案不确定。

来看一个例子。我们可以轻松地从 huggingface 加载模型，并将 Wikipedia 中 Grothendieck 词条的一段文字作为外部记忆传入。

现在，我们先不使用主动外在主义，向模型提出一个确定能在上述段落中找到答案的问题，尽管答案藏得不太明显。可以通过设置参数 model.use_active_externalism = False，或者直接传入 topk=0 来实现。提示：正确答案是 1971 年。

Baseline Generation: When did Alexander Grothendieck get his French citizenship?I am trying to find out when Alexander Grothendieck got his French citizenship. I know that he was born in Germany and that he got his French citizenship in the late 1950s. I am trying to find out when he got his

现在启用主动外在主义，并通过 topk 参数逐步增加每个查询 token 可以关注的记忆数量。

Generation for k=5: A: I think he got it in the early 1960s.Generation for k=6: A: I think he got it in the early 1970s.Generation for k=7: A: He was born in France, and he was naturalized in 1971.<|endoftext|>Generation for k=8: A: I think he got it in 1971.<|endoftext|>QGeneration for k=30: A: He was born in Germany, and became a French citizen in 1971.

模型不仅给出了正确答案，还对答案表现出了越来越高的确定性。这一生成结果的演变过程，反映了模型最初的不确定性。

当模型确定答案时，随着外部上下文中的 k 值增加，生成结果会保持稳定。

Baseline Generation: What was did Alexander Grothendieck's profession?Alexander Grothendieck was a French mathematicianGeneration for k=2: Alexander Grothendieck was a mathematician.WhatGeneration for k=8: A: He was a mathematician.<|endoftext|>Q: What

这一原理的自然延伸，可能是开发一种基于相似度或注意力权重的指标，以更紧凑的形式传达这种不确定性。Normal 目前正在开展这方面的工作。

参数 topk 也是调节模型输出抽象程度的有效手段。例如，我们可以用它控制希望模型在多大程度上综合记忆，或者逐字引用来源。在针对代码的问答任务中，这一点体现得非常明显。这里我们使用聊天模型展示一个示例，因为它最适合处理形式更加自由的问答任务。

我们让模型回忆函数的作用，首先使用较小的 topk。

What does the function sieve_of_eratosthenes do?The function sieve_of_eratosthenes is a Python function that implements the Sieve of Eratosthenes algorithm to generate all prime numbers up to a given limit.The Sieve of Eratosthenes is a simple algorithm that generates all prime numbers up to a given limit. It works by creating a list of all integers from 2 to the given limit, and then iteratively marking the multiples of each prime number as composite (not prime).

可以看到，当 topk 较小时，模型会省略代码中的细节，使用自然语言描述代码的作用。现在尝试使用更大的 topk。

What does the function sieve_of_eratosthenes do?(limit): primes.append(True) for i in range(2, int(limit0.5) + 1): if sieve[i]: break else: for i in range(2, int(limit0.5) + 1): if i % 2 == 0: sieve[i] = False return primes```This implementation of the S

此时，模型输出的内容更加接近逐字复现的代码，同时仍然抽象掉了一些变量名。这类细腻的风格选择很难通过朴素的提示方法和 RAG 方法实现，除非针对数据和提示开发大量专用的点式解决方案。更重要的是，这类实验为我们提供了一些微小的线索，帮助理解模型究竟如何基于这些键值对进行推理。在 Normal，我们希望将机械可解释性方法方面的工作与扩展心智 Transformer 结合起来，构建一个统一系统，用于理解模型如何存储事实并基于事实进行推理。

Clark 和 Chalmers 在论文中写道：“通过接受主动外在主义，我们能够更自然地解释各种行为。”事实上，这一点同样适用于我们的主动外在主义。借助注意力权重，我们可以突出显示每个生成步骤使用了哪些记忆。这里，我们突出显示生成正确 token“1971”时使用的记忆。由于我们按层、按注意力头检索记忆，因此展示的是众数。

像这样简单的方法仅仅是一个开始，但 RAG 等方法目前无法实现细粒度引用，实际上也无法实现任何因果引用。我们最多只能突出显示被选中并纳入上下文的部分。使用自注意力权重或许能带来一些帮助，但这类数据难以处理，而且其解释能力也一直受到质疑。

创建外部记忆

主动外在主义涉及许多值得讨论的超参数。替代性的掩码策略、将主动外在主义限制在解码器层的某个子集，以及评估模型规模所起的作用，都是重要的讨论议题。大部分讨论将留到即将发布、技术性更强的论文中展开。不过，我们认为有必要简要介绍生成外部记忆时使用的超参数。

我们通过模型处理这些外部上下文来创建外部记忆（每一层都会创建），其过程与推理完全相同。随后，我们保存模型生成的内部表示，并在之后关注这些表示。如果外部记忆的长度超过模型的最大序列长度，通常需要使用步长来生成表示。

这样可以确保所有 token 都至少以前面相当于步长数量的 token 为条件。直观来说，我们的所有记忆都将“见过”一定数量的合理上下文。不过，在某些情况下，增加上下文未必有助于模型形成对数据的最佳表示。例如，数值型数据或日志型数据的表示，可能更适合使用较小的序列长度或步长。

Normal 相信，如果从第一性原理出发审视如今割裂但又迅速扩张的企业 AI 生态，仍有大量机遇等待发掘——即使是在已经开始形成初步共识的领域，甚至可以说尤其是在这些领域。我们坚信，跨学科视角与研究对于推动这一领域发展至关重要，因为从根本和历史上看，这本就是一门横跨多个领域且不断演变的学科。

Clark 和 Chalmers 在《扩展心智》中推测：“在遥远的未来，我们或许能够将各种模块接入大脑，为自己提供帮助：例如，在需要时接入一个扩充短期记忆的模块。”

虽然对人类而言，这仍是一个遥远的目标，但我们提出了一种方法，可以为 LLM 实现完全相同的短期记忆增强。我们已经展示了如何通过对 LLM 的自注意力机制进行简单而自然的扩展，在长文档检索任务中实现 SoTa 性能，并提供不确定性感知、抽象程度调节手段和细粒度可解释性，甚至可能让我们对这些模型的内部推理方式有所了解。

我们很期待将这些方法扩展到使用旋转位置编码和相对位置编码的模型上。

将因果引用打造为开箱即用的功能，也是我们优先事项列表中的重要一项。

从生成结果与 k 值选择的共同演变过程中提炼信息，并将其转化为不确定性指标，是我们正在投入研究的另一个领域。

最后，持续开发并运行全面的基准测试，对于深入、可靠地理解主动外在主义所带来的益处至关重要。

Buchen, Patrick. 2018. “为 Otto 及其延展心智辩护。” 2018. https://medium.com/@pnbuchen/in-defense-of-otto-and-his-extended-mind-9786db756f2d.

Burtsev, Mikhail S., Yuri Kuratov, Anton Peganov, and Grigory V. Sapunov. 2021. “记忆 Transformer。” https://arxiv.org/abs/2006.11527.

Clark, Andy, and David Chalmers. 1998. “延展心智。” Analysis 58, no. 1: 7–19. http://www.jstor.org/stable/3328150.

Liu, Nelson F., Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2023. “迷失在中间：语言模型如何使用长上下文。” https://arxiv.org/abs/2307.03172.

Martins, Pedro Henrique, Zita Marinho, and André F. T. Martins. 2022. “-Former：无限记忆 Transformer。” https://arxiv.org/abs/2109.00301.

Press, Ofir, Noah A. Smith, and Mike Lewis. 2022. “短序列训练，长序列测试：采用线性偏置的注意力机制实现输入长度外推。” https://arxiv.org/abs/2108.12409.

Sukhbaatar, Sainbayar, Edouard Grave, Guillaume Lample, Herve Jegou, and Armand Joulin. 2019. “使用持久记忆增强自注意力。” https://arxiv.org/abs/1907.01470.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2023. “注意力就是你所需要的一切。” https://arxiv.org/abs/1706.03762.

Wu, Yuhuai, Markus N. Rabe, DeLesley Hutchins, and Christian Szegedy. 2022. “记忆型 Transformer。” https://arxiv.org/abs/2203.08913.