当前位置：首页 > article >正文

【大模型理论篇】RoPE旋转位置编码底层数学原理分析

article 2024/11/15 11:52:04

1. 位置编码对于NLP模型的作用

位置编码（Positional Encoding）在大模型（例如Transformer架构）中起到了非常重要的作用。没有位置编码的信息，模型会丧失序列的顺序信息，导致模型退化成一个简单的“词袋模型”（Bag of Words model）。

在Transformer中，自注意力（Self-Attention）机制是模型学习序列数据特征的核心部分。自注意力机制允许模型为序列中的每个词计算它与序列中其他词的注意力权重。但是，自注意力机制本身是“位置不敏感”的，因为它并不包含任何关于词语在序列中位置的信息。如果没有位置编码的信息引入，模型将无法区分句子中词语的不同排列。比如，“我养猫”和“猫养我”会被视为相同的输入，因为自注意力机制只关注词语之间的相对相似性，而不考虑它们在序列中的具体位置。

这种情况下，模型就会退化成一个词袋模型（Bag of Words Model），忽略词语顺序，只关心词语出现频率。在这种模型中，序列中词语的顺序信息被丢失，对于自然语言处理任务（如翻译、情感分析）等是非常不利的，因为这些任务通常对词语顺序极为敏感。

位置编码的核心作用是为每个词提供其在序列中的位置信息，这使得模型能够区分词语的不同排列方式。由于位置编码的引入，模型能够感知词元与词元之间的相对位置是不同的，从而学习到正确的上下文信息。通过位置编码，模型不仅能学习到每个词的含义，还能学会词之间的顺序关系。本质上，自然语言处理模型学习的正是词元的词义以及词元的顺序组合关系。

2. 常见的位置编码方法

【1】中整理了两种主流思路的位置编码，包括绝对位置编码以及相对位置编码。

针对绝对位置编码：绝对位置编码信息会加到输入中，在输入的第k个向量 $x_k$ 中加入位置向量 $p_k$ 变为 $x_k+p_k$ ，其中 $p_k$ 只依赖于位置编号 $k$ 。绝对位置编码主要有训练式、三角函数式、递归式等。在Transformer原论文《Attention is all you need》中，提出的三角函数式就是一种绝对位置编码。该论文中位置编码被设计成使用正弦和余弦函数的组合，这种编码方法是固定的，并不随着训练而变化。其公式为：

$\text{PE}_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad \text{PE}_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$

其中，pos 是位置，i 是维度索引， $d_{model}$ 是模型的维度。这种编码方式使得不同位置的编码在每个维度上具有不同的周期性，确保每个位置的编码都是唯一的，并能让模型一定程度上推断出相对位置关系。

针对相对位置编码：自然语言中的语义关系通常更多依赖于词与词之间的相对位置，而不是它们在句子中的绝对位置。例如，在句子“猫追狗”和“公园里猫追狗”中，“追”前后的词（“猫”和“狗”）的相对位置决定了句子的含义，绝对位置不重要。仅仅通过 “猫”“追”“狗” 这三个元素的相对位置就能准确传达意思。在 “公园里猫追狗” 中，增加了 “公园里” 这个表示地点的修饰成分，但 “猫”“追”“狗” 的相对位置依然决定了核心的追逐含义，只是多了地点的限定。无论在句子中的绝对位置如何变化，只要 “猫” 在 “追” 前，“狗” 在 “追” 后，这种追逐关系就不会改变。这种相对位置的稳定性使得能够快速理解句子的主要语义，而不受绝对位置变化的过多影响。【2,3,4,5,6,7】分别对相对位置编码给出了技术实现，有兴趣可以关注。

其实还存在第三种方式，也是本文主要关注的RoPE旋转位置编码【8】，它是将绝对位置编码与相对位置编码进行了巧妙的结合，既体现了词元在句子序列中的绝对位置信息，又体现了词元与词元之间的相对位置关系。

RoPE目前已经成为主流认同的大模型位置编码模块，见表1。除此之外，ALiBi（Attention with linear biases）位置编码【11】因为良好的外推性也被广泛应用。

【9】进行了使用RoPE与其他位置编码方式在效果上的对比。使用 GPT-Neox 代码库对旋转嵌入与 GPT-3 中使用的学习型绝对位置嵌入以及 T5 中使用的学习型相对位置嵌入（ RPE）进行了比较。比较是使用具有 1.25 亿参数的模型进行。模型在 OpenWebText2 上进行训练，可以看到RoPE训练和验证曲线的收敛速度更快，总体验证损失更低，而吞吐量仅有轻微下降。还使用 mesh-transformer-jax 代码库和具有 14 亿参数的模型进行了额外的大规模实验，与学习型绝对位置嵌入和 T5 RPE 的基线进行对比。使用了与 GPT3 的 13 亿参数模型类似的超参数，数据集为 The Pile。观察到与学习型绝对位置嵌入相比，RoPE收敛速度有约 30%提高，与 T5 相对位置编码相比有10% - 20%的改进，这表明RoPE可扩展到数十亿参数范围。在 enwiki8 上对 Performer 进行了小规模测试，针对具有 512 维、8 个头的 8 层基于字符的 Transformer。这些测试表明，将旋转嵌入代入 Performer 会导致验证损失急剧下降并快速收敛。从对比结果来看，RoPE在各个方面都带来了积极的增益。

3. 旋转位置编码数学原理分析

虽然在之前的文章《LLaMA3结构关键模块分析》已经对RoPE旋转位置编码做了介绍，但是对于该算法的数学原理讨论不够，因此本文将针对底层的数学原理进一步分析。接下来的分析会以原论文【8】和参考材料【10】为主来进行分析。

3.1 直觉理解

旋转位置编码在编码位置信息方面，不同于Transformer原论文模式，没有采用添加一个单独的位置编码向量，而是采用对现有的词嵌入进行旋转。旋转角度是序列中词的位置以及嵌入维度的函数【12】。

$q_m=f_q(x_m, m)$ 表示在位置 m 处的查询向量。
$k_n = f_k(x_n, n)$ 表示在位置 n 处的键向量。
$v_n = f_v(x_n, n)$ 表示在位置 n 处的值向量。

假设有一个词序列 w₁，w₂，…，wₗ，其中这个序列的长度为 L。每个词在高维空间中被转换为一个向量，称这个空间的维度为 |D|。在 Transformer 模型中，序列中的每个位置（m）都有一个查询向量（qₘ）和一个键向量（kₘ），这些向量是通过函数 $f_q$ 和 $f_k$ 进行创建。在计算注意力矩阵时，主要体现出两个关键信息：

词相似性：具有相似嵌入向量的词有更高的注意力得分。

位置接近性：在序列中位置更接近的词应该有更高的得分。

为了确定一个词元与另一个词元的注意力得分，需要引入相应的点积计算。对于任何一对位置 m 和 n，查询向量 $q_m$ 和键向量 $k_n$ 的点积需要体现出词的相似性和位置信息。在这个点积中：（q・k = ||q|| ||k|| cos (θ)）。q 和 k 的幅度相似性（表示为 ||q||・||k||）对应于词嵌入的相似性。q 和 k 的角度（ $\theta _m$ 和 $\theta _n$ ）有助于位置相似性。可以看到，角度分量仅取决于向量在序列中的位置，与实际的词嵌入无关。RoPE巧妙的利用了该特性，根据序列中的位置旋转查询和键向量。旋转在保持词的相似性的同时，在角度中编码位置信息。在计算旋转向量之间的点积时，向量之间的相对角度捕获位置相近程度。这种思路使得RoPE将词和位置信息无缝地集成到单个操作中。

3.2 复数知识

在理解旋转位置编码数学推导之前，还需要具备一些先验知识。

对于二维向量 $(x, y)$ ，可以用复数来表示： $z=x+iy$ ，其中 i 是虚数单位（ $i^2 = -1$ ）。
使用欧拉公式 $e^{ix} = \cos x + i\sin x$ ，可以用三角函数来表示复数的实部和虚部。
给定两个二维向量 $\mathbf{q}$ 和 $\mathbf{k}$ ，假设它们分别表示为复数 $q = x_q + iy_q$ 和 $k = x_k + iy_k$ 。其复数共轭的定义： $k^* = x_k - iy_k$ 。
要求两个向量的内积，可以借助复数的乘积来实现： $\text{Re}[qk^*] = x_q x_k + y_q y_k$ 。

3.3 公式推导

3.3.1 目标公式

有了前两部分的铺垫后，接下来我们就可以对原论文中的公式进行详细推导，并证明通过旋转之后的查询向量与键向量的内积计算，可以同时包含绝对位置信息与相对位置信息。也就是证明存在以下关系：

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

对公式进行解释：

该公式表示旋转后的查询向量 $f_q(x_m, m)$ 和旋转后的键向量 $f_k(x_n, n)$ 的内积等于一个函数 $g(x_m, x_n, m - n)$ ，其中：

$\langle \cdot, \cdot \rangle$ 表示内积运算。
$f_q(x_m, m)$ 和 $f_k(x_n, n)$ 是定义在输入词向量 $x_m$ 和 $x_n$ 以及位置参数 m 和 n 上的函数。
$g(x_m, x_n, m - n)$ 是一个函数，表示输入 $x_m$ 和 $x_n$ 以及位置之差 $m - n$ 的某种函数关系。

该公式表达了两个函数在特定位置上的内积与它们位置的差有关，而不是绝对位置本身。计算自注意力时，公式中查询向量与键向量的内积反映了词元之间的相似性。

3.3.2 寻找满足关系的函数定义

接下来，就是找到某种函数表达，能够满足上述关系，原论文中给出了如下函数形式定义：

$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) = \text{Re} \left[ (W_q x_m) (W_k x_n)^* e^{i(m - n)\theta} \right]$

公式解释：

公式中使用了复数表示的旋转编码，通过旋转因子 $e^{im\theta}$ 和 $e^{in\theta}$ 引入相对位置信息。

$f_q(x_m, m)$ 和 $f_k(x_n, n)$ 是经过线性变换后的输入词向量分别乘上一个复数旋转因子 $e^{im\theta}$ 和 $e^{in\theta}$ ，其中：

$W_q$ 和 $W_k$ 是查询（query）和键（key）的权重矩阵。

$x_m$ 和 $x_n$ 分别是输入的词元嵌入向量表示。

$\theta$ 是一个旋转角度，与位置 m 和 n 的差异相关。

$g(x_m, x_n, m - n)$ 表示函数 $f_q$ 和 $f_k$ 的内积，它包含了位置差 m - n 的信息。

3.3.3 旋转位置编码的定义

根据欧拉公式，对复数旋转因子进行变换，得到：

$e^{im\theta} = \cos(m\theta) + i \sin(m\theta)$

$e^{in\theta} = \cos(n\theta) + i \sin(n\theta)$

$e^{i(m-n)\theta} = \cos((m-n)\theta) + i \sin((m-n)\theta)$

我们以词嵌入向量维度为2维来进一步说明，后续多维的情况其实就是按照2维字空间进行处理即可。

首先，将查询向量 q_m 和键向量 k_n 转换为复数形式。查询向量 q_m 是一个二维向量，表示为：

$q_m = \begin{bmatrix} q_m^{(1)} \\ q_m^{(2)} \end{bmatrix} = W_q x_m$

其中， $W_q$ 是一个 $2 \times 2$ 的矩阵， $x_m$ 是一个二维向量。乘积 $q_m$ 也是一个二维向量。

然后，将 $q_m$ 转化为复数形式：

$q_m = q_m^{(1)} + i q_m^{(2)}$

旋转位置编码将向量与一个复数的指数形式相乘，用于表示位置信息。查询向量的旋转位置编码表示为：

$f_q(x_m, m) = (W_q x_m) e^{im\theta} = q_m e^{im\theta}$

$q_m e^{im\theta} = (q_m^{(1)} + i q_m^{(2)}) \cdot (\cos(m\theta) + i \sin(m\theta))$

根据复数乘法的性质：

$(a + ib) \cdot (c + id) = (ac - bd) + i(bc + ad)$

对查询向量处理：

$q_m e^{im\theta} = (q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta)) + i(q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta))$

重新表达为向量形式：

$q_m e^{im\theta} = \begin{bmatrix} q_m^{(1)} \cos(m\theta) - q_m^{(2)} \sin(m\theta) \\ q_m^{(2)} \cos(m\theta) + q_m^{(1)} \sin(m\theta) \end{bmatrix}$

上述公式，说明查询向量乘以了一个旋转矩阵，形式如下：

$\begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix} \cdot \begin{bmatrix} q_m^{(1)} \\ q_m^{(2)} \end{bmatrix}$

同样的，对键向量 $k_n$ 的旋转位置编码表示为：

$f_k(x_n, n) = (W_k x_n) e^{in\theta} = k_n e^{in\theta}$

其中：

$k_n e^{in\theta} = \begin{bmatrix} k_n^{(1)} \cos(n\theta) - k_n^{(2)} \sin(n\theta) \\ k_n^{(2)} \cos(n\theta) + k_n^{(1)} \sin(n\theta) \end{bmatrix}$

同样表示为一个旋转矩阵乘以键向量：

$\begin{bmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{bmatrix} \cdot \begin{bmatrix} k_n^{(1)} \\ k_n^{(2)} \end{bmatrix}$

函数 $g(x_m, x_n, m - n)$ 的定义为：

$g(x_m, x_n, m - n) = \text{Re}\left[(W_q x_m)(W_k x_n)^* e^{i(m-n)\theta}\right]$

其中， $\text{Re}$ 表示取实部，这个函数就是用于计算查询向量和键向量在旋转位置编码下的内积。

3.3.4 函数g形式的成立证明

接下来就是需要计算该函数g，首先，我们有查询向量和键向量的复数表示：

$W_q x_m = q_m = q_m^{(1)} + i q_m^{(2)}, \quad W_k x_n = k_n = k_n^{(1)} + i k_n^{(2)}$

以及它们的复共轭：

$(W_k x_n)^* = k_n^* = k_n^{(1)} - i k_n^{(2)}$

和复数 $e^{i(m - n)\theta}$ ：

$e^{i(m - n)\theta} = \cos((m - n)\theta) + i \sin((m - n)\theta)$

继续计算函数 $g(x_m, x_n, m - n)$ ：

$g(x_m, x_n, m - n) = \text{Re}\left[(W_q x_m)(W_k x_n)^* e^{i(m - n)\theta}\right]$

将复数形式代入：

$= \text{Re} \left[ (q_m^{(1)} + i q_m^{(2)}) (k_n^{(1)} - i k_n^{(2)}) (\cos((m - n)\theta) + i \sin((m - n)\theta)) \right]$

首先计算两个复数的乘积：

$(q_m^{(1)} + i q_m^{(2)}) (k_n^{(1)} - i k_n^{(2)}) = (q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)}) + i (q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)})$

接着，与 $e^{i(m - n)\theta}$ 相乘：

$\left( (q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)}) + i (q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)}) \right) \cdot (\cos((m - n)\theta) + i \sin((m - n)\theta)) \\ = (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) + i \left[ (q_m^{(1)} k_n^{(1)} + q_m^{(2)} k_n^{(2)}) \sin((m - n)\theta) + (q_m^{(2)} k_n^{(1)} - q_m^{(1)} k_n^{(2)}) \cos((m - n)\theta) \right]$

取上式的实部：

$\text{Re} \left[ (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) \right]$

这正是我们希望拿到的结果：

$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)$

接下来，我们继续通过内积的方式来证明函数g公式的成立性：

我们有位置 m 的 query 向量和位置 n 的 key 向量，分别为：

$f_q(x_m, m) = [q^{(1)}_m \cos(m \theta) - q^{(2)}_m \sin(m \theta), q^{(2)}_m \cos(m \theta) + q^{(1)}_m \sin(m \theta)]$

$f_k(x_n, n) = [k^{(1)}_n \cos(n \theta) - k^{(2)}_n \sin(n \theta), k^{(2)}_n \cos(n \theta) + k^{(1)}_n \sin(n \theta)]$

内积计算：

$\langle f_q(x_m, m), f_k(x_n, n) \rangle$

展开内积操作得到：

$\langle f_q(x_m, m), f_k(x_n, n) \rangle = (q^{(1)}_m \cos(m \theta) - q^{(2)}_m \sin(m \theta))(k^{(1)}_n \cos(n \theta) - k^{(2)}_n \sin(n \theta))+ (q^{(2)}_m \cos(m \theta) + q^{(1)}_m \sin(m \theta))(k^{(2)}_n \cos(n \theta) + k^{(1)}_n \sin(n \theta))$

整理得：

$= q^{(1)}_m k^{(1)}_n (\cos(m \theta) \cos(n \theta) + \sin(m \theta) \sin(n \theta))- q^{(1)}_m k^{(2)}_n (\cos(m \theta) \sin(n \theta) - \sin(m \theta) \cos(n \theta))- q^{(2)}_m k^{(1)}_n (\sin(m \theta) \cos(n \theta) - \cos(m \theta) \sin(n \theta))+ q^{(2)}_m k^{(2)}_n (\sin(m \theta) \sin(n \theta) + \cos(m \theta) \cos(n \theta))$

应用三角函数恒等式:

$\cos(a + b) = \cos a \cos b - \sin a \sin b$

$\cos(a - b) = \cos a \cos b + \sin a \sin b$

$\sin(a + b) = \sin a \cos b + \cos a \sin b$

$\sin(a - b) = \sin a \cos b - \cos a \sin b$

我们可以将内积表达式整理为：

$\langle f_q(x_m, m), f_k(x_n, n) \rangle = (q^{(1)}_m k^{(1)}_n + q^{(2)}_m k^{(2)}_n) \cos((m - n) \theta)+ (q^{(1)}_m k^{(2)}_n - q^{(2)}_m k^{(1)}_n) \sin((m - n) \theta)= (q^{(1)}_m k^{(1)}_n + q^{(2)}_m k^{(2)}_n) \cos((m - n) \theta)- (q^{(2)}_m k^{(1)}_n - q^{(1)}_m k^{(2)}_n) \sin((m - n) \theta)$

而之前我们已经得到函数 g的形式为:

$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)$

所以因此可以证明：

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

3.3.5 结论

因此，函数 $g(x_m, x_n, m - n)$ 的计算公式成立，这个公式通过使用旋转编码捕捉了查询和键向量之间的相对位置关系，并在实际应用中能够帮助模型更好地理解序列数据的结构。

3.3.6 扩展到多维向量场景

根据旋转矩阵的定义，RoPE 在处理查询和键向量的时候，将每对连续出现的两个元素视为一个子空间。因此，对于一个长度为H 的向量来说，将会形成 H / 2 个这样的子空间。在这些子空间中，每一个子空间 i ∈ { 1 , . . . , H / 2} 所对应的两个元素 $t \cdot \theta_i$ 都会根据一个特定的旋转角度进行旋转，其中t 代表位置索引，而 $\theta_i$ 表示该子空间中的基。有了上述的知识背景，再回过头来看旋转编码，就一清二楚了。