An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

Planteamiento del Problema
Los modelos de difusión (DM) han demostrado una capacidad excepcional para resolver problemas inversos de imagen computacional mediante el modelado de priors complejos. Sin embargo, su aplicación práctica se ve severamente limitada por la dependencia de conjuntos de datos de gran escala, de alta calidad y limpios para el preentrenamiento. En muchos escenarios del mundo real, particularmente en la imagen científica y biomédica (por ejemplo, TC de rayos X, criomicroscopía electrónica), la adquisición de tales datos limpios es inviable o prohibitivamente costosa. Por el contrario, las observaciones corrompidas (con ruido, desenfoque o enmascaramiento) con modelos directos diferenciables son abundantes. Esto crea un dilema de "el huevo y la gallina": entrenar un modelo de difusión preciso requiere imágenes limpias, pero reconstruir imágenes limpias a partir de observaciones corrompidas requiere un modelo de difusión bien entrenado. Los métodos existentes para aprender de datos corrompidos, como AmbientDiffusion o SURE-Score, a menudo tienen dificultades con tipos de corrupción específicos (por ejemplo, inpainting a gran escala o ruido significativo) o no logran generalizarse a través de diversos problemas inversos.

Metodología: EMDiffusion
Los autores proponen EMDiffusion, un nuevo marco de trabajo de Esperanza-Maximización (EM) diseñado para resolver conjuntamente problemas inversos de imagen y entrenar un modelo de difusión directamente a partir de observaciones corrompidas. El método opera a través de un ciclo iterativo de dos pasos:

1. Inicialización: El proceso comienza con un prior de difusión "vago" entrenado con un conjunto muy pequeño de imágenes limpias (tan pocas como 10–50 imágenes). Esta inicialización es crucial para evitar que el modelo colapse en una distribución definida únicamente por las entradas corrompidas. Los autores señalan que estas imágenes limpias iniciales no necesitan pertenecer a la misma distribución que la tarea objetivo (los datos fuera de la distribución pueden servir como inicializadores efectivos).

2. Paso E (Esperanza): Utilizando el modelo de difusión actual como un prior, el algoritmo realiza un Muestreo de Posterior de Difusión Adaptativo para reconstruir imágenes limpias a partir de las observaciones corrompidas.

   • El muestreo posterior estándar a menudo falla en las primeras iteraciones porque el prior es débil, lo que conduce a un "colapso de modo" donde las muestras ignoran la observación y revierten a la distribución de entrenamiento.
   • Para mitigar esto, EMDiffusion introduce un hiperparámetro $\lambda$ que reescala el término de verosimilitud en la ecuación diferencial estocástica (SDE) de tiempo inverso.
   • $\lambda$ se ajusta adaptativamente en cada paso E para minimizar la pérdida de datos ($\|y - A(\hat{x}_0)\|_2^2$), equilibrando la influencia del prior imperfecto frente a los datos observados. Esto permite la generación de reconstrucciones más limpias incluso cuando el prior es inicialmente pobre.

3. Paso M (Maximización): El algoritmo actualiza los pesos del modelo de difusión utilizando las imágenes reconstruidas (muestras posteriores) generadas en el paso E.

   • Este paso se asemeja al emparejamiento de puntuación de eliminación de ruido (denoising score matching) estándar, donde el modelo aprende a aproximar la función de puntuación de los datos "limpios" reconstruidos.
   • Para acelerar el entrenamiento y evitar el sobreajuste a las reconstrucciones ruidosas de las primeras etapas, el método emplea una estrategia de entrenamiento de dos fases: inicialmente, realiza un ajuste fino de los pesos heredados de la iteración anterior; una vez que la calidad de la reconstrucción mejora (indicado por la caída del $\lambda$ óptimo por debajo de 1), el modelo se reinicia y se vuelve a entrenar desde cero sobre el conjunto completo de muestras posteriores.

Contribuciones Clave

• Marco EM para Datos Corruptos: El artículo introduce el primer marco general para entrenar modelos de difusión basados en puntuación (score-based) a partir de observaciones arbitrariamente corrompidas utilizando un enfoque EM, rompiendo efectivamente la dependencia de grandes conjuntos de datos limpios.
• Muestreo Posterior Adaptativo: La propuesta de un factor de escala adaptativo ($\lambda$) para el término de verosimilitud en la SDE de tiempo inverso aborda el problema del colapso de modo inherente a las iteraciones de EM en etapas tempranas, permitiendo una reconstrucción estable incluso con priors débiles.
• Refinamiento Iterativo: El método demuestra que una pequeña cantidad de datos limpios puede actuar como un regularizador implícito, permitiendo que el modelo converja progresivamente hacia la verdadera distribución de datos limpios a través de pasos alternos de reconstrucción y entrenamiento.
• Generalizabilidad: A diferencia de trabajos previos adaptados a corrupciones específicas (por ejemplo, enmascaramiento), EMDiffusion proporciona un enfoque unificado aplicable al inpainting aleatorio, la eliminación de ruido (denoising) y el desenfoque (deblurring).

Resultados Experimentales
Los autores validan EMDiffusion en los conjuntos de datos CIFAR-10 y CelebA a través de tres tareas: inpainting aleatorio (enmascaramiento de píxeles al 60%), eliminación de ruido Gaussiano y desenfoque Gaussiano.

• Inpainting: EMDiffusion supera significativamente a los modelos de referencia como AmbientDiffusion y SURE-Score, logrando una calidad de reconstrucción (PSPS/LPIPS) comparable a la de Diffusion Posterior Sampling (DPS) que utiliza un prior entrenado con el conjunto completo de datos limpios.
• Eliminación de Ruido y Desenfoque: El método supera a SURE-Score y a los modelos de referencia autosupervisados (Noise2Self), recuperando detalles finos (por ejemplo, ruedas de coches, rasgos faciales) que otros métodos suavizan excesivamente.
• Calidad Generativa: Los modelos de difusión aprendidos alcanzan puntuaciones FID de vanguardia en comparación con otros métodos entrenados con datos corrompidos. Notablemente, un modelo aprendido de 50,000 imágenes corrompidas mediante EM funciona mejor que un modelo entrenado con aproximadamente 15,000 imágenes limpias.
• Eficiencia: El método logra estos resultados utilizando una arquitectura DDPM estándar (35.7M de parámetros), la cual es significativamente más pequeña que las arquitecturas utilizadas por los modelos de referencia competidores (por ejemplo, los 113M+ de parámetros de AmbientDiffusion).

Significancia y Reivindicaciones
El artículo afirma que EMDiffusion demuestra que es teóricamente posible aprender la distribución subyacente de datos limpios principalmente a partir de mediciones corrompidas. La importancia radica en:

1. Superar la Escasez de Datos: Ofrece una solución práctica para dominios donde los datos limpios no están disponibles, aprovechando la abundancia de observaciones corrompidas.
2. Inicialización Robusta: Establece que un conjunto mínimo de imágenes limpias (incluso fuera de la distribución) es suficiente para inicializar el manifold, evitando que el modelo colapse en una distribución sesgada de entradas corrompidas.
3. Rendimiento de Vanguardia: El método alcanza un nuevo estado del arte en tareas de imagen inversa sin requerir modelos preentrenados en grandes conjuntos de datos limpios, cerrando la breancia entre los priors teóricos y las restricciones prácticas de datos.

Los autores mantienen la modestia respecto al alcance, reconociendo que el método depende del supuesto de que el modelo directo es conocido y diferenciable, y que el trabajo futuro podría extenderse al aprendizaje de modelos directos desconocidos o estadísticas de ruido.