从 LLM 的 RoPE 编码到 PI 与 NTK-Aware

出发点

在 Transformer 早期，位置编码通常是绝对位置编码，直接加到 token embedding 上，但这类方法存在几个问题：

• 不显式依赖相对位置关系，模型必须自己推导出来，而大量任务本质上依赖的是相对距离而非绝对位置；
• 可学习 PE 长度外推能力差，正弦风格的 PE 虽然理论上可以外推，但是在注意力层中并不保持结构的一致性，也就是两个 token 相对位置关系相同的时候，不一定能够得到相同的 $Q{K}^T$；
• 破坏了注意力的平移不变性，也就是整个序列的平移，在我们看来 token 之间的关系不应该发生变化，但是绝对 PE 做不到这一点；

前置知识：向量的旋转

在二维空间中，旋转 $\theta$ 的线性变换是：

$$R(\theta )=\left(\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right)$$

对向量 $x=(x_1,x_2)$：

$$R(\theta )x=\left(\begin{array}{c}{x}_1\cos \theta -x_2\sin \theta \\ x_1\sin \theta +x_2\cos \theta \end{array}\right)$$

以上的操作，具有这样的性质：

1. 向量的长度不变
2. 在这两个向量进行相同的旋转后，两个向量的内积不会发生变化
3. 两个向量各自旋转了不同角度以后，它们的点积的变化，只取决于这两个角度的差，与各自转了多少无关。

第 3 点性质的推导：

设有两个原始向量 $q,k$，分别被旋转 ${\theta }_i$ 和 ${\theta }_j$：

$$q^{\prime }=R({\theta }_i)q,k^{\prime }=R({\theta }_j)k$$

它们的点积是：

$$\begin{array}{rl}(q^{\prime }{)}^{\mathrm{\top }}{k}^{\prime }& =(R({\theta }_i)q{)}^{\mathrm{\top }}(R({\theta }_j)k)\\ & =q^{\mathrm{\top }}R({\theta }_i{)}^{\mathrm{\top }}R({\theta }_j)k\end{array}$$

注意看，$R(\theta )$ 是由三角函数组成的，由于 $\cos \theta =\cos -\theta ,\sin -\theta =-sin\theta$，故满足性质 $R({\theta }_i{)}^{\mathrm{\top }}=R(-{\theta }_i)$：

$$(q^{\prime }{)}^{\mathrm{\top }}{k}^{\prime }=q^{\mathrm{\top }}R(-{\theta }_i)R({\theta }_j)k$$

旋转矩阵还满足性质：

$$R(-{\theta }_i)R({\theta }_j)=R({\theta }_j-{\theta }_i)$$

于是得到：

$$(q^{\prime }{)}^{\mathrm{\top }}{k}^{\prime }=(R({\theta }_i)q{)}^{\mathrm{\top }}(R({\theta }_j)k)=q^{\mathrm{\top }}R({\theta }_j-{\theta }_i)k$$

可以看到，是 $q^{T}k$ 在中间乘以一个 $R({\theta }_j-{\theta }_i)$，所以旋转后的内积变化只和两个向量旋转的角度差有关。

方法

相对位置的注意力建模

在 Attention 的计算中，有 Query 和 Key 的交互算子，即 $Q{K}^T$，此时 $Q$ 和 $K$ 都是由包含了位置编码的 embedding 计算来的。形式化的，对于位于 $m$ 位置的查询 $Q_i$ 和位于 $n$ 位置的键 $K_i$，注意力机制可以表示为：

$$\begin{array}{r}{Q}_i=f_q(x_i,i),K_j=f_k(x_j,j)\\ Q_i{K}_j^T=f_q(x_i,m)f_k(x_j,n{)}^T\end{array}$$

为了让位置注意力机制显式地考虑相对位置 $m-n$ ，我们可以将这个 $q_m{k}_n^T$ 计算抽象成一个函数：$g(x_m,x_n,m-n)$

$$g(x_m,x_n,m-n)=f_q(x_m,m)f_k(x_n,n{)}^T$$

现在目标就变成了，要找到函数 $g,f_q,f_k$ ，使其满足上面的关系。

RoPE 旋转矩阵

RoPE （Rotary Positional Embedding）提出的方法就是，用对 token 的旋转来加入位置信息。

$$f_q(x_i,i)=R_i{Q}_i.f_k(x_j,j)=R_j{Q}_j$$$$\begin{array}{c}{R}_i=\left(\begin{array}{ccccc}\cos {\theta }_{i,0}& -\sin {\theta }_{i,0}& & & \\ \sin {\theta }_{i,0}& \cos {\theta }_{i,0}& & & \\ & & \cos {\theta }_{i,1}& -\sin {\theta }_{i,1}& \\ & & \sin {\theta }_{i,1}& \cos {\theta }_{i,1}& \\ & & & & \ddots \end{array}\right)\end{array}$$$${\theta }_{i,k}=i\cdot {\omega }_k$$

引入频率 ${\omega }_k$ 来解决三角函数的周期性问题，其计算是：

$${\omega }_k={10000}^{-\frac{2k}{d}}$$

• $i$：token 的 绝对位置索引
• $k$：第 $k$ 个 2D 子空间（embedding 维度按 $(2k,2k+1)$ 分组）
• $d$：attention head 的维度（$Q$,$K$ 的维度）

也就是说，将 embedding 的每两维看成一个 2D 子空间进行旋转。对于 2D 旋转矩阵，其旋转角度受到当前位置 $i$、当前维度 2D 子空间序号（从 0 到 $d/2-1$）的影响，总共有 $d/2$ 个这样的 2D 子空间，所以总共有 $d/2$ 个 2D 旋转矩阵。由向量旋转的性质可得：

$$(R_i{Q}_i)(R_j{K}_j{)}^T=Q_{i}R({\theta }_{j,k}-{\theta }_{i,k})K_j^T$$

由此，RoPE 让注意力机制显式地考虑相对位置，相同的相对位置关系可以得到相同的内积变化。

RoPE 中 $w_k$ 的底数的影响

影响最长相对距离

$w_k$ 的底数（记作 $b$）直接影响了要覆盖的相对距离范围 $di{s}_{max}$。

我们知道，三角函数具有周期性，要保证两个不同的 $\theta$ 得到的结果不同，那 $\theta$ 的值域大小必须限制在 $2\pi$ 以内。

${\omega }_k={10000}^{-\frac{2k}{d}}$ 直接决定了三角函数的频率，其会将每个周期的长度变成 $2\pi /w_k$，在维度最高的位置，其频率最低。也就是 $w_k$ 的最小值 $w_{min}$ 就是

$${\omega }_{min}=b^{-\frac{d-2}{d}}\approx \frac{1}{b}(d\text{ 足够大})$$

在用最大的相对距离 $di{s}_{max}$ 乘以这个系数的时候，要做到：

$$di{s}_{max}\cdot {\omega }_{min}<2\pi$$

故 $di{s}_{max}<2\pi b$，直接受到了底数的影响。

注意，我们这里用的是最低频率，其思想是：至少要有维度能够区分开这两个相对位置，这样模型才有区分这两者的能力。我们要做的并不是保证 $di{s}_{max}$ 在所有维度都能区分开），而是要保证至少有维度可以区分开，允许有一些维度无法区分长程距离。

影响短、中、长程的效果

在维度低的地方，频率高，相对位置的细微变化，会让旋转矩阵的值发生很大的变化。

在维度高的地方，频率低，相对位置的巨大变化，才会让旋转矩阵的值发生很大的变化。

这就决定了：

• 如果 $b$ 的值变大，那同样的相对距离，其旋转矩阵的差异会变小，梯度信号变弱，只有更长的相对距离才能产生与先前相同的梯度信号，使得中短程变弱。
• 如果 $b$ 的值变小，那同样的相对距离，其旋转矩阵的差异会变大，能建模的最长距离变短，强行建模长距离将使其出现大量重叠，模型丢失长程的分辨能力。

RoPE 计算简化

将矩阵乘法转化为 element-wise 乘法和加法：

$${\tilde{\mathit{x}}}_i={\mathit{x}}_i\odot \cos ({\mathit{\theta }}_i)+\text{rotate\_half}({\mathit{x}}_i)\odot \sin ({\mathit{\theta }}_i)$$

其中：

$${\mathit{\theta }}_i=[{\theta }_{i,0},{\theta }_{i,0},{\theta }_{i,1},{\theta }_{i,1},\dots ,{\theta }_{i,d/2-1},{\theta }_{i,d/2-1}{]}^{\mathrm{\top }}\in {\mathbb{R}}^d$$$$\text{rotate\_half}([x_0,x_1,x_2,x_3,\dots ])=[-x_1,x_0,-x_3,x_2,\dots ]$$

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """
    对每对维度旋转90°:
    x: (..., d)
    返回: (..., d)
    """
    x1 = x[..., ::2]  # 偶数索引
    x2 = x[..., 1::2] # 奇数索引
    x_rot = torch.stack([-x2, x1], dim=-1)  # 堆叠成2维对
    return x_rot.flatten(-2)  # 展开回原来的最后一维

def rope(x: torch.Tensor, theta: torch.Tensor) -> torch.Tensor:
    """
    应用 RoPE
    x: (..., d) token 向量
    theta: (..., d) 对应每个维度的角度 theta_{i,k} 展开成 d 维
    返回: (..., d) 旋转后的向量
    """
    return x * torch.cos(theta) + rotate_half(x) * torch.sin(theta)

上下文窗口扩展

RoPE 的外推性问题

RoPE 主要解决的是相对位置建模的问题，但是其外推性交叉。

举例来说，如果模型训练时最大只见过 $m\in [0,2048]$，当推理时突然输入 $m=3000$，这个 $m\cdot {\theta }_i$ 的旋转角度超出了模型见过的域。模型不仅无法识别这个新的绝对位置，甚至连相对位置的距离感知也会失效。这就是所谓的外推性差。

Positional Interpolatiopn (PI)

将更大的位置索引，压缩回到模型训练阶段的位置索引范围内。

假设我们要把上下文窗口扩大 $s$ 倍，PI 的做法非常简单粗暴：将输入的位置索引 $i$ 除以 $s$。

$$i^{\prime }=\frac{i}{s}$$

原本训练范围：$[0,L]$，新的推理范围：$[0,sL]$，PI 变换后：新的范围 $[0,sL]$ 被线性映射回 $[0,L]$。

优点

• 非常稳定： 所有的旋转角度都落在了训练时的域内，Attention 分数不会崩坏。
• 无需重训练（或只需微调）： Meta 在论文中证明，只需极少量的微调（几百步），模型就能适应新的长度。

缺点：丢失了位置编码的分辨率信息。高频与低频以相同的倍数进行缩放，使得相对位置的感知显著减弱，尤其是对相邻 token 的感知减弱。

NTK-Aware Scaled RoPE

在神经网络中，neural tangent kernel (NTK) 是描述无限宽深度神经网络在梯度下降训练过程中演化的核。它最开始由Arthur Jacot, Franck Gabriel, 和Clément Hongler在2018年发表的一篇论文中引入 Neural Tangent Kernel: Convergence and Generalization in Neural Networks

NTK-Aware 不直接修改位置 $m$，而是修改 RoPE 的基数（Base） $b$，即那个 $10000$。NTK-Aware 的做法是：通过增大 $b$ 来拉伸周期。 如果我们要扩展 $s$ 倍上下文，我们将基数改为 $b^{\prime }$：

$$b^{\prime }=b\cdot s^{\frac{d}{d-2}}\approx b\cdot s$$

这样就变成了非线性的变换。当我们将 Base 变大时：

1. 对于高频分量（$i$ 很小）： ${\theta }_i$ 的变化非常微小。这意味着高频部分几乎没有被压缩，保留了区分相邻 Token 的能力；
2. 对于低频分量（$i$ 很大）： ${\theta }_i$ 变化显著，周期被拉长了。这意味着低频部分被有效地内插了，从而适应更长的上下文窗口。