layout	title	date	summary	categories
post	RoPE 位置编码算法详解	2024-10-24 13:00:00 -0700	旋转位置编码（Rotary Position Embedding，RoPE）是论文 Roformer Enhanced Transformer With Rotray Position Embedding 提出的一种能够将相对位置信息依赖集成到 self-attention 中并提升 transformer 架构性能的位置编码方式。	Transformer

相关 torch 知识
RoPE 算法推导
RoPE 实现
- RoPE 实现流程
- RoPE 代码
参考资料

旋转位置编码（Rotary Position Embedding，RoPE）是论文 Roformer: Enhanced Transformer With Rotray Position Embedding 提出的一种能够将相对位置信息依赖集成到 self-attention 中并提升 transformer 架构性能的位置编码方式。

和相对位置编码相比，RoPE 具有更好的外推性，目前是大模型相对位置编码中应用最广的方式之一。这里的外推性实质是一个训练和预测的文本长度不一致的问题。具体来说，不一致的地方有两点：

预测的时候用到了没训练过的位置编码（不管绝对还是相对）；
预测的时候注意力机制所处理的 token 数量远超训练时的数量。

RoPE 的核心思想是将位置编码与词向量通过旋转矩阵相乘，使得词向量不仅包含词汇的语义信息，还融入了位置信息，其具有以下优点：

相对位置感知：RoPE 能够自然地捕捉词汇之间的相对位置关系。
无需额外的计算：位置编码与词向量的结合在计算上是高效的。
适应不同长度的序列：RoPE 可以灵活处理不同长度的输入序列。

三角函数、旋转矩阵、欧拉公式、复数等数学背景知识可以参考这篇文章学习。

RoPE 算法推导

PE 和 Self-Attention 概述

设 $q_m$ 表示第 $m$ 个 token 对应的词向量 $x_m$ 集成位置信息 $m$ 之后的 $query$ 向量；$k_n$ 和 $v_n$ 则表示词向量 $x_n$ 集成其位置信息 $n$（第 $n$ 个 token）之后的 key 和 value 向量，$q_m、k_n、v_n$ 的表达用如下公式:

$$q_m = f_q(x_m, m) \tag{1} \\ k_n = f_k(x_n, n) \\ v_n = f_v(x_n, n) $$

注意，这里的 $f_q$ 其实是把 $\text{embedding}_\text{vector} \times W_q$ 的矩阵乘法过程包含进去了，至于为什么要这样构造，下文会讲。

其中函数 $f_q、f_k、f_v$ 正是我们需要构造的位置编码函数。有了 query、key 和 value 向量表达式，接着就可以利用查询和键的值来计算注意力权重（$softmax(qk^T)$），输出则是对 $v_n$ 的加权求和。

$$ a_{m,n} = \frac{\exp\left(\frac{q_m^T k_n}{\sqrt{d}}\right)}{\sum_{j=1}^{N} \exp\left(\frac{q_m^T k_j}{\sqrt{d}}\right)} \\ o_m = \sum_{n=1}^{N} a_{m,n} v_n \quad (2)$$

方程 (1) 的一种常见选择是：

$$f_t:t∈{q,k,v}(x_i, i) := W_{t}(x_i + p_i)，\quad (3)$$

其中，$p_i \in \mathbb{R}^d$ 是与 token $x_i$ 的位置相关的 $d$ 维向量。Devlin 等人 [2019]、Lan 等人 [2020]、Clark 等人 [2020]、Radford 等人 [2019]、Radford 和 Narasimhan [2018] 使用了一组可训练向量 $p_i \in {p_t}_{t=1}^L$ ，其中 $L$ 表示最大序列长度。Vaswani 等人 [2017] 则提出了通过正弦函数来生成 $p_i$ 的方法:

$$p_{i,2t} = \sin\left(\frac{k}{10000^{2t/d}}\right) \\ p_{i,2t+1} = \cos\left(\frac{k}{10000^{2t/d}}\right)\quad (4)$$

其中， $p_{i,2t}$ 是 $p_i$ 的第 $2t$ 个维度。下一节会描述 RoPE 与这种基于正弦函数的直觉之间的关系。但是，RoPE 并不是直接将位置信息 $p_i$ 和嵌入向量元素 $x_i$ 相加，而是通过与正弦函数相乘的方式引入相对位置信息。

2D 的 RoPE 算法

RoPE 论文提出为了能利用 token 之间的相对位置信息（$m-n$），假定 query 向量 $q_m$ 和 key 向量 $k_n$ 之间的内积操作可以被一个函数 $g$ 表示，该函数 $g$ 的输入是词嵌入向量 $x_m$、$x_n$ 以及它们之间的相对位置 $m - n$，公式表达如下所示：

$$\langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m - n) \quad (5)$$

注意，这里只有 $f_q(x_m, m)$, $f_k(x_n, n)$ 是需要求解的函数，$\langle \rangle$ 表示内积操作，而对于 $g$，我们要求是表达式中有 $x_m, x_n, (m-n)$，也可以说是 $q_m, k_n$ 的内积会受相对位置 $m-n$ 影响。

接下来的目标就是找到一个等价的位置编码方式 $f$，从而使得上述关系成立，函数 $f_q$ 包含了位置编码和 $W_q \times q$（嵌入向量转换为 $q$ 向量）过程。

假设现在词嵌入向量的维度是两维 $d=2$，这样就可以利用上 $2$ 维度平面上的向量的几何性质，然后论文中提出了一个满足上述关系的 $f$ 和 $g$ 的形式如下:

$$ f_q(x_m, m) = (W_q x_m) e^{im\theta} \\ f_k(x_n, n) = (W_k x_n) e^{in\theta} \\ g(x_m, x_n, m - n) = Re \left[ (W_q x_m)(W_k x_n)^* e^{i(m-n)\theta} \right] \quad (6)$$

其中 ( Re ) 表示复数的实部，( (W_k x_n)^* ) 表示 ( (W_k x_n) ) 的共轭复数。

$f_q、f_k$ 的推导需要基于三角函数定理、欧拉公式等，推导过程参考这里，本文直接给出结论：

1，$f_q(x_m, m)$ 其实等于 query 向量乘以了一个旋转矩阵，即:

$$f_q(x_m, m) = \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} q_m^{(1)} \\ q_m^{(2)} \end{pmatrix} \quad (7)$$

2，$f_k(x_n, n)$ 其实等于 key 向量乘以了一个旋转矩阵，即:

$$f_k(x_n, n) = \begin{pmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \\ k_n^{(2)} \end{pmatrix} \quad (8)$$

3，同样可得 $g(x_m, x_n, m - n)$ 等于 $q_m^T$ 乘以旋转矩阵再乘以 $k_n$，即:

$$\langle f_q(x_m, m), f_k(x_n, n) \rangle = \mathbf{q}_m^T R(m - n) \mathbf{k}_n \quad (9)$$

$$\begin{aligned} g(x_m, x_n, m - n) &= (q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)}) \cos((m - n)\theta) - (q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)}) \sin((m - n)\theta) \\ &= \begin{pmatrix} q_m^{(1)} & q_m^{(2)} \end{pmatrix} \begin{pmatrix} \cos((m - n)\theta) & -\sin((m - n)\theta) \\ \sin((m - n)\theta) & \cos((m - n)\theta) \end{pmatrix} \begin{pmatrix} k_n^{(1)} \\ k_n^{(2)} \end{pmatrix} \\ &= \mathbf{q}_m^T R(m - n) \mathbf{k}_n \end{aligned} \quad(10)$$

公式（9）的证明可通过旋转矩阵性质得到，先将公式 (9) 抽象成 $\langle R_a X, R_b Y \rangle = \langle X, R_{b-a} Y \rangle$（$R$ 表示旋转矩阵，$X、Y$ 表示向量）, 该等式的证明过程如下：

$$\begin{aligned} \langle R_a X, R_b Y \rangle &= (R_aX)^T R_bY \\ &= X^T R_a^T R_bY \\ &= X^T R(-a)R_bY \\ &= X^T R_{(b-a)}Y = \langle X, R_{(b-a)}Y \rangle\\ \end{aligned} \quad(11)$$

上述推导过程分别应用了：展开内积、矩阵乘法的结合律、旋转矩阵性质1、旋转矩阵性质2。

多维的 RoPE 算法

前面的公式推导，是假设的词嵌入维度是 2 维向量，将二维推广到任意维度，$f_{{q,k}}$ 可以表示如下：

$$f_{{q,k}}(x_m, m) = R_{\Theta, m}^d W_{{q,k}} x_m \tag{12}$$

其中，$R_{\Theta, m}^d$ 为 $d$ 维度的旋转矩阵，表示为：

$$R_{\Theta, m}^d = \begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_1 & \cos m\theta_1 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2-1} & -\sin m\theta_{d/2-1} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2-1} & \cos m\theta_{d/2-1} \end{pmatrix} \tag{13}$$

$R_{\Theta, m}^d$ 的形状是 [sqe_len, dim//2]。$可以看出，对于 $d >= 2$ 的通用情况，则是将词嵌入向量元素按照两两一组分组，每组应用同样的旋转操作且每组的旋转角度计算方式如下：

$$ \Theta = \left{ \theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, \dots, d/2] \right} $$

将 RoPE 应用到前面公式（2）的 Self-Attention 计算，可以得到包含相对位置信息的Self-Attetion：

$$q_m^T k_n = \left( R_{\Theta, m}^d W_q x_m \right)^T \left( R_{\Theta, n}^d W_k x_n \right) = x_m^T W_q R_{\Theta, n-m}^d W_k x_n \tag{14}$$

其中， $$R_{\Theta, n-m}^d = \left( R_{\Theta, m}^d \right)^T R_{\Theta, n}^d$$

Rotary Position Embedding(RoPE) 实现的可视化如下图所示:

最后总结结合 RoPE 的 self-attention 操作的流程如下：

首先，对于 token 序列中的每个词嵌入向量，都计算其对应的 query 和 key 向量;
然后在得到 query 和 key 向量的基础上，应用公式（7）和（8）对每个 token 位置都计算对应的旋转位置编码；
接着对每个 token 位置的 query 和 key 向量的元素按照两两一组应用旋转变换；
最后再计算 query 和 key 之间的内积得到 self-attention 的计算结果。

RoPE 实现

RoPE 实现流程

先看旋转矩阵用于旋转一个二维向量过程示例：

但是 Llama 模型的嵌入维度高达 4096，比二维复杂得多，如何在更高维度的嵌入上应用旋转操作呢？通过 RoPE 算法原理我们知道，嵌入向量的旋转实际是将每个嵌入向量元素位置 $m$的值与每一对嵌入维度对应的 $\theta$ 相乘，过程如下图所示：

RoPE 通过实现旋转矩阵，是既捕获绝对位置信息，又结合相对位置信息的方式（论文公式有更详细体现）。

图中每组的旋转角度计算方式如下：

$$ \Theta = \left{ \theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, \dots, d/2] \right} $$

在实现 RoPE 算法之前，需要注意：为了方便代码实现，在进行旋转之前，需要将旋转矩阵转换为极坐标形式，嵌入向量（$q$、$k$）需要转换为复数形式。完成旋转后，旋转后的嵌入需要转换回实数形式，以便进行注意力计算。此外，RoPE 仅应用于查询（Query）和键（Key）的嵌入，不适用于值（Value）的嵌入。

RoPE 代码

通过仔细阅读和一步步分析了 llama 官方代码后，会发现作者直接转化为复数相乘形式来计算 $f_q(x_m, m) = (W_q x_m) e^{im\theta}$，旋转矩阵定义和各种变换完全没用上（有种前面推导了个寂寞的感觉），但是没办法，虽然直接使用旋转矩阵，更符合线性代数的常规思路，但是这需要手动处理每个维度的旋转矩阵，代码稍微繁琐而且计算效率不如复数乘法高效。

所以，作者在 RoPE 算法实现中，没有使用矩阵相乘的形式，而是把旋转角度张量和 $W_qx_mk$ 转为复数形式再相乘，即直接实现公式（6）中的 $f_q(W_q x_m) e^{im\theta}$，因此 RoPE 算法实现的流程和代码理解的难点如下：

如何生成旋转角度 $\theta$ 向量, $\Theta = { \theta_i = 10000^{-2(i-1)/d}, i \in \left [1, 2, \dots, d/2 \right ] }$;
如何将旋转角度和 token 位置索引相乘，并构造一个矩阵，该矩阵包含了每个位置和每个维度对应的旋转角度。
得到所有 token 位置和其对应旋转角度后，如何转换为复数形式 $e^{im\theta}$的旋转矩阵；
如何对 RoPE 函数的输入参数 x_q 做形状变换，并将实数张量转换为复数张量形式；
两个复数张量相乘（应用旋转操作）后，转换回实数域，并恢复原始形状。

最后，如果你直接看 pytorch 代码，其实很难理解 rope 是如何应用相对位置信息的，这个只能通过前面的公式推导才能理解。

结合 llama 官方实现代码，下述是我修改优化和添加注释后的代码，更容易看懂:

def compute_theta(dim: int, base: float = 10000.0, device: torch.device = torch.device('cpu')) -> torch.Tensor:
    """
    计算旋转位置编码中的 Theta 角度值。

    参数：
    - d (int): 嵌入向量的维度（必须为偶数）。
    - base (float): 基础频率参数, 默认为10000.0。
    - device (torch.device): 计算设备, 默认为CPU。

    返回：
    - torch.Tensor: 包含Theta值的1D张量, 形状为 [d/2]。
    """
    if dim % 2 != 0:
        print("嵌入维度 dim 必须为偶数")
    i = torch.arange(1, (dim//2) + 1, dtype=torch.float32, device=device)
    theta_i = base ** (-2*(i - 1) / dim)

    return theta_i

def precompute_freqs_cis(dim: int, seq_len: int, base: float = 10000.0, device: torch.device = torch.device('cpu')):
    theta = compute_theta(dim, base, device) # theta 角度值序列，向量, 大小为 dim // 2
    m = torch.arange(seq_len, device=device) # # token 位置值序列，向量，大小为 seq_len
    m_theta = torch.outer(m, theta) # 所有 token 位置的所有 Theta 值范围, 矩阵，尺寸为 [seq_len, dim // 2]
    freqs_cis = torch.polar(torch.ones_like(m_theta), m_theta) # e^{i*m*\theta}，本质上是旋转矩阵
    return freqs_cis

def reshape_for_broadcast(freqs_cis, x):
    ndim = x.ndim
    assert ndim >= 2
    assert freqs_cis.shape == (x.shape[1],x.shape[-1]), "the last two dimension of freqs_cis, x must match"
    shape = [d if i==1 or i==ndim-1 else 1 for i,d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, device: torch.device = torch.device('cpu')):
    """
    参数:
        - x_q(torch.Tensor): 实际上是权重 W_q * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]
        - x_k(torch.Tensor): 实际上是权重 W_k * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]
        - freqs_cis (torch.Tensor): 频率复数张量, 形状为 [max_seq_len, head_dim]
    返回:
        - Tuple[torch.Tensor, torch.Tensor]: 旋转编码后的查询和键张量
    """
    # 实数域张量转为复数域张量
    xq_reshape = xq.reshape(*xq.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] 
    xk_reshape = xk.reshape(*xk.shape[:-1], -1, 2) # [batch_size, seq_len, dim] -> [batch_size, seq_len, dim//2, 2] 
    xq_complex = torch.view_as_complex(xq_reshape) # 复数形式张量
    xk_complex = torch.view_as_complex(xk_reshape) # 复数形式张量

    # 旋转矩阵（freqs_cis）的维度在序列长度（seq_len，维度 1）和头部维度（head_dim，维度 3）上需要与嵌入的维度一致。
    # 此外，freqs_cis 的形状必须与 xq 和 xk 相匹配，因此我们需要将 freqs_cis 的形状从 [max_seq_len, head_dim] 调整为 [1, max_seq_len, 1, head_dim]。
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_complex) # [max_seq_len, 1, 1, dim // 2]

    # 应用旋转操作，并将结果转回实数域。# flatten(2) 将后面两个维度压成一个维度
    xq_out = torch.view_as_real(xq_complex * freqs_cis).flatten(3) 
    xk_out = torch.view_as_real(xk_complex * freqs_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

结合 rope 位置编码的 attention 结构的完整代码在这里，运行后，单元测试结果如下所示：

test_compute_theta passed.
test_precompute_freqs_cis passed.
test_apply_rotary_emb passed, xq_out and xq [0][0][0][0]: -1.3532123565673828 -1.3532123565673828.
test_attention passed.

transformers 库提供的 llama rope 实现在这里

参考资料

RoFormer: Enhanced Transformer with Rotary Position Embedding
十分钟读懂旋转编码（RoPE）
一文看懂 LLaMA 中的旋转式位置编码（Rotary Position Embedding）

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

RoPE位置编码算法详解.md

RoPE位置编码算法详解.md

相关 torch 知识

RoPE 算法推导

PE 和 Self-Attention 概述

2D 的 RoPE 算法

多维的 RoPE 算法

RoPE 实现

RoPE 实现流程

RoPE 代码

参考资料

Files

RoPE位置编码算法详解.md

Latest commit

History

RoPE位置编码算法详解.md

File metadata and controls

相关 torch 知识

RoPE 算法推导

PE 和 Self-Attention 概述

2D 的 RoPE 算法

多维的 RoPE 算法

RoPE 实现

RoPE 实现流程

RoPE 代码

参考资料