田渊栋团队新作：通过位置插值来扩展大语言模型的上下文窗口

科技

2023-07-04 04:07

©PaperWeekly 原创 · 作者 | 陈首元、王胜铭等

单位 | Meta Platform Inc

论文标题：

Extending Context Window of Large Language Models via Positional Interpolation

论文作者：

陈首元, 王胜铭, 陈梁坚, 田渊栋

作者单位：

Meta Platform Inc

论文链接：

https://arxiv.org/abs/2306.15595

方法动机

大型语言模型（LLM）通常具有预定义的上下文窗口大小。例如，LLaMA 的输入必须少于 2048 个 Token。在进行长时间对话、总结长文档或执行长期规划等应用中，经常超过这个预设的上下文窗口限制。对于这些应用，拥有更长上下文窗口的 LLM 更受青睐。然而，从头开始训练具有较长上下文窗口的 LLM 需要巨大的投入。这自然引出一个问题：我们能否扩展现有预训练 LLM 的上下文窗口？

一种直接的方法是对现有预训练 Transformer 进行微调，使其具有更长的上下文窗口。然而，根据经验，我们发现通过这种方式训练的模型对于较长的上下文窗口适应速度很慢。在训练了 10000 个批次以上后，有效上下文窗口的增加非常小，这表明这种方法不适用于显著延长上下文窗口。

尽管一些技术（如 ALiBi 和 LeX）可以实现 Transformer 的长度外推，即在短上下文窗口上进行训练，然后在更长的上下文窗口上进行推理，但许多现有的预训练 LLM，包括 LLaMA（Touvron 等人，2023），使用具有较弱外推性质的位置编码（例如 RoPE）。因此，这些技术在扩展这些 LLM 的上下文窗口大小方面的适用性仍然有限。

在这项工作中，我们引入了位置插值技术，以实现对某些现有预训练 LLM（包括 LLaMA）的上下文窗口扩展。关键思想是，我们不进行外推，而是直接降低位置索引，使最大位置索引与预训练阶段的先前上下文窗口限制相匹配。见图 1 进行说明。

▲ 图1

换句话说，为了容纳更多的输入标记，我们在相邻整数位置上插值位置编码，利用位置编码可以应用于非整数位置的事实，而不是在训练位置之外进行外推，这可能导致灾难性的值。我们在理论上验证了我们的方法，通过展示插值的注意力得分具有更小的上界。

理论方法

尽管 RoPE 中的注意力得分仅取决于相对位置，这正是我们所希望的，但其外推性能并不出色。特别是，当直接扩展到训练中未见过的更大上下文窗口时，困惑度 (perpelexity) 可能会飙升到非常高的数值（即大于），与未经训练的模型相当。

理想情况下，我们希望在大小为 L = 2048 的上下文窗口上训练的模型在更长的上下文窗口上仍能表现得相当不错。例如，为了回答位于 3000 位置的问题，训练在最大窗口大小 L = 2048 上的模型无法利用位于位置 0 处的信息，但仍可以利用位于位置 2900 处的证据。然而，在现实中，我们观察到灾难性的行为，即使信息位于位置 2900，位于位置 3000 的问题也无法被正确回答。

这是什么原因？根据 (Su et al., 2021) 第 3.4.3 节的说明，如果注意力得分随着相对距离的增加而衰减，为什么会发生这种情况呢? 原来，在（Su 等人，2021 年）的第 3.4.3 节中得出的上界可能过于宽松：尽管它确实随着的增大而减小，但上界仍然可能相当大（即上界可能严重依赖于的大小），因此无用。事实上，如果我们将所有三角函数视为基函数（即，并将方程 2 看作以下基函数展开：

感谢基础函数的平滑性，插值更加稳定，不会导致异常值。因此，我们可以将注意力得分外推范围扩展到 s > L 的情况下，定义一个注意力得分，其中是更长的上下文窗口。形式上，我们将 RoPE f 替换为以下定义的 f'：

我们将这种对位置编码的转换称为位置插值。在这一步中，我们将位置索引从减小到，以匹配计算 RoPE 之前的原始索引范围。因此，作为 RoPE 的输入，任意两个标记之间的最大相对距离从减小到。由于我们在扩展之前和之后对位置索引和相对距离的范围进行了对齐，减轻了上下文窗口扩展对注意力得分计算的影响，这使得模型更容易适应。为了进一步证明这一点，下面的定理表明插值后的注意力得分具有良好的性质：

根据实证结果，我们发现位置插值非常有效和高效，仅需要非常短的微调期间，模型就能完全适应大幅扩展的上下文窗口。我们通过使用位置插值，在 7B 到 65B 的 LLaMA 模型中将上下文窗口从初始的 2048 扩展到 32768，并呈现了实验结果。我们的结果表明：

位置插值可以轻松实现非常长的上下文窗口（例如 32768），只需在 Pile 上进行 1000 步的微调即可获得良好的质量。与预训练成本相比，微调的成本微不足道。这验证了我们的假设，即模型相对容易适应插值位置编码。位置插值生成强大的模型，能够有效利用更大范围的上下文窗口。

我们展示了通过位置插值扩展的模型在文本建模中从大幅扩展的上下文窗口中获得显著的困惑度提升，并且我们展示了随着上下文窗口的扩大，困惑度逐渐降低。我们还在长文本摘要任务中应用了位置插值，并展示了竞争性的性能。位置插值相对保持了原始上下文窗口大小的模型质量。

我们对扩展的 LLaMA 模型在原始 LLaMA 基准测试中进行了各种评估结果。与原始 LLaMA 模型相比，在 2048 标记限制内的几个标准基准测试中，扩展的 LLaMA 模型的性能略有下降。

实验结果

我们在表 1 和表中 2 中报告了我们模型和基线模型的困惑度结果。从结果中，我们发现通过我们的方法扩展的模型在更长的上下文窗口大小下享受到了显著改善的困惑度 (perplexity)。通过将上下文窗口大小从 2048 增加到 16384，我们观察到在两个数据集上扩展 LLaMA 7B 模型的困惑度降低了 0.28 和 0.5，扩展 LLaMA 13B 模型的困惑度降低了 0.27 和 0.48，以及扩展 LLaMA 33B 模型的困惑度降低了 0.14 和 0.42。对于 LLaMA 65B 模型，我们观察到通过扩展到 8192 的上下文窗口大小，困惑度降低了 0.12 和 0.3。

总体而言，我们观察到我们的模型在更长的上下文窗口下能够取得更好的困惑度。这表明我们的模型能够有效利用更长的上下文窗口来更好地预测语言建模任务中的下一个标记。此外，我们发现这种趋势在 LLaMA 7B 和 13B 模型的 PG19 数据集上延伸到了 32768 的窗口大小，而不会降低。这表明我们的方法可能使得扩展到更长的上下文窗口成为可能。

相比之下，我们观察到通过直接微调方法扩展的模型在更长的上下文窗口上显示出了困惑度的回归（最高达到+0.48）或轻微改善（最高达到-0.12）。这表明通过这种方式扩展的模型在超过它们预训练设置的上下文窗口长度时能力有限。

在某些情况下，我们在扩展模型的原始上下文窗口 2048 上观察到困惑度的轻微降低。例如，在 Proof-pile 数据集上，我们发现使用位置插值扩展的所有模型的性能下降在 0.01 到 0.05 之间。原始评估上下文窗口内的性能轻微下降是可以预期的，因为位置插值将原始上下文窗口中的位置编码限制在一个更窄的区域，可能对语言模型的性能产生负面影响。