Mode-seeking for inverse problems with diffusion models

Sai bharath chandra Gutha, Ricardo Vinuesa, Hossein Azizpour

Keywords: Inverse Problems, Diffusion Models, Mode-seeking loss, MAP estimation

TL;DR: Inference-time guidance method using a mode-seeking loss function for solving inverse problems with pre-trained diffusion models

Abstract: A pre-trained unconditional diffusion model, combined with posterior sampling or maximum a posteriori (MAP) estimation techniques, can solve arbitrary inverse problems without task-specific training or fine-tuning. However, existing posterior sampling and MAP estimation methods often rely on modeling approximations and can be computationally demanding. In this work, we propose the variational mode-seeking loss (VML), which, when minimized during each reverse diffusion step, guides the generated sample towards the MAP estimate. VML arises from a novel perspective of minimizing the Kullback-Leibler (KL) divergence between the diffusion posterior $p(\mathbf{x_0}|\mathbf{x_t})$ and the measurement posterior $p(\mathbf{x_0}|\mathbf{y})$, where $\mathbf{y}$ denotes the measurement. Importantly, for linear inverse problems, VML can be analytically derived and need not be approximated. Based on further theoretical insights, we propose VML-MAP, an empirically effective algorithm for solving inverse problems, and validate its efficacy over existing methods in both performance and computational time, through extensive experiments on diverse image-restoration tasks across multiple datasets.

Primary Area: generative models

Submission Number: 24649