Reformer 模型 - 突破语言建模的极限

Reformer 如何在不到 8GB 的内存上训练 50 万个词元

Kitaev、Kaiser 等人于 20202 年引入的 Reformer 模型 是迄今为止长序列建模领域内存效率最高的 transformer 模型之一。

最近，人们对长序列建模的兴趣激增，仅今年一年，就涌现出了大量的工作，如 Beltagy 等人的工作 (2020) 、Roy 等人的工作 (2020) 、Tay 等人的工作 以及 Wang 等人的工作 等等。长序列建模背后的动机是，N​​LP 中的许多任务 (例如 摘要、问答 ) 要求模型处理更长的序列，这些序列长度超出了 BERT 等模型的处理能力。在需要模型处理长输入序列的任务中，长序列模型无需对输入序列进行裁剪以避免内存溢出，因此已被证明优于标准的 BERT 类模型 ( 见 Beltagy 等人 2020 年的工作)。

Reformer 能够一次处理多达 50 万个词元，从而突破了长序列建模的极限 (具体可参见本 笔记本)。相形之下，传统的 bert-base-uncased 模型最长仅支持 512 个词元。在 Reformer 中，标准 transformer 架构的每个部分都经过重新设计，以最小化内存需求，并避免显著降低性能。

内存的改进来自于 Reformer 作者向 transformer 世界引入的 4 大特性:

1. Reformer 自注意力层 - 如何在不受限于本地上下文的情况下高效地实现自注意力机制？
2. 分块前馈层 - 如何更好地对大型前馈层的时间和内存进行权衡？
3. 可逆残差层 - 如何聪明地设计残差架构以大幅减少训练中的内存消耗？
4. 轴向位置编码 (Axial Positional Encodings) - 如何使位置编码可用于超长输入序列？

本文的目的是 深入 阐述 Reformer 的上述四大特性。虽然这四个特性目前是用在 Reformer 上的，但其方法是通用的。因此，读者不应被此束缚，而应该多思考在哪些情况下可以把这四个特性中的某一个或某几个应用于其他的 transformer 模型，以解决其问题。

下文四个部分之间的联系很松散，因此可以单独阅读。

Reformer 已集成入