Reformer:高效的Transformer

导读：Transformer模型已经在许多任务上取得了令人欣喜的成绩，但是当扩展到更长的上下文窗口时会遇到许多限制。更大的上下文窗口意味着模型能力变得更加强大，但也会让其变得不够高效，消耗更多的内存。来自Google AI团队最新的Reformer模型结合了两个至关重要的技术来解决限制Transformer应用到长上下文窗口的注意力和内存分配的问题。

作者：Nikita Kitaev, Lukasz Kaiser
编译：DestinedAI

理解序列数据——如语言，音乐或视频——是一项非常具有挑战性的任务，尤其当有大量依赖前后上下文的情况下。比如，如果视频中的一个人或一个物体从视野中消失了，很久以后又重新出现了，许多模型会忘记它是什么样子。在语言领域，长短程记忆（LSTM）神经网络能够覆盖到足够多的上下文信息来逐句进行翻译[1]。在这种情况下，上下文窗口（如翻译时需要考虑的数据范围）从十几个到上百个词不等。最新的Transformer模型[2]不仅仅提升了逐句翻译的性能，而且能够被用来通过多文档摘要生成整个Wikipedia文章[3]。这是可能的，因为Transformer使用的上下文窗口可能扩展到上千个词。有了这么大的上下文窗口，Transformer能够被用于文本之外的应用，包括像素或音符，使得它能够生成音乐[4]和图片[5]。

但是，把Transformer扩展到更大的上下文窗口会遇到许多限制。Transformer的能力来自于注意力，它通过这个过程考虑到窗口大小内所有可能的单词对，以理解它们之间的联系。所以，在一个包含100K个单词的文本中，这需要评估100Kx100K个单词对，或每一步100亿个单词对，这是不切实际的。另一个问题在于存储每个模型层的输出的标准实践。对于那些需要使用大的上下文窗口的应用，存储多个模型层的输出的所需的内存需求很快变得十分庞大（从只有几层的GB字节到有数千层的模型[6]的TB字节）。这意味着使用多个层的Transformer模型只能用于几段文本或生产简短的音乐片段。

今天，我们引入Reformer[7]，它是一个Transformer模型，被设计用于处理至多处理100万单词的上下文窗口，所有的这些只在单个加速器上，只使用了16GB内存。Reformer模型结合了两个至关重要的技术来解决限制Transformer应用到长上下文窗口的注意力和内存分配的问题。Reformer使用位置敏感哈希[8]（Locality-Sensitive-Hashing, LSH)来减少关注长序列的复杂度，使用可逆残差层[9]来更加高效地利用可用的内存。

注意力

当将Transformer应用到一个非常大的文本序列时，首要的挑战就是如何处理注意力层。LSH通过计算哈希函数[10]解决了这个问题，LSH把相似的向量匹配在一起，而不是在所有的可能的向量对上进行搜索。例如，在机器翻译任务中，来自网络的第一层的每个向量代表了一个单词（在后续的层中甚至有更大的上下文），不同语言中相同单词对应的向量可能会得到相同的哈希值。在下面的图中，不同的颜色描绘了不同的哈希值，相似的单词有相同的颜色。当哈希值被分配时，序列会被重新排列，以把具有相同哈希值的元素放在一起，然后被划分成片段（或块）以支持并行处理。接着，在这些更短的块（和这些块的相邻块以覆盖溢出）中应用注意力，从而大幅度地减少计算负载。
图：位置敏感哈希：Reformer接收一个输入keys序列，其中每个key是一个代表第一个层中的独立单词的向量（或图像中的像素），以及后续层中更大的上下文。LSH应用于序列，然后按key的哈希和块排序。注意力只应用于单个块及其近邻。

内存问题

尽管LSH解决了注意力问题，但是仍然有内存问题。一个网络的单层通常需要的内存达到几GB，并且通常适用于单个GPU，所以即使一个具有长序列的模型只有一层，也可以被执行。但是，当训练一个具有梯度下降[11]的多层模型，每一层的**函数值[12]需要被保存下来以便在反向传播中使用。一个典型的Transformer模型有12个或更多的层，所以如果用来存储每一层的**函数值，内存会很快消耗完。

在Reformer中实现的第二种新颖的方法就是在反向传播期间按需重新计算每一层的输入，而不是将其存储在内存中。这可以通过使用可逆层来实现，其中来自网络最后一层的**函数值用来恢复任何中间层的**函数值，这相当于反向运行网络。在一个典型的残差网络中，栈中的每一层不断地添加通过网络的向量。相反，可逆层对每个层都有两组**函数值。一个遵循刚刚描述的标准过程，从一个层到下一层逐步更新，但是另一个只捕获对第一个层的更改。因此，为了反向运行网络，只需减去应用于每一层的**函数值。

图：可逆层：(A)在标准残差网络中，每一层的**用于更新到下一层的输入。(B)在可逆网络中，维持两组**，每层后只更新其中一组。©这种方法使反向运行网络以恢复所有中间值成为可能。

Reformer应用

在Reformer中，这两个方法的新颖应用使得其非常高效，使得它能够在单个加速器中仅用16GB的内存就能处理多大100万个单词长度的文本序列。因为Reformer有这样高的效率，它能够被直接应用在长下文窗口比几乎所有当前最先进的文本域数据集大得多的数据。也许Reformer能够处理如此大的数据集的能力会激励社区创建它们。

不缺大上下文数据的一个领域是图像生成，所以我们在图像上实验Reformer。在[13]中，我们呈现了一些关于如何使用Reformer来”补全”部分图像的例子。从下图最上面一行的图像片段开始，Reformer 可以逐像素地生成全帧图像(下面一行)。

图：顶部：图像片段用作Reformer的输入。底部：“完成”的全帧图像。原始图像来自Imagenet64数据集[14]

尽管Reformer在图像成像和视频任务的应用显示了其巨大的潜力，在文本上的应用则甚至更加令人激动。Reformer能够一次性地在单个设备上处理整个小说。在[13]中，演示了在单个训练样本中处理整部小说”Crime and Punishment“。将来，当有更多的长序列文本需要训练时，诸如Reformer之类的技巧可能使得生成连贯的长文本成为可能。

结论

我们相信Reformer为未来使用Transformer模型奠定了基础，既能用于长文本，又能用于自然语言处理之外的应用。根据公开研究的传统，我们已经开始探索如何在更长的序列上应用Reformer以及如何改进位置编码的处理。阅读 Reformer 的论文(被选为 ICLR 2020 的口头报告)，探索我们的代码并开发你自己的想法。目前在深度学习中广泛使用的长上下文数据集还很少，但在现实世界中，长上下文却无处不在。