现有人工智能研究的一个挑战是对复杂数据（如图像，视频或声音）进行长期、微妙的相互依赖建模。以前，在这些数据上使用的模型是针对一个域专门设计的，难以扩展到超过几千个元素的序列。

相比之下，新模型可以使用数百个层对具有数万个元素的序列进行建模，从而在多个域中实现最先进的性能。OpenAI正在使用它来帮助构建具有更强的理解世界能力的AI系统。

OpenAI开发了一种深度神经网络Sparse Transformer，可以预测序列中接下来的内容，无论是文本，图像还是声音。它使用改进的注意机制算法，可以从比以前长30倍的序列中提取模式。

深度注意力

在Transformer中，每个输出元素都连接到每个输入元素，它们之间的权重是根据环境动态计算的，这个过程即为注意力。虽然这允许Transformer比具有固定连接模式的模型更灵活，但在实践中它需要为每一层和注意力头创建一个N×N注意力矩阵，当应用于具有许多元素的数据类型（如图像或原始音频）时，会占用大量内存。



当矩阵存储在内存中或在后向传递期间重新计算时，使用深度Transformer占用注意力内存（64层，4个注意力头）。作为参考，用于深度学习的标准GPU通常配置为12-32GB的内存。

减少这种情况的一种方法是在反向传播期间从检查点重新计算注意力矩阵，这是深度学习中的一种成熟技术，以更多计算为代价来减少内存使用。当完成注意力矩阵时，这意味着最大的内存成本与层数无关，让我们比以前更深入地训练网络。在实践中，我们发现深度高达128层的Transformer在基准任务（如CIFAR-10）上的表现优于较浅的网络。

为了更深入地训练这些模型，团队对Transformer中的操作顺序进行了几次调整，并修改了初始化方案。完整的细节可以在论文中看到。

Sparse注意力

然而，对于非常大的输入，即使计算单个注意矩阵也可能变得不切实际。相反，我们使用稀疏注意模式，其中每个输出位置仅从输入位置的子集计算权重。当子集相对于整个输入集较小时，由此产生的注意力计算即使对于很长的序列也变得易于处理。

为了评估该方法的可行性，首先将深度Transformer在图像上的学习注意模式可视化，发现许多模型表现出可解释和结构化的稀疏模式。下面的每个图像显示给定的注意头处理哪些输入像素（以白色突出显示）以便预测图像中的下一个值。当输入部分聚焦在小的子集上并显示出高度的规则性时，该层使易于稀疏化的。对于CIFAR-10图像上的128层模型，下面显示了它们的样本：

第19层

第20层

学习的128层CIFAR-10网络的多个层的注意模式（白色高亮）。这些层学会了将注意力分散在两个维度上。第19层总结了每一行的信息，第20层按列聚合这些摘要，从而使全面注意力操作有效分解。

第6层

第36层

一些层学会了访问位置存储器，无论输入数据或时间步长如何，通常都会访问类似的位置（第6层）。其他层学习了高度依赖数据的访问模式（第36层）。

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



第一个版本，大步幅注意力，大致相当于关注该行和列的每个位置，与上面网络学习的注意力模式类似（注意，列注意力可以等效地表示为关注转置矩阵的行）。

第二个版本，固定注意力，注意固定列和最新列元素之后的元素，这种模式在数据不适合二维结构（如文本）时很有用。

实验结果

Sparse Transformer为CIFAR-10，Enwik8和Imagenet 64的密度估计创造了最新记录。





在各种基准数据集上，密度建模性能以比特/字节或dim为单位。M表示网络中使用的数百万个参数，W表示网络的宽度，L表示层数，H表示头数。

除了明显变得更快之外，稀疏注意力的损失低于完全注意力。这可能表明稀疏模式存在有用的归纳偏差，或者是密集关注的潜在优化问题。

生成图片

使用稀疏注意力的Transformer似乎有一个全局结构的概念，可以通过查看图像完成来定性评估。在这里，对64×64 ImageNet上训练的模型进行可视化：

Prompt

Completions

Ground truth

团队还用未经调整的softmax temperature of 1.0生成了完全无条件的样品，这些模型使用最大似然目标进行训练，众所周知，目标是覆盖所有数据模式（包括可能不存在的数据），而不是增加较小部分数据的保真度。从这些具有未调整温度的模型中采样，可以看到模型认为存在于世界中的图像的完整分布。结果，一些样本看起来很奇怪。

生成原始音频波形

Sparse Transformers也可以通过简单地改变位置嵌入以适应地生成原始音频而不是图像。随着深度学习扩展到新型数据类型，这类网络易于指定归纳偏差，将是很实用的工具。

该模型用原始古典音乐剪辑训练，并使用稀疏注意力生成长度为65000的序列。这相当于5秒的原始音频。

代码开源

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

开源：github.com/openai/sparse_attention

未来的工作和局限

我们引入的稀疏注意力模式只是长序列有效建模方向的第一步。我们认为探索稀疏性的不同模式和组合是有用的，并且学习稀疏模式对于下一代神经网络体系结构来说是一种非常有前途的研究方向。

即使做了上述改进，自回归序列生成对于非常高分辨率的图像或视频仍然是不切实际的。然而，我们引入的优化注意力操作可能是有用的一环，可以与其他建模高维数据的方法结合使用。

博客：
openai.com/blog/sparse-transformer/

论文：
d4mucfpksywv.cloudfront.net/Sparse_Transformer/sparse_transformers.pdf