Table of Contents
Fetching ...

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.

Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel

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 . 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 . 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.
Paper Structure (29 sections, 10 theorems, 78 equations, 12 figures, 5 tables, 1 algorithm)

This paper contains 29 sections, 10 theorems, 78 equations, 12 figures, 5 tables, 1 algorithm.

Key Result

Proposition 1

Let $X,Z \in \mathbb{R}^d$ be independent random variables and $Y \in \mathbb R^n$ be another random variable. Assume that $T\colon\mathbb{R}^d\times\mathbb{R}^n\to \mathbb{R}^d$ fulfills $P_{T(Y,Z),Y}=P_{X,Y}$. Then, it holds $P_Y$-almost surely that $T(\cdot,y)_\#P_Z=P_{X|Y=y}.$

Figures (12)

  • Figure 1: Generated posterior samples, mean image and pixel-wise standard deviation for limited angle computed tomography using conditional MMD flows.
  • Figure 2: Class-conditional samples of MNIST, FashionMNIST and CIFAR10 and average class conditional FIDs. Note that these FID values are not comparable to unconditional FID values. A more detailed version is given in Table \ref{['table:FID_classcond']}.
  • Figure 3: Image inpainting and superresolution for different data sets.
  • Figure 4: Two different posterior samples and pixel-wise standard deviation for superresolution using conditional MMD flows.
  • Figure 5: Generated mean image, error towards ground truth and pixel-wise standard deviation for low dose computed tomography using conditional MMD flows.
  • ...and 7 more figures

Theorems & Definitions (16)

  • Proposition 1
  • Theorem 2
  • Theorem 3
  • Theorem 4
  • Lemma 5
  • Example 6: Theorem \ref{['thm:fundamental']} fails for non-compactly supported measures
  • Remark 7: Relation between joint and conditioned distributions for different distances
  • Lemma 8
  • Lemma 9: Stability under Pushforward
  • proof
  • ...and 6 more