Эффективное масштабирование диффузионных трансформаторов с использованием μ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

Аннотация

Диффузионные трансформеры стали основой для генеративных моделей в области компьютерного зрения, но их масштабируемость ограничена высокой стоимостью настройки гиперпараметров (HP) на больших масштабах. Недавно была предложена Максимальная Параметризация Обновлений (muP) для стандартных трансформеров, которая обеспечивает стабильный перенос HP от небольших к крупным языковым моделям и значительно снижает затраты на настройку. Однако остается неясным, применима ли muP для стандартных трансформеров к диффузионным трансформерам, которые отличаются как архитектурно, так и по целям. В данной работе мы обобщаем стандартную muP для диффузионных трансформеров и подтверждаем ее эффективность в ходе масштабных экспериментов. Во-первых, мы строго доказываем, что muP для основных диффузионных трансформеров, включая DiT, U-ViT, PixArt-alpha и MMDiT, согласуется с muP для стандартного трансформера, что позволяет напрямую применять существующие методологии muP. Используя этот результат, мы систематически демонстрируем, что DiT-muP обладает устойчивой переносимостью HP. В частности, DiT-XL-2-muP с перенесенным коэффициентом обучения достигает сходимости в 2.9 раза быстрее, чем оригинальный DiT-XL-2. Наконец, мы подтверждаем эффективность muP в задаче генерации изображений по тексту, масштабируя PixArt-alpha с 0.04B до 0.61B и MMDiT с 0.18B до 18B. В обоих случаях модели с использованием muP превосходят свои базовые версии при минимальных затратах на настройку: всего 5.5% от одного цикла обучения для PixArt-alpha и 3% от затрат, требуемых экспертами для MMDiT-18B. Эти результаты подтверждают muP как принципиальный и эффективный фреймворк для масштабирования диффузионных трансформеров.

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.