AIGC拾遗：大模型中的位置编码

前言

最近在准备秋招，正好趁此机会写一下位置编码的相关内容。
首先，我们考虑不加位置编码的情况：
假设我们的输入是\([a, b, c] \in \mathbb{R}^{d \times 3}\)，我们输出为\([o_{a}, o_{b}, o_{c}]\)；此时，我们调换输入的顺序为\([b, a, c]\)，输出相应变为\([o_{b}, o_{a}, o_{c}]\)。由此可见，attention操作只能够捕捉到不同token之间的关联，但是无法识别不同token的位置。因为当输入变为\([b, a, c]\)时，我们期望输出为\([g_{b}, g_{a}, g_{c}]\)，而非\([o_{b}, o_{a}, o_{c}]\)。

绝对位置编码

从形式看，绝对位置编码直接与输入相加。从实现方式看，绝对位置编码主要有两大类：学习式和固定式。其中，最经典和常用的是固定式，尤其是基于正余弦函数的编码。添加位置编码\(p\)后，\(q_{i}, k_{j}, v_{j}\)的计算方式为

\[q_{i} = (x_{i}+p_{i})W_{Q}, k_{j}=(x_{j}+p_{j})W_{K}, v_{j}=(x_{j}+p_{j})W_{V} \]

\(q_{i}, k_{j}\)的内积为

\[\begin{align} q_{i}k_{j}^{T} = (x_{i}+p_{i})W_{Q}W_{K}^{T}(x_{j}+p_{j})^{T}=x_{i}W_{Q}W_{K}^{T}x_{j}^{T}+x_{i}W_{Q}W_{K}^{T}p_{j}^{T}+p_{i}W_{Q}W_{K}^{T}x_{j}^{T}+p_{i}W_{Q}W_{K}^{T}p_{j}^{T} \notag \end{align} \tag{eq. 1} \label{eq. 1} \]

学习式绝对位置编码 (Learned Positional Embeddings)

这种方式最简单直观。模型为序列中的每个位置都分配一个随机初始化的、可学习的向量。这些向量在模型训练的过程中，与其他权重一起通过反向传播进行更新。因此，绝对位置编码矩阵\(P \in \mathbb{R}^{N \times d}\)，在做attention之前直接与输入\(x\)相加。通过这种方式实现的位置编码往往不具备外推能力，例如训练时最大序列长度为512时，推理时往往不能处理1024长度的序列。

固定式绝对位置编码 (Fixed Positional Encodings)

这种方式不依赖于学习，而是使用预先定义的数学函数来生成位置编码。最常见的就是正余弦编码。

\[PE(pos, 2i) = \sin(\frac{pos}{10000^{2i/d}}) \]

\[PE(pos, 2i+1) = \cos(\frac{pos}{10000^{2i/d}}) \]

这种正余弦形式的位置编码具有一定的相对位置表征的能力，对于固定的相对距离\(k\)，\(pos+k\)的位置编码可以用\(pos\)的位置编码进行表示，例如

\[\begin{align} PE(pos+k, 2i) &= \sin(\frac{pos+k}{10000^{2i/d}}) \notag \\ &=\sin(\frac{pos}{10000^{2i/d}})\cos(\frac{k}{10000^{2i/d}})+\cos(\frac{pos}{10000^{2i/d}})\sin(\frac{k}{10000^{2i/d}}) \notag\\ &= PE(pos, 2i)u+PE(pos, 2i+1)v \notag\\ \end{align} \tag{eq. 2} \]

\[\begin{align} PE(pos+k, 2i+1) &= \cos(\frac{pos+k}{10000^{2i/d}}) \notag \\ &=\cos(\frac{pos}{10000^{2i/d}})\cos(\frac{k}{10000^{2i/d}})-\sin(\frac{pos}{10000^{2i/d}})\sin(\frac{k}{10000^{2i/d}}) \notag\\ &= PE(pos, 2i+1)u-PE(pos, 2i)v \notag\\ \end{align} \tag{eq. 3} \]

由于正余弦位置编码采用固定函数生成，因此具备一定的外推能力

相对位置编码

相对位置编码保证不同位置之间的相对距离是同样的。传统的相对位置编码考虑到\eqref{eq. 1}中的后三项均和位置有关，直接将其替换为相对位置编码矩阵\(R(i, j)\)，其中，\(R(i, j)\)与相对距离\(i-j\)相关。

旋转位置编码 (RoPE)

上述基础上我们能够发现相对位置编码的添加位置均在\(q\)和\(k\)的内积部分。因此，我们将添加位置编码的函数记为\(f(x, pos)\)，则有\(q_m = f(q, m), k_n=f(k, n)\)。我们希望的是，能够实现\(<q_{m}, k_{n}>=g(q, k, n-m)\)。考虑到复数的内积，可以假设函数\(f(x, pos)=xe^{ipos\theta}\)，此时\(<q_{m}, k_{n}>=q^{T}ke^{i(n-m)\theta}\)满足我们的需求。由于复数符号的存在，需要将其转化为实数域的旋转运算。对于\(2\)维变量\(q\)，令

\[q_m = f(q, m) = \begin{bmatrix} cos(m\theta) & -sin(m\theta) \\ sin(m\theta) & cos(m\theta) \end{bmatrix} \begin{bmatrix} q^{0}\\ q^{1} \end{bmatrix} =R(m)q= \begin{bmatrix} Re((q^{0}+jq^{1})e^{jm\theta}) \\ Im((q^{0}+jq^{1})e^{jm\theta}) \end{bmatrix} \]

可以验证

\[R^{T}(m)R(n) = \begin{bmatrix} cos(m\theta) & sin(m\theta) \\ -sin(m\theta) & cos(m\theta) \end{bmatrix} \begin{bmatrix} cos(n\theta) & -sin(n\theta) \\ sin(n\theta) & cos(n\theta) \end{bmatrix} = \begin{bmatrix} cos((n-m)\theta) & -sin((n-m)\theta) \\ sin((n-m)\theta) & cos((n-m)\theta) \end{bmatrix} =R(n-m) \]

\[q_{m}^{T}k_{n} = q^{T}R^{T}(m)R(n)k = q_{m}R(n-m)k_{n} \]

因此，旋转位置编码能够实现相对位置编码。

代码实现

def rope_params(max_seq_len, dim, theta=10000):
    assert dim % 2 == 0
    freqs = torch.outer(
        torch.arange(max_seq_len),
        1.0 / torch.pow(theta,
                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
    freqs = torch.polar(torch.ones_like(freqs), freqs)
    return freqs

def apply_rotary_emb(x, freqs):
    x = x.to(torch.float64).reshape(*x.shape[:-1], -1, 2)
    x = torch.view_as_complex(x)
    x = torch.view_as_real(x * freqs).flatten(2)
    return x.float()

首先计算出不同位置的频率分量\(freqs \in \mathbb{R^{N \times \frac{d}{2}}}\)

\[freqs = \begin{bmatrix} \frac{0}{10000^{\frac{0}{d}}} & \frac{0}{10000^{\frac{2}{d}}} & \frac{0}{10000^{\frac{4}{d}}} & \dots & \frac{0}{10000^{\frac{2(i-1)}{d}}} & \dots & \frac{0}{10000^{\frac{d-2}{d}}} \\ \frac{1}{10000^{\frac{0}{d}}} & \frac{1}{10000^{\frac{2}{d}}} & \frac{1}{10000^{\frac{4}{d}}} & \dots & \frac{1}{10000^{\frac{2(i-1)}{d}}} & \dots & \frac{1}{10000^{\frac{d-2}{d}}} \\ \dots & \dots & \dots & \dots & \dots & \dots & \dots \\ \frac{N-1}{10000^{\frac{0}{d}}} & \frac{N-1}{10000^{\frac{2}{d}}} & \frac{N-1}{10000^{\frac{4}{d}}} & \dots & \frac{N-1}{10000^{\frac{2(i-1)}{d}}} & \dots & \frac{N-1}{10000^{\frac{d-2}{d}}} \end{bmatrix} = \begin{bmatrix} \theta_{0, 0} & \theta_{0, 1} & \theta_{0, 2} & \dots & \theta_{0, i} & \dots & \theta_{0, \frac{d}{2}-1} \\ \theta_{1, 0} & \theta_{1, 1} & \theta_{1, 2} & \dots & \theta_{1, i} & \dots & \theta_{1, \frac{d}{2}-1} \\ \dots & \dots & \dots & \dots & \dots & \dots & \dots \\ \theta_{N-1, 0} & \theta_{N-1, 1} & \theta_{N-1, 2} & \dots & \theta_{N-1, i} & \dots & \theta_{N-1, \frac{d}{2}-1} \\ \end{bmatrix} \]

然后取复数

\[freqs = \begin{bmatrix} e^{j\theta_{0, 0}} & e^{j\theta_{0, 1}} & e^{j\theta_{0, 2}} & \dots & e^{j\theta_{0, i}} & \dots & e^{j\theta_{0, \frac{d}{2}-1}} \\ e^{j\theta_{1, 0}} & e^{j\theta_{1, 1}} & e^{j\theta_{1, 2}} & \dots & e^{j\theta_{1, i}} & \dots & e^{j\theta_{1, \frac{d}{2}-1}} \\ \dots & \dots & \dots & \dots & \dots & \dots & \dots \\ e^{j\theta_{N-1, 0}} & e^{j\theta_{N-1, 1}} & e^{j\theta_{N-1, 2}} & \dots & e^{j\theta_{N-1, i}} & \dots & e^{j\theta_{N-1, \frac{d}{2}-1}} \\ \end{bmatrix} \]

对于输入\(x \in \mathbb{R^{N \times d}}\)先将其reshape到\([N, \frac{d}{2}, 2]\)维度，再转换成复数形式\([x_{1}+jx_{2}, x_{3}+jx_{4}, \dots]\)，得到复数张量\(x \in \mathbb{R^{N \times \frac{d}{2}}}\)。\(x\)与\(freqs\)相乘后再转为实部\([[x_{1}, x_{2}], [x_{3}, x_{4}], \dots]\)再展平。