Reformer：一个高效的 Transformer

文 / Nikita Kitaev , 学生研究员 , 加利福尼亚大学伯克利分校与  Łuk asz Kaiser,  研究科学家,  G oogle Research 

理解语言、音乐或视频之类序列数据 (Sequential Data) 的任务充满挑战，尤其是当这些数据依赖于时间跨度较广的上下文时。

例如，在一个视频中，如果某个人或某个物体在消失很长一段时间后才重新出现，许多模型会忘记它的样子。而在语言领域，长短期记忆 (LSTM) 神经网络会考虑足够多的上下文来执行 逐句翻译 。在这种情况下，上下文窗口 (Context Window) 的范围（即翻译中考虑的数据范围）会涵盖数十到上百个单词。

最新的 Transformer 模型 不仅提高了逐句翻译的性能，而且可以用于通过多文档摘要生成 完整的 Wikipedia 文章 。由于 Transformer 使用的上下文窗口扩展到了数千个单词，使上述情景得以实现。凭借宽泛的上下文窗口，Transformer 的应用范围从文本扩展到了包括像素、音符在内的其他场景，可以用于生成 音乐 和 图像 。

但是，Transformer 的上下文窗口有限制范围。

Transformer 的强大来源于 注意力 (Attention) 机制 ，通过这一机制，Transformer 将上下文窗口内所有可能的单词对纳入考虑，以理解它们之间的联系。因此，如果文本包含 10 万个单词，Transformer 将需要评估 100 亿单词对（10 万 x 10 万），这显然不切实际。 另一个实践问题是 如何保存每个模型层的输出 。对于使用大型上下文窗口的应用来说，存储多个模型层输出的内存需求会迅速变得过大（从几层模型层的数 G 级别到 数千层 模型层的数 T 级别）。这意味着，实际使用大量层的 Transformer 模型只能用于生成几小段落的文本或一小段的音乐。

今天，我们推出了 Reformer ，一个设计为处理多达 100 万单词的上下文窗口的 Transformer 模型，所有工作在单个加速器上进行且仅使用 16 GB 内存。此模型将综合运用两种关键技术来解决 Transformer 在长上下文窗口的注意力和内存分配问题的应用限制。Reformer 使用局部敏感哈希 (Locality-Sensitive-Hashing, LSH) 来降低长序列的处理复杂度和 可逆残差层 ，从而更有效地使用可用内存。

注意力问题

将 Transformer 模型应用于宽范围的文本序列时，第一个挑战是如何处理注意力层。

局部敏感哈希的挑战在于通过计算向量匹配哈希函数，而不是搜索所有可能的向量。例如，在翻译任务中，来自网络首层的每个向量表示一个单词（甚至是后续层中更大量的上下文），而不同语言中相同单词的向量可获得相同的哈希值。

在下图中，不同的颜色表示不同的哈希值，相似的词则具有相同的颜色。分配哈希值后，序列重新排列，将具有相同哈希值的元素放在一起，再分为多个片段（或多个区块）以实现并行处理。然后在这些短得多的区块（及其相近邻块以覆盖溢出）内应用注意力，从而大大降低计算负载。

局部敏感哈希：Reformer 接收 Key 的序列，其中每个键值代表首层中每个单词（在图像的情况下为像素）和后续层中大量的上下文向量。将局部敏感哈希应用到序列之后，依照键值的哈希值对其 排序 并分块。注意力仅针对于单个区块及其直接邻块

内存问题

虽然局部敏感哈希解决了注意力问题，但仍然存在内存问题。

一个单层的网络一般需要多达几 GB 的内存以及单个 GPU，因此，单层能够执行长序列的单模型。但是，当使用梯度下降训练多层模型时，由于需要保存每一层的激活（函数），以用于执行逆推。一个传统的 Transformer 模型具有十几个或更多的层，通过缓存这些层的值，内存将会很快用完。

第二种 Reformer 中实现的新方法是在反向传播期 (Back-Propagation) 间，按需重新计算每个层的输入，而不是将其保存在内存中。这是使用 可逆层 来实现的，其中来自网络最后一层的激活 (Activation) 用于还原来自任何中间层的激活 (Activation)，这相当于反向运行网络。

在典型的残差网络中，通过网络传递的输入将会向堆栈中的每一层不断添加至向量。相反，可逆层中每个层有两组激活。一组遵循刚才描述的标准过程，从一层逐步更新到下一层，但是另一组仅捕获第一层的变更。因此，若要反向运行网络，只需简单地减去每一层应用的激活。

可逆层：（图 a）在标准残差网络中，每一层的激活用于将输入更新到下一层。（图 b）在可逆网络中，将会维持两组激活，只有其中一组逐层更新。（图 c）这种方法可使网络反向运行，以还原所有中间值

Reformer 的应用

这两种新方法能让 Reformer 具有更高的效率，可以在仅有 16GB 内存的单个加速器上处理长达一百万个单词的文本序列。正因为 Reformer 的高效率，可以被直接应用于上下文窗口远大于任何最新 (SOTA) 文本域数据集的数据。借助 Reformer 能处理如此大规模数据集的能力，可促进社区创建大规模数据集的速度。

在不缺乏大规模上下文数据的领域中，我们选择图像生成作为我们实验 Reformer 的对象。在这份 Colab 笔记 [1] 中，我们提供了一些示例，说明如何使用 Reformer 将部分图像变得“完整”。从下图上行所示的图像片段开始，Reformer 可以逐像素生成全画幅图像（下行）。

上行：用于输入 Reformer 的图像片段。下行：“已完成的”全画幅图像。原始图像来源于 Imagenet64 数据集

Reformer 不仅仅在处理图像和视频任务中的应用有巨大的潜力，其在文本中的应用更加令人兴奋。Reformer 可以在单一设备上一次处理完整部小说。这份 Colab 笔记 [2] 演示了在单个训练示例中处理《罪与罚》 (Crime and Punishment)  的全部内容。将来会有更多的带有长文本的数据集需要训练，诸如 Reformer 之类的技术让生成条理清楚的长篇作品成为可能。

结论

我们相信， 无论是在长文本还是自然语言处理之外的应用， Reformer 为将来使用 Transformer 模型提供了基础。

遵循 公开研究 的传统，我们已经开始探索如何将 Reformer 应用于 更长的序列 以及如何 改善位置编码的处理 。阅读这篇 Reformer 论文 （选自 ICLR 2020 上的口头演讲），探索 我们的代码 并生成自己的创意。

虽然深度学习中尚未广泛使用长上下文数据集，但在现实世界中比比皆是。也许这可以激发您找到 Reformer 的新应用，这份 Colab 笔记 [3] 可作为您的起点，如果遇到任何问题或疑问，欢迎 与我们沟通 ！

致谢

此项研究由 Nikita Kitaev、Łukasz Kaiser 和 Anselm Levskaya 共同完成。另外还要感谢 Afroz Mohiuddin、Jonni Kanerva 和 Piotr Kozakowski 在 Trax 上的工作，并感谢整个 JAX 团队的支持。