Effiziente Skalierung von Diffusions-Transformatoren mittels μP

Scaling Diffusion Transformers Efficiently via μP

May 21, 2025

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

cs.AI

Zusammenfassung

Diffusion Transformers haben sich als Grundlage für visuelle generative Modelle etabliert, doch ihre Skalierbarkeit wird durch die hohen Kosten der Hyperparameteroptimierung (HP) in großem Maßstab eingeschränkt. Kürzlich wurde die Maximal Update Parametrisierung (muP) für Standard-Transformers vorgeschlagen, die einen stabilen HP-Transfer von kleinen zu großen Sprachmodellen ermöglicht und die Optimierungskosten drastisch reduziert. Es bleibt jedoch unklar, ob muP für Standard-Transformers auch auf Diffusion Transformers anwendbar ist, die sich sowohl architektonisch als auch in ihren Zielsetzungen unterscheiden. In dieser Arbeit verallgemeinern wir die Standard-muP auf Diffusion Transformers und validieren ihre Wirksamkeit durch groß angelegte Experimente. Zunächst beweisen wir rigoros, dass die muP von gängigen Diffusion Transformers, einschließlich DiT, U-ViT, PixArt-alpha und MMDiT, mit der des Standard-Transformers übereinstimmt, was die direkte Anwendung bestehender muP-Methoden ermöglicht. Basierend auf diesem Ergebnis zeigen wir systematisch, dass DiT-muP eine robuste HP-Übertragbarkeit aufweist. Insbesondere erreicht DiT-XL-2-muP mit übertragener Lernrate eine 2,9-mal schnellere Konvergenz als das ursprüngliche DiT-XL-2. Schließlich validieren wir die Wirksamkeit von muP bei der Text-zu-Bild-Generierung, indem wir PixArt-alpha von 0,04B auf 0,61B und MMDiT von 0,18B auf 18B skalieren. In beiden Fällen übertreffen Modelle unter muP ihre jeweiligen Baselines bei geringen Optimierungskosten, die nur 5,5 % eines Trainingslaufs für PixArt-alpha und 3 % des Verbrauchs durch menschliche Experten für MMDiT-18B betragen. Diese Ergebnisse etablieren muP als ein prinzipielles und effizientes Framework für die Skalierung von Diffusion Transformers.

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.