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

选自 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-十、Enwik8 和 Imagenet 64 数据集上刷新了当前最优密度估计分数。

在 CIFAR-十、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 开源了这些核并提供了稀疏注意力函数的示例：github.com/openai/spar…

将来发展和限制

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

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

论文：Generating Long Sequences with Sparse Transformers

论文连接：d4mucfpksywv.cloudfront.net/Sparse_Tran…

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

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

参考连接： openai.com/blog/sparse…