O paper original do DiT (Peebles & Xie, 2023) mostrou que simplesmente substituir a U-Net por um Transformer padrão e escalá-lo produzia melhor qualidade de imagem. O Transformer processa patches de imagem (similar a Vision Transformers) com condicionamento adicionado do timestep de difusão e rótulos de classe. A descoberta chave: DiT segue leis de escalonamento claras — modelos maiores e mais computação produzem imagens previsivelmente melhores, assim como com LLMs.
U-Nets processam imagens em múltiplas resoluções, reduzindo e depois aumentando a resolução com skip connections. Esse viés indutivo era útil quando a computação era limitada, mas introduz complexidade arquitetural e não escala tão limpo. Transformers, com sua arquitetura uniforme, são mais simples de escalar e se beneficiam mais de computação e dados adicionais. O trade-off: Transformers consomem mais memória devido à attention quadrática sobre todos os patches de imagem.
Stable Diffusion 3 e Flux usam MM-DiT (Multi-Modal DiT), que processa tokens de texto e imagem através de fluxos separados que interagem via cross-attention. Isso é mais eficaz que o condicionamento de texto mais simples usado no DiT original. O fluxo de texto usa um text encoder congelado (como T5 ou CLIP), e o fluxo de imagem usa o processo de difusão. Os dois fluxos trocam informação em cada bloco Transformer.