Transformer-XL解读（论文 + PyTorch源码）

前言

目前在NLP领域中，处理语言建模问题有两种最先进的架构：RNN和Transformer。RNN按照序列顺序逐个学习输入的单词或字符之间的关系，而Transformer则接收一整段序列，然后使用self-attention机制来学习它们之间的依赖关系。这两种架构目前来看都取得了令人瞩目的成就，但它们都局限在捕捉长期依赖性上。

为了解决这一问题，CMU联合Google Brain在2019年1月推出的一篇新论文《Transformer-XL：Attentive Language Models beyond a Fixed-Length Context》同时结合了RNN序列建模和Transformer自注意力机制的优点，在输入数据的每个段上使用Transformer的注意力模块，并使用循环机制来学习连续段之间的依赖关系。Transformer-XL在多种语言建模数据集（如单词级别的enwik8和字符级别的text8）上实现了目前的SoTA效果，且该模型在推理阶段速度更快，比之前最先进的利用Transformer进行语言建模的方法快300～1800倍。 同时，该论文也放出了其配套源码（包括TensorFlow和PyTorch的）、预训练模型及在各个数据集上训练的超参数，可以说是非常良心了～造福我等伸手党！

本文将主要针对模型原理及其PyTorch实现进行逐一对照解读，因笔者能力有限，如有不详尽之处，可移步文末的传送门进行详细阅读，并欢迎指出～

一. 回顾Transformer

在NLP领域中，一种对语言建模的最常用模型就是RNN，它可以捕捉单词之间的依赖关系。但因为梯度消失和爆炸的问题，RNN变得非常难以训练，LSTM单元和梯度裁剪方法的提出也不足以解决此类问题。同时RNN网络的计算速度往往很慢，其学习长期依赖的能力也较为有限（论文中提到，LSTM语言模型平均只能建模200个上下文词语）。

2017年6月，Google Brain在论文《Attention Is All You Need》中提出的Transformer架构，完全摒弃了RNN的循环机制，采用一种self-attention的方式进行全局处理。其接收一整段序列，并使用三个可训练的权重矩阵——Query、Key和Value来一次性学习输入序列中各个部分之间的依赖关系。Transformer网络由多个层组成，每个层都由多头注意力机制和前馈网络构成。由于在全局进行注意力机制的计算，忽略了序列中最重要的位置信息。Transformer为输入添加了位置编码（Positional Encoding），使用正弦函数完成，为每个部分的位置生成位置向量，不需要学习，用于帮助网络学习其位置信息。其示意如下图所示：

有关Transformer的更深入讨论，可参考笔者之前的博客：

Transformer（论文 + PyTorch源码解读）

二. vanilla Transformer

为何要提这个模型？因为Transformer-XL是基于这个模型进行的改进。

Al-Rfou等人基于Transformer提出了一种训练语言模型的方法（ https://arxiv.org/abs/1808.04444 ），来根据之前的字符预测片段中的下一个字符。例如，它使用$x_1,x_2,...,x_{n-1}$预测字符$x_n$，而在$x_n$之后的序列则被mask掉。论文中使用64层模型，并仅限于处理 512个字符这种相对较短的输入，因此它将输入分成段，并分别从每个段中进行学习，如下图所示。 在测试阶段如需处理较长的输入，该模型会在每一步中将输入向右移动一个字符，以此实现对单个字符的预测。

该模型在常用的数据集如enwik8和text8上的表现比RNN模型要好，但它仍有以下两个缺点：

a. 上下文长度受限：字符之间的最大依赖距离受输入长度的限制，模型看不到出现在几个句子之前的单词。
b. 上下文碎片：对于长度超过512个字符的文本，都是从头开始单独训练的。段与段之间没有上下文依赖性，会让训练效率低下，也会影响模型的性能。
c. 推理速度慢：在测试阶段，每次预测下一个单词，都需要重新构建一遍上下文，并从头开始计算，这样的计算速度非常慢。

三. Transformer-XL

Transformer-XL架构在vanilla Transformer的基础上引入了两点创新：循环机制（Recurrence Mechanism）和相对位置编码（Relative Positional Encoding），以克服vanilla Transformer的缺点。与vanilla Transformer相比，Transformer-XL的另一个优势是它可以被用于单词级和字符级的语言建模。

1. 引入循环机制

与vanilla Transformer的基本思路一样，Transformer-XL仍然是使用分段的方式进行建模，但其与vanilla Transformer的本质不同是在于引入了段与段之间的循环机制，使得当前段在建模的时候能够利用之前段的信息来实现长期依赖性。如下图所示：

在训练阶段，处理后面的段时，每个隐藏层都会接收两个输入：

1. 该段的前面隐藏层的输出，与vanilla Transformer相同（上图的灰色线）。
2. 前面段的隐藏层的输出（上图的绿色线），可以使模型创建长期依赖关系。

这两个输入会被拼接，然后用于计算当前段的Key和Value矩阵。对于某个段的某一层的具体计算公式如下：

其中，$\tau$表示第几段，$n$表示第几层，$h$表示隐层的输出。$SG(\cdot )$表示停止计算梯度，$[h_u\circ h_v]$表示在长度维度上的两个隐层的拼接，$W_{.}$是模型参数。乍一看与Transformer中的计算公式很像，唯一关键的不同就在于Key和Value矩阵的计算上，即$k_{\tau +1}^n$和$v_{\tau +1}^n$，它们基于的是扩展后的上下文隐层状态${\tilde{h}}_{\tau +1}^{n-1}$进行计算，$h_{\tau }^{n-1}$是之前段的缓存。

原则上只要GPU内存允许，该方法可以利用前面更多段的信息，测试阶段也可以获得更长的依赖。

在测试阶段，与vanilla Transformer相比，其速度也会更快。在vanilla Transformer中，一次只能前进一个step，并且需要重新构建段，并全部从头开始计算；而在Transformer-XL中，每次可以前进一整个段，并利用之前段的数据来预测当前段的输出。

2. 相对位置编码

在Transformer中，一个重要的地方在于其考虑了序列的位置信息。在分段的情况下，如果仅仅对于每个段仍直接使用Transformer中的位置编码，即每个不同段在同一个位置上的表示使用相同的位置编码，就会出现问题。比如，第$i-2$段和第$i-1$段的第一个位置将具有相同的位置编码，但它们对于第$i$段的建模重要性显然并不相同（例如第$i-2$段中的第一个位置重要性可能要低一些）。因此，需要对这种位置进行区分。

论文对于这个问题，提出了一种新的位置编码的方式，即会根据词之间的相对距离而非像Transformer中的绝对位置进行编码。在Transformer中，第一层的计算查询$q_i^T$和键$k_j$之间的attention分数的方式为：

其中，$E_{x_i}$是词$i$的embedding，$E_{x_j}$