基于RoPE外推的规模化定律

ICLR‘24来自上海AI实验室和复旦大学的论文“Scaling Laws Of Rope-Based Extrapolation”。

基于旋转位置嵌入（RoPE）的大语言模型 (LLM) (Su et al.，2021) 外推能力是目前备受关注的话题。解决使用 LLM 外推问题的主流方法是修改 RoPE，将原始 RoPE 中 θn = 10000^（−2n/d） 的旋转基数 10000 替换为更大的值，并提供更长的微调文本。这项工作用较小或较大的预训练上下文长度基对基于 RoPE 的 LLM 进行微调可以显著提高其外推性能。之后，提出基于 RoPE 外推的规模化定律，这是一个从周期性出发的统一框架，用于描述外推性能与基值以及上下文长度调整之间的关系。在此过程中，还通过外推的关键维度解释基于 RoPE 外推问题的起源。

观察

关于 RoPE 的外推问题，不同的研究提供了不同的解释和相应的解决尝试。这些研究可以分为两大流派。一派限制了自注意的范围 (Ratner et al., 2022; Han et al., 2023)，因为基于 RoPE 的自注意无法在训练上下文之外保持稳定，并且表现出注意得分爆炸以及单调熵增加 (Chen et al., 2023; Han et al., 2023)。另一派旨在捕捉具有较小旋转角和较长微调上下文的较长上下文 (Chen et al., 2023; Peng et al., 2023)。目前流行的方法，如动态 NTK (bloc97, 2023a) 和 Code LLaMA (Rozie`re et al., 2023)，主要来自第二种方法。这两种方法都将 RoPE 适应具有较大旋转基数的较长上下文。具体来说，动态 NTK（bloc97，2023a）用随推理长度增加的系数来调整基数，从而使基于 RoPE 的 LLM 能够同时适应更长的上下文，而Code LLaMA（Rozie`re，2023）直接将基数设置为 1000000，并在 16K 上下文长度上进一步训练，从而产生超过 100K 的上下文。

虽然最近的研究显示出不错的结果，但它们主要是特定基数值和上下文长度调整。这在理解基数值、调整长度和外推性能之间的关系方面留下了明显的空白。例如，虽然较大的基数可以改善 LLaMA2 等模型的外推能力（Touvron，2023a），但令人惊讶的是，在原始训练长度下使用较小基数进行微调也有利于 LLaMA2 的外推能力，这也在如图得到了证明：使用不同的外推方法，包括动态 NTK在 Books3（Presser，2020）上对 Perplexity 进行了测试；在原始训练长度 4K 或更长的上下文 16K 上使用较小或较大基数进行微调的 RoPE 可以胜过其他外推策略，并外推到 100K 上下文长度。

首先用更大的基数进行外推实验。如图所示，显然使用更大的基数进行调整可以显著提高 RoPE 的外推性能。

首先，更大的基数允许 LLaMA 2 (Touvron，2023b) 在其训练上下文长度之外进行外推，这与 （Rozie`re，2023) 的发现一致。

其次，使用更大基数的外推有一个明显的极限，即语言建模困惑度保持一致。超过这个极限，外推性能会显著下降。

此外，随着基数值的增加，LLaMA2 可以外推到更长的上下文。

最后，与动态 NTK (bloc97, 2023a) 相比，使用更大基数调整的 RoPE 在其外推上限之外的退化速度要快得多。

因此，对于使用较大基数进行微调，超出上限的性能可能会被动态 NTK 持续超越。尽管如此，在上限内，这种方法仍然比动态 NTK 有相当大的优势，当基数超过 600000 时，超过 100K 的上下文只需要 4K 的调整长度。

然后，用与较大基数相同的设置，对较小的基数进行外推实验。有趣的是，尽管这与常见的研究结果（bloc97，2023a；Rozie`re，2023）相悖，但在原始上下文长度上使用较小的基数对 RoPE 进行微调仍然可以增强外推，如图所示。它还将上下文窗口扩展到训练长度之外。

首先，具有较小基数的 RoPE 没有明显的外推上限。虽然困惑度会随着上下文长度的增加而恶化，但基数越小，这种下降就越温和。其次，具有较小基数的 RoPE 的外推增强并不均匀。在 10000 到 8000 的基数之间，外推性能表现出微小的改善。然后在 8000 和 2608 的基数之间，改进是适度的。之后，从 2608 到 1304 再到 652，改进变得更加明显。最后，当旋转基数为 500 时，外推曲线变得足够平滑，从而带来 48K 上下文长度上的强外推，同时优于动态 NTK（bloc97, 2023a）的性能。

结合这两个结果，观察到基于 RoPE外推中一个令人惊讶的奇怪现象，如图所示。具体而言，基数 10000 在微调时产生最差的外推性能，从而回答了 问题 1 。随着基数的减少或增加，性能显着提高。有趣的是，两个方向的改进不同。对于较大的基数，虽然性能稳步提高，但存在明显的外推上限。相反，对于较小的基数，虽然改进并不均匀，但得到的外推曲线没有明显的断点。在图中的热图中，较大基具有清晰连续的边界，较小基具有明显的过渡阶段。

定理 1 .（ 较小基数的规模化定律 ）对于用上下文长度 Ttrain 预训练的基于 RoPE LLM，如果将基数调整为 β < 10000 并仍然使用上下文长度 Ttrain 进行微调，则基于 RoPE LLM 的外推性能将得到改善。如果 β 减少到 β1、β2、β3（如下计算），则每个维的 cos 或 sin 输入将分别从 0 到 π/2、π、2π，从而带来更显著的改进。

与如500这些小得多的基数相比（注：其中 cos (t − s)θn 的每个周期都适合训练上下文），RoPE（Su，2021）中的默认基数为 10000，这会导致某些维度的周期超出训练上下文，如图所示。因此，对于基于 RoPE 的 LLM，存在一个特定的特征维 dextra。对于 dextra 之前的维，相应 θn 的周期保持小于 Ttrain，而对于 dextra 之后的维，周期超过 Ttrain。换句话说，本质上，dextra 是 cos (t − s)θn 和 sin (t − s)θn 在预训练或微调期间可以在一个周期内循环其值的维数。

引理 1 .（ 关键维 定义）对于以上下文长度 Ttrain 预训练的基于 RoPE LLM，假设自注意头的大小为 d，最多有前面dextra维数可以感知完整的周期信息，从而得到足够的训练进行外推，其形式化描述如下：

基于关键维的概念，可以阐明在原始上下文长度下使用较大基数对 RoPE 进行微调时的外推结果。对于 LLaMA2（Touvron，2023b），由于前 92 个维度的周期适合训练长度，因此这些特征维度从一开始就具有强大的微调基础，可以适应扩展上下文中位置嵌入的新周期性变化。因此，当使用较大的基值（如 1000000）对 RoPE 进行微调时，即使调整长度小于较大基数对应的扩展周期，这些维仍然可以正确表示位置信息。

然而，对于最后 36 个维，缺乏对周期性的充分理解会导致过拟合。此外，当基数扩大，周期延长时，这些维仍然无法捕获上下文长度内的全部位置信息。因此，只有当先前观察到 θn(t − s) 的值时，这些维才是可靠的。因此，可以将关键维的更新周期作为基于 RoPE LLM 外推的上限。结果得到 定理 2 ，即具有较大基数基于 RoPE 外推的规模化定律，该定律说明了基数与外推上限之间的关系，从而回答了 问题 2 。

定理 2 .（ 较大基数的规模化定律 ）对于用上下文长度 Ttrain 预训练的基于 RoPE LLM，如果将基数调整为 β > 10000 并仍然用上下文长度 Ttrain 进行微调，则基于 RoPE LLM 的外推性能将得到改善。具有较大基数 Textra、基于 RoPE LLM 的外推上限计算如下：