Skip to content

Latest commit

 

History

History
64 lines (35 loc) · 3.95 KB

File metadata and controls

64 lines (35 loc) · 3.95 KB

旋转位置编码(RoPE)之所以被设计成这种旋转矩阵的形式,其背后有着极其优雅的数学推导和深思熟虑的工程考量。

它最核心的设计初衷可以概括为一句话:通过注入“绝对位置”的操作,在注意力机制的内积计算中自然地体现出“相对位置”的依赖。

以下是其设计推导的逻辑脉络与工程考量:

1. 核心动机:寻找相对位置的完美表达

在 Transformer 的自注意力机制中,核心计算是 Query 和 Key 的内积 $\langle q_m, k_n \rangle$

如果不加任何位置编码,这个内积只和 token 本身的语义有关,对位置是“盲”的。

RoPE 的作者(苏剑林等)提出了一个设想:我们能不能找到一种位置编码函数 $f(\cdot, \cdot)$,使得带有位置 $m$ 的 Query 和带有位置 $n$ 的 Key 在做内积时,其结果仅仅依赖于这两个 token 的原始特征以及它们的相对位置 $m-n$

用数学语言表达,就是我们希望找到一个函数 $f$,满足:

$$\langle f_q(q, m), f_k(k, n) \rangle = g(q, k, m-n)$$

2. 数学推导:复数域的降维打击

为了求解上述等式,作者巧妙地借助了二维复数平面

如果我们把二维特征向量 $(x_0, x_1)$ 看作复数 $x = x_0 + i x_1$,那么上述的内积可以等价为复数乘法(取实部):

$$\langle q, k \rangle = \text{Re}(q \cdot k^*)$$

(注:$k^$ 是 $k$ 的共轭复数)*

结合我们的目标,如果我们将函数 $f$ 定义为在复平面上乘以一个与位置 $m$ 相关的旋转因子 $e^{im\theta}$

$$f_q(q, m) = q e^{im\theta}$$

$$f_k(k, n) = k e^{in\theta}$$

那么,它们在复数域的内积奇迹般地化简了:

$$\langle f_q(q, m), f_k(k, n) \rangle = \text{Re}\left( (q e^{im\theta}) \cdot (k e^{in\theta})^* \right)$$

$$= \text{Re}\left( q e^{im\theta} \cdot k^* e^{-in\theta} \right)$$

$$= \text{Re}\left( q k^* e^{i(m-n)\theta} \right)$$

结论:内积的结果完美地包含且仅包含了相对位置信息 $m-n$

将复数域的运算退回我们熟悉的实数域矩阵表示,复数乘法 $e^{im\theta}$ 就等价于乘以一个二维旋转矩阵:

$$\begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \ \sin(m\theta) & \cos(m\theta) \end{pmatrix}$$

这就是您图片中公式的根本来源。

3. 高维度的扩展与设计考量

在实际的大模型中,特征维度 $d$ 远大于 2。RoPE 的设计考量非常精妙:

  • 两两分组的正交子空间: RoPE 将 $d$ 维特征空间划分为 $d/2$ 个二维子空间,在每个二维平面上分别应用上述的旋转矩阵。
  • 频率的指数衰减(Distance Penalty): 对于不同的二维平面,设计了不同的旋转基频 $\theta_i$。通常采用 $\theta_i = 10000^{-2i/d}$。这意味着:
    • 靠前的特征维度(低频部分)旋转非常快,对局部位置变化极度敏感。
    • 靠后的特征维度(高频部分)旋转极其缓慢,捕捉全局的长距离依赖。
    • 远程衰减性:这种不同频率的组合使得随着相对距离 $m-n$ 的增加,内积的期望值呈现衰减趋势。这赋予了模型一种自然的“距离感”,距离越远的 token 注意力权重越倾向于下降。
  • 工程计算的高效性: 虽然公式里写的是矩阵乘法,但在实际 PyTorch 等框架的底层实现中,这其实是一个极度稀疏的块对角矩阵。我们根本不需要做 $O(d^2)$ 的稠密矩阵乘法,而是利用交换元素符号并相乘的技巧(例如 [-q1, q0]),将时间复杂度降到了 $O(d)$

RoPE 完美地统一了绝对位置和相对位置,这也是为什么 LLaMA、Qwen 等当今主流大模型几乎清一色采用 RoPE 的原因。

在学术界和工业界,关于 RoPE 目前最火热的研究方向其实是长度外推(Length Extrapolation),即如何在短文本上训练,但在长文本上推理(例如基于 NTK-Aware Scaling 或是 YaRN 的旋转角度缩放)。