OpenAI新研究补齐Transformer短板，将可预测序列长度提高30倍

2019 年 4 月 24 日 机器之心

选自 openai

机器之心编译

机器之心编辑部

Transformer是一种强大的序列模型，但是它所需的时间和内存会随着序列长度出现二阶增长。近日，OpenAI研究人员开发出了一种深度神经网络Sparse Transformer，该网络在预测长序列方面创造了新纪录——无论预测的是文本、图像还是声音。该神经网络利用注意力机制中的一种改进算法，可以从长度可能是之前30倍的序列中提取模式。

现在，AI 研究中的一项挑战是在图像、视频或声音等复杂数据中进行长序列的精细相关性建模。Sparse Transformer 合并了 O（N^2）Transformer 自注意力机制的 O(N√N) 重组以及其他一些改进，从而直接用于这些丰富的数据类型。以前，这些数据上所使用的模型是专为某个领域制作的，或者很难将序列扩展到包含几千个元素。

相比之下，OpenAI 开发的模型通过使用数以百计的层可以对包含上万个元素的序列进行建模，在诸多领域都取得了当前最佳的表现。OpenAI 研究人员利用该模型帮助创建能够更好地理解世界的 AI 系统。

深度注意力

在 Transformer 中，每一个输出元素与输入元素相连接，同时根据具体情况对它们之间的权重进行动态计算，这一过程被称为「注意力机制」。虽然人们相信这使得 Transformer 较那些具有固定连接模式的模型更为灵活，但实际操作中需要为每一层和注意力头创建一个 N×N 注意力矩阵，当应用于图像或原始音频等具有许多元素的数据类型时会消耗大量内存。

当矩阵存储在内存或在逆推计算过程中进行再计算时，深度 Transformer（64 层和 4 个头）的注意力内存使用情况。作为参考，用于深度学习的标准 GPU 内存通常是 12-32GB.

减少内存消耗的一种方法是在反向传播过程中从检查点处重新计算注意力矩阵，这是深度学习中的一种成熟的方法，以更多的计算来减少内存使用。

当 Transformer 中的注意力矩阵完成时，这意味着最大的内存消耗将不受层数的支配，使研究人员训练网络的深度大大超过从前。在实际操作中，研究人员发现在处理 CIFAR-10 等基准测试任务时，深度达 128 层的 Transformer 表现出的性能优于较浅的网络。

为了训练深度更大的模型，研究人员对 transformer 的操作顺序进行了几次调整，修改了初始化方法。详情参见论文。

稀疏注意力

然而，对于非常大的输入来说，甚至计算单个注意力矩阵都是不现实的。因此，OpenAI 使用了稀疏注意力模式，在这种模式中，每个输出位置仅从输入位置子集中计算权重。当子集相对于整个输入集较小时（如元素数量是√N 而不是 N），即使对于非常长的序列，注意力计算也会变得比较容易，算法复杂度为 O（N√N）而不是 O（N^2）。

为了评估该方法的可行性，研究人员首先可视化并学习了图像上深度 Transformer 的注意力模式，发现其中许多模式表现出了可解释和结构化的稀疏模式。以下每幅图像都显示了哪个输入像素（白色高亮标出）由一个给定的注意力头处理，以预测图像中的下一个值。当输入部分集中在小的子集上并显示出高度规律性时，该层就易于稀疏化。以下是 CIFAR-10 图像上 128 层模型的样本：

左：Layer 19，右：Layer 20。为一个 128 层的 CIFAR-10 网络的若干层学习注意力模式（白色高亮显示）。这些层学会了在两个维度上分割注意力。Layer 19 汇总每一行的信息，Layer 20 按列汇总这些信息，从而有效分解了全注意力运算。

为获取位置记忆而训练的层（左：Layer 6；右：Layer 36），它们通常关注类似的位置，不管输入数据或时间步长如何（Layer 6）。其他层学习高度依赖数据的访问模式（Layer 36）。

虽然许多层显示出稀疏的结构，但有些层清晰地显示出了动态注意力，这种注意力延伸到整个图像。为了保持网络学习这种模式的能力，研究人员实现了注意力矩阵的二维分解，其中网络可以通过两步稀疏注意力关注到所有位置。

第一版 strided attention 大概等同于每个位置处理自己的行和列，它与以上网络学得的注意力模式类似。（注意，列注意力可等同于处理转置矩阵的行）。第二版 fixed attention 在最新的列元素之后处理固定列和元素，研究者认为这个模式对于数据无法拟合二维结构（如文本）的情况很有用。

实验结果

Sparse Transformer 在 CIFAR-10、Enwik8 和 Imagenet 64 数据集上刷新了当前最优密度估计分数。

在 CIFAR-10、Enwik8 和 Imagenet 64 数据集上的密度估计性能（单位为 bits per byte／dim）。M 表示网络中使用的参数（单位为百万），W 表示网络宽度，L 表示层数，H 表示头数。

研究者还发现稀疏注意力比完整注意力的损失更低，且速度更快。这可能指向稀疏模式产生的有用归纳偏置，或者密集注意力的底层优化问题。

生成图像

使用了稀疏注意力的 Transformer 似乎有一种全局结构的概念，这可以通过观察图像补全（image completion）进行定性评估。下图可视化了一个在 64×64 ImageNet 上训练的模型：

损坏原图

修复图像

真实图像

研究人员还生成了完全无条件的样本，其中未调整的 softmax 温度为 1.0。这些模型使用最大似然目标进行训练，其覆盖了所有的数据模式（其中包括可能不存在的数据），而不是增强较小部分数据的保真度。从具有未调整温度的模型中取样，研究人员看到了该模型认为世界上存在的图像的完整分布。因此，一些样本看起来奇奇怪怪的。

模型示例

生成原始音频波形

通过简单改变位置嵌入，稀疏 Transformer 还能用来生成原始音频，而非图像。随着深度学习扩展到新的数据类型，用这类网络来指定归纳偏置也很容易。

该模型是在原始的古典音乐片段上训练的，并使用了稀疏注意力来生成长度为 65000 的序列。这相当于大约 5 秒长的原始音频，研究人员在下面的每个片段中将几个样本连接在一起。

代码公布

通常，实现稀疏注意力需要将查询和关键矩阵分割成块，因此为了简化实验，OpenAI 实现了一组块稀疏核，这些核在 GPU 上高效地执行这些操作。OpenAI 开源了这些核并提供了稀疏注意力函数的示例：

https://github.com/openai/sparse_attention

未来发展和限制

本文介绍的稀疏注意力模式只是对长序列进行高效建模的初步尝试。研究人员认为，探索稀疏注意力的不同模式和各种组合非常有用，而且学习稀疏模式对下一代神经网络架构来说也是一个很重要的研究途径。

即使有了上述改进，自回归序列生成对非常高分辨率图像和音频来说仍是不切实际的。但是，研究人员介绍的优化注意力操作可能有用，将它与其它方法（如多尺度方法）结合，可以建模高维数据。

论文：Generating Long Sequences with Sparse Transformers

论文链接：https://d4mucfpksywv.cloudfront.net/Sparse_Transformer/sparse_transformers.pdf

摘要：Transformer 是一种强大的序列模型，但是它所需的时间和内存会随着序列长度出现二阶增长。这篇论文介绍了注意力矩阵的稀疏因式分解，可以将其降低到 O(N√N)。该研究提出了 a）训练更深网络的架构和初始化变体；b）重新计算注意力矩阵以节省内存；c）用于训练的快速注意力内核。研究者将具备这些变化的网络称为 Sparse Transformer，并证明该网络可以使用数百个层来建模成千上万个时间步长的序列。

该网络在从原始字节中建模图像、音频和文本时使用的是同样的架构，在 Enwik8、CIFAR10 和 ImageNet-64 数据集上取得了当前最优的密度估计性能。研究者生成的无条件样本展示了全局一致性和极大的多样性，并证明原则上可以使用自注意力建模长度超百万的序列。