《Transformer Quality in Linear Time》论文解读

会议/期刊: ICML

年份: 2022

1. Vanilla Transformer Block（MHSA FFN）

原本的Transformer的Block遵循如下的设计范式：MHSA（多头自注意力）一层或者两层的FFN（全连接层），如下图所示。我们只考虑FFN的话，其数学表达式如下：T表示句子长度，d表示词向量维度（也表示模型隐藏层维度），e表示expanded intermediate 特征大小。

{cal O}=phi(X W_{u})W_{o}mathrm{~Where~},, Xinmathbb{R}^{Ttimes d}, W _{u}inmathbb{R}^{dtimes e},W_{o}inmathbb{R}^{etimes d}

2. 改进Transformer Block （MHSA GLU）

后面有工作对FFN做了改进，提出了GLU（Gated Linear Unit）结构，并且发现能有效提升模型性能。GLU结构大致如下图。简单理解就是有两个支路，两条支路都是全连接层加激活函数。两条支路的激活函数可以不同。最后两路的结果会做element-wise相乘，得到的结果会再经过一个全连接层进行处理。

上图左边的GLU结构的数学表达式如下：

begin{array}{l l}{{U=phi_{u}(X W_{u}),}}&{{V=phi_{v}(X W_{v})}}\ {{{O}=left(U odot Vright)W_{o}}}tag{1}end{array}

其中$U,Vinmathbb{R}^{Ttimes e}，Oinmathbb{R}^{Ttimes d}$

3. GAU（Gated Attention Unit）

上面的GLU和注意力模块是独立开的，GAU做了一个很巧的构思把二者融合到了一个模块，其结构和伪代码如下图所示

GAU的数学表达式如下：

{O}=(Uodothat{V})W_{o}quadmathrm{where}quadhat{V}=A V tag{2}

其中

begin{array}{l l}{{Z=phi_{z}(X W_{z})}}&{{qquadinmathbb{R}^{Ttimes s} }}\ {{A=operatorname{relu}^{2}left(mathcal{Q}(Z)mathcal{K}(Z)^{top} bright)}}&{{qquadinmathbb{R}^{Ttimes T} }}tag{3}end{array}

可以看到在计算注意力矩阵A用到的Q和K是基于共享的矩阵Z计算得到的，$mathcal{Q}(Z), mathcal{K}(Z)$都是对矩阵Z做per-dim的归一化，类似于LayerNorm。得到注意力A后，还要经过ReLU激活函数，然后取二次方，即$relu^2$，这个是在《Primer: Searching for Efficient Transformers for Language Modeling》论文中用NAS搜索出来的。

3.1 参数量比较

下面我们比较一下 MHSA MLP/GLU与 GAU 结构的参数量：

MHSA MLP/GLU
- MHSA: Q, K, V对应的映射模块权重均为hdd/h=dd,最后MHSA的Dense层的权重参数量也是dd，所以MHSA的参数量为4dd
- MLP: 通常是两个全连接层，每个的权重参数量为de,一般e=4d，所以MLP模块的权重参数量为 2 (de)=2 (d4d)=8d*d
- GLU: 如果采用GLU结构，那么权重参数量则为3de=12dd
- 总结：如果采用MHSA MLP，则参数量是12dd；如果采用MHSA GLU，则参数量是16dd
GAU参数量为3de ds。通常s会远远小于d,所以参数量近似为3de。改论文中，作者设置e=2d,那么GAU模块的参数量则为6d*d。换言之两个GAU级联后的参数量等价于MHSA MLP。

3.2 计算复杂度比较

对比GLU MHSA和GAU，我们可以看到GAU只有一个head，而且去掉了Softmax，而且实验结果显示GAU的表现和原来的MHSA MLP也不分伯仲，甚至更好

但是，仔细分析一下，我们会发现GAU的计算复杂度和原本的自注意力机制一样，仍旧是句子长度的二次方，即$O(T^2)$。

下面我们分析一下二次复杂度的来源，GAU和原始的自注意力机制的计算都可以用如下的数学公式表示：

A=phi(QK^T)V

在原始的自注意力机制中，激活函数$phi$是softmax，而在GAU中是$ReLU^2$。矩阵$Q, Kinmathbb{R}^{Ttimes d}$，二者矩阵乘法的复杂度为$O(Ttimes d times T)$,如果只考虑句子长度，我们可以将d忽视，所以复杂度为$O(T^2)$.

后续的一些尝试将复杂度降低至线性复杂度的方法的思路是这样的，

phi(QK^T)Vrightarrow(phi_q(Q)phi_k(K)^T)V=phi_q(Q)(phi_k(K)^TV)

简而言之就是尝试将矩阵$K^T$和$V$先做矩阵乘法，这样一来它们的复杂度则为$O(dtimes T times d)$，得到大小为$mathbb{R}^{dtimes d}$的矩阵，该矩阵再和$Q$相乘，计算复杂度同样是$O(dtimes T times d)$。

3.3 推理阶段的复杂度

我们接下来考虑推理时GAU的复杂度。

我们知道GAU会先算$M=K^TV$，然后再计算$QM$，所以我们先着重分析一下矩阵$M$的计算。

由于推理阶段采用的是自回归的解码方式，也就是说K和V的长度（即词数量）是从1逐渐增加到T的。考虑t时刻的情况，要得到矩阵$M_t$, 我们需要$O(dtd)$的计算复杂度，随着t逐渐从1增加到T，计算复杂度是不断增加的，换言之计算复杂度是$O(Td^2)$。

这里其实有一个计算上的技巧，即我们需要先存储上一次的结果$M{t-1}$。当到t时刻的时候，我们计算出新词的$K_t,V_tinmathbb{R}^{1times d}$向量，然后计算$K_t^TV_tinmathbb{R}^{dtimes d}$，最后将这个值和$M{t-1}$累加即可得到$M_t$，即

M_t=M_{t-1} K_t^TV_t tag{4}

简而言之，每个时刻（即有新的词输入的时候），只需要计算新词的$K_t^TV_t$即可，因此空间复杂度是$O(d^2)$，计算复杂度始终保持为$O(d^2)$，相比于原来的$O(Td^2)$计算复杂度有了明显改进。

上述这种计算技巧在推理阶段非常有效，可以很巧妙地降低计算复杂度。但是，在训练阶段就会有问题了，因为这个技巧是基于自回归的特点设计的，也就是说推理阶段就像RNN一样，每次只新增一个单词，无需考虑并行性。训练阶段输入的数据一般是大小为$btimes Ttimes d$的张量，如果想采用上面的计算技巧，那么训练阶段的输入就需要像推理阶段一样，显然这会得不偿失，因为这样无法并行计算了。

4. Mixed Chunk Attention

为了解决上面提到的推理计算技巧无法应用到训练阶段，本文作者提出了Mixed Chunk Attention方法，该方法将Partial Attention（简单理解就是只计算更重要部分的注意力，但是实际上这类方法的计算效率不高，因为计算是不规则和碎片化的）和Linear Attention的优点进行了结合。

上面图中每个圆圈代表一个单词的词向量，中间的正方形表示$Mt=M{t-1} K_t^TV_t$。

图（top）表示原始的注意力机制计算方法，每次计算注意力矩阵的复杂度是$O(T^2d)$。

图（middle）即表示通过公式(4)可以复用前一时刻的结果，将计算复杂度降低至$O(d^2)$，但是在这种类似RNN的计算方式缺乏并行性，很难在训练阶段使用

图（bottom）则做了这种，所以称作mixed chunk attention （MCA）。假设输入序列维度是$btimes Ttimes d$，后面为避免符号太多，我们省略batch size，即$b$。由图（bottom）可以看到，MCA其实就是将原来的一个句子划分成$G$个chunk，每个chunk包含$C$个单词（该论文取$C=256$），也就是说原来的句子长度$T=Gtimes C$。所以原本的输入序列$Ttimes drightarrow Gtimes Ctimes d$。原本的GAU模块转变成了如下图：

为方便理解，我们只考虑单个chunk，那么对于第$g$个chunk，则中间结果$U_ginmathbb{R}^{Ctimes e},V_ginmathbb{R}^{Ctimes e},Z_ginmathbb{R}^{Ctimes s},$其中Q，K矩阵是基于共享的$Z_g$采用不同的放射变化得到的，具体而言会有两套Q，K矩阵：

一套用于计算local Attention的复杂度为二次方的$Q_g^{quad},K_g^{quad}inmathbb{R}^{Ctimes s}$。如图5（bottom）最下面那一行圆圈所示，每两个圆圈之间会计算彼此之间的注意力矩阵，这其实可以理解成一种稀疏的注意力，其计算公式如下

hat{V}_g^{quad}=relu^2(Q_g^{quad}K_g^{quad} B)V_g tag{5}

单个chunk的local Attention的计算中的$Q_g^{quad}K_g^{quad}$计算复杂度为$O(C^2s)$，计算得到的结果与矩阵$V_g$相乘的计算复杂度为$O(C^2e)$，因为$s<<d$，另外$e$正比于$d$（例如$e=2d$），所以复杂度是$O(C^2s C^2e)=O(C^2s 2C^2d)=O(C^2d)$。另外由于总共有G个chunk，所以总的复杂度为$O(GC^2d)=O(TCd)$

另一套是用于计算global Attention的复杂度为线性的$Q_g^{lin},K_g^{lin}inmathbb{R}^{Ctimes s}$。我们其实可以将图5（bottom）最下面每两个圆圈视为一个圆圈，就像图5（middle）一样。此时计算global Attention可以分成两种情况：训练和推理，或者也可以称作Non-Causal和Causal。Causal表示因果，即下一个单词的预测依赖前面的输入，这就对应推理。两种情况的具体计算公式如下：
Non-Causal (训练):

hat{V}_g^{lin}=Q_g^{lin}(sum_{h=1}^G{K_h^{lin}}^TV_h) tag{6}

训练阶段其实可以不用像公式(6)那样分chunk的去计算，我们其实可以直接用完整的矩阵$Q^{lin},K^{lin}inmathbb{R}^{Ttimes s}$直接计算得到公式(6)右边的累加项。

我们再看看计算复杂度，${K_h^{lin}}^TV_h$的计算复杂度为$O(Cse)$，累加G个chunk，那么复杂度就是$O(GCse)=O(Tse)=O(Tsd)$。矩阵Q与KV计算的到矩阵相乘的复杂度为$O(Cse)=O(Csd)$。所以公式(6)的计算复杂度近似为$O(Tsd)$。

Causal (推理)：

hat{V}_g^{lin}=Q_g^{lin}(sum_{h=1}^{g-1}{K_h^{lin}}^TV_h) tag{7}

根据两套Q，K矩阵，我们可以分别求得$hat{V}_g^{quad},hat{V}_g^{lin}$，最后我们将二者相加得到混合注意力，最终第$g$个chunk的输出计算公式如下

O_g=[U_godot(hat{V}_g^{quad} hat{V}_g^{lin})]W_o tag{8}

Mixed Chunk Attention伪代码如下：

5. 论文中的一些讨论

5.1 Chunk是否需要overlap

前面提到将输入序列划分成多干个chunk，这些chunk彼此之间是没有overlap的。比如说这句话“今天我吃了好多好吃的，有龙虾、鲍鱼、海参和饺子。”，以non-overlap的划分方式将这个句子（总共24个字符）划分成三个chunk，则得到

今天我吃了好多好
吃的，有龙虾、鲍
鱼、海参和饺子。

那么，一个很自然的问题是如果overlap会怎么样呢？结果是否会更好？作者对这个做了测试，实验结果表明overlap的划分chunk的方式的确能够提升模型性能，但是引入了额外的计算成本。与其使用overlap 的chunk划分方式，还不如直接多加几层non-overlapping GAU模块。

5.2 局部和全局注意力的Ablation Study

原论文还做了消融实验，显示相对来说局部注意力比全局注意力更重要，而混合式的效果最好。下面实验中的MC-TFM 是指将Mixed Chunk Attention运用到Transformer 。MC-TFM 和FLASH一样都是线性复杂度，但是用的FFN。可以看到使用GAU的FLASH要明显优于MC-TFM 。

5.3 Chunk大小该如何选择

当C和句子长度一样时，此时等价于FLASH-Quad，即计算复杂度为二次方
当C=1时，则等价于Linear Attention，但是在做auto-regressive training的时候不够高效，缺少并行性
下图给出了在不同句子长度下，C取不同值「128, 256, 512, 1024」的效果，彼此差距不是很大，最终作者选取了256

huggingface-transformers

1 人点赞