Inference-Time Search using Side Information for Diffusion-based Image Reconstruction
Mahdi Farahbakhsh, Vishnu Teja Kunde, Dileep Kalathil, Krishna Narayanan, Jean-Francois Chamberland
TL;DR
This work tackles ill-posed inverse problems by incorporating side information at inference time using a training-free, modality-agnostic framework. A reward function $r({f x}_0;{f s})$ tilts the unconditional diffusion prior via $p_{0|S}({f x}_0|{f s}) \\propto p_0({f x}_0) \,e^{r({f x}_0;{f s})/\tau}$, and the conditional reverse-score includes a value term $\nabla_{{\bf x}_t} V_t^{\tau}({\bf x}_t;{f s},{\bf y})$ approximated with DPS-style methods to avoid expensive backprop through the denoiser. The authors introduce two inference-time search strategies, Greedy Search (GS) and Recursive Fork-Join Search (RFJS), which use group-based resampling schedules to balance exploration and exploitation in the particle set. Experiments across six inverse-problem tasks (e.g., box inpainting, super-resolution, motion/Gaussian/nonlinear/deblurring) and multiple side-information modalities (images, text) show that the proposed methods consistently outperform strong baselines, including reward-gradient-guided approaches, while maintaining computational efficiency. The results demonstrate that leveraging side information at inference time can substantially improve reconstruction fidelity and semantic alignment in diffusion-based solvers, with practical impact for multimodal and measurement-limited imaging problems.
Abstract
Diffusion models have emerged as powerful priors for solving inverse problems. However, existing approaches typically overlook side information that could significantly improve reconstruction quality, especially in severely ill-posed settings. In this work, we propose a novel inference-time search algorithm that guides the sampling process using the side information in a manner that balances exploration and exploitation. This enables more accurate and reliable reconstructions, providing an alternative to the gradient-based guidance that is prone to reward-hacking artifacts. Our approach can be seamlessly integrated into a wide range of existing diffusion-based image reconstruction pipelines. Through extensive experiments on a number of inverse problems, such as box inpainting, super-resolution, and various deblurring tasks including motion, Gaussian, nonlinear, and blind deblurring, we show that our approach consistently improves the qualitative and quantitative performance of diffusion-based image reconstruction algorithms. We also show the superior performance of our approach with respect to other baselines, including reward gradient-based guidance algorithms. The code is available at \href{https://github.com/mhdfb/sideinfo-search-reconstruction}{this repository}.