Une perspective variationnelle sur la résolution de problèmes inverses avec des modèles de diffusion

A Variational Perspective on Solving Inverse Problems with Diffusion Models

May 7, 2023

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

cs.AI

Résumé

Les modèles de diffusion sont devenus un pilier essentiel des modèles de base dans les domaines visuels. L'une de leurs applications critiques est de résoudre universellement différentes tâches inverses en aval via un seul a priori de diffusion sans réentraînement pour chaque tâche. La plupart des tâches inverses peuvent être formulées comme l'inférence d'une distribution a posteriori sur les données (par exemple, une image complète) étant donné une mesure (par exemple, une image masquée). Cependant, cela est difficile dans les modèles de diffusion car la nature non linéaire et itérative du processus de diffusion rend l'a posteriori intraitable. Pour faire face à ce défi, nous proposons une approche variationnelle qui, par conception, cherche à approximer la véritable distribution a posteriori. Nous montrons que notre approche conduit naturellement à une régularisation par le processus de diffusion de débruitage (RED-Diff) où les débruitages à différents pas de temps imposent simultanément différentes contraintes structurelles sur l'image. Pour évaluer la contribution des débruitages à différents pas de temps, nous proposons un mécanisme de pondération basé sur le rapport signal-sur-bruit (SNR). Notre approche offre une nouvelle perspective variationnelle pour résoudre les problèmes inverses avec des modèles de diffusion, nous permettant de formuler l'échantillonnage comme une optimisation stochastique, où l'on peut simplement appliquer des solveurs prêts à l'emploi avec des itérations légères. Nos expériences pour des tâches de restauration d'image telles que l'inpainting et la super-résolution démontrent les forces de notre méthode par rapport aux modèles de diffusion basés sur l'échantillonnage de pointe.

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.