Solving Diffusion Inverse Problems with Restart Posterior Sampling

TL;DR

This work addresses solving inverse problems with diffusion-prior models by introducing Restart for Posterior Sampling (RePS), a conditioned ODE-based sampler that avoids backpropagation through diffusion score networks. RePS integrates a MAP-based conditioning step into a DDIM-like ODE and employs a simple restart strategy that periodically re-injects noise to contract accumulated approximation errors, allowing efficient handling of both linear and nonlinear measurement models. Across FFHQ and ImageNet experiments, RePS demonstrates faster convergence and superior reconstruction quality relative to diffusion-based baselines, with ablations highlighting the benefits of the restart schedule and conditioning strategy. The approach offers a practical, computationally efficient pathway to leverage pre-trained diffusion priors for a wide range of inverse problems, potentially expanding diffusion-based plug-and-play imaging and signal reconstruction in real-world settings.

Abstract

Inverse problems are fundamental to science and engineering, where the goal is to infer an underlying signal or state from incomplete or noisy measurements. Recent approaches employ diffusion models as powerful implicit priors for such problems, owing to their ability to capture complex data distributions. However, existing diffusion-based methods for inverse problems often rely on strong approximations of the posterior distribution, require computationally expensive gradient backpropagation through the score network, or are restricted to linear measurement models. In this work, we propose Restart for Posterior Sampling (RePS), a general and efficient framework for solving both linear and non-linear inverse problems using pre-trained diffusion models. RePS builds on the idea of restart-based sampling, previously shown to improve sample quality in unconditional diffusion, and extends it to posterior inference. Our method employs a conditioned ODE applicable to any differentiable measurement model and introduces a simplified restart strategy that contracts accumulated approximation errors during sampling. Unlike some of the prior approaches, RePS avoids backpropagation through the score network, substantially reducing computational cost. We demonstrate that RePS achieves faster convergence and superior reconstruction quality compared to existing diffusion-based baselines across a range of inverse problems, including both linear and non-linear settings.