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

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 数据集上刷新了当前最优密度估计分数。