Modelos de difusión

Introducción

Los modelos de difusión se introdujeron mediante una conexión con la termodinámica del no equilibrio^{[sohl-dickstein2015deep]}. Añaden ruido a los datos para después aprender un proceso inverso que genera muestras a partir del ruido.

En concreto, se aplica un proceso de difusión para convertir gradualmente los datos observados $x_0$ en una versión con ruido $x_T$ haciendo pasar los datos a través de $T$ pasos de un encoder $q(x_t \mid x_{t-1})$. Después de suficientes pasos, se tendrá $x_T \approx \mathcal{N}(0, I)$, o cualquier otra distribución de referencia conveniente. Luego, se aprende un proceso inverso pasando el ruido a través de $T$ pasos de un decoder $p(x_{t-1} \mid x_t)$ hasta generar $x_0$ de nuevo.

¿Recuerdas la imagen del principio del curso? Ahora debería quedar mucho más clara.

En los modelos de difusión el proceso forward no se aprende, sino que sigue un procedimiento fijo, y el estado ruidoso suele tener la misma dimensionalidad que los datos originales. Por otro lado, el denoiser comparte parámetros entre pasos. Como no existe el cuello de botella latente aprendido característico de un VAE, el posterior collapse no aparece de la misma forma.

Varios modelos generativos basados en difusión han sido propuestos, incluidos diffusion probabilistic models^{[sohl2015deep]}, noise-conditioned score network^{[song2019generative]}, denoising diffusion probabilistic models^{[ho2020denoising]} y variational diffusion models^{[kingma2021variational]}.

Encoder (forward diffusion)

El forward diffusion es la parte fácil porque no conlleva entrenar ningún modelo. Es un procedimiento fijo en el que decidimos cómo añadir ruido.

El proceso se suele definir como un modelo lineal gaussiano simple:

$$q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}\!\left( \mathbf{x}_t \,\middle|\, \sqrt{1 - \beta_t}\, \mathbf{x}_{t-1},\, \beta_t \mathbf{I} \right),$$

donde los valores de $\beta_t \in (0, 1)$ se eligen según un noise scheduler.

Intuitivamente, $\beta_t$ controla cuánto ruido añadimos en cada paso $t$:

• si $\beta_t$ es pequeño, la imagen cambia muy poco;
• si $\beta_t$ es grande, añadimos mucho ruido de golpe.

El noise scheduler decide cómo crece esa cantidad de ruido a lo largo del proceso. No es lo mismo añadir mucho ruido desde el principio que añadirlo lentamente y acelerar al final.

La distribución conjunta sobre los estados latentes es:

$$q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) = \prod_{t=1}^{T} q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$$

La distribución $q(\mathbf{x}_t \mid \mathbf{x}_0)$ se conoce como diffusion kernel. Aplicar esto a la distribución de datos de entrada y luego calcular los marginales incondicionales resultantes es equivalente a la convolución gaussiana:

$$q(\mathbf{x}_t) = \int q_0(\mathbf{x}_0)\, q(\mathbf{x}_t \mid \mathbf{x}_0)\, d\mathbf{x}_0$$

A medida que $t$ aumenta, disminuye la relación señal-ruido y los marginales se acercan a una distribución sencilla. En imágenes naturales suele parecer que los detalles finos desaparecen antes que la estructura global, pero esto es una intuición ligada al espectro de las imágenes y al schedule elegido, no una propiedad determinista del ruido gaussiano.

Decoder (reverse diffusion)

El objetivo final de un modelo de difusión es que el proceso inverso (decoder) aprenda la distribución de datos originales, $p(x_0)$.

• Proceso gradual: al añadir ruido en $T$ pasos pequeños, se crea una secuencia de distribuciones intermedias $q(x_t \mid x_0)$, donde $x_t$ es el dato original $x_0$ con una cantidad creciente de ruido.
• El proceso progresivo descompone una tarea de generación muy compleja en una secuencia de tareas de denoising más sencillas y manejables.

Aquí está la clave, generar una imagen completa desde cero es difícil, así que reformulamos el problema. Ahora lo que queremos es 'únicamente' quitar un poco de ruido de una imagen ruidosa. Al hacerlo en varios pasos, convertimos una tarea complicada en muchas tareas sencillas.

En otras palabras, en el proceso de difusión inverso queremos invertir el proceso:

• Si se conoce $x_0$, puede calcularse la posterior condicional $q(x_{t-1}\mid x_t,x_0)$; esto no convierte el paso forward aleatorio en una función determinista invertible.
• Para generar una muestra nueva no conocemos $x_0$, así que entrenamos un modelo para aproximar la distribución inversa:

$$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = \mathcal{N}\!\left( \mathbf{x}_{t-1} \,\middle|\, \mu_\theta(\mathbf{x}_t, t),\, \Sigma_\theta(\mathbf{x}_t, t) \right)$$

La distribución conjunta correspondiente sobre todas las variables generadas está dada por:

$$p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1} \mid x_t), \quad \text{donde } p(x_T) = \mathcal{N}(0, I)$$

Para generar datos a partir del modelo, se muestrea un punto $x_T \sim \mathcal{N}(0, I)$ y luego se ejecuta la cadena de Markov hacia atrás, muestreando $x_t \sim p_\theta(x_t \mid x_{t+1})$ hasta obtener una muestra en el espacio de datos original, $x_0$.

Entrenamiento

Ya hemos definido los dos procesos que componen un modelo de difusión: el proceso forward añade ruido mediante transiciones conocidas y el proceso inverso intenta recuperar progresivamente los datos. Lo que falta es aprender los parámetros $\theta$ de cada transición inversa $p_\theta(x_{t-1}\mid x_t)$.

Durante el entrenamiento partimos de una muestra limpia $x_0$, por lo que podemos generar cualquier estado ruidoso $x_t$ y calcular cómo debería ser el paso hacia $x_{t-1}$. El objetivo consiste en conseguir que la transición aprendida se aproxime a ese proceso inverso. Esta idea se formula probabilísticamente mediante el ELBO y termina dando lugar a una pérdida sencilla de predicción de ruido.

El ELBO como punto de partida

Queremos maximizar la log-verosimilitud de los datos, $\log p_\theta(x_0)$, pero calcularla exige integrar todos los estados latentes $x_{1:T}$. Como en un VAE, utilizamos el proceso forward $q(x_{1:T}\mid x_0)$ como distribución variacional y minimizamos el negativo del ELBO:

$$\mathcal{L}_{\mathrm{VLB}} = \mathbb{E}_{q} \left[ -\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)} \right]$$

Esta expresión se puede descomponer en una suma de problemas locales:

$$\mathcal{L}_{\mathrm{VLB}} = \underbrace{D_{\mathrm{KL}}\!\left(q(x_T\mid x_0)\,\|\,p(x_T)\right)}_{\mathcal{L}_T} + \sum_{t=2}^{T} \underbrace{D_{\mathrm{KL}}\!\left(q(x_{t-1}\mid x_t,x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\right)}_{\mathcal{L}_{t-1}} + \underbrace{-\log p_\theta(x_0\mid x_1)}_{\mathcal{L}_0}$$

• $\mathcal{L}_T$ compara el último estado del forward process con el prior gaussiano y no depende de $\theta$ si el noise scheduler está fijado.
• Cada $\mathcal{L}_{t-1}$ compara el paso inverso verdadero, que podemos calcular durante el entrenamiento porque conocemos $x_0$, con el paso inverso aprendido.
• $\mathcal{L}_0$ corresponde al último paso de reconstrucción hacia los datos.

La distribución $q(x_{t-1}\mid x_t,x_0)$ es gaussiana y su media puede expresarse utilizando el ruido $\epsilon$ con el que construimos $x_t$:

$$x_t=\sqrt{\bar{\alpha}_t}\,x_0+\sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon\sim\mathcal{N}(0,I)$$

Si fijamos la varianza del proceso inverso y parametrizamos su media mediante una red $\epsilon_\theta(x_t,t)$, cada término KL que depende de $\theta$ se convierte, salvo constantes, en un error cuadrático de predicción de ruido:

$$\mathcal{L}_{\mathrm{VLB}} = \mathbb{E}_{x_0,\epsilon,t} \left[ w_t \left\| \epsilon-\epsilon_\theta(x_t,t) \right\|^2 \right] +C,$$

donde $w_t$ es un peso determinado por el noise scheduler y por la varianza elegida para $p_\theta(x_{t-1}\mid x_t)$.

Pérdida simplificada

Ho et al.^{[ho2020denoising]} observaron que eliminar los pesos $w_t$ y muestrear $t$ uniformemente producía mejores muestras. Esta variante se conoce como $\mathcal{L}_{\mathrm{simple}}$:

$$\boxed{ \mathcal{L}_{\mathrm{simple}}(\theta) = \mathbb{E}_{x_0,\epsilon,t} \left[ \left\| \epsilon-\epsilon_\theta(x_t,t) \right\|^2 \right] }$$

Esta es la función de pérdida que asumiremos en el resto del capítulo. Es una versión reponderada del objetivo variacional que mantiene el mismo problema de denoising, pero ya no es exactamente el ELBO original porque asigna el mismo peso explícito a todos los timesteps.

En resumen, durante el entrenamiento:

1. Tomamos una imagen $x_0$.
2. Elegimos un paso aleatorio $t$.
3. Añadimos una cantidad conocida de ruido para obtener $x_t$.
4. Le damos $x_t$ al modelo.
5. La red predice el ruido que hemos añadido y minimizamos $\mathcal{L}_{\mathrm{simple}}$.

Aquí $\epsilon$ es el ruido real añadido y $\epsilon_\theta(x_t, t)$ es el ruido que predice la red.

Ejemplo: Stable Diffusion

La familia original de Stable Diffusion^{[rombach2022high]} popularizó la difusión latente. En lugar de ejecutar el proceso directamente sobre los píxeles, lo hace sobre un espacio latente comprimido, lo que reduce de forma importante el coste computacional.

Esto lo convierte en un latent diffusion model.

La motivación es muy práctica. Una imagen de $512 \times 512$ con 3 canales tiene una dimensionalidad muy grande, por lo que hacer difusión sobre todos esos píxeles es caro. Stable Diffusion primero comprime la imagen a una representación más pequeña, realiza el proceso de denoising en ese espacio comprimido y, al final, decodifica el resultado para volver a píxeles.

Visión general

Podemos pensar en Stable Diffusion como tres bloques principales:

1. Text encoder: convierte el prompt de texto en una representación numérica.
2. Image generator: genera una representación latente de la imagen guiándose por el texto.
3. Image decoder: reconstruye la representación latente a la dimensión objetivo de la imagen.

Text encoder: CLIP

En primer lugar el modelo necesita convertir los prompts en vectores que una red neuronal pueda utilizar.

En el Stable Diffusion original, este bloque es el encoder de texto de CLIP^{[radford2021learning]}. CLIP entrena conjuntamente un encoder de imagen y otro de texto con pares imagen-texto y proyecta ambas modalidades a un espacio compartido.

Los embeddings de texto y de imagen se comparan mediante la similitud del coseno. Durante el entrenamiento se actualizan los pesos de ambos encoders y sus proyecciones para acercar pares compatibles y separar pares incompatibles.

Repitiendo esta operación, al final del entrenamiento se consigue que los encoders produzcan embeddings en los que, por ejemplo, una imagen de un pangolín y la frase «una foto de un pangolín» queden en zonas similares del espacio latente.

Stable Diffusion utiliza el text encoder de CLIP para convertir el prompt en una secuencia de embeddings:

$$c = \text{CLIP}_{\text{text}}(\text{prompt})$$

La secuencia $c$ contiene representaciones contextuales de los tokens. El generador la utilizará para guiar la aparición de objetos, estilos y relaciones.

Ejemplo: si el prompt dice "un perro con gafas de sol en la playa", el text encoder genera una representación que ayudará al image generator a saber que deberían aparecer conceptos como perro, gafas de sol, playa, luz exterior, arena, etc.

Image generator

El nombre "image generator" puede llevar a confusión porque, en esta familia de Stable Diffusion, el componente no produce píxeles directamente, sino una matriz en el espacio latente.

El bloque central es una U-Net. Esta red recibe:

• un latente con ruido $z_t$;
• el timestep $t$;
• el condicionamiento textual $c$ producido por CLIP.

Y predice el ruido que habría que restar:

$$\epsilon_\theta(z_t, t, c)$$

Después, un scheduler usa esa predicción para calcular un latente un poco menos ruidoso:

$$z_t \rightarrow z_{t-1}$$

Este proceso se repite hasta obtener un latente limpio $z_0$. En generación condicionada suele emplearse classifier-free guidance, que combina una predicción condicionada y otra no condicionada para ajustar la fidelidad al prompt frente a la diversidad.

La U-Net usa mecanismos de cross-attention para mezclar la información visual del latente con la información textual del prompt. Gracias a esto, el modelo puede decidir qué partes de la imagen deberían prestar atención a qué partes del texto.

Cuando termina el denoising, todavía no tenemos una imagen visible, sino un latente limpio. Para convertirlo a píxeles, Stable Diffusion usa el decoder de un VAE:

$$\hat{x} = D_{\text{VAE}}(z_0)$$

Entrenamiento

Durante el entrenamiento, el modelo aprende principalmente a quitar ruido en el espacio latente.

El proceso típico es:

1. Se toma una imagen real $x$ y su texto asociado.
2. Se codifica la imagen con el encoder del VAE para obtener un latente $z_0$.
3. Se codifica el texto con CLIP para obtener el condicionamiento $c$.
4. Se elige un timestep aleatorio $t$.
5. Se añade ruido al latente $z_0$ para obtener $z_t$.
6. La U-Net recibe $(z_t, t, c)$ y predice el ruido añadido.
7. Se compara el ruido real con el ruido predicho.

La pérdida es la misma $\mathcal{L}_{\mathrm{simple}}$ que acabamos de definir, pero aplicada en el espacio latente y condicionada por el texto:

$$\mathcal{L}_{\mathrm{simple}}^{\mathrm{latent}}(\theta) = \mathbb{E}_{z_0, \epsilon, t, c} \left[ \left\| \epsilon - \epsilon_\theta(z_t, t, c) \right\|^2 \right]$$

En muchas implementaciones, el VAE y el text encoder ya están preentrenados y se mantienen congelados durante el entrenamiento principal. La parte que aprende a generar es la U-Net de denoising.

Inferencia

Aquí queremos usar el modelo ya entrenado para generar una imagen nueva.

El flujo típico es:

1. Escribimos un prompt.
2. CLIP convierte el prompt en embeddings de texto $c$.
3. Muestreamos un latente inicial aleatorio $z_T \sim \mathcal{N}(0, I)$.
4. La U-Net predice ruido en el latente actual.
5. El scheduler actualiza el latente para hacerlo un poco menos ruidoso.
6. Repetimos los pasos 4 y 5 durante varios timesteps.
7. El VAE decoder convierte el latente final $z_0$ en una imagen.

Relación con Flow Matching

Los modelos de difusión y Flow Matching están muy relacionados. Ambos parten de una distribución simple, normalmente ruido gaussiano, y aprenden a transformarla en una distribución objetivo.

En sus formulaciones pedagógicas más habituales, los objetivos se presentan con targets distintos, aunque pueden transformarse entre sí bajo un camino de ruido especificado.

En un DDPM con parametrización de ruido, el modelo aprende a predecir el ruido:

$$\epsilon_\theta(x_t, t)$$

Es decir, recibe una muestra ruidosa $x_t$ y predice qué parte de esa muestra corresponde al ruido añadido.

En Flow Matching, el modelo aprende directamente una velocidad:

$$v_\theta(x_t, t)$$

Es decir, recibe un punto intermedio $x_t$ y predice hacia dónde debería moverse para avanzar desde la distribución inicial hacia la distribución de datos.

Durante el entrenamiento

En difusión, construimos ejemplos ruidosos añadiendo ruido conocido a datos reales:

$$x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon$$

Como conocemos $\epsilon$, entrenamos a la red para predecir ese ruido.

En Flow Matching, construimos puntos intermedios entre una muestra inicial $z$ y un dato $x$:

$$x_t = (1 - t)z + tx$$

Y entrenamos a la red para predecir la velocidad que seguiría ese camino:

$$u_t = x - z$$

En ambos casos fabricamos un problema supervisado: generamos artificialmente un estado intermedio y conocemos un objetivo equivalente. En difusión ese objetivo también puede expresarse como score, muestra limpia $x_0$ o velocidad $v$; las parametrizaciones se transforman entre sí para un schedule dado.

Durante la generación

En difusión, empezamos desde ruido y vamos aplicando pasos de denoising:

$$x_T \rightarrow x_{T-1} \rightarrow \cdots \rightarrow x_0$$

En Flow Matching, empezamos desde ruido y seguimos un campo de velocidades:

$$\frac{dx_t}{dt} = v_\theta(x_t, t)$$

La generación también avanza de ruido a datos, pero cambia el punto de partida matemático. La difusión puede implementarse como una cadena estocástica o mediante su probability-flow ODE; Flow Matching define y aprende directamente un campo de velocidades sobre un camino de probabilidad. Los dos marcos están estrechamente relacionados.