论文导读 | Transformer架构下面向长文本的语言模型

问题介绍

自BERT出现之后，基于Transformer[1]架构的语言模型在各项NLP任务上取得了统治性的表现。原始的Transformer架构可以分为编码器和解码器两端，如图所示。在底层输入中，作者使用了位置编码（positional embedding）以注入序列的顺序信息；在每层单元中，都使用了多头注意力机制、残差连接和层归一化的组成单元。在解码时，解码器的注意力头可以同时看见解码端之前的输出和编码器的隐层表示以计算注意力权重。

根据编码方式，主流的语言模型架构又可以分为三类：只使用编码器的自编码型语言模型，如BERT；只使用解码器的自回归型语言模型，如GPT；以及二者皆使用的编码器-解码器型语言模型，如BART。

在计算注意力权重时，可以注意到其中关键的权重计算需要一个O(N^2)规模的矩阵乘法QK^T。由于这个计算的运行规模随序列长度N二次增长，因此在现行的硬件设备上限制了Transformer的单次输入规模。所以，许多预训练语言模型显式地将最大长度限制设置成512。

 解决方案

那么，对于超出了这一长度限制的长文本，我们又能如何用语言模型编码呢？一个最直接的想法是，采用分治法的思想，将原始文本切分成符合长度限制的文本段分别编码，然后再将得到的表示拼接起来。但这种做法存在的问题也显而易见：在编码不同文本段时看不到别的部分的信息，从而产生了上下文碎片化（context fragmentation）的问题。

本文将介绍三类在架构层次改进的方法。

一、循环机制（recurrence）

针对上下文碎片化的问题，Transformer-XL[2]提出了一种循环机制加以克服。作者对原始的分治法做出了简单的改进：在编码每个文本段时，同时复用上一文本段的隐层状态，但是停止这部分的梯度回传。以SG指代stop gradient操作，Transformer-XL的编码过程为：

在复用文本段隐状态的计算下，原始的位置编码不能区分当前和上一文本段中同一位置的词序差别。为了克服这一问题，Transformer-XL同时引入了相对位置编码（relative positional encoding）。

总体来讲，Transformer-XL的解决方案不仅能够保证线性的计算复杂度，同时随着层数L的加深，上下文的感受野范围最远可达N*L的距离。但是，这种方案只适用于自回归型的语言模型，只有单向的信息传递，无法扩展到自编码型的语言模型中。

2、稀疏特征（sparse patterns）

稀疏特征方法的主要想法是，将原始的注意力权重矩阵稀疏化，从而减少计算规模。目前主流的策略可以分为两类：基于内容的稀疏策略，和固定特征的稀疏策略。我们将分别介绍两种代表性的方法。

2.1 基于内容的稀疏策略（content based）

Reformer[3]是此类中比较有代表性的一种方法。他的想法是，通过局部敏感哈希（locality sensitive hashing, LSH）技术，快速在高维空间中寻找近邻，从而将注意力计算约束在这些近邻上计算。因为注意力分数（在归一化的前提下）可以被认为是通过余弦相似度经过softmax操作获得，因此近邻上的注意力权重较大，而其余向量对于新表示的贡献小，从而可被近似忽略。

Reformer对于注意力操作的近似可如下描述：

（1）选取合适的哈希函数，满足如果x和y距离度量小则高概率有h(x)=h(y)；

（2）对Q和K矩阵使用LSH；

（3）只计算同一个桶内q和k的注意力加权结果。

具体实现上，为了加速并行化，作者将原始序列拆分成块，并每次在两块内计算注意力。如图所示。

2.2 固定特征的稀疏策略（pre-specified）

固定特征即为对任意输入，都只计算相同位置的注意力。在Longformer[4]中，作者在局部特征上使用了滑动窗口策略，在全局特征上选取了一些全局点（global token）计算与其余所有输入的注意力。这两种特征也被其他工作广泛使用。

与CNN、Transformer-XL类似，随着层数的加深，滑动窗口的感受野也会随之线性增加。

在具体实现上，由于固定特征难以写成单一的矩阵乘法操作，因此作者借助了TVM工具自定义了CUDA内核实现，在与完整矩阵乘法相近的时间复杂度上，大大减小了空间复杂度。并且，由于我们可以简单地将原始注意力操作替换成此种稀疏操作，因此Longformer很好地适用于各种语言模型，并在长文本相关的下游任务上取得了良好的表现。

3、低秩近似与核方法（low-rank/kernel based methods）

最后一类方法同样目标是简化注意力机制的计算复杂度。不同于在矩阵上寻找稀疏特征，此类方法希望通过核方法等工具近似注意力中的矩阵乘法。但是，虽然这类方法能取得良好的计算复杂度，在实际任务上可能表现反而不及原始的Transformer。

Linear Transformer[5]中，作者注意到原始的注意力公式等价于如下形式：

因此作者直接使用核函数φ的点积替代了sim函数：

并选取

，以避免梯度消失问题。利用上述形式计算，作者取得了线性的计算复杂度。

在Linformer[6]中，作者首先从理论和实验两方面证明，原始的注意力矩阵可以被低秩矩阵近似替代。利用这个结论，作者在计算注意力前，将K和V从[seq_len, d]投影到低维空间[k, d]中加速计算，使计算规模从O(N^2)缩减为O(Nk)。

总结

本文从架构设计层面介绍了一系列面向长文本的语言模型设计。总体而言，大多数工作的目标都是将计算限制在线性复杂度（w.r.t. 序列长度N）以内；但不足之处在于，受限于计算资源，许多工作在实验时，没有以“预训练-精调”的范式在下游NLP任务上验证模型效果，同时也只有少数工作（e.g. Longformer）提供了社区开源的预训练模型。在未来，统一长文本上的预训练语料和下游测试任务（类似于GLUE和CLUE的基准测试），是进一步全面准确评估长文本上的语言模型的关键。

参考文献

[1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems30 (2017).

[2] Dai, Zihang, et al. "Transformer-XL: Attentive Language Models beyond a Fixed-Length Context." Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. 2019.

[3] Kitaev, Nikita, Lukasz Kaiser, and Anselm Levskaya. "Reformer: The Efficient Transformer." International Conference on Learning Representations. 2019.

[4] Beltagy, Iz, Matthew E. Peters, and Arman Cohan. "Longformer: The long-document transformer." arXiv preprint arXiv:2004.05150 (2020).

[5] Katharopoulos, Angelos, et al. "Transformers are RNNs: Fast autoregressive transformers with linear attention." International Conference on Machine Learning. PMLR, 2020.

[6] Wang, Sinong, et al. "Linformer: Self-attention with linear complexity." arXiv preprint arXiv:2006.04768(2020).