阅读笔记 | REFORMER: THE EFFICIENT TRANSFORMER

阅读论文：

Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. “Reformer: The efficient transformer.” arXiv preprint arXiv:2001.04451 (2020).

背景与动机

这篇论文发表较早，主要关注Transformer的效率问题。标准的Transformer模型在许多自然语言处理任务上取得了最先进的结果，但它在长序列上的训练和推理代价非常大。主要的计算和内存瓶颈在于自注意力机制中的点积注意力，其复杂度为$\mathrm{O}(L^2)$，其中$L$是序列长度。随着$L$的增大，计算和内存需求急剧增加。因此，Transformer难以扩展到处理长序列的任务。

模型与方法

1. 引入基于局部敏感哈希（LSH）的近似注意力机制，将自注意力的复杂度从$\mathrm{O}(L^2)$降低到$\mathrm{O}(L\log L)$，大大减少了内存和计算需求。具体来说，文章首先使用随机投影作为敏感哈希函数。相似的query和key通过投影转换后可以映射到相同的哈希桶中。然后根据query和key的哈希值对序列进行排序。相似的query和key会聚集在一起。在排序后的序列上，将每个query块只与相邻的几个query块计算注意力。这样可以大约保证每个块内的query可以attend到相似的key。最后使用多轮不同的哈希函数和注意力计算，综合多个注意力输出，可以降低哈希误差。
2. 使用17年别人提出的一个可逆残差网络取代标准的前馈和残差连接以减少内存占用。这允许在训练过程中只需保存每层的一个激活副本，而不是N个副本，其中N是层数。
3. 将前馈网络中的激活拆分为多块分别处理，减小每层的内存占用。

实验

在拼写重复、机器翻译以及图像生成等多个任务上验证了Reformer的有效性。结果表明，Reformer模型比标准Transformer有相同的建模能力，但在长序列任务上速度更快，内存占用也少得多。例如，在64K词元的英文文本建模任务上，12层的Reformer只需要一个GPU就可以高效训练，而Transformer基本无法进行训练。

总的来说，Reformer是一个既高效又强大的长序列Transformer模型。它结合了可缩放的近似注意力机制与高效的网络结构，能够在单机单卡上处理长度达10万词元的任务，为Transformer的工业应用提供了可能性。

个人思考

• 文章中实现更高效自注意力的方式是寻找近似方法进行替代，这对于后面的改进是有启发的，因为类似的替代并非只有一种。而且这类改进存在一个好处就是可以在用于电负荷时序预测模型中去带一个计算与存储效率的点，这也是前面阅读的所有时序预测论文都在关注的一个点。
• 这篇论文也从网络结构方面对内存消耗进行了优化，对残差连接进行了优化。在Informer中也对结构进行了关注，具体是改进了decoder的自回归方式。可见结构方面进行考虑的优化也是大有可为的。此外，几个最近的Transformer时序预测模型在残差连接上基本没有考虑太多，用的是标准残差连接，可以将可逆残差连接放到里面试试效果。