NLP论文解读:Generating Long Sequences with Sparse Transformers

OpenAl提出了一种适用于文本、图像和语音的稀疏Transformer,将先前基于注意力机制的算法处理序列的长度提高了三十倍。

对复杂高维度的数据分布进行估计一直是非监督学习领域的核心问题,特别是针对像文本、语音、图像等长程、相关性数据更使得这一领域充满了挑战。同时,对于复杂数据的建模也是非监督表示学习的关键所在。

近年来,神经自回归模型在这一领域取得了一系列十分优秀进展,针对自然语言、原始音频和图像成功建模。这些方法将联合概率分布分解成了条件概率分布的乘积来解决。但由于数据中包含许多复杂、长程的依赖性并需要合适的表达模型架构来进行学习,使得对数据的条件概率分布建模依旧十分复杂。

基于卷积神经网络的架构在这个方向取得了一系列进展,但需要一定的深度来保证足够的感受野。为了解决这一问题,WaveNet引入了膨胀卷积(dilated conv)帮助网络在的对数数量层数下学习长程依赖性。于此同时Transformer由于可以利用一定的层数为任意的依赖性建模,在自然语言任务上显示出了强大的优势。由于每个自注意力层用于全局感受野使得网络可以将表示能力用于最有用的输入区域,对于生成多样性的数据具有更加灵活的特征。

但这种方法在处理序列时需要面临随着序列长度平方增长的内存与算力。对于过长的序列,现有的计算机将无法处理和实现。

为了解决这一问题,OpenAI的研究人员在最新的论文中为注意力矩阵引入了多种稀疏的分解方式,通过将完全注意力计算分解为多个更快的注意力操作,通过稀疏组合来进行稠密注意力操作,在不牺牲性能的情况下大幅降低了对于内存和算力的需求。

新提出了稀疏Transformer将先前Transforme的平方复杂度O(N^ 2) 降低为O(N N^1/2),通过一些额外的改进使得自注意力机制可以直接用于长程的语音、文本和图像数据。

原先的方法大多集中于一个特定的领域、并且很难为超过几千个元素长度的序列建模,而稀疏Transformer则可利用几百层的模型为上万个数据长度的序列建模,并在不同领域中实现了最优异的结果。

Sparse Transformer: 主要目的是减少传统Transformer 的时间复杂度和空间复杂度。 通过top-k选择,将注意退化为稀疏注意。这样,保留最有助于引起注意的部分,并删除其他无关的信息。这种选择性方法在保存重要信息和消除噪声方面是有效的。注意力可以更多地集中在最有贡献的价值因素上。

稀疏Attention

self attention 要对序列中任意两个向量计算相关度,得到一个n^2大小的 注意力矩阵。
NLP论文解读:Generating Long Sequences with Sparse Transformers

在上图中,左边显示了注意力矩阵,右变显示了关联性,这表明每个元素都跟序列内所有元素有关联。
所以,如果要节省显存,加快计算速度,那么一个基本的思路就是减少关联性的计算,也就是认为每个元素只跟序列内的一部分元素相关,这就是稀疏Attention的基本原理。

Atrous Self Attention

第一个要引入的概念是Atrous Self Attention,中文可以称之为“膨胀自注意力”、“空洞自注意力”、“带孔自注意力”等。这个名称跟后面的Local Self Attention一样,都是笔者根据它的特性自行命名的,原论文《Generating Long Sequences with Sparse Transformers》并没有出现过这两个概念,但我认为将它们单独引出来是有意义的。
NLP论文解读:Generating Long Sequences with Sparse Transformers

很显然,Atrous Self Attention就是启发于“膨胀卷积(Atrous Convolution)”,如下右图所示,它对相关性进行了约束,强行要求每个元素只跟它相对距离为k,2k,3k,…的元素关联,其中k>1是预先设定的超参数。从下左的注意力矩阵看,就是强行要求相对距离不是k的倍数的注意力为0(白色代表0)

由于现在计算注意力是“跳着”来了,所以实际上每个元素只跟大约n/k个元素算相关性,这样一来理想情况下运行效率和显存占用都变成了O(n2/k)O(n^2 /k),也就是说能直接降低到原来的1/k。

Local Self Attention

另一个要引入的过渡概念是Local Self Attention,中文可称之为“局部自注意力”。其实自注意力机制在CV领域统称为“Non Local”,而显然Local Self Attention则要放弃全局关联,重新引入局部关联。具体来说也很简单,就是约束每个元素只与前后k个元素以及自身有关联,如下图所示:
NLP论文解读:Generating Long Sequences with Sparse Transformers
从注意力矩阵来看,就是相对距离超过k的注意力都直接设为0。
其实Local Self Attention就跟普通卷积很像了,都是保留了一个2k+1大小的窗口,然后在窗口内进行一些运算,不同的是普通卷积是把窗口展平然后接一个全连接层得到输出,而现在是窗口内通过注意力来加权平均得到输出。对于Local Self Attention来说,每个元素只跟2k+1个元素算相关性,这样一来理想情况下运行效率和显存占用都变成了????((2k+1)n)∼????(kn)了,也就是说随着n而线性增长,这是一个很理想的性质——当然也直接牺牲了长程关联性。

Sparse Self Attention

到此,就可以很自然地引入OpenAI的Sparse Self Attention了。我们留意到,Atrous Self Attention是带有一些洞的,而Local Self Attention正好填补了这些洞,所以一个简单的方式就是将Local Self Attention和Atrous Self Attention交替使用,两者累积起来,理论上也可以学习到全局关联性,也省了显存。

(简单画个草图就可以知道,假如第一层用Local Self Attention的话,那么输出的每个向量都融合了局部的几个输入向量,然后第二层用Atrous Self Attention,虽然它是跳着来,但是因为第一层的输出融合了局部的输入向量,所以第二层的输出理论上可以跟任意的输入向量相关,也就是说实现了长程关联。)

但是OpenAI没有这样做,它直接将两个Atrous Self Attention和Local Self Attention合并为一个,如下图:

NLP论文解读:Generating Long Sequences with Sparse Transformers
从注意力矩阵上看就很容易理解了,就是除了相对距离不超过k的、相对距离为k,2k,3k,…的注意力都设为0,这样一来Attention就具有“局部紧密相关和远程稀疏相关”的特性,这对很多任务来说可能是一个不错的先验,因为真正需要密集的长程关联的任务事实上是很少的。

参考:
Sparse Transformer解读
sparse transformer
attention is all you need 解读