Google AI博客解读论文《Reformer: The Efficient Transformer》，百万量级注意力机制

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