选自 openai, 机器之心编译,机器之心编辑部。
Transformer是一种强大的序列模型,但是它所需的时间和内存会随着序列长度出现二阶增长。近日,OpenAI研究人员开发出了一种深度神经网络Sparse Transformer,该网络在预测长序列方面创造了新纪录——无论预测的是文本、图像还是声音。该神经网络利用注意力机制中的一种改进算法,可以从长度可能是之前30倍的序列中提取模式。
现在,AI 研究中的一项挑战是在图像、视频或声音等复杂数据中进行长序列的精细相关性建模。Sparse Transformer 合并了 O(N^2)Transformer 自注意力机制的 O(N√N) 重组以及其他一些改进,从而直接用于这些丰富的数据类型。以前,这些数据上所使用的模型是专为某个领域制作的,或者很难将序列扩展到包含几千个元素。
相比之下,OpenAI 开发的模型通过使用数以百计的层可以对包含上万个元素的序列进行建模,在诸多领域都取得了当前最佳的表现。OpenAI 研究人员利用该模型帮助创建能够更好地理解世界的 AI 系统。
深度注意力
在 Transformer 中,每一个输出元素与输入元素相连接,同时根据具体情况对它们之间的权重进行动态计算,这一过程被称为「 注意力机制 」。虽然人们相信这使得 Transformer 较那些具有固定连接模式的模型更为灵活,但实际操作中需要为每一层和注意力头创建一个 N×N 注意力矩阵,当应用于图像或原始音频等具有许多元素的数据类型时会消耗大量内存。
减少内存消耗的一种方法是在 反向传播 过程中从检查点处重新计算注意力矩阵,这是深度学习中的一种成熟的方法,以更多的计算来减少内存使用。
当 Transformer 中的注意力矩阵完成时,这意味着最大的内存消耗将不受层数的支配,使研究人员训练网络的深度大大超过从前。在实际操作中,研究人员发现在处理 CIFAR-10 等基准测试任务时,深度达 128 层的 Transformer 表现出的性能优于较浅的网络。
为了训练深度更大的模型,研究人员对 transformer 的操作顺序进行了几次调整,修改了初始化方法。详情参见论文。
稀疏注意力
然而,对于非常大的输入来说,甚至计算单个注意力矩阵都是不现实的。因此,OpenAI 使用了稀疏注意力模式,在这种模式中,每个输出位置仅从输入位置子集中计算权重。当子集相对于整个输入集较小时(如元素数量是√N 而不是 N),即使对于非常长的序列,注意力计算也会变得比较容易,算法复杂度为 O(N√N)而不是 O(N^2)。
为了评估该方法的可行性,研究人员首先可视化并学习了图像上深度 Transformer 的注意力模式,发现其中许多模式表现出了可解释和结构化的稀疏模式。以下每幅图像都显示了哪个输入像素(白色高亮标出)由一个给定的注意力头处理,以预测图像中的下一个值。当输入部分集中在小的子集上并显示出高度规律性时,该层就易于稀疏化。以下是 CIFAR-10 图像上 128 层模型的样本:
虽然许多层显示出稀疏的结构,但有些层清晰地显示出了动态注意力,这种注意力延伸到整个图像。为了保持网络学习这种模式的能力,研究人员实现了注意力矩阵的二维分解,其中网络可以通过两步稀疏注意力关注到所有位置。
第一版 strided attention 大概等同于每个位置处理自己的行和列,它与以上网络学得的注意力模式类似。(注意,列注意力可等同于处理转置矩阵的行)。第二版 fixed attention 在最新的列元素之后处理固定列和元素,研究者认为这个模式对于数据无法拟合二维结构(如文本)的情况很有用。
实验结果
Sparse Transformer 在 CIFAR-10、Enwik8 和 Imagenet 64 数据集上刷新了当前最优密度估计分数。
研究者还发现稀疏注意力比完整注意力的损失更低,且速度更快。这可能指向稀疏模式产生的有用归纳偏置,或者密集注意力的底层优化问题。
生成图像
使用了稀疏注意力的 Transformer 似乎有一种全局结构的概念,这可以通过观察图像补全(image completion)进行定性评估。下图可视化了一个在 64×64 ImageNet 上训练的模型:
研究人员还生成了完全无条件的样本,其中未调整的 softmax 温度为 1.0。这些模型使用最大似然目标进行训练,其覆盖了所有的数据模式(其中包括可能不存在的数据),而不是增强较小部分数据的保真度。从具有未调整温度的模型中取样,研究人员看到了该模型认为世界上存在的图像的完整分布。因此,一些样本看起来奇奇怪怪的。
生成原始音频波形
通过简单改变位置嵌入,稀疏 Transformer 还能用来生成原始音频,而非图像。随着深度学习扩展到新的数据类型,用这类网络来指定归纳偏置也很容易。
该模型是在原始的古典音乐片段上训练的,并使用了稀疏注意力来生成长度为 65000 的序列。这相当于大约 5 秒长的原始音频,研究人员在下面的每个片段中将几个样本连接在一起。