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

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

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

Deep Attention

在 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，主要关注固定列与最新列元素之后的元素，我们发现这种模式适用于数据无法同二维结构（例如文本）相匹配的情况。更多细节信息，请参阅我们的研究论文。

实验结果

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

密度建模性能，以各种基准数据集上的每字节（或 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/