Modelos de deep learning inspirados en la termodinámica

Los modelos de difusión inspirados en la física fueron introducidos por primera vez por Sohl-Dickstein et al en su artículo de 2015, "Deep Unsupervised Learning using Nonequilibrium Thermodynamics". Su algoritmo aplicó la dinámica de Langevin (enlace externo a ibm.com), un método para modelar el movimiento de sistemas moleculares, para respaldar la premisa básica de los modelos de difusión: convertir los datos en ruido, para que luego pueda convertir el ruido en datos.

Una nota sobre las funciones de densidad de probabilidad
Al igual que la mayoría de los modelos generativos, como los autocodificadores variacionales (VAE), el algoritmo de Sohl-Dickstein modeló la densidad de probabilidad: la probabilidad relativa de que una variable muestreada aleatoriamente, x, caiga dentro de un rango particular de valores. Básicamente, el modelado de una función de densidad de probabilidad para un conjunto de datos de entrenamiento permite a un algoritmo generar muestras que tienen muchas probabilidades de ajustarse a la distribución de los datos de entrenamiento. Al generar una nueva imagen, el modelo asume una alta probabilidad de que los valores de los píxeles se distribuyan de esa manera específica, basándose en la distribución de probabilidad que aprendió de los patrones de los datos de entrenamiento.

Lógicamente, las funciones de densidad de probabilidad exigen que la probabilidad de todas las posibilidades sume uno. Dicho de otro modo, el porcentaje de probabilidad de todas las posibilidades debe sumar exactamente el 100 %. En la práctica, esto suele requerir una constante de normalización: un valor incorporado en una función de probabilidad que reduce la probabilidad total a uno.

Calcular una constante de normalización que funcione para todos los valores variables posibles suele ser difícil de resolver: técnicamente tiene solución, pero requiere un tiempo infinito para calcularlo. En tales casos, los modelos basados en probabilidades deben restringirse a arquitecturas de modelos específicas o desarrollar soluciones inteligentes que se aproximen a la constante de normalización de una manera manejable.

Modelos generativos basados en puntuaciones
 

Independientemente del trabajo de Sohl-Dickstein, Yang Song y Stefano Ermon desarrollaron un tipo de modelo basado en la energía (enlace externo a ibm.com) llamado red de puntuación condicional de ruido en su artículo de 2019, “Generative Modeling by Estimating Gradients of the Data Distribution". Su algoritmo modeló el gradiente (∇_x) del logaritmo ($log$) de la función de densidad de probabilidad $p\left(X\right)$. El gradiente de la función de densidad de probabilidad logarítmica, escrito como ${\nabla }_X\mathrm{registro}p\left(X\right)$, se denomina puntuación de Stein o simplemente "función de puntuación".

A diferencia de las funciones de densidad de probabilidad convencionales, las funciones de puntuación no requieren una constante de normalización porque no modelan directamente la densidad de probabilidad (y por lo tanto no tienen que normalizar la probabilidad total a 1). En su lugar, se entrenan mediante la coincidencia de puntuaciones: parámetros del modelo de aprendizaje, θ, que produce un modelo p_θ(x) cuya puntuación (en otras palabras, su gradiente) coincide con la de la distribución de datos q(x) de los datos de entrenamiento.

Otra ventaja de estos modelos generativos basados en la puntuación (SGM) es que, a diferencia de los modelos basados en la verosimilitud, no imponen muchas restricciones a la arquitectura del modelo de p_θ(x).

Al explorar formas de mejorar el rendimiento de su modelo, Song y Ermon llegaron por casualidad a las mismas técnicas empleadas por Sohl-Dickstein et al. Su artículo señaló que "perturbar los datos con ruido gaussiano aleatorio hace que la distribución de datos sea más susceptible al modelado generativo basado en puntuaciones". Su modelo, construido utilizando la arquitectura U-Net desarrollada originalmente para la segmentación de imágenes, también aplicó la dinámica de Langevin para generar muestras.

Modelos probabilísticos de difusión de eliminación de ruido (DDPM)

En 2020, Ho et al propusieron utilizar el enfoque de Sohl-Dickstein para generar imágenes de alta calidad mediante el uso de la inferencia variacional en su artículo seminal, "Denoising difusion probabilistic models" (DDPM). Su artículo mostró que maximizar el límite inferior de evidencia (ELBO), una forma de reescribir los problemas de optimización basados en la probabilidad para que sean manejables, para entrenar modelos de difusión es esencialmente equivalente a la combinación de objetivos de coincidencia de puntuación utilizados para entrenar SGM.

Al implementar el enfoque de Sohl-Dickstein mediante la coincidencia de puntuaciones, Ho et al demostraron que los modelos probabilísticos de difusión pueden lograr una calidad de imagen competitiva con las GAN, que en ese momento eran lo último en tecnología. Song, Ermon, Sohl-Dickstein y otros, incluido Diederik P. Kingma, creador de la VAE- exploraron más a fondo estas conexiones en su artículo de 2021, "Score-Based Generative Modeling through Stochastic Differential Equations".

Más tarde ese mismo año, Dhariwal y Nichol, aprovechando los conocimientos del artículo mencionado anteriormente, publicaron "Diffusion Models Beat GANs on Image Synthesis", donde establecían firmemente los modelos de difusión como la nueva vanguardia.

Los DDPM, en lugar de los SGM, siguen siendo generalmente el modo dominante de los modelos de difusión, aunque con mejoras iniciadas a través de investigaciones posteriores. Por ejemplo, el influyente artículo de 2022 "High-Resolution Image Synthesis with Latent Diffusion Models".

El propósito del proceso de difusión directa es transformar los datos limpios del conjunto de datos de entrenamiento, como una imagen o una muestra de audio, en ruido puro. El método más común consiste en inyectar iterativamente ruido gaussiano hasta que toda la distribución de los datos sea gaussiana.

Matemáticamente, este proceso paso a paso se formula como una cadena de Markov: un tipo de proceso estocástico (un proceso aleatorio que sigue ciertas reglas probabilísticas) para modelar datos de series de tiempo secuenciales. En una cadena de Markov, el resultado en cada paso temporal sólo está influido por el paso temporal inmediatamente anterior. En pocas palabras: x_t, el estado de la cadena de Markov x en el paso temporal t, solo está influenciado directamente por x_t-1. La función matemática que define la transición desde cualquier x_t a x_t+1 se denomina núcleo de transición.

En cada paso temporal t, se añade una pequeña cantidad de ruido gaussiano a x_t-1 y, a continuación, se reescala la imagen para mantener un tamaño de imagen constante a pesar de la inyección continua de píxeles aleatorios. En esta formulación, x₀ es el punto de datos limpio original; x₁ es el punto de datos tras el primer paso temporal, con una pequeña cantidad de ruido gaussiano añadido; x_T es el estado final del proceso de difusión hacia delante. Si T es lo suficientemente grande, es decir, tras suficientes pasos, x_T convergerá a ruido gaussiano puro.

Definimos cada paso hacia delante como $q\left(X_t\right|X_{t-1}):$predecir el estado de la distribución de datos q (x_t), dado q(x_t-1). En un DDPM estándar, este proceso directo no implica ningún machine learning: el resultado final de la cadena de Markov siempre será una distribución gaussiana y, por lo tanto, no requiere optimización.

Añadir ruido

En un DDPM, el ruido gaussiano añadido en cada paso de la cadena de Markov no es constante ni arbitrario. En cambio, el ruido se deriva de la estructura de la imagen original y la velocidad a la que se agrega aumenta constantemente con cada paso consecutivo.

Variar la cantidad de ruido mejora la estabilidad del entrenamiento del modelo y mejora el rendimiento general al equilibrar dos prioridades contrapuestas. Como señaló Yang Song en su entrada de blog sobre modelos generativos basados en puntuaciones (enlace externo a ibm.com):

• Un ruido más grande mejora la capacidad del modelo para aprender con precisión en regiones de "baja densidad" de datos de entrenamiento (conceptos visuales y categories que tienen menos representación en los datos de entrenamiento) al rellenar esas regiones con datos ruidosos. Pero también puede alterar en exceso los datos, reduciendo la precisión global.
• Un ruido más pequeño provoca menos daños en la distribución de datos original, pero produce un rendimiento deficiente en las regiones de baja densidad.
• Por lo tanto, para lograr lo mejor de ambos mundos, los modelos de difusión utilizan múltiples escalas de ruido en el entrenamiento.

Recuerde que cualquier distribución gaussiana (normal) tiene una media, $\mu$ , y una variación, $\Sigma$ . Otro parámetro, β, sirve como factor de escala para la media y la varianza del ruido gaussiano en el núcleo de transición que define cada paso en la cadena de Markov hacia adelante. Cambiar el valor de β en un paso dado da como resultado un cambio en el ruido gaussiano agregado en ese paso. β₁ es la varianza en el paso de tiempo 1; β_t es la varianza en el paso de tiempo t, y así sucesivamente, hasta β_T.

El valor de la tasa de β en cada paso viene determinado, a su vez, por el programa de varianza.

Programación de variaciones

En cada paso t, la imagen se desplaza ligeramente de su iteración en el paso anterior (por la media) y se añade ruido a esta versión desplazada de la imagen (por la varianza). La magnitud de cada cambio y adición de ruido depende del valor de  ${\beta }_t$: como ${\beta }_t$ aumenta de acuerdo con el programa de varianza, la tasa de difusión aumenta constantemente como resultado. β siempre es un valor entre 0 y 1: por lo tanto, 0 < β₁ < β₂ < … < β_T < 1.

Elegir un calendario de varianzas específico para β es una consideración importante. Suele establecerse a mano como un hiperparámetro, ya sea fijado a un valor constante o siguiendo alguna fórmula con un valor inicial y un valor final predeterminados para β. En el documento DDPM, Ho et al utilizaron un programa lineal con 1000 pasos en el que β₁ = 10^-4 y β_T = 0,02. Investigaciones posteriores encontraron mejoras en el rendimiento y la eficacia con otros tipos de horarios, como un horario coseno,^[1 ] o haciendo que el propio horario fuera otro parámetro aprendido^[2].

El valor de β_t determina tanto la media como la varianza del ruido gaussiano añadido en el paso t.

• La media $\mu$ del ruido gaussiano añadido en el paso de tiempo t, ${\mu }_t$, se calcula como ${\mu }_t=\left(1-{\beta }_t\right)X_{t-1}$. En lenguaje sencillo, el promedio del ruido agregado en cada paso t es simplemente una versión escalada de la imagen del paso anterior, x_t-1. El tamaño de β_t determina cuánto se desvía esta media del paso anterior: cuando β_t es muy pequeño, este cambio es muy pequeño, porque $\left(1-{\beta }_t\right)\approx (1-0)\approx 1$y el ruido añadido se parecerá mucho a la imagen original. A medida que aumenta el valor de β_t , este cambio se vuelve más significativo.
  
  
• La varianza del ruido gaussiano añadido en el paso de tiempo t se calcula como ${\Sigma }_t={\beta }_{t}I$, donde $I$ es la matriz de identidad (enlace externo a ibm.com). Un β_t más grande da como resultado más ruido. Un β_t muy pequeño da como resultado un ruido insignificante.

En resumen, en cada paso t, la imagen se desplaza ligeramente de su iteración en el paso anterior (por la media) y se añade ruido a esta versión desplazada de la imagen (por la varianza). La magnitud de cada cambio y adición de ruido depende del valor de β_t: a medida que β_t aumenta de acuerdo con el programa de varianza, la tasa de difusión aumenta constantemente como resultado.

Dado que la adición de ruido gaussiano comienza gradualmente y que el propio ruido siempre se deriva de la estructura esencial de la imagen original en el paso anterior, las cualidades esenciales de la imagen original se conservan durante muchos pasos. Esto permite que el modelo aprenda de manera significativa los patrones y la estructura de la distribución de datos original durante el proceso de difusión inversa.

El truco de la reparametrización

Una deficiencia de ese proceso paso a paso es que es engorroso y costoso desde el punto de vista computacional: para cada imagen de un conjunto de datos de entrenamiento que podría contener miles o millones de imágenes, el proceso de avance requeriría docenas o cientos de pasos individuales.

En lugar de añadir ruido repetidamente, la fórmula del proceso directo puede reescribirse de forma inteligente reparametrizando la expresión de  $1–{\beta }_t$ como un nuevo parámetro,  ${\alpha }_t$. A través de una "buena propiedad" de las cadenas de Markov, este nuevo parámetro se puede ampliar aún más a un parámetro adicional, ${\overline{\alpha }}_t$, derivada de la multiplicación iterativa de ${\alpha }_t$ en cada paso progresivo de la cadena hasta ese punto. Este parámetro adicional refleja esencialmente la relación señal-ruido (SNR) de x_t: en otras palabras, cuánto de la imagen original permanece en el paso de tiempo t.

Por ejemplo, en x₁, se ha añadido una pequeña cantidad de ruido una vez. El valor de  ${\overline{\alpha }}_t$ está cerca de 1, lo que significa que la imagen aún conserva la mayor parte de su "señal" original. En un paso posterior, como x₅₀, se ha añadido ruido muchas veces. Porque ${\alpha }_t=1–{\beta }_t$, el valor de  ${\alpha }_t$ siempre es menor que 1. Desde ${\overline{\alpha }}_{50}={\alpha }_1\bullet {\alpha }_2\bullet ...\bullet {\alpha }_{49}\bullet {\alpha }_{50}$, el valor de  ${\overline{\alpha }}_t$ en el paso 50 estará mucho más cerca de 0, lo que significa que una mayor parte de la imagen original ha sido sustituida por ruido. En el paso de tiempo T, x_T es completamente ruido y el valor de  ${\overline{\alpha }}_t$ se aproxima a 0.

Aunque la derivación compleja de la ecuación está más allá del alcance de este artículo, hay dos conclusiones importantes para comprender la importancia de este truco de reparametrización:

• El estado de x en cualquier paso de tiempo t ahora se puede definir como $X_t=\sqrt{{\overline{\alpha }}_t}\bullet X_0+\sqrt{1-{\overline{\alpha }}_t}\bullet {\epsilon }_0$, donde ${\epsilon }_0$ es el ruido añadido en el primer paso. Como el valor de ${\overline{\alpha }}_t$ disminuye con cada paso, la influencia de x₀ disminuye y la influencia de ε₀ aumenta.
  
  
• Porque ${\overline{\alpha }}_t$ se deriva de $1–{\beta }_t$ y el valor de ${\beta }_t$ está determinada por el programa de varianza, reescribir la fórmula de esta manera permite el cálculo directo de x_t en cualquier paso de tiempo t sin tener que pasar por todo el proceso de avance paso a paso.

En los modelos de difusión, el proceso de difusión inversa es donde tiene lugar el machine learning real. Al aprender a realizar la inversa de los pasos de ruido del proceso directo, el modelo está esencialmente aprendiendo a eliminar el ruido gaussiano puro en una imagen limpia. Una vez que se ha entrenado la red neuronal, esta capacidad se puede utilizar para generar nuevas imágenes a partir del ruido gaussiano mediante difusión inversa paso a paso.

En teoría, la tarea del modelo puede considerarse como el simple reverso de la difusión directa. El proceso directo, que comienza con el punto de datos x₀ muestreado a partir de la distribución de datos reales q(x) del conjunto de datos de entrenamiento, se define como $q\left(X_t\right|X_{t-1})$: es decir, dado $q\left(X_{t-1}\right)$, calcular  $q\left(X_t\right)$Su opuesto, la difusión inversa, se definiría como $q\left(X_{t-1}\right|X_t)$. Pero en la práctica, la computación $q\left(X_{t-1}\right|X_t)$ es intratable.

En su lugar, la tarea de entrenamiento se formula a través de dos soluciones alternativas:

• Como se ha descrito anteriormente, q(x) se aproxima con una red neuronal p_θ(x) que constituye el propio modelo de difusión real. El objetivo del entrenamiento es aprender los parámetros del modelo θ que hacen que la salida de p_θ( $X_{t-1}|X_t$ ), haga coincidir la salida de q( $X_{t-1}|X_t$ ).
  
  
• Este modelo, p_θ(x), no predice directamente $X_{t-1}$, ni siquiera el ruido específico añadido entre x_t-1 y x_t. En cambio, predice todo el ruido presente en x_t y, a continuación, elimina una fracción de ese ruido (según el estado del programa de varianzas en ese paso) para llegar a x_t-1. El documento original del DDPM simplificaba aún más este proceso al estimar únicamente la media del ruido gaussiano, aunque los modelos posteriores suelen predecir también la varianza.

Recordemos de nuevo que la media del ruido gaussiano añadido en la difusión directa no es arbitraria: aunque es aleatoria, la estructura del ruido se deriva inicialmente de la estructura de la imagen original x₀. Por lo tanto, al aprender a predecir con precisión el ruido a través de la difusión inversa, el modelo no solo aprende a eliminar el ruido de la imagen, sino que también aprende de manera implícita la estructura de x₀.

Función de pérdida para el entrenamiento del modelo de difusión

El objetivo de entrenamiento específico utilizado para los modelos de difusión está estrechamente relacionado con el término de pérdida de reconstrucción utilizado para optimizar los autocodificadores variacionales (VAE). Al igual que los VAE, los modelos de difusión se optimizan maximizando el límite inferior variacional (VLB), también llamado límite inferior de evidencia (ELBO), de una combinación de múltiples términos de pérdida.

La maximización del VLB se utiliza en la inferencia variacional para aproximar la función de puntuación intratable ${\nabla }_X\mathrm{registro}\left(p\right(X\left)\right)$: en lugar de minimizar directamente el error, reformula la ecuación para maximizar la estimación mínima (o límite inferior) de la precisión de las predicciones del modelo.

Cada uno de los términos de pérdida utilizados refleja la divergencia Kullback-Leibler (o "divergencia KL", generalmente denotada como D_KL) entre los resultados de los pasos de difusión directa de q y los pasos inversos predichos por p_θ. La divergencia KL se utiliza para medir la diferencia entre dos distribuciones de probabilidad, por ejemplo, entre la distribución de valores de píxeles de una imagen y la distribución de valores de píxeles de otra.

En concreto, la función de pérdida para los modelos de difusión combina tres términos de pérdida: L_T, L_t y L₀.

•  _LT refleja la divergencia KL entre q$\left(X_T\right|X_0)$ y p_θ(x_T). En otras palabras, la diferencia entre el resultado final del proceso de reenvío q y el punto de partida del proceso inverso. Por lo general, este término puede ignorarse, porque x_T es gaussiano y q no tiene parámetros que se puedan aprender.
  
  
• L_t refleja la divergencia KL entre $q\left(X_{t-1}\right|X_t,X_0)$ y $p_{\theta }\left(X_{t-1}\right|X_t)$ en cada paso. En otras palabras, la precisión de cada una de las predicciones de eliminación de ruido de p_θdurante la difusión inversa en comparación con cada paso de ruido correspondiente durante el proceso de difusión directa para la imagen original, x₀.
  
  
• L₀ medidas $-\mathrm{registro}{p}_{\theta }\left(X_0\right|X_1)$. En otras palabras, L₀ refleja la probabilidad logarítmica negativa de la predicción del modelo de la imagen totalmente eliminada x₀. El gradiente de L₀ es el término de coincidencia de puntuaciones descrito anteriormente en este artículo. El término de pérdida es negativo, por lo que minimizar la función de pérdidas equivale a maximizar la probabilidad de las predicciones del modelo.
  

Aunque su compleja derivación matemática está fuera del alcance de este artículo, el VLB puede simplificarse en última instancia hasta el error cuadrático medio (MSE) entre el ruido predicho por el modelo, ${\epsilon }_{\theta }$ y el ruido real añadido en el proceso de avance, $\epsilon$, en cada paso de tiempo. Esto explica por qué el resultado del modelo es una predicción del ruido en cada paso, en lugar de la propia imagen eliminada.

Al calcular el gradiente de la función de pérdida durante la retropropagación y luego ajustar las ponderaciones del modelo para minimizar la función de pérdida a través del descenso del gradiente, las predicciones del modelo en todo el conjunto de datos de entrenamiento serán más precisas.

Una vez que el modelo de difusión ha aprendido a estimar con precisión el ruido que se restará en cada paso, se puede utilizar para generar nuevas imágenes mediante el muestreo de una imagen ruidosa aleatoria x_T a partir de la distribución de datos que ha aprendido y eliminarlo para T pasos. Al igual que los VAE, introducir un ligero elemento de aleatoriedad en el proceso de muestreo permite a los modelos de difusión producir nuevas imágenes que se parecen a los datos de entrenamiento, en lugar de reproducir directamente las imágenes de entrenamiento.

A diferencia del proceso de formación por difusión inversa, la cantidad de pasos en el proceso de generación no tiene por qué coincidir con la cantidad de pasos utilizados en el proceso de avance. Esto es posible porque el modelo está entrenado para predecir todo el ruido en cada paso de la imagen, en lugar de la cantidad específica de ruido que se eliminará en ese paso.

Un menor número de pasos conlleva una mayor velocidad y una menor demanda computacional, con una posible contrapartida en el detalle; un mayor número de pasos suele mejorar la precisión, aunque a costa de una disminución de la velocidad y un aumento de los costes computacionales.

Modelos de difusión guiada

Aunque un modelo de difusión estándar puede producir variaciones aleatorias de alta calidad de las imágenes de entrenamiento, la mayoría de los usos prácticos de un modelo de generación de imágenes requieren cierto control sobre la salida del modelo. Los modelos de difusión guiada permiten al usuario condicionar las imágenes generadas con una orientación específica.

La forma más común de modelo de difusión guiada es un modelo de difusión de texto a imagen que permite a los usuarios condicionar la salida con un mensaje de texto, como "una jirafa con sombrero de copa". Esto implica emparejar un modelo de difusión con un modelo de lenguaje de gran tamaño (LLM) separado para interpretar el mensaje de texto, introducido por primera vez por Google en el artículo "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding".

Los modelos de difusión estándar tienen como objetivo predecir la función de puntuación incondicional ${\nabla }_X\mathrm{registro}p\left(X\right)$: en otras palabras, el gradiente del logaritmo de la probabilidad de que la imagen x generada por el modelo p se ajuste a los datos de entrenamiento x. Los modelos de difusión guiada introducen una categoría visual específica, y, por ejemplo, “jirafa”, y predicen la función de puntuación condicional ${\nabla }_X\mathrm{registro}p\left(X\right|y)$: en otras palabras, la probabilidad de la imagen x, dado que debe ajustarse a la categoría y.

Los métodos de difusión guiada se pueden dividir en dos categorías:

• La difusión guiada por clasificador requiere un modelo de clasificador independiente para aprender incrustaciones vectoriales para cada categoría y para las que se entrenará el modelo de difusión para producir objetos visuales. Esta incrustación vectorial se utiliza para condicionar la salida en cada paso t. El modelo de difusión no requiere formación adicional, sino que solo puede condicionar los resultados a las específicas categories aprendidas por el clasificador.
  
  
• La orientación sin clasificador no requiere un modelo independiente, pero sí requiere un modelo de difusión de dos etapas que se entrene específicamente para la orientación condicional. Esto suele implicar un modelo de dos etapas: en la primera etapa, un algoritmo de incrustación como CLIP (enlace externo a ibm.com) genera una incrustación y para el prompt. En la segunda etapa, un modelo de difusión utiliza esa incrustación para condicionar su salida. A pesar de la sobrecarga de entrenamiento adicional, esto tiene la ventaja de permitir una guía zero-shot para categorías de imágenes no vistas.