自注意力机制中的位置编码

  本内容主要介绍自注意力（Self-Attention）机制中的位置编码。

1.1 为什么自注意力机制需要位置编码

1.1.1 自注意机制简介

  Google 在 Transformer 论文 Attention Is All You Need 中提出了自注意力机制。

  当一个有 $n$ 个元素的序列 $x=(x_1,\cdots ,x_n)$ 输入进一个自注意力模块时，需要计算一个新的序列 $z=(z_1,\cdots ,z_n)$。元素 $z_i$ 为输入元素线性变换后的加权和：

$$z_i=\sum _{j=1}^n{\alpha }_{ij}(x_j{W}^V)\text{\hspace{1em}(1.1)}$$

其中，权重系数 ${\alpha }_{ij}$ 由 Softmax 函数计算得到：

$${\alpha }_{ij}=\text{softmax}(e_{ij})=\frac{\exp e_{ij}}{{\sum }_{k=1}^n\exp e_{ik}}\text{\hspace{1em}(1.2)}$$

其中，参数 $e_{ij}$ 由两个输入元素计算得到：

$$e_{ij}=\frac{(x_i{W}^Q)(x_j{W}^K{)}^T}{\sqrt{d}}\text{\hspace{1em}(1.3)}$$

  上面公式中的 $W^Q$，$W^K$ 和 $W^V$ 是可训练的参数矩阵。

图 1.1 自注意力机制

1.1.2 为什么自注意力机制需要位置编码

  在一个句子里，各个词的前后关系对这个句子的意义是有影响的，这种影响不仅体现在语法方面，而且还体现在语义方面。

  从上一小节自注意力机制简介中，我们知道当一个序列输入进一个自注意力模块时，由于序列中所有的 Token 是同时进入并被处理的，如果不提供位置信息，那么这个序列的相同的 Token 对自注意力模块来说就不会有语法和语义上的差别，他们会产生相同的输出。比如，图 1.2 里的词 “我” 出现 2 次，它们的初始表达向量是一样的，如果不加入位置信息，那么自注意力模块产生的对它们的关注度是一样的，或者模型产生的新的表达式一样的。所以，我们需要在输入序列里人为加入每个 token 的位置信息。

图 1.2 自注意力机制本身并不理会输入序列里各个词的位置

1.2 位置编码

  为了给自注意力模块加入位置编码，我们大体有两中选择：

• 想办法将位置信息融入到输入中，这构成了绝对位置编码的一般做法。

• 想办法微调一下 Attention 结构，使得它有能力分辨不同位置的 Token，这构成了相对位置编码的一般做法。

1.2.1 绝对位置编码

  形式上来看，绝对位置编码是相对简单的一种方案。一般来说，绝对位置编码会加到输入中：在输入的第 $k$ 个向量 $x_k$ 中加入位置向量 $p_k$ 变为 $x_k+p_k$，其中 $p_k$ 只依赖于位置编号 $k$。

  那么式（1.1）和（1.3）分别修改为：

$$z_i=\sum _{j=1}^n{\alpha }_{ij}((x_j+p_j)W^V)\text{\hspace{1em}(1.4)}$$

$$e_{ij}=\frac{((x_i+p_i)W^Q)((x_j+p_j)W^K{)}^T}{\sqrt{d}}\text{\hspace{1em}(1.5)}$$

1.2.1.1 Transformer 中的静态绝对位置编码

  Google 在 Transformer 论文 Attention Is All You Need 中提出：

$$\{\begin{array}{c}\begin{array}{rl}& p_{k,2i}=\sin (k/10000^{2i/d})\\ & p_{k,2i+1}=\cos (k/10000^{2i/d})\end{array}\end{array}\text{\hspace{1em}(1.6)}$$

其中，$p_{k,2i}$ 和 $p_{k,2i+1}$ 分别是位置 $k$ 的编码向量的第 $2i$ 和 $2i+1$ 个分量，$d$ 是位置向量的维度。

1.2.1.2 BERT 中的动态绝对位置编码

  BERT 使用了绝对位置动态编码。其位置编码向量由一个 $L\times d$ 的矩阵提供（其中 $L$ 是序列长度，$d$ 为词嵌入维度），其通过训练产生。

1.2.2 相对位置编码

  顾名思义，这种编码是以输入序列的一个位置为参考点，为离这个位置的距离进行编码，而绝对位置编码是为每个位置进行编码。

1.2.2.1 经典相对位置编码

  Google 在论文 Self-Attention with Relative Position Representations 中提出了一种相对位置编码。引入了两个训练向量 $a_{ij}^V，a_{ij}^K$ 用来学习两个输入元素 $x_i$ 和 $x_j$ 之间的相对位置表示，这两个向量在每层的各个注意力头之间共享。

  首先，将式（1.1）修改为：

$$z_i=\sum _{j=1}^n{\alpha }_{ij}(x_j{W}^V+a_{ij}^V)\text{\hspace{1em}(1.7)}$$

  然后，将式（1.3）修改为：

$$e_{ij}=\frac{(x_i{W}^Q)(x_j{W}^K+a_{ij}^K{)}^T}{\sqrt{d_z}}\text{\hspace{1em}(1.8)}$$

  一般认为，距离过远的词之间在语法和语义上的关联很弱了，可以不考虑。这样，我们定义一个截断距离 $k$，因此只需要考虑 $2k+1$ 个相对位置编码。这样做，还有一个好处，模型可以很好的扩展到训练期间没遇到的序列长度。

$$\begin{array}{c}{a}_{ij}^K=w_{\text{clip}(j-i,k)}^K\\ a_{ij}^V=w_{\text{clip}(j-i,k)}^V\\ \text{clip}(x,k)=\max (-k,\min (k,x))\end{array}\text{\hspace{1em}(1.9)}$$

  然后我们学习相对位置表示 $w^K=(w_{-k}^K,\cdots ,w_k^K)$ 和 $w^V=(w_{-k}^V,\cdots ,w_k^V)$，其中 $w_i^K,W_i^V\in {\mathbb{R}}^{d_a}$。从上面可以看到 $a_{ij}^V，a_{ij}^K$ 依赖于相对距离 $j-i$。

  小结：一个长度为 $L$ 的序列，相对位置最多有 $2L-1$ 个。如果设置了截断距离 $k$，那么相对位置的数量就变成 $2k+1$，我们要为这些相对位置构建维度为 $d$ 的相对位置编码。我们以可训练的变量形式创建大小为 $(2k+1)\times d$ 的相对位置编码矩阵 $P$，以及为每对词之间的相对距离创建大小为 $L\times L$ 的相对距离矩阵 $D$。然后，以 $D$ 矩阵里的元素值 $d_{ij}$（即相对距离）为行索引，在 $P$ 矩阵里抓取一行作为词 $j$ 相对词 $i$ 的相对位置编码向量 $p_{ij}^K$，再把这个位置编码向量作为偏置向量加到词 $j$ 的 Key 向量上，与自注意力模块一起训练即可。为 Value 向量创建的相对位置编码向量 $p_{ij}^V$ 也是同样的过程。

1.2.2.2 Transformer-XL 和 XLNet 中的相对位置编码

  Transformer-XL 和 XLNet 的相对位置编码借鉴了经典相对位置编码，然后在绝对位置编码式（1.5）的基础上做了若干变化，从而得到最终的公式。

  首先，将式（1.5）的分子部分展开（在这里，我们不考虑分母的 $\sqrt{d}$，因为其对分析结果不会产生影响），得到：

$$\begin{array}{rl}{\text{A}}_{i,j}^{\text{abs}}=& \underset{(a)}{\underbrace{x_i{W}^Q{W^K}^{\top }{x}_j^{\top }}}+\underset{(b)}{\underbrace{x_i{W}^Q(W^K{)}^{\top }{p}_j^{\top }}}\\ & +\underset{(c)}{\underbrace{p_i{W}^Q(W^K{)}^{\top }{x}_j^{\top }}}+\underset{(d)}{\underbrace{p_i{W}^Q(W^K{)}^{\top }{p}_j^{\top }}}\end{array}\text{\hspace{1em}(1.10)}$$

  然后，对上式做一些变化，得到：

$$\begin{array}{rl}{\text{A}}_{i,j}^{\text{abs}}=& \underset{(a)}{\underbrace{x_i{W}^Q(W^{K,E}{)}^{\top }{x}_j^{\top }}}+\underset{(b)}{\underbrace{x_i{W}^Q(W^{K,R}{)}^{\top }{r}_{i-j}^{\top }}}\\ & +\underset{(c)}{\underbrace{u(W^{K,E}{)}^{\top }{x}_j^{\top }}}+\underset{(d)}{\underbrace{v(W^{K,R}{)}^{\top }{r}_{i-j}^{\top }}}\end{array}\text{\hspace{1em}(1.11)}$$

• 首先，是（b）和（d）中的绝对位置编码 $p_j$ 被换成相对位置编码 $r_{i-j}$，其中 $r$ 采用了 Transformer 中不需要训练的 sinusoid 编码矩阵。

• 第二，引入两个可学习的参数 $u$ 和 $v$ 来分别替换（c）和（d）中的 $p_i{W}^Q$。因为对所有的查询位置，查询向量是一样；即无论查询位置如何，对不同词的注意力偏置应保持一致。

• 最后，将 $W^K$ 拆分为两个权重矩阵 $W^{K,E}$ 和 $W^{K,R}$，即分别为基于内容的 key 向量和基于位置的 key 向量。

  通过修改之后，式（1.11）中每个部分都有了其含义：（a）项表示基于内容的寻址（没有考虑位置编码）；（b）项表示相对于内容的位置偏差；（c）项表示全局的内容偏置（从内容层面衡量键的重要性）；（d）项表示全局的位置偏差（从相对位置层面衡量键的重要性）。

参考：

[1] Attention Is All You Need

[2] BERT(Pre-training of Deep Bidirectional Transformers for Language Understanding)

[3] Self-Attention with Relative Position Representations

[4] Transformer-XL - Attentive Language Models Beyond a Fixed-Length Context

[5] 详解自注意力机制中的位置编码（第一部分）

[6] 详解自注意力机制中的位置编码（第二部分）

[7] 详解Transformer-XL

[8] 让研究人员绞尽脑汁的Transformer位置编码