提出μP方法,让扩散Transformer高效缩放,大幅降低调参成本。
Scaling Diffusion Transformers Efficiently via $μ$P
- 将μP扩展至扩散Transformer架构,实现参数稳定迁移。
- DiT-XL-2-μP收敛速度提升2.9倍,仅需原模型5.5%调参成本。
- 适用于大规模文本到图像生成,适合追求高效训练的研究者。
扩散Transformer已成为视觉生成模型的核心,但其在大规模下的可扩展性受限于高昂的超参数调优成本。近期提出的最大更新参数化(μP)适用于普通Transformer,可实现小规模到大规模语言模型的稳定超参数迁移并显著降低调参开销。然而,μP能否适用于结构与目标不同的扩散Transformer尚不明确。本文首次将标准μP推广至主流扩散Transformer架构(包括U-ViT、DiT、PixArt-α和MMDiT),并从理论上严格证明其与普通Transformer的μP一致性,支持直接应用现有μP方法。基于此,系统验证了DiT-μP的强泛化能力:以迁移学习率训练的DiT-XL-2-μP,收敛速度比原始模型快2.9倍。进一步在文本到图像生成任务中验证,将PixArt-α从0.04B扩展至0.61B,MMDiT从0.18B扩展至18B,二者均在μP框架下表现优于基线模型,且调优成本分别仅为单次训练的5.5%和人类专家消耗的3%。结果确立μP为扩散Transformer高效缩放的原理性框架。
原文摘要 · Abstract (English)
Diffusion Transformers have emerged as the foundation for vision generative models, but their scalability is limited by the high cost of hyperparameter (HP) tuning at large scales. Recently, Maximal Update Parametrization ($μ$P) was proposed for vanilla Transformers, which enables stable HP transfer from small to large language models, and dramatically reduces tuning costs. However, it remains unclear whether $μ$P of vanilla Transformers extends to diffusion Transformers, which differ architecturally and objectively. In this work, we generalize standard $μ$P to diffusion Transformers and validate its effectiveness through large-scale experiments. First, we rigorously prove that $μ$P of mainstream diffusion Transformers, including U-ViT, DiT, PixArt-$α$, and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing $μ$P methodologies. Leveraging this result, we systematically demonstrate that DiT-$μ$P enjoys robust HP transferability. Notably, DiT-XL-2-$μ$P with transferred learning rate achieves 2.9 times faster convergence than the original DiT-XL-2. Finally, we validate the effectiveness of $μ$P on text-to-image generation by scaling PixArt-$α$ from 0.04B to 0.61B and MMDiT from 0.18B to 18B. In both cases, models under $μ$P outperform their respective baselines while requiring small tuning cost, only 5.5% of one training run for PixArt-$α$ and 3% of consumption by human experts for MMDiT-18B. These results establish $μ$P as a principled and efficient framework for scaling diffusion Transformers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。