Efficient Transformers: A Survey

作者: cornbig | 来源:发表于2020-12-15 17:19 被阅读0次

论文

[2009.06732] Efficient Transformers: A Survey

1.Background on Transformers

Dense Attention： encoder和decoder最大的区别在于掩码不同，encoder中的token可以看见所有的，而decoder中的token只能看见它前面的，他后面的token被mask掩掉了，需要预测出来。

Transformers是将 Transformer block彼此堆叠形成的多层体系结构，Transformer block（基本单元块）的特点是多头自注意力机制（multi-head self-attention mechanism）、位置前馈网络（position-wise feed-forward network）、层归一化（layer normalization）和残差连接（residual connectors.）

Transfomers的输入是tensor，shape是 $R^B\times R^N$ ，B is batch size，N is sequence length。

（1）输入首先通过embedding layer，embedding layer将每个one-hot token 表示为d维embedding，例， $R^B \times R^N\times R^D$ ，之后新的tensor与positional encodings相加并通过多头自注意模块。Positional encodings可以采用正弦输入的形式或者可训练的embedding。

（2）multi-head self-attention module的输入和输出residual connectors和一层 normalization layer连接。multi-head self-attention module的输出传递到两层前馈网络，输入和输出类似地以残差方式通过层归一化连接，表示为。

$X = LayerNorm（F_{s}（X））+ X$

$F_{s}$ 是sub-layer module，（the multi-headed self-attention或者position-wise 前馈层）

1.1.Multi-Head Self-Attention（多头注意力机制）

Transformer模型利用多头注意力机制，其背后的机制是学习alignment（对比），序列中的每个element学习从其他tokens中收集信息。a single head 定义为：

$A_{h} = Softmax（\alpha Q_{h} K_{h}^T ） V_{h}$ ，

$Q_{h} = W_{q} X$ ， $K_{h} = W_{k}X$ ， $V_{h} = W_{v}X$ 是应用于输入序列的时间维度的线性变换。

$W_{q}$ ， $W_{k}$ ， $W_{v}$ $\in$ $R^{d\times \frac{d}{N_{h} } }$ 是query的权重矩阵（参数），key和values投影并且将输入X映射到x维的输出tensor， $N_{h}$ 是heads的个数。 $X \in R^N \times R^d$ ； $\alpha$ 是比例因子通常设置为 $\frac{1}{\sqrt{d} }$ 。 $N_{H}$ 表示heads的个数，heads的输出 $A_{1} ... A_{N_{H} }$ concatenated together 并且用过一个dense层。输出Y表示为：

$Y = W_{o}[A_{1} ...A_{N_{H} }]$ .这里 $W_{o}$ 是输出层映射。注意all heads的线性变换并行计算。

$A\in R^B \times R^N \times R^{N_{h} } \times R^{\frac{d}{N_{h} } }$ .

attention matrix $A = QK^T$ 主要学习序列中tokens之间的对比分数（alignment scores），在这个公式中，获取query（Q）中每个元素和key（K）之间的点积。这会推动self-attention中的self-alignment过程，从而使tokens彼此学习。

On the scalability of Self-Attention（自我注意的可扩展性）

很明显输入sequence 长度为N，the attention matrix A 内存和计算的复杂性是 N X N。特别的，仅 $Q K^T$ 矩阵的乘法运算就使用 $N^2$ 的时间和内存。这限制了self-attention在处理长序列的效果。

1.2 Position-wise Feed-forward Layers（位置前馈层）

self-attention的输出传递给具有ReLU激活的两层前馈网络，该前馈网络独立的在每个位置运行，因此称为“position-wise”，表达式为：

$F_{2}（ReLU（F_{1}（X_{A} ）））$ ，F1、F2 为形式是Wx+b的前馈函数。

1.3 Putting it all together

每个Transformer Block 表示为：

$X_{A}$ = $LayerNorm(MultiheadSelfAttention(X))$ + $X$

$X_{B}$ = $LayerNorm(PositionFFN(X_{A} ))$ + $X_{A}$

X 是Transformer Block 的输入， $X_{B}$ 是Transformer Block 的输出。不管哪一种类型的结构都是由多个transformer基本单元组成，基本单元包括多头注意力机制、前馈网络、层归一化和残差连接。多头注意力机制主要是将输入的embedding映射到不同的低维度空间上，通过内积计算不同token之间的相关性得分，再将相关性得分乘以对应的token embedding求和得到某一个token的embedding，最后把同一token在不同低维度空间上embedding拼接起来得到最终的向量表征。在这里，每个token都与其他的token进行一次attention计算，如果有序列长度为n，那么时间复杂度为o( $n^2$ ),这种attention方式称之为稠密注意力机制（dense attention mechanism）。

1.4 Transformer Mode

Transformer block的使用差异，主要有三种使用方式（1）encoder-only（分类），（2）decoder-only（语言模型）（3）encoder - decoder（机器翻译）。在encoder - decoder模式下，通常有多个multi-headed self-attention 模块。包括encoder-decoder中的标准self-attention。包括encoder-decoder cross attention允许decoder

3. A Survey of Efficient Transformer Models

3.1 Efficient Transformers 的分类

除了基于段的递归模型，这些模型的主要目标是近似quadratic-cost 注意力矩阵，每种方法都将稀疏性的概念应用稠密的注意力机制。

Table 1:FP（ Fixed Patterns ）= 固定模式 or 固定模式的组合；M = Memory； LP = Learnable Pattern (科学习模式)； LR = Low Rank ； KR = Kernel and RC = Recurrence ；通常n表示sequence的长度，b是local window （or block） size ，

n_{g}

表示global memory长度，

n_{c}

表示卷积压缩的长度。

（1）Fixed Patterns (FP)

通过限制attention的范围，从全局到局部，降低计算复杂度，包括下述3种方法：

Blockwise Patterns ，将输入序列切成多个block，attention只发生在block范围内，所以复杂度从 $O(N^2 )$ 降低到 $O(B^2 )$ ,B 是Block size，这种切割会导致序列不连贯，attention能力受限，效果应该不太行。

Strided Patterns ，采用滑动窗口的方法，每个token与相邻几个token做attention，相邻的token范围就是window size，这样是有一定道理的，因为自然语言在多数情况下都是局部相关的，所以在一个窗口范围内作attention往往不会丢失太多信息，相比之前的blockwise pattern应该好一些。

Compressed Patterns ，则通过卷积池化的方式对序列进行降采样，比如用kernel 为2，stride为2的CNN，将2个token表征成一个向量，然后在做attention，同样也能降低attention的计算复杂度。其实就相当于通过CNN对序列进行n-gram的切分。

（3）Learnable patterns(LP)

learnable patterns是对上文提到的fixed patterns的扩展，简单来说fixed pattern是认为规定好一些区域，让该区域的token进行注意力计算，而learnable patterns则是通过引入可学习参数，让模型自己找到划分区域。比如reformer引入基于哈希的相似度度量方法将输入序列切割；routing transformer对token向量进行k-means聚类来将整体序列分割成多个子序列。所以本质上说LP与FP是一致的，都是通过将整体序列切分成子序列，attention只在子序列中进行，从而降低计算开销，只不过LP的区域划分是通过模型学得，而FP则是人为定义。

（4）Memory

memory乍一听好像有点让人摸不着头脑，其实想法也很简单。最开始是在19年的set transformer上使用。一般来说做multihead self-attention时，Q=K=V=X（X为输入序列，长度为n），而在set transformer中，作者先单独设置了m个向量（m是超参数），然后这m个向量与X做multihead attention，得到m个临时向量（这些个临时变量我们就称作“temporary memory”），接着把X与这m个临时向量再做一次multihead attention得到输出。这个过程其实就是利用这m个向量将输入序列X的信息先通过attention进行压缩，再通过attention还原，达到抽取输入序列的特征的目的。但是在压缩编码解码的过程中肯定会有信息损失，所以后来的改进方法是引入全局记忆，即设置一些token，它们可以与所有的token进行注意力交互，比如bert的[CLS]，由于这些token的数目远小于序列长度，因此也不会给计算带来负担，而且往往携带了整个输入序列的信息，这就是我们可以在bert上用[CLS]作文本分类任务的原因。

（4）Low rank methods（低秩方法）

使用self-attention matrix的leveraging low-rank 近似，方法的关键是采用 $N\times N$ 矩阵的low-rank 结构，Linformer，简单来说，经过softmax的 $N \times N$ 矩阵不是满秩的，这意味着我们不需要计算一个完整的注意力矩阵，因此我们可以将 $n \times d$ 维（n表示序列长度，d表示模型向量维度）的K，V映射到 $k \times d$ 维空间。

E、F是projection layer，注意力矩阵变为

n \times k

， k是超参数，因此复杂度为

o(n)

(5) Kernels

以核函数变换的新形式取代原有的softmax注意力矩阵计算，将计算复杂度降至 $o(n)$ 范围内，比较代表的有Linear Transformers，下文有介绍。

(6) Recurrence

recurrence实际上也是上文提到的fixed patterns中blockwise的一种延伸。本质上仍是对输入序列进行区域划分，不过它进一步的对划分后的block做了一层循环连接，通过这样的层级关系就可以把一个长序列的输入更好的表征。以recurrence为代表的就是Transformer-XL。