← 全部论文 / 2020

Reformer: The Efficient Transformer

2020 稀疏化与内存高效扩展
下载 PDF

Reformer: The Efficient Transformer - 中文验证版

英文原文卡片：reformer_2020.md

状态：已翻译。

元数据

计算设置

论文说明所有实验的训练在 8 个设备上并行化，具体为"8 GPUs or 8 TPU v3 cores"。它未列出 GPU 型号、内存大小或 TPU v3 cores 以外的确切 TPU 切片。按项目规则，这是明确的 2020 年多加速器设置，但确切的非 TPU 设备 SKU 未指定。主要长序列实验使用 enwik8-64K 和 imagenet64。对于消融，论文使用 3 层模型，d_model = 1024、d_ff = 4096、8 个注意力头、总 batch size 8 序列，以及 Adafactor 优化器。WMT 英德实验遵循 Transformer 超参数。论文还强调大型 Reformer 模型可以装入单核，但该核的硬件型号未给出。

瓶颈

Reformer 将内存分析为复合瓶颈，而非仅注意力瓶颈。密集 QK^T 注意力具有形状 [batch, length, length]。在序列长度 64K 时，即使 batch size 1 也给出 64K x 64K 矩阵，float32 在计入值、梯度或其他激活之前已占用 16GB。引言给出第二个算术示例：64K token、嵌入大小 1024 和 batch size 8 的激活为 64K x 1K x 8，即约 0.5B 浮点数，需要另外 2GB。前馈层可能更糟，因为 d_ff 通常远大于 d_model；使用 d_ff = 4K 和 16 层，论文再次将 64K-token 内存需求称为不切实际。

实际后果是大型 Transformer 模型即使一次训练步骤也可能需要多加速器设置。当有用的注意力模式稀疏但实现仍形成所有成对分数时，标准注意力也浪费计算。

方法适配

Reformer 结合了三种计算适配。首先，局部敏感哈希注意力以近似最近邻注意力替代全对注意力。查询和键共享，向量分配到局部敏感哈希桶中，按桶和位置排序，注意力在排序后的块内计算。目标是在避免物化长度乘长度矩阵的同时，仍让相似的远处 token 可以互相关注。可以并行运行多轮哈希以减少遗漏邻居，以额外计算换取更好的近似。

其次，模型使用可逆残差层。无需为反向传播存储每层的激活，每层的输入可以从其输出和层参数重建。这消除了随层数增长的激活内存因子。第三，前馈分块每次处理序列的一段，因为前馈计算在位置上独立。这减少了来自 d_ff 的内存压力而不改变数值函数。论文还指出，在 batch x 序列乘积足够长时，将参数移入和移出 CPU 内存可以比在普通小 batch Transformer 中更有效地摊销。

因此该设计并不仅是一个稀疏注意力技巧。LSH 攻击 O(L^2) 注意力矩阵；可逆性攻击存储的层激活；分块攻击宽前馈中间量。

证据

论文的首个证据是受控消融。共享查询-键注意力不损害训练曲线，在 enwik8 上似乎稍快。可逆层具有与相同参数量的普通 Transformer 层几乎相同的学习曲线。在一个需要非局部查找的合成重复任务上，用 4 个哈希训练的 1 层 LSH 模型在 4 个哈希评估时达到 99.9% 准确率，8 个哈希评估时达到 100%，展示了哈希轮次的计算-准确率权衡。

长序列实验是核心证据。Reformer 在被分块为 2^16 = 64K-token 子序列的 enwik8 上和序列长度约 12K 的 imagenet64 生成上训练。论文报告，当总 token 保持固定时，常规注意力随序列长度增长变慢，而 LSH 注意力速度保持平坦。它还报告 20 层大型 Reformer 适配内存并在 enwik8 和 imagenet64 上训练，而可比的 Transformer 基线太慢且内存消耗太大。一个 12 层 enwik8 Reformer 训练 20K 步达到 1.19 bits/dim，更长的调优运行在 enwik8 测试集上达到 1.05 bits/dim。在翻译中，可逆 Transformer 模型在基础 100K 步达到 WMT14 En-De BLEU 27.6，基础 500K 步达到 28.0，大 300K 步模型达到 29.1，接近强 Transformer 基线。

历史影响

Reformer 将长上下文 Transformer 设计空间从仅稀疏模式扩展到完整的训练内存核算。其影响在于使可逆层、通过构造的激活重算、前馈分块和近似注意力成为同一对话的一部分。它还给了社区一个难忘的硬件规模论证：单个 64K 注意力矩阵可以在计算模型其余部分之前就消耗一块加速器的内存。

局限

论文自身证据显示若干局限。LSH 注意力是近似的，质量取决于哈希轮次数量；更多轮次提高准确率但花费更多计算。排序、分块、因果掩码和共享 QK 细节使实现比密集注意力更复杂。确切的 GPU 硬件未列出，限制了挂钟时间和内存比较的可复现性。最后，某些收益仅在非常长的序列长度上明显；对于短句级 WMT 示例，作者不应用 LSH 注意力，因为典型示例短于 128-token 的 LSH 块大小。

链接