Difusión Latente y el VAE

¿Por qué no difundir en el espacio de píxeles?

Una imagen RGB de 512×512 es un tensor de forma $512 \times 512 \times 3 = 786{,}432$ valores. En cada paso de eliminación de ruido, el modelo de difusión debe procesarlos todos: introducirlos, predecir el ruido, restarlo y repetir. ¿Qué tan mal se pone esto? Con 1000 pasos de eliminación de ruido (un schedule típico de DDPM), eso son 786 millones de operaciones a nivel de valor solo para una imagen. Escala la resolución a 1024×1024 y estamos en $1024 \times 1024 \times 3 = 3{,}145{,}728$ valores por paso — más de 3 mil millones de operaciones a lo largo de toda la cadena.

Pero aquí está la idea clave: la mayoría del detalle a nivel de píxel es perceptualmente redundante . Los píxeles vecinos en una fotografía están altamente correlacionados — un parche de cielo azul no contiene 64 piezas independientes de información solo porque abarca 64 píxeles. Los detalles de textura de alta frecuencia (grano de ruido exacto, gradientes de color sutiles) que consumen la mayoría de esos 786K valores contribuyen muy poco a lo que realmente percibimos como el "contenido" de la imagen.

Este es el mismo principio detrás de la compresión JPEG: una imagen puede comprimirse a una fracción de su tamaño sin pérdida perceptual apreciable, precisamente porque el espacio de píxeles es masivamente sobredeterminado. Entonces, ¿qué pasa si primero comprimimos la imagen a una representación mucho más pequeña, ejecutamos todo el proceso de difusión allí , y luego descomprimimos el resultado de vuelta a píxeles? Eso es exactamente lo que hace la difusión latente.

El Autoencoder Variacional (VAE)

El motor de compresión detrás de la difusión latente es un Autoencoder Variacional (VAE) (Rombach et al., 2022) . Un autoencoder es una red neuronal con dos mitades: un codificador que comprime una entrada en una representación interna compacta (el "cuello de botella"), y un decodificador que reconstruye la entrada original a partir de esa forma compacta. La parte "variacional" significa que el codificador no produce un único vector fijo sino los parámetros de una distribución de probabilidad (una media y varianza), de la cual muestreamos la representación del cuello de botella. Esto hace que el espacio latente sea suave y continuo — puntos cercanos decodifican a imágenes similares — lo cual es crucial para que la difusión funcione bien.

El VAE tiene dos componentes. El codificador $\mathcal{E}$ mapea una imagen a una representación latente:

Para una imagen de entrada de 512×512, el codificador típicamente produce un latente de $64 \times 64 \times 4$ — dimensiones espaciales reducidas por 8× en cada dirección, con 4 canales latentes en lugar de 3 canales de color. El decodificador $\mathcal{D}$ invierte esto:

La razón de compresión es dramática. La imagen original tiene $512 \times 512 \times 3 = 786{,}432$ valores. El latente tiene $64 \times 64 \times 4 = 16{,}384$ valores. Eso es una reducción de 48× . El modelo de difusión ahora opera sobre 16K valores por paso en lugar de 786K.

# Compression ratio: pixel space vs latent space
pixel_h, pixel_w, pixel_c = 512, 512, 3
latent_h, latent_w, latent_c = 64, 64, 4

pixel_values = pixel_h * pixel_w * pixel_c
latent_values = latent_h * latent_w * latent_c
ratio = pixel_values / latent_values

print(f"Pixel space:  {pixel_h}x{pixel_w}x{pixel_c} = {pixel_values:,} values")
print(f"Latent space: {latent_h}x{latent_w}x{latent_c} = {latent_values:,} values")
print(f"Compression ratio: {ratio:.0f}x fewer values")
print()

# At 1024x1024
pixel_1024 = 1024 * 1024 * 3
latent_1024 = 128 * 128 * 4
ratio_1024 = pixel_1024 / latent_1024
print(f"At 1024x1024:")
print(f"  Pixel space:  {pixel_1024:,} values")
print(f"  Latent space: {latent_1024:,} values")
print(f"  Compression ratio: {ratio_1024:.0f}x")

El VAE se entrena con una pérdida combinada que equilibra dos objetivos:

El primer término es la pérdida de reconstrucción : el error cuadrático medio entre la imagen original $x$ y la reconstrucción decodificada $\mathcal{D}(\mathcal{E}(x))$. Esto es lo que obliga al autoencoder a preservar información. Si este término es cero, el decodificador reconstruye perfectamente cada píxel. Si es grande, el cuello de botella es demasiado agresivo y pierde detalle crítico.

El segundo término es el regularizador de divergencia KL (ver teoría de la información para un tratamiento más profundo de la divergencia KL). Aquí $q(z|x)$ es la distribución que produce el codificador (una gaussiana con media y varianza aprendidas para cada entrada $x$), y $p(z)$ es un prior gaussiano estándar $\mathcal{N}(0, I)$. Este término penaliza al codificador por producir distribuciones latentes que se desvían de una gaussiana estándar.

¿Por qué importa esto? Sin regularización KL, el codificador puede "hacer trampa" esparciendo los códigos latentes lejos unos de otros en regiones arbitrarias del espacio, dejando vastas zonas muertas donde la decodificación produce basura. El término KL fuerza al espacio latente a ser compacto y con forma gaussiana, de modo que cualquier punto muestreado de $\mathcal{N}(0, I)$ decodifica a algo razonable. Verifiquemos el comportamiento en los límites: cuando $q(z|x)$ es exactamente igual a $p(z)$ (el codificador produce una gaussiana estándar independientemente de la entrada), $D_{\text{KL}} = 0$ — sin penalización, pero tampoco se codifica información sobre $x$, así que la reconstrucción es terrible. Cuando $q(z|x)$ es un pico muy estrecho lejos de cero (el codificador memoriza cada entrada como un punto único), $D_{\text{KL}}$ crece mucho — penalización fuerte. El óptimo está en el medio: codificar lo suficiente para reconstruir, pero permanecer cerca de la gaussiana.

En la práctica, el peso $\lambda$ se establece muy pequeño (alrededor de $10^{-6}$). La calidad de reconstrucción perceptual importa mucho más que la gaussianidad perfecta — el objetivo es un compresor de imágenes útil, no un modelo generativo en el sentido de VAE. El VAE de Stable Diffusion también añade una pérdida perceptual (comparando características VGG de imágenes reales vs reconstruidas) y una pérdida adversaria (un discriminador que penaliza reconstrucciones borrosas), pero la fórmula anterior captura la idea central.

Difusión en el espacio latente

Con un VAE entrenado en mano, todo el proceso de difusión se mueve del espacio de píxeles al espacio latente. En lugar de trabajar con $x_0 \in \mathbb{R}^{512 \times 512 \times 3}$, codificamos la imagen de entrenamiento y trabajamos con su latente:

El proceso directo añade ruido gaussiano a $z_0$ usando exactamente las mismas matemáticas que la difusión estándar (cubierta en el artículo 1), solo aplicadas al latente en lugar del tensor de píxeles:

El schedule de ruido $\bar{\alpha}_t$ funciona idénticamente: en $t=0$, $\bar{\alpha}_0 \approx 1$ así que $z_0$ está casi limpio; en $t=T$, $\bar{\alpha}_T \approx 0$ así que $z_T$ es casi ruido puro. El modelo $\epsilon_\theta(z_t, t)$ se entrena para predecir el ruido $\epsilon$ que se añadió, usando el mismo objetivo MSE:

En tiempo de inferencia, partimos de ruido puro $z_T \sim \mathcal{N}(0, I)$ en el espacio latente, iterativamente eliminamos ruido para recuperar $z_0$, y luego decodificamos de vuelta a píxeles:

Los beneficios de trabajar en el espacio latente van más allá de la velocidad bruta:

• 48× menos valores: cada pase hacia adelante a través de la red de eliminación de ruido es dramáticamente más barato. El entrenamiento converge más rápido y el muestreo requiere menos cómputo por paso.
• Espacio más suave y semántico: el codificador VAE elimina las correlaciones redundantes de píxeles perceptualmente. El espacio latente es más "denso en significado" — pequeños movimientos corresponden a cambios visuales significativos en lugar de desplazamientos imperceptibles de píxeles.
• División del trabajo: el VAE maneja detalles perceptuales de bajo nivel (texturas, patrones finos, valores exactos de píxeles) mientras el modelo de difusión se enfoca en la estructura semántica de alto nivel (composición, objetos, estilo). Cada componente hace lo que mejor sabe hacer.

El eliminador de ruido U-Net

La red neuronal $\epsilon_\theta$ que predice el ruido en el espacio latente es, en el Stable Diffusion original, una U-Net . La U-Net fue originalmente diseñada para segmentación de imágenes médicas (Ronneberger et al., 2015) , pero su estructura resultó ser ideal para la eliminación de ruido. Es una arquitectura codificador-decodificador con conexiones residuales.

La ruta del codificador reduce progresivamente las dimensiones espaciales a través de una serie de bloques de convolución residual seguidos de operaciones de submuestreo. Un latente de $64 \times 64$ podría pasar por resoluciones $64 \to 32 \to 16 \to 8$. A medida que la resolución espacial se reduce, el número de canales crece, de modo que la red captura características cada vez más globales y abstractas en cada nivel.

La ruta del decodificador revierte esto, aumentando progresivamente la resolución de vuelta a la resolución latente original: $8 \to 16 \to 32 \to 64$. Pero el sobremuestreo solo perdería detalle espacial fino. Ahí es donde entran las conexiones residuales : en cada escala, los mapas de características del codificador se concatenan con los del decodificador. El codificador dice "así es como se ve el detalle fino a esta resolución" y el decodificador usa eso junto con su contexto global sobremuestreado para producir una salida refinada.

Dos mecanismos adicionales son críticos. Primero, el condicionamiento por paso de tiempo : el modelo debe saber en qué paso de tiempo $t$ está eliminando ruido, porque el nivel de ruido (y por lo tanto la estrategia de eliminación de ruido) difiere dramáticamente entre $t=1000$ (casi ruido puro, enfoque en estructura global) y $t=10$ (casi limpio, enfoque en detalle fino). El paso de tiempo se codifica usando un embedding sinusoidal (similar a la codificación posicional en transformers), se proyecta a través de un MLP, y se añade a cada bloque residual.

Segundo, se insertan capas de auto-atención en ciertas resoluciones (típicamente $32 \times 32$, $16 \times 16$ y $8 \times 8$). Las convoluciones son operaciones locales — cada píxel de salida solo ve un vecindario pequeño. La auto-atención permite que cada posición espacial atienda a todas las demás, dando al modelo contexto global. Esto es esencial para la coherencia estructural: asegurar que un rostro tenga dos ojos, que las ventanas de un edificio estén espaciadas uniformemente, que la composición general sea consistente.

Condicionamiento de texto vía atención cruzada

Hasta ahora, el eliminador de ruido genera imágenes de forma incondicional — puede eliminar ruido, pero no tiene idea de lo que queremos que produzca. Para generar imágenes a partir de prompts de texto, necesitamos una forma de inyectar información lingüística en la U-Net. El mecanismo es la atención cruzada (ver atención codificador-decodificador para el concepto general).

Primero, el prompt de texto es procesado por un codificador de texto separado y congelado (como CLIP o T5) que convierte la cadena en una secuencia de vectores de embedding $\tau \in \mathbb{R}^{L \times d_{\text{text}}}$, donde $L$ es el número de tokens y $d_{\text{text}}$ es la dimensión del embedding. Estos embeddings de texto se inyectan luego en la U-Net en múltiples resoluciones vía atención cruzada:

donde $Q = W_Q \cdot \phi(z_t)$ proviene de las características de la imagen (las activaciones intermedias de la U-Net en una capa dada), y $K = W_K \cdot \tau$ y $V = W_V \cdot \tau$ provienen de los embeddings de texto. El denominador $\sqrt{d_k}$ evita que los productos punto crezcan demasiado a medida que la dimensión $d_k$ aumenta: sin él, para $d_k$ alto los productos punto serían tan grandes que softmax se satura a vectores casi one-hot, matando el flujo de gradientes. Con el escalado, la varianza de $QK^T$ se mantiene aproximadamente en 1 independientemente de $d_k$.

Lo que esto significa intuitivamente: cada posición espacial en las características de imagen consulta la descripción textual, preguntando "¿qué palabras son relevantes para lo que debo generar aquí?" Una posición espacial correspondiente a la región del cielo atenderá fuertemente a la palabra "atardecer"; una posición cerca de un rostro atenderá a "retrato" o "mujer". El modelo aprende qué tokens de texto importan para qué ubicaciones espaciales.

Las capas de atención cruzada aparecen en múltiples resoluciones en la U-Net (típicamente junto a las capas de auto-atención en $32 \times 32$, $16 \times 16$ y $8 \times 8$). Resoluciones más bajas capturan la alineación gruesa texto-composición ("un gato a la izquierda, un perro a la derecha"), mientras que resoluciones más altas capturan detalle más fino ("ojos azules", "pelaje rayado").

La elección del codificador de texto importa enormemente para la comprensión del prompt. Diferentes versiones de Stable Diffusion usan diferentes codificadores de texto:

• SD 1.x: CLIP ViT-L/14 — contexto máximo de 77 tokens, embeddings de 768 dimensiones. El original, pero limitado en vocabulario y longitud de contexto.
• SD 2.x: OpenCLIP ViT-H/14 — embeddings de 1024 dimensiones, entrenado en un dataset más grande. Mejor comprensión de texto, pero la comunidad encontró que era más difícil de dirigir mediante prompts debido a diferencias en los datos de entrenamiento.
• SDXL: codificadores de texto duales — CLIP ViT-L y OpenCLIP ViT-bigG, con sus salidas concatenadas para producir embeddings de 2048 dimensiones (Podell et al., 2023) . Dos codificadores capturan aspectos complementarios del texto: uno entrenado en pares imagen-texto curados (CLIP), el otro en un dataset más amplio (OpenCLIP).

import json, js

versions = [
    ("SD 1.5", "768", "77", "1"),
    ("SD 2.1", "1024", "77", "1"),
    ("SDXL",   "2048", "77", "2"),
]

js.window.py_table_data = json.dumps({
    "headers": ["Version", "Embed Dim", "Max Tokens", "Encoders"],
    "rows": [list(v) for v in versions]
})

print("SDXL's 2048-dim conditioning = 2.7x richer than SD 1.5's 768-dim.")

Stable Diffusion: la arquitectura completa

Uniendo todo, el pipeline de Stable Diffusion consiste en tres componentes entrenados independientemente que trabajan en secuencia:

• Codificador de texto (congelado): convierte el prompt de texto en una secuencia de vectores de embedding. Entrenado por separado (por ejemplo, CLIP fue entrenado en 400M de pares imagen-texto mediante aprendizaje contrastivo).
• Eliminador de ruido U-Net (el modelo central): toma un latente ruidoso $z_t$ y un paso de tiempo $t$, recibe embeddings de texto vía atención cruzada, y predice el ruido $\epsilon$. Este es el único componente que se entrena durante el proceso de entrenamiento de difusión latente.
• Decodificador VAE (congelado): convierte el latente final sin ruido $z_0$ de vuelta a una imagen en píxeles. Entrenado por separado como se describió anteriormente.

El pipeline de inferencia fluye así:

# Stable Diffusion inference (pseudocode)
# 1. Encode text prompt
text_embeddings = text_encoder(prompt)        # shape: (77, 768) for SD 1.5

# 2. Start from pure noise in latent space
z_T = torch.randn(1, 4, 64, 64)              # random latent for 512x512 output

# 3. Iterative denoising (e.g., 50 steps with a DDIM scheduler)
z_t = z_T
for t in reversed(scheduler.timesteps):       # T, T-1, ..., 1
    noise_pred = unet(z_t, t, text_embeddings)   # predict noise
    z_t = scheduler.step(noise_pred, t, z_t)     # remove predicted noise

# 4. Decode latent to pixel image
image = vae.decode(z_t)                       # shape: (1, 3, 512, 512)

Observa cómo los tres componentes tienen cantidades de parámetros y procedimientos de entrenamiento completamente diferentes, pero trabajan juntos en la inferencia:

import json, js

models = [
    ("SD 1.5", 860, 123, 84, "512x512", "~4 GB"),
    ("SD 2.1", 865, 354, 84, "768x768", "~5 GB"),
    ("SDXL",  2600, 817, 84, "1024x1024", "~7 GB"),
]

rows = []
for name, unet, text, vae, res, vram in models:
    total = unet + text + vae
    rows.append([name, f"{unet:,}", f"{text:,}", f"{vae}", f"{total:,}", res, vram])

js.window.py_table_data = json.dumps({
    "headers": ["Model", "U-Net (M)", "Text (M)", "VAE (M)", "Total (M)", "Native Res", "VRAM FP16"],
    "rows": rows
})

print("SD 1.5 runs on 8 GB consumer GPUs — this democratized image generation.")

¿Por qué Stable Diffusion democratizó la generación de imágenes? Antes de él, los modelos de texto a imagen como DALL-E 2 eran de código cerrado y solo accesibles por API. Stable Diffusion fue lanzado con pesos de código abierto , y debido a que la difusión latente redujo los requisitos de cómputo tan dramáticamente, el modelo podía ejecutarse en una GPU de consumidor con solo 8 GB de VRAM. Esto permitió una explosión de desarrollo comunitario: modelos personalizados ajustados, adaptaciones LoRA, ControlNet para condicionamiento estructural, e interfaces y flujos de trabajo completamente nuevos — nada de lo cual habría sido posible con un modelo de difusión cerrado, en espacio de píxeles, que requiriera cientos de gigabytes de VRAM.

Quiz

Pon a prueba tu comprensión de la difusión latente y el VAE.