Transformer XL原理介绍

1. 引言

在自然语言处理中，当前深度学习主流的结构是RNN和Transformer，因为这两种结构能够比较好地捕捉文本的上下文信息，但是，我们知道RNN主要的问题是梯度消失和梯度爆炸的问题，而且其捕捉上下文的长度没有Transformer那么强大，而Transformer虽然能力比较强，但是在预测时会受到训练时所设定的最大长度限制，因此，本文将介绍一个Transformer模型的变体，即Transformer XL(extra long)，该模型是CMU和谷歌大脑在2019年提出来的，通过引入了递归的机制和相对位置编码，解决了Transformer长度限制的问题，作者在实验中发现，Transformer XL能捕捉到的文本长度比RNN长80%。

• 论文地址：《Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context》

2. Transformer XL原理介绍

2.1 Vanilla Transformer

在介绍Transformer模型之前，先介绍另一个模型，该模型是2018年由Al-Rfou等人提出来的，作者称之为Vanilla Transformer，该模型其实本质上还是Transformer，只是将原来的句子进行切片，切成一个一个固定长度的短句子进行训练，如下图中的(a)所示，然后在预测时，采用训练时所设定的长度窗口，在句子上进行平移，从而不断获得句子的预测，如下图中的(b)所示，虽然这样操作可以解决句子的长度限制问题，但是存在两个主要的缺点：①每次平移后，模型都是从头开始计算，并没有利用前面计算得到的上文信息，因此，会出现一种“上下文断裂”的情况；②在预测时，当句子的长度比较长时，计算速度会非常慢，因为每次窗口只移动一个时间步，这样会导致很多信息重复计算。作者在实验中发现，采用Transformer XL模型，速度是Vanilla Transformer的1800多倍。

2.2 Segment-Level Recurrence

为了克服Vanilla Transformer模型的缺点，Transformer XL引入了一种递归的机制，即通用将句子按照固定的长度$L$分成若干个子句，然后在训练和预测时，依次将每个子句传入Transformer模型，并且将每个子句在Transformer中各层的输出传递给下一个子句，如下图(a)中绿色线，在下一个子句每一层的计算中，将上一个子句对应上一层的输出与当前子句对应层的输入进行拼接，这样在每一个子句的预测时，就可以考虑到前面各个子句的信息。

记两个连续的子句为${\mathbf{s}}_{\tau }=\left[x_{\tau ,1},\cdots ,x_{\tau ,L}\right]$，${\mathbf{s}}_{\tau +1}=\left[x_{\tau +1,1},\cdots ,x_{\tau +1,L}\right]$，记第$\tau$个子句第$n$层的输出为${\mathbf{h}}_{\tau }^n\in {\mathbb{R}}^{L\times d}$，其中$d$表示隐藏层的维度，则第$\tau +1$个子句${\mathbf{s}}_{\tau +1}$第$n$层的输出计算如下：

$$\begin{array}{l}{\tilde{\mathbf{h}}}_{\tau +1}^{n-1}=\left[\mathrm{S}\mathrm{G}\left({\mathbf{h}}_{\tau }^{n-1}\right)\circ {\mathbf{h}}_{\tau +1}^{n-1}\right]\\ {\mathbf{q}}_{\tau +1}^n,{\mathbf{k}}_{\tau +1}^n,{\mathbf{v}}_{\tau +1}^n={\mathbf{h}}_{\tau +1}^{n-1}{\mathbf{W}}_q^{\top },{\tilde{\mathbf{h}}}_{\tau +1}^{n-1}{\mathbf{W}}_k^{\top },{\tilde{\mathbf{h}}}_{\tau +1}^{n-1}{\mathbf{W}}_v^{\top }\\ {\mathbf{h}}_{\tau +1}^n=\text{ Transformer-Layer }\left({\mathbf{q}}_{\tau +1}^n,{\mathbf{k}}_{\tau +1}^n,{\mathbf{v}}_{\tau +1}^n\right)\end{array}$$

其中，$\mathrm{S}\mathrm{G}(\cdot )$表示stop-gradient，$\left[{\mathbf{h}}_u\circ {\mathbf{h}}_v\right]$表示将两个序列进行拼接，$\mathbf{W}$表示模型的参数。这里主要的不同点在于计算K和V时，使用的是${\tilde{\mathbf{h}}}_{\tau +1}^{n-1}$，其中，包含了来自上一个子句的信息${\mathbf{h}}_{\tau }^{n-1}$。由于每个子句都会递归地使用上一个子句的信息，因此，随着层数的加深，每个子句中每个时间步可以利用到的前面的序列信息长度应该是$O(N\times L)$，如下图中(b)所示。在实际训练中和预测中，对于每个子句，还可以考虑前面多几个子句的信息，而不是仅使用上一个子句。作者在实验中通过一个参数$M$来表示考虑的上文的长度信息，在训练时，令$M=L$，即只考虑一个子句信息，但是在预测时，可以将其扩增几个倍数。

Transformer XL除了可以捕捉更长的文本信息外，也提高了模型预测的速度，因为我们知道Transformer XL在预测时每次是平移一个子句的长度，而不是像Vanilla Transformer一样每次只平移一个时间步。

2.3 Relative Positional Encodings

虽然Transformer XL通过切分子句，并引入递归的思想解决了句子长度的限制问题，但是还忽略了一个问题，即positional encoding。在Transformer中，通过一个positional encoding来表征句子的时序信息，记为$\mathbf{U}\in {\mathbb{R}}^{L_{\max }\times d}$，其中，${\mathbf{U}}_i$表示第$i$个位置的绝对位置信息，$L_{max}$表示句子的最大长度，我们知道，Transformer中，会将词汇embedding信息与位置信息进行元素相加，然后作为整个模型的原始输入，那么，如果Transformer XL也采用这样的方式进行计算的话，其表达式如下：
$$\begin{array}{rl}{\mathbf{h}}_{\tau +1}& =f\left({\mathbf{h}}_{\tau },{\mathbf{E}}_{{\mathbf{s}}_{\tau +1}}+{\mathbf{U}}_{1:L}\right)\\ {\mathbf{h}}_{\tau }& =f\left({\mathbf{h}}_{\tau -1},{\mathbf{E}}_{{\mathbf{s}}_{\tau }}+{\mathbf{U}}_{1:L}\right)\end{array}$$

其中，${\mathbf{E}}_{{\mathbf{s}}_{\tau }}\in {\mathbb{R}}^{L\times d}$表示子句${\mathbf{s}}_{\tau }$的embedding信息，可以看到，两个子句${\mathbf{s}}_{\tau },{\mathbf{s}}_{\tau +1}$用的都是同一个位置信息，这会导致模型没法区分两个子句的顺序信息。

为了克服这个问题，作者引入了一种新的Positional Encoding方法，即Relative Positional Encodings。我们知道，Transformer在$q_i$和$k_j$的注意力时，根据矩阵的分配律，可以得到如下：
$$\begin{array}{rl}{\mathbf{A}}_{i,j}^{\mathrm{a}\mathrm{b}\mathrm{s}}& =\underset{(a)}{\underbrace{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{E}}_{x_j}}}+\underset{(b)}{\underbrace{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{U}}_j}}\\ & +\underset{(c)}{\underbrace{{\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{E}}_{x_j}}}+\underset{(d)}{\underbrace{{\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{U}}_j}}\end{array}$$

作者认为，其实在每次计算时，只要让模型可以知道句子的相对位置信息即可，因此，基于这个思想，作者将上式中(b)和(d)项中的${\mathbf{U}}_j$替换为${\mathbf{R}}_{i-j}$，用来表征相对位置信息，然后将©中的${\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }$采用一个可训练的$u\in {\mathbb{R}}^d$向量替代，对于(d)中的${\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }$，同样采用一个可训练的$v\in {\mathbb{R}}^d$替代，最后，对于${\mathbf{W}}_k$，采用两个矩阵进行替代${\mathbf{W}}_{k,E}$、${\mathbf{W}}_{k,R}$。最终修正后的表达式如下：
$$\begin{array}{rl}{\mathbf{A}}_{i,j}^{\mathrm{r}\mathrm{e}\mathrm{l}}& =\underset{(a)}{\underbrace{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_{k,E}{\mathbf{E}}_{x_j}}}+\underset{(b)}{\underbrace{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_{k,R}{\mathbf{R}}_{i-j}}}\\ & +\underset{(c)}{\underbrace{u^{\top }{\mathbf{W}}_{k,E}{\mathbf{E}}_{x_j}}}+\underset{(d)}{\underbrace{v^{\top }{\mathbf{W}}_{k,R}{\mathbf{R}}_{i-j}}}\end{array}$$

采用这种相对位置编码后，假设对于一个$N$层的Transformer，不失一般性，在只有一个multi-heads的情况下，Transformer XL每一层的计算如下：

其中，${\mathbf{h}}_{\tau }^0:={\mathbf{E}}_{{\mathbf{s}}_{\tau }}$表示词汇的embedding。虽然采用相对位置编码使得模型可以区分不同的子句相对位置，但是，我们可以注意到，在计算multi-heads attention时，此时${\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j}$的计算时间是二次方的，严重影响了模型的计算速度，因此，为了克服这个问题，作者通过推导将其计算时间转化为线性的。

我们知道，对于配对$(i,j)$，$i-j$的取值范围为$[0,M+L-1]$，因此，可以将${\mathbf{W}}_{k,R}{\mathbf{R}}_{i-j}$的计算转化为如下：

此时，对于每一个配对的$(i,j)$，$i-j$，都可以通过位置索引直接从矩阵$Q$中获取到，为了便于表示，我们记${\mathbf{Q}}_k={\mathbf{W}}_{k,R}{\mathbf{R}}_{M+L-1-k}$，此时，对于${\mathbf{q}}_{\tau ,i}^n\top {\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j}$的计算，可以表示如下：

因此，可以定义：

可以发现，对于$\mathbf{B}$中的每一行，其实就是$\tilde{\mathbf{B}}$中每一行分别向左移动$L、L-1、L-2、...$个单位，因此，通过这样的方式可以把计算时间转化为线性的。

同理，对于$v^{\top }{\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j}$的计算，采用同样的方法可以定义矩阵：

以及：

此时，对于$\mathbf{D}$中的每一行，其实就是$\tilde{\mathbf{d}}$向左移动一定单位后的结果。

3. 总结

以上就是Transformer XL的介绍，感觉还是有一定复杂度的。最后总结一下：

• Transformer XL由于加入递归机制和相对位置编码，因此，解决了文本长度的限制问题以及Vanllia Transformer中上下文信息断裂、预测速度慢的问题。
• 由于采用相对位置编码，因此，计算速度相比Transformer要慢一点。