Transformer-XL: 非固定长度上下文的注意力语言模型（Attentive Language Models Beyond a Fixed-Length Context）

Reference
1. Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context
2. Transformer-XL: Unleashing the Potential of Attention Models
3. Transformer-XL介绍

代码基于tensorlfow2.3实现，仓库地址：https://github.com/dwdb/transformer-xl

---

Transformers潜在地学习长期依赖，但受到上下文固定长度限制，当处理序列长度超过固定长度时，会出现以下问题：

• 训练阶段，需将输入文本分割成不同分段，不同分段独立训练，由于分段未考虑语义边界，可能造成模型缺乏上下文信息学习分段的完整表示，模型不易优化，这种现象称为上下文碎片问题；
• 预测阶段，每次移动一个输入单元，引入大量重复计算，预测效率低；

Transformer-XL使用两种技术： 循环分段 和 相对位置编码，克服vanilla Transformers的缺陷，并解决上下文碎片问题。使用循环分段机制，在处理当前分段时，重用之前分段的隐状态（缓存），可不是从头计算新分段的隐状态，不同分段之间不再独立，解决了上下文碎片的问题。使用相对位置编码，而非绝对位置编码，避免利用之前分段隐状态造成的时序混乱问题。

实验结果表明，Transformer-XL在学习上下文依赖上，比RNNs网络长0.8倍、比vanilla Transformers网络长4.5倍，Transformer-XL是首个在字和词级别上均优于RNNs的使用自我注意力的模型。

Vanilla Transformer Language Models

使用Transformer或者self-attention的语言模型，最核心问题是怎样把任意长、具有上下文信息的序列编码为固定长度的向量表示。有限的计算资源下，无法处理较长序列，可行地做法是将长序列分割成数个固定长度序列，各分段独立训练，忽略各分段间的语义关系，随意分割会造成上下文碎片。

Segment-Level Recurrence with State Reuse

训练阶段，缓存一定长度的之前片段的各层隐状态向量，在处理新分段时，缓存向量作为新分段的扩展上下文重用，使得模型可以学习长期依赖，以避免学习上下文碎片。

对于两个连续分段$s_{\tau }=[x_{\tau ,1},\cdots ,x_{\tau ,L}]$和$s_{\tau +1}=[x_{\tau +1,1},\cdots ,x_{\tau +1,L}]$，$s_{\tau }$在第$n$层的隐状态为${\mathbf{h}}_{\tau }^n\in {\mathbb{R}}^{L\times d}$，其中$d$是隐状态向量维度，则
$$\begin{array}{rl}& {\tilde{\mathbf{h}}}_{\tau +1}^{n-1}=[\text{SG}({\mathbf{h}}_{\tau }^{n-1})\circ {\mathbf{h}}_{\tau +1}^{n-1}]\\ & {\mathbf{q}}_{\tau +1}^n,{\mathbf{k}}_{\tau +1}^n,{\mathbf{v}}_{\tau +1}^n={\mathbf{h}}_{\tau +1}^{n-1}{W}_q^{\top },{\tilde{\mathbf{h}}}_{\tau +1}^{n-1}{W}_k^{\top },{\tilde{\mathbf{h}}}_{\tau +1}^{n-1}{W}_v^{\top }\\ & {\mathbf{h}}_{\tau +1}^n=\text{Transformer-Layer}({\mathbf{q}}_{\tau +1}^n,{\mathbf{k}}_{\tau +1}^n,{\mathbf{v}}_{\tau +1}^n)\end{array}$$

式中函数$\text{SG}(\cdot )$表示不计算梯度，$[{\mathbf{h}}_u\circ {\mathbf{h}}_v]$表示序列长度方向拼接两个隐藏状态序列。

循环分段状态重用与标准Tranformer最大的不同在于，利用当前分段的$\mathbf{q}$向量，以及之前分段和当前分段的$\mathbf{k}$和$\mathbf{v}$向量，计算当前分段的Transformer层输出，使得当前分段输出考虑到之前分段信息（self-attention注意之前分段）。

从图二左图中可看出，在训练阶段当仅利用前一个分段信息时，两个分段的不同层的隐状态${\mathbf{h}}_{\tau +1}^n$和${\mathbf{h}}_{\tau }^{n-1}$具有依赖关系，为保持时序信息，需考虑相对位置信息，下节介绍。

循环分段机制除能够学习长期依赖、解决上下文碎片化之外，对预测的性能上也有较大提高。模型通过学习转换矩阵得到固定嵌入，而不是直接学习嵌入，使得预测阶段可以学习更长期的依赖。。此外，在当GPU内存允许条件下，也可利用之前多个分段信息。

Relative Positional Encodings

在使用之前隐状态时，如何保证连贯的位置信息？ 传统的Transformer中，使用绝对位置编码$U\in {\mathbb{R}}^{L_{\max }\times d}$，其第$i$行表示分段中的第$i$个绝对位置的编码向量，$L_{\max }$表示最大编码长度，实际是直接将词向量和绝对位置编码向量按元素相加作为实际输入，各分段的处理方式相同。显然，不同分段同时参与运算时，会造成时序混乱。

避免时序混乱的思想是仅在隐状态中引入相对位置信息。位置嵌入目的是给予模型各输入点的时序线索或偏差，以决定如何收集信息，因此，可向每一层的注意力分数中注入相对时序信息，取代将位置编码直接加入初始词向量。

举例来说，对于计算查询向量${\mathbf{q}}_{\tau ,i}$在键向量${\mathbf{k}}_{\tau ,\le i}$的注意力，我们不需知道${\mathbf{k}}_{\tau ,j}$在时序中绝对位置，只需要知道其相对${\mathbf{q}}_{\tau ,i}$的时序偏差（相对查询的位置偏差）即可，如$i-j$。

因此，创建一系列相对位置编码向量$R\in {\mathbb{R}}^{L_{\max }\times d}$（正弦信号），其中$R_i$表示两位置相对距离为$i$对应的编码向量，通过在注意力分数中动态地引入相对距离，查询向量可通过相对距离的不同区分$x_{\tau ,j}$和$x_{\tau +1,j}$。

标准Transformer使用绝对位置编码，同一分段中查询向量${\mathbf{q}}_i$对键向量${\mathbf{k}}_j$的注意力分数为
$$\begin{array}{rl}{A}_{i,j}^{\text{abs}}& =W_q(E_{x_i}+U_i)\cdot W_k(E_{x_j}+U_j)\\ & =\underset{(a)}{\underbrace{E_{x_i}^{\top }{W}_q^{\top }{W}_k{E}_{x_j}}}+\underset{(b)}{\underbrace{E_{x_i}^{\top }{W}_q^{\top }{W}_k{U}_j}}+\underset{(c)}{\underbrace{U_i^{\top }{W}_q^{\top }{W}_k{E}_{x_j}}}+\underset{(d)}{\underbrace{U_i^{\top }{W}_q^{\top }{W}_k{U}_j}}\end{array}$$

使用相对位置编码，则同一分段中查询向量${\mathbf{q}}_i$对键向量${\mathbf{k}}_j$的注意力分数为
$$\begin{array}{rl}{A}_{i,j}^{\text{rel}}& =\underset{(a)}{\underbrace{E_{x_i}^{\top }{W}_q^{\top }{W}_{k,E}{E}_{x_j}}}+\underset{(b)}{\underbrace{E_{x_i}^{\top }{W}_q^{\top }{W}_{k,R}{R}_{i-j}}}+\underset{(c)}{\underbrace{u^{\top }{W}_{k,E}{E}_{x_j}}}+\underset{(d)}{\underbrace{v^{\top }{W}_{k,R}{R}_{i-j}}}\end{array}$$
相对位置编码的改动在于：

• 将用于计算键向量的绝对位置编码$U_j$，替换为不需要学习的正弦的相对位置编码$R_{i-j}$，使得扩展上下文长度在预测阶段可大于训练阶段；
• 引入参数$u$取代$(c)$项中的$W_q{U}_i$，此时$(c)$项仅与内容有关。为使序列中任意位置对其它任意位置的注意力偏差相同，引入参数$v$取代$(d)$项中的$W_q{U}_i$，此时$(d)$项仅与相对位置有关；
• 使用权重矩阵$W_{k,E}$生成基于内容的键向量（右乘词向量），$W_{k,R}$生成基于位置的键向量（右乘相对位置向量）；
• 四项意义：$(a)$项为内容表示，$(b)$项为依赖位置的内容偏差，$(c)$项为全局内容偏差，$(d)$项为全局位置偏差；

Transformer-XL的整体架构表示为
$$\begin{array}{rl}{\tilde{\mathbf{h}}}_{\tau }^{n-1}& =[\text{SG}({\mathbf{m}}_{\tau }^{n-1})\circ {\mathbf{h}}_{\tau }^{n-1}]\\ {\mathbf{q}}_{\tau }^n,{\mathbf{k}}_{\tau }^n,{\mathbf{v}}_{\tau }^n& ={\mathbf{h}}_{\tau }^{n-1}{W_q^n}^{\top },{\tilde{\mathbf{h}}}_{\tau }^{n-1}{W_{k,E}^n}^{\top },{\tilde{\mathbf{h}}}_{\tau }^{n-1}{W_v^n}^{\top }\\ A_{\tau ,i,j}^n& ={{\mathbf{q}}_{\tau ,i}^n}^{\top }{\mathbf{k}}_{\tau ,j}^n+{{\mathbf{q}}_{\tau ,i}^n}^{\top }{W}_{k,R}^n{R}_{i-j}+u^{\top }{\mathbf{k}}_{\tau ,j}+v^{\top }{W}_{k,R}^n{R}_{i-j}\\ {\mathbf{a}}_{\tau }^n& =\text{Masked-Softmax}(A_{\tau }^n){\mathbf{v}}_{\tau }^n\\ {\mathbf{o}}_{\tau }^n& =\text{LayerNorm}(\text{Linear}({\mathbf{a}}_{\tau }^n)+{\mathbf{h}}_{\tau }^{n-1})\\ {\mathbf{h}}_{\tau }^n& =\text{Positionwise-Feed-Forward}({\mathbf{o}}_{\tau }^n)\end{array}$$
式中，计算$A_{\tau }^n$，意味着需对所有位置对$(i,j)$计算$W_{k,R}^n{R}_{i-j}$，时间复杂度$O(n^2)$，优化后可降至$O(n)$。