Escalando Transformadores de Difusión de Manera Eficiente mediante μP

Scaling Diffusion Transformers Efficiently via μP

May 21, 2025

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

cs.AI

Resumen

Los Transformadores de Difusión han surgido como la base para los modelos generativos de visión, pero su escalabilidad está limitada por el alto costo del ajuste de hiperparámetros (HP) a gran escala. Recientemente, se propuso la Parametrización de Actualización Máxima (muP) para los Transformadores estándar, lo que permite una transferencia estable de HP desde modelos de lenguaje pequeños a grandes y reduce drásticamente los costos de ajuste. Sin embargo, sigue sin estar claro si la muP de los Transformadores estándar se extiende a los Transformadores de Difusión, que difieren tanto arquitectónica como objetivamente. En este trabajo, generalizamos la muP estándar a los Transformadores de Difusión y validamos su efectividad mediante experimentos a gran escala. Primero, demostramos rigurosamente que la muP de los principales Transformadores de Difusión, incluyendo DiT, U-ViT, PixArt-alpha y MMDiT, se alinea con la del Transformador estándar, permitiendo la aplicación directa de las metodologías muP existentes. Aprovechando este resultado, demostramos sistemáticamente que DiT-muP disfruta de una robusta transferibilidad de HP. Notablemente, DiT-XL-2-muP con una tasa de aprendizaje transferida logra una convergencia 2.9 veces más rápida que el DiT-XL-2 original. Finalmente, validamos la efectividad de muP en la generación de texto a imagen escalando PixArt-alpha de 0.04B a 0.61B y MMDiT de 0.18B a 18B. En ambos casos, los modelos bajo muP superan a sus respectivas líneas base mientras requieren un costo de ajuste mínimo, solo el 5.5% de una ejecución de entrenamiento para PixArt-alpha y el 3% del consumo por expertos humanos para MMDiT-18B. Estos resultados establecen a muP como un marco principiado y eficiente para escalar Transformadores de Difusión.

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.