新智元报道  

编辑：LRS

【新智元导读】挖掘大模型固有的长文本理解能力，InfLLM在没有引入额外训练的情况下，利用一个外部记忆模块存储超长上下文信息，实现了上下文长度的扩展。

大模型只能够记忆与理解有限的上下文已经成为大模型在真实应用中的能力瓶颈，例如对话式AI系统往往无法记忆你前一天对话的内容，利用大模型构建智能体会产生前后不一致的行为与记忆。

为了让大模型能够记忆并处理更长的上下文，来自清华大学、麻省理工学院和人民大学的研究人员联合提出无需额外训练的大模型长文本理解方法 InfLLM，利用少量计算和显存开销实现了 LLM的超长文本处理。

论文地址：https://arxiv.org/abs/2402.04617

代码仓库：https://github.com/thunlp/InfLLM

实验结果表明，InfLLM能够有效地扩展Mistral、LLaMA的上下文处理窗口，并在1024K上下文的海底捞针任务中实现100%召回。

研究背景

大规模预训练语言模型（LLMs）近几年在众多任务上取得了突破性的进展，成为众多应用的基础模型。

这些真实应用也给LLMs处理超长序列的能力提出了更高的要求，例如LLM驱动的智能体需要基于所有历史记忆持续处理从外部环境接收的信息，对话式AI需要更好地记忆与用户的对话内容从而产生个性化回答。

然而，现有的大模型往往仅在长度只有数千个Token的序列上进行预训练，将他们应用到超长文本中存在两大挑战：

1. 分布外长度：直接将LLMs应用到更长长度的文本中，往往需要LLMs处理超过训练范围的位置编码，从而造成Out-of-Distribution问题，无法泛化；

2. 注意力干扰：过长的上下文将使模型注意力被过度分散到无关的信息中，从而无法有效建模上下文中远距离语义依赖。

方法介绍

InfLLM示意图

为了高效地实现大模型的长度泛化能力，作者提出了一种无需训练的记忆增强方法，InfLLM，用于流式地处理超长序列。

InfLLM旨在激发LLMs的内在能力，以有限的计算成本捕获超长上下文中的长距离语义依赖关系，从而实现高效的长文本理解。

整体框架：考虑到长文本注意力的稀疏性，处理每个Token通常只需要其上下文的一小部分。

作者构建了一个外部记忆模块，用于存储超长上下文信息；采用滑动窗口机制，每个计算步骤，只有与当前Token距离相近的Tokens（Local Tokens）和外部记忆模块中的少量相关信息参与到注意力层的计算中，而忽略其他不相关的噪声。

因此，LLMs可以使用有限的窗口大小来理解整个长序列，并避免引入噪声。

然而，超长序列中的海量上下文对于记忆模块中有效的相关信息定位和记忆查找效率带来了重大挑战。

为了应对这些挑战，上下文记忆模块中每个记忆单元由一个语义块构成，一个语义块由连续的若干Token构成。

具体而言， (1)为了有效定位相关记忆单元，每个语义块的连贯语义比碎片化的Token更能有效满足相关信息查询的需求。

此外，作者从每个语义块中选择语义上最重要的Token，即接收到注意力分数最高的Token，作为语义块的表示，这种方法有助于避免在相关性计算中不重要Token的干扰。

(2)为了高效的内存查找，语义块级别的记忆单元避免了逐Token，逐注意力的相关性计算，降低了计算复杂性。

此外，语义块级别的记忆单元确保了连续的内存访问，并减少了内存加载成本。

得益于此，作者设计了一种针对上下文记忆模块的高效卸载机制（Offloading）。

考虑到大多数记忆单元的使用频率不高，InfLLM将所有记忆单元卸载到CPU内存上，并动态保留频繁使用的记忆单元放在GPU显存中，从而显著减少了显存使用量。

可以将InfLLM总结为：

1. 在滑动窗口的基础上，加入远距离的上下文记忆模块。

2. 将历史上下文切分成语义块，构成上下文记忆模块中的记忆单元。每个记忆单元通过其在之前注意力计算中的注意力分数确定代表性Token，作为记忆单元的表示。从而避免上下文中的噪音干扰，并降低记忆查询复杂度

实验分析

作者在 Mistral-7b-Inst-v0.2（32K） 和 Vicuna-7b-v1.5（4K）模型上应用 InfLLM，分别使用4K和2K的局部窗口大小。

与原始模型、位置编码内插、Infinite-LM以及StreamingLLM进行比较，在长文本数据 Infinite-Bench 和 Longbench 上取得了显著的效果提升。

超长文本实验

此外，作者继续探索了 InfLLM 在更长文本上的泛化能力，在 1024K 长度的「海底捞针」任务中仍能保持 100% 的召回率。

海底捞针实验结果

总结

在本文中，团队提出了 InfLLM，无需训练即可实现 LLM 的超长文本处理拓展，并可以捕捉到长距离的语义信息。

InfLLM 在滑动窗口的基础上，增加了包含长距离上下文信息的记忆模块，并使用缓存和offload 机制实现了少量计算和显存消耗的流式长文本推理。

参考资料：https://arxiv.org/abs/2402.04617