SITCOM: Step-wise Triple-Consistent Diffusion Sampling For Inverse Problems

Ismail Alkhouri, Shijun Liang, Cheng-Han Huang, Jimmy Dai, Qing Qu, Saiprasad Ravishankar, Rongrong Wang

Abstract: Diffusion models (DMs) are a class of generative models that allow sampling from a distribution learned over a training set. When applied to solving inverse problems, the reverse sampling steps are modified to approximately sample from a measurement-conditioned distribution. However, these modifications may be unsuitable for certain settings (e.g., presence of measurement noise) and non-linear tasks, as they often struggle to correct errors from earlier steps and generally require a large number of optimization and/or sampling steps. To address these challenges, we state three conditions for achieving measurement-consistent diffusion trajectories. Building on these conditions, we propose a new optimization-based sampling method that not only enforces standard data manifold measurement consistency and forward diffusion consistency, as seen in previous studies, but also incorporates our proposed step-wise and network-regularized backward diffusion consistency that maintains a diffusion trajectory by optimizing over the input of the pre-trained model at every sampling step. By enforcing these conditions (implicitly or explicitly), our sampler requires significantly fewer reverse steps. Therefore, we refer to our method as **S**tep-w**i**se **T**riple-**Co**nsistent Sa**m**pling (**SITCOM**). Compared to SOTA baselines, our experiments across several linear and non-linear tasks (with natural and medical images) demonstrate that SITCOM achieves competitive or superior results in terms of standard similarity metrics and run-time.

Lay Summary: Modern image recovery techniques are increasingly powered by diffusion models, a class of AI tools that can generate high-quality images. To apply these models to inverse problems (where we aim to recover an image from partial or corrupted measurements), researchers often adapt the standard generation process to account for the measurements. However, these adapted methods tend to be computationally expensive, especially when the measurements are noisy or the task is non-linear. In our work, we identify three key conditions that any measurement-guided sampling process should meet to stay consistent with both the data and the underlying model. Based on this, we introduce a new method called SITCOM (Step-wise Triple-consistent Sampling) that enforces these conditions during the sampling process. Our results demonstrate that SITCOM recovers images more reliably and efficiently—achieving high-quality results in fewer steps. We show its benefits across nine image recovery tasks, including real-world medical imaging.

Link To Code: https://github.com/sjames40/SITCOM

Primary Area: Applications->Computer Vision

Keywords: Diffusion Model; Inverse Problems; Image Restoration; MRI

Submission Number: 7438