简介
因为RoPE的优异性能，其已成为各种大模型中位置编码的首选，包括多模态模型；在一些多模态模型或视频理解模型中，甚至会用到多维度RoPE。虽然RoPE已广泛应用，之前也看了不少针对其原理解析的文章，特别是苏神的推导帖，可能是我个人能力问题，也可能是太过于追求原理的问题，对RoPE一直没有理解透彻，总感觉差那么一点。本文会适当弱化RoPE为什么会有效的原理推导，将重点放在RoPE应该如何实现以及具体的实现步骤和一些不同实现方式的解析。
RoPE的核心思想是通过旋转操作将位置信息融入到词向量。在复数域中，一个复数由实部和虚部组成，故一个复数可以看作复数域中的向量。将一个复数与复数
e
i
θ
=
cos
θ
+
i
sin
θ
e^{i\theta}=\cos{\theta}+i \sin{\theta}
eiθ=cosθ+isinθ相乘时，相当于在复平面上绕原点旋转了
θ
\theta
θ度角。假设有一个复数
v
=
x
+
i
y
v=x+iy
v=x+iy，按上述说明，将
v
v
v与
e
i
θ
e^{i\theta}
eiθ相乘就是将
v
v
v绕原点旋转了
θ
\theta
θ度角，可以通过公式推导验证：
v
′
=
v
∗
e
i
θ
=
(
x
+
i
y
)
∗
(
cos
θ
+
i
sin
θ
)
=
cos
θ
x
+
i
sin
θ
x
+
i
cos
θ
y
−
sin
θ
y
=
(
cos
θ
x
−
sin
θ
y
)
+
i
(
sin
θ
x
+
cos
θ
y
)
\begin{align*} v'&=v*e^{i\theta} \\ &=(x+iy)*(\cos{\theta}+i \sin{\theta}) \\ &=\cos{\theta}x+i\sin{\theta}x+i\cos{\theta}y-\sin{\theta}y \\ &=(\cos{\theta}x-\sin{\theta}y)+i(\sin{\theta}x+\cos{\theta}y) \tag{1} \end{align*}
v′=v∗eiθ=(x+iy)∗(cosθ+isinθ)=cosθx+isinθx+icosθy−sinθy=(cosθx−sinθy)+i(sinθx+cosθy)(1)
通过公式(1)有
x
′
=
cos
θ
x
−
sin
θ
y
,
y
′
=
sin
θ
x
+
cos
θ
y
x'=\cos{\theta}x-\sin{\theta}y,y'=\sin{\theta}x+\cos{\theta}y
x′=cosθx−sinθy,y′=sinθx+cosθy。从二维向量的角度看，变换前后的复数
v
,
v
′
v,v'
v,v′其实均可以看作列向量
[
x
y
]
，
[
x
′
y
′
]
\begin{bmatrix} x \\ y \end{bmatrix}，\begin{bmatrix} x' \\ y' \end{bmatrix}
[xy]，[x′y′]，公式(1)表示
v
′
v'
v′是由
v
v
v旋转而来，且恰好对应了二维向量空间的旋转矩阵，如下所示：
[
x
′
y
′
]
=
[
cos
θ
−
sin
θ
sin
θ
cos
θ
]
[
x
y
]
(2)
\begin{equation} \begin{bmatrix} x' \\ y' \end{bmatrix}=\begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix} \end{equation} \tag2
[x′y′]=[cosθsinθ−sinθcosθ][xy](2)
如公式(2)所示，在二维向量空间中，将一个向量绕原点旋转 θ \theta θ度就是乘以一个旋转矩阵，即 R ( θ ) = [ cos θ − sin θ sin θ cos θ ] R(\theta)=\begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} R(θ)=[cosθsinθ−sinθcosθ]。通过上述推导可知，二维向量空间中的旋转和复数乘法在进行旋转时是等价的。
RoPE是将隐向量序列中不同位置的token的特征向量与不同的旋转角对应的旋转矩阵相乘，即将不同位置token的特征向量旋转不同的角度。注意力计算的主要部分是token向量之间的点积运算，可以通过推导展示RoPE是如何在自注意力计算中引入相对位置信息的。
假设目前token的维度为2，有序列中m处的查询向量
q
m
=
[
q
m
,
1
q
m
,
2
]
q_m=\begin{bmatrix} q_{m,1} \\ q_{m,2} \end{bmatrix}
qm=[qm,1qm,2]和n处的键向量
k
n
=
[
k
n
,
1
k
n
,
2
]
k_n=\begin{bmatrix} k_{n,1} \\ k_{n,2} \end{bmatrix}
kn=[kn,1kn,2]，设旋转基础角度为
θ
\theta
θ，序列中不同位置token的旋转角由位置索引与基础角度相乘得到，即上述两个向量的旋转角度为别是
m
θ
,
n
θ
m\theta,n\theta
mθ,nθ，RoPE就是给不同位置的token向量与对应位置的旋转角的旋转矩阵相乘，基于前文的推理，有，
q
m
′
=
[
cos
(
m
θ
)
−
sin
(
m
θ
)
sin
(
m
θ
)
cos
(
m
θ
)
]
[
q
m
,
1
q
m
,
2
]
=
R
(
m
θ
)
q
m
(3)
\begin{equation} q'_m=\begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} q_{m,1} \\ q_{m,2} \end{bmatrix} = R(m\theta) q_m \end{equation} \tag3
qm′=[cos(mθ)sin(mθ)−sin(mθ)cos(mθ)][qm,1qm,2]=R(mθ)qm(3)
k
n
′
=
[
cos
(
n
θ
)
−
sin
(
n
θ
)
sin
(
n
θ
)
cos
(
n
θ
)
]
[
k
n
,
1
k
n
,
2
]
=
R
(
n
θ
)
k
m
(4)
\begin{equation} k'_n=\begin{bmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{bmatrix} \begin{bmatrix} k_{n,1} \\ k_{n,2} \end{bmatrix} = R(n\theta) k_m \end{equation} \tag4
kn′=[cos(nθ)sin(nθ)−sin(nθ)cos(nθ)][kn,1kn,2]=R(nθ)km(4)
旋转后的两个向量进行以下点积计算：
q
m
′
T
k
n
′
=
(
R
(
m
θ
)
q
m
)
T
⋅
R
(
n
θ
)
k
n
=
q
m
T
⋅
R
(
m
θ
)
T
⋅
R
(
n
θ
)
⋅
k
n
=
q
m
T
⋅
[
cos
(
m
θ
)
sin
(
m
θ
)
−
sin
(
m
θ
)
cos
(
m
θ
)
]
[
cos
(
n
θ
)
−
sin
(
n
θ
)
sin
(
n
θ
)
cos
(
n
θ
)
]
⋅
k
n
=
q
m
T
⋅
[
cos
(
m
θ
)
cos
(
n
θ
)
+
sin
(
m
θ
)
sin
(
n
θ
)
−
cos
(
m
θ
)
sin
(
n
θ
)
+
sin
(
m
θ
)
cos
(
n
θ
)
−
sin
(
m
θ
)
cos
(
n
θ
)
+
cos
(
m
θ
)
sin
(
n
θ
)
sin
(
m
θ
)
sin
(
n
θ
)
+
cos
(
m
θ
)
cos
(
n
θ
)
]
⋅
k
n
=
q
m
T
⋅
[
cos
(
(
n
−
m
)
θ
)
−
sin
(
(
n
−
m
)
θ
)
sin
(
(
n
−
m
)
θ
)
cos
(
(
n
−
m
)
θ
)
]
⋅
k
n
=
q
m
T
⋅
R
(
(
n
−
m
)
θ
)
⋅
k
n
\begin{align*} q_m^{'T}k'_n &= (R(m\theta) q_m)^T \cdot R(n\theta) k_n \\ &= q^T_m \cdot R(m\theta)^T \cdot R(n\theta) \cdot k_n \\ &= q^T_m \cdot \begin{bmatrix} \cos(m\theta) & \sin(m\theta) \\ -\sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{bmatrix} \cdot k_n \\ &= q^T_m \cdot \begin{bmatrix} \cos(m\theta)\cos(n\theta) + \sin(m\theta)\sin(n\theta) & -\cos(m\theta)\sin(n\theta) + \sin(m\theta)\cos(n\theta) \\ -\sin(m\theta)\cos(n\theta) + \cos(m\theta)\sin(n\theta) & \sin(m\theta)\sin(n\theta) + \cos(m\theta)\cos(n\theta) \end{bmatrix} \cdot k_n \\ &= q^T_m \cdot \begin{bmatrix} \cos((n - m)\theta) & -\sin((n - m)\theta) \\ \sin((n - m)\theta) & \cos((n - m)\theta) \end{bmatrix} \cdot k_n \\ &= q^T_m \cdot R((n - m)\theta) \cdot k_n \tag5 \end{align*}
qm′Tkn′=(R(mθ)qm)T⋅R(nθ)kn=qmT⋅R(mθ)T⋅R(nθ)⋅kn=qmT⋅[cos(mθ)−sin(mθ)sin(mθ)cos(mθ)][cos(nθ)sin(nθ)−sin(nθ)cos(nθ)]⋅kn=qmT⋅[cos(mθ)cos(nθ)+sin(mθ)sin(nθ)−sin(mθ)cos(nθ)+cos(mθ)sin(nθ)−cos(mθ)sin(nθ)+sin(mθ)cos(nθ)sin(mθ)sin(nθ)+cos(mθ)cos(nθ)]⋅kn=qmT⋅[cos((n−m)θ)sin((n−m)θ)−sin((n−m)θ)cos((n−m)θ)]⋅kn=qmT⋅R((n−m)θ)⋅kn(5)
上述公式(5)推导过程要使用高中数学知识–三角函数和差定理。可以看到，两个旋转不同角度后向量的点积与它们之间旋转的角度之差有关系，是一种相对位置的体现，即体现了RoPE给序列中不同位置token引入了相对位置信息。
注意：上述旋转均是逆时针旋转！
拆解
上述只是对RoPE的原理进行了简短解释，尽可能直白地说明RoPE原理，但上述讨论局限在二维空间内，而LLMs模型中token向量的维度都很大，接下来对常规的RoPE实现进行拆解，将整个实现过程解剖出来。
- 为了沿用二维旋转矩阵的性质，RoPE会将一个token的高维向量看作多个复数，假设向量特征维度为 d d d(基本所有LLMs中的特征维度数是偶数)，那么会将一个 d d d维的token特征向量视为 d / 2 d/2 d/2个复数
- 假设一个token向量为 [ 1 , 2 , 3 , 4 ] [1,2,3,4] [1,2,3,4]，那么 [ 1 , 2 ] [1,2] [1,2]是第一个复数， [ 3 , 4 ] [3,4] [3,4]是第二个复数
- 同一个token中的不同复数对的旋转角度不同
- 基础旋转角
θ
i
=
1000
0
−
2
i
/
d
=
1
1000
0
i
d
2
\theta_i=10000^{-2i/d}=\frac{1}{10000^{\frac{i}{\frac{d}{2}}}}
θi=10000−2i/d=100002di1；注意，同一token向量中不同位置复数的旋转角是由
i
i
i来决定的。上文已经说到，一个token向量会分成
d
/
2
d/2
d/2个复数，
i
∈
[
0
,
d
/
2
)
i \in [0,d/2)
i∈[0,d/2)，所以
θ
i
\theta_i
θi就对应第
i
i
i对复数的旋转角
- 10000是一个超参数
- 同一序列中不同位置的token向量的旋转角也不同，就是基础旋转角与token对应的序列索引乘积，即一个token在序列中的索引为m，那么其对应的最终的旋转角为 θ i ⋅ m \theta_i \cdot m θi⋅m。
- 实现时一般会先将 θ i \theta_i θi计算出来，再基于序列长度获取不同位置token索引的一维向量，然后两者相乘，得到不同位置token中不同复数维度的准确旋转角度
- 将输入沿着特征维度，按顺序每两个为一组作为一对复数进行拆分，向量维度从[batch_size, seq_len, dim]–>[batch_size, seq_len, dim//2, 2]
- 最后一个维度中，第一组向量为所有复数的实部，第二组向量为所有复数的虚部
- 按照二维矩阵中的计算方式，进行旋转矩阵乘法，然后将最终的实部和虚部重新合并为原来的维度[batch_size, seq_len, dim]，即完成了RoPE的应用
上述对RoPE的具体操作过程进行了叙述，以下是一种较原始的朴素实现：
import torch
def get_rotary_matrix(seq_len, dim, base=10000):
"""生成RoPE的旋转矩阵"""
# 生成不同频率的正弦和余弦值
theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) # shape为[dim//2]
# 生成位置索引
position = torch.arange(seq_len).float() # shape为[seq_len]
# 计算每个位置和维度对应的角度
theta = torch.outer(position, theta) # 计算外积，其中第(i, j)个元素是position[i] * theta[j]；shape为[seq_len, dim//2]
# 计算正弦和余弦值
cos = torch.cos(theta) # shape为[seq_len, dim//2]
sin = torch.sin(theta) # shape为[seq_len, dim//2]
return cos, sin
def apply_rotary_embedding(x, cos, sin):
"""应用旋转位置编码"""
# 假设x的形状为[batch_size, seq_len, dim]
# 将向量视为复数，每两个维度一组
x_reshape = x.view(*x.shape[:-1], -1, 2) # shape为[batch_size, seq_len, dim//2, 2]，即沿着特征维度拆分
# 构建正弦和余弦矩阵，使其与x_reshape形状匹配
cos_expanded = cos.view(1, cos.shape[0], cos.shape[1], 1) # shape为[1, seq_len, dim//2, 1]
sin_expanded = sin.view(1, sin.shape[0], sin.shape[1], 1) # shape为[1, seq_len, dim//2, 1]
# 旋转操作（复数乘法）
# [x_real, x_imag] * (cos + i*sin) = [x_real*cos - x_imag*sin, x_real*sin + x_imag*cos]
x_out_1 = x_reshape[:, :, :, 0:1] * cos_expanded - x_reshape[:, :, :, 1:2] * sin_expanded
x_out_2 = x_reshape[:, :, :, 0:1] * sin_expanded + x_reshape[:, :, :, 1:2] * cos_expanded
# 合并结果
x_out = torch.cat([x_out_1, x_out_2], dim=-1) # shape为[batch_size, seq_len, dim//2, 2]
return x_out.view(*x.shape)
# 示例用法
def apply_rope(x):
"""对输入向量应用RoPE位置编码"""
batch_size, seq_len, dim = x.shape
cos, sin = get_rotary_matrix(seq_len, dim)
return apply_rotary_embedding(x, cos, sin)
# 测试代码
if __name__ == "__main__":
# 创建一个随机输入张量
batch_size, seq_len, dim = 2, 10, 512
x = torch.randn(batch_size, seq_len, dim)
# 应用RoPE
x_with_rope = apply_rope(x)
print(f"输入形状: {x.shape}")
print(f"输出形状: {x_with_rope.shape}")
一些tricks
- 在对输入的嵌入维度上应用RoPE时，当特征维度很长，或者希望节约资源资源时，可以不对所有维度全部应用，而是可以设置一个比例，即只在前 d r o p e d_{rope} drope维度上应用RoPE，后面保持原始数值不变
- 在计算 θ i \theta_i θi时原始论文使用的base为10000，但可以对其进行缩放，实现长度外推；如增加base，可以处理训练过程中未见过的长序列
扫一扫
1905
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
