图解 Reformer: The efficient Transformer

翻译自：https://towardsdatascience.com/@alirezadir

如果你一直在开发用于处理顺序数据的机器学习算法，例如语言处理中的文本、语音信号或视频，那么你可能听说过或使用过Transformer 模型，并且你可能知道它与 twitter 认为的不同：

图1，开心一刻，来自于 Chris Manning 教授的推特

最近，Google 引入了 Reformer 架构，这是一种 Transformer 模型，旨在有效地处理非常长的数据序列（例如，在语言处理中多达 100 万字的序列）。Reformer 的执行需要非常少的内存消耗，即使只在一个 GPU 上运行也能获得惊人的性能。论文 Reformer: The efficient Transformer将在 ICLR2020 中展示（并在评论中获得接近完美的分数）。Reformer 模型有望超越语言应用（如音乐、语音、图像和视频生成）而对该领域产生重大影响。

在这篇文章中，我们将试着深入研究 Reformer 模型，并尝试用一些视觉指南来理解它。准备好了吗？

Why Transformer?

NLP 中有一类任务（如机器翻译、文本生成、问答）可以表示为序列到序列的学习问题。长短期记忆（LSTM）神经网络，后来配备了[注意力机制](http://neural machine translation by jointly learning to align and translate./)，是一个突出的架构，用于建立预测模型等问题-例如谷歌的神经机器翻译系统。然而，（从速度和消失梯度的角度看）LSTMs 中递归的内在序列性是数据序列并行计算的最大障碍，因此，这些体系结构无法利用长序列的上下文。

Transformer 模型在论文 Attention is all you need 中被提出，通过引入多头自我注意机制，在许多任务中实现最先进的性能。Transformer 的主要创新之处在于它的并行处理能力，它能够处理长序列（上下文窗口为数千字），从而称为最先进的模型，例如显著的 Open AI 的 GPT2 语言模型，训练时间更短。 Huggingface 的 Transformers 库有超过 32 个 100 多种语言的预先训练模型，TensorFlow 和 Pythorch 之间具有互操作性，是构建最先进的 NLP 系统的极好的开源工作。 Write with transformers 和 Talk to Transformers 是一些有趣的演示。Transformer 已经被用于文本以外的应用，如[生成音乐](Write with transformers and Talk to Transformers )和图像。

What’s missing from the Transformer?

在深入 reformer 之前，让我们回顾一下 transformer 模型的挑战。这需要对 transformer 体系结构本身有一些了解，我们在本文中无法详细介绍。然而，如果你还不知道的话，Jay Alamar 的 The Illustrated Transformer 是目前为止最好的视觉解释，我强烈建议你在阅读完这篇文章之前先阅读他的文章。

尽管 transformer 模型在越来越长的序列上得到了很好的结果——例如（Liu et al., 2018）中 11K 的长文本示例——许多这样的大型模型只能在大型工业计算平台上进行训练，即使是在单个训练步骤中，因为它们的内存需求也不能在单个 GPU 上进行微调。例如，完整的 GPT-2 模型由大约 15 亿个参数组成。在（Shazeer et al.，2018）中报告的最大配置中，每层的参数数量超过5 亿，而层的数量增加到 64 层（al-Rfou et al.，2018）。

让我们看一下 Transformer 模型的简化概述：

图2：标准 Transformer 模型的简化摘要

如果这个模型看起来不熟悉或者很难理解，我建议您在这里停下来回顾一下➡️The Illustrated Transformer。

你可能注意到图中有 3 种不同颜色的 。每一个唯一的 代表了 Transformer 模型的一部分，Reformer 作者认为它是计算和内存问题的根源：

Problem 1 (Red ): Attention computation

对长度为 $L$ 的序列的计算注意力为 $O(L^2)$（时间和空间复杂度）。想象一下如果我们有一个长度为 64K 的序列会发生什么。

Problem 2 (Black ): Large number of layers

具有 $N$ 层的模型比单层模型消耗的内存大 $N$ 倍，因为每层中的激活都需要存储以进行反向传播。

Problem 3 (Green ): Depth of feed-forward layers

中间前馈层的深度通常远大于注意力激活层的深度。

Reformer 模型解决了 Transformer 中上述三个主要的内存消耗源，并对它们进行了改进，使 Reformer 模型能够处理多达 100万字的上下文窗口，所有这些窗口都在一个加速器上，并且只使用 16GB 的内存。

简言之，Reformer 模型结合了两种技术来解决注意力和内存分配的问题：局部敏感哈希（LSH），以减少过长序列的复杂度，以及可逆残差层以更有效地使用可用内存。

下面我们将详细介绍。

1. Locality sensitive hashing (LSH) Attention

Attention and nearest neighbors

深度学习中的注意力是一种机制，它使网络能够根据上下文的不同部分与当前时间步的相对性，将注意力集中到上下文的不同部分。transformer 模型中存在 3 种类型的注意力机制：

图3：Transformer 模型中的三种注意力机制

Transformer 中使用的标准注意力是标量点积，公式如下：

从上面的方程和下图可以看出，乘法 $Q{K}^T$（形状[$L$，$L$]）的计算和存储成本都是 $O(L^2)$，这是主要的存储瓶颈。

图4：（左）：点积注意力机制中的主要计算，（右）注意力序列中其他 tokens 的子集的标记（‘the’、‘animal’、‘street’、‘it’、’’）的示例

❓但是否有必要计算和存储完整的矩阵 $Q{K}^T$？答案是否定的，因为我们只对 $softmax(Q{K}^T)$ 感兴趣，它由典型稀疏矩阵中的最大元素支配。因此，正如您在上面的示例中所看到的，对于每个查询 $q$，我们只需要关注最接近 $q$ 的键 $k$。例如，如果 $K$ 的长度为 64K，对于每个 $q$，我们只能考虑 32 个或 64 个最接近 $k$ 的一小部分。因此，注意力机制查找 $q$ 的最近邻居 $k$，但效率很低。这会让你想起最近邻居查找吗？

reformer 中的第一个新颖性来自于用局部敏感散列（LSH）代替点乘注意力来改变从 $O(L^2)$ 到 $O(L\log L)$ 的复杂度。

LSH for nearest neighbors search

LSH 算法是一种高效、近似的高维数据集最近邻搜索算法。LSH 背后的主要思想是选择散列函数，使得对于两点 “p” 和 “q”，如果 “q” 接近 “p”，那么我们有足够的概率得到 “hash(q) == hash§”。

实现这一点的最简单方法是通过随机超平面保持空间切割，并在每个点的散列码上附加 $sign(p^{T}H)$。让我们看一个下面的例子：

图5：近邻搜索的局部敏感哈希简化动画

一旦我们找到所需长度的散列码，我们就根据它们的散列码将这些点分成若干个桶-在上面的示例中，“a”和“b”属于同一个桶，因为 hash(a) == hash(b)。现在，查找每个点的最近邻的搜索空间从整个数据集中显著减少到它所属的 bucket 中。

角 LSH：平面 LSH 算法的一种变体，称为角 LSH，将点投影到一个单位球体上，该球体已被分成预定义的区域，每个区域都有一个不同的 code。然后一系列随机旋转的点定义了这些点所属的 bucket。让我们通过一个简化的二维示例来说明这一点，该示例取自 Reformer 的论文：

图 6：近邻搜索角度 LSH 的简化动画：两点相隔较远

这里我们有两个点投影到一个单位圆上，以不同的角度随机旋转 3 次。我们可以观察到，它们不太可能共享同一个散列桶。然而，在下一个示例中，我们看到彼此非常接近的两个点在 3 次随机旋转后将共享相同的散列桶：

图中第一次散列值应该是 （0，0），而不是（ 0，3）

图7：近邻搜索角度 LSH 的简化动画：两点接近

LSH attention

现在 LSH 注意力背后的基本思想如下。回顾上面的标准注意力公式，我们没有计算 $Q$ 和 $K$ 矩阵中所有向量的注意力，而是执行以下操作：

• 寻找矩阵 $Q$ 和 $K$ 的 LSH hashes
• 只计算同一散列桶中 $k$ 和 $q$ 向量的标准注意力

多轮 LSH 注意力：重复上述步骤几次，以增加类似项不会落入不同桶中的可能性。

下面的动画演示了一个基于该图的 LSH Attention 的简化版本。

图 8：LSH 注意力机制的简化说明

2. Reversible Transformer and Chunking

现在我们已经准备好解决 Transformer 中的第二和第三个问题，即深度的 ($N$ 层) 编码器和解码器层以及前馈层。

Reversible residual Network (RevNet)

注意图 2 中的编码器和解码器块，我们认识到每个注意层和前馈层被包装成一个残差块（类似于我们在图 9（左）中看到的）。论文介绍的残差网络（ResNets）是神经网络结构中用来解决深层网络（多层网络）中的消失梯度问题的强大组件。然而，ResNets 中的内存消耗是一个瓶颈，因为需要将每个层中的激活存储在内存中，以便在反向传播期间计算梯度。内存成本与网络中的单元数成正比。

为了解决这个问题，提出了由一系列可逆块组成的可逆残差网络（RevNet）。在 Revnet 中，每一层的激活都可以从下一层的激活中精确地重建，这使得我们能够执行反向传播，而不必将激活存储在内存中。图 9 说明残差块和可逆残差块。请注意，我们如何从块的输出 ($Y_1,Y_2$) 计算块的输入 $(X_1,X_2)$。

图 9：图解残差网络块（左）和可逆残差块（右）

Reversible Transformer

回到我们的第二个问题，这个问题是处理 $N$ 层 Transformer 网络的内存需求-潜在的相当大的 $N$。

Reformer 通过将 RevNet 块中的注意力层和前向层相结合，将 RevNet 的思想应用于 Transformer。在图 9 中，现在 $F$ 成为注意力层，$G$ 成为前向层：

$Y_1=X_1+Attention(X_2)$

$Y_2=X_2+FeedForward(Y_1)$

现在使用可逆残差层而不是标准残差层可以在训练过程中只存储一次激活，而不是 $N$ 次。

Chunking

Reformer 效率改进的最后一部分涉及第三个问题，即前馈层的高维中间向量-可以达到 4K 或更高的维数。

由于前馈层中的计算是独立于序列中的各个位置的，因此前向和后向过程的可逆计算都可以分成块。例如，对于前向过程，我们将有：

$$Y_2=\left[Y_2^{(1)};\dots ;Y_2^{(c)}\right]=\left[X_2^{(1)}+\text{ FeedForward }\left(Y_1^{(1)}\right);\dots ;X_2^{(c)}+\text{ FeedForward }\left(Y_1^{(c)}\right)\right]$$

前向传递计算中的分块[公式来自于 Reformer 论文]

Experimental Results

作者对图像生成任务 imagenet64（长度为12K 的序列）和文本任务 enwik8（长度为 64K 的序列）进行了实验，评价了可逆 Transformer 和 LSH 哈希对内存、精度和速度的影响。

可逆 Transformer 与基线匹配：他们的实验结果表明，可逆 Transformer 在不牺牲精度的情况下节省了内存：

可逆性对 enwik8 和 imagenet64 训练性能的影响[图片和标题取自 Reformer 论文]。

LSH注意力匹配基线：注意，LSH注意力是完全关注的近似值，其精度随着哈希值的增加而提高。当散列值为 8 时，LSH attention 几乎等于 full attention：

LSH 注意力作为哈希轮次函数对 imagenet64 的影响[图片和标题取自 Reformer 论文]。

他们还证明，随着序列长度的增加，传统注意力速度减慢，而 LSH 注意力速度保持稳定，它在 8GB GPU 上以常规速度运行在长度约 100K 的序列上：

Full 和 LSH 注意力的速度评估值与输入长度成函数关系[图片和标题取自 Reformer 论文]。

最终的 Reformer 模型与 Transformer 模型的性能相似，但在长序列上显示出更高的存储效率和更快的速度。

Trax: Code and examples

The code for the Reformer has been released as part of the new Trax library. Trax is a modular deep learning training and inference library which is aimed to allow you to understand deep learning from scratch. The Reformer code includes several examples that you can train and infer on image generation and text generation tasks.

Reformer 的代码已作为新 Trax 库的一部分发布。Trax 是一个模块化的深度学习训练和推理库，旨在让你从零开始理解深度学习。Reformer 程序代码包括几个示例，你可以对图像生成和文本生成任务进行训练和推理。

Acknowledgment

I would like to thank Łukasz Kaiser for his vivid presentation of the Reformer and providing supplementary material.

References and related links:

Reformer: The Efficient TransformerUnderstanding sequential data - such as language, music or videos - is a challenging task, especially when there is…ai.googleblog.com

Transformer: A Novel Neural Network Architecture for Language UnderstandingNeural networks, in particular recurrent neural networks (RNNs), are now at the core of the leading approaches to…ai.googleblog.com

• Reformer: The efficient Transformer
• Google/Trax deep learning library
• The Illustrated Transformer
• Huggingface/Transformers NLP library
• Attention is all you need
• Open AI’s GPT2 language model
• Write with transformers
• Talk to Transformers
• Google’s Neural Machine Translation System