Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency

TL;DR

This work proposes ReSample, an algorithm that can solve general inverse problems with pre-trained latent diffusion models and incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that is term as hard data consistency.

Abstract

Diffusion models have recently emerged as powerful generative priors for solving inverse problems. However, training diffusion models in the pixel space are both data-intensive and computationally demanding, which restricts their applicability as priors for high-dimensional real-world data such as medical images. Latent diffusion models, which operate in a much lower-dimensional space, offer a solution to these challenges. However, incorporating latent diffusion models to solve inverse problems remains a challenging problem due to the nonlinearity of the encoder and decoder. To address these issues, we propose \textit{ReSample}, an algorithm that can solve general inverse problems with pre-trained latent diffusion models. Our algorithm incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that we term as hard data consistency. Upon solving this optimization problem, we propose a novel resampling scheme to map the measurement-consistent sample back onto the noisy data manifold and theoretically demonstrate its benefits. Lastly, we apply our algorithm to solve a wide range of linear and nonlinear inverse problems in both natural and medical images, demonstrating that our approach outperforms existing state-of-the-art approaches, including those based on pixel-space diffusion models.

Figures (21)

• Figure 1: Example reconstructions of our algorithm (ReSample) on two noisy inverse problems, nonlinear deblurring and CT reconstruction, on natural and medical images, respectively.
• Figure 2: Overview of our ReSample algorithm during the reverse sampling process conditioned on the data constraints from measurement. The entire sampling process is conducted in the latent space upon passing the sample through the encoder. The proposed algorithm performs hard data consistency at some time steps $t$ via a skipped-step mechanism.
• Figure 3: Qualitative results of multiple tasks on the LSUN-Bedroom and CelebA-HQ datasets. All inverse problems have Gaussian measurement noise with variance $\sigma_{\bm{y}} = 0.01$.
• Figure 4: Qualitative results of CT reconstruction on the LDCT dataset. We annotate the critical image structures in a red box, and zoom in below the image.
• Figure 5: Effectiveness of our resampling technique compared to stochastic encoding. Results were demonstrated on the LSUN-Bedroom and CelebA-HQ datasets with measurement noise of $\sigma_{\bm{y}} = 0.05$ to highlight the effectiveness of stochastic resample.

Theorems & Definitions (15)

• Proposition 1: Stochastic Encoding
• Proposition 2: Stochastic Resampling
• Lemma 1
• Theorem 1
• Theorem 2
• Proposition 1: Stochastic Encoding
• proof
• Proposition 2: Stochastic Resampling
• proof
• Theorem 1