线性变换器

作者: Valar_Morghulis | 来源:发表于2020-12-16 16:51 被阅读0次

变换器在一些任务中效果很好，但是它的时间复杂度是二次的，因而有个局限性，就是输入的序列很长的时候，它的计算会非常慢。为了解决这个局限性，我们将自注意力表示为核特征图的线性点积，并利用矩阵乘积的结合性，将复杂度从 $O(N^2 )$ 降低至 $O(N )$ ，其中 $N$ 是序列长度。我们证明，这个公式可以迭代地实现，这样就可以极大地加速自回归变换器，我们还揭示了其余递归网络的关系。在序列非常长的自回归预测任务上，我们的线性变换器与普通的变换器性能相似，速度速度提高了4000倍。

https://fast-transformers.github.io/

https://linear-transformers.com/

https://github.com/idiap/fast-transformers

还有另外两篇，和这篇有关系：

Fast Transformers with Clustered Attention (2020.9) https://arxiv.org/abs/2007.04825 (和本文一作二作换个位置，需要关注，对语义分割可能非常适用，开源又是放在一起的)，这个是在NIPS2020里的，可关注两篇论文的团队leader的网址：https://fleuret.org/francois/

Efficient Attention: Attention with Linear Complexities v1:2018， v9:To appear at WACV 2021 https://arxiv.org/abs/1812.01243v9

隔壁家的Performer的实现：https://github.com/lucidrains/performer-pytorch

上面这个人好可怕，实现了好多：

https://github.com/lucidrains/linear-attention-transformer

https://github.com/lucidrains?tab=repositories

1. 引言

变换器最早由Vaswani等(2017)提出，是用在神经网络机器翻译(Sutskever et al.，2014；Bahdanau et al.，2015)中，并且在自然语言处理(Devlin et al.，2019)、音频(Sperber et al.，2018)和图像(Parmar et al.，2019)领域中的各种任务中展示了令人印象深刻的结果。除了在充分监督的任务上，当使用自回归(Radford et al.，2018；2019)或掩盖语言建模目标(Devlin et al.，2019；Yang et al.，2019；Song et al.，2019；Liu et al.，2020)的预训练时，变换器也能有效地将知识迁移到有限监督或无监督的任务上。

然而，其虽然性能好，但计算量和内存占用大。这一问题主要是自注意力的全局感受野造成的，其在处理 $N$ 个上下文输入时，内存和时间复杂度是二次的 $O(N^2 )$ 。因此在实际应用中，它的训练很慢，因而会限制它的上下文长度，但是这样又会破坏时序的一致性，并有损于长期依赖关系的捕获。Dai等人(2019)(Transformer-XL)通过关注先前上下文中的记忆来解决后一个问题，但这样计算效率低。

最近的研究有关注在不降低速度的情况下增加上下文长度。Child(2019)(Sparse Transformer)提出注意力矩阵的稀疏因子分解，将自注意力的复杂度降低至 $O(N\sqrt{N} )$ 。Kitaev等人(2020)(Reformer)使用局部敏感的哈希方法进一步将复杂度降低至 $O(N\log_{}N )$ ，这样就可以处理长序列。尽管这两个方法可以在长序列上快速训练，但是在自回归任务的推断上没有带来速度的提升。

本文提出线性变换器，可显著减小内存占用，而且其复杂度和上下文长度是线性关系。我们是通过基于核的自注意力公式，并利用矩阵乘积的结合性来计算自注意力权重（第3.2节）。我们还展示了，在使用线性变换器的基础上，复杂度为线性、内存占用为恒定值的因果掩膜（3.3节）。这揭示了我们方法和RNN的关系，这使得我们可以更快地执行巨量级的自回归推理（3.4节）。

我们在图像生成和语音识别上的实验表明，线性变换器可以达到传统变换器的性能，而推理速度可以提升3个数量级。

2. 相关工作

本节将概述与本文最相关的一些工作，这些先前工作旨在解决变换器的内存占用大、计算量大的问题。此外，我们还从理论上分析变换器的核心部分，即自注意力机制。最后，我们还提出另一个工作方向，即试图缓解softmax在自注意力中带来的计算量大的问题。

2.1 快速的变换器

现有的一些工作通过剪枝(Michel et al.，2019)、权值因子分解(Lan et al.，2020)、模型量化(Zafrir et al.，2019)、知识蒸馏来减小变换器的内存占用。Clark等人(2020)提出一种新的预训练目标，称为 replaced token detection，能使用更少的样本来训练，并减小了整体的计算量。Lample(2019)使用product-key注意力来提高任意层的容量，而带来的额外的计算代价可忽略不计。

使用这些方法可以减小内存占用、减小计算量、使训练得更快、推理得更快，但从根本上讲这些方法仍然是二次的，因而难以应用到长序列上。相比之下，我们证明我们的方法在理论上（3.2节）和实证上（4.1节）都降低了变换器的内存占用和时间复杂度。

另一个研究方向是旨在增加变换器中自注意力的上下文。上下文是指序列中用来计算自注意力的最远部分。Dai(2019)提出TransformerXL，通过学习固定长度以外的上下文依赖关系而不破坏时序一致性，达到语言建模当时的最佳性能。然而，在内存中保留先前的上下文会带来额外的计算代价。相比之下，Sukhbaatar(2019)通过学习每个注意力头的最佳注意力跨距，同时对历史内存和计算时间进行控制，显著增加了上下文长度。请注意，这两种方法的渐进时间复杂度和原初的变换器是相同的。相比之下，我们减小了渐进时间复杂度，这使得我们可以使用更远程的上下文。

和我们的工作更相关的是Child(2019)和Kitaev(2020)的工作，前者(SparseTransformer)提出注意力矩阵的稀疏因子分解，在长序列的生成任务上，将整体复杂度从二次降低到 $O(N\sqrt{N} )$ 。更晚近的，Kitaev(2020)提出Reformer，使用局部敏感散列(LSH)，这样可以计算更少的点积，从而将复杂度降低至 $O(N\log_{}N )$ 。请注意，为了能够使用LSH，Reformer增加了约束，也就是keys和queries要是相同的，因此，该方法不适用于解码任务，在解码任务中keys和queries是不同的。相比之下，我们的方法对keys和queries没有任何约束，而且其时间复杂度和序列长度是线性关系。此外，我们的方法在自回归任务上的推理速度是提高3个数量级，在验证困惑度上取得相当的性能。

2.2 理解自注意力

很少有人从理论分析上更好地理解自注意力机制。Tscai(2019)等人对变换器中注意力机制提出一个基于核的公式，该公式将注意力视为对输入做平滑的核函数，核得分为输入之间的相似程度。该公式是有助于更好地理解注意力机制和位置编码模块。相比之下，我们使用基于核的公式来加速注意力的计算过程并降低它的时间复杂度。另外，我们还发现，在keys和queries上使用相似度为正的核函数，线性注意力通常会收敛。

最近，Cordonnier等人(2020)从理论上和实验上证明了，多头注意力的头够多的话，可以表达任何卷积层。这里，我们证明了，用在自回归任务上的自注意力层可以看做是一个递归网络，而且还可以显著加快在自回归任务上变换器的推理速度。

2.3 将softmax线性化

多年来，softmax一直是类别很多的多分类任务的计算瓶颈(Goodman，2001；Morin&Bengio，2005；Mnih & Hinton，2009)。最近的研究(Blanc&Rendle，2017；Rawat et al.，2019)通过用特征图的线性点积来近似softmax，以及采样，来加快训练速度。受这些工作启发，我们将变换器中的注意力中的softmax线性化。在本文工作进展的同时，Shen(2020)探索了线性化注意力在图像中目标检测任务中的应用。相比之下，我们不仅将注意力计算线性化，而且建立了一个在训练阶段和推理阶段，复杂度为线性、内存占用为恒定值的自回归变换器模型。此外，我们还证明，通过核的映射(the lens of kernels)，每个变换器可看做是个递归网络。

3. 线性变换器

本节将我们提出的线性变换器公式化。我们展示了，通过将注意力从传统的softmax注意力转换为基于特征图点积的注意力，可以减小时间复杂度和内存占用，以及可以构建在耗时与序列长度呈线性关系的序列生成的因果模型，类似于递归网络。

首先，在3.1节中，我们会介绍Vaswani等(2017)提出的原初版的变换器的公式。然后，在3.2节和3.3节我们提出线性变换器。最后，在3.4节，我们将变换器改写为递归网络的形式。

3.1 原初版的变换器

令 $x\in R^{N\times F}$ 表示 $N$ 个维度为 $F$ 的特征向量。变换器是从 $R^{N\times F}\rightarrow R^{N\times F}$ 映射的函数，函数的计算为 $L$ 个变换器层的组成：

$T_{l} (x)=f_{l} (A_{l} (x)+x)$ $(1)$

其中 $f_{l} (\cdot )$ 单独地变换每个特征，通常都是简单地由两层前馈神经网络构成。 $A_{l} (\cdot )$ 是自注意力函数，是变换器中唯一执行序列间交叉作用的部分。

对于每个位置，自注意力函数计算所有其他位置的特征表示的加权平均值，权重与表示之间的相似性分数成比例。公式上，输入序列 $x$ 被3个矩阵权重投影： $W_{Q} \in R^{F\times D}$ 、 $W_{K} \in R^{F\times D}$ 、 $W_{V} \in R^{F\times D}$ ，分别对应 $Q$ 、 $K$ 、 $V$ 。所有位置的最终输出， $A_{l} (x )=V^{’}$ ，计算为：

$Q=xW_{Q}$ ， $K=xW_{K}$ ， $V=xW_{V}$ ，

$A_{l} (x )=V^{’}=softmax(\frac{QK ^ T }{\sqrt{D} } )V$ $(2)$

注意，在上述公式中，softmax函数按行用于 $QK ^ T$ 。通常地， $Q$ 、 $K$ 和 $V$ 被称为queries、keys和values。

公式(2)实现的是自注意力的一种特例，被称为softmax注意力，其中相似度分数是一个query和一个key的点积的指数。

令一个矩阵带一个下标 $i$ 表示该矩阵的第 $i$ 行的向量，我们可以写一个广义的注意力函数，适用于任意的相似度计算方式：

$V_{i}^{’} = \frac{\sum\nolimits_{j=1}^N sim(Q_{i}, K _{j})V_{j} }{\sum\nolimits_{j=1}^N sim(Q_{i}, K _{j})}$ $(3)$

如果我们令相似度函数为 $sim(q,k)=exp(\frac{q^T k}{ \sqrt{D} } )$ ，公式(3)和公式(2)就是等效的。

3.2 线性化的注意力

公式(2)中的注意力的定义是通用的，可以用来定义其它的注意力的实习，例如多项式注意力或者RBF核注意力(Tsai,2019)。要使公式(3)来定义一个注意力函数，唯一的约束是它的计算是非负的。这包括所有的核函数 $k(x,y):R^{2\times F}\rightarrow R_{+}$ 。

给定一个特征表示的核函数 $\phi (x)$ ，我们可以将公式(2)重写为：

$V_{i}^{’} = \frac{\sum\nolimits_{j=1}^N \phi( Q_{i}) ^T \phi( K_{j}) V_{j} }{\sum\nolimits_{j=1}^N \phi( Q_{i}) ^T \phi( K_{j} )}$ $(4)$

然后利用矩阵相乘的结合性，将上述公式进一步简化为：

$V_{i}^{’} = \frac{\phi( Q_{i}) ^T\sum\nolimits_{j=1}^N \phi( K_{j}) V_{j} ^T }{\phi( Q_{i}) ^T\sum\nolimits_{j=1}^N \phi( K_{j} )}$ $(5)$

我们将分子写成向量的形式，这个公式会更容易理解：

$(\phi (Q)\phi (K)^T)V= \phi (Q)(\phi (K)^T V)$ $(6)$

注意，特征映射 $\phi (\cdot )$ 是按行应用于矩阵 $Q$ 和 $K$ 的。

从公式(2)可以看出，softmax注意力的计算量是与 $O(N^2 )$ 成正比的，其中 $N$ 是序列长度。内存占用也是一样的，因为必须存储完整的注意力矩阵来计算关于查询、键和值的梯度。相比之下，我们从公式(5)中提出的线性变换器的时间复杂度和内存占用为 $O(N)$ ，因为我们可以计算 $\sum\nolimits_{j=1}^N \phi( K_{j}) V_{j} ^T$ 和 $\sum\nolimits_{j=1}^N \phi( K_{j} )$ ，然后在每个query上重复使用已计算结果。

3.2.1 特征映射和计算代价

对于softmax注意力，乘加的总的计算损失为 $O(N^2 max(D,M) )$ ，其中 $D$ 是查询和键的维度， $M$ 是值的维度。相反，对于我们提出的线性注意力，我们首先计算维度 $C$ 的特征图，然后计算新的值的乘加运算的复杂度为 $O(NCM)$ 。

前面的分析没有考虑核和特征函数的选择。注意，对应于指数核的特征函数是无穷维的，这使得精确的softmax注意力的线性化是不可行的。另一方面，例如，多项式核有一个精确的有限维度的特征映射，并且被证明与指数核或RBF核(Tsai et al.，2019)同样有效。一个线性化的度为2的多项式变换器的计算复杂度为 $O(ND^2 M)$ 。这样的话，当 $N>D^2$ 时，计算复杂度就非常的好。请注意，这个假设在实践中是正确的，因为我们希望能够处理上万个元素的序列。

在我们的实验中，当处理较小的序列时，我们使用一个特征映射，其使得相似度函数的计算结果总为正：

$\phi (x)=elu(x)+1$ $(7)$

其中 $elu(\cdot )$ 表示指数线性单应激活函数(Clevert et al.,2015)。相比 $relu$ ，我们更喜欢 $elu$ ，这样可避免当 $x$ 为负值时梯度为0。这种特征映射方式使得注意力函数的乘加运算的复杂度为 $O(NDM)$ 。在本文的实验部分，我们展示了公式(7)的性能与完整的变换器的性能相当，却显著降低了计算量和内存占用。

3.3 因果屏蔽

通过屏蔽注意力计算，变换器结构可用于高效地训练自回归模型，也就是第 $i$ 个位置的计算只能受它之前的位置的影响，也即位置 $j$ ， $j\leq i$ ，而不能受它之后的位置的影响。因果屏蔽将公式(3)改写为：

$V_{i}^{’} = \frac{\sum\nolimits_{j=1}^i sim(Q_{i}, K _{j})V_{j} }{\sum\nolimits_{j=1}^i sim(Q_{i}, K _{j})}$ $(8)$

按照第3.2节的推导，我们将带屏蔽机制的注意力线性化为：

$V_{i}^{’} = \frac{\phi( Q_{i}) ^T\sum\nolimits_{j=1}^i \phi( K_{j}) V_{j} ^T }{\phi( Q_{i}) ^T\sum\nolimits_{j=1}^i \phi( K_{j} )}$ $(9)$

令 $S_{i}$ 和 $Z_{i}$ 为：

$S_{i} =\sum\nolimits_{j=1}^i \phi( K_{j}) V_{j} ^T$ $(10)$

$Z_{i} =\sum\nolimits_{j=1}^i \phi( K_{j})$ $(11)$

我们可以将公式(9)进一步简化为：

$V_{i}^{’} = \frac{\phi( Q_{i}) ^T S_{i} }{\phi( Q_{i}) ^T Z_{i} }$ $(12)$

请注意， $S_{i}$ 和 $Z_{i}$ 可以从 $S_{i-1}$ 和 $Z_{i-1}$ 计算得到，计算耗时是恒定值，这使得带因果屏蔽机制的线性变换器的计算复杂度和序列长度成线性关系。

3.3.1 梯度计算

任何深度学习框架下，公式(12)的朴素实现都需要存储所有中间值 $S_{i}$ 以计算梯度。这导致内存占用以 $max(D,M)$ 成倍增长，从而阻碍因果模型在长序列或深网络上的使用。

为了解决这个问题，我们将公式(9)中的分子项的梯度导出为累加和。这样我们就可以在前向传播和反向传播中，以线性的时间复杂度和恒定的内存占用，来计算因果注意力。详细的推导过程在补充材料中。

给定分子项 $\tilde{V_{i} }$ ，以及标量损失函数对该分子项的梯度 $\nabla_{\tilde{V_{i} } } L$ ，此处省略一段话以及三个公式，回头再看，为因果模型而设计的，不一定用得上。

这样，公式(9，13-15)中的累加和项的计算的时间复杂度与序列长度呈线性关系，内存占用为恒定值。这样，给定一个维度为 $C$ 的特征图，算法的计算复杂度为 $O(NCM)$ ，内存占用为 $O(Nmax(C,M))$ 。算法1是分子项的前向传播和反向传播的伪代码实现。

算法1

3.3.2 训练和推理

当训练自回归变换器模型时，完整的序列真值是可得的。这使得，对于公式(1)中的 $f_{l} (\cdot )$ ，以及注意力的计算，在层的层面上的并行化是可行的。因此，变换器比递归网络训练得更快。另一方面，在推理阶段，时间步 $i$ 的输出是时间步 $i+1$ 的输入，这使得自回归模型无法并行化。此外，变换器的每个时间步的计算耗时不是不变的，而是和当前序列长度的平方成正比，因为必须计算所有以前的时间步的注意力。

我们提出的线性变换器结合了两者的优点。训练时，计算可以并行化，充分利用GPU或其它加速器的优势。推理时，每个时间步的预测的计算耗时和内存占用都是恒定的。这意味着我们可以简单地将矩阵 $\phi( K_{j}) V_{j} ^T$ 存储为内部状态，并像递归网络一样在每个时间步对其进行更新。这样我们的线性变换器要比其它的变换器模型要快千倍。

3.4 变换器就是递归网络

已有文献中，变换器和递归网络被认为是从根本上不同的。然而，从3.3节的因果屏蔽公式和上一节的讨论中可以明显看出，任何具有因果屏蔽的变换器层都可以表示成这样一个模型：给一个输入，修改内部状态，预测一个输出，这就是递归网络RNN。请注意，和Universal Transformers(Dehghani，2018)不同，我们说的递归是指计算时间，而不是网络层的深度。

在下面的公式中，我们将公式(1)中的变换器层写成递归网络的形式。得到的RNN有两个隐藏状态，称为注意力记忆 $s$ 和归一化记忆 $z$ 。我们用下标表示递归的时间步。

$s_{0} =0$ $(16)$

$z_{0} =0$ $(17)$

$s_{i} = s_{i-1} +\phi (x_{i} W_{K} ) (x_{i} W_{V} )^T$ $(18)$

$z_{i} = z_{i-1} +\phi (x_{i} W_{K} )$ $(19)$

$y_{i} = f_{l} (\frac{\phi (x_{i} W _{Q} )^T s_{i} }{\phi (x_{i} W _{Q} )^T z_{i} }+x_{i} )$ $(20)$

在上面这些公式中， $x_{i}$ 表示某个变换器层的第 $i$ 个输入， $y_{i}$ 表示该层的第 $i$ 个输出。请注意，我们的公式对特征函数没有任何的限定，它可以表示任何变换器模型，理论上甚至适用于softmax注意力。这些公式是为更好地理解变换器和常用的递归网络(Hochreiter&Schmidhuber，1997)之间的关系，以及存储和检索信息的过程，迈出的第一步。

4. 实验

本节从实验分析我们提的线性变换器的性能。首先，在4.1节，我们从计算量、内存占用，以及在合成数据上的收敛性等方面评估线性注意力。在4.2节和4.3节，分别在图像生成和语音识别两个实际应用中评估我们的模型，以进一步展示我们提出的线性变换器的性能。结果表明，我们的模型的性能与最先进的变换器模型相比是由竞争力的，而GPU内存占用和计算量上是大大减少了。

我们实验中用于对比的基线有两个，一个是带softmax注意力的完整变换器，另一个是Reformer(Kitaev et al., 2020)，后者是最先进的快速变换器。对于Reformer，我们使用已经开源的PyTorch复现代码。对于完整的变换器，我们使用默认的PyTorch实现。请注意，对于Reformer，我们不使用可逆层，但这不影响，因为我们只测量自注意力层的内存占用。在所有实验中，我们使用softmax表示标准的变换器(Vaswani et al., 2017)，用linear表示我们提出的线性变换器，用lsh-X表示Reformer(Kitaev et al., 2020)，其中X表示散列轮数。

在训练线性变换器中，我们使用公式(7)的特征映射方式。我们的PyTorch(Paszke et al., 2019)代码、文档、例子在https://linear-transformers.com/ 公式(13-15)的使梯度计算的内存占用为恒定值的实现是大概200行的CUDA代码。

4.1 合成任务

4.1.1 收敛性分析

为了检验线性变换器的收敛特性，我们在一个带因果屏蔽的人工复制任务上训练模型。也就是说，变换器必须复制一系列符号，类似于Kitaev等人的序列复制任务(2020)。我们使用一个最长为128的序列，包含10个不同的符号，符号由专门的分隔符分开。对于三种变换器，我们都构建4层的网络，8个注意力头，批量大小都是64，都是要RAdam优化器(Liu et al., 2019)，初始学习率都为10-3，更新3000次后学习率减小至10-4。图2所示是损失相对于梯度步数的曲线。我们观察到，线性变换器收敛得更平滑，并且比Reformer收敛到一个更小的损失值，这是因为线性变换器没有引入散列导致的噪声。线性变换器收敛到与完整变换器相同的损失值。

图2 在序列复制任务上的收敛性分析。线性变换器收敛地稳定，并最终性能与softmax变换器一样。

4.1.2 内存占用和计算量

本节比较变换器的计算量和内存占用。我们计算在合成数据上不同的序列长度 $N=\left\{ 2^9 , 2^{10} ... 2^{16} \right\}$ ，并衡量GPU占用的峰值，以及变换器的每个变换消耗的时间。我们将批量大小设置为与序列长度成反比，并报告批量中每个样本的时间消耗和内存占用。

每种变换器模型都被评估到适合GPU内存的最大序列长度。我们使用的是NVidia GTX 1080 Ti，内存为11GB。这使得softmax的最大序列长度为4096个元素，lsh-4和lsh-8的最大序列长度为16384个。正如预期的那样，softmax的时间复杂度和序列长度的二次方成正比。我们的方法比其它的方法都更快，且只需要更少的内存，如图1所示。我们注意到，线性注意力机制和Reformer的时间复杂度、内存占用都与序列长度成线性的正比关系。注意，尽管Reformer的渐进复杂度为 $O(N\log_{}N )$ ， $\log_{}N$ 是足够的小所以不影响计算耗时。

图1 前向-反向传播的耗时与内存占用。线性变换器与Reformer相对于序列长度是线性的，softmax是二次的

4.2 图像生成

变换器在一些条件的或者是非条件的自回归生成任务上展现出很好的性能(Radford et al.，2019；Child et al.，2019)。但是由于任务本身是按顺序的，而且内存占用和序列长度成正比，从变换器中采样是很慢的。本节，我们训练带因果屏蔽的变换器，以逐像素地生成图像。我们取得的性能(bits per dimension)与softmax注意力不相上下，而速度快千倍以上，而且每张图像从第1个像素到最后个像素的内存占用是不变的。在我们的补充材料中，有训练过程、生成图像的质量、生成单张图像的耗时等方面的比较。此外，我们还和faster softmax transformer进行比较，其不同于PyTorch实现，是在推理阶段读取keys和values的缓存。

4.2.1 MNIST

首先，我们在广泛使用的MNIST数据集(LeCun et al., 2010)用自回归变换器训练图像生成模型。本实验的网络结构是8个注意力层，每层8个注意力头。我们将embedding的长度设置为256，也就是每个注意力头是32维。我们的前馈的大小是embedding大小的4倍。我们使用Salimans等人(2017)提出的10种logistics的混合来建模我们的输出。我们使用RAdam优化器，学习率设置为10-4，所有模型训练250个epoch。对于Reformer的散列，我们使用1到4轮的散列。另外，按照Kitaev(2020，即Reformer作者)所提的，我们使用64个buckets以及大约32个元素为一个chunk。具体地是，我们将长度为783的输入序列划分为27个chunks，每个chunk有29个元素。由于序列长度实际上很小，即只有784个像素，为了避免因设置不同的批量大小而导致对比不公平，我们将3种变换器的批量大小都设为10。

表1是实验结果。我们观察到，就最终的性能上（就困惑度而言），线性变换器与softmax变换器性能几乎相同，而速度快300倍以上。这是因为我们的模型内存占用低，它能够用一个GPU同时生成10000个MNIST图像。特别的，无论当前序列长度多少，内存占用是恒定的，因为在像素之间唯一需要存储的东西是公式(18)和公式(19)中的 $s_{i}$ 和 $z_{i}$ 。相比之下，softmax变换器和Reformer的内存占用都会随着当前序列长度的增长而增加。

表1 在MINST图像上的自回归图像生成任务。线性变换器的性能与softmax变换器几乎相同，但是要快300倍多

图3所示是我们的模型的图像补全以及无条件图像生成的结果的样例。我们观察到，我们的线性变换器生成的样本是很逼真的，边界清晰，没有噪声。在图像补全的结果中，我们还观察到，我们的线性变换器的结果有与原图像相同的笔划样式和宽度，也就是能够有效地捕获远程关系。请注意，所有变换器产生的困惑度都是相当的，我们没有观察到不同变换器生成结果的质量上的差异。

图3 无条件生成样本和图像补全结果。(a)是被遮挡的样本，(b)是补全结果，(c)是原始图像，我们的模型与softmax变换器性能相当，但快300倍多，每秒生成142张图像。

4.2.2 CIFAR-10

当序列越长时，我们的线性变换器就越有优势。为了展示这一点，我们训练了16层的变换器来生成CIFAR-10图像(Krizhevsky,2019)。每层的设置与前一个实验相同。对于Reformer，我们再一次使用64 buckets and 83 chunks of 37 elements，这和论文中建议的32接近。由于序列长度几乎是前一个实验的4倍，在我们能使用的最大GPU上（即NVidia P40，24GB的内存）上，批量大小设置为1。线性变换器和Reformer的批量大小都设置为4。所有模型都训练7天。表2中是每维比特数和图像生成吞吐量。请注意，尽管这个实验不是重点关注最终的困惑度，但是很明显，序列越长，更快的变换器在每GPU小时上的效率的优势就越明显，比更慢的变换器模型的得分也更高。

表2 在CIFAR-10上生成图像，单块GPU，用7天训练自回归变换器模型。我们的线性变换器要比softmax训练3倍多的epochs，因而困惑度更好。我们的模型在生成图像上比基线快4000多倍

由于对于Reformer和softmax注意力来说，生成单个像素的内存占用和耗时，与像素的个数成二次的正比关系，因此线性变换器的的吞吐量增加更为明显。尤其是，当softmax变换器生成1张图像的时候，我们的线性变换器可以生成4460张图像。我们观察到，我们模型生成的图像是具有空间一致性的，并且生成结果逼真，不会显著阻碍图像分类识别。例如，图4b中，所有图像都成功地完成了狗的鼻子（第一排）或卡车的挡风玻璃（最后一行）。

图4 在CIFAR-10上生成图像的无条件生成样本和图像补全结果，(a)是被遮挡的原图，(b)是补全结果，(c)是原图。当序列越长时，线性变换器在速度上相比softmax变换器的优势就越能体现出来，我们的模型要快4000多倍，每秒生成1785张图像

4.3 自动语音识别(ASR)

为了证明我们的方法也可以用于非自回归任务，我们评估了线性变换器在端到端ASR任务上的性能，使用CTC损失(Connectionist Temporal Classification，2006)。我们以非自回归方式预测每个输入帧的音素分布。我们使用我们使用80小时的《华尔街日报》数据集（Paul&Baker，1992）和40维的不带temporal differences的mel-scale filterbanks作为特征。这个数据集的平均序列长度是800帧，最大长度是2400帧。本实验我们还对比了双向LSTM(Hochreiter&Schmidhuber，1997)，其有3个隐藏层，隐藏层size是320。我们使用Adam优化器，学习率设置为10-3，当验证集错误率不下降时学习率就减小。变换器模型有9层，每层6个注意力头，embedding的维度和图像的实验相同。使用RAdam优化器，初始学习率是10-4，当验证集错误率不下降时学习率减半。

所有模型都是在每个训练epoch上评估音素错误率(PER)。我们观察到，线性变换器在性能和速度方面都远远优于递归网络基线和Reformer，如表3所示。请注意，与所有基线相比，softmax变换器实现了更低的错误率，但速度很慢。尤其是，线性变换器在每个epoch上要比它快3倍。我们在补充材料中提供了训练过程的曲线图。

表3 在WSJ数据集上的语音识别。评估指标是音素错误率(PER)和每个epoch的训练耗时。我们的模型比LSTM和Reformer性能更好，且训练得更快，推理得更快。

5. 结论

我们提出线性变换器，大大减小原初变换器的内存占用和计算量。具体地是，利用矩阵相乘的结合性，我们使自注意力计算的时间复杂度、内存占用与序列长度成线性关系。我们证明我们的模型可以用于因果屏蔽任务，并仍然保持线性渐近的复杂度。最后，我们将变换器模型表示为递归网络的形式，这样我们就可以对自回归任务更快地推理。

这一特性为将来研究RNN和Transformer中的信息存储和检索开辟了许多方向。另一个有待探索的研究方向是选取线性注意力的特征映射方式。例如，用随机Fourier特征逼近RBF核可以使我们使用带有softmax注意力的预训练模型。

补充材料：

A. 梯度推导

略

B. 训练过程曲线

在图5中，我们展示了实验中所有变换器模型的训练过程曲线。对于MNIST实验（图5a），所有方法都是训练250个epoch。因为序列长度是很小的，所以不同方法的训练耗时不会有明显变化。我们还观察到，我们的方法和softmax注意力是收敛到相同的水平，显然优于两个Reformer模型。

图5 线性变换器总是收敛地比Reformer更快，在自回归任务上与softmax相当。在MNIST上都是训练250epoch，在CIFAR-10上都是训练7天，在语音识别上都是训练到收敛为止。实验细节见4.2.1节、4.2.2节、4.3节

在CIFAR-10上，见图5b，我们设置所有模型的训练耗时都为7天。我们观察到，线性变换器和lsh-1比另外两个模型能训练更多的epoch，并且取得更好的性能。如果序列更长的话，这一差距有望继续增大。

最后，在语音识别任务上，见图5c，softmax收敛时的性能明显优于Reformer和线性变换器。注意，线性变换器训练每个epoch的速度要快3倍，这意味着它比softmax训练了大约4倍多的epoch。尽管在这项任务中，softmax变换器是表现最好的，但我们观察到，和Reformer相比，线性变换器在收敛性和最终性能上都更好。

C. 图像生成吞吐量讨论

C.1. Stateful softmax 注意力

4.2节，我们报告了图像生成的吞吐量，并与softmax transformer和lsh进行了比较。在本节中，我们创建另一个基线，表示为stateful softmax，它是在自回归任务上将softmax变换器改成递归网络的计算方式，也就是说，所有的keys和values都会被保存下来，然后当预测序列的下一个元素时，将它们再传递给模型。这个递归模型的状态(state)是一组键和值，其规模与序列长度成正比。这与我们提出的模型有质的不同，我们提出的模型的状态的维度是固定不变的，并且是给定第 $i-1$ 个状态来计算第 $i$ 个状态，无论 $i$ 是多少，计算量都是不变的。

表4是结果。我们注意到，stateful-softmax比原初的变换器要快得多，然而它的复杂度仍然是序列长度的二次方，在CIFAR-10上，我们提的方法要比stateful-softmax快50倍以上。另外我们要说明的是，构建stateful-Reformer并不容易，因为每次提供新的输入时都需要执行排序(sorting)和分块( chunking )操作。

表4

C.2. 将批量大小都设为1

在前面的部分中，我们评估了用于自回归图像生成任务的所有变换器变量的吞吐量。然而，要考虑的另一个重要因素是延迟，即生成单个图像所需的总时间。为此，我们使用批处理大小为1并测量所有方法生成单个图像所需的时间。除了在GPU上运行推理，我们还评估了CPU上所需的时间。结果见表5。

表5

我们观察到，所有的方法都没有充分利用GPU，并且比表4的图像生成吞吐量小得多。我们所提的线性变换器是最快的，尤其是，它在CIFAR-10上比softmax变换器要几乎快6.6倍（笔者注，CIFAR-10的图像尺寸是32x32）。请注意，我们的线性自回归变换器是唯一一个，在各种设置下，在CPU上比在GPU上更快的。这是因为将注意力机制以RNN的形式的计算成本很低，以至于主要的计算瓶颈是序列上不可避免的外循环。

D 图像生成的定性分析

在本节中，我们将为我们的图像生成实验提供定性的结果。由于所有模型的困惑度大致相同，如预期的那样，质量的差异并不显著。然而，一个相当有趣的观察结果是，Reformer模型在其无条件生成样本中提供的变化要少得多。此外，我们观察到，图像完成是一项比无条件生成更容易的任务，因为所有模型都表现得更好。