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 | |
|---|---|---|---|---|
| MRI Reconstruction | fastMRI Brain (test) | SSIM0.953 | 18 | |
| Fluid Data Assimilation | vorticity fields 128x128 (test) | Relative L2 Error0.32 | 12 | |
| Black Hole Imaging | InverseBench Black-Hole Imaging (test) | PSNR26.21 | 9 | |
| Image Inpainting | FFHQ 95% missing random mask 100 images (val) | PSNR22.64 | 9 | |
| Image Deblurring | FFHQ Gaussian deblurring 100 images (val) | PSNR20.84 | 9 | |
| Image Super-resolution | FFHQ 4x super-resolution 100 images (val) | PSNR21.72 | 9 | |
| Navier–Stokes inverse problem | InverseBench 128 x 128 (test) | NRMSE (σnoise=0)0.12 | 6 | |
| JPEG restoration | FFHQ JPEG QF=5 100 images (val) | PSNR20.89 | 5 | |
| Black Hole Imaging | Black Hole Imaging 100% observation ratio | PSNR23.03 | 4 | |
| Black Hole Imaging | Black Hole Imaging 10% observation ratio | PSNR22.77 | 4 |