摘要

近年来，大型语言模型（LLMs）在各个领域取得了显著的成功。然而，由于时效性、Attention 机制内存消耗等问题，处理长上下文仍然是 LLMs 的一个重大挑战。本文介绍了 MemLong：用于长文本生成的记忆增强检索方法，是一种通过利用外部检索器进行历史信息检索来增强长上下文语言建模能力的方法。MemLong 结合了 ret-mem 模块和一个部分可训练的仅解码语言模型，并引入了一种细粒度、可控的检索注意力机制，该机制利用语义级别相关的块。在多个长上下文语言建模基准上的全面评估表明，MemLong 始终优于其他最先进的 LLMs。更重要的是，MemLong 可以在单个 3090 GPU 上将上下文长度从 4k 扩展到 80k。

介绍

大型语言模型在各个领域取得了显著的成功。然而，由于原始注意力机制的时间和空间复杂度为二次方，扩展上下文长度具有挑战性，这对涉及长序列任务的应用（如长文档摘要和多轮对话）构成了重大限制。因此，LLMs 通常被期望保持长时间工作的能力（即长上下文 LLMs）以有效处理这些苛刻的场景。一些工作尝试采用稀疏注意力操作来减少原始注意力机制的计算量，但通常会在模型能力上做出权衡。另一类工作转向了检索增强生成。这些工作通常引入检索机制以增强模型处理长文本的能力。然而，这类方法存在几个缺点。

• 由于训练期间模型参数的变化，存储在记忆中的信息可能会经历分布偏移。
• 这些方法通常需要重新训练，这在大模型时代是不切实际的。
• 这些模型往往在处理长文本输入时以牺牲预训练模型的原始能力为代价。

为了应对之前研究的局限性，本文思考此问题：我们能否利用检索器的显式检索能力来近似模型内部的隐式检索过程？

本文提出了 MemLong，一种高效且轻量的方法来扩展 LLMs 的上下文窗口。关键思想是将过去的上下文和知识存储在一个不可训练的记忆库中，并进一步利用这些存储的 embedding 来检索块级别的键值对以输入到模型中。MemLong 适用于任何仅解码预训练语言模型。

对比 RAG，MemLong 的记忆和检索过程如上图所示。在生成过程中，超过模型最大处理长度的文本被存储为记忆库中的上下文信息。随后，给定长文档中最近生成的文本块，本文使用检索器显式检索过去的信息，通过索引对齐获得额外的上下文信息。

MemLong 提供了几个好处：

• 分布一致性：与之前的模型在信息存储在记忆中时经历分布偏移不同，MemLong 确保缓存信息的分布保持一致。
• 训练高效：本文冻结了模型的下层，只需微调上层，大大降低了计算成本。
• 广泛的上下文窗口：由于只需记忆单层的 K-V 对，MemLong 能够在单个 3090 GPU 上轻松扩展上下文窗口到 80k tokens。

实验表明，MemLong 在多个方面表现优于其他领先的 LLMs。MemLong 在多个长上下文语言建模数据集上优于 OpenLLaMA 和其他基于检索的模型。在检索增强的上下文学习任务中，MemLong 相对于 OpenLLaMA 的性能提升高达 10.2 个百分点。

任务定义

语言模型旨在定义 token 序列上的概率分布，有效地预测给定语言中序列的可能性。给定这样一个序列 ，建模其概率 的标准方法是通过下一个 token 的预测：$p(x_1^*, \ldots, x_n^*) = \prod_{i=0}^{n} p_{\theta}(x_i | x_{<i})$，其中 $x_{<i}="" :="x_1^*," \ldots,="" x_{i-1}^*$="" 是="" $x_i$="" 之前的="" token="" 序列。不同于标准的语言建模目标，本文不仅使用当前上下文进行下一个="" 的预测，还利用外部检索获取相关信息，并在模型的上层进行知识融合。具体来说，给定一个包含="" $l$="" 个="" 的序列和每个="" chunk="" 的大小="" $\tau$，本文将其划分为一个长度为="" $\nu="\frac{l}{\tau}$" 的不重叠="" 序列，记为="" $c="(c_1^*," c_{\nu}^*)$。相应地，其文本形式被划分为="" $\nu$="" 个文本="" chunk，记为="" $t="(t_1^*," t_{\nu}^*)$。在每一步中，本文在下层对="" $c_i$="" 进行因果语言建模，而在上层，对="" $t_i$="" 进行细粒度的可控检索，以融合额外的信息。完成此操作后，本文的语言建模目标变为<="" p="">$$p(x_{1}, \ldots, x_{n}) = \sum_{i=0}^{n} p_{\theta}(x_{i} | \mathcal{R}(t_{i}), x_{<i}) \tag{1}="" $$="" 表示检索到的包含 的相邻 chunk。

如上图所示，Ret-Mem模块包括一个Retriever和一个Memory组件用于信息交换。首先，定义 Memory 组件为 M，Retriever 为 R，它们对应的操作分别为 和 。此外，本文指定模型的维度为 ，Retriever 的维度为 。Memory 模块包括两个部分：K-V 对和对应的表示嵌入。键和值的维度表示为 ，嵌入的维度表示为 。需要强调的是，实际的检索过程涉及表示块的嵌入，而不是 K-V 对。Retriever 本质上是一个预训练的密集嵌入器，具有出色的表示能力。MemLong 使用它将每个块编码为表示嵌入。由于它为一个块生成一个一维表示向量，即使内存大小很大，内存占用仍然很小。

每一步都涉及一个块 的输入，其中该块的原始文本为 。在模型被冻结的下层中，标准的 causal attention 应用于整个 。本文称下层的最后一层为记忆层。在每次遍历记忆层之后，会执行两个关键操作。第一个操作是检索，由红线表示，其中 被用来获取最相关的 K-V 对。第二个操作由蓝线表示，涉及将获取的 K-V 对及其相关的块表示进行缓存。在模型的上层中，检索到的 K-V 对与当前输入上下文相结合，随后调整模型参数以校准检索参考。

检索与动态内存管理

检索过程：由于目标是用显式检索替代基于 K-V 对的传统 kNN 检索，本文旨在每次模型输入之前尽可能预取所需信息。具体来说，对于每个潜在的查询块 及其对应的文本块 ，首先通过 Retriever 处理，然后获得表示 embedding ，其中 。随后，使用此表示 embedding 对 M 中的 embedding 进行检索，以获得所需的 k 个块级索引。本文计算检索表示 与存储在内存 M 中的 embedding 之间的余弦相似度。最后，本文得到 的前 k 个索引 ，其中 。由于块内的连续性，本文可以轻松扩展获得的索引以覆盖整个相关范围进行检索。最后，本文根据这些索引从内存中检索相应的 K-V 对 ，并用于 upper layer。值得注意的是，本文为 Memory 配备了计数机制，以记录其中每个索引的检索频率。此频率数据将作为动态内存更新的基础，允许优先处理更频繁检索的信息。

memory 过程：memory 过程同步存储来自内存层的 K-V 对和之前计算的用于检索的 representation embedding，确保 K-V 对的索引准确对应其表示 embedding。

对于每个可能的块内存 及其对应的文本块 ，本文将内存过程分为两部分：第一部分详细说明如何缓存 K-V 对，第二部分解释如何存储相应的表示。首先，本文将 输入到 MemLong 并从 memory layer 获取输出。值得注意的是，由于训练期间冻结了较低层，本文可以确保输出 K-V 对的分布一致。这种一致性对于避免之前在模型如 MemTrm 中观察到的分布偏移问题至关重要。本文的内存操作非常高效，因为它仅涉及存储检索所需的表示 ，从而避免了冗余。在所有块对的检索完成后，内存操作——记为 ——同步更新内存中的 Key-Value 对及其对应的表示。

▶动态内存更新

当内存溢出时，本文使用 Counter 智能更新内存。在实验中，本文保留最新的 10%内存内容，因为它们可能相关，丢弃最旧的 10%内容，因为它们可能过时，并根据检索频率优先处理中间 80%的内容，删除访问最少的条目，直到内存使用量降至 50%。这种选择性修剪平衡了新近性和相关性，保留有价值的信息并删除不太相关的数据。与传统的 FIFO 策略不同，本文的方法专注于检索频率，以高效修剪冗余信息，保持高质量的数据集。动态更新数据存储的决策是在有效性和效率之间的权衡。对于需要长期依赖的任务，存储所有信息可以增强全面处理，但对于短期任务，动态更新更为合适。动态更新控制内存大小以防止内存溢出问题，丢弃陈旧信息并减少检索开销，确保效率而不显著影响性能。

▶Attention Reformulation

在模型的可训练上层中，本文修正了注意力机制以融合长期记忆。如图 3 所示，不同于传统的使用多头注意力机制的 Transformer 解码层，本文提出了一种Retrieval Causal Attention，将其扩展为联合注意力机制，并提出了一种长期记忆融合过程，使每个 token 能够同时关注局部上下文和具有完整连续语义的块级过去上下文。通过从前一层输出的头部隐藏状态 和相应的检索键值对 ，下一层的输出隐藏状态 计算如下：

为了避免在训练初期由检索注意力得分 引起的干扰，本文采用了一种多头注意力机制，遵循 LLaMA-adapter 的方法：

最终，本文将 和 拼接以获得 ：

▶使用 MemLong 推理

当 MemLong 接收到超过长度的输入时，本文将其视为两个部分：前缀和主体。本文将分别描述在推理阶段长输入的编码和长输出的生成。当 MemLong 接收到长输入时，它首先将前缀划分为多个不重叠的块，并从其记忆层中计算，这确保了参与注意力的 tokens 数量等于块大小，这远小于输入的长度。需要注意的是，每个块是相互关联的（例如，第 t 个块需要处理前面 t-1 个块的内容）。

第二步是基于块级检索表示选择与主体最相关的 k 个块，并获取它们的 key 和 value 表示。之后，对于上层检索层，检索的注意力窗口相当于 ，这也小于输入长度。最后，长度受限的因果注意力和检索注意力都能高效地执行。

实验分析

本文在多项需要内存长上下文处理的任务上评估了所提出的 MemLong 模型：

• 长上下文语言建模和检索增强语言建模
• 可扩展的上下文学习，能够在内存中处理大量示例。

由于 K-V 缓存提供了显著的背景和上下文信息，MemLong 可以快速检索相关的 K-V 缓存并充分利用，从而增强模型在长上下文建模任务中的表现。

数据集。本文在四个广泛的文本基准数据集上评估了模型：英文书籍 PG-19 和 BookCorpus，维基百科文章 Wikitext-103，以及数学论文 Proof-Pile。实验结果表明，在所有数据集上困惑度显著改善。本文的模型在长度从 1024 到 32768 个 token 的范围内进行了测试。通过利用外部检索器和内存，本文的模型在所有数据集上都展示了显著的性能提升，同时内存开销最小。

设置。按照实验设置，本文计算了每个序列最后 2048 个 token 的困惑度。此实验设置旨在验证不同检索器大小对模型整体性能的影响。为了实现高效的细粒度检索，本文使用faiss工具包在 GPU 上构建一个精确搜索索引，以存储文本块的Representation Embeddings并进行高效检索。对于 MemLong，本文将 token 分割并放入 finetune-length = 1024 的 M 中用于进一步检索。

首先，本文在长上下文语言建模基准上评估了 MemLong，以评估其基本的语言建模能力。结果如上表所示。本文采用困惑度(PPL)作为语言模型的评估指标。较低的 PPL 表示更强的语言建模能力。

位置编码模型具有很强的泛化能力。然而，这些方法的性能只能保证在长距离生成时不退化。相比于这些方法，MemLong利用外部检索器处理更长的输入 token，并实现更好的困惑度改进。同时，由于高效的存储效率，MemLong 可以有效控制 GPU 的使用，避免 OOM 问题。

传统的上下文学习将少量非参数化的示例与查询一起输入模型。然而，这些方法通常受到模型输入长度的限制。在本实验中，由于 MemLong 可以在其内存中以参数化形式存储示例，本文主要研究 MemLong 是否可以有效利用存储在其内存中的知识来增强其新兴能力。与仅依赖非参数化知识的 OpenLLaMA 相比，在给定相同数量的上下文示例的情况下，MemLong 可以利用存储在其内存中的额外示例。随着内存中示例数量的增加，性能进一步提高或保持一致。在与 LongLLaMA 的对比分析中，观察到在相同的内存示例保留条件下，本文模型在大多数数据集上表现优于 LongLLaMA。需要强调的是，与 LongLLaMA 相比，本文模型的训练参数显著减少（200M 对比 0.3B）以及微调数据量（0.5B 对比 5B）。这突显了本文模型在利用外部检索器进行信息获取方面的效率，展示了在显著减少资源的情况下，合成和利用知识的卓越能力。

Conclusion

本文介绍了 MemLong，这是一种创新的方法，通过利用外部检索器显著增强了语言模型处理长文本的能力。MemLong 使用一个高效的检索器，以极小的内存开销快速且准确地访问与远距离上下文相关的文本。MemLong 成功地将模型的上下文窗口从 2k 扩展到 80k tokens。

本文展示了 MemLong 在长距离文本建模和理解任务中表现出显著的竞争优势。与全上下文模型相比，MemLong 的性能可以提高多达 10.4 个百分点。