ReGuidance: A Simple Diffusion Wrapper for Boosting Sample Quality on Hard Inverse Problems

TL;DR

ReGuidance is a simple, effective wrapper that enhances the realism and reward quality of diffusion-based solutions for challenging inverse problems by inverting solutions and refining them with diffusion models.

Contribution

The paper introduces ReGuidance, a novel method that improves sample quality and reward in diffusion-based inverse problem solving, with theoretical guarantees on certain data distributions.

Findings

ReGuidance significantly improves sample realism and measurement consistency.

It outperforms state-of-the-art baselines on hard inverse tasks like in-painting and super-resolution.

Theoretical analysis shows it boosts reward and data manifold proximity for multimodal distributions.

Abstract

There has been a flurry of activity around using pretrained diffusion models as informed data priors for solving inverse problems, and more generally around steering these models using reward models. Training-free methods like diffusion posterior sampling (DPS) and its many variants have offered flexible heuristic algorithms for these tasks, but when the reward is not informative enough, e.g., in hard inverse problems with low signal-to-noise ratio, these techniques veer off the data manifold, failing to produce realistic outputs. In this work, we devise a simple wrapper, ReGuidance, for boosting both the sample realism and reward achieved by these methods. Given a candidate solution $\hat{x}$ produced by an algorithm of the user's choice, we propose inverting the solution by running the unconditional probability flow ODE in reverse starting from $\hat{x}$, and then using the resulting latent as an initialization for DPS. We evaluate our wrapper on hard inverse problems like large box in-painting and super-resolution with high upscaling. Whereas state-of-the-art baselines visibly fail, we find that applying our wrapper on top of these baselines significantly boosts sample quality and measurement consistency. We complement these findings with theory proving that on certain multimodal data distributions, ReGuidance simultaneously boosts the reward and brings the candidate solution closer to the data manifold. To our knowledge, this constitutes the first rigorous algorithmic guarantee for DPS.