μP를 통해 확산 트랜스포머를 효율적으로 확장하기

초록

Diffusion Transformer는 비전 생성 모델의 기반으로 부상했지만, 대규모에서의 하이퍼파라미터(HP) 튜닝 비용이 높아 확장성이 제한적입니다. 최근, 일반적인 Transformer를 위한 Maximal Update Parametrization(muP)이 제안되었으며, 이는 소규모에서 대규모 언어 모델로의 안정적인 HP 전이를 가능하게 하고 튜닝 비용을 크게 줄입니다. 그러나 일반적인 Transformer의 muP가 아키텍처와 목적이 다른 diffusion Transformer에도 적용될 수 있는지는 여전히 불분명합니다. 본 연구에서는 표준 muP를 diffusion Transformer로 일반화하고 대규모 실험을 통해 그 효과를 검증합니다. 먼저, DiT, U-ViT, PixArt-alpha, MMDiT를 포함한 주류 diffusion Transformer의 muP가 일반적인 Transformer와 일치함을 엄밀히 증명하여, 기존 muP 방법론을 직접 적용할 수 있음을 보입니다. 이를 활용하여 DiT-muP가 강력한 HP 전이성을 갖는다는 것을 체계적으로 입증합니다. 특히, 전이된 학습률을 사용한 DiT-XL-2-muP는 원래의 DiT-XL-2보다 2.9배 빠른 수렴 속도를 달성합니다. 마지막으로, PixArt-alpha를 0.04B에서 0.61B로, MMDiT를 0.18B에서 18B로 확장하면서 텍스트-이미지 생성에서 muP의 효과를 검증합니다. 두 경우 모두, muP를 적용한 모델이 각각의 기준선을 능가하면서도 적은 튜닝 비용만을 요구합니다. PixArt-alpha의 경우 한 번의 학습 실행 비용의 5.5%, MMDiT-18B의 경우 인간 전문가가 소비하는 비용의 3%만 필요했습니다. 이러한 결과는 muP가 diffusion Transformer를 확장하기 위한 원칙적이고 효율적인 프레임워크임을 입증합니다.

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.

μP를 통해 확산 트랜스포머를 효율적으로 확장하기

Scaling Diffusion Transformers Efficiently via μP

초록

Support