Los modelos de difusión claramente explicados

  • Post author:
You are currently viewing Los modelos de difusión claramente explicados

Tiempo estimado de lectura: 10 minutos.

Los modelos de difusión son un tipo de modelo generativo de datos que está consiguiendo resultados asombrosos en aplicaciones de síntesis de imágenes.

Lo podemos ver en GLIDE de Open AI para generación, edición o modificación de imágenes a partir de texto.

También son la base de Palette de Google para la  reconstrucción de imágenes mediante coloreado (“colorization” en inglés), relleno de pixels (“inpainting” y “uncropping” en inglés) y en restauración  (“jpeg restoration” en inglés) como podemos ver en el vídeo de su página.

Contexto

Este tipo de modelos están inspirados en el término de difusión en termodinámica en “no-equilibrio” y su clave está en  generar la distribución de datos de la imagen objetivo eliminando el ruido de una imagen inicial.

Aunque por definición suene parecido a lo que hacen las redes generativas antagónicas (GAN), el proceso es muy diferente.

Por un lado, la generación del ruido en los modelos de difusión se hace en varios pasos.

Por el otro lado, se apoya en métodos bayesianos, con lo que son, en cierto modo, similares a los autocodificadores variacionales o VAE, otro método de generación de información.

Veamos su funcionamiento con un poco más de detalle.

Proceso de entrenamiento

Este tipo de modelos se entrena a través de dos procesos:

Entrenamiento de modelos de difusión

Proceso “hacia delante”

El primer proceso es el proceso hacia delante, que también nos podemos encontrar con la nomenclatura “q”, prior, o forward noising process (en inglés).

Su objetivo es generar los datos o muestras de entrenamiento.

Para ello, a partir de una imagen inicial (x0), se añade ruido en pequeños pasos a lo largo de una cadena de Markov, es decir, cada estado  en un momento de tiempo depende únicamente del anterior (q(xt-1|xt).

Las muestras se construyen paso a paso, tomando una serie de supuestos derivados de las propiedades de matrices y del teorema central del límite:

1) El ruido añadido es una gaussiana donde las variables son independientes. O, lo que es lo mismo, cada muestra generada es independiente del anterior: q(xt-1|xt) es independiente de q(xt-2|xt-1).

Este tipo de distribución recibe el nombre de gaussiana diagonal.

2) Si se añaden suficientes pasos (en la publicación original eran 1.000), la distribución de la muestra final (XT) es una gaussiana.

Y más concretamente, es una gaussiana isotrópica con la media en “0” y con varianza la matriz identidad (N(0 | I)).

En resumen: las muestras de los datos son independientes entre sí y al final es solo ruido.

Estos supuestos son lo que facilitan la estimación del proceso reverso de reconstrucción de la imagen como veremos a continuación.

Proceso reverso

El segundo proceso es el proceso reverso o “hacia atrás”, que también podemos encontrar como “pΘ“, posterior o reverse diffusion process (en inglés).

Su fin es generar la imagen objetivo gradualmente a partir del ruido.

Para ello, se entrena una red neuronal que aprende a generar la imagen objetivo a partir de los datos de entrenamiento generados en el proceso “hacia delante”.

Es a partir de aquí donde los supuestos con los que se ha generado la muestra de entrenamiento facilitan el aprendizaje:

  • Por un lado, se realiza paso a paso a través de otra cadena de Markov, de nuevo, mediante pasos de estado temporales independientes.
  • Por otro lado, como el ruido final generada en el proceso 1 (XT) es una normal o gaussiana, podemos generar la imagen original (x0) a partir del proceso para la distribución inversa q(xt-1|xt) (¿recuerdas la “q” del proceso “hacia delante”?).
  • Finalmente, como no conocemos toda la distribución inversa, la cadena de Markov en muchos pasos nos permite aproximar, de nuevo, a una gaussiana que llamamos “pΘ“.

El objetivo del entrenamiento consiste en que la distribución de pΘ(t-1|t) (aproximación) y q(xt-1|xt) (real), sean los más parecidas o en notación matemática, que el “variational lower bound” (distancia) entre p y q sea mínima.

Con esto, ya tenemos todos lo ingredientes para resolver nuestra ecuación que consiste en encontrar la media y la matriz de covarianza de la aproximación gaussiana “pΘ” en cada paso.

Esta media y covarianza serán los parámetros que, una vez entrenado, se utilizarán para la inferencia en el momento de la predicción.

El proceso se puede simplificar prediciendo el ruido que hay que sustraer en cada paso, en lugar de predecir la imagen que hay que generar, pero aquí, que cada científico utilice la mejor aproximación.

Modelo de difusión guiado

Recientes estudios han demostrado que los modelos de difusión pueden llegar a superar a otros modelos generativos de referencia, como las GAN.

El resultado mejora con algunos cambios en la arquitectura y guiando el proceso reverso con etiquetas (“classifier guidance” en inglés), en la fase de inferencia, a través la clase de la imagen (por ejemplo: “mujer”).

Proceso guiado por etiquetas o texto

Para guiar un modelo de difusión con etiquetas necesitaremos entrenar dos modelos, el de difusión y un clasificador que indica la etiqueta:

  • Objetivo de modelo 1 (difusión): que la imagen generada se parezca mucho a la real (p ≈ q).
  • Objetivo de modelo 2 (clasificador): que la clase de la foto generada coincida lo más posible con la del clasificador.

El proceso sería algo así:

  • A partir de la imagen con ruido, generamos la foto con el modelo de difusión (modelo 1).
  • Calculamos la clase  (modelo 2).
  • Hacemos un ajuste fino del resultado moviendo la imagen a lo largo del vector que representa la clase (proceso guiado).

Proceso guiado sin etiquetas

Pero, no es oro todo lo que reluce. El inconveniente de la guía con un clasificador es que es más lento y costoso porque hay que entrenar dos modelos (uno de difusión y uno de clasificación).

Para resolverlo, análisis alternativos proponen un proceso guiado sin necesidad de etiquetas y con ello evitando el modelo de clasificación.

Proceso de difusión guiado sin etiquetas
Proceso de difusión guiado sin etiquetas

En este caso, durante el entrenamiento del modelo de difusión, parte de las muestras se condicionan a un texto y parte no se condicionan.

La imagen final se “corrige” moviéndola a lo largo del vector diferencia [imagen condicionada – imagen sin condicionar].

En general, está técnica de guiado es un artificio algo extraño pero obtiene mejores imágenes porque aparentemente el vector de la clase o diferencia incluye más información semántica.

Comparativa con otros modelos generativos

En general, revertir el ruido de la imagen de forma progresiva tiene ventajas frente al salto único de las GAN, ya que hay menos posibilidad de desviarse con respecto a la realidad.

Además, se consigue un mejor compromiso entre diversidad de imágenes y fidelidad, al tiempo que se aumenta la estabilidad a la hora de entrenar el modelo, ambos, problemas comunes en las GAN.

Por otro lado, a pesar de que son fáciles de definir, los modelos de difusión son más lentos que las GAN, que hacen el proceso de generación de muestras en un solo paso, aunque, por supuesto, ya hay estudios tratando de optimizar los primeros.

Lo que es cierto es que cada vez se ven más y se están convirtiendo en un referente como vemos de su uso en las aplicaciones presentadas más recientemente para síntesis de imágenes.

Referencias

He intentado mantener la explicación lo más sencilla posible sin incluir las fórmulas complejas detrás de esta “magia”, pero como no quiero que te quedes con las ganas, te dejo una serie de referencias para que puedas entrar en todo el detalle que quieras.