Scalabilità Efficiente dei Diffusion Transformer tramite μP

Abstract

I Diffusion Transformer si sono affermati come fondamento per i modelli generativi visivi, ma la loro scalabilità è limitata dall'elevato costo della sintonizzazione degli iperparametri (HP) su larga scala. Recentemente, è stata proposta la Parametrizzazione ad Aggiornamento Massimo (muP) per i Transformer standard, che consente un trasferimento stabile degli HP da modelli linguistici piccoli a quelli grandi, riducendo drasticamente i costi di sintonizzazione. Tuttavia, rimane incerto se la muP dei Transformer standard si estenda ai Diffusion Transformer, che differiscono sia a livello architetturale che obiettivo. In questo lavoro, generalizziamo la muP standard ai Diffusion Transformer e ne validiamo l'efficacia attraverso esperimenti su larga scala. In primo luogo, dimostriamo rigorosamente che la muP dei principali Diffusion Transformer, tra cui DiT, U-ViT, PixArt-alpha e MMDiT, è allineata con quella del Transformer standard, consentendo l'applicazione diretta delle metodologie muP esistenti. Sfruttando questo risultato, dimostriamo sistematicamente che DiT-muP gode di una robusta trasferibilità degli HP. In particolare, DiT-XL-2-muP con tasso di apprendimento trasferito raggiunge una convergenza 2,9 volte più veloce rispetto al DiT-XL-2 originale. Infine, validiamo l'efficacia della muP nella generazione di immagini da testo scalando PixArt-alpha da 0,04B a 0,61B e MMDiT da 0,18B a 18B. In entrambi i casi, i modelli con muP superano i rispettivi baseline pur richiedendo un costo di sintonizzazione ridotto, solo il 5,5% di una singola esecuzione di addestramento per PixArt-alpha e il 3% del consumo da parte di esperti umani per MMDiT-18B. Questi risultati stabiliscono la muP come un framework principiato ed efficiente per la scalabilità dei 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.

Scalabilità Efficiente dei Diffusion Transformer tramite μP

Scaling Diffusion Transformers Efficiently via μP

Abstract

Support