⏶32
通过μP高效扩展扩散Transformer
发表
由
Chenyu Zheng 提交

作者:
Chenyu Zheng, Xinyu Zhang, Rongzhen Wang, Wei Huang, Zhi Tian, Weilin Huang, Jun Zhu, Chongxuan Li

摘要
扩散 Transformer 已成为视觉生成模型的基础,但其可扩展性受到大规模下高昂的超参数(HP)调优成本的限制。最近,为普通 Transformer 提出了最大更新参数化(muP),实现了从小到大语言模型的稳定 HP 迁移,并显著降低了调优成本。然而,普通 Transformer 的 muP 是否能扩展到扩散 Transformer 仍不清楚,因为它们在架构和目标上有所不同。在这项工作中,我们将标准 muP 推广到扩散 Transformer,并通过大规模实验验证了其有效性。首先,我们严格证明包括 DiT、U-ViT、PixArt-alpha 和 MMDiT 在内的主流扩散 Transformer 的 muP 与普通 Transformer 的 muP 一致,使得现有 muP 方法能够直接应用。利用这一结果,我们系统地证明 DiT-muP 具有强大的 HP 迁移能力。值得注意的是,具有迁移学习率的 DiT-XL-2-muP 比原始 DiT-XL-2 收敛速度快 2.9 倍。最后,我们通过将 PixArt-alpha 从 0.04B 扩展到 0.61B,将 MMDiT 从 0.18B 扩展到 18B,验证了 muP 在文本到图像生成上的有效性。在这两种情况下,使用 muP 的模型性能优于其各自的基线,同时只需要少量调优成本,仅为 PixArt-alpha 单次训练运行成本的 5.5%,以及 MMDiT-18B 人工专家消耗的 3%。这些结果将 muP 确立为扩展扩散 Transformer 的一种原则性且高效的框架。
太棒了!通过 μP 高效地将扩散 Transformer 扩展到 18B!