PaReprop：快速并行化可逆反向传播

数据集与深度学习模型规模的不断增长，使得更快速且内存高效的训练至关重要。可逆 Transformer 近期作为一种极具前景的极内存高效训练方法被提出，但其在反向传播阶段带来了激活重计算的额外计算开销。我们提出 PaReprop，一种快速并行化可逆反向传播算法，将可逆训练中额外的激活重计算开销与反向传播阶段本身的梯度计算并行化。我们通过在模型族（ViT、MViT、Swin 与 RoBERTa）、数据模态（视觉与 NLP）、模型规模（从小到巨大）及训练批大小上的广泛基准测试，展示了所提 PaReprop 算法的有效性。我们的实证结果显示，PaReprop 比原始可逆训练实现高达 20% 更高的训练吞吐量，大幅缓解了可逆训练中激活重计算带来的 25% 吞吐量降低的理论开销。项目页：https://tylerzhu.com/pareprop。