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

Authors: Hongkai Zheng, Wenda Chu, Austin Wang, Nikola Borislavov Kovachki, Ricardo Baptista, Yisong Yue

TMLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental We empirically study our En KG method on the classic image restoration problems and two scientific inverse problems. We view the scientific inverse problems as the more interesting domains for evaluating our method, particularly the Navier-Stokes equation where it is impractical to accurately compute the gradient of the forward model.
Researcher Affiliation Collaboration Hongkai Zheng EMAIL Caltech Wenda Chu EMAIL Caltech Austin Wang EMAIL Caltech Nikola B. Kovachki EMAIL NVIDIA Ricardo Baptista EMAIL Caltech Yisong Yue EMAIL Caltech
Pseudocode Yes Algorithm 1 Generic Guidance-based Method (ODE version) ... Algorithm 2 Our method: Ensemble Kalman Diffusion Guidance (En KG).
Open Source Code Yes We open-source our code at https://github.com/devzhk/enkg-pytorch.
Open Datasets Yes The pre-trained diffusion model for FFHQ 64 64 is taken unmodified from Karras et al. (2022). The model for FFHQ 256 256 is taken from Chung et al. (2022a) and converted to the EDM framework (Karras et al., 2022) using their Variance-Preserving (VP) preconditioning. ... We trained a diffusion model on the GRMHD dataset (Wong et al., 2022) with resolution 64 64 to generate black hole images from telescope measurements.
Dataset Splits Yes We sample 20,000 vorticity fields to train our diffusion model. Then, we independently sample 10 random vorticity fields as the test set.
Hardware Specification Yes Runtime is reported on a single A100 GPU. ... autograd encounters out-of-memory issues when the pseudospectral solver unrolls beyond approximately 6k steps on an A100-40GB GPU.
Software Dependencies No The model for FFHQ 256 256 is taken from Chung et al. (2022a) and converted to the EDM framework (Karras et al., 2022) using their Variance-Preserving (VP) preconditioning. ... Our experiments implement it as a pseudo-spectral solver (He & Sun, 2007). While these mention software/frameworks, they lack specific version numbers.
Experiment Setup Yes Measurement noise σ = 0.05 in all experiments except for superresolution on 64 64 images, where we set σ = 0.01. ... Gaussian deblurring, a blurring kernel of size 61 61 with standard deviation 3.0. ... Our experiments also reveal that a guidance scale of 2.0 yields the best average performance across our experiments. ... We use 1024 particles for FFHQ and black hole imaging, and 2048 particles for Navier-Stokes.