Ensemble Kalman Diffusion Guidance: A Derivative-free Method for Inverse Problems

TL;DR

Ensemble Kalman Diffusion Guidance (EnKG) is a novel derivative-free method that leverages pre-trained diffusion models and only requires forward model evaluations, enabling solutions to complex inverse problems without privileged information.

Contribution

EnKG introduces a derivative-free framework for inverse problems using diffusion models and ensemble Kalman methods, applicable to scientific scenarios with black-box forward models.

Findings

Effective in scientific inverse problems like fluid flow and astronomy

Operates without derivative or privileged model information

Open-source implementation available

Abstract

When solving inverse problems, one increasingly popular approach is to use pre-trained diffusion models as plug-and-play priors. This framework can accommodate different forward models without re-training while preserving the generative capability of diffusion models. Despite their success in many imaging inverse problems, most existing methods rely on privileged information such as derivative, pseudo-inverse, or full knowledge about the forward model. This reliance poses a substantial limitation that restricts their use in a wide range of problems where such information is unavailable, such as in many scientific applications. We propose Ensemble Kalman Diffusion Guidance (EnKG), a derivative-free approach that can solve inverse problems by only accessing forward model evaluations and a pre-trained diffusion model prior. We study the empirical effectiveness of EnKG across various inverse problems, including scientific settings such as inferring fluid flows and astronomical objects, which are highly non-linear inverse problems that often only permit black-box access to the forward model. We open-source our code at https://github.com/devzhk/enkg-pytorch.