Transformer性能优化:运算和显存

Transformer性能优化:运算和显存_第1张图片

©作者 | 王晗炜

单位 | 中科院信工所ASCII LAB

研究方向 | 自然语言处理

244b1eba7a4bacd45f090e4609b50297.png

概述

Transformer [1] 在如今的深度学习领域有着不可或缺的地位,它被广泛应用于自然语言处理、图像处理等领域,有着极大的影响力。自注意力机制作为 Transformer 模型的核心,其几乎不含归纳偏置的特性在足量数据的基础上带给了 Transformer 模型强大的建模能力,也给 Transformer 带来了一系列效率问题:运算和显存的限制使得其无法在长序列问题建模上应用。针对此问题,许多工作对 Transformer 的结构进行了魔改,使得其性能得到优化,下面将对部分有代表性的工作进行介绍。

b1d356997b19fc98318eb68745d6f491.png


经典Transformer结构

经典的 Transformer 结构包含 Encoder 和 Decoder 两个部分,主要的组成组件可以分为以下三个小块:多头自注意力(Multi-Head Self-Attention)、前馈神经网络(Position-wise Feed-forward)以及残差连接(Residual Connect)。

Transformer性能优化:运算和显存_第2张图片

▲ Transformer模型结构图

2.1 Multi-Head Self-Attention

ec72d62b72bf0a02b8601eb7247c9061.png

982d0e7ec7c8df4d690b4e92efacb6d0.png

自注意力机制是一种基于缩放点积的注意力机制,其将原始序列的输入向量投影至三个不同的空间,作为 query、key 和 value,每个序列中的输入都会对整个序列进行注意力计算,包括自身。在 Transformer 中这个过程还引入了多头机制,即对输入向量进行维度的切分,在每个头的空间中进行同样的操作,最后再将每个头的信息拼接起来,以便学到更丰富的信息。

2.2 Position-wise Feed-forward

在 Transformer 的每个小模块中,多头自注意力模块后面都会接上一个全连接的前向神经网络模块,其对输入向量的维度进行放大再缩小(一般为放大四倍再缩小,例如 BERT),以下为模块的公式描述:

3f4ad6c7af25f57feea70e3d0445341b.png

2.3 Residual Connect

为了使模型能够堆叠更多的子模块,完成深度网络的训练而避免梯度消失等问题,Transformer 内每一个自注意力模块和前向神经网络模块均会伴有残差连接模块:

faa1c9f1e5f32b55dba73daf389ce87f.png

2.4 复杂度分析

通过对以上三个模块的简单回顾,我们可以发现 Transformer 内主要运算资源消耗集中在 Mulit-Head Attention 模块和 FFN 模块,稍加分析我们可以得知其时间复杂度和空间复杂度如下:

Transformer性能优化:运算和显存_第3张图片

▲ self-attention 和 FFN 模块复杂度分析,T 代表序列长度,D 表示隐层向量的维度

从表中的分析可以看出,当输入序列较短的时候,模型的主要计算开销集中在 FFN 模块(复杂度与隐层向量维度的平方成正比),而当输入序列较长时,模型的计算开销则会转移至 Multi-Head Self-Attention 模块(复杂度与序列长度的平方成正比)。目前基于 Transformer 结构的 BERT 等预训练模型在一定维度下已经能够达到较好的效果,因此 FFN 模块的计算开销一般能够被承受,但是当输入序列过长时,模型则无法进行处理,这也是目前 Transformer 架构模型的一个通病。下面我们将介绍一些工作针对此问题对 Transformer 模型的改进。

cd7705ebf34dd3365802c5595fcfad09.png


基于递归连接的改进

使用经典的 Transformer 模型时,当遇到输入为长序列,很多模型采用的方式都是直接将序列进行截断,像 Bert 这种预训练模型一般只保留前 512 个字符,或者其他的模型将文本划分为多个 segments,训练的时候对每个 segment 单独处理,segments之间没有联系,这便会导致以下两个问题:

  • 文本最长语义依赖关系取决于 segment 的长度,不同 segment 之间没有关联

  • 分割出来的 segments 语义不完整,存在一句话分隔在两个 segment 之中的情况

针对这一问题,Transformer-XL [2] 对原始的 Transformer 结构做了以下两点改进:

  • 提出片段级递归机制,引入一个记忆模块,循环用来建模片段之间的联系。使得 Transformer 能够对长距离依赖进行建模,解决上下文碎片化问题

  • 提出相对位置编码机制,代替绝对位置编码。避免使用记忆模块时出现混淆

首先我们可以结合论文中的图片来分析经典的 Transformer 是如何对长文本编码的:

Transformer性能优化:运算和显存_第4张图片

▲ Transformer分段训练和测试

通过这张图可以清晰地看见在训练过程中每个 segment 会分别编码,相互之间没有任何联系。在测试时为了保证测试的完整性,其对 segment 的划分更为紧密,不同 segment 之间的重叠度很高,造成了运算资源开销巨大。

而为了解决这一问题,Transformer-XL 引入了一个 memory 状态,在对当前 segment 进行处理的时候,缓存并利用上一个 segment 中所有 layer 的隐向量序列,综合两个 segment 的信息进行 attention 等操作。为了节约计算资源的开销,上一个 segment 的所有隐层向量只参与前向计算,不进行反向传播,具体计算过程如下: 

Transformer性能优化:运算和显存_第5张图片

通过此机制,Transformer-XL 能够将序列的所有 segments 连接起来,在训练测试过程中保持前后文的关联。并且可以显著地提高测试时的效率,其过程可以如下图所示:

Transformer性能优化:运算和显存_第6张图片

▲ Transformer-XL的训练和测试过程

在引入片段级递归机制后,虽然能够建模出截断的文本的关联性,但是也带来了一个额外的问题:每个 segment 都添加相同的位置编码,多个 segments 之间无法区分位置关系。针对此问题,Transfomer-XL 放弃使用绝对位置编码,而是采用相对位置编码,在计算当前位置隐向量的时候,考虑与之依赖 token 的相对位置关系。

首先将包含绝对位置编码的 Attention 计算展开:

Transformer性能优化:运算和显存_第7张图片

可以发现展开后的表达式与位置信息相关的仅为 ,如果以一个 的相对位置视角来看, 对于所有位置均为一个定值,变化的只是 的值,因此在使用相对位置编码时我们可以将公式改写为以下的形势:

Transformer性能优化:运算和显存_第8张图片

其中 和 为随机初始化的可学习参数,在计算 self-attention 时,由于 query 所有位置对应的 query 向量是一样的,因此不管的 query 位置如何,对不同单词的 attention 偏差应保持相同。 为计算出来的位置向量编码,同 Transformer 里的计算方式一致。

在原论文中作者使用字符级的语言模型任务对 Transformer-XL 进行了评估,与当时的一些基线模型相比取得了最优的结果:

Transformer性能优化:运算和显存_第9张图片

▲ 部分语言模型实验结果实验

同时,如同前面分析的一致,Transformer-XL 在进行长文本推理时有着十分明显的速度优势,并且随着序列长度的增大愈发明显:

Transformer性能优化:运算和显存_第10张图片

▲ 推理加速同序列长度的关系

总体而言,Transformer-XL 通过引入递归连接机制和相对位置编码机制,使得 Transformer 模型有了能够处理长序列的能力,无论在模型效果还是性能上均有可观的提升,但是其对于长序列的解决思路是停留在分段这一个前提上的,并没有真正得一次性处理全部长度的序列,只不过通过增加一些额外的空间开销来让实验性能和效率有一定提升,也就是空间换性能和时间,从这个角度上来看 Transformer-XL 其实并不算优化了复杂度,反而增加了复杂度,提升的速度也只是相对于截断式的原始 Transformer 速度的提升。

f7c7bf33effe472bd86d60e70891fe61.png


基于稀疏注意力的改进

Transformer 强大的性能可以说很大程度上来源于其特殊的 Self-Attention 注意力,而其巨大的计算开销也来源于此,可谓祸福相依,接下来的改进均将围绕这一机制进行,即如何在保证模型性能的前提下尽量减少 Attention 过程的计算开销。


4.1 Sparse Transformer

引入稀疏注意力是很多工作的解决思路,其中 Sparse Transformer [3] 是较为早期的工作,其关注长序列的生成问题,针对 Transformer 因效率问题无法应用于长序列生成的任务,提出一种稀疏注意力机制来解决此问题,最终在图像、文本、音频三个模态上验证了其效果。

首先作者为引入稀疏注意力机制做了一些实验进行支撑:其构建了一个由原始 Transformer 模块深层全注意力网络,用于图像的像素级自回归生成任务,即把图像的像素点按照从上到下从左到右的方式当成一个序列,然后在序列上去做自回归。其在生成过程中对网络的注意力情况进行了观察,如下图所示,底部的黑色部分表示还没有生成的部分,白色凸显的部分则是注意力权重高的地方。下图是比较低的层次的注意力,可以看到,低层次的时候主要关注的还是局部区域的部分。

Transformer性能优化:运算和显存_第11张图片

▲ 浅层注意力

在网络的 19 层和 20 层,注意力的分布也发生了变化,开始关注横向和纵向的像素点:

Transformer性能优化:运算和显存_第12张图片

▲ 19层和20层的注意力

而还有一些层的注意力甚至学习到了物体本身的一些特征,例如它们的边缘:

Transformer性能优化:运算和显存_第13张图片

▲ 某些层的注意力

而我们通过对这些层注意力的观察可以得知不同层的注意力关注的区域是不一样的,但是无论其关注的是哪一部分的信息,注意力权重高的区域始终很小,即注意力是稀疏的,这为稀疏注意力机制的引入提供了实验依据。

作者在设计稀疏注意力之前先对 Attention 的计算过程进行了重新定义:

Transformer性能优化:运算和显存_第14张图片

这里同原始 Attention 范式的区别主要在于引入了 Attend 项,即当前位置关注的范围,这里 表示 需要关注的范围,对于自回归生成任务中的 Transformer,。根据这个范式,我们在定义不同的稀疏注意力时就是在定义 。

论文中定义了两种 Attention 机制,第一种为 Strided Attention,其描述如下:

5fe50a3791b9537e02e63ca38f658ec5.png

Transformer性能优化:运算和显存_第15张图片

▲ Strided Attention模式图

上面的图代表是二维的图像像素图,下面的图是序列的注意力矩阵图。这个注意力其实就分成两个部分,每个字符关注前面若干个字符,然后每隔相同的间距关注一个字符。对应到图像里面就是一个关注行,一个关注列。作者认为这种 attention 机制比较倾向于有规律的数据,比如图像音频,不适合文本。

另外论文中还设计了一种 Fixed Attention,描述如下:

5225abacc5dc61d3b0bb4f7fac8bb0ad.png

Transformer性能优化:运算和显存_第16张图片

▲ Fixed Attention模式图

这里 c 是一个超参数。既然是针对文本的,我们看这个 attention 其实可以把左边的二维图的每一行看成一个句子,其实就变成了每个 token 关注自己句子之前的 token,然后关注之前句子特定位置的 token。相当于一个是局部注意力,一个是全局注意力。

可以看出每一种 Attention 模式都包含了多种关注方式,因此论文给出了三种 attention 结合的方式:

  • 不同层使用不同的注意力

  • 将所有的注意力方式融合

  • 不同的头使用不同的注意力

论文直接表明第三种方式的效果最好。作者在三个模态的数据集上均进行了自回归任务的实验,均达到了较好的效果,甚至在速度和性能上都超过了原始的 self-attention。

Transformer性能优化:运算和显存_第17张图片

▲ 三个模态上Sparse Transformer的性能结果

Transformer性能优化:运算和显存_第18张图片

▲ 文本和图像上Sparse Transformer的运行效率

由实验结果还可以看出作者设计的注意力模式确实分别使用于图像和文本。至此为对 Sparse Transformer 关于稀疏注意力部分的简单介绍,该模型提出较早,虽然效果不如后面一些同样基于稀疏注意力的模型,但是其范式的总结对后续模型的发展有着深远的意义。

4.2 Longformer

Longformer [4] 和之前介绍的 Sparse Transformer 都可被视为同一类改动,也是设计不同的 Attention 范围,但是 Longformer是专门为 NLP 任务设计的,因此除了评测模型本身的性能,其也在很多 NLP 下游任务上进行了实验。

Longformer 设计了三种 Attention 模式,分为 Sliding Window、Dilated Sliding Window 和 Global Attention:

  • Sliding Window:对每个 token 设置一个 w 大小的窗口,每次进行 Attention 操作时仅对前后 w/2 个位置的 token

  • Dilated Sliding Window:为了能够注意到更远距离的 token,还可以对窗口进行挖洞,每隔 d 个位置关注一个字符

  • Global Attention:在捕获到局部信息的同时,也不能丢失全局的信息,因此存在一些字符能够关注到所有的字符,并且所有字符也能关注到此字符

Transformer性能优化:运算和显存_第19张图片

▲ Longformer中的Attention机制

除此之外论文中给出了三个使用这些 Attention 的细节:

  • 窗口的大小会随着模型层数的加深变大

  • 多头注意力中,可以另某些头使用 Sliding window,另外一些头使用 Dilated sliding window

  • 对于不同的任务,设计不同的位置使用 global attention

在实验部分,Longformer 首先和之前的模型一样在自回归任务上进行了实验:

Transformer性能优化:运算和显存_第20张图片

▲ Longformer 在字符级自回归任务上的性能

可以看出 Longformer 较之于之前的模型有一个更好的效果,而如前文所提到的,Longformer 是专门为 NLP 任务设计的 Transformer 模型,因此其在下游任务上的表现也很重要。因为只修改 Attention 的范围,不修改模型的结构,可以直接在预训练模型 Roberta 的 checkpoint 上继续训练并且运用到下游任务上,这里预训练的任务也是 MLM。将窗口大小设置为 512,添加了额外的 position embedding 到 4096 的大小,通过多次复制其 512 个位置嵌入来对其进行初始化。

随后其在三个下游任务上进行了 Finetuning:

  • Question answering:问题和候选答案使用 Global Attention

  • Coreference Resolution:不使用 Global Attention

  • Document Classification:[CLS] 使用 Global Attention

其实验结果如下所示:

Transformer性能优化:运算和显存_第21张图片

▲ Longformer在下游任务中的效果

可以看出 Longformer 在下游的任务上的表现基本都超过了 Roberta,而且下游任务的数据集文本越长 Longformer 的优势更为明显。


4.3 Big Bird

还是在这个稀疏注意力的模式上,谷歌又继续推出了 [5],除了之前的 window attention 和 global attention,还加入了一种 random attention,让模型不局限于人为增加的强先验。而对于 global attention 的使用,其还提出了两种方式:

  • Big Bird-ITC:选择已经存在于序列的词汇使用 global attention

  • Big Bird-ETC:添加额外的特殊字符使用 global attention

该模型没有单独对模型效率进行评估,而是在大量下游任务上进行了实验并取得了良好的效果:

Transformer性能优化:运算和显存_第22张图片

Transformer性能优化:运算和显存_第23张图片

▲ Big Bird在问答任务上的表现

总体而言,这三种都可以看成是一类稀疏 attention,都是通过引入一些先验知识来限定 attention 的范围,从而提升效率,看起来是直观有效的,但是存在两个问题:

  • 通用性降低:self-attention 相对于 CNN、RNN 这类模型一大优势就是具有更广泛的归纳偏置,也就是能够处理更一般化的信息,做了这些限定之后其实算是一种倒退

  • 稀疏 Attention 计算优化困难:这种稀疏 attention 的做法虽然看上去很简单,但是真正在实现的时候是没那么方便的,前面两种都设计到 cuda 内核的修改,后面的 big bird 设计了一种很复杂的分块计算方法,实际的运算优化度不如理论上高


4.4 Reformer

前面几种人为设计的稀疏 Attention 模式不一定能够满足待处理数据的特点,而且不具有普适性,而 Reformer [6] 使用了一种局部哈希的方式来获取每个字符需要关注的范围。在论文的开始,其提出了三个 Transformer 中存在的效率问题并给出了自己的解决思路:

  • 自注意力机制消耗的资源随着序列长度的增长而平方倍增长:局部哈希

  • 模型占用的内存随模型深度呈线性增长:可逆残差

  • Feed-Forward 部分参数过多:分块计算

首先我们可以介绍一下局部哈希算法,如下图所示,对于空间中的点,先将其投影到一个圆上,然后将分成八个区域,每个区域都代表一个独立的值。随机转动圆,记录下投影后的点所在区域的值;那么经过多次转动后,就为一个点得到了多个值,这些值就是点的哈希值。

Transformer性能优化:运算和显存_第24张图片

▲ 局部哈希示例

上图中的上半部分是两个不相似的点组成的例子,它们的哈希值差别很大。下半部分则是两个相似的点,它们的哈希值也是一样的。通过这个简单的示例我们就能发现,在空间中距离越近的点,其哈希完之后哈希值相同的概率也会越高,这就是局部敏感哈希的原理 [7]。

随后我们可以通过矩阵乘法随机构造一个哈希函数:

c96a70ff9a7cf528cdd635823bc53888.png

接着把原始 Attention 公式进行重写:

884f29b13ee66b84bc6002553bb26423.png

这里 表示 mask 项, 表示归一化项, 啧表示当前字符注意的区域,这里规定为当前字符的 query 的哈希值和序列中 key 的哈希值相同的位置进行 Attention 操作。经过这些定义,我们可以将原序列按哈希值进行排序,并分为 个桶:

Transformer性能优化:运算和显存_第25张图片

▲ 分桶哈希示意图

上图中左边为序列的表示,可以观察到其首先将序列进行哈希然后排序。这首先便会存在一个问题,就是 query 和 key 均通过一个哈希函数进行哈希,因此在一个桶中,可能会发生 query 很多但是 key 很少的问题,甚至,会有 qeury 很多而 key 不存在的问题。为了解决这个问题,这里让 query 和 key 在同一空间,即生成 Q 和 K 的矩阵是同一个。这样在避免了刚刚的问题的同时又引入了一个新的问题:同一个 token 的 key 值和 query 值的点积值会远远高于不同 token 之间的值,因此论文中的做法是不计算本身的 Attention 值。

看到这里 Reformer 对 Self-Attention 的优化思路已经很清晰了,就是讲相似的 token 置于一个桶中,仅对在桶中进行 Attention 操作。然而在实现的过程中,我们需要保证桶内的元素分布均衡,否则在进行计算优化的时候会带来很大的麻烦。在这里为了保证计算的均衡,采用的策略是分块,如果一个桶的元素跨块的话,则让后面的块再去 attend 到前一个块中的元素(仅考虑自回归任务时),如上图左图所示。由此可见 Reformer 的加速仅在模型结构上就能够实现,较之之前几种方式更为方便。

另外,在局部敏感哈希中,随机的次数越多,得到的哈希桶就越准确,所以哈希值可以做多轮。这样,P 就成为多次哈希的值的全集:

0a66143213da01407b991f3a5ce4a254.png

随后论文针对局部哈希 Attention 的优化做了一个很有意思的实验,该实验的任务是重复句子的单词预测,即将两个相同的句子通过特殊字符拼接在一起,以 的格式作为序列输入,每次通过前一个句子预测后面一个句子。其使用了一个单层的 Transformer 模型,通过训练和测试时不同轮数的局部哈希进行了效果对比,实验结果如下:

Transformer性能优化:运算和显存_第26张图片

▲ 重复句预测实验

可以发现无论在训练时还是测试时,哈希的轮数越多模型的效果越好,这证明了多轮哈希的合理性和重要性,同时通过这个实验我们也可以进一步认识到:局部哈希函数无需训练,只需要随机初始化,并且训练和测试过程的哈希轮数和函数也可以不一致,自由度非常高。

在完成了对 Self-Attention 的优化后,论文还提出了一种可逆残差的方式来解决模型过深带来的性能开销。在训练网络的时候,一般需要记录每层的激活值,用来在反向传播的时候进行计算。所以每增加一层,内存也会随之增长。一般的残差网络的结构如下:

56aa431110c583422ca8c00bd4f6e45a.png

我们在反向传播的过程中无法通过 y 复原 x 和 F(x) 的值,因此需要将其记录下来,而论文通过对 outside_default.png进行拆分巧妙地避免了这一问题:

ca782af2fc8466e50fb5ec1e8c68ea52.png

在 Transformer 的模块中 和 不用由 拆分,只需分别代表 Self-Attention 和 FeedForward 的输入即可:

bf7841c2652c921af75bf6f8df6d4cab.png

同时,对于 FeedForward 部分,论文也进行了优化。该模块的全称为 Position-wise FeedForward,如名字所示,改模块对序列中的每个 token 的操作是一样的,也可以理解为 token 之间不会进行交互,因此当输入序列过长时,完全可以讲序列分段输入,来缓解内存上的开销。

论文针对自己优化的模块在一些基础任务上进行了实验,来验证模型优化的有效性:

Transformer性能优化:运算和显存_第27张图片

▲ 验证共享q、k参数和可逆残差的效果

Transformer性能优化:运算和显存_第28张图片

▲不同轮数哈希取得的效果

Transformer性能优化:运算和显存_第29张图片

▲  不同轮数哈希对应的模型效率

总体而言,Reformer 通过实验证实了自己模型结构优化的有效性,多轮哈希在带来效果提升的同时也会一定程度上降低模型的效率。相较于之前介绍的集中稀疏 Attention 机制的模型,Reformer 不依赖人工设计先验知识,模型的适用性更广,并且能够在模型代码层面进行优化,不涉及到cuda内核的修改。但是在模型的设计和实现上过于复杂,在时间和空间复杂度上反复横跳,从最后的实验结果也可以看出模型只在超长序列上效果较为明显,能够应用的场景还是十分有限的。

8aae65ad1fd34a6621bd66d3a284a244.png


基于低秩分解的改进

同样是解决 Self-Attention 部分的复杂度问题,Linformer [8] 从另外一个角度给出了自己的优化方案。同 Sparse Transformer 一样,其也先对自己的理论进行了实验的论证。论文在在 Wiki103 和 IMDB 两个数据集上,用 Roberta-large 模型上计算出的 Attention 矩阵做了奇异值分解,然后从下图左两图中可以看出,前 128 维的奇异值累计值已经占了到了 0.9 左右。

而在右图中可以看到,越高层,128 个奇异值累积值就越高。在第 11 层,128个奇异值累积起来达到了0.96。因而说明了,虽然 Attention 的计算结果是一个 N x N 的矩阵,但其实一个低秩矩阵比如 N x 128 可能就已经足够存储 Attention 的所有信息。因此直接对矩阵进行合理降维就能够在保证模型效果的前提下完成 Self-Attention 过程复杂度从平方级到线性的转变。

作者在论文中花了很多篇幅证明降维的合理性,在这里我们不过多赘述,感兴趣的读者可以自行阅读原文。根据以上的结论,我们很容易想到在训练和测试的过程对 K、V 的矩阵进行 SVD 分解就能够解决效率问题,但实际上再测试过程中进行 SVD 分解还是比较麻烦,因此论文中给出了一种十分简单粗暴的方式来代替这一步骤:直接给 key 和 value 矩阵加一层线性变换,将其的长度变为一个定值,其模型结构如下所示:

Transformer性能优化:运算和显存_第30张图片

▲ Linformer 模型结构图

Transformer性能优化:运算和显存_第31张图片

由公式和模型图均可知 Linformer 在 V 和 K 参与 Self-Attention 计算之前,通过两个投影矩阵将其序列长度那一维降至一个固定值,因此无论原始序列的输入长度是多少,最后均会变成一个 k 长度的向量,最后 Self-Attention 的时间复杂度也会变成线性复杂度。

论文在 MLM 预训练任务和 NLP 下游任务上均做了实验:

Transformer性能优化:运算和显存_第32张图片

▲ MLM 预训练任务实验

Transformer性能优化:运算和显存_第33张图片

▲ 下游任务实验

总体而言 Linformer 无论是在运行效率还是下游任务的效果上都有着不错的表现,实现的方式也十分简单,但是这种投影的方式却有一个致命的缺陷,即无法做自回归的生成任务,因为投影相当于把序列信息都杂糅至一个定长的向量,模型无法通过之前 causal mask 的方式将信息给掩盖,这也是论文只进行了 MLM 预训练任务的原因。

300ff0d07450a4fcc876be23bfc9b10a.png


基于线性注意力的改进

在线性注意力的探索上,除了对 K、V 矩阵进行下采样之外,还有一些工作给出了另一种视角。让我们继续回到注意力的计算过程,其本质上是三个矩阵的连乘,而矩阵的连乘是满足结合律的,我们正常的计算顺序是 ,前一步我们会得到一个 的矩阵,这一步将导致我的时间复杂度为 ,而如果我们以顺序 计算,复杂度则会变成 。

Transformer性能优化:运算和显存_第34张图片

▲ Attention 计算顺序与复杂度

然而 softmax 的存在却让这一步的操作无法实现,因为其需要先将 值指数化并归一化,因此没有办法先计算后面两个矩阵的乘积。然而一定需要 softmax 吗?

首先我们先来思考一下 Self-Attention 的本质,其使用点积相乘并进行 softmax 其实想得到的仅仅是 token 之间的相似度,因此其计算过程可以如下描述:

Transformer性能优化:运算和显存_第35张图片

softmax 的目的其实只是使得注意力的值恒为正且满足归一化,而这两个条件通过其他的方式显然是可以实现的。Linear Transformer [9] 通过一种核函数的方式巧妙地替换了 softmax,其认为只需要找一个恒为正的核函数将 Q 和 K 的值进行映射便可完成以上过程:

Transformer性能优化:运算和显存_第36张图片

这里 代表核函数,其定义如下:

8598f4ce7cfe547c8dd97ae7b1e759b2.png

如此转换我们就可以将softmax函数删去从而使得整个运算的复杂度变为 。

之后谷歌推出的 Performer [10] 的解决思路其实也和 Linear Transformer 是一致的,其工作的最大亮点在于为 softmax 找到了一个更为优美的映射来替代并给出了理论证明,而不像之前 Linear Transformer 那样仅通过核函数的替换显得有些空穴来潮。其映射表示如下 [11]:

Transformer性能优化:运算和显存_第37张图片

由公式可知其使用了一种向量采样的方式来表示 ,随后为了使采样出的向量能够表征更多的信息,即让采出来的向量线性无关,其对采样出的向量使用了正交化技术。

Performer 也在一些任务上进行了测评,也是一些长序列相关的任务,我们可以简单了解一下其大致性能:

Transformer性能优化:运算和显存_第38张图片

▲ Performer 训练和推理速度

Transformer性能优化:运算和显存_第39张图片

▲ Performer 在蛋白质序列预测任务上的性能

模型综合评测

上文已经对 Transformer 的魔改模型分类别介绍了不少,通过这些介绍我们可能能够大致了解这些工作改进的思路和大概方法,但是很难知道这些模型在使用时的差异。针对这一问题,谷歌推出了一个评测框架 LRA [12] 对一众 Transformer 改进模型进行了统一评测。

LRA 旨在用一些(形式)简单、通用、有一定挑战性的多个任务对这些模型进行长序列任务的评测,其中包括:

  • Long Listops

  • Byte-Level Text Classification

  • Byte-Level Document Retrieval

  • Image Classification on Sequences of Pixels

  • Pathfinder (Long-Range Spatial Dependency)

  • Pathfinder-X (Long-Range Spatial Dependencies with Extreme Lengths)

以下为模型在这些任务上的表现:

Transformer性能优化:运算和显存_第40张图片

▲ LRA性能测试效果

Transformer性能优化:运算和显存_第41张图片

▲ LRA速度测试效果

由结果可以大体看出 Linformer、Linear Transformer、Performer 效率很高,Big bird Reformer 效率很低,但是 Big Bird 的性能还是不错的。总体的性能还可以由以下这张图看出,其中纵轴是效果,横轴是速度,圆圈的大小代表所需要的显存。理论上来说,越靠近右上方的模型越好,圆圈越小的模型越好。

Transformer性能优化:运算和显存_第42张图片

▲ Transformer 们的性能-速度-显存图


0810f90cf471761b1e3026666b013f2e.png


总结

本文大体介绍了若干种高效 Transformer 的改进版本,均集中在对长序列任务的处理上,里面有一些笔者对这些模型的思考,其中更多的细节感兴趣的读者可以自行阅读原文。

outside_default.png

参考文献

outside_default.png

[1] Transformer https://arxiv.org/abs/1706.03762

[2] Transformer-XL https://arxiv.org/abs/1901.02860

[3] Sparse Transformer https://arxiv.org/abs/1904.10509

[4] Longformer https://arxiv.org/abs/2004.05150

[5] Big Bird https://arxiv.org/abs/2007.14062

[6] Reformer https://arxiv.org/abs/2001.04451

[7] Reformer: 局部敏感哈希、可逆残差和分块计算带来的高效 https://mp.weixin.qq.com/s?__biz=MzI4ODg3NDY2NQ==&mid=2247483911&idx=1&sn=8d98a214d455a55650bb589830b08dae&chksm=ec368bc1db4102d7d54216e917ec22f83b47df55153ef4c3f83aaafd68a0dd60caf836fb712a&scene=178&cur_album_id=1464771644039610372#rd

[8] Linformer https://arxiv.org/abs/2006.04768

[9] Linear Transformer https://arxiv.org/abs/2006.16236

[10] Performer https://arxiv.org/abs/2009.14794

[11]《Performer:用随机投影将Attention的复杂度线性化 》 https://kexue.fm/archives/7921

[12] LRA https://arxiv.org/abs/2011.04006

更多阅读

Transformer性能优化:运算和显存_第43张图片

Transformer性能优化:运算和显存_第44张图片

Transformer性能优化:运算和显存_第45张图片

1bad51ba5c0a0d669e6ce77fc7412b72.gif

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。

总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。

 稿件基本要求:

• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 

• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题

• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算

 投稿通道:

• 投稿邮箱:[email protected] 

• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿

Transformer性能优化:运算和显存_第46张图片

△长按添加PaperWeekly小编

现在,在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

9eb27e10eaae113e6ecacd29a82fc534.png

你可能感兴趣的:(算法,机器学习,人工智能,深度学习,java)