Diffusion-based Models

Introducción

Los modelos de difusión se inspiran en la termodinámica del no equilibrio. Añaden ruido a los datos para después aprender a invertir el proceso para reconstruir las muestras a partir del ruido.

La intuición es muy visual: partimos de una imagen real, le añadimos un poco de ruido, luego un poco más, y así sucesivamente hasta que prácticamente solo queda ruido. Después entrenamos una red neuronal para hacer el camino contrario: mirar una imagen ruidosa y decir qué parte de esa imagen parece ruido para poder retirarlo.

En concreto, se puede aplicar 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 estocástico $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. Después, se aprende un proceso inverso para deshacer este proceso, 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 encoder se no se aprende si no que sigue un procedimiento fijo y la variable latente tiene la misma dimensionalidad que los datos originales (a diferencia de los VAEs). Por otro lado, el decoder se comparte en todos los pasos. Estas condiciones dan lugar a una función de error simple, que facilita el entrenamiento y evita el riesgo del posterior collapse.

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 se aprende: nosotros decidimos cómo añadir ruido. Es un procedimiento fijo.

El proceso de codificación hacia delante se define 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 el 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, los marginales se vuelven más simples. En imágenes, este proceso primero eliminará el contenido de alta frecuencia (es decir, detalles de bajo nivel, como la textura), y luego eliminará el contenido de baja frecuencia (información de alto nivel o "semántica", como la forma).

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.

Esta es la clave práctica: 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 pequeñas.

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

• Si se conoce la entrada $x_0$, se puede calcular la inversa matemáticamente.
• Si se quiere generar un nuevo punto de datos, aunque no se conozca $x_0$, se entrena al generador/decoder para que se aproxime a la distribución de los datos sobre el proceso inverso. Por lo tanto, sigue la forma:

$$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$.

Model fitting

Los modelos de difusión, al igual que los VAE, pueden utilizar el ELBO como función de error. En este contexto, el ELBO es una suma de términos que corresponden a la divergencia KL entre las distribuciones de ruido añadido en cada paso del proceso hacia delante y las distribuciones predichas por el proceso inverso. Al optimizar este ELBO, el modelo aprende a invertir el proceso de adición de ruido, aprendiendo efectivamente a eliminar el ruido de los datos en cada paso.

En lugar de entrenar el modelo para predecir la imagen con ruido anterior $x_{t-1}$ dado $x_t$, se entrena el modelo para predecir el ruido que hay en $x_t$.

Es decir, durante el entrenamiento hacemos algo así:

1. Tomamos una imagen real $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. El modelo intenta predecir el ruido que hemos añadido.

Como nosotros hemos añadido el ruido, sabemos cuál es la respuesta correcta. Por eso podemos entrenar el modelo con una pérdida sencilla, normalmente un error cuadrático entre el ruido real y el ruido predicho:

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

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

Los modelos de difusión están estrechamente relacionados con el score matching, una técnica en la que el modelo aprende el gradiente (score) de la densidad de datos logarítmica. El ELBO proporciona un método para estimar los parámetros del proceso inverso, que está relacionado con el score matching en el sentido de que se aprende la función de puntuación del proceso inverso.

Maximizar el ELBO equivale a minimizar la diferencia entre la distribución real de los datos y la distribución del modelo tras el proceso de difusión inverso. El ELBO suele incluir:

• Un término de reconstrucción que asegura que el modelo pueda reconstruir los datos a partir de una versión ruidosa.
• Un término de regularización (divergencia KL) que garantiza que el ruido añadido en el proceso hacia delante sigue la distribución esperada.

Ejemplo: Stable Diffusion

Stable Diffusion es uno de los modelos más icónicos que aplicaron la idea de difusión. Uno de los motivos en el salto de calidad en generación de imágenes respecto a modelos anteriores fue que la difusión no se hace directamente sobre los píxeles de la imagen si no que se hace sobre el espacio latente.

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.

Normalmente este encoder es un encoder de un transformer. Técnicamente, el encoder de texto de un modelo tipo CLIP. CLIP es una técnica para entrenar con pares de imágenes y textos para aprender representaciones donde una imagen y su descripción queden cerca en el espacio vectorial.

Los embeddings de texto y de imagen se comparan mediante la similitud del coseno. Al inicio del proceso de entrenamiento, la similitud será baja, aunque el texto describa correctamente la imagen simplemente porque los espacios vectoriales de cada encoder son diferentes. Lo que se hace entonces es actualizar los dos pesos de ambos modelos para que la próxima vez los embeddings resultantes sean similares.

Repitiendo esta operación, al final del entrenamiento se consigue que los encoders sean capaces de producir embeddings en las que una imagen de un perro y la frase «una foto de un perro» quedan 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}).$$

Ese vector $c$ es una representación semántica del texto: contiene información sobre objetos, estilos, relaciones y detalles que el modelo usará para guiar la generación.

Image generator

El nombre "image generator" puede llevar a confusión, porque en Stable Diffusion este componente la parte no produce píxeles directamente si no que genera un matriz en el espacio latente.

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

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

Y predice el ruido que hay que retirar:

$$\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 muchas veces hasta obtener un latente limpio $z_0$.

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 elege 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 se parece a la de los modelos de difusión que vimos antes, pero ahora se aplica en el espacio latente:

$$\mathcal{L}(\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

La inferencia es el proceso de 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.

Diferencias 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.

La diferencia principal está en qué aprende el modelo.

En difusión, el modelo aprende a quitar 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 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 sabemos cuál debería ser la respuesta correcta.

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 la interpretación cambia. Difusión lo ve como una cadena de eliminación de ruido. Flow Matching lo ve como una trayectoria continua guiada por velocidades.