上下文|Transformers 长程上下文综述 |综述|复杂度|

本文插图
【上下文|Transformers 长程上下文综述】

字幕组双语原文：Transformers 长程上下文综述
英语原文：A Survey of Long-Term Context in Transformers
翻译：雷锋字幕组（wiige）
多头自注意力开销很大，这已经不是什么秘密了——序列长度的O(n2)复杂度意味着不太可能让经典的Transformer模型处理超长文本序列。在过去的两年里， NLP社区已经开发出了名副其实的缝合怪方法来应付这种复杂性，但这篇文章将重点介绍在大规模下更有前景的方法。

稀疏Transformers
自适应Transformers
Transformer-XL
压缩Transformers
Reformer
寻觅式 Transformer

密集多头注意力的时间和空间复杂度
多头注意力对序列长度的伸缩性很差，原因有二: 首先是计算注意力矩阵所需的FLOPs按序列长度的平方暴增，导致对单个序列的自注意力操作的计算复杂度为O(hdn2) ，其中h是注意力头数， d是key和query的维度， n是文本序列的长度。第二点是，点积自注意力操作的空间复杂度也按序列长度的平方增大。计算注意力矩阵的空间复杂度为O(hdn+hn2)——第一项是存储key和query所需的内存，第二项是每个头得出的注意力标量值。

让我们向BERT-Base代入具体的数字来了解哪些项的开销占大头。BERT-Base使用的序列长度为512 ，隐藏层大小为768 ， 12个头，这意味着每个头的维度为64(768 / 12) 。在这种设置下，需要393216 floats(约1.5MB)(头12个*头大小64*序列长度512)来存储键和值，而存储所有头的注意力标量所需的空间为3145,728 floats (12*512*512)即约12MB的内存——几乎10倍于键存储的空间开销，这还仅仅是仅仅512个上下文词符。
由于在训练过程中必须缓存激活，以便进行梯度计算（除非使用梯度检查点这种激活再计算策略），因此仅存储这些12层的BERT-Base注意力矩阵，每个样本就要消耗约150MB的内存。当序列长度为1024时，这就上升为约600MB ，而序列长度为2048时，每个样本仅存储注意力矩阵就已经消耗了约2.4GB的内存。这意味着训练时的批处理尺寸(batch size)更小，并行性更差，进一步阻碍了模型利用长距离上下文的能力。
稀疏Transformers
"Generating Long Sequences with Sparse Transformers"由 Rewon Child, Scott Gray, Alec Radford 和 Ilya Sutskever 通过离散化(factorization)方法解决了自注意力的时间和空间复杂度中的O(n2)项问题。
离散注意力
在一个典型的自注意力操作中，输入序列中的每一项都会关注到输入序列中的所有其他项，从而形成如下的注意力模式:

本文插图

典型的自回归设置下的自注意力连接模式。深蓝色方块代表 "查询(query)" ，浅蓝色方块代表 "键(key)"
经典自注意力带来的益处是，其高连通性使得信息在词符之间很容易传递——只需要一层注意力就可以聚合任意两个词符的信息。但如果我们放宽这一约束，保证两层内信息可以在任意两个词符之间传递即可的话，就可以大幅降低按序列长度增长的复杂度。稀疏Transformer通过利用固定注意力模式的自定义核来实现这一目标。

本文插图

固定的稀疏Transformer 。深蓝色方块代表查询，中浅蓝色方块代表奇数层注意力键索引，最浅蓝色的方块代表偶数层关注的关键索引。
一半的头只关注短的局部上下文，而另一半则预分配均匀分布在整个序列中的索引。
通过这些聚合索引来传递信息，网络就能够从较远的词符获取信息并利用长距离上下文，同时将时间和空间复杂度降低到O(n√n) 。最重要的一点是，只需要两层隐藏层，任意词符就都可以整合来自其他词符的信息。
稀疏Transformers的实验结果
重要的是，离散化注意力结构似乎不会对语言模型的性能产生负面影响，它的每字符比特数比enwiki8上的密集注意力(令人惊讶地)更高，且可以在高达12228个词符的上下文上产生有效的注意力。

本文插图

可以看到，稀疏Transformer的注意结构之所以有效，部分原因是因为这种注意力模式与真正习得的密集注意力并没有什么不同。Kevin Clark, Urvashi Khandelwal, Omer Levy, 和 Christopher D. Manning在他们的 "What Does BERT Look At? An Analysis of BERT’s Attention"一文中探究了密集注意力所习得的模式，试图弄明白注意力在Transformer模型中承载了什么功能。他们发现注意力头倾向于关注紧接在前的词符（类似于稀疏注意力中的局部注意力模式），以及如[SEP]和句号的特定词符。所以，也许稀疏Transformer的注意力模式中包含的归纳偏差是有用而非有害的。

本文插图

例子: BERT习得的注意力模式
如果想在自己的项目中采用固定注意力核，可以查看OpenAI的blockparse库以及作者发布的配套示例。
自适应窗口Transformers

Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski,和 Armand Joulin在他们的文章 "Adaptive Attention Span in Transformers" 中对复杂性问题采取了不同的方法.他们与 "What Does Bert Look At?" 的作者得到了同样的观察，并指出虽然密集注意力允许每个头关注完整的上下文，但许多注意力头只关注局部上下文，而剩下的头才会关注整个文本序列。他们建议使用某种自注意力来利用这一现象，以此让模型可以选择其上下文尺寸。
自适应遮罩(Adaptive Masking)自适应窗口Transformer通过对序列进行遮罩，使每头习得的上下文之外的词符贡献迅速消失。遮罩(M)与softmax操作的对数相乘，使某些词符对当前隐状态x的贡献归零，其中超参数R控制最小窗口(span)大小。

本文插图

"Adaptive Attention Span in Transformers"所用的软遮罩函数。图片来自论文所附的Facebook AI博文
为鼓励模型只在有益的地方使用额外上下文，对习得的z值施加了?1惩罚。
注意力工作细探与实验结果在这些限制条件下，大部分的头都会选择关注&lt100个字符的上下文，只有少数的头（主要是在网络的后几层）会选择顶着?1?1的惩罚来关注&gt1000个字符的上下文。

本文插图

除了巧妙的缓存机制，这种对长上下文的惩罚使得自适应窗口高达8k个字符，同时仍然保持模型的整体计算成本可控。此外，它在基线上的性能仍然很高——在enwiki8上达到了0.98比特/字符， text8数据集上达到了1.07比特/字符。
然而，可变窗口大小在便于并行方面并不理想，我们通常希望密集的、尺寸规整的矩阵来实现最佳性能。虽然这种方法可以大幅减少预测时前向计算所需的flops数量，但作者只提供了模糊的性能估计，并表示自适应跨度的实现能够以2048个上下文词符的固定上下文大小模型的速度最高处理8192个词符。
Facebook AI Research还开源了他们的工作——代码和预训练的模型可以在github.com/facebookresearch/adaptive-spans上获得。
Transformer-XL
比起让密集注意力操作的成本更低廉， Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le和 Ruslan Salakhutdinov 则从RNNs中获得灵感，并在Transformer中引入了除自关注机制外的循环机制。他们的文章"Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" ，引入了两个新颖的概念——第一个是将之前 "段(segments)"的隐状态作为输入反馈给当前段层的组件，另一个是促进这一策略的相对位置编码方案。
段循环
使用具有固定上下文大小的标准变换器，处理长输入时需要将输入分割成块（或称段）分别处理。然而，这种方法有一个局限，即前一段的信息与当前段的词符无关。这种段独立性在某种程度上是有益的，因为它让我们能分批地有效处理每一段，但如果目标保持是长序列的一致性，这就会成为一个很大的限制。

本文插图

带密集注意力的自回归Transformer的词符注意力结构来自Transformer-XL 。
Transformer-XL通过强制进行分段串联来克服这一限制。在第一段之后，后续段中的词符将始终保持512个词符的上下文尺寸，之前段的激活作为上下文传递给后续段的注意力操作。这意味着来自NN上下文大小*LL层外的信息可以传播到一个给定的令牌。假设上下文尺寸为640 ，模型有16层，理论上Transformer-XL可以纳入10,240个词符的信息。

本文插图

Transformer-XL原文中的词符注意力模式。
为了避免存储所有段的激活，作者阻止了梯度流向之前的段。
纳入相对位置
Transformer-XL还引入了一种新颖的位置编码方案，称为 "相对位置编码" 。不是简单地将内容和绝对位置嵌入的和输入网络，而是将每层的注意力操作分解为基于内容的部分和基于相对位置的部分——如果一个段中的第512个词符要关注第511个词符，它的相对位置嵌入就是-1 。
为了使相对位置编码的使用变得简单，他们将从键和查询中产生注意力权重的操作进行了分解。对于一个典型的密集注意力操作， pre-softmax注意力权重可以如下分解:

本文插图

上式中， Exi是词符在位置i的基于内容的嵌入， Uj是词符j的位置嵌入。
(a) 将查询的内容与键的内容相关联
(b) 将查询的内容与键的位置相关联
(c) 将查询的位置与键的内容相关联
(d) 将查询的位置与键的位置相关联
当使用相对位置嵌入时，作者将式子改成如下形式:

本文插图

在b)和d)中， Uj已经被它的相对位置嵌入Ri-j替换.
对于包含查询位置的项，我们将矩阵UiUi替换为两个新学习参数uu和vv 。这些向量可以理解为两个不依赖于查询具体细节的偏置——cc鼓励多关注一些项， dd鼓励多关注相对位置。作出这种替换是因为对自身的查询其相对位置保持不变。
注意力工作细探与实验结果
要想让Transformer-XL模型利用这样的长程上下文，每一层至少有一个头要利用其注意力跨度的全部上下文。平均注意力权重图显示，每一层都有头主要关注先前的位置。

本文插图

Transformer-XL论文中的平均注意力权重图
此外Transformer-XL论文还测量了有效上下文长度对困惑度的影响，并发现增大上下文长度至~900个词符时会产生更好的困惑度得分——这进一步证明了循环机制在实践中的有效性，而非仅仅停留于理论。
Kimi Young的github提供了源代码， HuggingFace的实现可以帮助你搭建自己的Transformer-XL项目。
压缩Transformers
列表上的下一个模型是Transfromer ，它建立在Transformer-XL架构的基础上，但通过压缩损失扩展了该方法，以便纳入更长的文本序列。在来自DeepMind的Jack W. Rae、Anna Potapenko、Siddhant M. Jayakumar和Timothy P. Lillicrap的文章 "Compressive Transformers for Long-Range Sequence Modelling "中，详细介绍了一种能够关注与全书长度的序列模型架构。
压缩Transformer注意力

本文插图

" Compressive Transformers For Long Range Sequence Modeling "中的压缩记忆(Compressed Memory)示意图

按照Transformer-XL的做法，当前序列可以关注已存储的前段的激活。此外，在同样的多头注意操作中，当前段中的词符可以关注存储在 "压缩记忆 "中的第二组状态。
在每一步中，最早的压缩记忆被丢弃，压缩记忆向后移动一个索引。然后，来自正常记忆段中最早的nn个状态进行压缩，然后转移到压缩记忆新空出来的槽中。
DeepMind博客上的一个gif很好地说明了这个过程:

本文插图

将过去的记忆逐渐压缩成压缩记忆——图片由 DeepMind博客提供
DeepMind团队尝试了多种压缩操作（包括平均池化、最大池化和学习卷积等等），但最终决定训练一个二级网络来重建缩记忆中基于内容的注意力矩阵。
换句话说，他们学习了一个函数fc ，通过最小化压缩记忆的注意力（C-1=fc(Mold)）和正被压缩的正常记忆状态的注意力之差，将n个最早的记忆状态压缩为一个压缩记忆状态。 :
他们没有将这种压缩操作与主语言模型联合训练，而是选择在一个单独的优化循环中更新压缩网络。因为让注意力状态容易被压缩，对降低语言模型的损失会起到反作用。
实验结果
在他们的实验中，压缩记忆的尺寸为512 ，正常记忆尺寸为512 ，窗口尺寸为512 ，压缩率为2——意味着在压缩步骤中，最早的2个记忆状态将被压缩为1个状态。在这个设置下，他们在WikiText-103上实现了17.1的SOTA困惑度。
由于利用较长序列长度的收益通常是长尾的，他们特别研究了不同词频下的困惑度，发现在最稀有的词符上的收益尤其显著:

本文插图

按词频分类的困惑度。图片由DeepMind博客提供
虽然他们还没有公开源码，但DeepMind已经开放了PG-19的源码，这是他们在研究压缩Transformer时开发的数据集。 PG-19是Project Gutenberg的衍生品，旨在进一步研究长期注意力。
Reformer
接下来是Nikita Kitaev, ?ukasz Kaiser, Anselm Levskaya 的文章"Reformer: The Efficient Transformer".Reformer在增大序列长度方面采取了不同的策略——他们没有引入循环机制或压缩记忆，而是选择通过使用局部敏感哈希(LSH)技术来缩小每个词符的注意力范围.
局部敏感哈希是一系列将高维向量映射到一组离散值（桶/簇）的方法。它常常用在近似最近邻搜索中。

Reformer的作者对注意力操作的键和查询使用了同一个投影，并使用基于随机旋转的局部敏感哈希方法将共享的键/查询分到最多几百个词符的不同桶中。下图是该方法的一个示例:

本文插图

Reformer论文中的角度化LSH 。该图展示了一个有3次旋转的4桶哈希设置。下方图中的三个向量已经映射到同一个哈希桶中，因为其的原始向量就很接近，而上方的第一个图和最后一个展示了向量被映射到不同的哈希桶中。
它们计算每个桶内的注意力矩阵，然后取对应值的加权和。由于只关注一个给定的桶内的元素，所以如果桶的大小选择得当，这可以将注意力操作的整体复杂度从O(n2)降低到O(nlogn) 。由于分桶是随机的，而且旋转映射向量也是随机的，所以他们计算了多个哈希值，以确保具有相似的共享键-查询嵌入的词符最终以高概率落入同一个桶中。

本文插图

他们还采用了 " The Reversible Residual Network: Backpropagation Without Storing Activations"中介绍的技术来控制训练时的空间消耗。可逆残差层使用巧妙的架构结构，允许从层输出中轻松重构层输入，并在网络深度不变的情况下以空间复杂度降低了计算量。

通过局部敏感的哈希技巧来降低计算成本，以及可逆残差来降低内存消耗， Reformer架构能够在单个加速器上处理长达64000个词符的序列。
虽然在enwiki-8上得出的1.05比特/字符的得分落后于我们在这篇文章中所研究的其他一些模型，但Reformer是一个令人耳目一新的独特机制，我很期待看到这种方法如何继续扩大规模。
如果你有兴趣进一步探索Reformer架构，可以看看我最近关于这个主题的博文 "A Deep Dive into the Reformer" 。在google/jax Github仓库中，有一个Reformer的开源实现示例。Phil Wang也维护了一个PyTorch版本。
路径 Transformer
我原打算在这里就结束本文，但在Aran Komatsauzaki的建议下，我加入了最后一篇论文。 Aurko Roy、Mohammad Taghi Saffar、David Grangier和Ashish Vaswani的第二篇ICLR 2020的论文 "Efficient Content-Based Sparse Attention with Routing Transformers "与前述的Reformer有一些异曲同工之处。他们将问题框定为路径问题，旨在学习选择稀疏词符聚类簇Si和内容x的函数。
作者的在下图中说明了他们的方法。他们不是只关注局部元素或第n个元素来增加稀疏性，而是学习了（图c中用彩色表示的）需要关注的聚类。重要的是这些簇是每个键和查询的内容的函数，而不仅仅与绝对或相对位置相关。

本文插图

Routing Attention与Local和Strided Attention的比较，来自Routing Transformer论文
Routing Attention在确保每个键和查询向量都单位化后，他们使用尺寸为（Dk,Dk）的随机正交权重的共享矩阵来投影键和查询的值，其中Dk是键和查询的隐藏维度。

本文插图

然后，根据一组k-means中心将R中的向量分组到k个簇中，每个k-means中心通过每批使用k-means更新来习得，与梯度下降过程无关。
在一个给定的聚类簇Ci内，它们使用标准的加权求和方法计算了一组新的上下文嵌入，其中每个注意力值Ai都是使用标准的点积自注意力计算的。
由于密集注意力中的注意力模式通常由少数关键元素主导，而且聚类分配过程需要将具有高注意力权重的键和查询归入同一个聚类中，作者认为这个机制能够有效告知我们X′i进行了高开销的密集操作。
最后，他们选择了一些接近√n的聚类簇，这样基于稀疏内容的注意力机制的整体复杂度就变成了O(n√n) 。为了使整个过程易于并行化和统一矩阵尺寸，作者最接近每个中心点的top-k个项来代替真正的k-means聚类。

除了基于内容的路径注意力，路径注意力Transformer还在尺寸为256的上下文上执行局部注意力。
实验结果Routing Transformer在计算效率上的提升也使得其在Wikitext-103这个词级语言建模基准测试上的困惑度得到提升，它们的性能明显超过了之前描述的Transformer-XL模型。

本文插图

Wikitext-103测试集上的困惑度结果，来自Routing Transformer论文。
在enwiki-8上， Routing Transformer的表现也相当不错，不过他们的结果略微落后于Adaptive Span Transformer 。

本文插图

enwiki-8的测试集的每个字符比特数结果。
我本来没找到Routing Transformer的实现，但Aurko Roy很好心地给我了一份他们源码的压缩包，这份源码是提交给ICLR审阅的版本。
其他在Transformer中的长程上下文方法
如果你对在Transformer中加入长程上下文的其他方法感兴趣，那你还可以读一下:

Efficient Content-Based Sparse Attention with Routing Transformers
Adaptively Sparse Transformers
BP-Transformer: Modelling Long-Range Context via Binary Partitioning
Scaling Laws for Neural Language Models

雷锋字幕组是一个由 AI 爱好者组成的翻译团队，汇聚五百多位志愿者的力量，分享最新的海外AI资讯，交流关于人工智能技术领域的行业变革与技术创新的见解。
团队成员有大数据专家、算法工程师、图像处理工程师、产品经理、产品运营、IT咨询人、在校师生；志愿者们来自IBM、AVL、Adobe、阿里、百度等知名企业，北大、清华、港大、中科院、南卡罗莱纳大学、早稻田大学等海内外高校研究所。
如果，你也是位热爱分享的AI爱好者。欢迎与雷锋字幕组一起，学习新知，分享成长。