您本可以设计出最先进的位置编码

加尔定律
一个能工作的复杂系统，总是从一个能工作的简单系统演变而来。
约翰·加尔

本文将逐步引导您了解Transformer模型中最先进的位置编码的发现过程。我们将通过迭代改进位置编码方法来实现这一目标，最终得到最新LLama 3.2版本和大多数现代Transformer中使用的Rotary Positional Encoding (RoPE)。本文旨在限制所需的数学知识，但需要一些基本的线性代数、三角学和自注意力机制的理解。

问题陈述

知其所伴，便知其义
约翰·鲁珀特·弗斯

与所有问题一样，最好先精确地理解我们试图实现的目标。Transformer中的自注意力机制用于理解序列中词元之间的关系。自注意力是一种集合操作，这意味着它具有置换等变性。如果我们不为自注意力机制提供位置信息，许多重要的关系将无法确定。

这通过一个例子可以很好地说明。

启发性示例

考虑以下句子，其中同一个词出现在不同位置

$$\text{The dog chased another dog}$$

直观上，“dog”指的是两个不同的实体。让我们看看如果先对它们进行词元化，映射到Llama 3.2 1B的真实词元嵌入，并通过torch.nn.MultiheadAttention会发生什么。

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
    for param in mha.parameters():
        nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True

正如我们所看到的，在没有任何位置信息的情况下，（多头）自注意力操作的输出对于不同位置的相同词元是相同的，尽管这些词元显然代表着不同的实体。现在，让我们开始设计一种方法，用位置信息增强自注意力，使其能够根据词的位置确定它们之间的关系。

理想的位置编码方案应该如何表现？

理想特性

让我们尝试定义一些理想的特性，这将使优化过程尽可能简单。

特性1 - 每个位置的唯一编码（跨序列）

每个位置都需要一个唯一的编码，无论序列长度如何，该编码都应保持一致——在位置5的词元，无论当前序列长度是10还是10,000，都应该具有相同的编码。

特性2 - 两个编码位置之间的线性关系

位置之间的关系应该在数学上是简单的。如果我们知道位置$p$的编码，那么计算位置$p+k$的编码就应该很简单，这使得模型更容易学习位置模式。

如果您考虑我们如何在数轴上表示数字，就很容易理解5距离3有2步，或者10距离15有5步。我们的编码中应该存在相同的直观关系。

特性3 - 推广到比训练中遇到的序列更长的序列

为了增加模型在现实世界中的实用性，它们应该能够泛化到训练分布之外的数据。因此，我们的编码方案需要足够灵活，能够处理意外的输入长度，而不会违反任何其他理想特性。

特性4 - 由模型可以学习的确定性过程生成

如果我们的位置编码能够从确定性过程中获取，那将是理想的。这将允许模型有效地学习我们的编码方案背后的机制。

特性5 - 可扩展到多个维度

随着多模态模型成为常态，我们的位置编码方案能够自然地从$1D$扩展到$nD$至关重要。这将使模型能够处理图像或脑部扫描等数据，它们分别是$2D$和$4D$的。

现在我们知道了理想的特性（以下称为$Pr_n$），让我们开始设计和迭代我们的编码方案。

整数位置编码

首先想到的方法可能是简单地将词元位置的整数值添加到词元嵌入的每个分量中，其值范围从$0 \rightarrow L$，其中$L$是当前序列的长度。

在上面的动画中，我们为词元$\color{#699C52}\text{chased}$创建了位置编码向量，并将其添加到我们的词元嵌入中。这里的嵌入值是Llama 3.2 1B真实值的一个子集。我们可以观察到它们集中在0附近。这对于避免训练期间梯度消失或爆炸是理想的，因此是我们在整个模型中希望保持的。

显然，我们目前这种朴素的方法将导致问题。位置值的幅度远远超过我们输入的实际值。这意味着信噪比非常低，模型很难将语义信息与位置信息分离。

有了这些新知识，一个自然的后续想法可能是将位置值标准化为$\frac{1}{N}$。这会将值限制在0到1之间，但引入了另一个问题。如果我们将$N$选为当前序列的长度，那么对于不同长度的每个序列，位置值将完全不同，这违反了$Pr_1$。

有没有更好的方法来确保我们的数字在0和1之间？如果我们对此深思熟虑一段时间，我们可能会想到从十进制数切换到二进制数。

二进制位置编码

我们不必将（可能已标准化的）整数位置添加到嵌入的每个分量中，而是可以将其转换为二进制表示，并将其拉伸以匹配我们的嵌入维度，如下所示。

我们已将感兴趣的位置（252）转换为其二进制表示（11111100），并将每个位添加到词元嵌入的相应分量中。最低有效位（LSB）将在每个后续词元之间在0和1之间循环，而最高有效位（MSB）将在每$2^{n-1}$个词元循环，其中$n$是位数。您可以在下面的动画中看到不同索引的位置编码向量$[^1]$。

我们已经解决了值范围问题，现在我们有了在不同序列长度下保持一致的唯一编码。如果我们将词元嵌入的低维版本绘制出来，并可视化二进制位置向量对于不同值的添加会发生什么。

我们可以看到结果非常“跳跃”（正如我们可能从二进制的离散性质中预期的那样）。优化过程喜欢平滑、连续和可预测的变化。我们知道有哪些具有相似值范围的函数是平滑和连续的吗？

如果我们四处看看，我们可能会注意到$\sin$和$\cos$都符合要求！

正弦位置编码

上述动画可视化了我们的位置嵌入，其中每个分量都交替从$\sin$和$\cos$中提取，并且波长逐渐增加。如果将其与之前的动画进行比较，您会发现惊人的相似之处！

我们现在得到了正弦嵌入；最初定义在Attention is all you need论文中。让我们看看这些方程：

$$PE_{(pos,2i)} = \color{#58C4DD}\sin\left(\color{black}\frac{pos}{10000^{2i/d}}\color{#58C4DD}\right)\color{black} \\ \quad \\ PE_{(pos,2i+1)} = \color{#FC6255}\cos\left(\color{black}\frac{pos}{10000^{2i/d}}\color{#FC6255}\right)\color{black} \\$$

其中$pos$是词元的位置索引，$i$是位置编码向量中的分量索引，$d$是模型维度。$10,000$是基准波长（以下简称$\theta$），我们根据分量索引对其进行拉伸或压缩。我鼓励您代入一些实际值，以感受这种几何级数。

这个方程中有几部分乍一看令人困惑。作者是如何选择$10,000$的？为什么我们对偶数位置使用$\sin$，对奇数位置使用$\cos$呢？

似乎使用$10,000$作为基准波长是实验确定的$[^2]$。解释$\sin$和$\cos$共同使用的原因更为复杂，但对于我们理解的迭代方法至关重要。这里的关键是我们希望两个编码位置之间存在线性关系$Pr_2$。要理解$\sin$和$\cos$如何协同产生这种线性关系，我们必须深入研究一些三角学。

考虑一个正弦和余弦对序列，每个对都与频率$\omega_i$相关联。我们的目标是找到一个线性变换矩阵$\mathbf{M}$，它可以将这些正弦函数平移一个固定的偏移量$k$

$$\mathbf{M} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i(p + k)) \\ \cos(\omega_i(p + k)) \end{bmatrix}$$

频率$\omega_i$遵循几何级数，其随维度索引$i$而减小，定义为

$$\omega_i = \frac{1}{10000^{2i/d}}$$

为了找到这个变换矩阵，我们可以将其表示为一个具有未知系数$u_1$、$v_1$、$u_2$和$v_2$的通用2×2矩阵

$$\begin{bmatrix} u_1 & v_1 \\ u_2 & v_2 \end{bmatrix} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i(p+k)) \\ \cos(\omega_i(p+k)) \end{bmatrix}$$

通过将三角和定理应用于右侧，我们可以将其展开为

$$\begin{bmatrix} u_1 & v_1 \\ u_2 & v_2 \end{bmatrix} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i p)\cos(\omega_i k) + \cos(\omega_i p)\sin(\omega_i k) \\ \cos(\omega_i p)\cos(\omega_i k) - \sin(\omega_i p)\sin(\omega_i k) \end{bmatrix}$$

通过展开，我们可以通过匹配系数得到一个包含两个方程的系统

$$\begin{align} u_1\sin(\omega_i p) + v_1\cos(\omega_i p) &= \cos(\omega_i k)\sin(\omega_i p) + \sin(\omega_i k)\cos(\omega_i p) \\ u_2\sin(\omega_i p) + v_2\cos(\omega_i p) &= -\sin(\omega_i k)\sin(\omega_i p) + \cos(\omega_i k)\cos(\omega_i p) \end{align}$$

通过比较两边包含 $\sin(\omega_i p)$ 和 $\cos(\omega_i p)$ 的项，我们可以解出未知系数

$$\begin{align} u_1 &= \cos(\omega_i k) & v_1 &= \sin(\omega_i k) \\ u_2 &= -\sin(\omega_i k) & v_2 &= \cos(\omega_i k) \end{align}$$

这些解给了我们最终的变换矩阵 $\mathbf{M_k}$

$$\mathbf{M_k} = \begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix}$$

如果你之前做过游戏编程，你可能会注意到我们推导的结果出奇地熟悉。没错，它就是旋转矩阵！ $[^3]$。

因此，Noam Shazeer 在2017年的论文《Attention is all you need》中设计的编码方案，早在2017年就已经将相对位置编码为旋转了！尽管旋转的概念早已出现，但从正弦编码到RoPE，却又花了**四年**的时间……

绝对位置编码与相对位置编码

了解旋转在此处的重要性后，让我们回到最初的例子，尝试为下一次迭代寻找一些直觉。

$$\begin{align*} &\hspace{0.7em}0 \hspace{1.4em} 1 \hspace{2em} 2 \hspace{2.6em} 3 \hspace{2.4em} 4\\ &\text{The dog chased another dog} \\ \\ &\hspace{0.3em}\text{-2} \hspace{1.4em} \text{-1} \hspace{1.7em} 0 \hspace{2.6em} 1 \hspace{2.4em} 2\\ &\text{The dog \color{#699C52}chased \color{black}another dog} \end{align*}$$

上面，我们可以看到标记的绝对位置，以及从$\color{#699C52}\text{chased}$到其他每个标记的相对位置。通过正弦编码，我们生成了一个单独的向量来表示绝对位置，并使用一些三角函数技巧来编码相对位置。

当我们试图理解这些句子时，这个词是这篇博文中的第2157个词重要吗？还是我们更关心它与周围词语的关系？一个词的绝对位置很少影响其含义——重要的是词语之间的相互关系。

上下文中的位置编码

从现在开始，关键是要**在自注意力机制的上下文中**考虑位置编码。重申一下，自注意力机制使模型能够权衡输入序列中不同元素的重要性，并动态调整它们对输出的影响。

$$\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

在之前的所有迭代中，我们都生成了一个单独的位置编码向量，并在进行$Q$、$K$和$V$投影之前将其**添加**到标记嵌入中。通过将位置信息直接添加到标记嵌入中，我们**污染**了语义信息与位置信息。我们应该尝试在不修改范数的情况下编码信息。转向乘法是关键。

用字典类比来说，当我们在字典（键）中查找一个词（查询）时，附近的词应该比远处的词有更大的影响。一个标记对另一个标记的影响是由$QK^T$点积决定的——所以这正是我们应该关注位置编码的地方！

$$\vec{a} \cdot \vec{b} = |\vec{a}| |\vec{b}| \cos \theta$$

上面所示点积的几何解释给了我们一个了不起的洞察力。我们可以通过单纯增加或减少两个向量之间的角度来调节点积的结果。此外，通过旋转向量，我们完全不影响向量的范数，而向量的范数编码了我们标记的语义信息。

现在我们知道应该把注意力集中在哪里了，也从另一个角度明白了为什么旋转可能是一个编码位置信息的“通道”，接下来我们把它们综合起来！

旋转**位**置**编**码 (RoPE)

**旋转位置编码**（或称RoPE）在RoFormer论文中定义（苏剑林在他的博客这里和这里独立设计了它）。如果跳到最终结果，它可能看起来像巫术，但通过在自注意力机制（特别是点积）的背景下思考正弦编码，我们可以看到它是如何协同工作的。

与正弦编码非常相似，我们将向量$\mathbf{q}$或$\mathbf{k}$（而不是预投影的$\mathbf{x}$）分解成二维对/块。我们不是通过添加由缓慢减小频率的正弦函数生成的向量来直接编码*绝对*位置，而是直接编码*相对*位置，方法是**将每对与旋转矩阵相乘**。

令$\mathbf{q}$或$\mathbf{k}$为我们在位置$p$的输入向量。我们创建一个块对角矩阵，其中$\mathbf{M_i}$是该分量对所需旋转的相应旋转矩阵

$$R(\mathbf{q}, p) = \begin{pmatrix} \mathbf{M_1} & & & \\ & \mathbf{M_2} & & \\ & & \ddots & \\ & & & \mathbf{M_{d/2}} \end{pmatrix} \begin{pmatrix} q_1 \\ q_2 \\ \vdots \\ q_d \end{pmatrix}$$

与正弦编码（Sinusoidal Encoding）非常相似，$\mathbf{M_i}$ 简单来说就是

$$\mathbf{M_i} = \begin{bmatrix} \cos(\omega_i p) & \sin(\omega_i p) \\ -\sin(\omega_i p) & \cos(\omega_i p) \end{bmatrix}$$

在实践中，我们不使用矩阵乘法来计算 RoPE，因为对于如此稀疏的矩阵来说，计算效率会很低。相反，我们可以直接将旋转独立地应用于成对的元素，利用计算中的规律模式：

$$R_{\Theta,p}^d q = \begin{pmatrix} q_1 \\ q_2 \\ q_3 \\ q_4 \\ \vdots \\ q_{d-1} \\ q_d \end{pmatrix} \otimes \begin{pmatrix} \cos p\theta_1 \\ \cos p\theta_1 \\ \cos p\theta_2 \\ \cos p\theta_2 \\ \vdots \\ \cos p\theta_{d/2} \\ \cos p\theta_{d/2} \end{pmatrix} + \begin{pmatrix} -q_2 \\ q_1 \\ -q_4 \\ q_3 \\ \vdots \\ -q_d \\ q_{d-1} \end{pmatrix} \otimes \begin{pmatrix} \sin p\theta_1 \\ \sin p\theta_1 \\ \sin p\theta_2 \\ \sin p\theta_2 \\ \vdots \\ \sin p\theta_{d/2} \\ \sin p\theta_{d/2} \end{pmatrix}$$

就这么简单！通过巧妙地将旋转应用于 $\mathbf{q}$ 和 $\mathbf{k}$ 的 2D 块，并在点积之前从加法切换到乘法，我们可以大大提升评估性能 $[^4]$。

将 RoPE 扩展到 $n$ 维

我们已经探讨了 RoPE 的 $1D$ 案例，我希望您现在对这个一开始不太直观的 Transformer 组件有了直观的理解。最后，让我们探讨如何将其扩展到更高维度，例如图像。

一个自然的初步直觉可能是直接使用图像中的 $\begin{bmatrix} x \\ y \end{bmatrix}$ 坐标对。这可能看起来很直观，毕竟我们之前几乎是随意地配对了我们的组件。然而，这将是一个错误！

在 $1D$ 情况下，我们通过旋转输入向量中的值对来编码相对位置 $m - n$。对于 $2D$ 数据，我们需要独立编码水平和垂直相对位置，例如 $m - n$ 和 $i - j$。RoPE 的精妙之处在于它如何处理多个维度。它不是将所有位置信息编码在一个旋转中，而是将同一维度内的组件配对并旋转它们，否则我们会将 $x$ 和 $y$ 偏移信息混合在一起。通过独立处理每个维度，我们保持了空间的自然结构。这可以推广到所需的任意多个维度！

位置编码的未来

RoPE 是位置编码的终极形态吗？DeepMind 的这篇 最新论文 深入分析了 RoPE 并指出了一些基本问题。简而言之：RoPE 并非完美解决方案，模型主要关注低频部分，但该论文表明，移除（而非旋转）最低频率可以提高 Gemma 2B 的性能！

我预计未来会取得一些突破，或许会从信号处理中汲取灵感，例如小波或分层实现等思想。随着模型越来越多地被量化以进行部署，我也希望看到编码方案的创新，使其在低精度算术下仍能保持稳健。

总结

位置编码在 Transformer 中一直被视为一个事后考虑的因素。我认为我们应该以不同的眼光看待它——自注意力机制有一个阿喀琉斯之踵，它被反复修补。

我希望这篇博客文章能向你展示，尽管起初它可能不直观，但你也可以发现最先进的位置编码。在后续的文章中，我很乐意探讨 RoPE 的实际实现细节，以最大限度地提高性能。

这篇帖子最初发布于此地。

参考资料

• Transformer 架构：位置编码
• 旋转嵌入：一场相对的革命
• Transformer 中位置编码的工作原理？
• Attention Is All You Need
• 周而复始！是什么让旋转位置编码如此有用？
• RoFormer: 带有旋转位置嵌入的增强型 Transformer

[^1]：二进制和正弦动画是此视频中包含的动画的再现。

[^2]：使用 $\theta = 10000$ 可提供 $2 \pi \cdot 10000$ 个唯一位置，或理论上上下文长度的上限约为 63,000。

[^3]：此文章的部分内容基于 Amirhossein Kazemnejad 撰写的这篇精彩文章。

[^4]：关于实证证据，请参阅 EleutherAI 的这篇精彩文章。