Reading Notes Collections: Context Length Extrapolation
date
May 7, 2025
slug
ctx-extrapolate
status
Published
tags
NLP
summary
type
Post
TL; DR
本文是我在阅读了一系列长度外推的论文以及博客之后的总结，按照个人理解的逻辑链进行了组织，内容组织如下：
- 位置编码的本质：进制表示视角
- 长度外推是在做什么
- 直接外推
- 位置内插（Position Interpolation）
- NTK-aware scaled RoPE
- NTK-by-parts
- NTK-Dynamic Interpolation
- YaRN
位置编码：进制表示视角
位置编码的目标是把位置信息嵌入化编码，使得输入神经网络的信息带有位置信息，在训练/推理的时候可以利用这一信息。简单直接一点的说法就是：要把位置，也就是一个整数 以某种方式揉到输入的向量里面去。
一个简单的做法是直接用这个整数数值，但是根据 Neural Tangent Kernel theory，如果输入的维度很低，并且潜入缺乏高频分量，那么神经网络难以学习高频信号。因此，目前的方法是把这个整数 变成一个向量，比如广泛使用的 RoPE，在这个向量中包含不同频率的分量。
在苏剑林老师的博客 RoPE是一种β进制编码 中，以进制表示的视角来解释了 RoPE（或者更广泛地说，这一类基于三角函数表示的位置编码），下面是我对这个视角的一些理解：
- llm 在利用位置信息时，主要关注相对位置，需要一种层次化的编码方式，有高频分量（快速轮换的，对应进制里面的低位）来表示细粒度的差别，有低频分量（低速轮换的，对应进制里面的高位）来表示粗粒度的差别，这一点和进制表示是一致的。
- 考虑 RoPE 映射的第 个复数: , 其中 ，也就是说位置 在 RoPE 这个编码下第 个分量是 ，其中 是一个周期为 的函数。这和 的 进制表示是类似的，差别是对于进制表示 是模函数，但是在这里的差别和 LLM 利用位置编码的方式无关所以可以忽略。
因此总结来说，RoPE （以及类似的基于三角函数的位置编码）可以类比为一种进制表示编码，不同的维度包含了位置 的轮转频率不同的分量（类似进制表示里面的低位和高位，随着 的变化，低位轮转频率高，高位轮转频率低）。
长度外推：问题定义
长度外推指的是能够在长度超过预训练的最大上下文窗口的文本上做有效的推理，并且性能下降在可接受范围内。
长度外推性实际上是一个训练和预测面对的长度不一致的问题。具体到 RoPE 这种相对位置编码，实际上就是当文本很长的时候，有些相对距离 在训练时没有见过，而神经网络的参数没有针对这些值调整适应过，带来了比较大的性能下降。
在目标上，我们希望一个好的长度外推方法可以：
- 不需要微调或者少量微调，在长文本上取得好的效果。
- 不影响或者少影响原来训练过的文本长度上的效果。
- 不影响或者少影响推理效率。
在有了上面位置编码的进制表示视角之后，我们可以对长度外推的方法做个大致的分类：
- 直接外推：什么改动也不做，就直接用。
- 位置内插（Position Interpolation）：假设 ，原来位置编码是 , 现在用: , 由于 RoPE 这类函数式的位置编码的定义域本来就是连续的，非整数位置也可以作为输入，因此不需要做什么架构上的改动。
- 调整频率（及其变体）：这类方法类比“进制转换”，比如本来用十进制，2 位能表示的范围是 0 - 99，现在换成 16 进制，两位能表示的范围就变成了 0 - 255，而换进制在 RoPE 里面实际上就表现为改变频率。由于 RoPE 编码是连续的，并且 LLM 主要关注相对位置和相对距离，这种方法能够有一定的泛化能力。NTK-aware scaled RoPE, NTK-by-parts, Dynamic NTK interpolation 和 YaRN 都属于这一类方法。
直接外推
直接外推就是什么也不做，直接来多长文本就把练好的模型拿过来直接推理。RoPE 本身有一定的外推能力，但是窗口稍微长一点就会带来 perplexity 爆炸的问题。主要是因为如果窗口比较大，就会有很多相对距离训练时没见过，注意力计算的时候没办法比较好的处理位置信息。
位置内插(Position Interpolation)
思路很简单，就是现在的位置 直接使用原来的 位置的位置编码。个人理解这个能有效的原因是它还是保留了相对位置的信息，由于 RoPE 位置编码的平滑性带来的泛化能力。
但是这样导致位置之间的距离被压缩了可能导致 LLM 对于位置的分辨率下降，因此需要一定量的微调。在 Position Interpolation 的论文中指出用 1000 个 steps 左右就能达到比较好的效果。
NTK-aware Scaled RoPE
这个方法来自于 Reddit 网友 bloc97, 这个名字比较玄乎，来源于一些我看不懂的文章，但是主要的思路就是：越低频越内插，越高频越外推。这个方法不需要微调，就取得了比较好的效果。
通过进制转换的思路可以推导出类似的结果，来比较直观的理解这个方法。
我们期望最低频的分量完全近似于内插，也就是：, 解得:
而在高频分量，由于 比较大，所以 接近 1，因此就近似于外推。
由于这个方法实际上就是调整了 RoPE 的角频率，因此在一些其他文章中也被称为 RoPE-ABF (Adjust Base Frequency).
一个理解这类 “低频内插，高频外推” 的思路的视角是转圈的视角，如下图：
RoPE 的每一个部分可以看成复值函数 ，因此，仅考虑每个二维的部分，那么 .
其中后半部分就是在复平面的单位圆上的点，随着相对距离变大旋转。如果频率越高，那么随着相对距离旋转的越多，在训练了足够多的样本之后，我们可以认为单位圆上大多数点都被看到过了；而如果频率低，那么可能即使是最大窗口，也还有很多地方达不到。
在使用位置编码的时候，如果使用到的相对距离是见过的那些点，或者那些点附近的点，则比较好直接外推，如果离见过的点比较远，外推就会带来剧烈的性能下降，只能考虑内插。因此在高频部分可以偏向外推，在低频部分偏向内插。
NTK-by-parts
在 NTK-aware Scaled RoPE 中，我们只用了低频部分对其内插来得到了缩放系数，然后对所有的维度的频率做了一样的缩放，这种缩放越偏向高频的时候就越接近外推。
然而，这种渐变的方式还是导致了偏高频的成分有一定程度的内插，而高频成分可能对这种 interpolation 比较敏感，一个直观的理解是，在单位圆上的点本来已经非常密，再去压缩可能就导致分辨率下降让 LLM 难以区分。因此，一个比较容易想到的优化就是分段调整频率，对于非常高频的成分，不调整频率，对于非常低频的，完全内插，对于中间部分做渐变的 interpolation.
其中, ，也就是目标上下文窗口大小，和单位圆的一个周期对应的上下文窗口大小的比值。
Dynamic NTK Interpolation
在前面的几种方法中，都有使用到 ，也就是缩放比例，并且计算的时候都采用固定的目标窗口大小和原来的窗口大小的比值。这种方式带来的结果是，在原来训练过的文本长度上有一定程度的效果下降。
Dynamic Scaling 的思路是在自回归推理的时候动态计算这个 ，也就是当文本长度还没超过的时候就不做 interpolation，从而避免了这个针对长文本的效果开销。
但是这种方法的问题是没办法直接利用 KV Cache, 使得推理的开销变大。（暂时好像没想到什么比较直接的改造可以解决这个问题）
YaRN
结合了 scaled attention 和 NTK-by-parts, 仅使用了 400 steps 微调就取得了 SOTA 的效果。文中实验的缩放系数 (对于 llama2) 根据拟合取了