Transformer中的Positional Encoding

由于self-attention是没有位置信息的，如何在其中加入位置信息就变成了一个活跃的研究领域。 在原始的Transformer论文中，作者使用了一个\(d\)维向量来表示某个位置的信息。在位置\(t\)，

\[\begin{align*} p_t^{i} &= \left\{ \begin{aligned} \sin(\omega_k t), \quad if \quad i &= 2k\\ \cos(\omega_k t), \quad if \quad i &= 2k + 1 \end{aligned} \right. \\ \omega_k &= \frac{1}{10000^{2k/d}} \\ k &\in [0, d/2] \end{align*}\]

然后简单地把这个positional encoding加到embedding层的输出上，

\[out(w_t) = emb(w_t) + p_t\]

可见这其中必然有 \(d_{emb} = d_{pos}\)。使用\(\sin\)和\(\cos\)产生了一个结果，就是\(t\)个不同的position向量，向量中的值沿着\(t\)有一个平移的效果， 间隔一定的位置会出现重复的值，产生了相对位置编码的效果。

Rotary Positional Embedding (RoPE)

RoPE在计算上是一个绝对位置编码，但产生的效果却是相对位置编码的。 由于在Attention中，query向量和key向量内积产生了权重，RoPE直接面向这个内积操作设计了一个编码器，使得： 对于在位置\(m\)的query向量\(\vec{q_m}\)，以及在位置\(n\)的key向量\(\vec{k_n}\)，经过函数\(f\)变换后， 内积\(\langle f(\vec{q_m}), f(\vec{k_n}) \rangle\)能够表达\(q_m\)和\(k_n\)之间的相对位置信息\(m-n\)。

RoPE对query和key向量中的元素两个一组做变换，这个\(f\)变换如下，

\[\begin{equation} f(\vec{q}, m) = \begin{pmatrix} M_1 & & & \\ & M_2 & & \\ & & \ddots & \\ & & & M_{\frac{d}{2}} \end{pmatrix} \begin{pmatrix} q_1 \\ q_2 \\ \vdots \\ q_{d-1} \\ q_d \end{pmatrix} \triangleq \Theta \vec{q} \label{eq:rope-transform} \end{equation}\]

其中 \(M_j = \begin{pmatrix} \cos m\theta_j & -\sin m\theta_j \\ \sin m\theta_j & \cos m\theta_j \end{pmatrix}\)，组成了分块对角矩阵\(\Theta\)里的非零元素。\(\theta_j = 10000^{−\frac{2j}{d}}\)是预先定义的常量，带来一定的远程衰减性。 对于key向量\(\vec{k}\)也做同样的变换。

那为什么这样变换之后就可以携带相对位置信息呢？简单起见，我们可以先考虑二维的情况，

\[M_j \vec{q_j} = \begin{pmatrix} \cos m\theta_j & -\sin m\theta_j \\ \sin m\theta_j & \cos m\theta_j \end{pmatrix} \begin{pmatrix} q_{2j-1} \\ q_{2j} \end{pmatrix}\]

我们知道，在复数的语境里，以下这些表达形式是相通的：

  • 复数的实部加虚部表达\(q_{2j-1} + q_{2j} \cdot i\)
  • 复数可以被视为一个二维向量 \(\begin{pmatrix} q_{2j-1} \\ q_{2j} \\ \end{pmatrix}\)
  • 复数的矩阵表达 \(\begin{pmatrix} q_{2j-1} & -q_{2j} \\ q_{2j} & q_{2j-1} \end{pmatrix}\)
  • 复数的极坐标表达\(\|\vec{q_j}\| e^{i \cdot \Theta(\vec{q_j})}\)，其中 \(\vec{q_j} = \begin{pmatrix} q_{2j-1} \\ q_{2j} \\ \end{pmatrix}\)，\(\Theta(\vec{q_j})\)是\(\vec{q_j}\)在极坐标下的幅角。
复数的矩阵表达

复数系\(\mathbb{C}\)和它的矩阵表达\(\mathbb{R}^{2\times 2}\)本质上是两个同态（homomorphism）的代数结构。定义复数\(a + bi\)到矩阵的映射\(f\)：

\[f(a + bi) = \begin{pmatrix} a & -b \\ b & a \end{pmatrix}\]

矩阵保留了复数的加法结构：

\[\begin{align*} f(a + bi + x + yi) &= f((a + x) + (b + y) i) \\ &= \begin{pmatrix} a + x & -b-y \\ b + y & a + x \end{pmatrix} \\ &= \begin{pmatrix} a & -b \\ b & a \end{pmatrix} + \begin{pmatrix} x & -y \\ y & x \end{pmatrix} \end{align*}\]

以及乘法结构：

\[\begin{align*} f((a + bi) (x + yi)) &= f((ax - by) + (bx + ay) i) \\ &= \begin{pmatrix} ax - by & -bx-ay \\ bx + ay & ax - by \end{pmatrix} \\ &= \begin{pmatrix} a & -b \\ b & a \end{pmatrix} \begin{pmatrix} x & -y \\ y & x \end{pmatrix} \end{align*}\]

值得注意的是，另一方面，复数的向量表达\(\mathbb{R}^2\)虽然保留了加法结构，但向量空间\(\mathbb{R}^2\)是不存在对应的乘法结构的。

于是我们可以得到位置\(m\)的\(\vec{q_j}\)，

\[\begin{align*} M_j \vec{q_j} &= \begin{pmatrix} \cos m\theta_j & -\sin m\theta_j \\ \sin m\theta_j & \cos m\theta_j \end{pmatrix} \begin{pmatrix} q_{2j-1} \\ q_{2j} \end{pmatrix} \\ &\equiv \begin{pmatrix} q_{2j-1} & -q_{2j} \\ q_{2j} & q_{2j-1} \end{pmatrix} \begin{pmatrix} \cos m\theta_j & -\sin m\theta_j \\ \sin m\theta_j & \cos m\theta_j \end{pmatrix} &\qquad（将二维向量视作复数，矩阵表达保持乘法结构）\\ &\equiv \| \vec{q_j} \| e^{i \cdot \Theta(\vec{q_j})} (\cos m\theta_j + i \cdot \sin m\theta_j) &\qquad（矩阵的复数实/虚部及极坐标表达）\\ &= \| \vec{q_j} \| e^{i \cdot \Theta(\vec{q_j})} \cdot e^{im\theta_j} &\qquad（欧拉公式） \\ &= \| \vec{q_j} \| e^{i (\Theta(\vec{q_j}) + m\theta_j)} \end{align*}\]

对位置\(n\)的\(\vec{k_j}\)做同样的变换之后，求两者的内积（在复数语义下，内积等于\(\vec{q}\)乘以\(\vec{k}\)的共轭），

\[\begin{align} \langle M_j \vec{q_j}, N_j \vec{k_j} \rangle &\equiv \| \vec{q_j} \| e^{i (\Theta(\vec{q_j}) + m\theta_j)} \| \vec{k_j} \| \overline{e^{i (\Theta(\vec{k_j}) + n\theta_j)}} \\ &= \| \vec{q_j} \| \| \vec{k_j} \| e^{i ( \Theta(\vec{q_j}) + m\theta - \Theta(\vec{k_j}) - n\theta_j ) } \\ &= \| \vec{q_j} \| \| \vec{k_j} \| e^{i ( \Theta(\vec{q_j}) - \Theta(\vec{k_j}) ) } e^{ i (m-n)\theta_j } \\ &\equiv \langle \vec{q_j}, \vec{k_j} \rangle \cdot e^{i (m-n) \theta_j} \label{eq:rope-inner-prod} \end{align}\]

可见，这样对query和key做变换之后，仍然保留了内积的语义，又嵌入了\(m-n\)的信息。 更妙的是，(\(\ref{eq:rope-inner-prod}\))式仍然是一个复数极坐标的表达方式，也就是说，RoPE结果是一个带上了相对位置信息的向量。

实际操作中，由于(\(\ref{eq:rope-transform}\))式中的\(\Theta \triangleq \begin{pmatrix} M_1 & & & \\ & M_2 & & \\ & & \ddots & \\ & & & M_{\frac{d}{2}} \end{pmatrix}\)是一个稀疏矩阵，直接计算会非常低效。通常会该用以下方式，

\[\begin{pmatrix} q_1 \\ q_2 \\ q_3 \\ q_4 \\ \vdots \\ q_{d-1} \\ q_d \end{pmatrix} \otimes \begin{pmatrix} \cos m\theta_1 \\ \cos m\theta_1 \\ \cos m\theta_2 \\ \cos m\theta_2 \\ \vdots \\ \cos q_{d/2-1} \\ \cos q_{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 m\theta_1 \\ \sin m\theta_1 \\ \sin m\theta_2 \\ \sin m\theta_2 \\ \vdots \\ \sin q_{d/2-1} \\ \sin q_{d/2} \end{pmatrix}\]

参考文献

RoPE作者苏剑林老师的博文Transformer升级之路：博采众长的旋转式位置编码 以及Roformer论文有更详细的推导过程。 关于复数的各种表示形式，这篇文章 有比较清楚的解释。 对Transformer原始论文中位置编码的解释和效应，Transformer Architecture: The Positional Encoding 这篇博文讲解得非常具体。