解析Transformer-XL

Link Start

在很久之前的文章理解Bert 中我们介绍了Transformer，它作为CNN和RNN之后一个划时代意义的特征处理器是无比强大的，但实际用起来，还是会发现它存在着一些较为尴尬的缺陷，而为了解决其中的某些缺陷，便发展出了Transformer-XL.

Transformer存在的问题

众所周知，Transformer相对于RNN的一个重大优势在于它能捕捉长距离的依赖关系。但是，我们在实际应用的过程中会发现，多头注意力机制和RNN不同，不是循环共享参数这么简单的，它为了保存这种依赖关系，需要存储大量的参数，并且随着序列长度的增加，参数数量以平方倍的速度在增加。这对于我们的GPU来说是一个灾难，随着序列长度的扩张，且不论因增加的计算量而导致的延迟，OOM是迟早的事情。这也是为什么那些预训练的Bert模型的长度基本在512以内的原因。
我们在处理长文本的时候，通常是将其切分成更小的片段（segment），然后对每个片段进行单独的处理。而且为了保持切分的容错率，尽可能的不丢失语义信息，我们在且分的时候还会设置一个1/3左右的重叠区，避免某些词汇别割裂开而导致的语义变化。

长序列Segment切分

但是这样子切分，还是无法避免前面segment之间无交集部分的关系依赖。举个例子：

“小明考上了清华大学实现了他小时候的梦想。”

然后这句话中的小明和他被切分到了两个segment中：

句子切分

此时，他的encoder向量将无法获取到任何小明的信息，因此便无法得知他指代的是谁，这对于相关任务而言可能十分不利。不仅如此，对于大篇幅的文章的某些判断，通过切分成segment之后进行的断章取义的分析再组合回去是否还能还原文章原汁原味的特征这一点也很值得怀疑。

盲人摸象

因此，我认为，使用某种方式，来保留整体的特征信息是很有必要的，这也是Transformer-XL主要解决的问题。

Transformer-XL是如何解决超长序列关系依赖的

Transformer-XL 采用了一种与RNN类似的思想。RNN是将前一个时间步的状态信息往下一个时间步传递，那么我们为何不能也这样搞？——将前一个时间步的segment的状态信息传往下一个segment。
于是它真的就这么干了：

Transformer-XL记忆传递

哈?... 说起来轻巧，你..你这是纸...纸上谈兵 ..说起来轻巧，关键是这玩意要怎么实现呢？

当然就是RNN是怎么搞的，我就依葫芦画瓢呗！
原始的encoder过程是这样的，原始输入加上位置编码，经过多层的encoder block：

transformer encoder

我们将这整个东西作为一个循环体，然后将每一步的输出特征复制一份传给下一个segment的循环：

Transformer-XL记忆传递细节

可以说是跟RNN一样一样的...

connect

我们再来看一下每一处的连接细节。
其实没什么特别的，就是简单的拼接即可：
假设前一个segment的输出tensor是 $h_{t-1,i}$ ,当前segment的输出tensor是 $h_{t,i}$ , 它们的形状都是[batch_size, segment_len, emb_dim], 把它们拼接起来之后就是[batch_size, segment_len*2, emb_dim] ,举个实际的例子：

张量拼接

我们将这一个操作记作

cat\_h_{t,i}

表示第t个segment第i个block的连接操作。

需要注意的是，在训练期间，我们对记忆的那部分取消了反向传播。否则，我们仍然会有关于计算资源的原始问题。而实际上 $h_{t-1,i}$ 和 $cat\_h_{t,i}$ 都是注意力计算中的重要组成部分。

注意力计算的过程和Transformer并没什么区别，出了形状上，每个query需要额外计算前面传过来向量的的注意力，也就是形状变成了(batch_size, segment_length, 2*segment_length)：

mask

mask的方式也只需对当前输入的矩阵segment_length到2*segment_length这部分进行mask即可，因为前文信息是可见的。

然后再说一下位置编码，在transformer的原始论文中的位置信息通过正余弦函数来获取，通过这种方式的话位置不同segment的位置信息会重置，绝对编码会有问题，因此需要改成相对编码的方式。
但实际Bert使用时位置信息是通过神经网络来计算的，我想Transformer-XL中也可以通过神经网络来模拟位置信息，只要根据sengment序号截取即可，就不存在位置编码冲突的问题。想了解具体相对位置编码实现的同学可以看一下原始论文Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context.

以上..Transformer-XL初看起来有点迷...但是仔细研究一下它的结构发现还是蛮简单的。接下来大概会抽个时间说一说它的TensorFlow实现。

参考资料

Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context
Dissecting Transformer-XL