拡散トランスフォーマーの効率的なスケーリングをμPによって実現

Scaling Diffusion Transformers Efficiently via μP

May 21, 2025

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

cs.AI

要旨

Diffusion Transformersは視覚生成モデルの基盤として登場しましたが、大規模なハイパーパラメータ（HP）チューニングのコストが高いため、そのスケーラビリティは制限されています。最近、標準的なTransformerに対してMaximal Update Parametrization（muP）が提案され、小規模から大規模な言語モデルへの安定したHP転移を可能にし、チューニングコストを劇的に削減しました。しかし、標準的なTransformerのmuPが、アーキテクチャ的および目的的に異なるdiffusion Transformerにも適用可能かどうかは不明でした。本研究では、標準的なmuPをdiffusion Transformerに一般化し、大規模な実験を通じてその有効性を検証します。まず、DiT、U-ViT、PixArt-alpha、MMDiTなどの主流のdiffusion TransformerのmuPが標準的なTransformerのそれと一致することを厳密に証明し、既存のmuP手法を直接適用可能であることを示します。この結果を活用し、DiT-muPが堅牢なHP転移性を享受することを体系的に実証します。特に、学習率を転移したDiT-XL-2-muPは、元のDiT-XL-2よりも2.9倍速く収束します。最後に、テキストから画像生成におけるmuPの有効性を検証するため、PixArt-alphaを0.04Bから0.61Bに、MMDiTを0.18Bから18Bにスケーリングします。どちらの場合も、muP下のモデルはそれぞれのベースラインを上回りながら、PixArt-alphaでは1回のトレーニング実行の5.5%、MMDiT-18Bでは人間の専門家による消費量の3%という小さなチューニングコストしか必要としませんでした。これらの結果は、muPがdiffusion Transformerをスケーリングするための原理的かつ効率的なフレームワークであることを確立します。

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 (muP) 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 muP of vanilla Transformers extends to diffusion Transformers, which differ architecturally and objectively. In this work, we generalize standard muP to diffusion Transformers and validate its effectiveness through large-scale experiments. First, we rigorously prove that muP of mainstream diffusion Transformers, including DiT, U-ViT, PixArt-alpha, and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing muP methodologies. Leveraging this result, we systematically demonstrate that DiT-muP enjoys robust HP transferability. Notably, DiT-XL-2-muP with transferred learning rate achieves 2.9 times faster convergence than the original DiT-XL-2. Finally, we validate the effectiveness of muP on text-to-image generation by scaling PixArt-alpha from 0.04B to 0.61B and MMDiT from 0.18B to 18B. In both cases, models under muP outperform their respective baselines while requiring small tuning cost, only 5.5% of one training run for PixArt-alpha and 3% of consumption by human experts for MMDiT-18B. These results establish muP as a principled and efficient framework for scaling diffusion Transformers.