Scaling Diffusion Transformers Efficiently via $μ$P

扩散变换器（Diffusion Transformers）已成为视觉生成模型的基础，但其可扩展性受到大规模超参数（HP）调优高成本的限制。最近，针对普通变换器提出了最大更新参数化方法（$\mu$P），该方法能够实现从小型到大型语言模型的稳定超参数迁移，并显著降低调优成本。然而，尚不清楚普通变换器的$\mu$P是否可以扩展到架构和目标不同的扩散变换器。在本工作中，我们将标准的$\mu$P推广到扩散变换器，并通过大规模实验验证了其有效性。首先，我们严格证明了主流扩散变换器（包括DiT、U-ViT、PixArt-$\alpha$和MMDiT）的$\mu$P与普通变换器一致，从而可以直接应用现有的$\mu$P方法。基于这一结果，我们系统地展示了DiT-$\mu$P具有强大的超参数迁移能力。特别地，使用迁移学习率的DiT-XL-2-$\mu$P比原始的DiT-XL-2收敛速度快2.9倍。最后，我们通过将PixArt-$\alpha$从0.04B扩展到0.61B，以及将MMDiT从0.18B扩展到18B，验证了$\mu$P在文本到图像生成中的有效性。在这两种情况下，采用$\mu$P的模型均优于各自的基线模型，同时调优成本极低，仅占PixArt-$\alpha$一次训练运行的5.5%，以及MMDiT-18B由人类专家调优消耗的3%。这些结果确立了$\mu$P作为扩展扩散变换器的原理性和高效框架的地位。

该论文试图解决扩散Transformer模型在大规模应用中的超参数（HP）调优成本过高的问题。这是一个长期存在的挑战，尤其是在生成式模型领域。

论文的关键思路是将Maximal Update Parametrization ($\mu$P) 方法从传统的Transformer模型扩展到扩散Transformer模型。通过证明主流扩散Transformer（如DiT、U-ViT等）与传统Transformer的$\mu$P一致性，论文使得已有$\mu$P方法可以直接应用于扩散Transformer，从而大幅减少超参数调优的成本和时间。

1. 系统性验证了DiT-$\mu$P的鲁棒性，展示了其更快的收敛速度（例如DiT-XL-2-$\mu$P比原始模型快2.9倍）。2. 在文本到图像生成任务中，成功将PixArt-$\alpha$从0.04B扩展到0.61B，MMDiT从0.18B扩展到18B，并显著优于基线模型。3. 调优成本极低，仅需5.5%的一次训练开销或3%的人工专家干预。4. 论文提供了大规模实验支持，并可能包含开源代码（未明确提及但符合社区惯例）。未来可进一步研究其他架构下的$\mu$P适配性。

相关研究包括：1. 原始$\mu$P方法在语言模型中的应用（如“MuP: A Parameterization for Scalable Learning of Neural Networks”）。2. 扩散模型的基础研究（如“Denoising Diffusion Probabilistic Models”）。3. 扩散Transformer的具体实现（如“DiT: Diffusion Models as Strong Baselines for Image Generation”、“PixArt-$\alpha$: Large-Scale Text-to-Image Generation”）。