Fast Samplers for Inverse Problems in Iterative Refinement Models

TL;DR

This paper introduces a plug-and-play framework with Conditional Conjugate Integrators for efficient inverse problem sampling using diffusion and flow models, significantly reducing the number of steps needed for high-quality results.

Contribution

It presents a novel method that leverages the structure of inverse problems to enable fast sampling with pre-trained models, outperforming existing approaches in efficiency and quality.

Findings

High-quality 4x super-resolution with only 5 steps

Outperforms baselines requiring 20-1000 steps

Effective on multiple datasets and inverse tasks

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