Phase 04 - Lesson 10

Generación de Imágenes — Modelos de Difusión

Un modelo de difusión aprende a quitar ruido. Entrénalo para eliminar una pequeña cantidad de ruido de una imagen ruidosa, repite eso hacia atrás mil veces y tendrás un generador de imágenes.

Tipo: Build Lenguajes: Python Requisitos previos: Fase 4 Lección 07 (U-Net), Fase 1 Lección 06 (Probabilidad), Fase 3 Lección 06 (Optimizadores) Tiempo: ~75 minutos

Objetivos de Aprendizaje

• Derivar el proceso de adición de ruido hacia adelante x_0 -> x_1 -> ... -> x_T y explicar por qué la forma cerrada q(x_t | x_0) se cumple para cualquier t
• Implementar un objetivo de entrenamiento al estilo DDPM que hace regresión del ruido añadido en cada paso, y un muestreador que retrocede desde el ruido puro hasta una imagen
• Construir una U-Net condicionada al tiempo (lo bastante pequeña para entrenar en CPU) que predice el ruido para cualquier timestep
• Explicar la diferencia entre el muestreo DDPM y DDIM, y cuándo cada uno es apropiado (la Lección 23 cubre flow matching y rectified flow en profundidad)

El Problema

Las GAN generan de un solo disparo: entra ruido, sale imagen, una pasada hacia adelante. Son rápidas y difíciles de entrenar. Los modelos de difusión generan de forma iterativa: parten del ruido puro, quitan ruido en pequeños pasos, la imagen emerge. Son lentos y fáciles de entrenar. En los últimos cinco años esa última propiedad ha dominado: cualquier equipo pequeño puede entrenar un modelo de difusión y obtener muestras razonables; entrenar una GAN es un oficio que se aprende a lo largo de años de ejecuciones fallidas.

Más allá de la estabilidad del entrenamiento, la estructura iterativa de la difusión es lo que desbloquea todo lo que hace la generación moderna de imágenes: condicionamiento por texto, inpainting, edición de imágenes, superresolución, estilo controlable. Cada paso del bucle de muestreo es un lugar para inyectar una nueva restricción. Ese gancho es la razón por la que Stable Diffusion, Imagen, DALL-E 3, Midjourney y todo modelo de imagen controlable que usarás están basados en difusión.

Esta lección construye el DDPM mínimo: adición de ruido hacia adelante, eliminación de ruido hacia atrás, bucle de entrenamiento. La siguiente lección (Stable Diffusion) lo conecta a un sistema de producción con un VAE, un codificador de texto y classifier-free guidance.

El Concepto

El proceso hacia adelante

Toma una imagen x_0. Añade una pequeña cantidad de ruido gaussiano para obtener x_1. Añade un poco más para obtener x_2. Continúa por T pasos hasta que x_T sea casi indistinguible del ruido gaussiano puro.

q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1},  beta_t * I)

beta_t es una pequeña planificación de varianza, típicamente lineal de 0.0001 a 0.02 a lo largo de T=1000 pasos. Cada paso encoge ligeramente la señal e inyecta ruido nuevo.

El salto en forma cerrada

Añadir ruido un paso a la vez es una cadena de Markov, pero las matemáticas se condensan: puedes muestrear x_t directamente desde x_0 en un solo paso.

Define alpha_t = 1 - beta_t
Define alpha_bar_t = prod_{s=1..t} alpha_s

Then:
  q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0,  (1 - alpha_bar_t) * I)

Equivalently:
  x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
  where epsilon ~ N(0, I)

Esta única ecuación es toda la razón por la que la difusión es práctica. Durante el entrenamiento eliges un t aleatorio, muestreas x_t directamente desde x_0 y entrenas en un solo paso — sin necesidad de simular toda la cadena de Markov.

El proceso inverso

El proceso hacia adelante es fijo. El proceso inverso p(x_{t-1} | x_t) es lo que la red neuronal aprende. Los modelos de difusión no predicen x_{t-1} directamente; predicen el ruido epsilon añadido en el paso t, y las matemáticas derivan x_{t-1} a partir de él.

flowchart LR
    X0["x_0<br/>(imagen limpia)"] --> Q1["q(x_t|x_0)<br/>añade ruido"]
    Q1 --> XT["x_t<br/>(ruidosa)"]
    XT --> MODEL["model(x_t, t)"]
    MODEL --> EPS["epsilon predicho"]
    EPS --> LOSS["MSE contra<br/>epsilon verdadero"]

    XT -.->|muestreo| STEP["p(x_{t-1}|x_t)"]
    STEP -.-> XT1["x_{t-1}"]
    XT1 -.->|repite 1000x| X0S["x_0 (muestreado)"]

    style X0 fill:#dcfce7,stroke:#16a34a
    style MODEL fill:#fef3c7,stroke:#d97706
    style LOSS fill:#fecaca,stroke:#dc2626
    style X0S fill:#dbeafe,stroke:#2563eb

La función de pérdida del entrenamiento

Para cada paso de entrenamiento:

1. Muestrea una imagen real x_0.
2. Muestrea un timestep t uniformemente de [1, T].
3. Muestrea ruido epsilon ~ N(0, I).
4. Calcula x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon.
5. Predice epsilon_theta(x_t, t) con la red.
6. Minimiza || epsilon - epsilon_theta(x_t, t) ||^2.

Eso es todo. La red neuronal aprende a predecir el ruido en cualquier timestep. La pérdida es MSE. No hay juego adversarial, ni colapso, ni oscilación.

El muestreador (DDPM)

Para generar: parte de x_T ~ N(0, I) y retrocede un paso a la vez.

for t = T, T-1, ..., 1:
    eps = model(x_t, t)
    x_{t-1} = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * eps) + sqrt(beta_t) * z
    where z ~ N(0, I) if t > 1, else 0
return x_0

La clave es que, aunque la condicional inversa no se conoce en forma cerrada en general, para este proceso hacia adelante gaussiano específico sí. Los coeficientes de aspecto feo son lo que te da la regla de Bayes.

Por qué 1000 pasos

La planificación de ruido hacia adelante se elige de modo que cada paso añada solo el ruido suficiente para que el paso inverso sea casi gaussiano. Muy pocos pasos y el paso inverso queda lejos de ser gaussiano, la red no puede modelarlo bien. Demasiados pasos y el muestreo se vuelve caro con ganancia decreciente. T=1000 con una planificación lineal es el valor por defecto del DDPM.

DDIM: muestreo 20x más rápido

El entrenamiento es el mismo. El muestreo cambia. DDIM (Song et al., 2020) define un proceso inverso determinista que salta timesteps sin reentrenar. Muestrear en 50 pasos con DDIM da una calidad cercana a la del DDPM de 1000 pasos. Todo sistema de producción usa DDIM o una variante aún más rápida (DPM-Solver, Euler ancestral).

Condicionamiento al tiempo

La red epsilon_theta(x_t, t) necesita saber qué timestep está quitando de ruido. Los modelos de difusión modernos inyectan t mediante embeddings de tiempo sinusoidales (la misma idea que la codificación posicional en los transformers) que se añaden a los mapas de características en cada nivel de la U-Net.

t_embedding = sinusoidal(t)
feature_map += MLP(t_embedding)

Sin condicionamiento al tiempo la red tiene que adivinar el nivel de ruido a partir de la propia imagen, lo cual funciona pero es mucho menos eficiente en muestras.

Constrúyelo

Paso 1: Planificación de ruido

import torch

def linear_beta_schedule(T=1000, beta_start=1e-4, beta_end=2e-2):
    return torch.linspace(beta_start, beta_end, T)


def precompute_schedule(betas):
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    return {
        "betas": betas,
        "alphas": alphas,
        "alphas_cumprod": alphas_cumprod,
        "sqrt_alphas_cumprod": torch.sqrt(alphas_cumprod),
        "sqrt_one_minus_alphas_cumprod": torch.sqrt(1.0 - alphas_cumprod),
        "sqrt_recip_alphas": torch.sqrt(1.0 / alphas),
    }

schedule = precompute_schedule(linear_beta_schedule(T=1000))

Precalcula una vez, recupera por índice durante el entrenamiento y el muestreo.

Paso 2: Difusión hacia adelante (q_sample)

def q_sample(x0, t, noise, schedule):
    sqrt_a = schedule["sqrt_alphas_cumprod"][t].view(-1, 1, 1, 1)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"][t].view(-1, 1, 1, 1)
    return sqrt_a * x0 + sqrt_one_minus_a * noise

Forma cerrada de una línea. t es un lote de timesteps, uno por imagen en el lote.

Paso 3: Una pequeña U-Net condicionada al tiempo

import torch.nn as nn
import torch.nn.functional as F
import math

def timestep_embedding(t, dim=64):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    emb = torch.cat([args.sin(), args.cos()], dim=-1)
    return emb


class TinyUNet(nn.Module):
    def __init__(self, img_channels=3, base=32, t_dim=64):
        super().__init__()
        self.t_mlp = nn.Sequential(
            nn.Linear(t_dim, base * 4),
            nn.SiLU(),
            nn.Linear(base * 4, base * 4),
        )
        self.t_dim = t_dim
        self.enc1 = nn.Conv2d(img_channels, base, 3, padding=1)
        self.enc2 = nn.Conv2d(base, base * 2, 4, stride=2, padding=1)
        self.mid = nn.Conv2d(base * 2, base * 2, 3, padding=1)
        self.dec1 = nn.ConvTranspose2d(base * 2, base, 4, stride=2, padding=1)
        self.dec2 = nn.Conv2d(base * 2, img_channels, 3, padding=1)
        self.time_proj = nn.Linear(base * 4, base * 2)

    def forward(self, x, t):
        t_emb = timestep_embedding(t, self.t_dim)
        t_emb = self.t_mlp(t_emb)
        t_proj = self.time_proj(t_emb)[:, :, None, None]

        h1 = F.silu(self.enc1(x))
        h2 = F.silu(self.enc2(h1)) + t_proj
        h3 = F.silu(self.mid(h2))
        d1 = F.silu(self.dec1(h3))
        d2 = torch.cat([d1, h1], dim=1)
        return self.dec2(d2)

U-Net de dos niveles con condicionamiento al tiempo inyectado en el cuello de botella. Aumenta la profundidad y el ancho para imágenes reales.

Paso 4: Bucle de entrenamiento

def train_step(model, x0, schedule, optimizer, device, T=1000):
    model.train()
    x0 = x0.to(device)
    bs = x0.size(0)
    t = torch.randint(0, T, (bs,), device=device)
    noise = torch.randn_like(x0)
    x_t = q_sample(x0, t, noise, schedule)
    pred = model(x_t, t)
    loss = F.mse_loss(pred, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

Ese es el bucle de entrenamiento completo. Sin juego de GAN, sin pérdida especializada, una sola llamada de MSE.

Paso 5: Muestreador (DDPM)

@torch.no_grad()
def sample(model, schedule, shape, T=1000, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    betas = schedule["betas"].to(device)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"].to(device)
    sqrt_recip_alphas = schedule["sqrt_recip_alphas"].to(device)

    for t in reversed(range(T)):
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        coef = betas[t] / sqrt_one_minus_a[t]
        mean = sqrt_recip_alphas[t] * (x - coef * eps)
        if t > 0:
            x = mean + torch.sqrt(betas[t]) * torch.randn_like(x)
        else:
            x = mean
    return x

1000 pasadas hacia adelante para producir un lote de muestras. En código real lo cambiarías por un muestreador DDIM de 50 pasos.

Paso 6: Muestreador DDIM (determinista, ~20x más rápido)

@torch.no_grad()
def sample_ddim(model, schedule, shape, steps=50, T=1000, device="cpu", eta=0.0):
    model.eval()
    x = torch.randn(shape, device=device)
    alphas_cumprod = schedule["alphas_cumprod"].to(device)

    ts = torch.linspace(T - 1, 0, steps + 1).long()
    for i in range(steps):
        t = ts[i]
        t_prev = ts[i + 1]
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        a_t = alphas_cumprod[t]
        a_prev = alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0, device=device)
        x0_pred = (x - torch.sqrt(1 - a_t) * eps) / torch.sqrt(a_t)
        sigma = eta * torch.sqrt((1 - a_prev) / (1 - a_t) * (1 - a_t / a_prev))
        dir_xt = torch.sqrt(1 - a_prev - sigma ** 2) * eps
        noise = sigma * torch.randn_like(x) if eta > 0 else 0
        x = torch.sqrt(a_prev) * x0_pred + dir_xt + noise
    return x

eta=0 es totalmente determinista (la misma entrada de ruido siempre produce la misma salida). eta=1 recupera el DDPM.

Úsalo

Para trabajo de producción, usa diffusers:

from diffusers import DDPMScheduler, UNet2DModel

unet = UNet2DModel(sample_size=32, in_channels=3, out_channels=3, layers_per_block=2)
scheduler = DDPMScheduler(num_train_timesteps=1000)

La biblioteca trae schedulers listos para usar (DDPM, DDIM, DPM-Solver, Euler, Heun), U-Nets configurables, pipelines para texto-a-imagen e imagen-a-imagen, y utilidades de fine-tuning con LoRA.

Para investigación, k-diffusion (Katherine Crowson) tiene las implementaciones de referencia más fieles y las mejores variantes de muestreo.

Entrégalo

Esta lección produce:

• outputs/prompt-diffusion-sampler-picker.md — un prompt que elige DDPM / DDIM / DPM-Solver / Euler según la meta de calidad, el presupuesto de latencia y el tipo de condicionamiento.
• outputs/skill-noise-schedule-designer.md — una skill que produce una planificación beta lineal, coseno o sigmoide dado T y el nivel de corrupción objetivo, además de gráficos de diagnóstico de la relación señal-ruido a lo largo del tiempo.

Ejercicios

1. (Fácil) Visualiza el proceso hacia adelante: toma una imagen y grafica x_t en t in [0, 100, 250, 500, 750, 1000]. Verifica que x_1000 parezca ruido gaussiano puro.
2. (Medio) Entrena la TinyUNet en el dataset de círculos sintéticos por 20 épocas y muestrea 16 círculos. Compara el muestreo DDPM (1000 pasos) y DDIM (50 pasos) — ¿producen imágenes similares a partir de la misma semilla de ruido?
3. (Difícil) Implementa una planificación de ruido por coseno (Nichol & Dhariwal, 2021): alpha_bar_t = cos^2((t/T + s) / (1 + s) * pi / 2). Entrena el mismo modelo con planificaciones lineal y coseno y demuestra que el coseno da mejores muestras con un número bajo de pasos.

Términos Clave

Lectura Adicional

• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — el artículo que hizo práctica la difusión y superó a las GAN en FID
• Improved DDPM (Nichol & Dhariwal, 2021) — planificación por coseno y parametrización v
• DDIM (Song, Meng, Ermon, 2020) — el muestreador determinista que hizo posible la inferencia en tiempo real
• Elucidating the Design Space of Diffusion (Karras et al., 2022) — una visión unificada de cada decisión de diseño de difusión; la mejor referencia actual