PETRA：基于可逆结构的并行端到端训练
机器学习
2025-05-20 v2 机器学习
摘要
可逆架构已被证明能够与非可逆架构性能相当，广泛应用于深度学习中的内存节省和生成模型。在本工作中，我们展示可逆架构如何解决深度模型训练中的并行化挑战。我们引入 PETRA，一种用于并行化梯度计算的新型反向传播方法。PETRA 通过使各阶段（即层组）在不同设备上独立计算，仅在彼此之间通信激活和梯度，从而实现有效的模型并行。通过解耦前向和反向传播并保持单个更新后的参数版本，无需参数堆叠（weight stashing）这一需求也被消除。我们开发了用于 PETRA 的自定义 autograd-like 训练框架，并在 CIFAR-10、ImageNet32 和 ImageNet 上Demonstrate其有效性，采用 ResNet-18、ResNet-34 和 ResNet-50 模型，实现了与反向传播相当的竞争性能。
引用
@article{arxiv.2406.02052,
title = {PETRA: Parallel End-to-end Training with Reversible Architectures},
author = {Stéphane Rivaud and Louis Fournier and Thomas Pumir and Eugene Belilovsky and Michael Eickenberg and Edouard Oyallon},
journal= {arXiv preprint arXiv:2406.02052},
year = {2025}
}