Diffusion models desde cero: DDPM, DDIM y Karras

1. DDPM: noising hacia adelante, denoising hacia atrás

Denoising Diffusion Probabilistic Models (Ho et al., 2020) simplificó una idea que técnicamente existía desde 2015 (Sohl-Dickstein et al., inspirados en termodinámica de no-equilibrio) pero era difícil de entrenar. El proceso hacia adelante (forward) agrega ruido gaussiano a una imagen en pasos sucesivos, según un schedule de varianza creciente:

β = torch.linspace(β_min, β_max, n_steps)  # variance schedule, lineal: 1e-4 → 0.02 en el paper
α = 1. - β
ᾱ = torch.cumprod(α, dim=0)                # producto acumulado

# noisificar la imagen x0 hasta el timestep t, en un solo paso:
xt = ᾱ[t].sqrt()*x0 + (1-ᾱ[t]).sqrt()*ruido

La genialidad matemática del paper es que ᾱ (el producto acumulado de todos los pasos hasta t) permite saltar directo a cualquier nivel de ruido en un solo cálculo, sin simular los t pasos intermedios uno por uno — esto es lo que hace viable entrenar en la práctica: durante el entrenamiento, se elige un timestep al azar por ejemplo del batch, se noisifica directo a ese nivel, y se le pide a una red neuronal que prediga qué ruido se agregó.

Por qué funciona el salto: un solo paso del forward es x_t = √α_t · x_{t−1} + √(1−α_t) · ε_t, con cada ε_t gaussiano independiente. Si metés la expresión de x_{t−1} adentro, el término de señal queda √(α_t·α_{t−1}) · x_{t−2}, y los dos términos de ruido son gaussianas independientes de varianzas α_t(1−α_{t−1}) y (1−α_t). La suma de dos gaussianas independientes es una gaussiana cuya varianza es la suma: α_t − α_t·α_{t−1} + 1 − α_t = 1 − α_t·α_{t−1}. Repetido t veces, la señal se multiplica por √ᾱ_t (el producto de todos los √α) y el ruido acumulado es una única gaussiana de varianza 1 − ᾱ_t. De ahí que x_t = √ᾱ_t · x0 + √(1−ᾱ_t) · ε sea exacto, no una aproximación.

Fijate también que con β lineal de 1e-4 a 0.02 y 1000 pasos, ᾱ_999 queda en el orden de 4·10⁻⁵: la "imagen" al final del forward es ruido con una traza prácticamente nula de señal, que es exactamente lo que se necesita para poder arrancar el sampling desde torch.randn puro.

2. Entrenar un DDPM es sorprendentemente simple

Todo el entrenamiento se reduce a una MSELoss entre el ruido real agregado y el ruido que predice un U-Net (arquitectura que se cubre en la próxima unidad):

class DDPMCB(TrainCB):
    def before_batch(self, learn):
        x0 = learn.batch[0]
        n = len(x0)
        ε = torch.randn(x0.shape)               # ruido real: el target
        t = torch.randint(0, self.n_steps, (n,))     # un timestep al azar POR EJEMPLO del batch
        ᾱ_t = self.ᾱ[t].reshape(-1,1,1,1)          # [n] → [n,1,1,1] para que broadcastee con [n,c,h,w]
        xt = ᾱ_t.sqrt()*x0 + (1-ᾱ_t).sqrt()*ε   # noisificar en un paso
        learn.batch = ((xt, t), ε)              # input: (imagen ruidosa, t) — target: el ruido

learn = Learner(model, dls, nn.MSELoss(), cbs=[DDPMCB(n_steps=1000, ...)])
learn.fit(epochs)

Vale la pena detenerse en por qué se predice el ruido y no la imagen limpia, si son matemáticamente intercambiables (x0 = (x_t − √(1−ᾱ_t)·ε) / √ᾱ_t). Ho et al. probaron ambas y predecir ε entrena mucho mejor con una MSE simple: el ruido tiene siempre la misma escala (varianza 1) sea cual sea t, mientras que la imagen limpia es casi imposible de adivinar a t alto y trivial a t bajo — con el target en x0 la loss queda dominada por los timesteps difíciles salvo que la reponderes por ᾱ_t. Parametrizar en ε es una forma de reponderación implícita que "sale gratis".

El otro detalle es el t distinto por ejemplo: cada batch cubre así un rango amplio de niveles de ruido, y el gradiente promedia sobre todos. Si usaras un solo t para todo el batch, cada paso de optimización vería un único nivel de ruido y el entrenamiento se volvería mucho más ruidoso.

3. Samplear con DDPM: mil pasos hacia atrás

La generación (sampling) es el proceso inverso: arrancar de ruido puro (x_T = torch.randn(...)) y, para cada timestep de 999 a 0, pedirle al modelo que prediga el ruido presente, restarlo con la fórmula del paper, y repetir — mil pasos, cada uno una pasada por la red completa. Este es el Algoritmo 2 de Ho et al., tal cual:

@torch.no_grad()
def sample_ddpm(model, sz, n_steps, α, ᾱ, β):
    x_t = torch.randn(sz)                                     # x_T: ruido puro
    for t in reversed(range(n_steps)):                        # T-1, T-2, ..., 0
        t_batch = torch.full((sz[0],), t, dtype=torch.long)
        ε_hat = model((x_t, t_batch))                         # 1) predecir el ruido presente
        mean = (x_t - β[t]/(1-ᾱ[t]).sqrt() * ε_hat) / α[t].sqrt()   # 2) media de p(x_{t-1} | x_t)
        z = torch.randn_like(x_t) if t > 0 else torch.zeros_like(x_t)  # 3) ruido fresco, salvo en el último paso
        x_t = mean + β[t].sqrt() * z                          #    σ_t = √β_t en la variante simple del paper
    return x_t

Cada iteración hace tres cosas: predice ε, calcula la media de la distribución del paso anterior (que es x_t menos una fracción del ruido predicho, reescalada), y le suma ruido gaussiano nuevo de desvío σ_t. Ese ruido fresco no es un detalle: es lo que hace al proceso estocástico y evita que los errores del modelo se acumulen de manera determinista. La excepción es t=0: ahí ya estás en la imagen final y sumarle ruido solo la ensuciaría — por eso el z = 0.

El notebook 15_DDPM escribe la misma media de otra forma, pasando por la estimación de la imagen limpia (x_0_hat = (x_t − √(1−ᾱ_t)·ε_hat) / √ᾱ_t, opcionalmente recortada a [−1, 1]) y combinándola con x_t usando ᾱ_{t−1}. Es álgebra equivalente — la versión con x_0_hat es la que después se generaliza a DDIM. En el notebook, generar imágenes reconocibles de Fashion-MNIST toma exactamente esos 1000 pasos secuenciales; no hay forma de paralelizarlos porque cada paso depende del resultado del anterior.

4. DDIM: el mismo modelo, 10× menos pasos

Mil pasos secuenciales por imagen es exactamente el motivo por el que la generación con difusión temprana era lenta. Denoising Diffusion Implicit Models (Song et al., 2020) resuelve esto sin reentrenar nada: reformula el proceso de sampling para que sea determinista (dado el mismo ruido inicial, siempre da la misma imagen) y permite saltar pasos. Un paso de DDIM va de un t a un t−1 que puede estar a 10 o 20 timesteps de distancia:

def ddim_step(x_t, noise, abar_t, abar_t1, eta):
    # abar_t: ᾱ del paso actual; abar_t1: ᾱ del paso siguiente (más limpio); noise: ε_hat del modelo
    bbar_t, bbar_t1 = 1 - abar_t, 1 - abar_t1                      # "beta barra": la varianza del ruido acumulado
    sig = eta * ((bbar_t1/bbar_t).sqrt() * (1 - abar_t/abar_t1).sqrt())   # cuánto ruido fresco se reinyecta
    x_0_hat = (x_t - bbar_t.sqrt()*noise) / abar_t.sqrt()               # 1) estimar la imagen limpia
    x_t = abar_t1.sqrt()*x_0_hat + (bbar_t1 - sig**2).sqrt()*noise         # 2) re-noisificarla al nivel t-1
    if eta > 0: x_t = x_t + sig * torch.randn_like(x_t)                # 3) ruido fresco (solo si eta>0)
    return x_t

La lectura en palabras: primero se usa el ruido predicho para estimar cómo sería la imagen limpia (x_0_hat); después se la vuelve a ensuciar, pero al nivel de ruido del paso destino, reusando la misma dirección de ruido noise en vez de tirar una nueva. Como el destino puede ser cualquier ᾱ de la tabla, nada obliga a avanzar de a un timestep.

El parámetro eta interpola entre DDPM puro (eta=1, estocástico: sig coincide con el σ_t del paso de DDPM cuando los timesteps son consecutivos) y DDIM puro (eta=0: sig=0, no se agrega ruido, y la trayectoria queda completamente determinada por x_T). Ese determinismo es útil por sí mismo: permite interpolar entre dos imágenes interpolando sus x_T, y hace reproducible cualquier generación con solo fijar la semilla. El notebook mide el efecto real de reducir pasos usando FID (Fréchet Inception Distance) y KID (Kernel Inception Distance) — métricas estándar que comparan la distribución estadística de imágenes generadas contra imágenes reales, más bajo es mejor:

La conclusión práctica, medida directamente: se puede bajar de 1000 a ~100 pasos sin pérdida de calidad perceptible — una mejora de 10× en velocidad de generación sin reentrenar el modelo, solo cambiando cómo se samplea. Esto es, literalmente, la diferencia entre Stable Diffusion tardando minutos o segundos en generar una imagen.

from diffusers import DDPMScheduler, DDIMScheduler

sched = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")
# entrenamiento: equivale al before_batch de §2
t  = torch.randint(0, 1000, (n,))
xt = sched.add_noise(x0, ε, t)

# sampling con DDIM en 50 pasos: equivale al loop de ddim_step
sched = DDIMScheduler(num_train_timesteps=1000, beta_schedule="linear")
sched.set_timesteps(50)                       # elige 50 t's espaciados de 999 a 0
x = torch.randn(1, 3, 64, 64)
for t in sched.timesteps:
    ε_hat = unet(x, t).sample
    x = sched.step(ε_hat, t, x, eta=0.0).prev_sample   # DDPMScheduler.step no tiene eta: siempre estocástico

5. Karras: preconditioning en vez de un schedule fijo

El paper de Karras et al. (2022), "Elucidating the Design Space of Diffusion-Based Generative Models", replantea el problema desde otro ángulo: en vez de un schedule de ruido fijo e indexado por timestep entero, muestrea directamente el nivel de ruido (sigma) de una distribución log-normal continua, y reescala el input/output del modelo con tres coeficientes derivados analíticamente (c_skip, c_out, c_in) para que la red siempre vea activaciones de escala consistente sin importar cuánto ruido tenga la imagen:

sig_data = 0.5   # desvío estándar del dataset (Fashion-MNIST normalizado, en el notebook)

def scalings(sig):
    totvar = sig**2 + sig_data**2
    return sig_data**2/totvar, sig*sig_data/totvar.sqrt(), 1/totvar.sqrt()   # c_skip, c_out, c_in

def noisify(x0):
    sig = (torch.randn([len(x0)])*1.2-1.2).exp()   # log-normal, no lineal
    sig = sig.reshape(-1,1,1,1)
    c_skip,c_out,c_in = scalings(sig)
    noised_input = x0 + torch.randn_like(x0)*sig
    target = (x0 - c_skip*noised_input)/c_out
    return (noised_input*c_in, sig.squeeze()), target

Los tres coeficientes tienen una lectura directa. c_in normaliza la entrada: una imagen con ruido de desvío sig tiene varianza total sig² + sig_data², y dividir por su raíz la deja en varianza 1 — la misma preocupación por la escala de las activaciones de la Unidad 2, ahora aplicada al input. c_skip y c_out deciden qué tiene que predecir la red: con mucho ruido (sig grande) c_skip ≈ 0 y la red predice esencialmente la imagen limpia; con poco ruido c_skip ≈ 1 y la red predice solo la corrección chica que falta. Es una interpolación continua entre "predecir x0" y "predecir ε", eligiendo automáticamente la que está mejor condicionada en cada nivel de ruido.

La idea central que sobrevive el detalle matemático: en vez de tratar "cuánto ruido hay" como un índice discreto en una tabla fija (como hacía DDPM con sus 1000 timesteps), tratarlo como una cantidad continua que se puede escalar analíticamente da un espacio de diseño mucho más flexible — es la base conceptual de los samplers modernos (Euler, Heun, LMS) que reemplazaron por completo al sampler original de DDPM en herramientas como Stable Diffusion y ComfyUI. Cuando en esas herramientas ves "Karras" como opción de schedule, es el espaciado de sigma de este paper (en diffusers, use_karras_sigmas=True en los schedulers que lo soportan).

6. Errores comunes

Una implementación de difusión desde cero falla casi siempre por escalas o por broadcasting, no por la red.

7. Resumen

1. DDPM noisifica una imagen hacia adelante con un schedule fijo, y entrena una red para predecir el ruido agregado en cualquier paso — una MSE, nada más.
2. El truco matemático de ᾱ (producto acumulado) permite noisificar a cualquier nivel en un solo cálculo, sin simular pasos intermedios: la suma de gaussianas independientes es una gaussiana con las varianzas sumadas.
3. El sampling de DDPM va de T−1 a 0: predecir ε, calcular la media, sumar ruido fresco — salvo en el último paso.
4. DDIM reformula el sampling para ser determinista (eta=0) y saltar pasos — de 1000 a ~100 sin pérdida perceptible de calidad, medido con FID/KID. En diffusers es cambiar la clase del scheduler.
5. Karras reemplaza el schedule discreto por un nivel de ruido continuo con preconditioning analítico (c_in, c_skip, c_out) — la base de los samplers modernos.
6. Nada de esto requirió cambiar la arquitectura de la red — son todas mejoras al proceso matemático de noising/denoising alrededor de un modelo que sigue prediciendo ruido.

8. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

9. Para practicar

import torch, math, torch.nn as nn
torch.manual_seed(0)

# dataset: 4096 puntos sobre un círculo de radio 1
theta = torch.rand(4096) * 2 * math.pi
x0_all = torch.stack([theta.cos(), theta.sin()], 1)

T = 200
beta = torch.linspace(1e-4, 0.02, T); alpha = 1 - beta; abar = alpha.cumprod(0)

model = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
opt = torch.optim.Adam(model.parameters(), 2e-3)
def pred_eps(x, t): return model(torch.cat([x, (t.float()/T)[:, None]], 1))   # t entra como una feature más

for step in range(3000):
    x0 = x0_all[torch.randint(0, 4096, (256,))]
    t = torch.randint(0, T, (256,))                 # un t por ejemplo
    eps = torch.randn_like(x0)
    a = abar[t][:, None]                            # [256] → [256,1]: el reshape que en imágenes es (-1,1,1,1)
    xt = a.sqrt()*x0 + (1-a).sqrt()*eps
    loss = ((pred_eps(xt, t) - eps)**2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
print("loss final:", round(loss.item(), 4))

@torch.no_grad()
def sample_ddpm(m=512):
    x = torch.randn(m, 2)
    for t in reversed(range(T)):
        e = pred_eps(x, torch.full((m,), t))
        mean = (x - beta[t]/(1-abar[t]).sqrt()*e) / alpha[t].sqrt()
        x = mean + (beta[t].sqrt()*torch.randn_like(x) if t > 0 else 0)
    return x

xs = sample_ddpm()
r = xs.norm(dim=1)
print(f"radio medio {r.mean():.3f} ± {r.std():.3f}  (ideal: 1.000 ± 0)")

def ddim_step(x_t, noise, abar_t, abar_t1, eta):
    bbar_t, bbar_t1 = 1 - abar_t, 1 - abar_t1
    sig = eta * ((bbar_t1/bbar_t).sqrt() * (1 - abar_t/abar_t1).sqrt())
    x_0_hat = (x_t - bbar_t.sqrt()*noise) / abar_t.sqrt()
    x_t = abar_t1.sqrt()*x_0_hat + (bbar_t1 - sig**2).sqrt()*noise
    if eta > 0: x_t = x_t + sig * torch.randn_like(x_t)
    return x_t

@torch.no_grad()
def sample_ddim(m=512, steps=20, eta=0., seed=0):
    torch.manual_seed(seed)
    x = torch.randn(m, 2)
    ts = torch.linspace(T-1, 0, steps).long()          # 20 timesteps espaciados, de 199 a 0
    for i, t in enumerate(ts):
        abar_t  = abar[t]
        abar_t1 = abar[ts[i+1]] if i+1 < steps else torch.tensor(1.)   # el último salto va a "sin ruido"
        e = pred_eps(x, torch.full((m,), int(t)))
        x = ddim_step(x, e, abar_t, abar_t1, eta)
    return x

a = sample_ddim(eta=0.); b = sample_ddim(eta=0.)
print("determinista:", torch.allclose(a, b))
for eta in (0., 0.5, 1.):
    r = sample_ddim(eta=eta).norm(dim=1)
    print(f"eta={eta}: radio medio {r.mean():.3f} ± {r.std():.3f}")

10. Fuentes y para seguir

• fastai/course22p2 — notebooks 15_DDPM.ipynb, 20_DDIM.ipynb y 23_karras.ipynb; el código completo, con U-Net, del que esta unidad muestra solo el núcleo.
• Deep Unsupervised Learning using Nonequilibrium Thermodynamics (Sohl-Dickstein et al., 2015) — el origen de la idea; leelo después de DDPM para ver qué simplificó Ho.
• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — los Algoritmos 1 y 2 de la página 4 son exactamente §2 y §3; la sección 3.2 explica por qué predecir ε.
• Denoising Diffusion Implicit Models (Song et al., 2020) — la ecuación (12) es ddim_step y la (16) define sig en función de eta.
• Elucidating the Design Space of Diffusion-Based Generative Models (Karras et al., 2022) — la Tabla 1 resume todos los samplers anteriores en un mismo marco; c_skip, c_out, c_in están en la sección 5.
• Documentación de diffusers — la referencia de schedulers para mapear cada fórmula de esta unidad a la clase que la implementa.