Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast Samplers for Inverse Problems in Iterative Refinement Models

Kushagra Pandey, Ruihan Yang, Stephan Mandt

TL;DR

This work proposes Conditional Conjugate Integrators, which leverage the specific form of the inverse problem to project the respective conditional diffusion/flow dynamics into a more amenable space for sampling and evaluates the proposed method's performance on various linear image restoration tasks across multiple datasets.

Abstract

Constructing fast samplers for unconditional diffusion and flow-matching models has received much attention recently; however, existing methods for solving inverse problems, such as super-resolution, inpainting, or deblurring, still require hundreds to thousands of iterative steps to obtain high-quality results. We propose a plug-and-play framework for constructing efficient samplers for inverse problems, requiring only pre-trained diffusion or flow-matching models. We present Conditional Conjugate Integrators, which leverage the specific form of the inverse problem to project the respective conditional diffusion/flow dynamics into a more amenable space for sampling. Our method complements popular posterior approximation methods for solving inverse problems using diffusion/flow models. We evaluate the proposed method's performance on various linear image restoration tasks across multiple datasets, employing diffusion and flow-matching models. Notably, on challenging inverse problems like 4x super-resolution on the ImageNet dataset, our method can generate high-quality samples in as few as 5 conditional sampling steps and outperforms competing baselines requiring 20-1000 steps. Our code will be publicly available at https://github.com/mandt-lab/c-pigdm

Fast Samplers for Inverse Problems in Iterative Refinement Models

TL;DR

This work proposes Conditional Conjugate Integrators, which leverage the specific form of the inverse problem to project the respective conditional diffusion/flow dynamics into a more amenable space for sampling and evaluates the proposed method's performance on various linear image restoration tasks across multiple datasets.

Abstract

Constructing fast samplers for unconditional diffusion and flow-matching models has received much attention recently; however, existing methods for solving inverse problems, such as super-resolution, inpainting, or deblurring, still require hundreds to thousands of iterative steps to obtain high-quality results. We propose a plug-and-play framework for constructing efficient samplers for inverse problems, requiring only pre-trained diffusion or flow-matching models. We present Conditional Conjugate Integrators, which leverage the specific form of the inverse problem to project the respective conditional diffusion/flow dynamics into a more amenable space for sampling. Our method complements popular posterior approximation methods for solving inverse problems using diffusion/flow models. We evaluate the proposed method's performance on various linear image restoration tasks across multiple datasets, employing diffusion and flow-matching models. Notably, on challenging inverse problems like 4x super-resolution on the ImageNet dataset, our method can generate high-quality samples in as few as 5 conditional sampling steps and outperforms competing baselines requiring 20-1000 steps. Our code will be publicly available at https://github.com/mandt-lab/c-pigdm
Paper Structure (51 sections, 5 theorems, 103 equations, 12 figures, 8 tables, 2 algorithms)

This paper contains 51 sections, 5 theorems, 103 equations, 12 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Fast Samplers for Inverse Problems using Diffusions/Flows.
  3. Background and Problem Statement
  4. Conditional Conjugate Integrators
  5. Conjugate Integrators
  6. Conjugate Integrators for Inverse Problems
  7. Practical Design Choices
  8. Theoretical Aspects
  9. Extension to Noisy and Non-Linear Inverse Problems
  10. Experiments
  11. Models and Dataset:
  12. Tasks and Metrics:
  13. Methods and Baselines:
  14. Quantitative Results
  15. Qualitative Results
  16. ...and 36 more sections

Key Result

Proposition 1

(Extended $\Pi$GDM) For the conditional diffusion dynamics defined in Eqn. eqn:cond_diff, introducing a diffeomorphism, $\bar{{\mathbf{x}}}_t={\bm{A}}_t{\mathbf{x}}_t$, where, induces the following projected diffusion dynamics, where ${\bm{H}}^\dag = {\bm{H}}^\top({\bm{H}}{\bm{H}}^\top)^{-1}$ and ${\bm{P}}={\bm{H}}^\top({\bm{H}}{\bm{H}}^\top)^{-1}{\bm{H}}$ represent the pseudoinverse and the ort

Figures (12)

  • Figure 1: Illustration of Conditional Conjugate Integrators for Fast Sampling in Inverse Problems. Given an initial sampling latent ${\mathbf{x}}_{t_s}$ at time $t_s$, our sampler projects the diffusion/flow dynamics to a more amenable space for sampling using a projector operator $\Phi$ which is conditioned on the degradation operator ${\bm{H}}$ and the sampling guidance scale $w$. The diffusion/flow sampling is then performed in the projected space. Post completion, the generated sample in the projected space is transformed back into the original space using the inverse of the projection operator, yielding the final generated sample. We define the form of the operator $\Phi$ in Section \ref{['sec:cci_diffusion']}. Conditional Conjugate Integrators can significantly speed up sampling in challenging inverse problems and can generate high-quality samples in as few as 5 NFEs as compared to existing baselines, which require from 20-1000 NFEs (see Section \ref{['sec:experiment']}).
  • Figure 2: Qualitative comparison between C-$\Pi$G(D/F)M and $\Pi$G(D/F)M baselines on five different datasets. (\ref{['fig:afhq']}, \ref{['fig:lsun']}, \ref{['fig:celebahq']}) Inpainting, De-blurring, and 4x Super-resolution with C-$\Pi$GFM, respectively. (\ref{['fig:imagenet']},\ref{['fig:ffhq']}) 4x Image Super-resolution and De-blurring with C-$\Pi$GDM, respectively. ($\sigma_y=0$, NFE=5)
  • Figure 3: Impact of $\lambda$ and $w$ on sampling quality. Red curves and labels represent the LPIPS scores, while blue curves and labels indicate the FID scores.
  • Figure 4: Qualitative comparison between $\Pi$GDM and C-$\Pi$GDM at NFE=5 for the ImageNet dataset on the 4x Superresolution task. C-$\Pi$GDM can generate high-frequency details even for a low compute budget as compared to the baseline $\Pi$-GDM (Best Viewed when zoomed in)
  • Figure 5: Qualitative comparison for different sampling budgets for the ImageNet dataset on the 4x Superresolution task. C-$\Pi$GDM can generate high-quality samples in just 5 steps (Best Viewed when zoomed in)
  • ...and 7 more figures

Theorems & Definitions (10)

  • Proposition 1
  • Proposition 2
  • proof
  • Proposition
  • proof
  • Proposition
  • proof
  • proof
  • Proposition
  • proof