Mise à l'échelle efficace des Transformers de Diffusion via μP

Scaling Diffusion Transformers Efficiently via μP

May 21, 2025

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

cs.AI

Résumé

Les Transformers de diffusion sont devenus la base des modèles génératifs visuels, mais leur scalabilité est limitée par le coût élevé du réglage des hyperparamètres (HP) à grande échelle. Récemment, la Paramétrisation à Mise à Jour Maximale (muP) a été proposée pour les Transformers classiques, permettant un transfert stable des HP des petits aux grands modèles de langage et réduisant considérablement les coûts de réglage. Cependant, il reste incertain si la muP des Transformers classiques s'étend aux Transformers de diffusion, qui diffèrent à la fois sur le plan architectural et objectif. Dans ce travail, nous généralisons la muP standard aux Transformers de diffusion et validons son efficacité à travers des expériences à grande échelle. Tout d'abord, nous prouvons rigoureusement que la muP des Transformers de diffusion dominants, incluant DiT, U-ViT, PixArt-alpha et MMDiT, s'aligne avec celle du Transformer classique, permettant l'application directe des méthodologies muP existantes. En exploitant ce résultat, nous démontrons systématiquement que DiT-muP bénéficie d'une robuste transférabilité des HP. Notamment, DiT-XL-2-muP avec un taux d'apprentissage transféré atteint une convergence 2,9 fois plus rapide que le DiT-XL-2 original. Enfin, nous validons l'efficacité de la muP sur la génération texte-image en faisant évoluer PixArt-alpha de 0,04B à 0,61B et MMDiT de 0,18B à 18B. Dans les deux cas, les modèles sous muP surpassent leurs bases de référence respectives tout en nécessitant un faible coût de réglage, seulement 5,5 % d'une session d'entraînement pour PixArt-alpha et 3 % de la consommation par des experts humains pour MMDiT-18B. Ces résultats établissent la muP comme un cadre à la fois rigoureux et efficace pour la mise à l'échelle des Transformers de diffusion.

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.