通过 $\mu$P 高效扩展扩散变换器

扩散变换器已成为视觉生成模型的基础,但其可扩展性受限于大规模超参数(HP)调谈的高昂成本。最近提出的最大更新参数化($\mu$P)用于标准变换器,可从小规模语言模型稳定 transferring 超参数,显著降低调谈成本。然而,是否可以将 $\mu$P 应用于扩散变换器尚不明确,因为扩散变换器在架构和目标上与标准变换器不同。本文将标准 $\mu$P 推广到扩散变换器上,并通过大规模实验验证其有效性。首先,我们严格证明了主流扩散变换器(包括 U-ViT、DiT、PixArt-$\alpha$ 和 MMDiT)的 $\mu$P 与标准变换器一致，可直接应用现有 $\mu$P 方法。基于此结果,我们系统性地表明 DiT-$\mu$P 具备稳健的超参数 transferability。值得注意的是,使用 transfer learning rate 的 DiT-XL-2-$\mu$P 实现了原始 DiT-XL-2 的 2.9 倍加速收敛。最后,我们在文本到图像生成中验证了 $\mu$P 在扩大 PixArt-$\alpha$ 从 0.04B 到 0.61B 以及 MMDiT 从 0.18B 到 18B 中的有效性。在两种情况下,在 $\mu$P 下的模型均优于基线,仅需对 PixArt-$\alpha$ 进行 5.5% 的一次训练运行,对 MMDiT-18B 只需 3% 的人类专家消耗。这一结果确立了 $\mu$P 作为扩散变换器可扩展的原则性且高效的框架。