1. 原理讲解#

假设我们在二维平面上有一个向量q=(x, y), 如果想要把它逆时针旋转一个角度 $\theta$，在线性代数中，这相当于左乘一个二维旋转矩阵 $R(\theta)$。 但这个矩阵具体长什么样？

我们高中的时候学过，在二维平面中，只需要两个向量就可以把平面内的任意向量表示出来，例如X和Y轴的单位向量，他们在矩阵论中被叫做基向量。

实际上只要确定两个基向量旋转后落在哪里，平面内任意向量都可以通过左乘一个旋转向量来解决。

假设我们用的是二维直角坐标系，我们定义两个单位向量分别是：

• $e_1 = (1, 0)^T$
• $e_2 = (0, 1)^T$

先看 $e_1 = (1, 0)^T$，它指向 x 轴正方向。逆时针转 $\theta$ 后，落点仍在单位圆上，对应的角度就是 $\theta$。单位圆上角度为 $\theta$ 的点的坐标就是：

再看 $e_2 = (0, 1)^T$，它指向 y 轴正方向，本身就是 $90^\circ$ 方向。再逆时针转 $\theta$，落点就在 $90^\circ + \theta$ 的角度上：

两个基向量的像都求出来了，按列排开，第一列放 $e_1$ 的像，第二列放 $e_2$ 的像：

向量左乘这个矩阵，就可以得到旋转后的新的向量 $q'$:

2. RoPE在大模型的示例讲解#

现在，我们有一段文本，每个 Token 都有一个物理上的绝对位置编号 $m$（比如第 $m=0, 1, 2...$ 个词）。

RoPE 的做法是我在第 $m$ 个位置上，就把向量旋转 $m \cdot \theta$ 这么多角度。 也就是说，向量在这个句子中越靠后，它被旋转的角度就越大。

此时，位置 $m$ 处的旋转矩阵变成了：

把 Query 向量放在位置 $m$ 旋转，把 Key 向量放在位置 $n$ 旋转：

• $\tilde{q}_m = R(m\theta) \cdot q$
• $\tilde{k}_n = R(n\theta) \cdot k$

到这一步为止，我们赋予 $q$ 和 $k$ 的都只是绝对位置信息。

Transformer的核心是Attention，Attention的核心就是计算Query和Key的内积来决定注意力分数。

我们算一下，带有绝对位置 $m$ 的 $\tilde{q}_m$ 和带有绝对位置 $n$ 的 $\tilde{k}_n$ 的内积是什么：

根据旋转矩阵的几何性质，一个旋转矩阵的转置，等于反向旋转，即 $R(m\theta)^T = R(-m\theta)$。 而两个旋转矩阵相乘，等于它们的角度相加：

所以，上面的内积公式化简为：

可以看到，现在它们只和两个Token的相对距离有关。

我们接着来思考多维向量的情况。 刚才讲的都是 2 维向量。但大模型里，每个 Attention 头（Head）的维度通常是 128 维（记为 head_size），而参与旋转的 rotary_dim 通常就等于它。

RoPE 的做法是把 rotary_dim 个维度按前后对半分成两段（各 rotary_dim/2 维），让前半的第 $i$ 维与后半的第 $i$ 维配成一对，并把这一对放到它们张成的二维平面里旋转 $\theta_i$。

以 rotary_dim = 128 为例，前 64 维与后 64 维一一配对：

• 第 $0$ 维与第 $64$ 维配成一对，按 $\theta_0$ 旋转；
• 第 $1$ 维与第 $65$ 维配成一对，按 $\theta_1$ 旋转；
• …
• 第 $63$ 维与第 $127$ 维配成一对，按 $\theta_{63}$ 旋转。

这 rotary_dim/2 个二维旋转子空间互相独立，整体相当于一个 rotary_dim 维的块对角旋转矩阵.只不过这些 2×2 块在矩阵中的位置是前后错开的，而不是相邻的。

不同的维度，旋转的速度（频率）不一样 每个维度对用自己的频率，RoPE 沿用原版 Transformer 的频率衰减公式：

代码实现如下：

1
inv_freq = 1.0 / (base ** (arange(0, rotary_dim, 2) / rotary_dim))

• 索引 $i$ 小的维度对：$\theta_i$ 接近 $1$，位置每加 1 角度就大幅前进，转得极快。这就像时钟的秒针，稍微错开一个位置，角度就差很多，用来感知局部、近距离的依赖关系。
• 索引 $i$ 大的维度对：$\theta_i$ 接近 $0$，越过成百上千个 Token 角度也没怎么变，转得很慢。这就像时钟的时针，用来感知全局、长距离的上下文关系。

回过头再看RoPE的代码，就可以很好的理解。

3. 代码实现#

整个模块只有三部分：__init__ 预计算查找表，_rotate 做二维旋转，forward 查表并把旋转施加到 Q/K 上。

1
import torch
2
from torch import nn
3

4
class RotaryEmbedding(nn.Module):
5
    def __init__(self, head_size: int, rotary_dim: int, max_position_embedding: int, base: float) -> None:
6
        super(RotaryEmbedding, self).__init__()
7
        if rotary_dim % 2 != 0:
8
            raise ValueError("rotary_dim must be even")
9
        self.head_size = head_size
10
        self.rotary_dim = rotary_dim
11

12
        inv_freq = 1.0 / (
13
            base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)
14
        )
15
        t = torch.arange(max_position_embedding, dtype=torch.float)
16
        freqs = torch.outer(t, inv_freq)
17
        self.register_buffer(
18
            "cos_sin_cache",
19
            torch.cat((freqs.cos(), freqs.sin()), dim=-1),
20
            persistent=False,
21
        )

1
def _rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
2
    x1, x2 = x.chunk(2, dim=-1)
3
    return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)

1
    def forward(self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
2
        flat = positions.reshape(-1)
3
        cos_sin = self.cos_sin_cache.index_select(0, flat).to(dtype=query.dtype)
4
        cos, sin = cos_sin.chunk(2, dim=-1)
5
        cos = cos.unsqueeze(-2)
6
        sin = sin.unsqueeze(-2)
7

8
        def _apply(x: torch.Tensor) -> torch.Tensor:
9
            shape = x.shape
10
            x = x.reshape(flat.numel(), -1, self.head_size)
11
            x_rot = _rotate(x[..., : self.rotary_dim], cos, sin)
12
            x_pass = x[..., self.rotary_dim:]
13
            x_out = torch.cat([x_rot, x_pass], dim=-1)
14
            return x_out.reshape(shape)
15

16
        return _apply(query), _apply(key)

1
if __name__ == "__main__":
2
    head_size = 64
3
    rotary_dim = 32
4
    max_pos = 128
5
    base = 10000.0
6
    rope = RotaryEmbedding(head_size, rotary_dim, max_pos, base)
7

8
    q = torch.randn(1, 8, 4, head_size)
9
    k = torch.randn_like(q)
10
    pos = torch.arange(8).unsqueeze(0)
11

12
    q_rot, k_rot = rope.forward(pos, q, k)
13
    print("q shape", q_rot.shape)
14
    print("k shape", k_rot.shape)