RoPE

Starx_Smile
0 阅读3分钟

RoPE (Rotary Position Embedding) 不是把“位置向量”加到输入上，而是把位置 pp 映射成一个旋转矩阵，作用到每个位置的 QQ 和 KK 上；这样注意力分数

(RoPE(qm,m))⊤(RoPE(kn,n))\big(\mathrm{RoPE}(q_m,m)\big)^\top \big(\mathrm{RoPE}(k_n,n)\big)

会自然只依赖于相对位置差 n−mn-m，从而把相对位置信息直接注入到 self-attention 中。

1. 先回顾 Transformer 注意力公式

单头自注意力里，给定输入 X∈RL×dmodelX\in\mathbb{R}^{L\times d_{\text{model}}}：

Q=XWQ,K=XWK,V=XWVQ = XW_Q,\qquad K = XW_K,\qquad V = XW_V

其中 Q,K,V∈RL×dhQ,K,V\in\mathbb{R}^{L\times d_h}（单头时）。 第 mm 个位置和第 nn 个位置的注意力打分是：

sm,n=qm⊤kndhs_{m,n}=\frac{q_m^\top k_n}{\sqrt{d_h}}

softmax 后得到权重：

am,n=exp⁡(sm,n)∑j=1Lexp⁡(sm,j)a_{m,n}=\frac{\exp(s_{m,n})}{\sum_{j=1}^{L}\exp(s_{m,j})}

输出：

om=∑n=1Lam,nvno_m=\sum_{n=1}^{L} a_{m,n} v_n

问题在于：如果不加位置编码，qm⊤knq_m^\top k_n 只反映“内容相似性”，不反映“位置关系”。

2. RoPE 的基本想法：对 QQ 和 KK 做“位置相关旋转”

2.1 从二维向量的旋转开始

对二维向量 x=(x1,x2)⊤x=(x_1,x_2)^\top，旋转角 θ\theta 的矩阵是：

R(θ)=[cos⁡θ−sin⁡θsin⁡θcos⁡θ]R(\theta)= \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}

旋转后：

R(θ)x=[x1cos⁡θ−x2sin⁡θx1sin⁡θ+x2cos⁡θ]R(\theta)x= \begin{bmatrix} x_1\cos\theta - x_2\sin\theta \\ x_1\sin\theta + x_2\cos\theta \end{bmatrix}

RoPE 的做法是： 对于位置 pp，不用“加一个位置向量”，而是对向量做一个与 pp 相关的旋转。

3. 关键公式：为什么它能得到“相对位置”？

设第 mm 个位置的 query 是 qmq_m，第 nn 个位置的 key 是 knk_n。 对它们分别做位置旋转：

q~m=R(mθ) qm,k~n=R(nθ) kn\tilde q_m = R(m\theta)\, q_m,\qquad \tilde k_n = R(n\theta)\, k_n

则它们的内积为：

q~m⊤k~n=(R(mθ)qm)⊤(R(nθ)kn)\tilde q_m^\top \tilde k_n = \big(R(m\theta) q_m\big)^\top \big(R(n\theta) k_n\big)

利用矩阵性质 (Ax)⊤y=x⊤A⊤y(Ax)^\top y = x^\top A^\top y，有：

q~m⊤k~n=qm⊤R(mθ)⊤R(nθ)kn\tilde q_m^\top \tilde k_n = q_m^\top R(m\theta)^\top R(n\theta) k_n

又因为旋转矩阵满足：

R(α)⊤=R(−α),R(α)⊤R(β)=R(β−α)R(\alpha)^\top = R(-\alpha),\qquad R(\alpha)^\top R(\beta)=R(\beta-\alpha)

所以：

q~m⊤k~n=qm⊤R((n−m)θ) kn\tilde q_m^\top \tilde k_n = q_m^\top R((n-m)\theta)\, k_n

这一步是 RoPE 的核心。 它说明：注意力分数不再只看内容 qm,knq_m,k_n，还会显式依赖于相对位移 n−mn-m，而不是绝对位置 m,nm,n 分别是多少。

4. 推广到高维：把向量按两维一组分块旋转

Transformer 的每个 head 维度一般是 dhd_h，通常为偶数。 RoPE 将 dhd_h 维向量分成 dh/2d_h/2 个二维子空间：

x=[x0,x1,x2,x3,…,xdh−2,xdh−1]x = [x_0, x_1, x_2, x_3, \dots, x_{d_h-2}, x_{d_h-1}]

分成：

(x0,x1), (x2,x3), …, (xdh−2,xdh−1)(x_0,x_1),\ (x_2,x_3),\ \dots,\ (x_{d_h-2},x_{d_h-1})

对于第 ii 个二维对子，使用一个不同的频率 ωi\omega_i。 在位置 pp 上，该对子旋转角度为：

θp,i=p⋅ωi\theta_{p,i}=p\cdot \omega_i

所以第 ii 对的旋转矩阵是：

R(θp,i)=[cos⁡(pωi)−sin⁡(pωi)sin⁡(pωi)cos⁡(pωi)]R(\theta_{p,i}) = \begin{bmatrix} \cos(p\omega_i) & -\sin(p\omega_i) \\ \sin(p\omega_i) & \cos(p\omega_i) \end{bmatrix}

5. RoPE 的完整分量公式

设某个向量 x∈Rdhx\in\mathbb{R}^{d_h}，位置为 pp。 RoPE 作用后的结果记为 RoPE(x,p)\mathrm{RoPE}(x,p)。 对每个 i=0,1,…,dh2−1i=0,1,\dots,\frac{d_h}{2}-1，有：

x2i′=x2icos⁡(pωi)−x2i+1sin⁡(pωi),x2i+1′=x2isin⁡(pωi)+x2i+1cos⁡(pωi).\begin{aligned} x'_{2i} &= x_{2i}\cos(p\omega_i) - x_{2i+1}\sin(p\omega_i), \\ x'_{2i+1} &= x_{2i}\sin(p\omega_i) + x_{2i+1}\cos(p\omega_i). \end{aligned}

于是：

RoPE(x,p)=[x0′,x1′,…,xdh−1′]\mathrm{RoPE}(x,p) = [x'_0,x'_1,\dots,x'_{d_h-1}]

频率 ωi\omega_i 的常见定义

通常设：

ωi=θi=10000− 2idh\omega_i = \theta_i = 10000^{-\,\frac{2i}{d_h}}

所以旋转角度是：

pθi=p⋅10000− 2idhp\theta_i = p \cdot 10000^{-\,\frac{2i}{d_h}}

这和经典 sinusoidal PE 的频率设计是一致的： 低维对子频率高，高维对子频率低，形成多尺度位置表示。

6. 用块对角矩阵写得更紧凑

把整个 dhd_h 维的 RoPE 写成一个块对角矩阵：

Rp=diag(R(pω0), R(pω1), …, R(pωdh2−1))\mathcal{R}_p = \mathrm{diag}\Big( R(p\omega_0),\, R(p\omega_1),\, \dots,\, R(p\omega_{\frac{d_h}{2}-1}) \Big)

则：

RoPE(x,p)=Rpx\mathrm{RoPE}(x,p)=\mathcal{R}_p x

对于 query/key：

q~m=Rmqm,k~n=Rnkn\tilde q_m = \mathcal{R}_m q_m,\qquad \tilde k_n = \mathcal{R}_n k_n

注意力打分变成：

sm,n=q~m⊤k~ndh=qm⊤Rm⊤Rnkndhs_{m,n} = \frac{\tilde q_m^\top \tilde k_n}{\sqrt{d_h}} = \frac{q_m^\top \mathcal{R}_m^\top \mathcal{R}_n k_n}{\sqrt{d_h}}

由于每个二维块都满足旋转矩阵的差角性质，因此：

Rm⊤Rn=Rn−m\mathcal{R}_m^\top \mathcal{R}_n = \mathcal{R}_{n-m}

于是：

sm,n=qm⊤Rn−mkndhs_{m,n} = \frac{q_m^\top \mathcal{R}_{n-m} k_n}{\sqrt{d_h}}

这说明整个高维版本同样只依赖于相对位置差 n−mn-m。

7. 用复数表示会更优雅

把每一对维度 (x2i,x2i+1)(x_{2i},x_{2i+1}) 看作一个复数：

zi=x2i+jx2i+1z_i = x_{2i} + j x_{2i+1}

旋转就等价于乘上一个单位复数：

zi′=zi⋅ejpωiz_i' = z_i \cdot e^{j p \omega_i}

因为：

ejpωi=cos⁡(pωi)+jsin⁡(pωi)e^{j p\omega_i} = \cos(p\omega_i)+j\sin(p\omega_i)

这和二维旋转完全等价。

对于 query 和 key：

z~m,i(q)=zm,i(q)ejmωi,z~n,i(k)=zn,i(k)ejnωi\tilde z^{(q)}_{m,i}=z^{(q)}_{m,i} e^{j m\omega_i},\qquad \tilde z^{(k)}_{n,i}=z^{(k)}_{n,i} e^{j n\omega_i}

在复数内积的意义下，会出现因子：

ej(n−m)ωie^{j(n-m)\omega_i}

所以相对位置差自然出来了。

8. RoPE 如何与 Transformer 融合？

8.1 插入位置

RoPE 不是加在输入 embedding 上，而是插在 线性投影后的 QQ、KK 上：

  1. 输入隐状态 XX
  2. 线性映射：
    Q=XWQ,K=XWK,V=XWVQ=XW_Q,\quad K=XW_K,\quad V=XW_V
  3. 对每个位置 pp：
    Q~p=RoPE(Qp,p),K~p=RoPE(Kp,p)\tilde Q_p=\mathrm{RoPE}(Q_p,p),\qquad \tilde K_p=\mathrm{RoPE}(K_p,p)
  4. 用 Q~,K~,V\tilde Q,\tilde K,V 做注意力：
    Attention(Q,K,V)=softmax(Q~K~⊤dh)V\mathrm{Attention}(Q,K,V) =\mathrm{softmax}\left(\frac{\tilde Q \tilde K^\top}{\sqrt{d_h}}\right)V

注意：通常只对 Q,KQ,K 做 RoPE，不对 VV 做。 因为位置信息主要需要进入“匹配分数” QK⊤QK^\top。

9. 多头注意力里的公式

多头注意力中，每个 head 有自己的 Q(h),K(h),V(h)Q^{(h)},K^{(h)},V^{(h)}，维度为 dhd_h。 RoPE 对每个 head 独立应用：

Q~p(h)=RpQp(h),K~p(h)=RpKp(h)\tilde Q^{(h)}_p = \mathcal{R}_p Q^{(h)}_p,\qquad \tilde K^{(h)}_p = \mathcal{R}_p K^{(h)}_p

然后第 hh 个头的注意力：

head(h)=softmax(Q~(h)(K~(h))⊤dh)V(h)\mathrm{head}^{(h)} = \mathrm{softmax}\left( \frac{\tilde Q^{(h)} (\tilde K^{(h)})^\top}{\sqrt{d_h}} \right)V^{(h)}

最后拼接各头结果并线性映射：

MHA(X)=Concat(head(1),…,head(H))WO\mathrm{MHA}(X)=\mathrm{Concat}(\mathrm{head}^{(1)},\dots,\mathrm{head}^{(H)})W_O

10. 为什么它比“直接加位置向量”更自然？

对比两种方式：

10.1 绝对位置编码（加法）

把位置向量 pmp_m 加到输入上：

xm′=xm+pmx_m' = x_m + p_m

然后经过线性层变成 qm,kmq_m,k_m。 这里位置信息是“混入内容表示里”的。

10.2 RoPE（乘法/旋转）

直接在 Q,KQ,K 上做：

q~m=Rmqm,k~n=Rnkn\tilde q_m = \mathcal{R}_m q_m,\qquad \tilde k_n = \mathcal{R}_n k_n

此时内积变成：

q~m⊤k~n=qm⊤Rn−mkn\tilde q_m^\top \tilde k_n = q_m^\top \mathcal{R}_{n-m} k_n

可以看到，相对位置差直接进入了 attention kernel，这是 RoPE 更“贴近注意力结构”的地方。

11. 一个小的二维展开示例

假设只有一个二维对子，query 和 key 分别是：

q=[q1q2],k=[k1k2]q=\begin{bmatrix}q_1\\q_2\end{bmatrix},\qquad k=\begin{bmatrix}k_1\\k_2\end{bmatrix}

位置分别为 m,nm,n，频率为 ω\omega。 则：

q~=R(mω)q,k~=R(nω)k\tilde q = R(m\omega)q,\qquad \tilde k = R(n\omega)k

内积：

q~⊤k~=q⊤R((n−m)ω)k\tilde q^\top \tilde k = q^\top R((n-m)\omega) k

将 R((n−m)ω)R((n-m)\omega) 展开：

R(Δω)=[cos⁡(Δω)−sin⁡(Δω)sin⁡(Δω)cos⁡(Δω)],Δ=n−mR(\Delta\omega)= \begin{bmatrix} \cos(\Delta\omega) & -\sin(\Delta\omega) \\ \sin(\Delta\omega) & \cos(\Delta\omega) \end{bmatrix}, \quad \Delta = n-m

于是：

q~⊤k~=q1(k1cos⁡Δω−k2sin⁡Δω)+q2(k1sin⁡Δω+k2cos⁡Δω)\tilde q^\top \tilde k = q_1(k_1\cos\Delta\omega - k_2\sin\Delta\omega) + q_2(k_1\sin\Delta\omega + k_2\cos\Delta\omega)

整理得：

q~⊤k~=(q1k1+q2k2)cos⁡(Δω)+(q2k1−q1k2)sin⁡(Δω)\tilde q^\top \tilde k = (q_1k_1+q_2k_2)\cos(\Delta\omega) + (q_2k_1-q_1k_2)\sin(\Delta\omega)

这个公式很有意思，它表明注意力分数中会出现：

  • 一个乘 cos⁡(Δω)\cos(\Delta\omega) 的项
  • 一个乘 sin⁡(Δω)\sin(\Delta\omega) 的项

所以分数会随着相对位置 Δ=n−m\Delta=n-m 发生规律变化。

12. 实际实现公式：常见的“rotate half”写法

在工程实现里，通常不显式构造大矩阵 Rp\mathcal{R}_p，而是直接按偶/奇位做操作。

设：

xeven=[x0,x2,… ],xodd=[x1,x3,… ]x_{\text{even}} = [x_0,x_2,\dots],\qquad x_{\text{odd}} = [x_1,x_3,\dots]

定义一个“旋转半维”的操作：

rotate(x)=[−x1,x0,−x3,x2,… ]\mathrm{rotate}(x) = [-x_1, x_0, -x_3, x_2, \dots]

那么 RoPE 可以简写为：

RoPE(x,p)=x⊙cos⁡p+rotate(x)⊙sin⁡p\mathrm{RoPE}(x,p)=x\odot \cos_p + \mathrm{rotate}(x)\odot \sin_p

其中 cos⁡p,sin⁡p\cos_p,\sin_p 是按维度广播后的余弦/正弦表。 更严格地写，对每个二维对子 (x2i,x2i+1)(x_{2i},x_{2i+1})：

[x2i′x2i+1′]=[x2ix2i+1]cos⁡(pωi)+[−x2i+1x2i]sin⁡(pωi)\begin{bmatrix} x'_{2i}\\ x'_{2i+1} \end{bmatrix} = \begin{bmatrix} x_{2i}\\ x_{2i+1} \end{bmatrix} \cos(p\omega_i) + \begin{bmatrix} -x_{2i+1}\\ x_{2i} \end{bmatrix} \sin(p\omega_i)

这个形式和前面完全等价。

13. 训练与推理时的使用

训练时

对整段序列每个位置 p=0,1,…,L−1p=0,1,\dots,L-1 预先算好：

cos⁡(pωi),sin⁡(pωi)\cos(p\omega_i),\qquad \sin(p\omega_i)

然后对所有 token 的 Q,KQ,K 施加 RoPE。

自回归推理 + KV Cache

假设当前解码到位置 tt：

  • 历史位置 0,…,t−10,\dots,t-1 的 K,VK,V 已缓存
  • 新 token 只需生成自己的 qt,kt,vtq_t,k_t,v_t

然后应用：

q~t=Rtqt,k~t=Rtkt\tilde q_t=\mathcal{R}_t q_t,\qquad \tilde k_t=\mathcal{R}_t k_t

缓存 k~t,vt\tilde k_t,v_t，并让 q~t\tilde q_t 和所有历史 K~\tilde K 做注意力即可。 因为每个位置的旋转是独立确定的，所以与 KV cache 天然兼容。

14. RoPE 为什么有一定长度外推能力？

因为它的位置函数不是查表，而是连续函数：

cos⁡(pωi),sin⁡(pωi)\cos(p\omega_i),\qquad \sin(p\omega_i)

理论上 pp 可以取任意大，不需要新增参数。 但注意：可计算 ≠ 一定泛化得好。 当 pp 很大时，某些高频分量会快速振荡，可能导致模型在超长上下文下性能下降。因此后续有很多 RoPE scaling、NTK-aware scaling、YaRN、LongRoPE 等改进。

15. 最后给一个“最标准”的总公式

如果你想在脑子里只记一个公式，可以记这个：

定义位置旋转矩阵

Rp=diag(R(pω0),R(pω1),…,R(pωdh/2−1))\mathcal{R}_p = \mathrm{diag}\Big( R(p\omega_0),R(p\omega_1),\dots,R(p\omega_{d_h/2-1}) \Big)

作用到 Query / Key

q~m=Rmqm,k~n=Rnkn\tilde q_m=\mathcal{R}_m q_m,\qquad \tilde k_n=\mathcal{R}_n k_n

注意力分数

sm,n=q~m⊤k~ndh=qm⊤Rn−mkndhs_{m,n} = \frac{\tilde q_m^\top \tilde k_n}{\sqrt{d_h}} = \frac{q_m^\top \mathcal{R}_{n-m} k_n}{\sqrt{d_h}}

这个最后等式就是 RoPE 的核心价值： 绝对位置旋转，得到相对位置感知的注意力。

avatar
文章
阅读
粉丝