LM-Infinite: 一种简单有效的大模型即时长度泛化，解决更长文本推理问题

深度学习自然语言处理 原创 作者：qazw

摘要

近年来，LLM在各个领域的性能都有了显著的进展。随着这些LLM被用于越来越复杂的任务，它们经常需要更长的推理过程或理解更大的上下文。在这些情况下，LLM在长序列上的长度泛化失败问题变得更加突出。

大多数预训练方案会将训练序列截断到固定长度，但是即使使用了相对位置编码来应对这个问题，LLMs在更长的上下文之后仍然难以生成流畅的文本。

为了更有效地利用现有LLM的生成能力，我们在理论和实证上调查了导致此问题的OOD主要因素，提出了一种简单而有效的解决方案，用于即时长度泛化，即LM-Infinite。

该解决方案只涉及Λ形状的注意力掩码和距离限制，无需进行参数更新或学习。我们发现它适用于使用相对位置编码方法的各种LLM。

LM-Infinite具有高效的计算性能，时间和空间复杂度为O(n)，在ArXiv和OpenWebText2数据集上，在达到32k个标记的情况下，表现出持续的流畅性和生成质量，并具有2.72倍的解码加速。

LLM中OOD因素诊断

我们的主要灵感来自一个假设，即预训练的LLM中的相对位置编码已经具备处理相对位置的能力。然而，当应用于更长的序列时，内部特征（如注意力权重和隐藏状态）变得对LLM来说是“陌生的”，即超出了训练分布。通过消除这些因素，我们可以将内部特征转移回训练分布，这对LLM来说是“舒适区”，LLM将能够以其原始质量生成。在本节中，我们寻找这样的因素，主要是寻找可能超出分布的内部特征并验证其存在。

OOD 因素 1：看不见的距离

回想一下，在相对位置编码中，两个标记之间的注意权重取决于它们的距离。如果文本变得太长，一些距离将增加到一个未见过的大数，最终超过预训练中见过的距离。随着长度的增加，为了相对位置编码能够识别新的距离，注意力权重将不得不无限增大。

我们通过对LLaMA在最长的ArXiv数据集上进行实证验证，将其截断为32k个标记。为了清晰可视化，我们选择每个Transformer层中的第0个注意头，并在图 1(a)中绘制自32k个标记到所有前一个标记的注意权重。可以看到，在较远的距离上，注意权重的绝对值振荡到比训练长度为4k时更大的值。

我们得出的结论是，相对位置编码要么无法识别未见过的距离，要么逻辑回归将增加到无穷大。后一种情况会产生ODD（Out-of-Distribution）的逻辑回归值，这对LLM来说是“陌生”的，并且可能导致不规则的结果。

OOD 因素 2: 看不见的数字 or tokens

另一个我们注意到的潜在影响因素是token数量。当文本变得更长时，将需要注意更多的标记，这可能会稀释注意力权重，使得注意分布更加扁平化，导致在注意力中丢失信息。在这里，我们研究注意力的熵，这是衡量分布信息量的理论度量指标。

这一发现提示我们限制token数量，以便LLM可以在一个熟悉的子特征空间上运行。经过对这两个因素的分析，人们可能会倾向于提出一个简单的解决方案：强制每个token只关注最近的几个token，在注意力机制中忽略更远的token。然而，我们发现这种方法并不奏效，LLM在短文本上的性能实际上会下降。这意味着XPos的外推能力在很大程度上依赖于对其进行明确的训练，并且不能直接适用于其他LLM，这一现象表明存在另一个OOD因素。

OOD 因素 3: 隐式编码的绝对位置

在本节中，我们将展示一种反直觉的现象，即使绝对位置信息没有明确地编码在计算图中，注意力机制仍然能够隐式地编码它。我们猜测这实际上就是在相对位置编码中的Transformer中发生的情况，以上定理证明了这个事实。

在以上定力所提供的构造中，初始标记的信号强度比尾随标记更强，更容易区分。如果这是真的，那么它构成了另一个OOD因素。当长度较短时，LLM隐含地编码了初始标记的位置信息。然而当长度超过训练语料库时，由于OOD FACTOR1和2，初始标记被处理不当，它们的绝对位置信息会变形或丢失。这个定理是存在性的：它仅证明了这种隐含的绝对位置编码是可能的，但不能保证它确实在实际的LLM中发生了。作为实证验证，采用LLaMA的第一层输出的隐藏状态，并在图中绘制其PCA投影到一个二维平面上，图中的点对应于序列中的前4096个标记，蓝色点对应初始标记，红色点对应尾随标记。从图中可以看出，处于不同位置的标记在特征空间中占据不同的子空间，即使没有显式编码绝对位置信息。这解释了为什么之前提到的简单解决方案会失败：当序列变长时，直接限制注意力窗口将消除初始标记，从而使特征子空间变得不可见。我们推测，保留这些起始的几个标记对于LLM的正常运行很重要。

在确定了这些OOD因素之后，我们已经找到了长度泛化问题背后的缺失要素。接下来，我们提出了我们的解决方案LM-Infinite，并绘制了一个描述相对位置编码工作原理的概念模型。

LM-Infinite

一般原则

基于前文的分析，我们提出了解决方案LM-Infinite，这是一种简单的即时技术，用于在LLM中进行长度泛化，并使用相对位置编码。需要注意的是，这些编码有不同的实现方式，因此LM-Infinite提供了一组不仅局限于单个LLM的高级原则。

前图(a)展示了LM-Infinite的概述。该解决方案由两个组成部分组成：一个Λ-形状的注意力掩码和一个距离限制。Λ-形状的注意力掩码有两个分支：左侧的全局分支和右侧的局部分支。全局分支允许每个标记在当前标记之前出现时关注最初的global个标记，局部分支允许每个标记在local范围内关注之前的标记。在注意力过程中忽略这两个分支之外的任何其他标记。这个设计基于上述OOD FACTOR 2和3，旨在控制要关注的标记数量，同时确保包含初始标记。

距离限制涉及限制“有效距离”d在Lpretrain范围内。这只影响位于全局分支的标记。具体而言，在相对位置编码中，注意力权重原本的形式为w(q,k,d)，其中d是两个标记之间的距离。现在我们将其修改为w(q,k,min(d,Lpretrain))。这个设计是基于OOD FACTOR 1的动机，确保LLM不会接触到在预训练期间没有见过的距离。

实现细节

LM-Infinite中的原则适用于大多数相对位置编码。由于本文的重点是解决LLM的长度泛化问题，我们将评估LM-Infinite的性能，评估对象包括三个最先进的开源LLM系列：LLaMA系列 （LLaMA和Llama-2）、 MPT-7B系列和GPT-J系列。LLaMA和GPT-J都使用RoPE编码，而MPT-7B使用Alibi编码。这些原则可以很容易地推广到其他相对位置编码。

相对位置注意力的概念模型

我们基于OOD FACTOR诊断和LM-Infinite的设计原则，描述了相对位置编码功能的概念模型。该图展示了在生成下一个标记时的观点，即每个标记都关注所有前面的标记。最后一个标记通过自注意力机制收集信息，计算输出特征并预测下一个标记。在这个解释中，一个较长的上下文可以粗略地分成3个部分：

1. 初始标记主要编码其绝对位置信息，它们是注意力层的重要组成部分，因为它们的特征占据特征空间中的一个区域。如果该区域缺失或者使用了一个未见过的大距离进行关注，则会产生异常FACTOR 3。
2. 最接近最后一个标记的尾标记。这里相对位置更为重要, 尾标记对于注意力层的正确功能至关重要。
3. 中间标记编码的信息与位置关联性较小。这一区域的注意力权重要么是爆炸性的，要么是注意力熵过高（OOD FACTOR1和2）的。因此，对于长度泛化而言，它带来的坏处多于好处，所以在长度超过训练数据的序列中，在LM-Infinite中将其移除。

评估

在本节中，我们对LM-Infinite的性能进行了实证评估。我们从Pile 数据集中选择了ArXiv和OpenWebText2语料库，该数据集包括了来自 ArXiv的预印论文和从2005年到 2020 年 4 月的 Reddit 提交。我们使用了 LLaMA-7B、Llama-2-7b-Chat、MPT-7B和GPT-J-6B 作为评估的 LLMs。其中一个基准是经过长序列微调的 MPT-7B-Storywriter。

流畅度

我们首先使用广泛采用的困惑度度量标准来评估LM-Infinite的流畅度。在正式评估概率模型M对分布D进行建模时，困惑度被定义为平均负对数似然（NLL）的指数：PPL(D,M)=exp−Ex∈D ln M(x)。在ArXiv数据集上绘制了NLL曲线，如图3所示。注意，由于Llama-2在过长的序列上输出NaN，所以曲线相对较短。可以看到，LM-Infinite成功地将困惑度曲线扁平化到比其训练输入长度长得多的程度。这表明在长序列上具有连贯而无损的流畅性。

我们还在ArXiv和OpenWebText2数据集上在一些里程碑长度（2k、4k、8k、16k和32k）上对困惑度分数进行了数值记录，如表1所示。OpenWebText2数据集中的数据在长度32k以上非常少，因此我们省略了该列。注意，借助LM-Infinite的帮助，语言模型成功地实现了长度的泛化，并在9个列中的7个上取得了SoTA分数。这是一个令人鼓舞的结果，考虑到LM-Infinite与众多强基线相比不需要任何参数更新。作为直接对比，MPT-7B+LM-Infinite的得分仅略低于其微调版本MPT-7B-Storywriter。这表明LM-Infinite是资源消耗大的精细调整的高 效替代品。

模型性能

由于困惑度是LLM的内部度量指标，我们采用BLEU和ROUGE指标评估LM-Infinite 在ArXiv和OpenWebText2上生成质量上的优势。简单来说，这两个指标评估生成文本与参考文本之间n-gram的重叠情况，其 中BLEU侧重于精确度，ROUGE侧重于召回率。

我们让LLM在每个长度之后生成100个token，并将原始文本中的接下来的100个token作为参考。由于生成需要很长时间，我们在每个数据集中随机选择100个长序列进行评估，结果列在表2中。与上一节的趋势类似，LM-Infinite成功地使LLM能够将其生成质量延伸到比训练更长的长度，与微调的效果相当，但无需参数更新。需要注意的是，LM-Infinite在不同的LLM上会有稍微不同的效果。例如，在LLaMA和GPT-J-6B上，质量在更长的位置上得到了更好的保持，而在Llama-2上，质量在更接近的位置上更好。

我们还在附录 D中评估了长度为32k时的计算效率，其中在编码 时，LM-Infinite表现出3.16倍的加速，在解码时表现出2.72倍的加速。

任务解决

最后，由于LLM经常用于解决下游任务，我们评估了LM-Infinite在长输入任务上的性能。我们遵循Mohtashami&Jaggi (2023) 的方法，并使用了一个token检索任务。该任务在长的干扰文本中随机隐藏了一个口令，并最终询问该口令是什么。作为一个合成数据集，可以更灵活地控制输入长度以进行细粒度分析。我们在图4中绘制了答案的正确率。可以看到，与无论长度如何都会立即失败的普通模型相比，LM-Infinite可以使LLM在训练长度之外的较长长度上保持较慢的准确率下降。由于token检索是一项信息敏感的任务，曲线显示出，即使LM-Infinite在理论上可以从长度为nlayerLpretrain的上下文中获取信息，但它并不能完美地保持信息感知能力。将此归因于Λ-mask，它在信息和流畅性之间进行权衡。将更高的信息敏感度保持问题留给了未来的工作。

总结

在本文中，我们针对使用相对位置编码的基于Transformer的LLM中的长度泛化问题提供了一个解释和一个简单的即时解决方案。我们通过对可能导致长度泛化失败的OOD FACTOR进行理论和实证分析。基于这些直觉，我们提出了LM-Infinite，这是一个无需任何参数更新的即插即用解决方案。实证评估显示，我们可以让多个开源的SoTA LLM保持其原始的生成质量，类似于显式微调后的性能。LM-Infinite还将任务解决能力扩展到比训练样本更长的序列。未来的工作可以探索如何让LM-Infinite更好地感知遮蔽关注区域中的信息。我们希望LM-Infinite的计算效率和易用性能够使没有极大计算资源的研究人员也可以在长序列上使用LLM。