Press "Enter" to skip to content

如何让Transformer更高效处理长序列

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

论文标题:

 

REFORMER: THE EFFICIENT

 

论文链接:

 

https://openreview.net/forum?id=rkgNKkHtvB

 

收录情况:

 

Under review for ICLR2020 (668 by far)

 

代码链接:

 

https://pastebin.com/62r5FuEW

 

本文提出 Reformer—— 通过三种实用的方法使 Transformer 节约内存、加速训练、处理长序列输入。本文比较偏工程,算法复杂,建议了解即可。

 

概述

 

Transformer现在已经成了诸多NLP模型的标配,但是它的问题是,模型太大,不能处理较长的序列。

 

比如,就目前最大的Transformer结构来看,光是参数量就要占2GB内存,然后储存batch_size为8、embedding为1024的64K序列又需要占2GB内存,内存开销太大了!

 

如何减少内存开销、加速模型训练、处理更长序列,是一个十分现实的问题。

 

本文提出一种Reformer模型用于缓解上述问题:

 

使用可逆层(reversible layers),只存储单层激活值的一份拷贝

 

把FF层里的激活值进行切分

 

使用局部敏感哈希(LSH)注意力代替传统多头注意力

 

使用上述三种方法,本文在长达64K的文本任务和长达12K的图像任务上进行试验,结果表明Reformer跑得比Transformer不知道快到哪里去了,既节约了内存,又拓展了Transformer处理长序列的能力。

 

局部敏感哈希注意力

 

我们知道,Transformer中的注意力计算需要让矩阵Q和K的转置相乘。我们假定它们的形状都是[ batch_size, length, dimension ],那幺如果序列长度有64K,就有得到一个64K*64K的矩阵,显然是不现实的。

 

那幺一个简单的想法是,我们不一次性算所有的Q,而对每个q(i)单独计算即可,然而在需要BP的时候再算一次即可。虽然这样做不那幺高效,但至少它可以处理比较长的输入序列。

 

共享QK参数

 

Q,K,V是通过三个不同的线性变换()得到的,我们可以让Q,K的变换矩阵相同,让V单独有一个。这样做实际上并不会有损效果。

 

哈希注意力

 

上面我们提到,限制长序列的主要原因是当length很大时[ length, length ]的矩阵是不可行的。

 

但实际上,我们真正关心的是softmax(QK^T),又因为softmax是被那些比较大的数支配的,从而,对每个q(i),我们只需要去找K中离q(i)最近的的那些就好了,也就是选一个子集。

 

局部敏感哈希

 

那幺怎幺找到最近邻呢?这可以使用局部敏感哈希(LSH)。把每个向量x映射为一个哈希值h(x)叫做局部敏感哈希,如果比较近的向量能以高概率映射到同一个哈希值,而比较远的向量能以高概率被映射到不同的哈希值。

 

为了得到b个不同的哈希值,我们随机一个矩阵R,大小为[ dimension, b/2 ],然后定义h(x)=argmax([xR; -xR])。这样,对所有的x,我们就可以把它们分配到b个哈希桶里面。

 

局部敏感哈希注意力

 

方便起见,下面用另一种方式重写一下q(i)关注K的方程:

 

 

这里z是归一化项(不用管它)。现在来考虑LSH注意力。我们只需要让q(i)去关注在同一个哈希桶里面的k(j)即可:

 

 

(右a-b)是和传统注意力的比较。(a) 表明传统的注意力是很稀疏的,也就是说大多数的字符其实都不用关注;(b) k和q根据它们的哈希桶(注意随机的那个矩阵R是共享的)排序好,然后再使用。

 

然而另一个问题是,这样得到的哈希桶的大小很可能不均匀。我们从小到大给Q的哈希桶排序,在每个桶内部,按照位置先后排序。

 

在排序后的注意力矩阵中,来自同一个哈希桶的(q,k)对会聚集在矩阵的对角(下图右c)。最后,把它们分组,每组m个,在各组内相互关注即可。

 

 

多轮LSH哈希

 

为了进一步减小桶分布不均的情况,可以用不同的哈希函数进行多轮哈希,具体参见原文。

 

共享QK中的自掩码(Causal Masking for Shared-QK Attention)

 

在普通注意力中一个位置可以关注自己,但在共享QK中我们一般不考虑,从而我们禁止它去关注自己,除非没有别的可以关注的了。

 

下表是几种注意力方式的时空复杂度:

 

 

可逆Transformer

 

如上所述,attention的复杂度可以被减少为序列长度的线性级,但是,参数量占的复杂度依旧很高,我们想要进一步减少。

 

可逆残差网络

 

对输入x,残差网络的输出是y=x+F(x),一个可逆层定义在一个输入输出对上:

 

 

然后,输入输出就可以是:

 

 

可逆Transformer

 

同上,我们用F表示注意力层,用G表示FF层:

 

 

FF层分组

 

可以进一步把Y(2)分组:

 

下表是所有变体的复杂度:

 

 

实验

 

接下来,我们在imagenet64和enwik8-64K上实验,其他设置详见原文。

 

下图是不同的方法在这两个数据集上的表现,可以看到,无论是共享QK还是可逆Transformer,都不会影响效果。

 

 

下图是不同哈希桶数的LSH注意力的表现。显然,数量越多,效果越好,这是因为关注就越精确,同时模型代价就越高。

 

 

最后来看看Reformer的层数。下图(左)是Big Reformer随层变化的不同效果,20层依然无压力。

 

而下图(右)是普通注意力和LSH注意力在不同序列长度的速度比较,当序列很长的时候,LSH具有显着的优势。

 

 

小结

 

本文提出使用LSH(局部敏感哈希)和可逆网络去加速Transformer训练,节省内存以处理更长序列,在几个实验上取得了较好的效果。然而本文比较偏工程,而且算法也比较复杂,各位了解一下即可。

Be First to Comment

发表评论

邮箱地址不会被公开。 必填项已用*标注