Diffusie Transformers efficiënt schalen via μP

Samenvatting

Diffusion Transformers zijn opgekomen als de basis voor visuele generatieve modellen, maar hun schaalbaarheid wordt beperkt door de hoge kosten van hyperparameter (HP) afstemming op grote schaal. Recentelijk werd Maximal Update Parametrization (muP) voorgesteld voor standaard Transformers, wat een stabiele HP-overdracht van kleine naar grote taalmodel mogelijk maakt en de afstemkosten aanzienlijk reduceert. Het blijft echter onduidelijk of muP van standaard Transformers ook van toepassing is op diffusion Transformers, die zowel architectonisch als doelgericht verschillen. In dit werk generaliseren we standaard muP naar diffusion Transformers en valideren we de effectiviteit ervan via grootschalige experimenten. Eerst bewijzen we rigoureus dat muP van mainstream diffusion Transformers, waaronder DiT, U-ViT, PixArt-alpha en MMDiT, overeenkomt met die van de standaard Transformer, waardoor bestaande muP-methodologieën direct kunnen worden toegepast. Gebruikmakend van dit resultaat, tonen we systematisch aan dat DiT-muP robuuste HP-overdraagbaarheid geniet. Opmerkelijk is dat DiT-XL-2-muP met een overgedragen leercurve 2,9 keer sneller convergeert dan de originele DiT-XL-2. Tot slot valideren we de effectiviteit van muP op tekst-naar-beeldgeneratie door PixArt-alpha te schalen van 0,04B naar 0,61B en MMDiT van 0,18B naar 18B. In beide gevallen presteren modellen onder muP beter dan hun respectievelijke baselines, terwijl ze slechts minimale afstemkosten vereisen: slechts 5,5% van één trainingsrun voor PixArt-alpha en 3% van het verbruik door menselijke experts voor MMDiT-18B. Deze resultaten vestigen muP als een principieel en efficiënt raamwerk voor het schalen van 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.

Diffusie Transformers efficiënt schalen via μP

Scaling Diffusion Transformers Efficiently via μP

Samenvatting

Support