Escalonamento Eficiente de Transformadores de Difusão via μP

Resumo

Os Transformers de Difusão emergiram como a base para modelos generativos de visão, mas sua escalabilidade é limitada pelo alto custo de ajuste de hiperparâmetros (HP) em grandes escalas. Recentemente, a Parametrização de Atualização Máxima (muP) foi proposta para Transformers convencionais, permitindo a transferência estável de HP de modelos de linguagem pequenos para grandes e reduzindo drasticamente os custos de ajuste. No entanto, ainda não está claro se a muP de Transformers convencionais se estende aos Transformers de Difusão, que diferem arquitetural e objetivamente. Neste trabalho, generalizamos a muP padrão para Transformers de Difusão e validamos sua eficácia por meio de experimentos em grande escala. Primeiro, provamos rigorosamente que a muP dos principais Transformers de Difusão, incluindo DiT, U-ViT, PixArt-alpha e MMDiT, está alinhada com a do Transformer convencional, permitindo a aplicação direta das metodologias muP existentes. Aproveitando esse resultado, demonstramos sistematicamente que o DiT-muP possui uma robusta transferibilidade de HP. Notavelmente, o DiT-XL-2-muP com taxa de aprendizado transferida alcança uma convergência 2,9 vezes mais rápida que o DiT-XL-2 original. Por fim, validamos a eficácia da muP na geração de texto para imagem ao escalar o PixArt-alpha de 0,04B para 0,61B e o MMDiT de 0,18B para 18B. Em ambos os casos, os modelos sob muP superam suas respectivas linhas de base enquanto exigem um custo de ajuste mínimo, apenas 5,5% de uma execução de treinamento para o PixArt-alpha e 3% do consumo por especialistas humanos para o MMDiT-18B. Esses resultados estabelecem a muP como uma estrutura eficiente e fundamentada para escalar Transformers de Difusão.

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.

Escalonamento Eficiente de Transformadores de Difusão via μP

Scaling Diffusion Transformers Efficiently via μP

Resumo

Support