Transformer家族纵览

paper：https://arxiv.org/abs/2106.04554.pdf

前言

        转眼间Transformer的提出已经有四年之久了，Transformer依靠弱归纳偏置和易于并行的特点，在人工智能领域大放异彩，尤其在NLP和CV领域。从Transformer论文《Attention is All You Need》的题目来看，有些言过其实了，事实上论文《Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth》阐明了纯自注意力网络（SAN）的操作和归纳偏置，并发现纯粹注意力随着深度以双倍指数方式衰减。论文提出了一种通过路径分解（path decomposition）分析 SAN 的新方法，将 SAN 表示为一个浅层网络的集合（ensemble），如下图Figure1所示：

       并且论文使用Bert、Albert、XLNet三个模型进行随深度增加的训练前和训练后的相对范数变化的实验，实验结果如下图Figure2所示：

纯注意力（SAN）网络的输出迅速收敛到秩为1的矩阵（绿色点虚线）。

添加MLP模块和跳过连接可得到一个Transformer（紫色实线）。

实验结论：跳过连接（绿色点横虚线）在缓解秩崩溃（即零残差）方面起着至关重要的作用，MLP（绿色横虚线）可以减缓收敛速度。

    尽管Transformer已经被证明有很好的通用性，但它也存在一些明显的问题，例如：

1）核心模块自注意力对输入序列长度有平方级别的复杂度，这使得Transformer对长序列应用不友好。例如一个简单的32x32图像展开就会包括1024个输入元素，一个长文档文本序列可能有成千上万个字，因此有大量现有工作提出了轻量化的注意力变体（例如稀疏注意力），或者采用『分而治之』的思路（例如引入recurrence）；

2）与卷积网络和循环网络不同，Transformer结构几乎没有什么归纳偏置。这个性质虽然带来很强的通用性，但在小数据上却有更高的过拟合风险，因此可能需要引入结构先验、正则化，或者使用无监督预训练。

      之前Google也发表过关于Transformer的综述性论文《Efficient Transformers: A Survey》，但是此论文主要关注Attention效率的问题；本次主要分享一下本月复旦大学邱锡鹏教授团队的佳作《A Survey of Transformer》

一、Vanilla Transformer

       在介绍Transformer的变种之前，再回顾一下经典的Vanilla Transformer，结构如下图所示：

        Vanilla Transformer是一个seq2seq结构，包括encoder和decoder结构，每个结构都是有L个子block stack的结果，其中Encoder包括多头注意力和FFN这两个module，每个module都是通过残差连接，并且通过LN进行归一化；而Decoder也有多头注意力和FFN，并且在其中插入了cross-Attention来建立Encoder和Decoder的联系，多头注意力是自回归，因此是下三角模型。下面对Vanilla Transformer的主要组件进行简要介绍：

1.1、Attention Modules

self-Attention的公式如下：

是注意力矩阵，softmax是row-wise的；

是为了缓解softmax梯度消失的。

在Transformer中一般都是多头的，其实就是计算完每个头之后，再进行concatenate起来，这时一般每个头的维度会是向量维度除以头数。

在Transformer中有三种注意力机制，分别是：

Self-Attention：在encoder中，Q=K=V=X，X表示前一层的输出；

Masked self-Attention：在decoder中，每个位置的Q只能注意到它前面的K、V的值和位置，因此需要在原来的注意力矩阵的基础上添加一个mask矩阵，mask矩阵在mask的位置都设置为-∞；

Cross-Attention：是encoder的K和V与decoder的Q进行计算Attention。

 

1.2、Position-wise FFN

FFN其实就是经过一个ReLU的激活函数后再经过一个普通的全连接，公式如下：

1.3、Residual Connection and Normalization

1.4、Position Encodings

       由于Transformer不像RNN或者CNN，它本身没有位置信息，然而通常模型是需要考虑tokens的顺序问题，比如NLP或者CV的像素，因此需要额外补充位置信息。

Model Usage

Encoder-Decoder：可以理解为seq2seq模型的使用；

Encoder only：把一个输入句子进行表示学习，可用于分类或者序列标注问题

Decoder only：需要把Decoder的cross-Attention移除，比如GPTs

Model Analysis

T表示输入序列的长度，D表示隐藏层大小

从Tabel1看出，当输入序列长度较短时，模型的瓶颈主要在FFN；当输入序列长度较长时，模型的瓶颈主要在self-Attention。

二、Taxonomy of Transformer

        目前Transformer的变种大致分为如下四类：Module-level、Arch-level、PTMs和Applications，如下图Fig2所示：

也可以进一步按照代表模型来进行细化分类，如下图Fig3所示：

下面逐一介绍这些变种的原理以及思路：

三、Attention

       注意力机制在Transformer中起着举足轻重的作用，然而在实际应用中，发现有两个缺点：

1）处理长文本的复杂度高；

2）模型对输入数据没有结构偏差，对于不是预训练好的Transformer，在小数据集上容易产生过拟合

注意力的优化可以从以下六个方面去考虑：

3.1、Sparse Attention

       标准的self-Attention是每个token都与所有的tokens进行连接，然而从已经训练好的Transformer观察注意力矩阵A，发现大部分的点是比较稀疏的，因此可以引入结构偏差来限制每个query所注意的query-key pairs，形式如下：

通常−∞不会被存储到内存中，这样可以减少内存的开销。

       从稀疏连接的实现方式不同，可以分为position-based sparse attention和content-based sparse attention。

3.1.1、position-based sparse attention

       position-based sparse attention是通过设置一些patterns对注意力矩阵进行限制从而实现稀疏化，然而这些pattern可以分解成一些原子级的稀疏pattern，这些原子级的pattern又可以组成新的pattern，下面进行详细介绍：

3.1.1.1、Atomic Sparse Attention

原子级的稀疏pattern大致分为如下图的五大类：

1）Global Attention

这种方式设置少数几个node作为全局node与所有node进行连接，其他非全局node之间不连接，这种方式可以缓解稀疏注意力在长依赖问题导致的建模能力退化问题，如图Fig4（a）所示；

2）Band Attention

这种方式是通过一个窗口来控制注意力的范围，每个node只与窗口范围内的其他node计算注意力，如图Fig4（b）所示；

3）Dilated Attention

这种方式类似Dilated CNNs；和Band Attention相比，每个窗口内的node是间隔几个node进行注意力计算的，如图Fig4（c）所示；

4）Random Attention

这种方式比较好理解，每个node都是随机选择和其他node进行注意力计算，如图Fig4（d）所示；

5）Block Local Attention

这种方式是把注意力矩阵划分为多个block，然后每个block内部进行注意力计算，如图Fig4（e）所示；

3.1.1.2、Compound Sparse Attention

       组合稀疏注意力，顾名思义就是把多个原子级的稀疏注意力组合起来，常见的有如下四种方式：

3.1.2、content-based sparse attention

       content-based sparse attention是query和最大化内积的keys来进行计算，从而实现稀疏化。

1）Routing Transformer：对query和keys进行k-means聚类，然后每个query只和同一个类内的keys进行注意力计算，训练的时候，类中心的向量使用指数移动平均方式更新的，公式如下：

 

2）Reformer：使用locality-sensitive hashing（LSH）对query和key进行hash计算，然后每个bucket里面的query和key进行注意力计算，hash公式如下：

3.2、Linearized Attention

这种方式主要是对标准的self-Attention中的QK计算进行拆分，如图Fig7所示：

标准的注意力计算如下：

如果使用核函数来计算q和k相似度的话，公式可以写成

       这样的话，S和u就可以向RNN一样进行计算了，这样也就说明S和u有记忆功能，注意力计算就抽象为 (1)feature map 𝜙 (·)和 (2) aggregation rule这关键的两部分了。

3.2.1、feature map

Linear Transformer使用𝜙𝑖 (x) = elu(𝑥𝑖 )+1作为feature map；

 

Performer采用随机feature map，公式如下：

Performer第一个版本：

      尽管三角函数随机feature map是无偏估计，但是它不能保证注意力得分是非负的，导致模型不稳定。

Performer第二个版本：

采用如下公式来改善第一个版本的缺点

3.2.2、aggregation rule

      标准的聚合规则𝜙 (k)𝑗 ⊗ v𝑗仅仅是通过求和聚合到记忆矩阵S的，然而更好的方式可能是在加入新关联的时候记忆矩阵可以选择性的删除的删除关联。

RFA提出了一种门控机制来对历史关联进行指数衰减，从而优先考虑当前内容。

3.3、Prototype and Memory Compression

       此方法主要是通过减少计算key-value pairs来减少注意力的计算复杂度的，其中有两个方法可以实现：query prototyping和memory compression

3.3.1、Attention with Prototype Queries

       这种方法主要是选择几个query作为原型来计算注意力分布，模型要么把分布复制到这几个代表的query对应的位置，要么使用离散均匀分布进行填充，Fig8（a）介绍了实现流程。

       Informer使用query注意力分布和离散均匀分布的最近似的KL散度来从query中选择原型，只有这几个query参与注意力的计算，其他的query都是离散具有分布。

3.3.2、Attention with Compressed Key-Value Memory

      该方法是把key和value的数量压缩了，代表模型有SetTransformer、Linformer和Poolingformer

3.4、Attention with Prior

        一般来说，Vanilla Transformer的Attention是由输入数据来生成的，然而其实可以加入一些先验的Attention来提升Attention的能力，可以把这两者的注意力进行混合，如下图Fig9所示：

根据先验Attention的方式不同，有如下几种形式：

3.4.1、Prior that Models locality

该方法引入高斯分布来对每个词的中心位置进行预测。

Yang提出了一种前馈神经网络来实现

Pi表示第i个query的位置

Gaussian Transformer提出了另一种方法：

𝑤 ≥ 0, 𝑏 ≤ 0

 

3.4.2、Prior from Lower Modules

       通过分析Transformer的架构发现，Attention的分布与相邻层的分布是比较相似的，那么其实可以把前一层的Attention作为先验，公式如下：

Predictive Attention Transformer提出对前一层Attention使用二维卷积进行计算，g(.)是卷积层，𝑤1 = 𝛼, 𝑤2 = 1 − 𝛼

Realformer是直接把前一层的Attention加入到后一层中，𝑤1 = 1, 𝑤2 = 1

Lazyformer提出在几个相邻的Attention层之间进行共享注意力，𝑤1 = 0, 𝑤2 = 1或者𝑤1 = 1, 𝑤2 = 0，这样可以减少计算量。

3.5、Improved Multi-Head Mechanism

       多头注意力机制是从不同的子空间表示来获取不同的信息，然而到目前为止没有一种理论说明不同的头可以获取不同的特征。

3.5.1、Head Behavior Modeling

       Li等人提出在损失函数中引入了一个辅助分歧正则化项，以鼓励不同注意头之间的多样性。两个正则化项分别是使输入子空间和输出表示的余弦距离最大化，而最后一个正则化项是用相应注意矩阵的元素相乘来分散多个头部的位置。

3.5.2、Multi-head with Restricted Spans

标准的Vanilla Attention 的mask值是恒等于1，如图Fig10（a）；

Sukhbaatar提出可以使用可以学习的标量z和超参数R来表示mask，如图Fig10（b）；

Multi-Scale Transformer提出在一个固定的窗口大小来注意力计算，这个窗口在高layer是相对大一些，在低layer是相对小一些。

3.5.3、Multi-head with Refined Aggregation

        这种方法的产生认为原始的Vanilla Transformer的多头注意力简单求和的方式不能充分挖掘多头注意力的信息，因此Gu和Feng以及Li提出了使用路由的方式来实现（最早有胶囊网络提出），注意力头的输出首先被转换为输入胶囊，通过路由迭代生成输出胶囊，多个输出胶囊concatenate起来作为多头的输出。

四、Position Representations

       由于Transformer中的Attention没有位置信息，然而对于尤其像文本的数据，在建模的时候是需要考虑tokens前后关系的，因此需要额外的增加位置编码信息作为模型的辅助输入。

4.1、Absolute Position Representations

        Vanilla Transformer使用手工规则的sin/cos来对位置进行编码，编码信息添加到token embedding中；还有一种方式是让模型去学习这个位置编码，这种方式比较灵活，但是学习出来的embedding是训练预料中最大的句子长度。Wang提出了使用数据中的词频𝜔𝑖来作为位置信息的参数，公式如下：

4.2、Relative Position Representations

这种方法考虑的不是单个token的位置编码，而是tokens之间的相对位置，Shaw提出了如下方式：

r是相对位置，K表示embedding的最大偏置。

Transformer-XL使用如下计算方式来得到相对位置：

R是Vanilla Transformer的位置编码方式，W，u是学习的参数。

DeBERTa使用Show的方式，风格类似Transformer-XL

五、Layer Normalization

5.1、Placement of Layer Normalization

       在Vanilla Transformer中，LN layer使用的是post-LN，如图Fig11（a），

Xiong从理论上说明在post-LN的Transformer中，模型训练开始的时候越接近输出层的梯度越大，这就说明在post-LN Transformer中不使用warm-up会导致训练不稳定，然而pre-LN不会出现类似的问题。尽管如此，但是实际使用中发现post-LN比pre-LN的性能好，进一步分析后发现训练不稳定是由于残差连接导致的，通过引入辅助参数来解决残差依赖问题。

5.2、Substitutes of Layer Normalization

       Xu通过大量实验观察LN模型的可学习参数发现其实这些参数是不起作用的，甚至会导致过拟合风险。他们进一步发现LN可以起作用的不是forward normalization，而是派生出来的均值、方差以及对梯度的归一化来实现的，因此提出不包含可训练参数的AdaNorm，公式如下：

C，k是超参数，⊙是element-wise相乘，𝜇 和𝜎是输入样本的均值和方差。

Nguyen和Salazar使用L2正则化来代替LN模块，使用公式如下：

g是一个可训练的标量，在机器翻译数据上表现比LN好，尤其是在低资源的情况下。

Shen分析了BN在Transformer中处理文本数据表现差的原因，从而提出了PowerNorm（PN），对BN做了如下修改：

1）不强制0均值正则化；

2）使用平方均值来代替方差；

3）使用平方均值来替换batch的统计

公式如下：

0 < 𝛼 < 1，是移动平均系数；𝛾, 𝛽是BN的训练参数。

5.3、Normalization-free Transformer

        除了LN，ReZero通过构建可学习的残差来替换LN模型也可以构建很深的网络，在Transformer中应用ReZero收敛的速度更快。

六、Position-wise FFN

FFN看起来比较简单，然而是不可或缺的，Dong尝试去掉FNN，只使用self-Attention堆叠会引起秩坍塌问题，从而导致归纳偏差，而FFN会解决这个问题。

6.1、Activation Function in FFN

       Vanilla Transformer使用的是ReLU激活函数，其实也有一些其他的激活函数可以替换，比如GPT使用的GELU，Gated Linear Units（GLU）等等

6.2、Adapting FFN for Larger Capacity

       这个基本想法是通过更多参数的网络来替换FFN来扩充FFN层的能力。

Lample使用product-key memory层来替换FFN，这个网络包括：一个query网络，包括两个sub-keys的key选择模型和一个value lookup table。模型首先使用query网络把输入映射到隐空间，然后将生成的查询与key选择模块中两组sub-key的笛卡尔乘积的key进行比较，得到𝑘 近邻，最后使用k近邻的key从value lookup table中查询对应的值，并且聚合后作为输出。整个过程集成了注意力机制，生成的query与所有的key-value pair进行注意力计算。

6.3、Dropping FFN Layers

在特殊场景下，FFN是可以删除的，从而简化模型。