L'article original DiT (Peebles & Xie, 2023) a montré que simplement remplacer le U-Net par un Transformer standard et le mettre à l'échelle produisait une meilleure qualité d'image. Le Transformer traite des patches d'image (similaire aux Vision Transformers) avec un conditionnement ajouté par le pas de temps de diffusion et les étiquettes de classe. La découverte clé : DiT suit des lois d'échelle claires — des modèles plus grands et plus de calcul produisent des images prévisiblement meilleures, exactement comme avec les LLM.
Les U-Nets traitent les images à plusieurs résolutions, sous-échantillonnant puis sur-échantillonnant avec des connexions de raccourci. Ce biais inductif était utile quand le calcul était limité, mais il introduit une complexité architecturale et ne passe pas à l'échelle aussi proprement. Les Transformers, avec leur architecture uniforme, sont plus simples à mettre à l'échelle et bénéficient davantage de calcul et de données supplémentaires. Le compromis : les Transformers consomment plus de mémoire à cause de l'attention quadratique sur tous les patches d'image.
Stable Diffusion 3 et Flux utilisent MM-DiT (Multi-Modal DiT), qui traite les tokens de texte et d'image à travers des flux séparés qui interagissent via l'attention croisée. C'est plus efficace que le simple conditionnement textuel utilisé dans le DiT original. Le flux texte utilise un encodeur de texte gelé (comme T5 ou CLIP), et le flux image utilise le processus de diffusion. Les deux flux échangent de l'information à chaque bloc Transformer.