近日，OpenAI开发出一种深层神经网络Sparse Transformer，能够预测序列中接下来将要出现的内容——包括文本、图像以及声音等。其利用attention机制的改进算法，能够在相当于以往最大周期的30倍的长序列当中提取模式。

AI研究中的一大现实挑战，在于如何面向复杂数据（例如图像、视频或者声音）对其中距离较远且极为微妙的相互依赖性进行建模。Sparse Transformer采用O(N \sqrt{N})O(NN)这一形式对O(N^2)O(N2) Transformer自关注机制进行进行重述，同时也提供其它几项改进，使其能够直接应用于更为丰富的数据类型。以往，面向这些数据使用的模型往往针对某一领域而专门设计，因此很难被扩展至拥有成千上万元素的长序列当中。相比之下，我们的模型能够利用数百个层对拥有数万个元素的序列进行建模，并在多个领域实现最强大的性能。在OpenAI，我们正在利用它构建出能够更好理解真实世界的AI系统。

Deep Attention

数据类型	存储	重新计算
1024文本标记 (几段)	1.0 GB	16 MB
32x32x3 像素 (CIFAR-10图像)	9.6 GB	151 MB
64x64x3 像素 (Imagenet 64图像)	154 GB	2.4 GB
24000 样本 (长度约2秒的12 kHz音频)	590 GB	9.2GB

在Transformers中，每一个输出元素都被接入每个输入元素，它们之间的权重则根据环境进行动态计算。这一过程被称为attention。虽然人们相信这使得Transformers拥有高于固定连接模式模型的灵活性，但这同时意味着我们必须在实践中为每个层及attention head创建M\乘以NN×N attention matrix。因此，在处理拥有多种元素的数据类型时——例如图像或原始音频——这种机制往往会消耗大量内存。

当指标被存储在内存中，或者在反向传播期间进行重新计算时，Deep Transformer（64层，4 head）attention的内存占用量表格。作为参考，用于深度学习的标准GPU通常拥有12到32 GB内存。

缓解这种情况的有效方法之一，是在反向传播期间立足checkpoints重新计算attention matrix，这在深度学习当中已经是一项成熟的技术，可用于利用更多计算资源减少内存的使用量。当Transformers中的attention matrix完成后，即意味着后续最大内存占用量与具体层数无关，我们将能够借此构建起规模远超以往的深层训练网络。在实践当中，我们发现深度高达128万的Transformers在基准任务（例如CIFAR-10）上的表现要优于深度较浅的网络。

为了更深入地训练这些模型，我们对Transformers中的操作顺序进行了多次调整，并修改了初始化方式。更多细节信息，请参阅我们的研究论文。

Sparse Attention

然而，即使只计算单个attention matrix，如果输入量特别巨大，其可行性仍然会受到严重影响。在这类场景下，我们改为使用sparse attention模式，其中每个输出位置仅计算来自输入位置中某一子集的权重。当该子集在整体输入量中所占的比例相对较小时（例如由NN个元素减少至\sqrt{N}N个元素），那么即使面对非常长的序列，相关attention计算也将变得易于处理，具体算法复杂度由O(N^2)O(N2)变成了O(N \sqrt{N})O(NN)。
为了评估这种方法的可行性，我们首先利用Deep Trasformers对图像中的attention模式进行可视化学习。结果显示，很多图像都表现出了可解释的结构化稀疏模式。在以下几张图片中，我们可以看到特定attention head如何附加各输入像素（以白色高亮显示）以预测图像中的下一个值。当输入部分聚焦在较小的子集上并显示出明确的规则性时，则代表该层易于实现稀疏化。对于CIFAR-10图像中的一套128层模型，我们可以看到以下样本：

第19层

第20层

学习自这个128层CIFAR-10网络中几个层的attention模式（白色高亮）。这些层学会了如何将attention分散在两个维度之上。第19层汇总了每一行中的信息，第20层则按列聚合这些信息，从而实现了整体attention操作的有效分解。

第6层

第36层

某些层学会了访问位置内存，因此无论输入数据或者时间步长如何，它们通常都会访问类似的位置（第6层）。其它一些层则学会了具有高度数据依赖性的访问模式（第36层）。

虽然很多层都显示出稀疏结构，但也有部分层清晰地显示出对整体图像的动态attention。为了保持我们的网络学习这些模式的能力，我们对attention matrix采取了二维分解方法，其中网络可以通过两个sparse attention步骤关注到所有位置。

常规transformer、步进attention和固定attention

第一个版本是步进（strided）attention，大致相当于附加至各行与各列的每一个位置，而且类似于前文说提到的网络学习到的attention模式。（请注意，列attention也可以等效地表示为转置后matrix中的对应行。）第二个版本为固定（fixed）attention，主要关注固定列与最新列元素之后的元素，我们发现这种模式适用于数据无法同二维结构（例如文本）相匹配的情况。更多细节信息，请参阅我们的研究论文。

实验结果

CIFAR10	每DIM Bit
PixelCNN++ (Oord等，2016年)	2.92
Image Transformer (Parmar等，2018年)	2.90
PixelSNAIL (Chen等，2017年)	2.85
Sparse Transformer 59M (256W, 128L, 2H)	2.80

Sparse Transformers在CIFAR-10、Enwik8以及Imagenet 64的密度估算当中获得了远超同侪的分数表现。

ENWIK8	每字节 Bit
Deeper Self-Attention (Al-Rfou等， 2018年)	1.06
Transformer-XL 88M (Dai等2018年)	1.03
Transformer-XL 277M (Dai等， 2018年)	0.99
Sparse Transformer 95M (512W, 30L, 8H)	0.99

IMAGENET 64X64	每DIM Bit
PixelCNN++ (Oord等，2016年)	3.57
Parallel Multiscale (Reed等，2017年)	3.7
SPN 150M (Menick & Kalchbrenner, 2018年)	3.52
Sparse Transformer 152M (512W, 48L, 16H)	3.44

密度建模性能，以各种基准数据集上的每字节（或dim）bit为单位。M表示网络当中所使用的数百万个参数，W表示网络宽度，L表示层数，H表示head数。
我们还发现，除了速度明显提升之外，sparse attention在损失方面也低于full attention（请参阅我们的比较文章）。这可能意味着我们的稀疏模式存在着有益的归纳偏差，或者在密集attention方面存在潜在的优化问题。

生成图像

利用sparse attention的Transformers似乎拥有一种全局结构概念，其能够通过查看图像来完成定性评估。在这里，我们对一个由64\乘以 6464×64 ImageNet训练出的模型进行可视化：

提示图像、补完图像和实际图像

我们还生成了完全无条件的样本，其中未经调整的softmax色温为1.0。这些模型以最大似然目标进行训练，且目标涵盖所有数据模式（包括可能并不存在的数据），而非增加其中某部分数据的保真度。从这些未调整色温的模型当中采样，我们得以看到该模型认为的、真实世界中图像的完整延伸。结果是，某些样本呈现出非常奇怪的状态。

模型样本、真实数据

生成原始音频波形

Sparse Transformers也可以通过简单的嵌入位置变更生成原始音频，而非仅限于生成图像。在尝试将深度学习扩展至新的数据类型中之后，我们认为这类网络能够轻松接受归纳偏差的特性将推动其发展为一种重要的工具。

该模型利用古典音乐原始片段进行训练，并使用sparse attention生成长度为65000的序列。这一长度相当于约5秒长的原始音频，我们随后在每个片段中将多个样本连接在一起。

代码发布

一般来讲，实现sparse attention往往需要对block中的查询与关键指标进行拆分。为了简化实验过程，我们构建了一组block-sparse内核，这些内核能够在GPU之上高效执行这些操作。我们对这些内核进行了开源，并在库中提供关于sparse attention函数的示例。

未来工作与当前局限

我们引入的sparse attention模式只是实现长序列高效建模道路上的起步阶段。我们认为接下来有必要对稀疏性的不同模式与组合进行深入探索，而学习稀疏模式则有望成为下一代神经网络架构极具发展潜力的研究方向。
虽然实现了前文所提到的改进，但就目前而言，利用自回归序列生成超高分辨率图像或视频仍然不具备现实可行性。不过，我们引入的attention优化操作可能将发挥重要作用，并与其它方法（例如多尺度方法）结合起来以实现高维数据建模。

原文链接：
https://openai.com/blog/sparse-transformer/

创作场景

OpenAI 新发现：利用 Sparse Transformers 实现文本、图像、声音预测