大模型基础之位置编码RoPE

RoPE

RoPE(Rotary Position Embedding，旋转式位置编码)是一种配合Attention机制能达到“通过绝对位置编码的方式实现相对位置编码”的设计。在2021年2月由苏剑林提出，是现在的大模型最常用的位置编码。

记$\mathbf{q}$和$\mathbf{k}$为self-attention的q和k向量，$j$是元素索引。假设有$0<\theta \le \frac{\pi }{2N}$，N是最大序列长度。 i表示复数的虚数单位， $⟨⟩$为内积符号， $\overline{z}$表示复数z的共轭复数(注：有些地方会将共轭复数记作$z^{\ast }$）。RoPE可表示为下述过程：
$$\begin{array}{rl}\text{RoPE}(\mathbf{x},m)& =\mathbf{x}{e}^{im\theta }\\ ⟨\text{RoPE}(q_j,m),\text{RoPE}(k_j,n)⟩& =⟨q_j{e}^{im\theta },k_j{e}^{in\theta }⟩\\ & =q_j{k}_j{e}^{im\theta }\overline{e^{in\theta }}\\ & =q_j{k}_j{e}^{i(m-n)\theta }\\ & =\text{RoPE}(q_j{k}_j,m-n)\end{array}$$

RoPE的示意图如下图(来自RoFormer论文)

准备知识

• 复数的笛卡尔积形式(Cartesian form): $z=a+ib$

• 复数的极坐标形式(polar form): $z=r(\cos \theta +i\sin \theta )$， 其中 $r=|z|=\sqrt{a^2+b^2}$， $\theta =\arg (x)={\tan }^{-1}\frac{b}{a}$

• 欧拉公式： $e^{ix}=\cos (x)+i\sin (x)$

• 复数的指数形式(exponential form)： $z=r{\text{e}}^{i\theta }$， 其中 $r=|z|=\sqrt{a^2+b^2}$， $\theta =\arg (x)={\tan }^{-1}\frac{b}{a}$

• 在复数的指数形式下复数$z=r{\text{e}}^{i\theta }$的共轭复数为 $\overline{z}=r{\text{e}}^{-i\theta }$， 两个复数$z=r{\text{e}}^{i\theta }$和$w=t{\text{e}}^{i\varphi }$的乘积为$zw=rt{\text{e}}^{i(\theta +\varphi )}$

• 在复数的极坐标形式下两个复数$z=r(\cos \theta +i\sin \theta )$和$w=t(\cos \varphi +i\sin \varphi )$的乘积为$zw=rt(\cos (\theta +\varphi )+i\sin (\theta +\varphi ))$

• 复数$z=a+ib$表示成矩阵时的形式为$\left(\begin{array}{cc}a& -b\\ b& a\end{array}\right)$，而旋转矩阵形式为$\left[\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]$，所以复数乘法的几何意义为将向量逆时针旋转$\theta$（也可以从复数极坐标形式的乘法来理解旋转）

  

RoPE的推导

假设有函数f(x, l)给位置l处的元素x添加绝对位置信息，对于$\mathbf{q}$和$\mathbf{k}$用函数f编码后有：
$${\tilde{\mathbf{q}}}_m=f(\mathbf{q},m),{\tilde{\mathbf{k}}}_n=f(\mathbf{k},n)(1)$$
即我们希望为$\mathbf{q}$和$\mathbf{k}$设计函数$f(\cdot ,m)$和$f(\cdot ,n)$，使用函数编码后，${\tilde{\mathbf{q}}}_m$和${\tilde{\mathbf{k}}}_n$带有了位置m和n的绝对位置信息。因为self-attention的核心运算是内积，我们希望$\mathbf{q}$和$\mathbf{k}$内积的结果有相对位置信息，即我们假设存在有如下恒等关系：
$$⟨f(\mathbf{q},m),f(\mathbf{k},n)⟩=g(\mathbf{q},\mathbf{k},m-n)(2)$$
我们的目标是求出这个恒等式的一个解。先将求解过程中的初始条件设为$f(\mathbf{q},0)=\mathbf{q}$， $f(\mathbf{k},0)=\mathbf{k}$，这可以理解为没有位置编码信息加入的情形。

我们先考虑二维情形，并借助复数来求解。用复数的指数形式来表示函数:

$$\begin{array}{rl}f(\mathbf{q},m)& =R_f(\mathbf{q},m)e^{i{\Theta }_f(\mathbf{q},m)}(3a)\\ f(\mathbf{k},n)& =R_f(\mathbf{k},n)e^{i{\Theta }_f(\mathbf{k},n)}(3b)\\ g(\mathbf{q},\mathbf{k},m-n)& =R_g(\mathbf{q},\mathbf{k},m-n)e^{i{\Theta }_g(\mathbf{q},\mathbf{k},m-n)}(3c)\end{array}$$

上式中的$R_f$，$R_g$ 是f和g的径向分量(radial component)，${\Theta }_f$，${\Theta }_g$是f和g的幅角分量(angular components)。将它们代到恒等表达式(2）中，可以得到：

$$\begin{array}{rl}{R}_f(\mathbf{q},m)R_f(\mathbf{k},n)& =R_g(\mathbf{q},\mathbf{k},m-n)(4a)\\ {\Theta }_f(\mathbf{q},m)-{\Theta }_f(\mathbf{k},n)& ={\Theta }_g(\mathbf{q},\mathbf{k},m-n)(4b)\end{array}$$

由初始条件$f(\mathbf{q},0)=\mathbf{q}$， $f(\mathbf{k},0)=\mathbf{k}$我们有（$||\mathbf{q}||$，$||\mathbf{k}||$ 和 ${\theta }_q$，${\theta }_k$是向量$\mathbf{q}$和$\mathbf{k}$在二维平面上的径向和幅角分量）：
$$\begin{array}{r}\mathbf{q}=||\mathbf{q}||e^{i{\theta }_q}=R_q(\mathbf{q},0)e^{i{\Theta }_q(\mathbf{q},0)}\\ \mathbf{k}=||\mathbf{k}||e^{i{\theta }_k}=R_q(\mathbf{k},0)e^{i{\Theta }_k(\mathbf{k},0)}\end{array}(5)$$
设m=n，以及考虑到初始条件$f(\mathbf{x},0)=\mathbf{x}$, 由式(4a)可以得到：
$$R_f(\mathbf{q},m)R_f(\mathbf{k},m)=R_g(\mathbf{q},\mathbf{k},0)=R_f(\mathbf{q},0)R_f(\mathbf{k},0)=\parallel \mathbf{q}\parallel \parallel \mathbf{k}\parallel (6)$$

由上式我们可以有$R_f(\mathbf{q},m)=R_f(\mathbf{q},0)=||\mathbf{q}||$， $R_f(\mathbf{k},m)=R_f(\mathbf{k},0)=||\mathbf{k}||$，$R_g(\mathbf{q},\mathbf{k},m-n)=R_g(\mathbf{q},\mathbf{k},0)=||\mathbf{q}||||\mathbf{k}||$即$R_f$和$R_g$不依赖于位置信息。

类似地，设m=n，以及考虑到初始条件$\Theta (\mathbf{x},0)=\Theta (\mathbf{x})$，由式(4b)可以得到（$\Theta (\mathbf{q})$ 和$\Theta (\mathbf{k})$是向量$\mathbf{q}$和$\mathbf{k}$的幅角）：
$${\Theta }_f(\mathbf{q},m)-{\Theta }_f(\mathbf{k},m)={\Theta }_g(\mathbf{q},\mathbf{k},0)={\Theta }_f(\mathbf{q},0)-{\Theta }_f(\mathbf{k},0)=\Theta (\mathbf{q})-\Theta (\mathbf{k})(7)$$
将上式的第一项和最后一项移位我们可得${\Theta }_f(\mathbf{q},m)-\Theta (\mathbf{q})={\Theta }_f(\mathbf{k},m)-\Theta (\mathbf{k})$，所以${\Theta }_f(\mathbf{q},m)-\Theta (\mathbf{q})$是一个只与m有关与$\mathbf{q}$无关的函数，将其记为$\phi (m)$，则有${\Theta }_f(\mathbf{q},m)=\Theta (\mathbf{q})+\phi (m)$。令n=m-1，将其代入到式(4b)并移项可得
$$\phi (m)-\phi (m-1)={\Theta }_g(\mathbf{q},\mathbf{k},1)+\Theta (\mathbf{k})-\Theta (\mathbf{q})$$
因为上式右侧与m无关，所以上式左侧也必须与m无关，因此$\phi$是一个等差数列(arithmetic progression)，如果我们设等差数列的初始值$\phi (0)=0$，$\phi (1)=\theta$，那么就可得$\phi (m)=m\theta$。

综上，我们得到了二维情况下用复数表示的RoPE, 它是满足恒等式2的一个解：
$$f(\mathbf{q},m)=R_f(\mathbf{q},m)e^{i{\Theta }_f(\mathbf{q},m)}=||\mathbf{q}||e^{i(\Theta (\mathbf{q})+m\theta )}=\mathbf{q}{e}^{im\theta }$$
根据复数乘法的几何意义，这个变换对应着向量的旋转，所以RoPE作者将其称之为”旋转式位置编码“。

将上式表示成矩阵形式：
$$\mathbf{f}(\mathbf{q},m)=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\genfrac{}{}{0ex}{}{q_0}{q_1}\right)$$
由于内积满足线性叠加性，所以任意偶数维的RoPE都可以表示为二维情形的拼接。
$$f(\mathbf{q},m)=\left(\begin{array}{cccc}{M}_0& & & \\ & M_1& & \\ & & \ddots & \\ & & & M_{d/2-1}\end{array}\right)\left(\begin{array}{c}{q}_0\\ q_1\\ ⋮\\ q_{d-1}\end{array}\right)={\mathbf{R}}_{\mathbf{m}}{\mathbf{q}}_{\mathbf{m}}={\mathbf{R}}_{\mathbf{m}}{\mathbf{W}}_{\mathbf{q}}{\mathbf{x}}_{\mathbf{m}}$$
上式中$M_j=\left(\begin{array}{cc}\cos m{\theta }_j& -\sin m{\theta }_j\\ sinm{\theta }_j& \cos m{\theta }_j\end{array}\right)$， Θ = { θ i = 1000 0 − 2 ( i − 1 ) / d , i ∈ [ 0 , 1 , 2 , … , d / 2 − 1 ] } \Theta = \{\theta_i=10000^{-2(i-1)/d}, i \in[0,1,2, \ldots, d/2-1] \} Θ={θi​=10000−2(i−1)/d,i∈[0,1,2,…,d/2−1]}， ${\mathbf{R}}_{\mathbf{m}}$是对角旋转矩阵，它是一个正交矩阵，${\mathbf{W}}_{\mathbf{q}}$是待学习的query权重，${\mathbf{x}}_{\mathbf{m}}$则是m处的token的embedding。

也就是说，给位置m的向量$\mathbf{q}$乘上矩阵${\mathbf{R}}_{\mathbf{m}}$，位置为n的向量$\mathbf{k}$乘上矩阵${\mathbf{R}}_{\mathbf{n}}$， 用变换后的序列做Attention，Attention就自动包含相对位置了，因为有如下恒等式：
$$({\mathbf{R}}_{\mathbf{m}}\mathbf{q}{)}^T({\mathbf{R}}_{\mathbf{n}}\mathbf{k})={\mathbf{q}}^{\mathbf{T}}{\mathbf{R}}_{\mathbf{m}}^{\mathbf{T}}{\mathbf{R}}_{\mathbf{n}}\mathbf{k}={\mathbf{q}}^{\mathbf{T}}{\mathbf{R}}_{\mathbf{m}-\mathbf{n}}\mathbf{k}$$
因为${\mathbf{R}}_{\mathbf{m}}$的稀疏性，直接用矩阵乘法来实现很浪费算法，所以RoPE作者推荐用如下方式来实现RoPE:
$$\left(\begin{array}{c}{q}_0\\ q_1\\ q_2\\ q_3\\ ⋮\\ q_{d-2}\\ q_{d-1}\end{array}\right)\otimes \left(\begin{array}{c}\cos m{\theta }_0\\ \cos m{\theta }_0\\ \cos m{\theta }_1\\ \cos m{\theta }_1\\ ⋮\\ \cos m{\theta }_{d/2-1}\\ \cos m{\theta }_{d/2-1}\end{array}\right)+\left(\begin{array}{c}-q_1\\ q_0\\ -q_3\\ q_2\\ ⋮\\ -q_{d-1}\\ q_{d-2}\end{array}\right)\otimes \left(\begin{array}{c}\sin m{\theta }_0\\ \sin m{\theta }_0\\ \sin m{\theta }_1\\ \sin m{\theta }_1\\ ⋮\\ \sin m{\theta }_{d/2-1}\\ \sin m{\theta }_{d/2-1}\end{array}\right)$$
上式中的$\otimes$是逐位对应相乘，是Numpy等中的*运算。 从这个实现也可以来解释为什么RoPE可以视为是乘性位置编码的变体。

RoPE的代码实现

按照上面RoPE作者推荐的方式实现RoPE的示例如下(来自参考资料4)

import torch
import math

def rotary_position_embedding(q, k):
    """
    Rotary Position Embedding (RoPE) for queries and keys.
    
    Args:
        q: tensor for queries of shape (batch_size, num_heads, seq_len, dim)
        k: tensor for keys of shape (batch_size, num_heads, seq_len, dim)
        
    Returns:
        Rotated queries and keys
    """
    batch_size, num_heads, seq_len, dim = q.size()
    
    # Begin of sinusoidal_position_embedding content
    # 序列对应的位置序号
    position = torch.arange(seq_len, dtype=torch.float).unsqueeze(-1).to(q.device)
    # q维度上的theta值
    div_term = torch.exp(torch.arange(0, dim, 2, dtype=torch.float) * -(math.log(10000.0) / dim)).to(q.device)
    
    pos_emb = position * div_term
    pos_emb = torch.stack([torch.sin(pos_emb), torch.cos(pos_emb)], dim=-1).flatten(-2, -1)
    pos_emb = pos_emb.unsqueeze(0).unsqueeze(1)
    pos_emb = pos_emb.expand(batch_size, num_heads, -1, -1)
    # End of sinusoidal_position_embedding content

    # Extract and duplicate cosine and sine embeddings
    cos_emb = pos_emb[..., 1::2].repeat_interleave(2, dim=-1)
    sin_emb = pos_emb[..., ::2].repeat_interleave(2, dim=-1)

    # Create alternate versions of q and k
    q_alternate = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1).reshape(q.size())
    k_alternate = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1).reshape(k.size())

    # Rotate queries and keys
    q_rotated = q * cos_emb + q_alternate * sin_emb
    k_rotated = k * cos_emb + k_alternate * sin_emb

    return q_rotated, k_rotated

llama实现RoPE的方式是先将向量转到复数域，再对两个向量进行旋转，接着将向量转回到实数域。

# 以长度为4，dim维度为6的q示意llama是如何实现RoPE的
q = torch.tensor([[1, 2, 4, 5, 6, 7], [1, 2, 5, 6, 7, 8], [2, 5, 4, 6, 7, 8], [1, 3, 5, 6, 7, 9]])
seq_len, dim = q.shape  # [4,6]

# 将q在其embedding维度分为一对一对的形式
q_per_token_split_into_pairs = q.float().view(q.shape[0], -1, 2)
q_per_token_split_into_pairs.shape  # [4,3,2]

# 计算复数域的cos和sin的频率
zero_to_one_split_into_dim_parts = torch.tensor(range(dim//2))/(dim//2)
rope_theta = 10000.0
freqs = 1.0 / (rope_theta ** zero_to_one_split_into_dim_parts)
freqs_for_each_token = torch.outer(torch.arange(seq_len), freqs)
freqs_cis = torch.polar(torch.ones_like(freqs_for_each_token), freqs_for_each_token)
freqs_cis

# 将q转到复数域
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)

# 进行dot product来按位置旋转q向量
q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis
q_per_token_as_complex_numbers_rotated

# 将旋转了的q向量转回到实数域
q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)
# 将维度还原
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q.shape)

参考资料

1. RoPE论文：Su, Jianlin, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. 2021. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” Cornell University - arXiv,Cornell University - arXiv, April.

2. RoPE作者苏剑林的博客： 让研究人员绞尽脑汁的Transformer位置编码 ，Transformer升级之路：2、博采众长的旋转式位置编码

3. eleuther.ai关于RoPE的博客

4. 知乎文章 位置编码｜RoPE｜ALiBi

5. 复数 wikipedia, 旋转矩阵wikipedia

6. rotary_embedding-torch github

7. llama3 from scratch