Diffusion Transformer (DiT)

Un Diffusion Transformer (DiT) es un modelo de difusión cuya red de eliminación de ruido es un transformer en lugar de una U-Net convolucional. Divide el latente de una imagen ruidosa en parches, los procesa como una secuencia de tokens mediante autoatención, del mismo modo que un Vision Transformer (ViT) procesa una imagen, y predice el ruido que debe eliminar en cada paso. William Peebles y Saining Xie lo presentaron en Scalable Diffusion Models with Transformers (ICCV 2023), donde DiT-XL/2 alcanzó una distancia de Inception de Fréchet (FID) de 2,27 en ImageNet 256×256 con condicionamiento por clase, superando el mejor resultado anterior de 3,60 de la difusión latente. El artículo demostró que la calidad de imagen mejora de forma constante a medida que aumenta el cómputo de un DiT, y desde entonces los transformadores de eliminación de ruido han sustituido a las U-Net en la mayoría de los grandes generadores de imágenes y vídeo, incluidos Stable Diffusion 3, FLUX.1 y Sora de OpenAI.

Cómo funciona un Diffusion Transformer#

Un DiT funciona en el espacio latente del autocodificador variacional (VAE) preentrenado que utiliza Stable Diffusion, cuyo codificador reduce la resolución por un factor de 8, de modo que una imagen de 256×256×3 se convierte en un latente de 32×32×4. Una capa de división en parches separa el latente ruidoso en parches de p×p y proyecta linealmente cada uno para convertirlo en un token; los tamaños de parche 2, 4 y 8 producen 256, 64 y 16 tokens. Las incrustaciones posicionales seno-coseno indican de dónde procede cada token, una pila de bloques transformer procesa la secuencia y una capa lineal final convierte cada token en un parche de ruido y varianza predichos, que se vuelven a ensamblar con la forma del latente. El paso temporal de difusión y la etiqueta de clase, o la instrucción de texto en un modelo de texto a imagen, condicionan cada bloque. La generación repite esta predicción desde ruido puro hasta el último paso, y el decodificador VAE convierte el latente limpio en píxeles:

El artículo comparó cuatro formas de incorporar el condicionamiento en cada bloque: tokens adicionales en la secuencia, atención cruzada, normalización adaptativa de capa (adaLN) y adaLN-Zero. adaLN-Zero obtuvo el FID más bajo y también fue la opción más eficiente en cómputo. Una red pequeña lee la suma de las incrustaciones del paso temporal y de la clase, y predice un desplazamiento y una escala que se aplican después de cada normalización de capa, además de una compuerta que escala cada rama residual. Las compuertas empiezan en cero, así que cada bloque comienza el entrenamiento como una función identidad. Este fragmento de código de PyTorch divide un latente en parches y aplica la rama de atención de un bloque:

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Imprime tokens (1, 256, 384), block is identity at initialization: True. Los 256 tokens de anchura 384 coinciden con el modelo más pequeño, DiT-S/2, y la compuerta a cero deja pasar la entrada sin cambios hasta que el entrenamiento la modifica. Un bloque DiT completo predice un segundo desplazamiento, escala y compuerta para su rama MLP.

Escalado con tokens y cómputo#

El artículo entrenó cuatro tamaños de modelo, desde DiT-S (12 capas, anchura 384) hasta DiT-XL (28 capas, anchura 1.152), cada uno con tamaños de parche 8, 4 y 2. En los 12 modelos, el FID se correspondía con el cómputo de la pasada hacia delante en Gflops, tanto si el cómputo procedía de una red más profunda o más ancha como de un mayor número de tokens. Reducir a la mitad el tamaño del parche cuadruplica el número de tokens y el cómputo, y apenas añade parámetros. DiT-XL/2 tiene 675 millones de parámetros, sin contar los 84 millones del VAE, y usa 119 Gflops por pasada hacia delante a 256×256. A 512×512 procesa 1.024 tokens de un latente de 64×64×4 con 525 Gflops y alcanza un FID de 3,04. Como la autoatención compara cada token con todos los demás, las resoluciones más altas y los vídeos más largos elevan su coste cuadráticamente. Ultralytics YOLO26 evita este coste al aplicar atención solo en su etapa a escala 1/32, que tiene 400 posiciones para una imagen de 640×640, y mantener las convoluciones en el resto para la detección en tiempo real.

DiT frente a U-Net y ViT#

• Eliminación de ruido con U-Net: Las U-Net de DDPM y Stable Diffusion 1.x y 2.x son codificadores-decodificadores convolucionales que funcionan a varias resoluciones, con conexiones de salto entre niveles equivalentes y atención añadida a resoluciones más bajas. Un DiT mantiene una única resolución de tokens, aplica atención global en cada bloque y tiene menos decisiones de diseño específicas para imágenes, por lo que escala como los transformers de los grandes modelos de lenguaje.
• Vision Transformer: ViT codifica una imagen para tareas como la clasificación. DiT reutiliza su incrustación de parches y sus bloques transformer, añade condicionamiento por el paso temporal mediante adaLN-Zero y genera una predicción para cada parche.
• Difusión latente y coincidencia de flujo: Un modelo de difusión latente es el marco general para eliminar ruido en el espacio latente de un VAE, y la difusión o la coincidencia de flujo es el objetivo de entrenamiento. DiT es la red de eliminación de ruido dentro de ese marco, y los generadores actuales combinan arquitecturas DiT con objetivos de coincidencia de flujo.

DiT de texto a imagen y de vídeo#

• DiT multimodal (MMDiT): Stable Diffusion 3 procesa tokens de imagen y de texto con pesos separados para cada modalidad y los combina en una única operación de atención, de modo que la información fluye en ambos sentidos entre la instrucción y la imagen. Esto mejoró la tipografía y el seguimiento de instrucciones, y Stable Diffusion 3.5 Large amplía el diseño hasta unos 8.000 millones de parámetros.
• FLUX.1: El modelo FLUX.1 de Black Forest Labs es un transformer de flujo rectificado con 12.000 millones de parámetros para la generación de texto a imagen.
• Vídeo: El informe técnico de Sora de OpenAI describe un transformer de difusión que comprime los vídeos en un espacio latente y los divide en parches espaciotemporales, de modo que los clips de distintas duraciones, resoluciones y relaciones de aspecto se convierten en secuencias de tokens de distintas longitudes. Para la generación de vídeo, el modelo Movie Gen más grande de Meta es un transformer de 30.000 millones de parámetros cuyo contexto de 73.000 tokens de vídeo corresponde a 16 segundos de vídeo a 16 fotogramas por segundo.