更快更强的Transformer！Google AI发布Reformer，百万量级注意力机制

2020 年 1 月 17 日 专知

【导读】ICLR 2020 Oral论文《Reformer: The Efficient Transformer》提出了一种能利用16GB内存在百万级大小上下文上执行注意力机制的Transformer，它的性能优化是由哈希和可逆残差来实现的。

理解序列数据——如语言、音乐或视频——是一项具有挑战性的任务，特别在当它依赖于大量的上下文的情况下。例如，如果一个人或一个物体在视频中消失，很久以后又重新出现，许多模型就会忘记它的样子。在语言领域，长短时记忆(LSTM)神经网络覆盖了足够的上下文来逐句翻译。在这种情况下，上下文窗口（翻译时考虑的数据跨度）覆盖了几十到上百个词。近期的Transformer模型不仅改进了逐句翻译的性能，还可以通过多文档摘要生成整个Wikipedia文章。这样可行的原因是Transformer的上下文窗口可以覆盖上千个词。用这样的大上下文窗口，Transformer可以被应用于文本以外的数据，如像素、音符等，这样它可以被用来生成音乐、图像等。

然而，更大的上下文窗口会对Transformer造成限制。Transformer的强大来自于注意力机制，在这个过程中，它考虑上下文窗口中所有可能的单词对，以理解它们之间的联系。因此，对于包含10万个单词的文本，它每次需要评估10万 x 10万个（即100亿个）单词对，这是不切实际的。另外，每层输出的存储也是个问题。对于使用大型上下文窗口的应用，存储多个模型层的输出的内存需求很快会变得非常大（几层需要GB级内存，上千层的模型需要TB级内存）。这意味着，使用一些层的Transformer模型，只能被用于少数几段文本或生成简短的音乐片段。

本文介绍Reformer，一个可用于应对百万级大小上下文窗口的Transformer，它可以在单一的加速器上运行且仅依赖16GB的内存。它结合了两种关键技术来解决注意力和内存分配问题，这些问题限制了Transformer不能被应用于长上下文窗口。Reformer使用局部敏感哈希（LSH）来降低长序列中注意力的复杂程度，并使用可逆残差层来更高效地利用内存。

注意力机制

将Transformer模型应用于非常大的文本序列时，第一个挑战是如何处理注意力层。LSH通过计算一个哈希函数来实现这一点，该哈希函数将类似的向量匹配在一起，而不是搜索所有可能的向量对。例如，在翻译任务中，网络第一层的每个向量表示一个单词(后续层中可能表示更大的上下文)，不同语言中相同单词对应的向量可能得到相同的哈希。下图中，不同的颜色描绘了不同的哈希，相似的单词有相同的颜色。当哈希值被分配时，序列会被重组，将具有相同哈希值的元素放在一起，并被分成片段(或块)，以支持并行处理。之后，注意力机制被应用于这些更短的块中(以及相邻块，以覆盖溢出)，从而大幅度减少了计算负载。

内存问题

虽然LSH解决了注意力的问题，但仍然存在一个内存问题。单个网络层通常需要几GB的内存，并且通常适用于一个GPU，所以即使是一个处理长序列的模型的单层网络也是可以接受的。但当用梯度下降训练一个多层模型时，为了反向传播，模型需要存储每层的激活值。典型的Transformer一般有几十或更多层，如果缓存每层的值，内存会很快被消耗完。

Reformer中实现的第二个新方法是在反向传播期间按需重新计算每个层的输入，而不是将其存储在内存中。这是通过使用可逆层来实现的，其中网络最后一层的激活值被用来恢复来自任何中间层的激活值，这相当于反向运行网络。在典型的残差网络中，栈中的每一层都不断地向传播信息中添加向量。每个可逆层有两组激活值，一个遵循刚才描述的标准过程逐层更新，但另一个只捕获变化。因此，要反向运行网络，只需在每层激活值间做简单的减法。

Reformer的应用

Reformer中这两种方法的引用使得它非常高效，以至于它可以在单个16GB的加速器上处理包含100万词的长文本序列。由于Reformer非常高效，它可以直接应用于上下文窗口远大于当前最先进文本领域数据集的上下文窗口的数据。也许Reformer能够处理如此大数据集的能力将刺激社区创建它们。

图像生成任务中不缺乏长上下文数据。下面链接的colab代码中展现了Reformer如何被用于补全残缺图片。它可以基于第一行残缺的图片来像素级地生成第二行中完整的图片。colab地址为：

https://colab.research.google.com/github/google/trax/blob/master/trax/models/reformer/image_generation.ipynb

目前论文《Reformer: The Efficient Transformer》已经被ICLR 2020接收为Oral论文。更多细节可以参考原始论文：

https://arxiv.org/abs/2001.04451

完整PPT下载

请关注专知公众号（点击上方蓝色专知关注）

后台回复“REFM20” 就可以获取完整博客和论文PDF的下载链接~

REFORMER: THE EFFICIENT TRANSFORMER

作者：Nikita Kitaev、Lukasz Kaiser、Anselm Levskaya
论文链接：https://openreview.net/pdf?id=rkgNKkHtvB

摘要： 大型的 Transformer 往往可以在许多任务上实现 sota，但训练这些模型的成本很高，尤其是在序列较长的时候。在 ICLR 的入选论文中，我们发现了一篇由谷歌和伯克利研究者发表的优质论文。文章介绍了两种提高 Transformer 效率的技术，最终的 Reformer 模型和 Transformer 模型在性能上表现相似，并且在长序列中拥有更高的存储效率和更快的速度。论文最终获得了「8，8，6」的高分。在最开始，文章提出了将点乘注意力（dot-product attention）替换为一个使用局部敏感哈希（locality-sensitive hashing）的点乘注意力，将复杂度从 O(L2 ) 变为 O(L log L)，此处 L 指序列的长度。此外，研究者使用可逆残差（reversible residual layers）代替标准残差（standard residuals），这使得存储在训练过程中仅激活一次，而不是 n 次（此处 n 指层数）。最终的 Reformer 模型和 Transformer 模型在性能上表现相同，同时在长序列中拥有更高的存储效率和更快的速度。