Una perspectiva variacional sobre la resolución de problemas inversos con modelos de difusión

A Variational Perspective on Solving Inverse Problems with Diffusion Models

May 7, 2023

Autores: Morteza Mardani, Jiaming Song, Jan Kautz, Arash Vahdat

cs.AI

Resumen

Los modelos de difusión han surgido como un pilar fundamental de los modelos base en dominios visuales. Una de sus aplicaciones críticas es resolver universalmente diferentes tareas inversas mediante un único prior de difusión, sin necesidad de reentrenar para cada tarea. La mayoría de las tareas inversas pueden formularse como la inferencia de una distribución posterior sobre los datos (por ejemplo, una imagen completa) dado una medición (por ejemplo, una imagen enmascarada). Sin embargo, esto es desafiante en los modelos de difusión, ya que la naturaleza no lineal e iterativa del proceso de difusión hace que la posterior sea intratable. Para abordar este desafío, proponemos un enfoque variacional que, por diseño, busca aproximar la verdadera distribución posterior. Demostramos que nuestro enfoque conduce naturalmente a una regularización mediante el proceso de difusión de eliminación de ruido (RED-Diff), donde los eliminadores de ruido en diferentes pasos temporales imponen simultáneamente diferentes restricciones estructurales sobre la imagen. Para evaluar la contribución de los eliminadores de ruido en diferentes pasos temporales, proponemos un mecanismo de ponderación basado en la relación señal-ruido (SNR). Nuestro enfoque proporciona una nueva perspectiva variacional para resolver problemas inversos con modelos de difusión, permitiéndonos formular el muestreo como una optimización estocástica, donde se pueden aplicar soluciones estándar con iteraciones ligeras. Nuestros experimentos en tareas de restauración de imágenes, como la reconstrucción de imágenes y la superresolución, demuestran las fortalezas de nuestro método en comparación con los modelos de difusión basados en muestreo más avanzados.

English

Diffusion models have emerged as a key pillar of foundation models in visual domains. One of their critical applications is to universally solve different downstream inverse tasks via a single diffusion prior without re-training for each task. Most inverse tasks can be formulated as inferring a posterior distribution over data (e.g., a full image) given a measurement (e.g., a masked image). This is however challenging in diffusion models since the nonlinear and iterative nature of the diffusion process renders the posterior intractable. To cope with this challenge, we propose a variational approach that by design seeks to approximate the true posterior distribution. We show that our approach naturally leads to regularization by denoising diffusion process (RED-Diff) where denoisers at different timesteps concurrently impose different structural constraints over the image. To gauge the contribution of denoisers from different timesteps, we propose a weighting mechanism based on signal-to-noise-ratio (SNR). Our approach provides a new variational perspective for solving inverse problems with diffusion models, allowing us to formulate sampling as stochastic optimization, where one can simply apply off-the-shelf solvers with lightweight iterates. Our experiments for image restoration tasks such as inpainting and superresolution demonstrate the strengths of our method compared with state-of-the-art sampling-based diffusion models.