参考

论文

  1. Attention is All You Need NeurIPS 2017
    [1706.03762] Attention Is All You Need

  2. RoFormer: Enhanced Transformer with Rotary Position Embedding
    [2104.09864] RoFormer: Enhanced Transformer with Rotary Position Embedding

  3. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation ICLR 2022
    [2108.12409] Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

  4. Extending Context Window of Large Language Models via Positional Interpolation
    [2306.15595] Extending Context Window of Large Language Models via Positional Interpolation

  5. A Length-Extrapolatable Transformer
    [2212.10554] A Length-Extrapolatable Transformer

  6. YARN: Efficient Context Window Extension of Large Language Models
    [2309.00071] YARN: Efficient Context Window Extension of Large Language Models

  7. Qwen2-VL: Enhancing Vision-Language Model’s Perception of the World at Any Resolution
    [2407.10671] Qwen2-VL: Enhancing Vision-Language Model’s Perception of the World at Any Resolution

  8. NTK-Aware Scaled RoPE
    NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation

  9. Chain-of-Thought Prompting Elicits Reasoning in Large Language Models NeurIPS 2022
    [2201.11903] Chain-of-Thought Prompting Elicits Reasoning in Large Language Models

  10. Self-Consistency Improves Chain of Thought Reasoning in Language Models
    [2203.11171] Self-Consistency Improves Chain of Thought Reasoning in Language Models

文档

  1. 通俗理解RoPE、2D-RoPE、M-RoPE

  2. 再论大模型位置编码及其外推性(万字长文)


1. 背景

必要性

Transformer模型的自注意力机制缺乏对序列顺序的处理能力。与RNN/LSTM通过循环结构天然保留顺序信息不同,Transformer并行处理所有token,因此需要显式地引入位置信息来区分序列中元素的顺序。

在自然语言处理中,词序对语义至关重要。例如”我吃苹果”和”苹果吃我”包含相同词汇但顺序不同,语义完全不同。如果模型无法感知位置差异,就无法正确理解句子的含义。

外推性挑战

外推性挑战的核心问题:例如,如果一个模型在训练时只使用了512个 token 的文本,那么在预测时如果输入超过512个 token,模型可能无法正确处理。这就限制了大模型在处理长文本或多轮对话等任务时的效果。如何确保模型在推理阶段能够处理远超预训练时所见文本长度的序列,已成为当前大模型面临的核心问题之一。我们将这个问题视为大模型的外推挑战

目前许多大模型已支持长文本推理,如GPT-4 Turbo可处理超过128k tokens,Baichuan2可处理高达192K tokens。然而,由于内存资源限制,这些模型在预训练阶段通常只涉及约4k tokens的训练。

理想的位置编码应满足三个条件:

  1. 唯一性:每个位置输出唯一的编码
  2. 外推能力:具有良好的外推能力,能够处理超出训练长度的序列
  3. 相对位置一致性:任意两个位置之间的相对距离在不同长度的句子中应保持一致

2. 绝对位置编码 (Absolute Positional Encoding)

简介

绝对位置编码使用正弦和余弦函数来编码每个位置的绝对位置,最早在Google的”Attention is All You Need”论文中提出,用于Transformer的位置编码。

计算步骤

Sinusoidal位置编码是Transformer中常用的位置编码方法。位置编码是一个向量,被添加到输入词嵌入中,其计算方法如下:

对于位置 $p$ 和维度 $i$:

$$PE(p, 2i) = \sin\left(\frac{p}{10000^{2i/d}}\right)$$

$$PE(p, 2i+1) = \cos\left(\frac{p}{10000^{2i/d}}\right)$$

其中:

  • $p$ 是位置索引
  • $i$ 是维度索引
  • $d$ 是嵌入的总维度

频率递减机制:通过指数递增的分母来实现。对于较高的维度 $2i$,分母的值更大,从而使得频率较低。这使得高维度的编码可以捕捉长距离的信息,而低维度的编码则捕捉短距离的信息。


核心性质

相对位置表达能力

  1. $PE(i+k)$ 可以表示成 $PE(i)$ 的线性函数
    Sinusoidal编码具有相对位置表达能力。对于固定位置距离 $k$,$PE(i+k)$ 可以表示成 $PE(i)$ 的线性函数。这表明Sinusoidal编码不仅能表示绝对位置,还能隐式地学习相对位置关系。
理解基本符号和定义
  • $i$ 是位置索引(比如第0个位置、第1个位置等)
  • $k$ 是相对距离(比如 $k=2$ 表示向后数2个位置)
  • $PE(i)$ 表示位置 $i$ 的编码向量(整个向量)
  • $PE(i, 2i)$ 表示位置 $i$ 的编码向量的第 $2i$ 维(偶数维度,用sin计算)
  • $PE(i, 2i+1)$ 表示位置 $i$ 的编码向量的第 $2i+1$ 维(奇数维度,用cos计算)
  • $w_i = \frac{1}{10000^{2i/d}}$ 是一个常数,只和维度索引 $i$ 有关(注意:这个 $i$ 和位置索引 $i$ 是不同的!为了区分,有时候会用 $j$ 表示维度索引)

为什么PE有两个参数?

用一个例子来说明:

位置编码是一个向量。假设维度 $d=8$(也就是这个向量有8个数字)。

对于位置 $i=0$(第0个词),它的位置编码向量 $PE(0)$ 有8个数字:

$$PE(0) = [PE(0,0), PE(0,1), PE(0,2), PE(0,3), PE(0,4), PE(0,5), PE(0,6), PE(0,7)]$$

这里:

  • 第一个参数($i=0$):表示这是第几个位置的编码
  • 第二个参数($0,1,2,3,4,5,6,7$):表示这是向量的第几个维度(第几列)

所以:$PE(0, 0)$ = 位置0的编码向量的第0维(第一个数字), $PE(0, 1)$ = 位置0的编码向量的第1维(第二个数字), $PE(0, 2)$ = 位置0的编码向量的第2维(第三个数字)

  • …以此类推

为什么写成 $PE(i, 2i)$ 和 $PE(i, 2i+1)$?

这是为了表示偶数维度奇数维度的模式:
如果 $d=8$,那么:

  • 偶数维度:$PE(i, 0), PE(i, 2), PE(i, 4), PE(i, 6)$ 用 $\sin$ 函数计算
  • 奇数维度:$PE(i, 1), PE(i, 3), PE(i, 5), PE(i, 7)$ 用 $\cos$ 函数计算

为什么要一个用sin,一个用cos?

这是Transformer论文的精心设计,原因如下:

  • 互补性(Complementarity): $\sin$ 和 $\cos$ 是互补的三角函数:
    • 它们永远不会同时为0,也不会同时为1,这样设计可以确保相邻的两个维度(比如第0维和第1维)包含不同的信息,避免信息重复。
  • 相对位置关系(Relative Position)
    • 使用 $\sin$ 和 $\cos$ 的配对设计,使得我们可以利用三角函数的性质来表示相对位置关系:
    • $\sin(x+y) = \sin(x)\cos(y) + \cos(x)\sin(y)$
    • $\cos(x+y) = \cos(x)\cos(y) - \sin(x)\sin(y)$
    • 这让我们可以将 $PE(i+k)$ 表示为 $PE(i)$ 的线性函数,这正是我们要证明的核心性质!
  • 周期性和唯一性
    • $\sin$ 和 $\cos$ 都是周期函数,但它们的相位相差90度
    • 通过在不同维度使用不同的频率(通过 $10000^{2i/d}$ 实现),每个位置都能得到唯一的编码
    • 即使某些维度值相同,其他维度的sin/cos组合也能区分不同位置
  • 对称性和稳定性
    • sin和cos的配对使用使得位置编码具有良好的数学性质,比如:
    • 对称性:$PE(t+k) \cdot PE(t) = PE(t) \cdot PE(t-k)$
    • 相对位置不变性:内积只与相对位置 $k$ 有关,与绝对位置 $t$ 无关
$PE(i+k)$ 可以表示成 $PE(i)$ 的线性函数推导【三角函数和角公式】

1.位置编码的公式

根据定义,对于任意位置 $p$ 和维度索引 $i$:

$$PE(p, 2i) = \sin\left(\frac{p}{10000^{2i/d}}\right) = \sin(p \cdot w_i)$$

$$PE(p, 2i+1) = \cos\left(\frac{p}{10000^{2i/d}}\right) = \cos(p \cdot w_i)$$

注意:这里我们用 $w_i = \frac{1}{10000^{2i/d}}$ 来简化写法

2.写出位置 $i$ 和 $i+k$ 的编码

对于位置 $i$:

  • 偶数维度:$PE(i, 2i) = \sin(i \cdot w_i)$
  • 奇数维度:$PE(i, 2i+1) = \cos(i \cdot w_i)$

对于位置 $i+k$:

  • 偶数维度:$PE(i+k, 2i) = \sin((i+k) \cdot w_i)$
  • 奇数维度:$PE(i+k, 2i+1) = \cos((i+k) \cdot w_i)$

3.三角函数的和角公式(这是关键!)

$$\sin(A+B) = \sin(A)\cos(B) + \cos(A)\sin(B)$$

$$\cos(A+B) = \cos(A)\cos(B) - \sin(A)\sin(B)$$

这两个公式告诉我们,$\sin(A+B)$ 可以写成 $\sin(A)$ 和 $\cos(A)$ 的”线性组合”(就是加减乘除的组合)

4.应用和角公式到 $PE(i+k, 2i)$

我们有:$PE(i+k, 2i) = \sin((i+k) \cdot w_i)$
应用和角公式:
$$= \sin(i \cdot w_i)\cos(k \cdot w_i) + \cos(i \cdot w_i)\sin(k \cdot w_i)$$

【关键】注意到:

  • $\sin(i \cdot w_i) = PE(i, 2i)$(这是位置 $i$ 的偶数维度)
  • $\cos(i \cdot w_i) = PE(i, 2i+1)$(这是位置 $i$ 的奇数维度)

所以:

$$PE(i+k, 2i) = PE(i, 2i) \cdot \cos(k \cdot w_i) + PE(i, 2i+1) \cdot \sin(k \cdot w_i)$$

重要观察:因为 $k$ 是固定的,所以 $\cos(k \cdot w_i)$ 和 $\sin(k \cdot w_i)$ 都是常数(不依赖位置 $i$)

类似地,我们有:$PE(i+k, 2i+1) = \cos((i+k) \cdot w_i)$

应用和角公式:

$$PE(i+k, 2i+1) = \cos(i \cdot w_i)\cos(k \cdot w_i) - \sin(i \cdot w_i)\sin(k \cdot w_i)$$

再次代入:

  • $\cos(i \cdot w_i) = PE(i, 2i+1)$
  • $\sin(i \cdot w_i) = PE(i, 2i)$

得到:

$$PE(i+k, 2i+1) = PE(i, 2i+1) \cdot \cos(k \cdot w_i) - PE(i, 2i) \cdot \sin(k \cdot w_i)$$

5.总结

我们证明了:

  • $PE(i+k, 2i) = PE(i, 2i) \cdot C_1 + PE(i, 2i+1) \cdot C_2$(其中 $C_1 = \cos(k \cdot w_i)$,$C_2 = \sin(k \cdot w_i)$)

  • $PE(i+k, 2i+1) = PE(i, 2i+1) \cdot C_3 - PE(i, 2i) \cdot C_4$(其中 $C_3 = \cos(k \cdot w_i)$,$C_4 = \sin(k \cdot w_i)$)

因为 $k$ 是固定的,所以所有系数 $C_1, C_2, C_3, C_4$ 都是常数!

结论:$PE(i+k)$ 的每一维都可以写成 $PE(i)$ 的对应维度乘以一些常数,然后相加。这就是”线性函数”的意思!所以对于固定的相对距离 $k$,$PE(i+k)$ 可以表示为 $PE(i)$ 的线性函数。


  1. 两个位置向量的内积只和相对位置 $k$ 有关。

Attention中的重要操作就是内积。计算两个位置的内积 $PE(t) \cdot PE(t+k)$ 如下:

$$PE(t) \cdot PE(t+k) = \sum_{i=0}^{d/2-1} PE(t, 2i) \cdot PE(t+k, 2i) + \sum_{i=0}^{d/2-1} PE(t, 2i+1) \cdot PE(t+k, 2i+1)$$

$$= \sum_{i=0}^{d/2-1} \sin(t \cdot w_{2i}) \sin[(t+k) \cdot w_{2i}] + \sum_{i=0}^{d/2-1} \cos(t \cdot w_{2i}) \cos[(t+k) \cdot w_{2i}]$$

$$= \sum_{i=0}^{d/2-1} \cos(k \cdot w_{2i})$$

其中 $w_{2i} = 1 / 10000^{2i/d}$。

可以看到,最终的结果是关于 $k$ 的一个常数。这表明两个位置向量的内积只和相对位置 $k$ 有关,与绝对位置 $t$ 无关。

对称性

Sinusoidal编码具有对称性。通过计算,很容易得到 $PE(t+k) \cdot PE(t) = PE(t) \cdot PE(t-k)$,这表明Sinusoidal编码具有对称性。

远程衰减

随着 $k$ 的增加,内积的结果会直接减少,即会存在远程衰减。这意味着距离越远的位置,其位置编码的相似度越低,使得模型能够区分不同距离的依赖关系。


优缺点

优点

  1. 捕捉序列顺序:位置编码使Transformer能够区分序列中词的位置,有助于上下文理解,这对自然语言处理至关重要。

  2. 保持全局一致性:位置编码对不同位置是唯一的,有助于模型在整个序列中保持全局一致性。

  3. 增强模型能力:位置编码使自注意力机制能够基于位置对信息进行加权,从而提升模型处理长距离依赖的能力。

缺点

  1. 固定的位置编码:标准的Sinusoidal位置编码是固定的且不可学习。这意味着它无法适应特定任务或数据变化,可能在特定应用中无法提供最优的位置信息。

  2. 有限的序列长度:Sinusoidal位置编码对序列长度存在限制。对于超长序列,固定的位置编码可能无法完全表示序列的全部信息。它缺乏外推能力,因为长度在预设后是固定的。

  3. 位置编码的线性组合:位置编码与词嵌入的线性组合可能导致信息丢失。位置编码的信息可能被嵌入的信息所掩盖,影响模型对位置数据的敏感性。


代码实现

原生实现(使用Embedding层)

1
2
3
4
5
6
# 初始化
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)

# 计算
position_embeddings = self.position_embeddings(position_ids)
embeddings += position_embeddings

手动实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch
import math

def manual_sinusoidal_encoding(max_len: int, d_model: int) -> torch.Tensor:
"""
手动实现 Transformer 正弦位置编码
参数:
max_len: 最大序列长度
d_model: 模型维度 (需为偶数)
返回:
pe: 位置编码矩阵 [max_len, d_model]
"""

if d_model % 2 != 0:
raise ValueError("模型维度 d_model 必须是偶数")

# 初始化位置编码矩阵
pe = torch.zeros(max_len, d_model)

# 生成位置序列 [max_len, 1]
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

# 计算频率项 (分母部分)
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)

# 计算位置编码
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)

return pe

3. 旋转位置编码 RoPE (Rotary Position Embedding)

简介

RoPE是一种位置编码方法,在论文”RoFormer: Enhanced Transformer with Rotary Position Embedding”中提出。其目的是将相对位置信息集成到自注意力中,提升Transformer架构的性能。

RoPE的核心思想是将位置编码嵌入到旋转变换中,使得位置之间的相对关系能够直接影响最终的Attention权重。具体来说,RoPE使用旋转矩阵来变换每个位置的编码,使得在计算attention时,相对位置 $m-n$ 能够被自然地考虑进去。

核心思想

RoPE的关键在于通过绝对位置编码来实现相对位置编码。具体来说,将位置编码表示为二维空间中的旋转操作。对于一个二维向量 $\mathbf{x}$,将其绕原点逆时针旋转 $m$ 弧度后,可以表示为矩阵乘法 $R(m) \cdot \mathbf{x}$,其中旋转矩阵 $R(m)$ 的形式如下:

$$R_m = \begin{bmatrix} \cos m & -\sin m \\ \sin m & \cos m \end{bmatrix}$$

这一变换保留了原始向量的性质,但通过旋转引入了位置信息。如果一个二维向量旋转了 $m$ 弧度,另一个旋转了 $n$ 弧度,那么这两个向量之间的夹角就是 $|m - n|$。基于这样的思想,我们可以有效地引入相对位置的信息。

旋转矩阵的推导:为什么是逆时针旋转?如何实现相对位置编码?

1. 旋转方向:逆时针旋转

在数学中,标准的旋转矩阵表示逆时针旋转(counterclockwise rotation)。这是数学和物理学的约定:

  • 逆时针旋转角度为正
  • 顺时针旋转角度为负

所以 $R_m$ 表示将向量逆时针旋转 $m$ 弧度

2. 旋转矩阵的推导过程

假设我们有一个二维向量 $\mathbf{x} = [x_0, x_1]^T$,在极坐标下表示为:

  • 长度:$r = \sqrt{x_0^2 + x_1^2}$
  • 角度:$\alpha = \arctan(x_1/x_0)$

将这个向量逆时针旋转 $m$ 弧度后,新向量的角度变为 $\alpha + m$,长度不变。

新向量的坐标:

  • $x’_0 = r \cos(\alpha + m)$
  • $x’_1 = r \sin(\alpha + m)$

利用三角函数的和角公式:

  • $\cos(\alpha + m) = \cos\alpha \cos m - \sin\alpha \sin m$
  • $\sin(\alpha + m) = \sin\alpha \cos m + \cos\alpha \sin m$

因为 $x_0 = r\cos\alpha$,$x_1 = r\sin\alpha$,所以:

$$x’_0 = r(\cos\alpha \cos m - \sin\alpha \sin m) = x_0 \cos m - x_1 \sin m$$

$$x’_1 = r(\sin\alpha \cos m + \cos\alpha \sin m) = x_1 \cos m + x_0 \sin m$$

写成矩阵形式:

$$\begin{bmatrix} x’_0 \\ x’_1 \end{bmatrix} = \begin{bmatrix} \cos m & -\sin m \\ \sin m & \cos m \end{bmatrix} \begin{bmatrix} x_0 \\ x_1 \end{bmatrix}$$

这就是旋转矩阵 $R_m$ 的推导过程。

3. 为什么这样能实现相对位置编码?

关键思想:如果我们对query向量旋转 $m$ 弧度,对key向量旋转 $n$ 弧度,那么:

  • 旋转后的query:$q’ = R_m \cdot q$
  • 旋转后的key:$k’ = R_n \cdot k$

计算它们的内积:

$$q’^T k’ = (R_m q)^T (R_n k) = q^T R_m^T R_n k$$

旋转矩阵的一个重要性质是:$R_m^T R_n = R_{m-n}$(转置矩阵与旋转矩阵的乘积等于角度差的旋转矩阵)

证明:$R_m^T = \begin{bmatrix} \cos m & \sin m \\ -\sin m & \cos m \end{bmatrix}$,所以:

$$R_m^T R_n = \begin{bmatrix} \cos m & \sin m \\ -\sin m & \cos m \end{bmatrix} \begin{bmatrix} \cos n & -\sin n \\ \sin n & \cos n \end{bmatrix} = \begin{bmatrix} \cos(m-n) & -\sin(m-n) \\ \sin(m-n) & \cos(m-n) \end{bmatrix} = R_{m-n}$$

所以:

$$q’^T k’ = q^T R_{m-n} k$$

关键观察:内积只依赖于角度差 $m-n$(或 $n-m$,取决于定义),也就是相对位置!与绝对位置 $m$ 和 $n$ 无关。

4. 直观理解

想象两个向量在二维平面上:

  • 位置 $m$ 的向量旋转了 $m$ 弧度
  • 位置 $n$ 的向量旋转了 $n$ 弧度
  • 它们之间的相对角度是 $|n-m|$,这正是我们想要的相对位置信息。

通过旋转操作,我们巧妙地将绝对位置信息转换成了相对位置信息,这正是RoPE的核心创新。

对于位置 $m$ 的查询向量和位置 $n$ 的键向量,RoPE使用:

$$f(q, m) = q \cdot e^{im\theta}$$

$$f(k, n) = k \cdot e^{in\theta}$$


计算步骤

对于高维向量,RoPE将向量分组处理。对于维度为 $d$ 的向量,将其分为 $d/2$ 组,每组2维,分别应用旋转矩阵。

对于位置 $m$,旋转矩阵 $R_m$ 是一个块对角矩阵,由 $d/2$ 个2×2旋转块组成:

$$R_m = \begin{bmatrix}
R_0 & 0 & \cdots & 0 \\
0 & R_1 & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & \cdots & R_{d/2-1}
\end{bmatrix}$$

其中每个 $R_i$ 是一个2×2旋转矩阵:

$$R_i = \begin{bmatrix}
\cos(m\theta_i) & -\sin(m\theta_i) \\
\sin(m\theta_i) & \cos(m\theta_i)
\end{bmatrix}$$

或者更简洁地,对于第 $i$ 组(对应维度 $2i$ 和 $2i+1$),旋转矩阵为:

$$\begin{bmatrix}
\cos(m\theta_i) & -\sin(m\theta_i) \\
\sin(m\theta_i) & \cos(m\theta_i)
\end{bmatrix}$$

其中 $\theta_i = 10000^{-2i/d}$。

RoPE的自注意力操作流程:

  1. 对于token序列中的每个词嵌入向量,首先计算其对应的query和key向量。
  2. 计算每个token位置对应的旋转位置编码。
  3. 对每个token位置的query和key向量的元素进行旋转变换,”两两一组”。
  4. 计算query和key之间的内积,得到自注意力计算结果。

旋转变换 $R_{\Theta,m}^d x$ 的详细公式:

$$R_{\Theta,m}^d x = \begin{bmatrix} x_0 \\ x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{d-2} \\ x_{d-1} \end{bmatrix} \otimes \begin{bmatrix} \cos m\theta_0 \\ \cos m\theta_0 \\ \cos m\theta_1 \\ \cos m\theta_1 \\ \vdots \\ \cos m\theta_{d/2-1} \\ \cos m\theta_{d/2-1} \end{bmatrix} + \begin{bmatrix} -x_1 \\ x_0 \\ -x_3 \\ x_2 \\ \vdots \\ -x_{d-1} \\ x_{d-2} \end{bmatrix} \otimes \begin{bmatrix} \sin m\theta_0 \\ \sin m\theta_0 \\ \sin m\theta_1 \\ \sin m\theta_1 \\ \vdots \\ \sin m\theta_{d/2-1} \\ \sin m\theta_{d/2-1} \end{bmatrix}$$


计算示例

假设有一句话”浙江大学在杭州”,其中 $x$=”杭”为需要处理的词。$m$ 是 $x$ 在sequence中的位置(在本句中为6),假设 $x$ 的embedding维度为512,我们可以通过以下公式计算 $\theta$ 值(其中 $i$ 是embedding中的维度位置,$d$ 是embedding的维度数):

$$\theta_i = 10000^{-2i/d}$$

接下来,我们对词向量的每两维进行相同角度的旋转。例如,$x[0]$ 和 $x[1]$ 乘以以下矩阵:

$$\begin{bmatrix} \cos(m\theta_0) & -\sin(m\theta_0) \\ \sin(m\theta_0) & \cos(m\theta_0) \end{bmatrix}$$


RoPE的作用

  1. 引入相对位置信息:RoPE的设计使得模型不仅能够保留绝对位置信息,还能直接学习位置之间的相对距离和关系。这对于捕捉长距离依赖和相对位置信息至关重要。

  2. 增强模型的表示能力:通过用旋转矩阵对每个位置进行编码,RoPE使得向量的旋转角度能够反映位置之间的相对关系。这使得模型在attention计算过程中能够自然地融入相对位置信息,提升其对序列内相对位置关系的敏感性和处理能力。

  3. 提升模型泛化能力:RoPE对位置编码进行旋转变换,使得位置关系能够被直接学习和调整。这种设计有助于模型更好地泛化到不同长度和类型的序列,因为它直接捕捉和利用了相对位置关系。

对比

特性 RoPE旋转位置编码 绝对位置编码
位置信息捕捉 捕捉token之间的相对位置关系 只捕捉每个token的绝对位置
序列长度适应性 在处理长序列时表现优异,特别是长距离依赖 对于长序列不够有效,位置信息可能被稀释
数学操作 通过旋转操作引入位置信息的相位偏移 直接嵌入正弦余弦函数
注意力机制的效果 能够在attention机制中自然地处理相对位置信息 无法直接捕捉相对位置,适合较短序列

4. 多维旋转位置编码 (2D/3D RoPE & M-RoPE)

背景

传统的旋转位置嵌入只能捕捉一维序列的位置信息,而M-RoPE通过将原始旋转嵌入分解为代表时间、高度和宽度的三个部分,使得大规模语言模型能够同时捕捉和整合一维文本序列、二维视觉图像以及三维视频的位置信息。这一创新赋予了语言模型强大的多模态处理和推理能力,能够更好地理解和建模复杂的多模态数据。

简介

Qwen2-VL在架构上的另一重要创新则是多模态旋转位置编码(M-RoPE),将原始旋转嵌入分解为代表时间、高度和宽度的三个部分,使得大规模语言模型能够同时捕捉和整合一维文本序列、二维视觉图像以及三维视频的位置信息。

一维RoPE回顾

RoPE的关键在于通过绝对位置编码来实现相对位置编码。具体来说,将位置编码表示为二维空间中的旋转操作。对于一个二维向量 $\mathbf{x}$,将其绕原点旋转 $m$ 弧度后,可以表示为矩阵乘法 $R(m) \cdot \mathbf{x}$,其中旋转矩阵 $R(m)$ 的形式如下:

$$R_m = \begin{bmatrix} \cos m & -\sin m \\ \sin m & \cos m \end{bmatrix}$$

对于高维向量,1D RoPE将维度为 $d$ 的向量分为 $d/2$ 组,每组2维,构造块对角矩阵:

$$R_n = \begin{bmatrix}
R_0 & 0 & \cdots & 0 \\
0 & R_1 & \cdots & 0 \\
\vdots & \vdots & \ddots & \vdots \\
0 & 0 & \cdots & R_{d/2-1}
\end{bmatrix}$$

其中每个 $R_i$ 是一个2×2旋转矩阵:

$$R_i = \begin{bmatrix}
\cos(n\theta_i) & -\sin(n\theta_i) \\
\sin(n\theta_i) & \cos(n\theta_i)
\end{bmatrix}$$

其中 $\theta_i = 10000^{-2i/d}$。

2D RoPE

核心思想:RoPE从1维扩展到2维的一个简单结论是:针对一个位置 $(x, y)$,对维度为 $d$ 的输入向量分成两半,前一半向量用 $x$ 的一维RoPE矩阵 $(R_x)$ 处理,后一半向量用 $y$ 的一维RoPE矩阵 $(R_y)$ 处理,然后再将两半处理后的结果拼接在一起,就完成了2维的RoPE处理。

计算过程

假设一个图像的patch的embedding是512维,位置是 $(x, y)$,其中 $x$ 表示Height方向索引,$y$ 表示Width方向索引。类似于一维情况,我们为每个坐标计算角度 $\theta$,计算公式与一维相同:

$$\theta_i = 10000^{-2i/d}$$

对于二维情况的旋转矩阵,假设embedding的维度为512维,分为四个一组。对于位置 $(x, y)$ 的embedding:

  • $x[0], x[1]$ 乘以下面的矩阵:
    $$\begin{bmatrix} \cos(x\theta_0) & -\sin(x\theta_0) \\ \sin(x\theta_0) & \cos(x\theta_0) \end{bmatrix}$$

  • $x[2], x[3]$ 乘以下面的矩阵:
    $$\begin{bmatrix} \cos(y\theta_0) & -\sin(y\theta_0) \\ \sin(y\theta_0) & \cos(y\theta_0) \end{bmatrix}$$

2D RoPE矩阵形式

词向量的512维embedding可以每四个一组,乘以下面的矩阵 $R_{x,y}$:

$$R_{x,y} = \begin{bmatrix}
\cos(x\theta) & -\sin(x\theta) & 0 & 0 \\
\sin(x\theta) & \cos(x\theta) & 0 & 0 \\
0 & 0 & \cos(y\theta) & -\sin(y\theta) \\
0 & 0 & \sin(y\theta) & \cos(y\theta)
\end{bmatrix}$$

总结

  • RoPE-1D:两两一组,每组乘以2×2旋转矩阵
  • RoPE-2D:四个一组,每组乘以4×4块对角矩阵(包含两个2×2旋转块)

3D RoPE & M-RoPE

核心思想:对于位置 $(x, y, z)$,将维度为 $d$ 的输入向量分成三部分。第一部分用 $x$ 的一维RoPE矩阵 $(R_x)$ 处理,中间部分用 $y$ 的一维RoPE矩阵 $(R_y)$ 处理,最后一部分用 $z$ 的一维RoPE矩阵 $(R_z)$ 处理,然后将三部分处理后的结果拼接在一起。

在Qwen2-VL中的应用

在Qwen2-VL中使用3D RoPE,类似于2D RoPE,但RoPE-3D是每六个元素一组。假设embedding维度为512,则分为6组:

  • $x[0], x[1]$ 使用旋转矩阵:
    $$\begin{bmatrix} \cos(x\theta_0) & -\sin(x\theta_0) \\ \sin(x\theta_0) & \cos(x\theta_0) \end{bmatrix}$$

  • $x[2], x[3]$ 使用旋转矩阵:
    $$\begin{bmatrix} \cos(y\theta_0) & -\sin(y\theta_0) \\ \sin(y\theta_0) & \cos(y\theta_0) \end{bmatrix}$$

  • $x[4], x[5]$ 使用旋转矩阵:
    $$\begin{bmatrix} \cos(z\theta_0) & -\sin(z\theta_0) \\ \sin(z\theta_0) & \cos(z\theta_0) \end{bmatrix}$$

M-RoPE在Qwen2-VL中的机制

M-RoPE(Multimodal Rotary Position Embedding)将位置分解为三个维度:

  • 时间维度(Time):第一个维度
  • 高度维度(Height):第二个维度
  • 宽度维度(Width):第三个维度

M-RoPE有三个ID:

  • 时间ID(Temporal ID)
  • 高度ID(Height ID)
  • 宽度ID(Width ID)

根据输入类型的行为

  1. 纯文本输入:M-RoPE等价于1D-RoPE
  2. 文本+图像输入:时间ID保持不变,只有高度ID和宽度ID变化
  3. 视频输入:时间ID变为帧ID

优点

  1. 在下游任务中表现更好,特别是视频基准测试
  • 相比1D-RoPE,使用M-RoPE在下游任务中实现了更好的性能,特别是在视频基准测试中(PerceptionTest、NextQA、STAR等)。
  1. 对超出训练长度文本的泛化能力更强
  • M-RoPE带来的外推性使得Qwen2-VL在面对超过训练长度的文本时也有不错的泛化性。实验表明,即使在序列长度从8K扩展到80K的情况下,模型精度仍能保持相对稳定(约68%-71%)。

5. 位置编码扩展方法

为了解决位置编码的外推性问题,研究者们提出了多种扩展方法,使模型能够处理超出训练时所见长度的序列。这些方法主要包括:位置线性内插(PI)、NTK-Aware缩放RoPE、ALiBi线性偏差注意力和YARN等方法。

PI (Position Interpolation) 位置线性内插

简介

PI(Position Interpolation)是一种用于扩展Transformer模型上下文窗口的扩展技术。其核心思想是线性缩放位置索引,使模型能够处理超出训练长度的序列。

数学本质:在保持位置编码相对顺序的同时,将超出训练长度的位置映射回原始编码范围内。

计算步骤

关键公式

对于目标长度 $L$ 和原始最大长度 $N$($L > N$),调整后的位置索引为:

$$i’ = i \times \frac{N}{L}$$

其中 $i \in [0, L-1]$ 是实际位置,$i’ \in [0, N-1]$ 是缩放后的位置。

实现步骤

  1. 原始旋转位置编码计算
1
2
3
4
def compute_rope(pos, dim, base=10000):
theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
pos_theta = pos * theta
return torch.cat([pos_theta.cos(), pos_theta.sin()], dim=-1)
  1. 位置缩放
1
2
3
def scale_position(L, N, current_length):
scale_factor = N / current_length
return torch.arange(0, current_length) * scale_factor
  1. 应用缩放后的位置编码
1
2
3
4
5
def apply_pi_rope(q, k, L, max_train_length=2048):
scaled_pos = scale_position(L, max_train_length, L)
q_rot = compute_rope(scaled_pos, q.size(-1))
k_rot = compute_rope(scaled_pos, k.size(-1))
return q * q_rot, k * k_rot
PI原理详解:为什么线性插值能够扩展上下文窗口?

1. 问题背景

假设模型在训练时使用的最大序列长度为 $N = 2048$,现在需要处理长度为 $L = 8192$ 的序列(4倍扩展)。直接使用原始位置编码会遇到两个问题:

  • 位置 $2048$ 到 $8191$ 的位置编码从未在训练中见过
  • 模型可能无法正确理解这些”新”位置的语义

2. PI的核心思想:将新位置”压缩”到已知范围

PI的关键洞察是:将超出训练长度的位置映射回训练时见过的位置编码范围

数学推导

对于位置 $i \in [0, L-1]$,我们定义缩放后的位置 $i’$:

$$i’ = i \times \frac{N}{L}$$

关键性质

  • 当 $i = 0$ 时,$i’ = 0$(起始位置不变)
  • 当 $i = L-1$ 时,$i’ = (L-1) \times \frac{N}{L} \approx N-1$(末尾位置映射到训练范围末尾)
  • 对于中间位置,线性映射保持了相对顺序

3. 为什么线性插值有效?

相对位置保持不变

考虑两个位置 $i$ 和 $j$($i < j$),它们的相对距离为 $\Delta = j - i$。

缩放后的相对距离:

$$\Delta’ = j’ - i’ = (j - i) \times \frac{N}{L} = \Delta \times \frac{N}{L}$$

关键观察:相对距离按相同比例缩放,相对位置的相对大小关系保持不变

例如,如果位置 $i$ 和 $j$ 在原始序列中距离为 $\Delta = 100$,在缩放后距离为 $\Delta’ = 100 \times \frac{2048}{8192} = 25$。虽然绝对值缩小了,但相对关系($i < j$)和比例关系($\Delta’/N = \Delta/L$)保持不变。

4. 与RoPE的结合

对于RoPE,位置 $i$ 的旋转角度为 $\theta_i = i \times \theta_0$(其中 $\theta_0$ 是基础频率)。

使用PI后,位置 $i$ 实际使用的角度为:

$$\theta’_i = i’ \times \theta_0 = i \times \frac{N}{L} \times \theta_0$$

这相当于将旋转频率降低了 $L/N$ 倍,使得”新位置”的编码落在模型训练时见过的角度范围内。

5. 直观理解

想象位置编码空间是一个长度为 $N$ 的”已知区域”。PI将长度为 $L$ 的序列”压缩”到这个已知区域内:

  • 原本位置 $0$ 到 $L-1$ 需要 $L$ 个单位长度
  • 现在压缩到 $0$ 到 $N-1$ 的 $N$ 个单位长度
  • 压缩比为 $N/L < 1$,所有位置编码都在训练范围内

6. 局限性分析

位置分辨率降低

原本每个位置间隔为 $1$,缩放后间隔变为 $\frac{N}{L} < 1$。这意味着:

  • 原本可以区分 $L$ 个不同位置
  • 现在只能区分约 $N$ 个不同的位置编码值
  • 分辨率降低了 $L/N$ 倍

例如,从 $2048$ 扩展到 $8192$ 时,分辨率降低 $4$ 倍。这意味着原本距离为 $4$ 的两个位置,现在可能使用相同或非常接近的位置编码值。

为什么仍然有效?

虽然分辨率降低,但模型主要依赖相对位置关系而非绝对位置。只要相对位置关系保持(通过线性插值保证),模型就能理解序列结构。

优缺点

优势

  1. 支持8-32倍长度扩展:例如从2K扩展到32K
  2. 几乎不损失原始长度性能:精度下降小于1%
  3. 无需微调即可应用:可以直接应用到预训练模型

局限

  1. 超长距离位置分辨率降低:在非常长的距离上,位置分辨率会降低
  2. 对绝对位置敏感的任务效果下降:对绝对位置敏感的任务性能可能下降
  3. 需要重新计算注意力矩阵:需要重新计算attention矩阵

NTK-Aware Scaled RoPE (非线性内插)

简介

NTK-Aware Scaled RoPE是对旋转位置编码(RoPE)的改进,旨在增强模型在处理长序列时的外推能力。其核心思想基于神经切线核(Neural Tangent Kernel, NTK)理论,通过动态缩放RoPE的频率基来平衡高频和低频位置信息的分布,从而缓解外推到更长序列时的性能下降。

RoPE回顾

RoPE通过旋转矩阵将位置编码集成到attention计算中。对于位置 $m$ 和维度 $i$,旋转角度 $\theta_i$ 为:

$$\theta_i = \frac{1}{10000^{2i/d}}$$

这在复数域中生成旋转矩阵。内积 $q_m^T k_n$ 隐式编码了相对位置 $|m-n|$。

NTK理论启发

问题:当外推到长序列时,RoPE的高频分量(对应较大的 $i$)衰减过快,导致位置区分度丧失。

解决方案:NTK-Aware方法调整频率基,使高频分量衰减更慢,从而保留更多位置信息。

计算步骤

步骤1:调整频率基

原始RoPE基:base = 10000
NTK-Aware将其缩放到更大的值:base_new = base × s

其中 $s$ 是缩放因子,通常与序列长度相关。

步骤2:计算缩放后的角度

对于维度 $i$,新的旋转角度 $\theta’_i$ 为:

$$\theta’_i = \frac{1}{(base \times s)^{2i/d}}$$

效果:这等价于降低频率,防止高频分量过早衰减。

步骤3:动态调整缩放因子

在推理时,$s$ 根据当前序列长度 $L$ 和训练时使用的最大长度 $L_{\text{train}}$ 动态计算:

$$s = \left(\frac{L}{L_{\text{train}}}\right)^\alpha$$

其中 $\alpha$ 是超参数,控制缩放的强度。通常设置为 $\alpha = \frac{d}{d-2}$,$d$ 为头维度。

步骤4:应用旋转矩阵

使用调整后的 $\theta’i$ 生成旋转矩阵 $R{\theta’, m}$,其余计算与原始RoPE方法保持一致。

NTK-Aware Scaled RoPE原理详解:为什么调整频率基能改善外推?

1. NTK理论背景

神经切线核(Neural Tangent Kernel, NTK)理论描述了无限宽神经网络在训练过程中的行为。关键洞察是:在NTK机制下,神经网络的行为类似于核方法,其中频率分布对模型性能至关重要。

2. RoPE的频率分布问题

回顾RoPE的角度公式:

$$\theta_i = \frac{1}{10000^{2i/d}}$$

对于维度索引 $i = 0, 1, 2, \ldots, d/2-1$:

  • 小的 $i$(低维度):$\theta_i$ 较大,对应低频(旋转慢)
  • 大的 $i$(高维度):$\theta_i$ 较小,对应高频(旋转快)

问题分析

当位置 $m$ 很大时,高频分量(小的 $\theta_i$)的旋转角度 $m \times \theta_i$ 可能会:

  • 超过 $2\pi$,导致周期重复
  • 使得不同位置的角度变得相似,丧失区分度

3. NTK-Aware的核心思想:降低高频,保留低频

关键观察:外推到长序列时,高频分量衰减过快,而低频分量仍然有效。

解决方案:通过增大 base(即减小 $\theta_i$ 的变化率),使得高频分量的频率降低,从而:

  • 防止高频分量过早衰减
  • 保持更长的位置区分度

4. 数学推导

原始RoPE

对于位置 $m$ 和维度 $i$,旋转角度为:

$$\phi_{m,i} = m \times \theta_i = m \times \frac{1}{10000^{2i/d}}$$

NTK-Aware调整

base 从 $10000$ 调整为 $10000 \times s$,其中 $s > 1$ 是缩放因子:

$$\theta’_i = \frac{1}{(10000 \times s)^{2i/d}} = \frac{1}{10000^{2i/d}} \times \frac{1}{s^{2i/d}} = \theta_i \times s^{-2i/d}$$

新的旋转角度:

$$\phi’_{m,i} = m \times \theta’_i = m \times \theta_i \times s^{-2i/d}$$

关键观察

  • 对于低维度(小的 $i$):$s^{-2i/d} \approx 1$,角度几乎不变(低频保留)
  • 对于高维度(大的 $i$):$s^{-2i/d} < 1$,角度减小(高频降低)

5. 动态缩放因子 $s$ 的设定

公式

$$s = \left(\frac{L}{L_{\text{train}}}\right)^\alpha$$

其中:

  • $L$ 是当前序列长度
  • $L_{\text{train}}$ 是训练时的最大长度
  • $\alpha$ 是超参数,通常设为 $\alpha = \frac{d}{d-2}$

推导逻辑

假设我们需要将上下文窗口从 $L_{\text{train}}$ 扩展到 $L$。

  • 如果 $L = L_{\text{train}}$,则 $s = 1$(不调整)
  • 如果 $L = 2L_{\text{train}}$,则 $s = 2^\alpha > 1$(需要调整)

$\alpha$ 的选择

根据NTK理论和实验,$\alpha = \frac{d}{d-2}$ 能够在不同头维度 $d$ 下提供较好的平衡。例如:

  • $d = 64$ 时,$\alpha = \frac{64}{62} \approx 1.03$
  • $d = 128$ 时,$\alpha = \frac{128}{126} \approx 1.016$

6. 频率分布调整的效果

原始频率分布base = 10000):

对于位置 $m = 8000$,不同维度的旋转角度:

  • $i = 0$:$\phi_{8000,0} = 8000 \times \frac{1}{10000^0} = 8000$(很大,可能超过 $2\pi$ 多次)
  • $i = 31$(假设 $d=64$):$\phi_{8000,31} = 8000 \times \frac{1}{10000^{62/64}} \approx 8000 \times 0.0001 = 0.8$

NTK-Aware调整后(假设 $s = 2$):

  • $i = 0$:$\phi’_{8000,0} = 8000 \times \frac{1}{20000^0} = 8000$(不变,低频保留)
  • $i = 31$:$\phi’_{8000,31} = 8000 \times \frac{1}{20000^{62/64}} \approx 8000 \times 0.00005 = 0.4$(降低,高频衰减变慢)

7. 为什么有效?

关键原理

  1. 低频分量(低维度):主要负责长距离依赖,需要保留
  2. 高频分量(高维度):主要负责短距离细节,但容易在长序列中失效

通过降低高频分量的频率(增大 base),我们:

  • 保持了低频分量的有效性(长距离信息)
  • 减缓了高频分量的衰减速度(短距离信息在更长序列中仍然有效)
  • 在长序列上保持了更好的位置区分度

8. 与PI的对比

  • PI:线性压缩所有位置到训练范围(简单但分辨率降低)
  • NTK-Aware:非线性的频率调整(更精细,保持更好的分辨率)

NTK-Aware通过调整频率分布而非位置索引,理论上能够更好地保持位置信息的质量。

优缺点

优点

  1. 无需重新训练:可以直接应用于预训练模型,提升长序列外推能力
  2. 计算效率高:只需修改频率基,几乎不增加额外计算成本
  3. 兼容性强:适用于所有基于RoPE的模型(如LLaMA、GPT-NeoX等)

缺点

  1. 超参数敏感:缩放因子 $s$ 和 $\alpha$ 需要根据具体任务调整
  2. 理论依赖:依赖于NTK假设,对于某些模型架构可能不成立

ALiBi (Attention with Linear Biases) 线性偏差注意力

简介

ALiBi(Attention with Linear Biases)是一种改进的Transformer位置编码方法,旨在解决传统位置编码(如绝对位置编码或旋转位置编码)在处理长序列外推时的局限性。

核心思想:在attention计算中引入一个线性偏置项,显式建模相对位置关系,而无需显式的位置嵌入。

关键思想

对于每个attention头,在计算query和key的点积后,添加一个线性偏置项。这个偏置项与相对位置成反比,距离越远的key会收到更大的负偏置,从而降低其attention权重。

外推能力:通过简单的线性衰减,模型在未见过的更长序列上表现更好。例如,在序列长度为512上训练的模型可以扩展到2048。

计算步骤

标准Attention计算

输入序列的attention分数矩阵公式为:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + B\right)V$$

其中:

  • $Q$ 是Query矩阵
  • $K$ 是Key矩阵
  • $V$ 是Value矩阵
  • $B$ 是ALiBi偏置矩阵
  • $d_k$ 是key的维度

1. 偏置矩阵 $B$ 的构造

对于每个query位置 $i$ 和key位置 $j$,相对位置偏移定义为 $|i - j|$。偏置项 $B_{i,j}$ 为:

$$B_{i,j} = -m \times |i - j|$$

其中 $m$ 是与attention头相关的斜率系数,不同头使用不同的 $m$ 值。

2. 斜率 $m$ 的设定

对于 $n$ 个attention头,斜率按几何级数分布。

示例:对于8个头,$m$ 值可能为:
$$m = \left{\frac{1}{2^1}, \frac{1}{2^2}, \ldots, \frac{1}{2^8}\right}$$

公式化表示:对于第 $k$ 个头,斜率 $m_k$ 为:
$$m_k = \frac{1}{2^{8/k}}$$

3. 计算流程

  1. 计算标准点积attention分数:$\frac{QK^T}{\sqrt{d_k}}$
  2. 生成偏置矩阵 $B$ 并加到attention分数上
  3. 应用Softmax和加权求和得到输出
ALiBi原理详解:为什么线性偏置能够实现外推?

1. 传统位置编码的问题

传统的绝对位置编码(如Sinusoidal PE或学习的PE)在训练时只能学习到固定长度范围内的位置信息。当序列长度超出训练范围时:

  • 新的位置编码是未见的,模型无法正确理解
  • 位置编码的分布可能发生变化,导致性能下降

2. ALiBi的核心创新:不使用位置嵌入,使用位置偏置

关键思想:不在输入中添加位置嵌入,而是在attention计算中直接引入相对位置偏置

标准Attention计算

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

ALiBi的Attention计算

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + B\right)V$$

其中 $B$ 是偏置矩阵,$B_{i,j} = -m \times |i - j|$。

3. 线性偏置的作用机制

偏置项的定义

对于query位置 $i$ 和key位置 $j$:

$$B_{i,j} = -m \times |i - j|$$

其中 $m > 0$ 是斜率系数。

关键性质

  • $B_{i,j} \leq 0$(总是非正)
  • $|i - j|$ 越大,$B_{i,j}$ 越负(距离越远,偏置越大)
  • 偏置是线性的(与距离成正比)

4. Softmax中的偏置效果

在Softmax中,偏置项会直接影响注意力权重:

$$\text{softmax}(x_i + b_i) = \frac{e^{x_i + b_i}}{\sum_j e^{x_j + b_j}} = \frac{e^{x_i} \cdot e^{b_i}}{\sum_j e^{x_j} \cdot e^{b_j}}$$

由于 $e^{b_i} < 1$(因为 $b_i < 0$),偏置项相当于对原始分数进行衰减

具体效果

对于query位置 $i$ 和key位置 $j$,attention分数为:

$$s_{i,j} = \frac{q_i^T k_j}{\sqrt{d_k}} - m \times |i - j|$$

经过Softmax后,注意力权重为:

$$w_{i,j} = \frac{\exp(s_{i,j})}{\sum_k \exp(s_{i,k})} = \frac{\exp\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right) \cdot \exp(-m \times |i - j|)}{\sum_k \exp\left(\frac{q_i^T k_k}{\sqrt{d_k}}\right) \cdot \exp(-m \times |i - k|)}$$

关键观察

  • 距离越远($|i - j|$ 越大),$\exp(-m \times |i - j|)$ 越小
  • 这意味着远距离的key会收到指数衰减的权重
  • 但偏置本身是线性的(与距离成正比),衰减模式简单且可预测

5. 为什么线性偏置能实现外推?

关键优势:不依赖绝对位置,只依赖相对距离

传统位置编码的问题:

  • 位置 $1000$ 的编码在训练时可能未见过
  • 模型无法理解这个新位置

ALiBi的优势:

  • 偏置项 $B_{i,j} = -m \times |i - j|$ 只依赖于相对距离 $|i - j|$
  • 即使位置 $i$ 和 $j$ 都超出了训练范围,只要它们的距离 $|i - j|$ 在训练时见过,偏置项就是有意义的
  • 例如,如果训练时见过距离 $1, 2, \ldots, 512$,那么测试时的距离 $513, 514, \ldots$ 可以自然地通过线性外推得到

数学证明

假设训练时最大序列长度为 $L_{\text{train}}$,那么训练时见过的相对距离为 $d \in {0, 1, 2, \ldots, L_{\text{train}}-1}$。

对于任意距离 $d$,偏置项为 $B_d = -m \times d$,这是一个线性函数

当测试时遇到新距离 $d’ > L_{\text{train}}-1$ 时:

$$B_{d’} = -m \times d’$$

这个公式在训练范围内和范围外形式完全相同,因此可以自然地外推到任意距离。

6. 多头的斜率设计

为什么不同头使用不同斜率?

ALiBi为不同的attention头分配不同的斜率 $m$,通常按几何级数分布:

对于 $n$ 个头:$m_k = \frac{1}{2^{k}}$,其中 $k = 1, 2, \ldots, n$

设计原理

  • 不同头关注不同距离范围

    • 小斜率(如 $m = 1/256$):衰减慢,允许关注远距离
    • 大斜率(如 $m = 1/2$):衰减快,主要关注近距离
  • 多样性:不同头使用不同斜率,增加了模型的表达能力和鲁棒性

7. 与RoPE和传统PE的对比

传统绝对位置编码

  • 位置 $p$ 的编码是固定的向量 $PE(p)$
  • 超出训练范围的位置编码是未见的
  • 外推能力差

RoPE

  • 通过旋转矩阵编码位置
  • 相对位置信息通过角度差体现
  • 有一定外推能力,但在极长序列上可能失效

ALiBi

  • 不使用位置嵌入,只使用线性偏置
  • 偏置只依赖于相对距离,与绝对位置无关
  • 外推能力最强:线性函数可以自然扩展到任意距离

8. 实际效果

实验表明,在序列长度为 $512$ 上训练的模型,使用ALiBi可以扩展到 $2048$ 甚至更长,性能下降很小。这证明了线性偏置方法的有效性。

ALiBi的特点

ALiBi不需要位置嵌入,而是通过简单的线性偏置直接建模相对位置关系。这种方法在长序列外推方面表现优异,因为线性衰减模式可以自然地扩展到更长的序列。


YARN (Yet Another RoPE Extension)

简介

YARN(Yet Another RoPE Extension)是一种基于改进的相对位置编码(RoPE)机制来扩展上下文长度的方法。原始的RoPE机制在Transformer模型中用于捕捉序列内的位置关系,YARN通过动态扩展位置编码能力来优化这一机制,以支持处理更长的上下文。

YARN原理

1. 相对位置编码扩展:RoPE编码器原本设计用于固定长度的相对位置信息。YARN调整并扩展RoPE的编码方法,以适应更长的序列长度。

2. 分段编码:将长序列划分为多个较小的段。每个段进行独立的RoPE编码,同时保持这些段之间的位置信息关系。

3. 全局与局部结合:在编码过程中,YARN同时考虑序列的局部和全局信息。这确保了模型在扩展上下文长度的同时,能够有效捕捉长距离依赖。

YARN原理详解:如何通过改进RoPE实现上下文窗口扩展?

1. YARN的核心思想

YARN(Yet Another RoPE Extension)是在RoPE基础上的改进,旨在解决RoPE在极长序列上的外推问题。与PI和NTK-Aware不同,YARN采用了更复杂的策略,结合了分段编码频率调整

2. RoPE的外推问题回顾

标准RoPE的角度公式:

$$\theta_i = \frac{1}{10000^{2i/d}}$$

对于位置 $m$,旋转角度为 $m \times \theta_i$。

问题:当 $m$ 非常大时:

  • 高频分量(大的 $i$,小的 $\theta_i$)的旋转角度 $m \times \theta_i$ 可能超过 $2\pi$ 多次
  • 导致位置编码的周期性重复,丧失区分度
  • 不同位置的编码变得相似

3. YARN的解决方案:分段频率调整

YARN的核心思想是:对于不同距离范围,使用不同的频率调整策略

关键公式

将序列分为多个”段”(segment),对于每个段内的位置,使用调整后的频率:

对于位置 $m$ 和维度 $i$,调整后的角度为:

$$\theta’{m,i} = \begin{cases}
m \times \theta_i & \text{if } m \leq L
{\text{train}} \\
L_{\text{train}} \times \theta_i + (m - L_{\text{train}}) \times \theta_i \times \alpha & \text{if } m > L_{\text{train}}
\end{cases}$$

其中 $\alpha < 1$ 是频率衰减因子。

4. 分段编码策略

策略1:固定分段

将长序列划分为固定长度的段(如每段 $L_{\text{train}}$ 长度):

  • 段内:使用标准RoPE编码
  • 段间:通过额外的偏置或缩放因子保持相对位置关系

策略2:频率分段衰减

对于超出训练长度的位置,使用衰减的频率:

如果 $m > L_{\text{train}}$,定义缩放因子:

$$s(m) = 1 + \alpha \times \frac{m - L_{\text{train}}}{L_{\text{train}}}$$

调整后的角度:

$$\theta’_{m,i} = \frac{m \times \theta_i}{s(m)}$$

5. 数学推导:为什么分段有效?

关键观察:RoPE的位置编码主要依赖于相对位置差,而非绝对位置。

对于两个位置 $m$ 和 $n$($m < n$),它们的相对角度差为:

$$\Delta\theta = (n - m) \times \theta_i$$

在训练范围内($m, n \leq L_{\text{train}}$):

相对角度差 $\Delta\theta$ 是模型学习过的,模型能够正确理解。

超出训练范围($m, n > L_{\text{train}}$):

如果直接使用原始 $\theta_i$,角度差可能过大。YARN通过频率调整,使得超出范围的角度差与训练范围内的角度差更相似。

6. 频率调整的具体实现

YARN通常使用幂律衰减的频率调整:

对于位置 $m$,定义缩放因子:

$$s(m) = \begin{cases}
1 & \text{if } m \leq L_{\text{train}} \\
\left(\frac{m}{L_{\text{train}}}\right)^\beta & \text{if } m > L_{\text{train}}
\end{cases}$$

其中 $\beta$ 是衰减指数,通常 $0 < \beta < 1$(如 $\beta = 0.5$)。

调整后的角度:

$$\theta’_{m,i} = \frac{m \times \theta_i}{s(m)}$$

效果分析

  • 当 $m = L_{\text{train}}$ 时,$s(m) = 1$,角度不变
  • 当 $m = 2L_{\text{train}}$ 时,$s(m) = 2^\beta > 1$,角度减小,频率降低
  • 当 $m \to \infty$ 时,$s(m) \to \infty$,角度趋于 $0$,频率趋于 $0$

7. 与PI和NTK-Aware的对比

PI(Position Interpolation)

  • 策略:线性压缩所有位置到训练范围
  • 优点:简单直接
  • 缺点:分辨率降低,所有位置都受影响

NTK-Aware Scaled RoPE

  • 策略:全局调整频率基
  • 优点:理论支撑,计算简单
  • 缺点:对所有位置使用相同的缩放策略

YARN

  • 策略:分段或位置相关的频率调整
  • 优点:更精细的控制,训练范围内的位置不受影响
  • 缺点:实现复杂,需要额外的参数和计算

8. 实际应用中的YARN变体

YARN有多个变体,常见的包括:

YARN (base):基础的频率调整

YARN (finetune):结合微调的版本,在扩展后的数据上进行进一步训练

YARN (mirostat):结合mirostat采样策略的版本

关键设计选择

  • 衰减指数 $\beta$ 的选择:通常在 $0.1$ 到 $0.5$ 之间
  • 是否对所有维度使用相同的衰减:可以维度相关的衰减
  • 分段策略:固定分段 vs. 自适应分段

优缺点

优点

  1. 计算效率高:通过使用基于RoPE的相对位置编码,YARN避免了传统全attention机制的 $O(n^2)$ 计算复杂度,显著减少计算量
  2. 支持长上下文:YARN优化了位置编码的生成,使长序列的高效处理成为可能,支持数万到数十万的上下文长度
  3. 与现有模型兼容:YARN可以通过修改原始RoPE编码层集成到现有模型架构中,兼容性好

缺点

  1. 实现复杂:YARN需要对RoPE机制进行修改,实现相对复杂的编码扩展逻辑,增加了代码复杂度和维护难度
  2. 实验验证不足:目前关于YARN的公开信息有限,其具体性能和效果仍需要进一步的实验验证

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch
import torch.nn as nn
import math

class YARN(nn.Module):
def __init__(self, embed_dim, max_relative_length=1024):
super(YARN, self).__init__()
self.embed_dim = embed_dim
self.max_relative_length = max_relative_length

# 扩展RoPE参数
inv_freq = 1.0 / (10000 ** (torch.arange(0, embed_dim, 2).float() / embed_dim))
self.register_buffer('inv_freq', inv_freq)

def forward(self, x, attention_mask=None):
"""
前向传播
:param x: 输入序列,形状为 [batch_size, seq_len, embed_dim]
:param attention_mask: 注意力掩码
:return: 处理后的序列,形状为[batch_size, seq_len, embed_dim]
"""
batch_size, seq_len = x.size()[:2]

# 扩展相对位置编码
pos = torch.arange(0, seq_len, device=x.device).float()
sin = torch.sin(pos[:, None] * self.inv_freq[None, :])
cos = torch.cos(pos[:, None] * self.inv_freq[None, :])

# 应用RoPE变换
# ... (具体实现)

return x

6. 如何提升外推能力

提升大模型外推能力的方法可以归纳为三个方面:

位置编码外推

问题:Transformer的绝对位置编码对于超长序列失效。

解决方案

  • ALiBi:通过在attention计算中添加线性偏置项,显式建模相对位置关系
  • RoPE扩展方法:通过PI、NTK-Aware、YARN等方法扩展RoPE的适用范围

长度泛化技术

NTK-Aware Scaling:动态调整RoPE的旋转角度,平衡高频和低频位置信息的分布。

关键公式
$$\theta_j = \frac{10000^{-2j/d}}{\text{scaling_factor}}$$

推理策略增强

  1. Chain-of-Thought (CoT):通过思维链引导模型进行逐步推理
  2. Self-Consistency:采用多路径推理投票来选择最优解决方案

7. 应用场景

超长文本生成

输入:10k tokens的文档
挑战:传统Transformer在4k tokens之后质量下降
方案:使用YARN(Yet Another RoPE Extension)位置编码

数值推理外推

训练数据:数字范围在0-1000的数学问题
测试数据:数字范围在1000-10000的数值计算
表现:GPT-4在加法任务上达到>95%的准确率(在 $10^{18}$ 范围内)

多模态理解

场景:处理文本、图像和视频的多模态数据
方案:使用M-RoPE(多维旋转位置编码),能够同时捕捉1D文本序列、2D视觉图像和3D视频的位置信息


总结

位置编码是Transformer模型的核心组件之一,其发展经历了从绝对位置编码到相对位置编码,再到多维旋转位置编码的演进。随着大模型对长序列处理需求的不断增长,位置编码的外推能力变得越来越重要。

  • 绝对位置编码(Sinusoidal):虽然简单有效,但缺乏外推能力
  • 旋转位置编码(RoPE):通过旋转变换实现相对位置编码,在长序列处理方面表现优异
  • 多维RoPE(M-RoPE):扩展了RoPE到多模态场景,支持图像和视频的位置编码
  • 扩展方法(PI、NTK-Aware、ALiBi、YARN):通过不同的技术手段提升模型的外推能力

这些技术的发展使得大模型能够处理越来越长的序列,为更复杂的任务和应用场景提供了支持。