文献阅读笔记：Transformer-XL : Attentive Language Models Beyond a Fixed-Length Context

0.背景

机构：CMU、谷歌大脑
作者：Zihang Dai、Zhilin Yang
发布地方：arxiv
面向任务：Language Understanding
论文地址：https://arxiv.org/abs/1901.02860
论文代码：https://github.com/kimiyoung/transformer-xl

0-1 摘要

Transformer具有学习长程依赖关系的潜力，但是受到语言建模中上下文长度固定的限制。为此，本文提出一种新的神经网络架构Transformer-XL，该网络结构能够在不破坏时间一致性的情况下，学习到超越固定长度的依赖性。该网络结构由片段级的循环机制（segment-level recurrence）和全新的位置编码策略（positional encoding scheme）组成。其优点是不仅可以捕获更长的依赖关系，还可以解决上下文碎片化（context fragmentation）的问题。从实验结果上来看，Transformer-XL 学习到的依赖性比 RNN 学习到的长 80%，比标准 Transformer 学到的长 450%，无论在长序列还是短序列中都得到了更好的结果，而且在评估时比标准 Transformer 快 1800+ 倍。值得一提的是，Transformer-XL还刷新了 bpc 和perplexity（困惑度）的当前最佳结果：在 enwiki8 上 bpc 从 1.06 提升至 0.99，在 text8 上从 1.13 提升至 1.08；在 WikiText-103 上困惑度从 20.5 提升到 18.3，在 One Billion Word 上从 23.7 提升到 21.8，在宾州树库（不经过微调的情况下）上从 55.3 提升到 54.5。本文模型的代码、预训练模型以及超参数在 TensorFlow 和 PyTorch 中都可以使用。

1. 介绍

语言建模需要对长期依赖性进行建模，它成功应用了无监督的预训练方法 (Peters et al., 2018; Devlin et al., 2018)。但要让神经网络对序列数据的长期依赖性建模一直都是一项挑战。RNN网络，特别是LSTM是一个标准的方案，它可以在多个benchmarks上获得健壮的结果（strong results）。尽管其使用广泛，但是RNNs由于梯度消失和梯度爆炸问题的存在，难以优化。纵使引入一些门限和梯度裁剪技术，仍然不足以完全解决该问题。此前的工作已经表明LSTM平均可以捕获200个word的上下文信息，这也指出了进一步改进的空间。

另一方面，通过attention机制直接连接长程word pairs可以缓解优化问题，同时习得长程依赖（即原始的Transformer工作）。近来Al-Rfou 等人（2018）设计了一组辅助损失来训练深度 Transformer 网络进行字符级（character-level）语言建模，其结果远超LSTM。虽然已经取得成功，但是 Al-Rfou 等人（2018）的语言模型是在长度固定的几百个字符片段上独立训练的，没有任何跨片段的信息流（即多个segments，每个segment的长度固定，由数百个characters组成。但是segments之间没有信息交流）。由于上下文的长度是固定的，因此模型无法捕获任何超过预定义上下文长度的长程依赖。此外，长度固定的segments都是在不考虑句子或其它语义边界的情况下，通过选择连续的符号块来创建的。因此，模型缺乏必要的上下文信息来很好地预测前几个符号，这就导致模型的优化效率和性能低下。我们将这个问题称为上下文碎片化（context fragmentation）。

为了解决上文提到的上下文固定长度的限制，本文提出了一种叫做Transformer-XL（超长）的新架构。我们将循环（recurrence）概念引入了深度自注意力网络。我们不再从头计算每个新segment的隐藏状态，而是复用从之前segments中获得的隐藏状态。被复用的隐藏状态视为当前segment的memory，而当前的segment为segments之间建立了循环连接（recurrent connection）。因此，超长依赖性建模成为了可能，因为信息可以通过循环连接来传播。同时，从之前的segment传递信息也可以解决上下文碎片化的问题。更重要的是，本文展示了使用相对位置而不是用绝对位置进行编码的必要性，这样做可以在不造成时间混乱（temporal confusion）的情况下，实现状态的复用。因此，作为额外的技术贡献，文本引入了简单但有效的相对位置编码公式，它可以泛化至比在训练过程中观察到的长度更长的注意力长度。

从单词级（word-level）到字符级（character level）的五个语言建模数据集上，Transformer-XL都获得了很好的结果。Transformer-XL在仅基于100M tokens训练的基础上也可以生成相对连贯的长文本文章（参见附录E）。

本文的主要贡献包括：(1)在纯粹的自注意力模型中引入了recurrence的概念，即循环连接。(2)推导了一种新的位置编码方案。这两种技术构成了一组完整的解决方案，因为其中任何一种单独都不能解决上下文长度固定的问题。Transformer-XL是首个从实质上不管是character-level还是word-level都比RNN更优秀的自注意力模型。

2. 相关工作

近年来语言建模领域的进展大概有以下几个方向：
(1)发明新的架构，以更好地对上下文进行编码
(2)改善正则化方法和优化算法
(3)加速Softmax的计算
(4)丰富输出分布的种类
针对语言建模过程中如何捕获长程关系，其中的一个方案是将更为宽广的上下文表征直接输入到网络结构中，作为一个额外的输入。现有的工作还包括手动定义上下文表征和依赖于从数据中学到的文档级主题。
更广泛地说，在一般序列建模中，如何捕获长程依赖关系一直是一个长期的研究问题。从这个角度来看，随着LSTM在这一方面的推广，学者们的研究重点主要集中在消除梯度消，包括如何获得更好的初始化，额外的loss信息，增强结构的记忆能力及其演变出其他RNNs的变种结构以简化优化。与先哲们不同的是，本文的工作是基于Transformer，并通过实验表明将语言建模作为现实世界的任务，能够从学习长程依赖关系中获益。

3. 模型

对于给定的语料tokens $\mathbf{x}=\left(x_1,\dots ,x_T\right)$，语言建模的任务是估计其联合概率$P(\mathbf{x})$。该联合概率可以分解为：$P(\mathbf{x})={\prod }_{t}P\left(x_t|{\mathbf{x}}_{<t}\right)$，所以问题可以简化为估计每个条件概率。在本文的工作中我们仍然坚持使用标准神经网络方法对上述条件概率进行建模。通过神经网络，将上下文${\mathbf{x}}_{<\mathbf{t}}$编码为一个固定长度的隐状态，该隐状态再与word embedding相乘得到logits。该logits再输入到Softmax函数中，得到下一个token的概率分布。

3.1 标准的Transformer语言模型

用 Transformer 或自注意力机制进行语言建模的核心问题在于，如何将Transformer训练为可以把任意长度的上下文有效编码为长度固定的表征。在给定无限内存和计算资源的情况下，一种简单的方法是用无条件的Transformer解码器处理整个上下文序列，这和前馈（feed-forward）神经网络相似。但在实际中，资源都是有限的，因此这种方法不可行。AI-Rfou等人于2018年提出一种可行，但是粗略的近似方法。该方法是将整个语料切分为更短的segments，再用每个segment训练模型，这种方式忽略掉历史过往的segments的上下文信息。我们称上述方法为vanilla model。其具体如下图所示：

vanilla model在训练阶段，不论是前向传播还是反向传播，不同segments之间信息是没有交流的。
这种采用固定上下文的方案存在2个主要局限性：
(1)长程关系依赖的最大长度是以segment长度为上限，而segment的长度只有几百个characters。尽管自注意力机制相比于RNNs受梯度消失的影响较少，但是vanilla model不能够完全利用这个优化上的优势。
(2)尽管可以在句子或者其他语义分界上使用padding，但是在实践过程中，为了提高效率，简单地将长文本分割成固定长度的segments已经成为标准做法。然而，简单地将序列分成固定长度的segments将导致上下文碎片问题。

在评估阶段，vanilla model在每个Step都消费与训练期间相同长度的一个segment，但是在最后一个位置只是进行一次预测。然后，在下一个Step中，这个segment只向右移动一个位置，就形成一个新的segment。这个新的segment必须从头开始处理，如上图Figure 1b所示。该方式能够确保每个预测都使用了训练中暴露的尽可能长的文本，还可以缓解训练中遇到的文本碎片问题。但是这种评估方式的计算代价很高。实验结果表明，本文提出的方法能够大大提高计算速度。

3.2 segment-level Recurrence with State Reuse

为了解决使用固定长度上下文的局限性，本文提出在 Transformer 架构中引入了循环机制（recurrence mechanism）。在训练过程中，前一个segment计算的隐藏状态序列是固定的，并被缓存起来，并在模型处理后面的新segment时作为扩展上下文重复使用，如图 2a 所示。

虽然梯度仍然保留在一个segment中，但是这个额外的输入允许网络利用历史上的信息，从而能够对长程依赖关系建模并避免上下文碎片化。从形式上，两个长度为L的连续segments分别记为：${\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]$。 ${\mathbf{h}}_{\tau }^n\in {\mathbb{R}}^{L\times d}$表示第$\tau$个segment（记为${\mathbf{s}}_{\tau }$）对应的第$n$层隐状态序列，其中$d$是隐状态的维度。第$n$层下一个segment ${\mathbf{s}}_{\tau +1}$ 的隐状态计算如下：

$$\begin{array}{l}{\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 },{\mathbf{h}}_{\tau +1}^{n-1}{\mathbf{W}}_k^{\top },{\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 )$表示梯度停止，$\left[{\mathbf{h}}_u\circ {\mathbf{h}}_v\right]$表示2个隐序列的拼接，$\mathbf{W}$表示模型的参数。对比标准的Transformer，本文模型的主要不同在于key值${\mathbf{k}}_{\tau +1}^n$和value值${\mathbf{v}}_{\tau +1}^n$，本文的key和value受拓展的上下文${\tilde{\mathbf{h}}}_{\tau +1}^{n-1}$影响，因此${\mathbf{h}}_{\tau }^{n-1}$需要从之前的segment中缓存下来。

在语料中的每两个segments之间引入 recurrence 机制，本质上是在隐藏层创建了一个segment-level的recurrence。其结果是，能够被使用到的有效上下文可以远远超出两个segments。需要注意的是，${\mathbf{h}}_{\tau +1}^n$ 和 ${\mathbf{h}}_{\tau }^{n-1}$之间的循环依赖是每segment向下移动一层，这与传统的RNN-LMs在同一层的循环是不同的。因此，本文模型依赖的最长值是与层数和segment的长度线性相关的，即$O(N\times L)$，这可以从Figure 2b中的阴影部分看出端倪。这种方式类似截断的BPTT(Mikolov et al., 2010)，BPTT是一项用以改善RNN-LMs的技术。本文与BPTT的不同在于，本文缓存的是一序列的隐状态，而不是单单最后一个。当然，这需要搭配后续介绍的相对位置编码（详见3.3章节）一起使用。

采用本文的方法除了可以获得更长的上下文和解决上下文碎片问题之外，还可以显著提升评估的速度。评估阶段，历史segments的表征是可以被复用的，而像vanilla model那样从头计算。最后，注意到recurrence机制并不局限于前一个segment。在理论上，在GPU容量范围内，我们是可以尽可能地缓存多的历史segments，并在处理当前segment时，将历史信息进行复用以作为额外的上下文。鉴于此，我们可以预设一个横跨多个segments的历史隐状态的长度为$M$，并将这些历史信息记为${\mathbf{m}}_{\tau }^n\in {\mathbb{R}}^{M\times d}$。在实验中，在训练阶段$M$值等于segment的长度，在评估阶段则是增大为数倍。

3.3 相对位置编码

尽管在上一小节中提出的想法极具吸引力，但为了复用隐藏状态，我们还有一个重要技术问题尚未解决。那就是，当复用隐藏状态时，要如何保证位置信息的一致性呢？回顾下之前，在标准的Transformer中，序列的位置信息是通过位置编码$\mathbf{U}\in {\mathbb{R}}^{L_{\max }\times d}$，其中${\mathbf{U}}_i$表示第$i$行，其对应的是一个segment中第$i$个绝对位置的编码信息；$L_{\max }$是规定的建模的最大长度。实际输入到Transformer的是按元素顺序（element-wise）添加的word embeddings和位置编码。如果简单地将这种位置编码引入到本文的recurrence机制中，则隐状态序列的计算如下所示：
$$\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}$是segment ${\mathbf{s}}_{\tau }$的word embedding，$f$是转换函数（transformation function）。注意到：${\mathbf{E}}_{{\mathbf{s}}_{\tau }}$ and ${\mathbf{E}}_{{\mathbf{s}}_{\tau +1}}$都是与同一位置编码${\mathbf{U}}_{1:L}$相关联。其结果是，该模型无法区分$x_{\tau ,j}$和$x_{\tau +1,j，};j=1,\dots ,L$在位置上的差异，从而导致纯粹的性能损失。

进一步说明就是，循环机制引入了新的挑战——原始位置编码将每个segment分开处理，因此，来自不同segment的表征会具有相同的位置编码。例如，第一segment和第二segment的第一个表征将具有相同的编码，虽然它们的位置和重要性并不相同（比如第一个segment中的第一个表征可能重要性低一些）。这种混淆可能会错误地影响网络。

解决这个问题的基本思想是只编码隐藏状态中的相对位置信息。观念上，位置编码给予模型时序上的信息或者信息收集过程中的“偏见”，即需要更被注意到的地方。同样为了这个目的，我们不再静态地将偏置bias引入到初始的embedding中，而是将相同的信息注入到每一层的注意分值中。更重要的是，用相关的方法来定义时间偏差更直观、更普遍。例如，一个query vector $q_{\tau ,i}$关注到key vectors ${\mathbf{k}}_{\tau ,\le i}$，在识别segment的时间顺序时，是无需知晓每个key vector的绝对位置；而只需要知道每个key vector $k_{\tau ,j}$和自身$q_{\tau ,i}$之间的相对位置，即$i-j$。实际上，我们可以创建一序列的相对位置编码$\mathbf{R}\in {\mathbb{R}}^{L_{\max }\times d}$，其中第$i$行的值${\mathbf{R}}_i$表示2个位置之间相对距离为i。通过动态地将相对距离引入到attention score中，query vector从不同的距离很容易就区分出$x_{\tau ,j}$和$x_{\tau +1,j}$的表征，这也使得状态的复用成为可能。与此同时，由于绝对位置可以从相对距离递归恢复，模型并没有丢失任何的时序信息。

在此之前，相对位置编码的想法已经在机器翻译（Shaw et al.，2018）和音乐生成（Huang et al.，2018）中探索过。我们在此进行了不同的推导，得到了相对位置编码的新形式，它不仅与绝对位置编码有一对一的对应关系，而且从经验上讲它的泛化效果更好。首先，在标准的Transformer中，一个segment中，query $q_i$ 和 key vector $k_j$ 之间的attention score可以通过以下进行分解：
$$\begin{array}{rl}{\mathbf{A}}_{i,j}^{\mathrm{a}\mathrm{b}\mathrm{s}}& =\underset{(a)}{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{E}}_{x_j}}+\underset{(b)}{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{U}}_j}\\ & +\underset{(c)}{{\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{E}}_{x_j}}+\underset{(d)}{{\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_k{\mathbf{U}}_j}\end{array}$$
基于只依赖相对位置编码的想法，本文提出重新参数化上述4项：
$$\begin{array}{rl}{\mathbf{A}}_{i,j}^{\mathrm{r}\mathrm{e}\mathrm{l}}& =\underset{(a)}{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_{k,E}{\mathbf{E}}_{x_j}}+\underset{(b)}{{\mathbf{E}}_{x_i}^{\top }{\mathbf{W}}_q^{\top }{\mathbf{W}}_{k,R}{\mathbf{R}}_{i-j}}\\ & +\underset{(c)}{u^{\top }{\mathbf{W}}_{k,E}{\mathbf{E}}_{x_j}}+\underset{(d)}{v^{\top }{\mathbf{W}}_{k,R}{\mathbf{R}}_{i-j}}\end{array}$$

• 第一个改动替换掉所有使用绝对位置编码${\mathbf{U}}_j$的项。在(b)项和(d)项中计算key vectors的时候用相对位置${\mathbf{R}}_{i-j}$代替。这本质上反映了一个前提，即只有相对的距离才会影响到注意到的地方。需要注意的一点是，$\mathbf{R}$是正弦信号编码矩阵，并不是可学习到的参数。

• 第二点不同是，在©项中引入一个可训练的参数项$u\in {\mathbb{R}}^d$，以替换掉query ${\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }$。由于query vector对于所有的查询位置都是相同的。这表明，无论查询位置如何，对不同词语的注意偏差，都应保持一致。出于类似的考虑，在(d)项中用$v\in {\mathbb{R}}^d$替换掉${\mathbf{U}}_i^{\top }{\mathbf{W}}_q^{\top }$。

• 最后，本文刻意设计2个权重矩阵${\mathbf{W}}_{k,E}$ 和 ${\mathbf{W}}_{k,R}$以分别生成content-based的key vectors和location-based的key vectors。

在新的参数化方式下，每一项都有其直观的意义：
(a)项表示content-base的定位
(b)项表示捕获content-dependent的位置bias
(c)项控制全局的content bias
(d)项表示全局位置bias的编码

在本文的循环机制中引入上述的相对位置编码，即为本文提出的Transformer-XL架构。出于完整性的考虑，总结下$N$层 单attention head的Transformer-XL的计算流程：
$$\begin{gathered} {\mathbf{h}}_{\tau }^{n-1}=\left[\mathrm{S}\mathrm{G}\left({\mathbf{m}}_{\tau }^{n-1}\right)\circ {\mathbf{h}}_{\tau }^{n-1}\right] \\ {\mathbf{q}}_{\tau }^n,{\mathbf{k}}_{\tau }^n,{\mathbf{v}}_{\tau }^n={\mathbf{h}}_{\tau }^{n-1}{{\mathbf{W}}_q^n}^{\top },{\mathbf{h}}_{\tau }^{n-1}{\mathbf{W}}_{k,E}^{n\top },{\mathbf{h}}_{\tau }^{n-1}{\mathbf{W}}_v^{n\top } \\ {\mathbf{A}}_{\tau ,i,j}^n={\mathbf{q}}_{\tau ,i}^n\top {\mathbf{k}}_{\tau ,j}^n+{{\mathbf{q}}_{\tau ,i}^n}^{\top }{\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j} \\ +u^{\top }{\mathbf{k}}_{\tau ,j}+v^{\top }{\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j} \\ \begin{array}{l}{\mathbf{a}}_{\tau }^n=\text{ Masked-Softmax }\left({\mathbf{A}}_{\tau }^n\right){\mathbf{v}}_{\tau }^n\\ {\mathbf{o}}_{\tau }^n=\text{ LayerNorm }\left(\mathrm{Linear}\left({\mathbf{a}}_{\tau }^n\right)+{\mathbf{h}}_{\tau }^{n-1}\right)\\ {\mathbf{h}}_{\tau }^n=\text{ Positionwise-Feed-Forward }\left({\mathbf{o}}_{\tau }^n\right)\end{array} \end{gathered}$$
其中$n=1,...,N$，${\mathbf{h}}_{\tau }^0:={\mathbf{E}}_{{\mathbf{s}}_{\tau }}$设为word embedding序列。值得一提的是，如果采用简单的方式计算矩阵$A$，由于其组成是需要计算所有$(i,j)$对的${\mathbf{W}}_{k,R}^n{\mathbf{R}}_{i-j}$，则时间复杂度是序列长度的二次方。但是，需要注意的是，$i-j$的值范围在0和序列长度之间，本文采用一种更为灵巧的方式（附录B）计算获得该矩阵，其计算复杂度与序列长度是线性关系。

4. 实验

细节请参考原文，此处略

4.2 消融研究

本文还在循环机制和相对位置上做了消融研究。

4.3 相对有效的上下文长度

Khandelwal et al. (2018)提出用有效的上下文长度（Effective Context Length，ECL）以评估序列模型。ECL是能够超过阈值增益的最大上下文跨度。但是ECL忽略了一个事实，当一个模型的perplexity已经比较低时，采用一个更短的上下文是很难将其提升的。因此，在多个模型之间如此比较其实是有失公允的。鉴于此，本文提出一个新的评估准则：Relative Effective Context Length (RECL)，在这里姑且称其为相对有效的上下文长度。RECL是基于一组模型，而不是单个模型。上下文长度所带有的增益是对比best short context model的相对改进来衡量的。model 群组是共享同一baseline，如此可以公平比较。RECL有一个参数$r$，表示约束在$top-r$ hard examples上的比较，细节可以参考附录C。

4.4 生成文本

仅在中等数据规模的WikiText-103训练后，Transformer-XL在没有人工干预下，能够产生相对连贯的文章。具体参见附件E。

4.5 评估速度

最后，我们将模型的评估速度与vanilla Transformer模型进行了比较。如表9所示，与Al-Rfou et al. (2018).的架构相比，由于state reuse方案，Transformer-XL的评估速度加快了高达1874倍。

5. 结论

本文提出了一种新的架构，Transformer-XL，这是一个超出了固定长度的上下文限制的self-attention的语言建模架构。本文主要技术贡献在于在纯粹的 self-attentive模型中引入递归的概念，并推导出一种新的位置编码方案。这两种技术形成了一套完整的解决方案，因为它们中的任何一种单独都不能解决固定长度上下文的问题。

Transformer-XL是第一个在字符级和单词级语言建模方面都取得了比RNN更好结果的self-attention模型。Transformer-XL还能够建模比RNN和Transformer更长期的依赖关系，并且与vanilla Transformers相比在评估过程中取得了显著的加速。