ICLR 2020 | reformer高效处理长序列，单机能跑，计算资源贫困人士的福音

机构：Google Research 、U.C. Berkeley

作者：Nikita Kitaev、Łukasz Kaiser、Anselm Levskaya

论文地址：

收录会议：ICLR2020

论文代码：

那训练 Transformer 模型是否真需要很多资源且很低效？以现有的最大 Transformer 层为例，该 Transformer 层中参数量是 0.5B，这需要 2GB 的内存。（1M=1024KB,1KB=1024Byte。所以1GB=1024M=1024x1024KB=1024x1024x1024Byte=1073741824Byte。float 占用 4 个 Byte。0.5B 即 5 亿参数，需要的内存量为 5 亿 *4 字节=20 亿字节。这差不多是 1.86GB 即约为 2GB）对于由 64Ktokens 组成的序列，如果嵌入层的尺寸是 1024，batch size 是 8，那么激活值需要 64K * 1K * 8=0.5B 个浮点数来存储，这又需要 2GB 的内存。如果每层的内存占用只有上述提到的这些的话，那么在单加速器上使用Transformer 处理 64K长度的序列也是轻而易举。此外，如此前提下训练 BERT 的整个语料库也只需 17GB 的内存。然而，现实并非如此，真实环境下为何甚至不能在单台机器上对这些模型进行微调呢？

这是因为上述仅仅考虑单层参数的内存占用和输入激活值的内存消耗，而忽略了 Transformer 在内存占用上的主要问题：

- 需要存储激活值用于反向传播，那么 N 层模型内存占用是单层的 N 倍；

- 由于中间全连接层的深度 d_{ff} 通常远大于注意力激活层的深度 d_{model}，而这需要占用很大的内存；

- 长度为 L 的序列的 attention 的时间和空间复杂度是 O(L^2)，那么对于 64K tokens 的序列就会耗尽内存。

为此，本文提出 Reformer 模型以解决上述问题，具体采用如下方案：

- 可逆层（Reversible layer），在整个模型中只使用单个副本，可以消除层数因子 N。

- 前馈层（feed-forward layer）分开激活和分块处理，从而消除 d_{ff} 因子的影响，降低前馈层的内存占用。

- 采用基于局部敏感哈希（locality-sensitive hashing，LSH）的近似注意力计算，让注意力层的 O(L^2) 因子变为 O(L log L) ，这使得在长序列上的处理成为可能。

上述的 attention 操作并行地进行 h 次，再输出维度为 d_v 的输出结果。再将这些结果拼接，再做一次投射操作得到最终的结果。即所谓的多头 attention。

先来算下上述 attention 机制消耗的内存。假设 Q，K，V 的尺寸为 [batch_size，length，d_model]。QK^T 的尺寸为 [batch_size，length，length]。当 length=64k，即使 batch_size=1，那么 64k*64k 大小的矩阵，如果用 32 位浮点数来存储的话，需要 16GB 内存。鉴于此，在长序列上使用 Transformer 显得不切实际。但是需要注意的是，QK^T 矩阵可以不必全部放在内存中，可以对每个 query 分别计算 attention。反向传播计算梯度时再重新计算一次。这种方式计算 attention 虽然低效，但是所占用的内存与 length 成正比。这种方法在本文这里作为一种全 attention 的 baseline。

上述讨论了 Q、K、V，但是一般我们只会得到大小为 [batch_size，length，d_model] 的激活值 A，这些值是 token 的嵌入所组成的句向量。那么为了从 A 中得到Q、K、V，Transformer 使用了 3 个不同的线性层(参数不同)将 A 投射为 Q、K、V。对于使用局部敏感哈希 attention 的模型，我们希望 queries 和 keys(即 Q 和 K)相同。只需要 A 投射到 Q 和 A 投射到 K 时采用相同线性变换参数即可，而 A 投射到 V 时采用不同参数。这种方式成为共享 QK-Transformer。实验表明共享 QK 并不会影响 Transformer 的性能，即使添加一项 d_k 的归一化项。

在 LSH attention 中，假设 Q、K、V 的尺寸为 [batch_size，length，d_model]，同时仍然使用此前介绍的多头 attention 机制。那么 QK^T 的尺寸为 [batch_size，length，length]。由于 softmax(QK^T) 的计算结果主要取决于值最大的部分，对于每个 query 只需关注 K 中与 query 最接近的点。当 K 的长度是 64k，那么对个每个 query，本文仅仅考虑其最近的的 32 或 64 个 keys。如此会更加高效，那么如何找寻最近的那些 keys 呢？

在高纬空间中找寻最近邻可以使用局部敏感哈希(LSH)。将每个向量 x 通过 hash 函数h(x) 进行映射，如果近处的向量获得相同的 hash，且具有高概率，而远处的向量没有，那么这样的 hash 称为位置敏感型 hash。在此处例子中，我们实际上只要求近邻的向量以高概率具有相同的 hash 值，并且 hash 桶也以高概率具有相同的大小。

具体是使用如 Figure 1 所示的随机投射方法：

图 a：常规的 attention 机制中，黑点代表的是 softmax 中占主导的位置。注意这边的 attention 使用的是 encoder 的 attention，否则 q_3 无法 attend 到 k_6。另外，这种全 attention(即 encoder 中的 attention)的 attention 矩阵一般是稀疏的，但计算中并没有利用这种稀疏性，所以可以利用这个降低时间空间复杂度。

图 b：计算 query 和 key 所归属的 hash 桶。再按照桶进行排序，同一个桶又按照原本的位置进行 排序 得到图 b。可以看到，同一个桶，可以出现多个 query 但 keys 很少的情况，例如图中蓝色的桶 query 有 3 个，都 attend 到同一个 key 中。由于相似的 item 很有可能落在同一个桶里，所以只在每个桶内部进行 attention 就可以近似全 attention。

图 c：为了缓解桶中 q 和 k 不均衡问题，本文通过令 $k_{j}=\frac{q_{j}}{\left\|q_{j}\right\|}$ 使得 h(k_j)=h(q_j)，即使用了 share-QK attention。然后先按照桶序号对 queries 排序，每个桶中，仍按照原本的 position 位置大小排序。得到图 c。对比 b 图和 c 图可以看出，纵轴的 k 已经变成了 q。这时就能保证对角线都是 attend 到的而且 q 和 k 在桶中的个数一样（因为 Q=K）。排序后的 attention 矩阵，相同桶的值会在对角线附近聚集。注意到图中对角线的点为空心，这是因为虽然在正常情况下，q 会 attend to 本身位置的 value，但是在 share-QK 的实现下，如果 attend to 本身，会导致其值特别大，其他的值特别小，经过 softmax 之后，其他都是 0，就自己本身是 1。所以为了避免这种情况，q 不会去 attend 自身位置的值，除非只有自己本身可以 attend。

图 d：即使 Q=K，还是会出现一个问题：有的桶中个数多，有的桶中个数少。比如一个极端情况，2 个桶，其中一个桶占据了所有的 keys，另一个桶为空，那么 LSH attention 就没有起作用。于是在图 c 的基础上，增加了 chunk 的操作。对输入进行排序之后(即图 c 中先桶排序，同个桶内按照 token 的 position 排序）得到新的序列顺序 s_i，比如图中原来的序列顺序是 [q_1,q_2,q_3,q_4,q_5,q_6]，新的序列顺序是[q_1,q_2,q_4,q_3,q_6,q_5] 。每个 chunk 内 query 的上限个数为 $m=\frac{2 l}{n_{\text {buckets}}}$, (l 为输入 query 的长度) ，每个桶平均大小为 $m=\frac{l}{n_{\text {buckets}}}$，这里假设桶中数量增加到均值两倍的概率足够低。对于桶中的每个 query，都可以 attend to 自己以及前一个桶中相同 hash 值的 key。

小结下，LSH attention 做了以下两个事情：

第一，找到 Q、K 矩阵的 LSH hashes。

第二，在同一个 hash 桶内计算 k 和 q 向量的标准 attention。

更具体来说可分为以下 5 个步骤：

第一，令输入序列 queries=keys

第二，做 LSH bucketing，即进行 hash 计算，得到每个 query 和 key 所归属的桶(不同颜色表示不同的桶)。

第三，根据桶编号对 query 进行排序，同个桶中，按照 query 原本的位置进行排序。

第四，对于排序后的新序列，进行 chunk 拆分

第五，对于每个 query 只 attend 自己以及自己之前的 chunk，对于这些候选集中相同桶的 key 进行 attend。

如上所述，attention 的复杂度可以被减少为与序列长度成线性正比，但是，参数量占的复杂度依旧很高，如何进一步减少呢？这里就开始尝试解决前文介绍部分所提到的第二和第三个问题，即大量的 encoder 和 decoder 层、全连接层 FFN 的深度问题。

如此，使用可逆的 Transformer 在每一层中就无需存储激活值，也就避免了 n_l 这一项。可逆层代替标准的残差层，可以在训练过程中只存储一次激活，而不是 N 次。

另外，可逆操作和反向传播操作也分块处理。除 FFN 之外，对于词汇量大的模型(单词类型>d_{model})，还对输出处的 log- probability 分块，并一次计算序列各部分的损失。

实验结果

对图像生成任务 imagenet64(序列长度为 12K)和文本任务 enwik8-64K(即序列长度为64K)进行了实验，评价了可逆层、共享 query-key、LSH attention 对内存、精度和速度的影响。

可逆层和共享 query-key 的影响：

Figure 3 中的右部分验证的是可逆层的影响。实验中对比的可逆层和常规 Transformer 参数量相同，且学习曲线看起来也几乎相同。这些结果表明，可逆 Transformer 在节省内存的同时并不会牺牲精度。

作者：刘杰鹏（微信号：onepieceand）

毕业院校：华中科技大学

研究方向：机器阅读理解、文本生成等。