Modelos de Difusión

¿Qué son los modelos de difusión?

Los modelos de difusión (Diffusion Probabilistic Models) son una familia de modelos generativos que aprenden a crear datos revirtiendo un proceso gradual de destrucción. La idea es conceptualmente elegante: primero destruyes los datos añadiendo ruido poco a poco hasta convertirlos en ruido puro; luego entrenas una red neuronal para que aprenda a deshacer cada paso de ruido, reconstruyendo los datos originales desde la nada.

A diferencia de las GANs, que usan un entrenamiento adversarial (generador vs discriminador), y los Autoencoders Variacionales (VAEs), que optimizan una ELBO, los modelos de difusión optimizan una suma ponderada de términos de denoising, lo que resulta en un entrenamiento más estable y una calidad de generación superior.

Historia de los modelos de difusión

Aunque los modelos de difusión conquistaron la IA generativa a partir de 2020, sus raíces se remontan a la termodinámica y la física estadística. El camino hasta la generación de imágenes fotorrealistas fue largo y multidisciplinar.

Intuición: destruir para crear

La idea central de los modelos de difusión se puede entender con una analogía visual. El proceso forward destruye la imagen original añadiendo ruido gaussiano en \(T\) pasos hasta obtener ruido puro. El proceso reverse es una red neuronal que aprende a revertir cada paso, reconstruyendo la imagen desde el ruido.

Lo elegante es que la red neuronal \(\epsilon_\theta\) se entrena con un objetivo muy simple: predecir el ruido que se añadió en cada paso. Durante la generación, empezamos con ruido puro \(\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I})\) y aplicamos el modelo \(T\) veces para obtener una imagen limpia.

Relación con otros modelos generativos

Los modelos de difusión no surgieron en el vacío. Son parte de un ecosistema más amplio de modelos generativos, y tienen conexiones profundas con otros enfoques que hemos estudiado en este curso.

Inspiración termodinámica

El nombre «difusión» no es casual. El proceso se inspira directamente en la termodinámica de no-equilibrio: así como las moléculas de un gas se difunden hasta alcanzar el equilibrio térmico (entropía máxima), el proceso forward lleva los datos hacia una distribución de máxima entropía (ruido gaussiano).

El paper original de Sohl-Dickstein et al. (2015) se titula literalmente «Deep Unsupervised Learning using Nonequilibrium Thermodynamics». La clave es que, mientras el proceso de «equilibrio» (forward) es fácil de definir matemáticamente, el proceso inverso (reverse) necesita una red neuronal para aprenderlo. Esta asimetría es lo que hace poderoso al modelo.

Taxonomía de los modelos de difusión

El campo ha evolucionado rápidamente, y hoy existen múltiples formulaciones y variantes. Estas son las familias principales:

Proceso Forward: destruir los datos

El proceso forward es la parte «fácil»: definimos una cadena de Markov que añade ruido gaussiano gradualmente a lo largo de \(T\) pasos. No hay parámetros que aprender aquí — todo está predefinido por un noise schedule \(\{\beta_1, \beta_2, \ldots, \beta_T\}\).

El truco de la reparametrización

Una propiedad crucial: no necesitamos simular los \(T\) pasos secuencialmente. Podemos saltar directamente a cualquier paso \(t\) desde \(\mathbf{x}_0\):

donde definimos:

Noise Schedules

La elección de \(\beta_t\) determina cuán rápido se destruyen los datos. Se han propuesto varios tipos de schedule:

Proceso Reverse: aprender a denoizar

El corazón del modelo de difusión es el proceso reverse: una cadena de Markov paramétrica que aprende a revertir cada paso de ruido.

La parametrización \(\epsilon\)

Ho et al. (2020) descubrieron que es más eficiente parametrizar la red para que prediga el ruido \(\boldsymbol{\epsilon}\) en lugar de la media \(\boldsymbol{\mu}\). La relación es:

Así, la red \(\boldsymbol{\epsilon}_\theta\) simplemente predice «¿qué ruido se añadió?» — un objetivo intuitivo y fácil de entrenar.

La arquitectura: U-Net para difusión

La red \(\boldsymbol{\epsilon}_\theta\) necesita una arquitectura que reciba una imagen ruidosa y devuelva un tensor del mismo tamaño (el ruido predicho). La elección natural es la U-Net, originalmente diseñada para segmentación de imágenes.

La U-Net de difusión tiene tres modificaciones clave respecto a la U-Net clásica de segmentación:

La ELBO y la función de pérdida

¿Cómo entrenamos un modelo de difusión? La conexión con los VAEs es profunda: ambos optimizan una Evidence Lower Bound (ELBO) sobre la log-verosimilitud de los datos.

La pérdida simplificada

Ho et al. (2020) demostraron que podemos simplificar toda la ELBO a una pérdida de predicción de ruido increíblemente simple:

En palabras:

Score Matching: la otra cara de la moneda

Existe una formulación equivalente de los modelos de difusión basada en estimación de scores (gradientes de la log-densidad).

La conexión elegante es que predecir ruido ≡ estimar el score:

Song et al. (2021) unificaron ambas visiones en un marco continuo basado en Stochastic Differential Equations (SDEs), donde el proceso forward es una SDE y el reverse es otra SDE con el score aprendido.

Algoritmo de entrenamiento

El entrenamiento de un modelo de difusión es sorprendentemente simple comparado con las GANs: no hay juego adversarial, no hay mode collapse, y la función de pérdida es un simple MSE.

Sampling: DDPM (estocástico)

Una vez entrenado el modelo, ¿cómo generamos imágenes nuevas? El sampling de DDPM sigue el proceso reverse estocástico:

DDIM y samplers acelerados

Song et al. (2021) descubrieron que se puede hacer el sampling determinista, saltando pasos sin perder calidad. Esto dio lugar a una familia de samplers rápidos.

Guidance: generación condicional

En la práctica, queremos controlar qué genera el modelo (e.g., «un gato en la nieve»). Existen dos enfoques principales:

Widget: simulador de denoising

Explora cómo el proceso de denoising reconstruye una señal a partir de ruido puro. Ajusta el número de pasos y el noise schedule para ver el efecto.

Implementación: DDPM simplificado

Veamos una implementación completa de un modelo de difusión simplificado. El código incluye el noise schedule, la U-Net simplificada, el loop de entrenamiento y el sampling.

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

# ─── Noise Schedule ───
def linear_beta_schedule(T, beta_start=1e-4, beta_end=0.02):
    return torch.linspace(beta_start, beta_end, T)

def cosine_beta_schedule(T, s=0.008):
    steps = torch.arange(T + 1, dtype=torch.float64)
    alphas_cumprod = torch.cos(((steps / T) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0.0001, 0.9999).float()

# ─── Sinusoidal Time Embedding ───
class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None] * emb[None, :]
        return torch.cat([emb.sin(), emb.cos()], dim=-1)

# ─── Simple U-Net Block ───
class Block(nn.Module):
    def __init__(self, in_ch, out_ch, time_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.gn1 = nn.GroupNorm(8, out_ch)
        self.gn2 = nn.GroupNorm(8, out_ch)
        self.time_mlp = nn.Linear(time_dim, out_ch)
        self.res_conv = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x, t_emb):
        h = self.gn1(F.silu(self.conv1(x)))
        # Inject time embedding
        h = h + self.time_mlp(F.silu(t_emb))[:, :, None, None]
        h = self.gn2(F.silu(self.conv2(h)))
        return h + self.res_conv(x)

# ─── Simplified U-Net ───
class SimpleUNet(nn.Module):
    def __init__(self, in_channels=1, base_channels=64, time_dim=256):
        super().__init__()
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.SiLU(),
        )
        # Encoder
        self.enc1 = Block(in_channels, base_channels, time_dim)
        self.enc2 = Block(base_channels, base_channels * 2, time_dim)
        self.pool = nn.MaxPool2d(2)

        # Bottleneck
        self.bot = Block(base_channels * 2, base_channels * 2, time_dim)

        # Decoder
        self.up = nn.Upsample(scale_factor=2)
        self.dec2 = Block(base_channels * 4, base_channels, time_dim)  # skip concat
        self.dec1 = Block(base_channels * 2, base_channels, time_dim)

        self.out = nn.Conv2d(base_channels, in_channels, 1)

    def forward(self, x, t):
        t_emb = self.time_mlp(t)
        # Encoder
        e1 = self.enc1(x, t_emb)
        e2 = self.enc2(self.pool(e1), t_emb)
        # Bottleneck
        b = self.bot(self.pool(e2), t_emb)
        # Decoder + skip connections
        d2 = self.dec2(torch.cat([self.up(b), e2], dim=1), t_emb)
        d1 = self.dec1(torch.cat([self.up(d2), e1], dim=1), t_emb)
        return self.out(d1)

# ─── Diffusion Model ───
class DDPM:
    def __init__(self, model, T=1000, device='cuda'):
        self.model = model.to(device)
        self.T = T
        self.device = device

        # Precompute schedule
        betas = cosine_beta_schedule(T).to(device)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)

        self.betas = betas
        self.sqrt_alphas = torch.sqrt(alphas)
        self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
        self.posterior_variance = betas * (1.0 - torch.cat([torch.tensor([1.0], device=device), alphas_cumprod[:-1]])) / (1.0 - alphas_cumprod)

    def q_sample(self, x_0, t, noise=None):
        """Forward process: add noise to x_0 at timestep t"""
        if noise is None:
            noise = torch.randn_like(x_0)
        return (self.sqrt_alphas_cumprod[t][:, None, None, None] * x_0 +
                self.sqrt_one_minus_alphas_cumprod[t][:, None, None, None] * noise)

    def train_loss(self, x_0):
        """Compute training loss: simple MSE on noise prediction"""
        t = torch.randint(0, self.T, (x_0.shape[0],), device=self.device)
        noise = torch.randn_like(x_0)
        x_t = self.q_sample(x_0, t, noise)
        predicted_noise = self.model(x_t, t.float())
        return F.mse_loss(predicted_noise, noise)

    @torch.no_grad()
    def sample(self, shape):
        """Generate samples via reverse process (DDPM sampling)"""
        x = torch.randn(shape, device=self.device)
        for t in reversed(range(self.T)):
            t_batch = torch.full((shape[0],), t, device=self.device, dtype=torch.float)
            predicted_noise = self.model(x, t_batch)

            # Compute x_{t-1}
            coef1 = 1.0 / self.sqrt_alphas[t]
            coef2 = (1 - self.betas[t] / self.sqrt_one_minus_alphas_cumprod[t])
            mean = coef1 * (x - coef2 * predicted_noise)

            if t > 0:
                noise = torch.randn_like(x)
                x = mean + torch.sqrt(self.posterior_variance[t]) * noise
            else:
                x = mean
        return x

# ─── Training Loop ───
# model = SimpleUNet(in_channels=1)
# ddpm = DDPM(model, T=1000)
# optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
#
# for epoch in range(100):
#     for batch in dataloader:
#         x_0 = batch.to('cuda')  # Normalize to [-1, 1]
#         loss = ddpm.train_loss(x_0)
#         optimizer.zero_grad()
#         loss.backward()
#         optimizer.step()
#
#     # Generate samples
#     samples = ddpm.sample((16, 1, 28, 28))  # e.g., MNIST

import tensorflow as tf
import numpy as np
import math

# ─── Noise Schedule ───
def cosine_beta_schedule(T, s=0.008):
    steps = np.arange(T + 1, dtype=np.float64)
    alphas_cumprod = np.cos(((steps / T) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return tf.constant(np.clip(betas, 1e-4, 0.9999), dtype=tf.float32)

# ─── Sinusoidal Time Embedding ───
class SinusoidalPosEmb(tf.keras.layers.Layer):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def call(self, t):
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = tf.exp(tf.range(half_dim, dtype=tf.float32) * -emb)
        emb = t[:, None] * emb[None, :]
        return tf.concat([tf.sin(emb), tf.cos(emb)], axis=-1)

# ─── Simple U-Net (Keras) ───
def build_unet(image_size=28, channels=1, base_ch=64, time_dim=256):
    # Inputs
    x_input = tf.keras.Input(shape=(image_size, image_size, channels))
    t_input = tf.keras.Input(shape=())

    # Time embedding
    t_emb = SinusoidalPosEmb(time_dim)(t_input)
    t_emb = tf.keras.layers.Dense(time_dim, activation='swish')(t_emb)

    # Encoder
    h = tf.keras.layers.Conv2D(base_ch, 3, padding='same', activation='swish')(x_input)
    t_proj = tf.keras.layers.Dense(base_ch)(t_emb)
    h = h + t_proj[:, None, None, :]
    skip1 = h

    h = tf.keras.layers.MaxPooling2D(2)(h)
    h = tf.keras.layers.Conv2D(base_ch * 2, 3, padding='same', activation='swish')(h)
    t_proj2 = tf.keras.layers.Dense(base_ch * 2)(t_emb)
    h = h + t_proj2[:, None, None, :]
    skip2 = h

    # Bottleneck
    h = tf.keras.layers.MaxPooling2D(2)(h)
    h = tf.keras.layers.Conv2D(base_ch * 2, 3, padding='same', activation='swish')(h)

    # Decoder
    h = tf.keras.layers.UpSampling2D(2)(h)
    h = tf.keras.layers.Concatenate()([h, skip2])
    h = tf.keras.layers.Conv2D(base_ch, 3, padding='same', activation='swish')(h)

    h = tf.keras.layers.UpSampling2D(2)(h)
    h = tf.keras.layers.Concatenate()([h, skip1])
    h = tf.keras.layers.Conv2D(base_ch, 3, padding='same', activation='swish')(h)

    output = tf.keras.layers.Conv2D(channels, 1)(h)
    return tf.keras.Model([x_input, t_input], output)

# ─── DDPM Training ───
class DDPMTrainer:
    def __init__(self, model, T=1000):
        self.model = model
        self.T = T
        betas = cosine_beta_schedule(T)
        alphas = 1.0 - betas
        self.alphas_cumprod = tf.math.cumprod(alphas)
        self.sqrt_alphas_cumprod = tf.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = tf.sqrt(1.0 - self.alphas_cumprod)

    @tf.function
    def train_step(self, x_0, optimizer):
        batch_size = tf.shape(x_0)[0]
        t = tf.random.uniform((batch_size,), 0, self.T, dtype=tf.int32)
        noise = tf.random.normal(tf.shape(x_0))

        sqrt_alpha = tf.gather(self.sqrt_alphas_cumprod, t)
        sqrt_one_minus = tf.gather(self.sqrt_one_minus_alphas_cumprod, t)
        x_t = sqrt_alpha[:, None, None, None] * x_0 + sqrt_one_minus[:, None, None, None] * noise

        with tf.GradientTape() as tape:
            pred_noise = self.model([x_t, tf.cast(t, tf.float32)], training=True)
            loss = tf.reduce_mean(tf.square(noise - pred_noise))

        grads = tape.gradient(loss, self.model.trainable_variables)
        optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
        return loss

# model = build_unet()
# trainer = DDPMTrainer(model, T=1000)
# optimizer = tf.keras.optimizers.Adam(2e-4)
# for epoch in range(100):
#     for batch in dataset:
#         loss = trainer.train_step(batch, optimizer)

Trucos de entrenamiento

Aunque el entrenamiento de difusión es más estable que el de GANs, hay varias prácticas que mejoran significativamente los resultados:

Aplicaciones en generación de imágenes

Los modelos de difusión se han convertido en el estándar para la generación de imágenes, superando a las GANs tanto en calidad como en diversidad. Estas son las aplicaciones principales:

Text-to-Image: la pipeline completa

La generación texto-a-imagen es la aplicación más impactante de los modelos de difusión. Veamos cómo funciona la pipeline completa, tomando Stable Diffusion como ejemplo:

Más allá de las imágenes

Los modelos de difusión no se limitan a imágenes. El mismo principio (destruir con ruido + reconstruir) se ha aplicado con éxito a una variedad sorprendente de dominios:

Diffusion Transformers (DiT)

Una de las evoluciones más importantes del campo: reemplazar la U-Net por un Transformer como backbone de denoising. Esto combina la potencia de la difusión con la escalabilidad de los Transformers.

Generación de vídeo con difusión

La extensión de difusión a vídeo es uno de los avances más impresionantes del campo. El principio es el mismo: añadir ruido a secuencias de frames y aprender a denoizarlas, pero con atención temporal adicional para mantener la coherencia.

Difusión en ciencia y medicina

Los modelos de difusión han encontrado aplicaciones transformadoras en la ciencia, especialmente en biología y química computacional:

DDPM y DDIM: los fundadores

DDPM (Ho et al., 2020) demostró que los modelos de difusión pueden generar imágenes con calidad comparable a las GANs. DDIM (Song et al., 2021) resolvió el principal problema práctico: la velocidad de sampling.

Stable Diffusion: la revolución open-source

Stable Diffusion (Rombach et al., 2022) es posiblemente el modelo de difusión más influyente de la historia. Su innovación clave: operar la difusión en un espacio latente comprimido (Latent Diffusion Model, LDM), reduciendo el coste computacional en un factor de ∼50×.

DALL-E 2/3 e Imagen

Mientras Stable Diffusion dominó el espacio open-source, OpenAI y Google desarrollaron sus propios modelos text-to-image con enfoques ligeramente diferentes.

Consistency Models: difusión en un paso

El principal problema de los modelos de difusión sigue siendo la velocidad: incluso con DDIM, necesitamos ∼20-50 pasos. Los Consistency Models (Song et al., 2023) proponen una solución radical: generar en un solo paso.

Panorama completo de modelos

El ecosistema de modelos de difusión ha crecido rápidamente. Aquí tienes un resumen de los modelos más importantes:

GANs vs Modelos de difusión: la gran comparativa

La rivalidad entre GANs y modelos de difusión es uno de los debates más interesantes del deep learning moderno. En 2021, Dhariwal & Nichol demostraron que los modelos de difusión superaban a las GANs en FID en ImageNet, marcando un punto de inflexión. Pero la historia no es tan simple: cada paradigma tiene fortalezas únicas.

El problema de la velocidad

El principal inconveniente de los modelos de difusión es su lentitud. Mientras una GAN genera una imagen en un solo paso forward (~20ms), un modelo de difusión necesita múltiples pasos de denoising:

La brecha se está cerrando rápidamente. Técnicas como destilación (progressive distillation), consistency training, y flow matching están acercando la velocidad de difusión a la de las GANs, sin sacrificar demasiada calidad.

Modelos híbridos: lo mejor de ambos mundos

Una tendencia creciente es combinar las fortalezas de GANs y difusión en arquitecturas híbridas:

Papers fundamentales

Una selección curada de los artículos más importantes que han definido el campo de los modelos de difusión:

Repositorios y herramientas

El ecosistema open-source es una de las grandes fortalezas de los modelos de difusión. Estas son las herramientas imprescindibles:

El ecosistema generativo: VAE → GAN → Difusión

Los modelos de difusión no son un reemplazo de las GANs ni de los VAEs, sino una evolución que se construye sobre los hombros de todos los enfoques anteriores. El ecosistema generativo moderno es una red de conexiones:

Como muestra el diagrama, Stable Diffusion es un ejemplo perfecto de convergencia: usa un VAE para comprimir, un modelo de difusión para generar, y compite con (y se inspira de) las GANs. Las últimas generaciones (SD3, Flux) reemplazan la U-Net por un Transformer, cerrando el círculo con la otra gran revolución del deep learning.

Resumen: el legado de los modelos de difusión