概述

相对于RNN这样的序列模型来说，Transformer可并行是一个很大的优势，但可并行性带来一个问题，由于不是从前到后，所以模型对于位置信息是不敏感的。于是在Transformer最早提出时就定义了位置编码（Positional Encodings）的概念，本文章旨在介绍常见位置编码方式。

绝对位置编码

绝对位置编码（Absolute Positional Encodings）也在早期的Transformer架构中被广泛使用。
Transformer中的位置编码如下图所示：

最初Transformer中的位置编码采用三角式绝对位置编码，计算方式如下：
$P{E}_{(k,2i)}=sin(\frac{k}{10000^{2i/d}})$
$P{E}_{(k,2i+1)}=cos(\frac{k}{10000^{2i/d}})$
其中，$2i$和$2i+1$代表向量的第几维，$k$代表token在序列中的位置，$d$是向量的维度。
实现代码：

import torch
import math
seq_len = 128 # 序列长度
d_model = 8   # 向量维度
pe = torch.zeros(seq_len, d_model)
position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1) # shape：[128, 1, 8]

相对位置编码

绝对位置编码的缺点很明显：

1. 可扩展性不足，只能处理训练数据中看过的长度，如基础BERT模型只能够处理512个token，这显然是不合理的。所以越来越多的模型使用相对位置编码的方法来增加模型对于长度的泛化性。
2. 每一个位置编码都是独立的。换而言之，位置编码和距离没有关系，不管距离多远都是一样的。

在正式介绍相对位置编码之前回顾一下原始论文中 self-attention的计算方式：
$q_i=(x_i+p_i)W_q$
$k_j=(x_j+p_j)W_k$
$v_k=(x_k+p_k)W_v$
${\alpha }_{i,j}=softmax(q_i{k}_j^T)$
$o_i={\sum }_j{\alpha }_{i,j}{v}_j$

我们把$q_i$和$k_j$中与位置编码的项去掉，再引入一个相对位置编码项。 对$q_i{k}_j^T$进行展开，得到以下等式：
$q_i{k}_j^T=(x_i{W}_q)(x_j{W}_k{)}^T\oplus R_{i,j}$

其中$R_{i,j}$代表相对位置的编码方式，$\oplus$代表融入相对位置编码的方式（如：相乘或相加）。

T5 Bias

T5模型直接使用一个可学习矩阵来进行相对位置编码的学习：
$q_i{k}_j^T=(x_i{W}_q)(x_j{W}_k{)}^T+{\beta }_{i,j}$
其中${\beta }_{i,j}$为一个可学习的、共享的bias值，这种方法的缺点很明显，需要训练矩阵的参数，训练会变慢。

ALiBi

基于T5 bias的缺点，Press[3]提出了ALiBi算法。该方法通过预定义相对位置编码，通过在$q_i{k}_j^T$后增加一个偏置项来达到相对位置编码的目的：

其中$m$是和注意力头有关的斜率，对于$n$个注意力头，其斜率集合从$2^{\frac{-8}{n}}$开始，计算示例如下图所示：

从原始论文可以看到，对比T5 Bias，输入长度越长优势明显。

RoPE

RoPE[5]提出的出发点是“通过绝对位置编码的方式实现相对位置编码”。如前所述，self-attention核心是计算是内积，所以问题就变成了内积前带有绝对位置编码信息，内积之后带有相对位置编码信息。假设存在这样的位置函数记为$f(\mathbf{x},m)$，则有：
$⟨f(\mathbf{q},m),f(\mathbf{k},m)⟩=g(\mathbf{q},\mathbf{k},m-n)$

所以现在的目标变成了找到这样一个函数使得上述等式成立。论文就提出了一个可以满足上述等式的函数。为了简化推理，我们假设向量的维度为2，对于第$m$个位置的token，对应的旋转矩阵为：
$R_{\theta ,m}=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)$

其中$\theta$是一个预设值，经过旋转之后的矩阵我们可以写成：
$q^{\prime }=R_{\theta ,m}{q}_m=\left(\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\end{array}\right)$

对于二维向量，原始论文给出的RoPE实现过程如下图所示：

RoPE一般性实现如下：

$$\begin{array}{rl}{q}_m^T{k}_n& =(R_{\theta ,m}{q}_m{)}^T(R_{\theta ,n}{k}_n)\\ & =q^T{R}_{\theta ,m}^T{R}_{\theta ,n}{k}_n(R_{\theta ,m}^T{R}_{\theta ,n}=R_{\theta ,n-m})\\ & =q^T{R}_{(\theta ,n-m)}{k}_n\end{array}$$
可以看出RoPE巧妙地结合了绝对位置编码和相对位置编码：
● $\theta$只与token当前的位置有关，且对于所有token都是一样的
● $n-m$和两个token之间的距离有关，token距离越小，旋转的角度越小，反之，旋转角度越大。
RoPE实现代码[8]如下：

import torchfrom torch import nndef get_rotary_matrix(context_len: int, embedding_dim: int) -> torch.Tensor:"""Generate the Rotary Matrix for ROPEArgs:context_len (int): context lenembedding_dim (int): embedding dimReturns:torch.Tensor: the rotary matrix of dimension context_len x embedding_dim x embedding_dim"""R = torch.zeros((context_len, embedding_dim, embedding_dim), requires_grad=False)positions = torch.arange(1, context_len+1).unsqueeze(1)# Create matrix theta (shape: context_len  x embedding_dim // 2)slice_i = torch.arange(0, embedding_dim // 2)# 原始论文旋转角度计算方式theta = 10000. ** (-2.0 * (slice_i.float()) / embedding_dim) m_theta = positions * theta# Create sin and cos valuescos_values = torch.cos(m_theta)sin_values = torch.sin(m_theta)# 按照公式进行赋值R[:, 2*slice_i, 2*slice_i] = cos_valuesR[:, 2*slice_i, 2*slice_i+1] = -sin_valuesR[:, 2*slice_i+1, 2*slice_i] = sin_valuesR[:, 2*slice_i+1, 2*slice_i+1] = cos_valuesreturn Rbatch_size = 1
context_len = 8
embedding_dim = 2
ff_q = nn.Linear(embedding_dim, embedding_dim, bias=False)
ff_k = nn.Linear(embedding_dim, embedding_dim, bias=False)
x = torch.randn((batch_size, context_len, embedding_dim))# 初始化矩阵：Q、K
queries = ff_q(x)
keys = ff_k(x)# 获取旋转矩阵
R_matrix = get_rotary_matrix(context_len, embedding_dim)queries_rot = (queries.transpose(0,1) @ R_matrix).transpose(0,1)
keys_rot = (keys.transpose(0,1) @ R_matrix).transpose(0,1)# ... Compute the score in the attention mechanism using the rotated queries and keys

目前RoPE技术已经在LLAMA系列、GLM系列、Baichuan系列、Qwen系列等广泛使用。

参考资料

1. Zhao et al. Length Extrapolation of Transformers: A Survey from the Perspective of Positional Encoding.
2. Su etal. ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING.
3. 科学空间《让研究人员绞尽脑汁的Transformer位置编码》
4. Press et al. Train short, test long: Attention with linear biases enables input length extrapolation.
5. Transformer升级之路：2、博采众长的旋转式位置编码
6. Su et al. RoFormer: Enhanced Transformer with Rotary Position Embedding
7. https://github.com/meta-llama/llama3/blob/main/llama/model.py
8. https://afterhoursresearch.hashnode.dev/rope-rotary-positional-embedding