通过 μP 高效扩展扩散变换器

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