Ensemble Kalman Diffusion Guidance: A Derivative-free Method for Inverse Problems
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Posterior Sampling | Linear Gaussian n=2 | SWD1.752 | 24 | |
| Posterior Sampling | Linear Gaussian n=80 | SWD4.938 | 24 | |
| Posterior Sampling | Linear Gaussian n=400 | SWD7.432 | 24 | |
| MRI Reconstruction | fastMRI Brain (test) | SSIM0.953 | 18 | |
| Black Hole Imaging | InverseBench Black-Hole Imaging (test) | PSNR26.21 | 17 | |
| Phase Retrieval | FFHQ 256x256 (test) | LPIPS0.232 | 16 | |
| Inverse Problem | InverseBench Navier-Stokes (sigma_noise=1.0) (test) | Relative L2 Error0.191 | 13 | |
| Inverse Problem | InverseBench Navier-Stokes (sigma_noise=2.0) (test) | Relative L2 Error0.294 | 13 | |
| Inverse Problem | InverseBench Navier-Stokes (sigma_noise=0) (test) | Relative L2 Error0.12 | 13 | |
| Fluid Data Assimilation | vorticity fields 128x128 (test) | Relative L2 Error0.32 | 12 |