Archivo del Autor: ADRIAN NOGALES HURTADO

Variations on autoencoder structures: Diffusion Models, modelos generativos que aprenden a crear datos mediante denoising

Adrián Nogales, Daniel Naval, Álvaro Jiménez y Daniel Nava

Introducción

Los modelos de difusión han ganado gran relevancia siendo la base de sistemas como Stable Diffusion, DALL·E o Imagen. Su intuición es sencilla: si aprendemos cómo se degrada una imagen añadiéndole ruido, podemos entrenar una red para hacer el proceso inverso [2].

Para explorarlo en la práctica, se ha desarrollado una implementación en PyTorch sobre MNIST con el objetivo de entender el proceso, no de maximizar la calidad generativa.

Estado del arte

El trabajo DDPM de Ho, Jain y Abbeel demostró que los modelos de difusión podían generar imágenes de alta calidad entrenando una red para predecir el ruido añadido en cada paso temporal [1][5]. Frente a las GANs, ofrecen un entrenamiento más estable basado en denoising, aunque su principal desventaja es la lentitud al requerir múltiples pasos iterativos [2][3]. Posteriormente, los Latent Diffusion Models mejoraron la eficiencia trabajando en espacios latentes comprimidos, permitiendo el desarrollo de sistemas como Stable Diffusion [3].

Fundamentos de los modelos de difusión

Los modelos de difusión se basan en dos procesos: el proceso forward y el proceso reverse.

El proceso forward toma una imagen real y le añade ruido progresivamente hasta convertirla en ruido gaussiano puro [4]. El proceso reverse hace lo contrario: parte de ese ruido e intenta eliminarlo paso a paso hasta obtener una imagen coherente. Como este proceso inverso no se conoce directamente, se entrena una red neuronal para aproximarlo [4].

Proceso forward: añadir ruido

El proceso forward es la parte en la que se añade ruido a la imagen original. Se denota con esta notación:

q(xt|xt−1)q(x_t \mid x_{t-1})

para

x0→x1→x2→...→xt→x500x₀ → x₁ → x₂ → … → xₜ → x₅₀₀ 

En el proyecto se utiliza una secuencia de 500 pasos temporales:

T=500T = 500

[4]

Esto significa que una imagen puede pasar por 500 niveles distintos de ruido. Para controlar cuánto ruido se añade en cada paso, se define una secuencia de valores beta que en nuestro caso, aumentan linealmente.

βt=βmin+(t/(T−1))(βmax−βmin)βₜ = βₘᵢₙ + (t/(T−1))(βₘₐₓ − βₘᵢₙ) 

con:

βmin=10−4βmax=0.02t∈0,1,2,...,T−1βₘᵢₙ = 10⁻⁴ βₘₐₓ = 0.02 t ∈ {0,1,2,…,T−1} 

Al principio, beta es muy pequeño, por lo que apenas se añade ruido. En los pasos finales, beta es mayor y la imagen se degrada mucho más. A partir de beta se calculan los valores alpha:

αt=1−βtαₜ = 1 − βₜ 

También se calcula el producto acumulado de los alpha, llamado ᾱₜ:

α‾t=α1·α2·α3·...·αtᾱₜ = α₁ · α₂ · α₃ · … · αₜ 

Este valor indica cuánta información de la imagen original se conserva en un determinado instante temporal.

La fórmula utilizada para obtener una imagen ruidosa en el instante t es:

xt=√α‾t·x0+√(1−α‾t)·εxₜ = √ᾱₜ · x₀ + √(1 − ᾱₜ) · ε 

con:

x₀: imagen original xₜ: imagen ruidosa en el instante t

ε ∼ N(0,I): ruido gaussiano aleatorio  (I es una matriz Identidad)

Esta fórmula permite obtener directamente una versión ruidosa de la imagen original en cualquier paso temporal, sin calcular todos los pasos anteriores.

En el código, esto se implementa en la función forward_diffusion, que recibe una imagen y un valor de t, y devuelve la imagen con ruido y el ruido añadido.

Distintos timesteps de la función forward en un elemento del dataset MNIST, timesteps [0, 50, 150, 300, 450]

Proceso reverse: eliminar ruido

El proceso reverse es la parte generativa del modelo: partiendo de ruido puro, la red elimina ruido paso a paso hasta obtener una imagen con estructura visual. En cada paso recibe la imagen actual y el instante temporal, predice el ruido presente y calcula una versión ligeramente menos ruidosa. Cada paso se representa con la siguiente notación:

pθ(xt−1|xt)p_{\theta}(x_{t-1} \mid x_t)

[4]

x500→x499→x498→...→x2→x1→x0x₅₀₀ → x₄₉₉ → x₄₉₈ → … → x₂ → x₁ → x₀ 

Este proceso se repite hasta llegar al paso 0. Si el modelo ha aprendido correctamente, la imagen final debería parecerse a una muestra real del dataset. En este caso, debería parecer un dígito manuscrito de MNIST.

La fórmula simplificada del proceso inverso es:

xt−1=(1/√αt)·(xt−((1−βt)/√(1−α‾t))·εθ(xt,t))xₜ₋₁ = (1/√αₜ) · (xₜ − ((1−βₜ)/√(1−ᾱₜ)) · εθ(xₜ,t)) 

con:

xₜ : imagen ruidosa en el instante t xₜ₋₁ : imagen tras eliminar una etapa de ruido

εθ(xₜ,t) : ruido estimado por la red neuronal

αt=1−βtα‾t=α0·α1·...·αtαₜ = 1 − βₜ ᾱₜ = α₀ · α₁ · … · αₜ 

En el código, la función sample crea una imagen aleatoria y recorre los pasos temporales de 499 a 0. En cada paso predice el ruido, actualiza la imagen y, mientras t>0, añade una pequeña cantidad de ruido adicional para mantener el carácter probabilístico del proceso.

El entrenamiento se basa en el ELBO como cota variacional, ya que maximizar directamente la probabilidad del proceso inverso es intratable. Tras simplificaciones, la red aprende εθ(xₜ, t) minimizando la diferencia entre el ruido real y el predicho:

L=||ε−εθ(xt,t)||2L = || ε − εθ(xₜ,t) ||²

De este modo, la conexión con el proceso reverse es directa: si la red sabe estimar qué parte de xₜ corresponde al ruido, entonces puede restarlo progresivamente para obtener xₜ₋₁, repitiendo este procedimiento hasta generar una imagen final x₀. [4]

Materiales y Herramientas

El dataset utilizado es MNIST, un conjunto de imágenes en escala de grises de 28×28 píxeles con dígitos manuscritos del 0 al 9. Es ampliamente usado en aprendizaje automático por su simplicidad y ligereza, lo que permite verificar rápidamente si un modelo aprende correctamente.

Para el desarrollo se han utilizado librerías como PyTorch, que permite construir la red neuronal y gestionar el entrenamiento, y Torchvision, que facilita la descarga y preprocesamiento de MNIST.

Desarrollo

Tras tener ya claro el dataset que íbamos a utilizar y cómo íbamos a abordar el cuaderno, lo primero fue cargar el dataset normalizando las imágenes al rango [−1, 1].

Proceso de difusión discreta

Lo primero que implementamos fue el proceso de difusión directa (forward process): definimos un calendario lineal de T = 500 pasos con varianzas que van de β₁ = 1×10⁻⁴ a β_T = 0.02. Con esta función podemos obtener la imagen ruidosa en cualquier paso t directamente desde x₀ con una única operación:

Arquitectura

Mini U-Net con time embedding: La red recibe como entrada la imagen ruidosa en el paso t y el propio valor de t, y devuelve una predicción del ruido añadido.

Entrenamiento 

Utilizamos el optimizador Adam con una tasa de aprendizaje de 1×10⁻⁴ y como función de pérdida el error cuadrático medio, MSELoss.

Durante cada época, se toma un lote de imágenes originales x₀ del dataloader y se selecciona aleatoriamente un instante temporal t para cada imagen del lote. A continuación, se aplica la función forward_diffusion, que genera la imagen ruidosa xₜ y devuelve también el ruido real ε que se ha añadido. Después, la red recibe como entrada xₜ y t, y produce una predicción del ruido:

εθ(xt,t)εθ(xₜ,t)

La pérdida se calcula comparando el ruido predicho por la red con el ruido real generado durante el proceso forward:

L=MSE(εθ(xt,t),ε)L = MSE(εθ(xₜ,t), ε)

De esta forma, el modelo no aprende directamente a reconstruir la imagen limpia, sino a identificar qué parte de la imagen ruidosa corresponde al ruido.

Sampling

Una vez entrenado el modelo, utilizamos la función sample para generar nuevas imágenes a partir de ruido puro. Esta función representa la aplicación práctica del proceso reverse, ya que no parte de una imagen real del dataset, sino de una imagen aleatoria generada con ruido gaussiano:

x∼N(0,I)x ∼ N(0,I)

Resultados

Realizamos distintas pruebas con distinto número de épocas, parámetros, e incluso arquitecturas de red neuronal, y los probamos generando 9 samples.

Epochs:T:Red:
50300SimpleCNN (3 capas)

Epochs:T:Red:
25300MiniUnet

Epochs:T:Red:
75300MiniUnet

Epochs:T:Red:
150500MiniUnet

Mejor versión de las pruebas

Discusión

El proyecto muestra el potencial y las limitaciones de los modelos de difusión. Su principal ventaja es un entrenamiento sencillo y estable: la red aprende a predecir el ruido dado xₜ y t. Como limitaciones destacan el alto coste computacional, la lentitud al recorrer 500 pasos inversos, la calidad visual limitada y la ausencia de condicionamiento por clase.

Por tanto, el proyecto cumple su objetivo de ilustrar cómo funciona un DDPM desde cero, aunque para resultados más realistas sería necesario ampliar la arquitectura.

Conclusiones

En este proyecto se ha implementado un DDPM en PyTorch sobre MNIST para comprender cómo una red neuronal aprende a invertir la degradación por ruido. Los resultados muestran un aprendizaje parcial del denoising, suficiente para entender por qué estos modelos se han vuelto tan relevantes en la IA generativa. La misma idea subyace en sistemas mucho más avanzados capaces de generar imágenes realistas o crear escenas a partir de texto.

Referencias

[1] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. https://arxiv.org/abs/2006.11239

[2] Weng, L. (2021). What are Diffusion Models? https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

[3] Dieleman, S. (2023). Perspectives on Diffusion Models. https://sander.ai/2023/07/20/perspectives.html

[4] LearnOpenCV. Denoising Diffusion Probabilistic Models Explained. https://learnopencv.com/denoising-diffusion-probabilistic-models/

[5] Vines. Variational Autoencoders (VAE) — base teórica. https://vinesmsuic.github.io/paper-vae/index.html[6] Comunidad Wikipedia. Diffusion Model.               https://en.wikipedia.org/wiki/Diffusion_model

About ADRIAN NOGALES HURTADO