Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel
Paul Hagemann, Johannes Hertrich, Fabian Altekrüger, Robert Beinert, Jannis Chemseddine, Gabriele Steidl
TL;DR
The paper tackles the problem of posterior sampling for ill-posed imaging inverse problems by introducing conditional MMD flows based on the energy distance with kernel $K(x,y)=-\|x-y\|$. It develops a theoretical framework linking joint and posterior distributions through Wasserstein gradient flows of conditioned MMD functionals and provides a practical generative scheme that approximates the gradient flow with neural networks to produce posterior samples $P_{X|Y=y}$. Key contributions include finite-sample error bounds between posterior and joint distributions in the MMD metric, a proof that the conditional particle flow is a Wasserstein gradient flow of a conditioned functional, and an end-to-end conditioned generative MMD flow that can handle high-dimensional inverse problems. The method is validated on class-conditional image generation and inverse problems such as inpainting, superresolution, and low-dose/limited-angle CT, showing competitive performance and meaningful uncertainty quantification. This work offers a scalable, principled approach to uncertainty-aware posterior sampling in high-dimensional imaging tasks and bridges joint-distribution approximations with conditional posteriors.
Abstract
We propose conditional flows of the maximum mean discrepancy (MMD) with the negative distance kernel for posterior sampling and conditional generative modeling. This MMD, which is also known as energy distance, has several advantageous properties like efficient computation via slicing and sorting. We approximate the joint distribution of the ground truth and the observations using discrete Wasserstein gradient flows and establish an error bound for the posterior distributions. Further, we prove that our particle flow is indeed a Wasserstein gradient flow of an appropriate functional. The power of our method is demonstrated by numerical examples including conditional image generation and inverse problems like superresolution, inpainting and computed tomography in low-dose and limited-angle settings.
