Reformer

论文标题：REFORMER: THE EFFICIENT TRANSFORMER

论文链接：https://arxiv.org/abs/2001.04451

提出机构：U.C. Berkeley & Google Research

收录情况：ICLR2020

---

Reformer的优化目标：使Transformer节约内存、加速训练、处理长序列输入。

性能结论：

Reformer相比Transformer的效率优化与sequence length有关，因此在长文本任务上效果才比较明显（论文实验在64K的长文本任务上进行），短文本任务上加速效果不明显。

实现细节：

• 使用局部敏感哈希（LSH）注意力代替传统多头注意力

• 使用可逆层（reversible layers），只存储单层激活值的一份拷贝

• 把FF层里的激活值进行切分

使用上述三种方法，本文在长达64K的文本任务和长达12K的图像任务上进行试验，结果表明Reformer既节约内存提升效率，又拓展了Transformer处理长序列的能力。

局部敏感哈希注意力LSH

Transformer中的注意力计算需要让矩阵Q和K的转置相乘。我们假定它们的形状都是[ batch_size, length, dimension ]，那么如果序列长度有64K，就有得到一个64K*64K的矩阵，显然是不现实的。

对于局部敏感哈希注意力而言，需要 Q=K，以及 V，它们的 shape 都是 [batch size,length,d_model]，而重点关注的是 QK^T，有着 [batch size,length,length] 的 shape。进一步来说，对于每个 q_i，实际需要关注的是它们在 key 的接近值。例如，如果 K 是 64K，对于每个 q_i，只需要考虑一小部分，如 32 个到 64 个最接近的 keys。

这样一来就需要找到最近邻的值，这就需要局部敏感哈希（LSH）了，它能够快速在高维空间中找到最近邻。一个局部敏感哈希算法可以将每个向量 x 转换为 hash h(x)，和这个 x 靠近的哈希更有可能有着相同的哈希值，而距离远的则不会。在这里，研究者希望最近的向量最可能得到相同的哈希值，或者 hash-bucket 大小相似的更有可能相同。

可逆Transformer

虽然 LSH 提升了时间效率，但仍然存在一个内存的问题。当训练一个具有梯度下降的多层模型时，需要保存每一层的激活值，以便在向后传递中使用。一个典型的 Transformer 模型有 12 个或更多的层，因此，如果用来缓存来自每个层的值，那么内存很快就会用完。

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

实验结果

下图是不同的方法在这两个数据集上的表现，可以看到，无论是共享QK还是可逆Transformer，都不会影响效果。

下图是不同哈希桶数的LSH注意力的表现。显然，数量越多，效果越好，这是因为关注就越精确，同时模型代价就越高。

最后是Reformer的层数对于性能的影响。下图（左）是Big Reformer随层变化的不同效果，20层依然无压力。而下图（右）是普通注意力和LSH注意力在不同序列长度的速度比较，当序列很长的时候，LSH具有显著的优势。