发布人：Avinava Dubey，Google Research 研究员
基于 Transformer 的自然语言处理 (NLP) 模型，如 BERT、RoBERTa、T5 或 GPT3，能够胜任多种任务，已成现代 NLP 研究的主流。Transformer 的通用性和稳健性使其能够适应各种基于序列的任务，也是其被广泛应用的主要因素：作为 seq2seq 模型用于 翻译、摘要、生成等，或者作为独立编码器用于情感分析、POS 标记、机器阅读理解等。Transformer 的关键创新在于引入自注意力 (self-attention) 机制，得以计算输入序列中所有位置对的相似度得分并对输入序列中的每个标记进行并行评估，避免循环神经网络的顺序依赖关系，使 Transformer 远优于先前的 LSTM 等序列模型。
然而，现有的 Transformer 模型及其衍生模型的局限性在于，完整自注意力的计算和内存要求与输入序列长度成平方关系。在当前常见的硬件和模型规模下，这通常会将输入序列限制为大约 512 个标记，这使得 Transformer 无法直接应用于需要更大上下文的任务，如 问答、文档摘要或 基因组片段分类。随之而来的是两个下意识的问题：1) 我们能在计算性能和内存大小的限制下通过稀疏模型以线性增长处理序列长度的增长，以实现 Transformer 原本在平方增长下的先验优势？2) 能否从理论上证明这些线性 Transformer 保留了 Transformer 在完整二次方下的完整表现力和灵活性？
我们的最近几篇论文都提到了这两个问题。在 EMNLP 2020 会议上发表的 “在 Transformers 中对长输入和结构化输入进行编码（ETC: Encoding Long and Structured Inputs in Transformer）”中，我们介绍了 Extended Transformer Construction (ETC)，这是一种基于稀疏 Attention 的新方法，可通过结构信息限制相似度得分的计算样本对数量。这就将对输入长度所需的计算量从平方降低为线性，并在 NLP 领域取得了大量的成果。然后，在 NeurIPS 2020 会议 上发表的“Big Bird：处理长序列的 Transformer（Big Bird: Transformer for Longer Sequence s）”中，我们介绍了另一种名为 BigBird 的稀疏 Attention 方法，将 ETC 扩展到更通用的场景。之前在这些场景中，关于源数据中存在的结构化的先验知识无法被有效利用。此外，我们还证明，从理论上讲，我们提出的稀疏 Attention 机制保留了二次方下完整 Transformer 的表现力和灵活性。我们提出的方法在挑战性长序列任务上达到了顶尖水平，包括问答、文档摘要和 基因组片段分类。
注意力即计算图
Transformer 模型中使用的 Attention 模块为输入序列中所有位置对计算相似度得分。实用的做法是将 Attention 机制视为有向计算图，用节点表示标记，用一条边表示一对标记之间算得的相似度得分。由此，完整 Attention 模型是一个完整的计算图。方法背后的核心思想是精心设计稀疏计算图，从而只计算相似度得分的线性数字。
完整 Attention 可被视为完整的计算图
Extended Transformer Construction (ETC)
对于需要较长结构化输入的 NLP 任务，我们提出了一种结构化稀疏 Attention 机制，即 Extended Transformer Construction (ETC)。为了实现 self-Attention 的结构化稀疏化，我们开发了全局 - 局部 Attention 机制。Transformer 的输入在这里被分成两部分：全局输入，其中标记的 Attention 不受限制；长输入，其中标记只能 Attention 全局输入或局部近邻。这实现了 Attention 的线性扩展，使 ETC 可以显著扩展输入长度。
为了进一步利用长文档的结构，ETC 还结合了其他想法：以相对的方式表示标记的位置信息，而不是使用序列中的绝对位置；在 BERT 等模型中使用的常规掩码语言模型 (MLM) 之外，还使用额外的训练目标；以及标记的灵活掩码，控制哪些标记可以 Attention 其他标记。例如，给定一个长的文本选择，对每个句子应用全局标记，连接到句子内的所有标记，同时对每个段落也应用全局标记，连接到同一段落内的所有标记。
基于文档结构的 ETC 模型稀疏 Attention 的示例。全局变量用 C（蓝色）表示段落，S（黄色）表示句子，局部变量用 X（灰色）表示对应长输入的标记
通过这种方法，我们在需要长或结构化输入的五个挑战性 NLP 数据集中获得了最先进的结果：TriviaQA、自然问题 (NQ)、 HotpotQA、WikiHop 和 OpenKP。
问答的测试集结果。对于经过验证的 TriviaQA 和 WikiHop，使用 ETC 取得了最优结果
BigBird
我们在 ETC 的基础上提出了 BigBird- 这种稀疏 Attention 机制的标记数量也呈线性，并且可以替代 Transformer 中使用的 Attention 机制。与 ETC 相比，BigBird 不需要任何有关源数据中存在的结构的先决知识。BigBird 模型中的稀疏 Attention 主要由三部分组成：
- 一组 Attention 输入序列的所有部分的全局标记
- 所有 Attention 一组局部相邻标记的标记
- 所有 Attention 一组随机标记的标记
BigBird 稀疏 Attention 可以看作是在 Watts-Strogatz 计算图上添加很少的全局标记
在 BigBird 论文中，我们解释了为什么稀疏 Attention 足以逼近二次方 Attention，部分说明了 ETC 成功的原因。一个关键的观察结果是，在计算的相似度得分与不同节点的信息流之间存在着内在张力（即标记相互影响的能力）。全局标记充当信息流的管道，我们证明了具有全局标记的稀疏 Attention 机制可以与完整 Attention 模型一样强大。我们还特别证明了 BigBird 与原始的 Transformer 一样具有表现力，在计算上通用（Yun 等人和 Perez 等人），并且是连续函数的通用逼近器。此外，我们的证据表明，使用随机计算图可以进一步简化信息流 - 促进使用随机 Attention 组件。
对于结构化和非结构化任务，这种设计都可扩展到更长的序列长度。通过使用梯度检查点，以序列长度换取训练时间，可以实现进一步扩展。这让我们的高效稀疏 Transformer 覆盖了需要编码器和解码器的生成任务，例如长文档摘要，并在此基础上取得新的先进成果。
长文档摘要 ROUGE 得分。我们对 BigPatent 和 ArXiv 数据集都实现了新的先进结果
此外，由于 BigBird 是通用的替代方案，这也使它可以在没有预先存在的域知识的情况下扩展到新的域。特别是，我们介绍了基于 Transformer 模型的新型应用，其中长上下文是有益的 - 提取基因组序列 (DNA) 的上下文表示。借助较长的掩码语言模型预训练，BigBird 可以在下游任务上实现最先进的性能，如启动子区域预测 和 染色质分布预测。
在多个基因组学任务上，如启动子区域预测 (PRP)，染色质分布预测（包括转录因子 (TF)，组蛋白标记 (HM) 和 DNase I 超敏性 (DHS) 检测），我们的结果均优于基线。此外，我们的结果表明，Transformer 模型可以应用于目前尚未充分开发的多个基因组学任务。
主要实现思路
大规模采用稀疏 Attention 的主要障碍之一是，在现代硬件中，稀疏运算效率很低。在 ETC 和 BigBird 的背后，我们的一个重要创新是对稀疏 Attention 机制的高效实现。由于 GPU 和 TPU 等现代硬件加速器擅长使用合并式内存运算，一次加载连续的字节块，因此由滑动窗口（局部 Attention）或随机元素查询（随机 Attention）引起的少量零星查询效率较低。相反，我们将稀疏局部 Attention 和随机 Attention 转化为密集的张量运算，充分利用现代单指令多数据 (SIMD) 硬件的优势。
为此，我们首先“块化” Attention 机制，从而更好地利用将在模块上运行的 GPU/TPU。然后，通过一系列简单的矩阵运算，如重塑、滚动和聚集，将稀疏 Attention 机制计算转换为密集的张量积，如以下动画所示。
如何利用滚动和重塑高效计算稀疏窗口 Attention，并且不使用较小的零星查找
最近，“Long Range Arena:A Benchmark for Efficient Transformer”提供了六个需要较长上下文的任务的基准，并对所有现有的长距离 Transformer 进行了基准测试。结果表明，BigBird 模型与其同类不同，可在不牺牲性能的前提下，明显降低内存消耗。
结论
我们的结果表明，精心设计的稀疏 Attention 可以像原始的完整 Attention 模型一样具有足够的表现力和灵活性。除了理论上的保证，我们还提供了一个非常高效的实现，能够扩展到更长的输入。由此，我们在问答、文档摘要和基因组片段分类方面取得了最先进的结果。考虑到我们稀疏 Attention 的通用性，该方法应该适用于多种其他任务，如程序合成 和 长格式的开放域问答。ETC (GitHub) 和 BigBird (GitHub) 的代码均已开源，都能在 GPU 和 TPU 的长序列上高效运行。
致谢
这项研究是与 Amr Ahmed、Joshua Ainslie、Chris Alberti、Vaclav Cvicek、Avinava Dubey、Zachary Fisher、Guru Guruganesh、Santiago Ontañón、Philip Pham、Anirudh Ravula、Sumit Sanghai、Qifan Wang、Li Yang、Manzil Zaheer 合作的结果，这些人员也是 EMNLP 和 NeurIPS 论文的合著者。
原文：Constructing Transformers For Longer Sequences with Sparse Attention Methods
中文：TensorFlow 公众号