Transformer组件(三):位置编码
paperreading
本文字数:9.2k 字 | 阅读时长 ≈ 39 min

Transformer组件(三):位置编码

paperreading
本文字数:9.2k 字 | 阅读时长 ≈ 39 min

1. 为什么需要位置编码?

1.1 一个直观的例子

假设:没有位置编码,下面这句话

“The cat sat on the mat.”(猫坐在垫子上。)

在 Transformer 的自注意力机制中,这些单词会被映射到词向量,例如:

"The"  -> [0.1, 0.3, 0.5, ...]
"cat"  -> [0.2, 0.4, 0.6, ...]
"sat"  -> [0.3, 0.5, 0.7, ...]
"on"   -> [0.4, 0.6, 0.8, ...]
"the"  -> [0.1, 0.3, 0.5, ...]   (与第一个 "The" 相同)
"mat." -> [0.5, 0.7, 0.9, ...]

在没有位置编码的情况下,Transformer 只能基于词的语义计算自注意力,而无法区分 “The” 是句首的 “The”,还是 “the mat” 里的 “the”,因为它们的向量是一样的。最终,模型无法理解单词之间的顺序关系。

在没有顺序概念的情况下,以下两个句子对模型来说(特别是计算attention的时候)可能是一样的:

  1. “The cat sat on the mat.”(猫坐在垫子上)
  2. “On mat the sat cat the.”(语法完全混乱)

为什么顺序错了对于 Transformer 来说没区别?

这可以从自注意力机制的计算方式来解释。当前词汇的最终注意力分数是所有词汇注意力分数的加权和,某个词汇对其他词汇的注意力分数是固定的,即不管顺序是什么,只要词汇的向量表示不变,那么最终的注意力分数就不会变。

1.2 加入位置编码

通过引入位置编码,我们可以让模型知道单词的顺序。例如,我们使用正弦位置编码:

"The"  (pos=0) -> [0.1, 0.3, 0.5, ...] + [0.99, 0.01, 0.98, ...] = [1.09, 0.31, 1.48, ...]
"cat"  (pos=1) -> [0.2, 0.4, 0.6, ...] + [0.87, 0.12, 0.94, ...] = [1.07, 0.52, 1.54, ...]
"sat"  (pos=2) -> [0.3, 0.5, 0.7, ...] + [0.76, 0.24, 0.91, ...] = [1.06, 0.74, 1.61, ...]
...

由于每个单词的向量都加上了不同的位置编码,这样他们的向量表示就不同了,现在 Transformer 能够区分 “The” 是在句首还是在句中,并能学习到句子的结构。

2. 位置编码的实现

位置编码分为绝对位置编码相对位置编码

绝对位置编码

相对位置编码

下面我们对这四个位置编码分别介绍

2.1 正余弦位置编码

常见的位置编码(Positional Encoding)——从 “Attention Is All You Need” 讲起

在 2017 年 Google 的论文 “Attention Is All You Need” 中,模型的核心机制是自注意力(Self-Attention),但自注意力本身是无序的。为了让 Transformer 能感知输入序列中 token 的顺序,作者引入了位置编码(Positional Encoding, PE)。论文使用的是固定的正弦-余弦位置编码(Sinusoidal Positional Encoding)

正弦-余弦位置编码公式

对于输入序列中第 pos 个 token,在隐藏维度 2i 和 2i+1 处的位置编码计算如下:

$$
\begin{aligned}
& PE(pos, 2i) = \sin \left( \frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}} \right) \\
& PE(pos, 2i+1) = \cos \left( \frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}} \right)
\end{aligned}
$$

其中:

优点:

例如,假设 $d_{\text{model}} = 4$,我们计算前 5 个位置的 PE:

pos PE(pos,0) (sin) PE(pos,2) (sin) PE(pos,4) (sin) PE(pos,6) (sin) PE(pos,8) (sin) PE(pos,10) (sin) PE(pos,12) (sin) PE(pos,14) (sin) PE(pos,16) (sin) PE(pos,18) (sin) PE(pos,20) (sin) PE(pos,1) (cos) PE(pos,3) (cos) PE(pos,5) (cos) PE(pos,7) (cos) PE(pos,9) (cos) PE(pos,11) (cos) PE(pos,13) (cos) PE(pos,15) (cos) PE(pos,17) (cos) PE(pos,19) (cos) PE(pos,21) (cos)
0 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1
1 0.841471 0.540302 0.821856 0.569695 0.801962 0.597375 0.781887 0.62342 0.76172 0.647906 0.74154 0.670909 0.721414 0.692504 0.701404 0.712764 0.681561 0.731761 0.661933 0.749563 0.642557 0.766238
2 0.909297 -0.416147 0.936415 -0.350895 0.958144 -0.286285 0.974888 -0.222695 0.987046 -0.160436 0.995011 -0.0997625 0.999164 -0.0408767 0.999871 0.0160656 0.99748 0.0709483 0.992321 0.12369 0.984703 0.174241
3 0.14112 -0.989992 0.245085 -0.969501 0.342782 -0.939415 0.433643 -0.901085 0.517306 -0.855801 0.593584 -0.804772 0.662436 -0.749118 0.723941 -0.689862 0.778273 -0.627927 0.825682 -0.564136 0.866477 -0.499217
4 -0.756802 -0.653644 -0.657167 -0.753745 -0.548606 -0.836081 -0.434205 -0.900814 -0.316715 -0.948521 -0.19853 -0.980095 -0.081685 -0.996658 0.032127 -0.999484 0.141539 -0.989933 0.245481 -0.969401 0.343152 -0.93928
5 -0.958924 0.283662 -0.993855 0.110692 -0.998229 -0.0594936 -0.975027 -0.222086 -0.927709 -0.373303 -0.859975 -0.510337 -0.77557 -0.631261 -0.678143 -0.73493 -0.571127 -0.820862 -0.457675 -0.88912 -0.340605 -0.940206
6 -0.279415 0.96017 -0.475221 0.879866 -0.644029 0.765001 -0.781498 0.623908 -0.885421 0.46479 -0.9554 0.295316 -0.992486 0.122357 -0.998839 -0.0481801 -0.977396 -0.211416 -0.931594 -0.363502 -0.865121 -0.501564
7 0.656987 0.753902 0.452392 0.891819 0.228775 0.973479 0.00062462 1 -0.21963 0.975583 -0.421997 0.906597 -0.599031 0.800726 -0.74573 0.666248 -0.859313 0.511449 -0.938902 0.344185 -0.985172 0.171572
8 0.989358 -0.1455 0.990673 0.136263 0.917358 0.398064 0.782276 0.622932 0.600822 0.799383 0.389156 0.921172 0.162824 0.986655 -0.0642208 0.997936 -0.280228 0.959933 -0.47594 0.879478 -0.644631 0.764494

在上述位置编码中,低维度变化快(高频),高维度变化慢(低频)

(1) 相同维度 i 下的不同 pos

(2) 公式解释
$$
PE(pos, 2i) = \sin \left( \frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}} \right)
$$

其中:

(3) 直观类比

  1. 为什么 Transformer 需要高低频编码?

Transformer 通过 低维高频和高维低频的组合,使得模型能够同时捕捉:
1. 短期依赖(短语或局部关系):由低维高频信号负责。
2. 长期依赖(句子或文档级别关系):由高维低频信号负责。

这种设计的好处是:

  1. 结论

✅ 低维度(小 i)变化快,是高频信号,适合建模短距离关系。
✅ 高维度(大 i)变化慢,是低频信号,适合建模长距离关系。
✅ Transformer 结合高低频信息,使得它能够兼顾局部和全局的位置信息,提高语言理解能力。

位置编码信息如何加入?

如下所示,在 Transformer 论文 “Attention Is All You Need” 中,正弦-余弦位置编码(Sinusoidal Positional Encoding) 是 加到输入词嵌入(word embeddings)之上的,然后输入到编码器(Encoder)和解码器(Decoder)。

Transformer 处理输入序列的流程如下:

  1. 将 token 转换为 word embeddings(维度为 $d_{\text{model}}$)。
  2. 加上位置编码(Positional Encoding, PE):
    $$
    X{\prime} = X + PE
    $$
  3. 将 $X{\prime}$ 送入 Transformer 编码器或解码器。

因此,位置编码是直接加在 word embeddings 上的,而不是作用于注意力计算或前馈网络。

为什么用加法?

2.2 可学习位置编码

可学习位置编码公式

可学习的位置编码使用 一个独立的可训练参数矩阵,类似于词嵌入(word embeddings):

$$
PE = \text{Embedding}(\text{position}, d_{\text{model}})
$$

加入位置?

在 Transformer 变体(如 BERT, GPT-2)中,可学习位置编码的加入方式如下:

Transformer 处理输入序列的流程如下:

  1. 将 token 转换为 word embeddings(维度为 $d_{\text{model}}$)。
  2. 加上可学习的位置编码(Positional Encoding, PE):
    $$
    X{\prime} = X + PE
    $$
  3. 将 $X{\prime}$ 送入 Transformer 编码器或解码器。

与正余弦位置编码的对比

可学习位置编码的优缺点

特点 正弦-余弦位置编码 可学习位置编码
是否可训练
是否能外推到更长序列
适用任务 Transformer (2017), 适用于标准 NLP 任务 BERT, GPT-2, 适用于固定长度任务
计算复杂度 低(固定计算) 略高(额外可训练参数)
学习灵活性

2.3 绝对位置编码的局限性

绝对位置编码只能够理解 token 在句子中的绝对位置,而无法理解 token 之间的关系。下面进行分析:

假设我们有两个句子:

从语义上看,这两个句子表达的是相同的意思(猫坐在垫子上),只是单词顺序不同。但在 Transformer 里:绝对位置编码 认为 “The” 在第 1 个位置,而 “On” 在第 1 个位置的情况下,它们的表示方式不同。而相对位置编码关注的是 token 之间的关系,而不是 token 在序列中的具体位置,因此能够更容易捕捉到两个句子的相似性

假设 Transformer 采用正弦-余弦位置编码(绝对位置编码),并为每个 token 赋予一个固定的位置编码:

Token 句子 1(位置) 句子 2(位置)
The 1 5
cat 2 6
sat 3 7
on 4 1
the 5 2
mat 6 3

在绝对位置编码下:

这样一来,Transformer 会认为这些两个句子完全不同,因为相同的词在不同的位置,它们的表示(embedding)也不同,从而影响句子整体的表征。尽管这两个句子具有相同的语义,绝对位置编码会导致模型学习到截然不同的特征表示。

使用相对位置编码改善这个问题

相对位置编码不直接编码 token 在句子中的绝对位置,而是编码 token 之间的相对关系。

引入相对位置,在相对位置编码的思路下,我们不关心 token 具体在第几位,而关心:

无论 token 的绝对位置如何变化,它们的相对关系是不变的,因此Transformer 仍然能够理解句子的结构。

示例:

Token 句子 1 中的相对位置 句子 2 中的相对位置
The → cat +1 +1
cat → sat +1 +1
sat → on +1 +1
on → the +1 +1
the → mat +1 +1

可以看到,在相对位置编码下,这两个句子的表示方式是完全相同的,这意味着:即使单词顺序改变,Transformer 仍然可以正确捕捉它们之间的关系,并理解两者语义相同

直观解释

想象一个人走在一条路上:

如果你换了一条更长的楼梯:

因此,相对位置编码让 Transformer 具备更好的泛化能力,能够在不同的句子结构中仍然保持对语义的理解。

总结

编码方式 是否受绝对位置影响 是否可以识别相似语义 适用场景
绝对位置编码 否(受位置影响) 适用于标准 Transformer
相对位置编码 是(关注 token 关系) 适用于长文本、结构变化的 NLP 任务

核心区别

相对位置编码的优势

2.4 相对位置编码

相对位置编码(Relative Positional Encoding, RPE) 旨在解决传统绝对位置编码(如正弦-余弦编码、可学习位置编码) 的局限性,即绝对位置编码仅能表示 token 在序列中的绝对位置,而无法直接捕捉 token 之间的相对关系。

核心思想

在自注意力计算时,引入相对位置偏差信息,让 attention 机制感知 token 之间的相对距离。下面先介绍标准的自注意力机制,然后介绍如何引入相对位置偏差信息。

标准自注意力机制(Self-Attention):
$$
A_{i, j} = \frac{(X_i W_q) (X_j W_k)^T}{\sqrt{d_k}}
$$
其中:

下面是几个引入相对位置编码的方法

(1)引入相对位置偏差:
$$
A_{i, j} = \frac{(X_i W_q) (X_j W_k)^T}{\sqrt{d_k}} + b_{i - j}
$$

(2)相对位置嵌入(Relative Positional Embeddings)

在 Attention 计算时,直接将相对位置嵌入(RPE)引入 Key 和 Query 的计算中:
$$
A_{i, j} = \frac{(X_i W_q) (X_j W_k + R_{i - j})^T}{\sqrt{d_k}}
$$

(3)Transformer-XL 的相对位置编码

Transformer-XL 采用可学习的相对位置编码,它的 attention 计算公式如下:
$$
A_{i, j} = \frac{(X_i W_q) (X_j W_k + R_{i - j})^T}{\sqrt{d_k}} + U X_i^T + V X_j^T
$$

(4)T5 的相对位置编码

T5 采用了一种 离散相对位置编码(Discretized Relative Positional Encoding),核心思想是:

优点:

  1. PyTorch 实现
import torch
import torch.nn as nn

class RelativePositionEncoding(nn.Module):
    def __init__(self, max_len, d_model):
        super(RelativePositionEncoding, self).__init__()
        self.relative_embeddings = nn.Embedding(2 * max_len, d_model)  # 允许负数位置
        self.max_len = max_len

    def forward(self, seq_len):
        position_ids = torch.arange(seq_len).unsqueeze(0)  # 生成绝对位置
        relative_position = position_ids - position_ids.T  # 计算相对位置
        relative_position += self.max_len  # 确保索引为正数
        return self.relative_embeddings(relative_position)  # 查表获取嵌入

max_len = 512
d_model = 768
seq_len = 10
relative_pe = RelativePositionEncoding(max_len, d_model)
relative_embedding_matrix = relative_pe(seq_len)

2.5 旋转位置编码

旋转位置编码(Rotary Positional Embedding, RoPE) 是一种高效的相对位置编码方法,用于增强 Transformer 的相对位置感知能力

旋转位置编码的数学原理

RoPE 的核心思想是利用旋转矩阵,让 token 的 Query 和 Key 在计算注意力时隐式编码相对位置信息。

(1)RoPE 如何作用于自注意力

在标准自注意力中,我们计算 Query-Key 之间的相似度:$A_{i,j} = Q_i K_j^T$,其中:

RoPE 通过对 Q 和 K 施加旋转变换,使得相对位置信息被编码进 Query 和 Key 本身,如下
$$
A_{i,j} = (R_{\theta_i} Q) \cdot (R_{\theta_j} K)^T
$$
其中:

(2)RoPE 旋转矩阵公式

假设 Query 和 Key 向量的维度为 d,我们对每个偶数维度上的分量(这里的 “对每个偶数维度上的分量进行旋转” 更准确的说法应该是 “成对的维度进行旋转”,即 RoPE 是将向量的相邻两个维度(通常是偶数索引和奇数索引的维度))作为一个二维平面进行旋转。
$$
R_{\theta} \begin{bmatrix} q_{2i} \ q_{2i+1} \end{bmatrix} =
\begin{bmatrix}
\cos(\theta) & -\sin(\theta) \\
\sin(\theta) & \cos(\theta)
\end{bmatrix}
\begin{bmatrix} q_{2i} \ q_{2i+1} \end{bmatrix}
$$
其中:

这样,我们可以在不增加额外参数的情况下,将相对位置信息直接编码进 Attention 计算

RoPE 的计算流程

RoPE 的计算可以分为以下几个步骤:
(1)计算旋转角度 $\theta$:
• 角度由 token 位置 pos 和维度索引 i 决定:
$$
\theta = \frac{pos}{10000^{2i/d}}
$$
(2)对 Query 和 Key 进行旋转变换:
• 每个偶数维度对的数值 ($q_{2i}, q_{2i+1}$) 进行旋转:
$$
\begin{aligned}
& q{\prime}_{2i} = q_{2i} \cos(\theta) - q_{2i+1} \sin(\theta) \\
& q{\prime}_{2i+1} = q_{2i} \sin(\theta) + q_{2i+1} \cos(\theta)
\end{aligned}
$$

(3)进行标准的 Attention 计算:
$$
A_{i,j} = (Q{\prime} W_q) (K{\prime} W_k)^T
$$
其中,$Q{\prime}$ 和 $K{\prime}$ 是旋转后的 Query 和 Key 向量。

RoPE 的 PyTorch 实现

import torch
import torch.nn as nn

def rotary_embedding(x, positions):
    """ 旋转位置编码 RoPE """
    dim = x.shape[-1] // 2  # 假设输入维度是偶数
    theta = 10000 ** (-2 * (torch.arange(dim, device=x.device) / dim))
    theta = positions[:, None] * theta  # 计算每个 token 的角度

    # 计算 cos 和 sin
    cos_theta = torch.cos(theta)
    sin_theta = torch.sin(theta)

    # 拆分偶数和奇数维度
    x1, x2 = x[..., ::2], x[..., 1::2]
    
    # 旋转
    x_rotated = torch.cat([x1 * cos_theta - x2 * sin_theta,
                           x1 * sin_theta + x2 * cos_theta], dim=-1)
    return x_rotated

# 示例
batch_size, seq_len, d_model = 2, 10, 512
positions = torch.arange(seq_len).expand(batch_size, -1)
x = torch.randn(batch_size, seq_len, d_model)

output = rotary_embedding(x, positions)

旋转矩阵公式的通俗理解

在 旋转位置编码(RoPE) 中,核心的数学操作是 旋转矩阵(Rotation Matrix),它的作用是让 Query 和 Key 向量在计算 Attention 之前按照特定的规则进行旋转,从而自然地编码相对位置信息。下面我们详细解析旋转矩阵的数学原理、作用、计算过程。

基本数学原理:给定一个 2D 旋转矩阵(Rotation Matrix):
$$
R_{\theta} = \begin{bmatrix} \cos\theta & -\sin\theta \\
\sin\theta & \cos\theta \end{bmatrix}
$$

这个矩阵的作用是将一个 2D 向量绕原点旋转角度 $\theta$,假设原始向量是:
$$
v = \begin{bmatrix} x \ y \end{bmatrix}
$$

应用旋转矩阵后,得到新的向量:
$$
v{\prime} = R_{\theta} v = \begin{bmatrix} \cos\theta & -\sin\theta \\
\sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x \ y \end{bmatrix}
$$

计算结果:
$$
x{\prime} = x \cos\theta - y \sin\theta
y{\prime} = x \sin\theta + y \cos\theta
$$

这表示:

在 Transformer 中,Query 和 Key 向量是高维的,我们希望它们能够编码 token 之间的相对位置信息。如果使用传统的位置编码(如 Learnable PE 或 Sinusoidal PE),位置信息和 token 语义是分离的,RoPE 通过旋转矩阵让位置编码与 token 语义紧密结合(RoPE 对 Query 和 Key 进行旋转变换,确保它们在 Attention 计算时隐式地包含相对位置关系),在 Attention 计算中隐式融入相对位置关系。

更简单的例子,我们讲解一个旋转90°的示例
$$
v = \begin{bmatrix} 1 \\ 0 \end{bmatrix}
$$

如果我们 顺时针旋转 90°($\theta = 90^\circ$):
$$
R_{90^\circ} = \begin{bmatrix} \cos 90^\circ & -\sin 90^\circ \\
\sin 90^\circ & \cos 90^\circ \end{bmatrix}
= \begin{bmatrix} 0 & -1 \\ 1 & 0 \end{bmatrix}
$$

计算:
$$
v{\prime} = R_{90^\circ} v =
\begin{bmatrix} 0 & -1 \\
1 & 0 \end{bmatrix} \begin{bmatrix} 1 \\ 0 \end{bmatrix}
= \begin{bmatrix} 0 \\ 1 \end{bmatrix}
$$

结果: (1, 0) 旋转 90° 后变成 (0, 1)。

基于上面理解,我们讲解一个旋转30°的示例

假设 两个 token 的 Query 向量如下:
• Token 1 的 Query 向量(不旋转):
$$
Q_1 = \begin{bmatrix} 1.0 \\ 0.5 \end{bmatrix}
$$
• Token 2 的 Query 向量(在 RoPE 作用下旋转 30°):
$$
Q_2 = R_{30^\circ} Q_1
$$
计算:
$$
R_{30^\circ} = \begin{bmatrix} \cos 30^\circ & -\sin 30^\circ \\
\sin 30^\circ & \cos 30^\circ \end{bmatrix} =
\begin{bmatrix} 0.866 & -0.5 \\
0.5 & 0.866 \end{bmatrix}
$$

应用到 $Q_1$:
$$
Q_2 = \begin{bmatrix} 0.866 & -0.5 \\
0.5 & 0.866 \end{bmatrix} \begin{bmatrix} 1.0 \\ 0.5 \end{bmatrix}
= \begin{bmatrix} (1.0 \times 0.866) + (0.5 \times -0.5) \\
(1.0 \times 0.5) + (0.5 \times 0.866) \end{bmatrix}
= \begin{bmatrix} 0.616 \\
0.933 \end{bmatrix}
$$

这样,Token 2 的 Query 向量经过旋转后,它与 Token 1 的相对关系就被编码进去了!

为什么 RoPE 能编码相对位置信息?

在 Transformer 的 Attention 计算 中,我们计算 Query 和 Key 的点积:
$$
A_{i,j} = Q_i K_j^T
$$

在 RoPE 作用下,$Q_i$ 和 $K_j$ 被旋转过不同的角度 $\theta$,但它们仍然可以保持相对关系:
$$
Q{\prime}_i \cdot K{\prime}_j = \cos(\theta_i - \theta_j) (Q_i \cdot K_j)
$$

这意味着:
• Token 之间的相对角度 $\theta_i - \theta_j$ 自然地编码了它们的相对位置!
• Attention 计算隐式地学习到了 token 之间的相对位置关系,而不需要额外的存储和参数。

一些疑问

在 RoPE 中,每个 token 的 Query 和 Key 向量都会被旋转到不同的方向,从而编码了它们之间的相对位置关系。但是,这种编码方式是否能够保留 token 的长度和形状(信息保留)?答案是肯定的,因为 RoPE 的旋转操作是线性的,不会改变向量的模长(长度)。

在 RoPE 中,我还有一个疑问就是,我们知道旋转之后的模长和方向,旋转过程中模长是不变的,也就是说我们也知道旋转之前的模长,那么旋转之前的方向模型是否知道?答案也是肯定的。

因为在 RoPE 中每个位置的旋转角度是固定的,旋转角度是根据位置计算出来的,所以模型知道旋转之前的方向。在 RoPE 中:每个序列位置 i 对应一个固定的旋转角度 $\theta_i$,对应的旋转矩阵 $R(\theta_i)$ 也是固定的。具体来说:

也就是说:每个 token 在特定位置 i 时,都会被 “固定旋转”成某个方向。而不同位置的 token 会被旋转到不同的方向(方向之间的差异体现了它们的相对位置)。这种设计让模型自然“知道” token 的顺序和位置差异。

我们可以总结成一句话:ROPE 中,旋转之前的方向和 token 的相对位置,都是隐式编码在旋转后的向量中的

解释一下:

总结:RoPE 是一种非常巧妙的隐式编码方式,通过旋转把 原始语义+位置信息 融合到一个向量中,不需要显式加位置向量。

下面是一个具体的代码,模拟了 ROPE 的旋转情况

import torch
import torch.nn.functional as F
import math
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tokens = torch.tensor([
    [1.0, 0.2],
    [0.3, 1.0],
    [-0.9, 0.5]
], device=device)

W_q = torch.tensor([[1.0, 0.8], [-0.8, 1.0]], device=device)
W_k = torch.tensor([[-1.0, 0.6], [0.6, 1.0]], device=device)
W_v = torch.eye(2, device=device)

Q = F.linear(tokens, W_q)
K = F.linear(tokens, W_k)
V = F.linear(tokens, W_v)

Q_before, K_before = Q.clone(), K.clone()

def apply_rope(QK, theta_list):
    N, D = QK.shape
    assert D % 2 == 0
    x = QK.clone()
    theta = torch.tensor(theta_list, device=QK.device)
    cos = torch.cos(theta)
    sin = torch.sin(theta)
    x_ = x.view(N, D // 2, 2)
    x_new = torch.zeros_like(x_)
    x_new[:, :, 0] = cos.unsqueeze(-1) * x_[:, :, 0] - sin.unsqueeze(-1) * x_[:, :, 1]
    x_new[:, :, 1] = sin.unsqueeze(-1) * x_[:, :, 0] + cos.unsqueeze(-1) * x_[:, :, 1]
    return x_new.view(N, D)

theta_list_deg = [10, 20, 30]
theta_list_rad = [math.radians(t) for t in theta_list_deg]

Q_after = apply_rope(Q, theta_list_rad)
K_after = apply_rope(K, theta_list_rad)

def draw_vectors_v4(Q_before, Q_after, K_before, K_after, theta_deg):
    q_colors = ['#E74C3C', '#C0392B', '#FF7F50']
    k_colors = ['#3498DB', '#1F618D', '#76D7C4']
    plt.figure(figsize=(7, 7))
    plt.axhline(0, color='gray', linestyle='--')
    plt.axvline(0, color='gray', linestyle='--')
    plt.arrow(-100, -100, 0.3, 0.0, head_width=0, color=q_colors[0], linestyle='--', alpha=0.8, label='Q before')
    plt.arrow(-100, -100, 0.3, 0.0, head_width=0, color=q_colors[0], linestyle='-', alpha=1, label='Q after')
    plt.arrow(-100, -100, 0.3, 0.0, head_width=0, color=k_colors[0], linestyle='--', alpha=0.7, label='K before')
    plt.arrow(-100, -100, 0.3, 0.0, head_width=0, color=k_colors[0], linestyle='-', alpha=0.95, label='K after')
    for i in range(Q_before.shape[0]):
        plt.arrow(0, 0, Q_before[i,0].item(), Q_before[i,1].item(),
                  head_width=0.06, color=q_colors[i], linestyle='--', alpha=0.8)
        plt.arrow(0, 0, Q_after[i,0].item(), Q_after[i,1].item(),
                  head_width=0.06, color=q_colors[i], linestyle='-', alpha=1)
        plt.arrow(0, 0, K_before[i,0].item(), K_before[i,1].item(),
                  head_width=0.04, color=k_colors[i], linestyle='--', alpha=0.7)
        plt.arrow(0, 0, K_after[i,0].item(), K_after[i,1].item(),
                  head_width=0.04, color=k_colors[i], linestyle='-', alpha=0.95)
        plt.text(Q_after[i,0].item(), Q_after[i,1].item(),
                 f'Q{i}\n({theta_deg[i]}°)', fontsize=10, ha='left', va='bottom', color=q_colors[i], weight='bold')
        plt.text(K_after[i,0].item(), K_after[i,1].item(),
                 f'K{i}\n({theta_deg[i]}°)', fontsize=10, ha='right', va='top', color=k_colors[i], weight='bold')
    plt.xlim(-2, 2)
    plt.ylim(-1, 1.5)
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('Q and K Vectors: before & after ROPE')
    plt.legend(loc='upper right')
    plt.grid()
    plt.show()

draw_vectors_v4(Q_before, Q_after, K_before, K_after, theta_list_deg)

d_k = Q.shape[-1]
scores = Q_after @ K_after.T / math.sqrt(d_k)
attn_weights = F.softmax(scores, dim=-1)
output = attn_weights @ V

print("Attention Weights:\n", attn_weights)
print("Self-Attention Output:\n", output)

RoPE 的旋转情况

2.6 🌟 RoPE原文解读 🌟

上面的 RoPE 我们并没有结合原文,这里结合原文进行解读一下

📚 1️⃣ 背景:为什么需要 RoPE?

Transformer 中,序列的顺序信息(即位置)必须编码进模型:

📚 2️⃣ RoPE 的设计思想:乘法旋转编码

RoPE 不采用加法,而是通过 旋转矩阵乘法 将位置信息融入内容:

📚 3️⃣ 高维空间推广

对于偶数维度 $d$,RoPE 将向量拆成 $d/2$ 对:
$$
R^\Theta_{d,m} = \text{block-diag}(R(m\theta_1), R(m\theta_2), \ldots, R(m\theta_{d/2}))
$$
其中:$\theta_i = 10000^{-2(i-1)/d}$

最终:$q_m^T k_n = x_m^T Trans((W^q)) R^\Theta_{d,n-m} W^k x_n$

RoPE 的高维引入

✅ 避免内容和位置混淆(乘法而非加法)
✅ 支持不同序列长度、线性Attention
✅ 解释性强:旋转模拟相对距离衰减

与传统加法位置编码对比如下:

方法 原理 相对位置信息 长文本支持
sinusoidal 加法 隐含 不灵活
learnable 加法 隐含 不灵活
RoPE 乘法(旋转) 自然引入 ( n-m ) 灵活支持

RoPE 的旋转情况

2.7 RoPE 拓展到2维

一维 RoPE(1D RoPE)

对一维序列中的位置 m,n:
$$
\begin{aligned}
f_q(x_m, m) = (W^q x_m) e^{i m \theta} \\
f_k(x_n, n) = (W^k x_n) e^{i n \theta}
\end{aligned}
$$
其中:

内积(注意 RoPE 引入了位置差):
$$
g(x_m, x_n, m-n) = \Re\big[(W^q x_m) (W^k x_n)^* e^{i(m-n)\theta}\big]
$$

$(W^q x_m) e^{im\theta}$ 对 Q 向量的旋转编码,$(W^k x_n) e^{in\theta}$ 对 K 向量的旋转编码,内积中相乘结果引入了相对位置 $m-n$

二维 RoPE(2D RoPE)

对二维位置 $(x_m, y_m)$、$(x_n, y_n)$:
$$
\begin{aligned}
f_q(x_m, y_m) = (W^q x_m) e^{i (x_m \theta_x + y_m \theta_y)} \\
f_k(x_n, y_n) = (W^k x_n) e^{i (x_n \theta_x + y_n \theta_y)}
\end{aligned}
$$
其中:

内积公式:
$$
g(x_m,x_n,x_m-x_n,y_m-y_n) = \Re\big[(W^q x_m)(W^k x_n)^* e^{i((x_m-x_n) \theta_x + (y_m-y_n) \theta_y)}\big]
$$

解释:

总结对比表

特性 一维 RoPE 二维 RoPE
位置索引 $m, n$ $x_m,y_m, x_n,y_n$
Q 编码 $W^q x_m e^{i m \theta}$ $W^q x_m e^{i (x_m \theta_x + y_m \theta_y)}$
K 编码 $W^k x_n e^{i n \theta}$ $W^k x_n e^{i (x_n \theta_x + y_n \theta_y)}$
内积公式 $\Re[(W^q x_m)(W^k x_n)^* e^{i(m-n)\theta}]$ $\Re[(W^q x_m)(W^k x_n)^* e^{i((x_m-x_n)\theta_x + (y_m-y_n)\theta_y)}]$
相对位置依赖 $m-n$ $x_m-x_n, y_m-y_n$

优势

在 RoPE 中,$\theta$ 其实不是一个固定值,而是 对不同 embedding 维度采用不同频率的角度因子。

对于 embedding 维度 d:

因此,对于一维序列中的位置 m,不同维度上:
$$
m \theta_i = m \cdot \frac{1}{10000^{2i/d}}
$$
这表示位置 m 对不同维度的旋转角度是 不同的(低维旋转快,高维旋转慢)。

$\theta$ 大小的确定公式(详解)

更准确地,一维 RoPE 的不同维度上:
$$
\theta_i = \frac{1}{10000^{2i/d}} = e^{-\frac{2i}{d} \log 10000}
$$

所以:
$$
m\theta_i = m \cdot e^{-\frac{2i}{d}\log 10000}
$$

这体现了:

二维 RoPE 中的 $\theta_x,\theta_y$

对于二维 RoPE,水平和垂直方向也采用类似的频率定义:
$$
\theta_{x,i} = \frac{1}{10000^{2i/d_x}}, \quad \theta_{y,i} = \frac{1}{10000^{2i/d_y}}
$$
其中 $d_x,d_y$ 是 embedding 中分配给 x、y 的维度(一般各取 embedding 的一半)。

因此:
$$
x_m\theta_{x,i} = x_m \cdot \frac{1}{10000^{2i/d_x}}, \quad y_m\theta_{y,i} = y_m \cdot \frac{1}{10000^{2i/d_y}}
$$

总结

特性 公式
一维 $\theta_i$ $ \theta_i = \frac{1}{10000^{2i/d}}$
一维旋转角 $m\theta_i = m \cdot \frac{1}{10000^{2i/d}}$
二维 $\theta_{x,i},\theta_{y,i}$ $ \theta_{x,i}=\frac{1}{10000^{2i/d_x}}, \theta_{y,i}=\frac{1}{10000^{2i/d_y}}$
二维旋转角 $x_m \theta_{x,i}, y_m \theta_{y,i}$

3. 各种位置编码方法对比

下表列出了常见的 位置编码方式(Positional Encoding)及其优缺点,包括 正弦-余弦位置编码(Sinusoidal PE)、可学习位置编码(Learnable PE)、相对位置编码(Relative PE) 和 旋转位置编码(RoPE)。

位置编码方法对比表

位置编码方式 是否可训练 是否支持长序列外推 是否编码相对位置 计算复杂度 存储开销 优点 缺点 适用场景
正弦-余弦位置编码 (Sinusoidal PE) ❌ 否 ✅ 是 ❌ 否 ($O(n)$) 无额外存储 - 无需训练,位置编码固定
- 支持长序列外推,适用于不同长度的输入
- 无法编码相对位置信息,仅编码绝对位置
- 泛化能力有限,对长文本推理可能有影响
标准 Transformer,适用于自然语言处理 (NLP)、计算机视觉 (CV) 任务
可学习位置编码 (Learnable PE) ✅ 是 ❌ 否 ❌ 否 ($O(n)$) 需要额外参数 ($O(d \cdot n)$) - 可以优化以适应特定任务,能捕捉训练数据的特定模式
- 简单易实现
- 不能外推到更长的序列,需要重新训练
- 不能编码相对位置信息
适用于固定长度的输入,如 BERT 预训练,文本分类任务
相对位置编码 (Relative PE) ✅ 是 ✅ 是 ✅ 是 中等 ($O(n^2)$) 需要存储相对位置信息 ($O(n^2)$) - 能编码相对位置,适合长文本建模
- 泛化能力强,适用于不同长度输入
- 能增强模型的局部模式学习能力
- 计算复杂度较高,比绝对位置编码更慢
- 存储需求较大,需要存储相对位置信息
适用于需要相对位置信息的任务,如 NLP 机器翻译 (MT)、长文本理解
旋转位置编码 (RoPE) ❌ 否 ✅ 是 ✅ 是 ($O(n)$) 无额外存储 - 高效,计算复杂度低
- 支持长文本外推,可处理比训练时更长的序列
- 直接作用于 Attention,不会额外增加存储
- 依赖 Query 和 Key 旋转计算,较难直观理解
- 对特定任务需要额外调优
适用于 LLM(大语言模型),如 GPT 系列、Qwen、ChatGLM,尤其是需要处理长序列任务