Phase 08 - Lesson 02

Autoencoders & Autoencoders Variacionales (VAE)

Un autoencoder común comprime y luego reconstruye. Memoriza. No genera. Agrega un truco — forzar el código a verse gaussiano — y obtienes un muestreador. Ese único truco, la reparametrización de z = μ + σ·ε, es la razón por la cual cada modelo de imagen de difusión latente y de coincidencia de flujo (flow-matching) que utilizas en 2026 tiene un VAE en la entrada.

Tipo: Construir Lenguajes: Python Prerrequisitos: Fase 3 · 02 (Backprop), Fase 3 · 07 (CNNs), Fase 8 · 01 (Taxonomía) Tiempo: ~75 minutos

El Problema

Comprime un dígito MNIST de 784 píxeles a un código de 16 números y luego reconstruye. Un autoencoder común superará con éxito el MSE de reconstrucción, pero el espacio de código será un desastre irregular. Elige un punto aleatorio en el espacio de código, decodifícalo y obtendrás ruido. No tiene muestreador. Es un modelo de compresión disfrazado.

Lo que realmente deseas es: (a) que el espacio de código sea una distribución limpia y suave de la que puedas muestrear — por ejemplo, una gaussiana isotrópica N(0, I), (b) que decodificar cualquier muestra produzca un dígito plausible, y (c) que el codificador y el decodificador sigan comprimiendo bien. Tres objetivos, una arquitectura, una pérdida.

El VAE de Kingma de 2013 resuelve esto entrenando al codificador para que genere una distribución q(z|x) = N(μ(x), σ(x)²), atrayendo esa distribución hacia la priori N(0, I) mediante una penalización KL, y luego muestreando z de q(z|x) antes de decodificar. En el momento de la inferencia, descarta el codificador, muestrea z ~ N(0, I) y decodifica. La penalización KL es lo que obliga a que el espacio de código esté estructurado.

En 2026, los VAE rara vez se implementan de forma independiente —han sido superados por la difusión en cuanto a calidad bruta de imagen—, pero son el codificador preferido para cada modelo de difusión latente (SD 1/2/XL/3, Flux, AudioCraft). Aprende sobre el VAE y aprenderás la primera capa invisible de cada pipeline de imagen que utilices.

El Concepto

Autoencoder. z = encoder(x), x̂ = decoder(z), loss = ||x - x̂||². Espacio de código no estructurado.

Codificador VAE. Genera dos vectores: μ(x) y log σ²(x). Estos definen q(z|x) = N(μ, diag(σ²)).

Truco de reparametrización. Muestrear de q(z|x) no es diferenciable. Reescribe la muestra como z = μ + σ·ε donde ε ~ N(0, I). Ahora z es una función determinista de (μ, σ) más un ruido no parametrizado — los gradientes fluyen a través de μ y σ.

Pérdida. Evidence Lower BOund (ELBO), dos términos:

loss = reconstruction + β · KL[q(z|x) || N(0, I)]
     = ||x - x̂||²  + β · Σ_i ( σ_i² + μ_i² - log σ_i² - 1 ) / 2

La reconstrucción empuja x̂ hacia x. El KL empuja q(z|x) hacia la priori. Hay un compromiso mutuo (trade-off). Un β pequeño (<1) = muestras más nítidas, espacio de código menos gaussiano. Un β grande (>1) = espacio de código más limpio, muestras más borrosas. El β-VAE (Higgins 2017) hizo famoso este ajuste e inició la investigación de desentrelazamiento (disentanglement).

Muestreo. En la inferencia: extrae z ~ N(0, I), pasa hacia adelante por el decodificador. Un solo paso hacia adelante (forward pass) — sin muestreo iterativo como en la difusión.

Constrúyelo

code/main.py implementa un VAE diminuto sin numpy ni torch. La entrada son datos sintéticos de 8 dimensiones extraídos de una mezcla gaussiana de 2 componentes en 8-D. El codificador y el decodificador son MLP de una sola capa oculta. Implementamos la activación tanh, el paso hacia adelante (forward pass), la pérdida y un paso de retropropagación (backward pass) escrito a mano. No es para producción — es pedagógico.

Paso 1: paso hacia adelante del codificador

def encode(x, enc):
    h = tanh(add(matmul(enc["W1"], x), enc["b1"]))
    mu = add(matmul(enc["W_mu"], h), enc["b_mu"])
    log_sigma2 = add(matmul(enc["W_sig"], h), enc["b_sig"])
    return mu, log_sigma2

log σ² en lugar de σ para que la salida de la red sea sin restricciones (el softplus de σ es una trampa — los gradientes mueren en σ ≈ 0).

Paso 2: reparametrizar y decodificar

def reparameterize(mu, log_sigma2, rng):
    eps = [rng.gauss(0, 1) for _ in mu]
    sigma = [math.exp(0.5 * lv) for lv in log_sigma2]
    return [m + s * e for m, s, e in zip(mu, sigma, eps)]

def decode(z, dec):
    h = tanh(add(matmul(dec["W1"], z), dec["b1"]))
    return add(matmul(dec["W_out"], h), dec["b_out"])

Paso 3: el ELBO

def elbo(x, x_hat, mu, log_sigma2, beta=1.0):
    recon = sum((a - b) ** 2 for a, b in zip(x, x_hat))
    kl = 0.5 * sum(math.exp(lv) + m * m - lv - 1 for m, lv in zip(mu, log_sigma2))
    return recon + beta * kl, recon, kl

KL de forma cerrada exacta porque ambas distribuciones son gaussianas. No integres numéricamente. La gente todavía distribuye código con estimaciones de KL de Monte Carlo en 2026 — es 3 veces más lento sin ninguna razón.

Paso 4: generar

def sample(dec, z_dim, rng):
    z = [rng.gauss(0, 1) for _ in range(z_dim)]
    return decode(z, dec)

Ese es el modelo generativo. Cinco líneas.

Dificultades

• Colapso del posterior (Posterior collapse). El término KL lleva q(z|x) → N(0, I) de forma tan agresiva que z no contiene información sobre x. Solución: recocido de β (beta-annealing, comenzar con β=0 y aumentar a 1), bits libres (free bits) o ignorar el KL en dimensiones inactivas.
• Muestras borrosas. La verosimilitud del decodificador gaussiano implica una reconstrucción por MSE, que es Bayes-óptima para L2 (la media) — la media de un conjunto de dígitos plausibles es un dígito difuso. Solución: decodificador discreto (VQ-VAE, NVAE), o usar el VAE solo como codificador y apilar la difusión sobre las variables latentes (esto es lo que hace Stable Diffusion).
• β demasiado grande, demasiado pronto. Ver colapso del posterior. Comienza en β≈0.01 y auméntalo.
• Dimensión latente demasiado pequeña. 16-D funciona para MNIST, 256-D para ImageNet 256², 2048-D para ImageNet 1024². El VAE de Stable Diffusion comprime 512×512×3 → 64×64×4 (factor de submuestreo de 32x en el área espacial, 32x en los canales).

Casos de Uso

La pila (stack) de VAE en 2026:

Un modelo de difusión latente es un VAE con un modelo de difusión situado entre el codificador y el decodificador. El VAE realiza la compresión gruesa, mientras que el modelo de difusión hace el trabajo pesado. Mismo patrón para video (VAE + DiT de difusión de video) y audio (Encodec + transformer MusicGen).

Envíalo

Guarda outputs/skill-vae-trainer.md.

La habilidad (skill) toma: perfil del conjunto de datos + objetivo de dimensión latente + uso posterior (reconstrucción, muestreo o entrada de difusión latente) y produce: elección de arquitectura (común/β/VQ/RVQ), cronograma de β, dimensión latente, verosimilitud del decodificador (gaussiana vs categórica) y plan de evaluación (MSE de reconstrucción, KL por dimensión, distancia de Fréchet entre q(z|x) y N(0, I)).

Ejercicios

1. Fácil. Cambia β en code/main.py a 0.01, 0.1, 1.0, 5.0. Registra el MSE de reconstrucción final y el KL. ¿Qué β es el mejor en el sentido de Pareto (Pareto-best) para tus datos sintéticos?
2. Medio. Reemplaza la verosimilitud del decodificador gaussiano con una verosimilitud de Bernoulli (pérdida de entropía cruzada). Compara la calidad de la muestra en una versión binarizada de los mismos datos sintéticos.
3. Difícil. Extiende code/main.py a un mini VQ-VAE: reemplaza el z continuo con una búsqueda del vecino más cercano en un libro de códigos (codebook) de K=32 entradas. Compara el MSE de reconstrucción e informa cuántas entradas del libro de códigos se utilizan (el colapso del libro de códigos es real).

Términos Clave

Nota de producción: el VAE es la ruta más crítica en un servidor de difusión

En un pipeline de Stable Diffusion / Flux / SD3, el VAE se llama dos veces por solicitud: una para codificar (si se hace img2img / inpainting) y otra para decodificar. A 1024², el paso del decodificador suele ser el pico de memoria de activación individual más grande en todo el pipeline porque sobredimensiona (upsamples) latentes de 128×128×16 de regreso a 1024×1024×3. Dos consecuencias prácticas:

• Fracciona o divide en bloques (slice/tile) la decodificación. diffusers expone pipe.vae.enable_slicing() y pipe.vae.enable_tiling(). La división en bloques (tiling) cambia un pequeño artefacto de costura por memoria O(tile²) en lugar de O(H·W). Esencial para 1024²+ en GPU de consumo.
• Decodificador bf16, cálculos fp32 para el redimensionamiento final. El VAE de SD 1.x se lanzó en fp32 y silenciosamente produce NaNs cuando se convierte a fp16 en 1024²+. El SDXL distribuye madebyollin/sdxl-vae-fp16-fix — prefiere siempre la variante fp16-fix o usa bf16.

Lecturas Recomendadas

• Kingma & Welling (2013). Auto-Encoding Variational Bayes — el artículo de VAE.
• Higgins et al. (2017). β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework — β-VAE desentrelazado.
• van den Oord et al. (2017). Neural Discrete Representation Learning — VQ-VAE.
• Vahdat & Kautz (2021). NVAE: A Deep Hierarchical Variational Autoencoder — VAE de imagen en el estado del arte.
• Rombach et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models — Stable Diffusion; VAE como codificador.
• Défossez et al. (2022). High Fidelity Neural Audio Compression — Encodec, el estándar de VAE de audio.