El paper original de DiT (Peebles & Xie, 2023) mostró que simplemente reemplazar la U-Net con un Transformer estándar y escalarlo producía mejor calidad de imagen. El Transformer procesa parches de imagen (similar a los Vision Transformers) con condicionamiento añadido del timestep de difusión y las etiquetas de clase. El hallazgo clave: DiT sigue leyes de escalado claras — modelos más grandes y más cómputo producen imágenes predeciblemente mejores, igual que con los LLMs.
Las U-Nets procesan imágenes a múltiples resoluciones, reduciendo y luego aumentando la resolución con conexiones residuales. Este sesgo inductivo era útil cuando el cómputo era limitado, pero introduce complejidad arquitectónica y no escala tan limpiamente. Los Transformers, con su arquitectura uniforme, son más simples de escalar y se benefician más del cómputo y datos adicionales. La compensación: los Transformers consumen más memoria debido a la atención cuadrática sobre todos los parches de imagen.
Stable Diffusion 3 y Flux usan MM-DiT (Multi-Modal DiT), que procesa tokens de texto e imagen a través de flujos separados que interactúan mediante atención cruzada. Esto es más efectivo que el condicionamiento simple de texto usado en el DiT original. El flujo de texto usa un encoder de texto congelado (como T5 o CLIP), y el flujo de imagen usa el proceso de difusión. Los dos flujos intercambian información en cada bloque de Transformer.