Modelos de difusión: introducción

CATEGORÍA

Rubén Rodríguez Abril

Los modelos de difusión generan imágenes mediante un proceso iterativo que transforma gradualmente una composición de ruido casi puro en una imagen plenamente estructurada. Su introducción ha supuesto una extraordinaria mejoría de los sistemas de IA generativos, que son ahora mucho más creativos y producen imágenes de mejor nitidez y calidad. Son la base de célebres algoritmos como Stable Diffusion, Midjourney y las versiones más avanzadas de DALL-E.

Evolución de los modelos de IA generativa

Durante la década de los 2010s, el panorama de la síntesis de imágenes a través de IA fue dominado por dos grandes modelos: Las redes generativas antagónicas (GANs) y los autocodificadores variacionales (VAEs). Ambos fueron explorados en sendos artículos de la srie de Visión Artificial. Estos modelos operan mediante la acción coordinada de dos redes, que combinan capas convolucionales y de aumento/reducción de resolución.

En las GANs, la primera de las redes (red generadora) crea imágenes a partir de ruido mientras que la segunda, denominada red discriminadora, clasifica las imágenes que se le presentan, distinguiendo las sintéticas de las naturales. En un VAE, un codificador comprime una imagen en una representación latente consistente en una distribución gaussiana multidimensional. Desde esta distribución se realiza el muestreo aleatorio (sampling) de un vector que a su vez es transformado en una imagen por el descodificador.

La revolución de los modelos de difusión

Los modelos de difusión fueron introducidos por el trabajo Deep Unsupervised Learning Using Nonequilibrium thermodynamics, del año 2015, y hoy en día son la base de los sistemas generativos más avanzados, entre los que podemos citar a Stable Difussion y a DALL-E en su versiones segunda y posteriores.

Básicamente la misión de un modelo de difusión es tomar una imagen de puro ruido e ir transformándola progresivamente en una imagen nítida, siendo guiado a menudo por una descripción textual (prompt) o por cualquier otra condición auxiliar.

Antes del entrenamiento, se agrega progresivamente ruido blanco gaussiano a una imagen original (x₀) , hasta transformarla en ruido casi puro (x_T). Este procedimiento se realiza en un número de pasos T y es conocido como proceso de difusión (de ruido). Las representaciones producidas en los pasos intermedios se llaman datos latentes.

Durante el entrenamiento modelo debe aprender a revertir este proceso y a ir retirando ruido de la imagen, paso a paso, hasta alcanzar una nitidez completa. Este segundo procedimiento recibe el nombre de proceso de limpieza de ruido o proceso de reversión.

Figura 1. El proceso de difusión corre de derecha a izquierda. El de reversión, de izquierda a derecha. Las transiciones q(x_t|x_t-1) difunden ruido y destruyen estructuras en la imagen. Las transiciones p_θ(x_t-1|x_t) representan el procedimiento inverso y el modelo de aprendizaje profundo debe aprender a realizarlas. Los datos de cada uno de los pasos intermedios reciben el nombre de datos latentes. Fuente: Denoising Diffusion Probabilistic Models.

Proceso de difusión de ruido

Antes del entrenamiento, se generan imágenes latentes o intermedias que representan estados intermedios entre la imagen original (x₀) y la imagen final consistente en ruido puro (x_T). En cada paso los datos latentes se obtienen mediante la siguiente ecuación:

x_t = \sqrt{\alpha_t} \cdot x_0 + \sqrt{1 – \alpha_t} \cdot \epsilon

donde x_t son los datos latentes correspondiente al paso t (numero situado entre 1 y T), x₀ la imagen original, α_t el coeficiente de ruido (un hiperparámetro no aprendido, inicialmente cercano a 0 y que progresivamente se acerca a 1) y ε un número obtenido del muestreo de una distribución normal N(0,1).

Los datos latentes pueden ser los componentes de un vector de un espacio de características o los píxeles de una imagen. A cada componente/píxel se aplica la citada ecuación.

Figura 2. El proceso de difusión corre de derecha a izquierda. El de reversión, de izquierda a derecha. Las transiciones q(x_t|x_t-1) difunden ruido y destruyen estructuras en la imagen. Las transiciones p_θ(x_t-1|x_t) representan el procedimiento inverso y el modelo de aprendizaje profundo debe aprender a realizarlas. Los datos de cada uno de los pasos intermedios reciben el nombre de datos latentes. Fuente: Denoising Diffusion Probabilistic Models.

Proceso de limpieza de ruido

La realización de este proceso es el objetivo principal de los modelos de difusión. Se trata de la transformación progresiva de una imagen compuesta de ruido casi blanco en una represetnaicón nítida y bien estructurada. Se realiza por una arquitectura de aprendizaje profundo, siendo U-Net una de las más empleadas.

En cada paso, U-Net es alimentada con los datos latentes correspondientes al paso t y debe producir como salida los del paso anterior (t-1). Este proceso se aplica recursivamente, avanzando progresivamente desde el estado final T hasta el estado inicial 0.

En la generación de imágenes, es común que para cada pixel se genere una distribución gaussiana y de ella se extraigan aleatoriamente los datos latentes. La labor de U-Net es determinar la media y la desviación estándar de las distribuciones, sobre las cuales se realiza un proceso de muestreo. Matemáticamente, ello es equivalente a la siguiente ecuación:

x_{t-1} = \mu_\theta(x_t, t) + \sigma_t z

Donde x_t-1 son los datos latentes en el paso t-1, μ_θ(x_t,t) y σ_t son la media y la desviación estándar calculadas por la U-Net, mientras que z es un número muestreado de una distribución normal N(0,1).

Para garantizar que el modelo se comporte de una manera diferente a cada paso, junto a la imagen latente U-Net recibe como entradael número de paso t. Este valor se codifica en un vector específico de embedding, que es inyectado en diferentes partes de la arquitectura, permitiendo que la red adapte su comportamiento dinámicamente a medida que avanza el proceso.

Función de pérdida

La función de pérdida, cuyo gradiente debe retropropagarse a lo largo de todo el modelo, es la siguiente:

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

donde ε es el ruido introducido en el paso t durante el proceso de difusión y ε_θ(x_t,t) el ruido estimado por el modelo (y que se retira de la imagen durante el proceso de reversión). E es la media, calculada a lo largo de todo el lote, para un x₀, t y ε específicos.

Se trata de una función de pérdida basada en la norma L2 que penaliza la reconstrucción errónea del ruido por parte del modelo en cuestión (en el ejemplo que hemos puesto, U-Net).

ELBO

Este pequeño apartado tiene una cierta complejidad matemática. Por ello, los lectores que lo desee puede saltar a la sección siguiente.

El modelo busca maximizar una función objetivo, denominada ELBO (Evidence Lower Bound Objetive, Evidencia de Límite Inferior), que se expresa mediante la siguiente ecuación:

L := \mathbb{E}_q \left[ -\log p(x_T) – \sum_{t \geq 1} \log \frac{p_\theta(x_{t-1} \mid x_t)}{q(x_t \mid x_{t-1})} \right]

El término log p(x_T) es mínimo cuando los datos finales (x_T) representan ruido puro, es decir, una distribución normal estándar N(0,1). La expresión fraccionaria p_θ(x_t-1|x_t)/q(x_t|x_t-1) tiende a 1 (y su logaritmo a 0) cuando ambas distribuciones, la del proceso de reversión y la del proceso de difusión, son iguales. Ésta última, además, es matemáticamente muy próxima a la función de pérdida definida en la sección anterior, toda vez que el logaritmo de una función gaussiana es, precisamente, un término cuadrático.

Comparación con las GANs

El generador de una GAN, al igual que un modelo de difusión, transforma un vector de ruido en datos estructurados como imágenes. Sin embargo, en una GAN el proceso se realiza en una única pasada, mientras que un modelo de difusión realiza su tarea en varias etapas, en cada una de las cuales es alimentado con la salida anterior.

Esta característica provoca que las GAN consuman mucho menos poder computacional y que por lo tanto sean idóneas para realizar tareas rápidas como la edición de fotos. No obstante, presentan ciertos inconvenientes. Uno de los más comunes el el colapso de modo, donde el generador no explora todas las posibilidades que ofrece el espacio latente y se centra únicamente en la generación de un reducido conjunto de imágenes que puedan engañar al discriminador. Este fenómeno restringe la diversidad y por lo tanto penaliza la creatividad. Además, la dinámica antagonista de dos redes contrapuestas puede provocar problemas de oscilación, impidiendo que el modelo en su conjunto converja a un estado estable.

Figura 3. Figura producida por Stable Diffusion en respuesta a la descripción «pájaros célticos entonan salmos en la presencia de San Brandán el Navegante y sus compañeros». La imagen tiene una variada paleta de colores y en ella se combinan estructuras arbóreas, siluetas de pájaros e incluso patrones que recuerdan al arte de La Tène.

En los modelos de difusión, por el contrario, su procedimiento iterativo y el uso de muestreos aleatorios en cada etapa provoca que haya más entropía presente en las distribuciones de salida y en las imágenes generadas, y que por lo tanto el modelo en su conjunto sea más creativo y más apropiado para tareas artísticas que una GAN.

Los generadores de las VAEs (autocodificadores variacionales) incluyen, al igual que los modelos de difusión, elementos estocásticos. Sin embargo, la inyección de la entropía y de bits aleatorios no se realiza directamente sobre los datos, sino sobre los espacios latentes del VAE. Esto es lo que explica que a menudo las imágenes generadas por los VAEs sean más borrosas y de peor calidad que en el caso de los modelos de difusión.