Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel
Authors: Paul Hagemann
Johannes Hertrich
Fabian Altekrüger
Robert Beinert
Jannis Chemseddine
Gabriele Steidl
Venue: ICLR 2024 poster (2024)
Keywords: Bayesian inverse Problems
MMD
Gradient Flows
Deep Learning
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We establish a negative distance kernel MMD flow to the joint distribution, which allows for posterior sampling in Bayesian inverse problems.

Abstract: We propose conditional flows of the maximum mean discrepancy (MMD) with the negative distance kernel for posterior sampling and conditional generative modelling. 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.

Published as a conference paper ICLR 2024
POSTERIOR SAMPLING BASED ON GRADIENT FLOWS OF THE MMD WITH NEGATIVE DISTANCE KERNEL
P. Hagemann', J. Hertrich2, F. Altekrügger",, Beinert', J. Chemseddine', G. Steidl Technische Universität Berlin, 2 University College London," 3 Humboldt-Univerrsstät zu Berlin Correspondence to: hagemann@math. tu-berlin.dee
ABSTRACT
We propose conditional flows of the maximum mean discrepancy (MMD) with the neg- ative distance kernel for posterior sampling and conditional generative modelling. This MMD, which is also known 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 Wasserstein gradient flow of an appropriate functional. The power of method is demonstrated by numerical examples including conditional image generation and in- verse problems like superresolution, inpainting and computed tomography in low-dose and limited-angle settings.
1 INTRODUCTION
The tremendous success of generative models led to a rising interest in their application for inverse problems in imaging. Here, an unknown image x has to be recovered from noisy observation y == +.. Since the forward operator I' is usually ill-conditioned, such reconstructions include uncertainties and are usually not unique. As remedy, we take Bayesian viewpoint and consider and y as samples from random variables X and Y. and assume that we are given training samples from their joint distribution Px,y. In order represent the uncertainties in the reconstruction, we aim to find process sample from the posterior distributions PxIy= This allows not only to derive different possible predictions, but also to consider pixel- wise standard deviations for identifying highly vague image regions. Figure visualizes this procedure on an example for limited angle computed tomography.
Nowadays, generative models like (Wasserstein) GANs (Arjovsky et al. 2017 Goodfellow et al. 2014) and VAEs (Kingma & Welling, 2014) have turned out to be suitable tool for approximating probability distri- butions. In this context, the field of gradient flows in measure spaces received increasing attention. |Welling & Tch 2011) proposed to apply the Langevin dynamics in order to generate samples from a known poten- tial, which corresponds to simulating Wasserstein gradient flow with respect the Kullback-Leibler (KI.) divergence, see Jordan et al. (1998). Score- and diffusion models extend this approach by estimating the gradients of the potential from training data, see (De Bortoli et al. 2021 Ho et al. 22000] Song & Ermon 2020 [Song et al. (021) and achieved state-of-the-art results. The simulation of Wasserstein gradient flows with other functionals than KL, based on the JKO scheme, was considered in Altekrüger et al. (003c); Alvarez-Melle et al. (2022); : Fan et al. (2022); Mokrov et al. (202I).
this paper, we focus on gradient flows with respect to MMD with negative distance kernel K 3) which is also known energy distance, see (Sejdinovic al. 2013 Székely 2002 Székely
& Rizzo 2009 2013). While MMDs have shown great success at comparing two distributions in general, see (Gretton et al. 2012 Székely & Rizzo 2005 Gretton et al. 2006), their combination with the negative

---

Published as a conference paper ICLR 2024
distance kernel results in many additional desirable properties as translation and scale equivariance (Székely & Rizzo 2013). efficient computation (Hertrich et al. 2024). dimension independent sample complexity
of O( (/2) (Gretton et al. 2012) and unbiased sample gradients Bellemare BeallmattBBeleemaree et al. (017).
To work with probability distributions in high dimensions, RRabin al. (012) proposed to slice them. Ap- plied on gradient flows, this leads to a significant speed-up. see (Du et al. 2023 Kolouri et al. 2019b] Liutkus et al. 2019). In particular, for MMD with negative distance kernel slicing does not change the met- ric itself and reduces the time complexity of calculating gradients from ( to ( log N) for measures with N support points, see Hertrich et al. (2024).
In order to use generative models for inverse problems, an additional conditioning parameter was added to the generation process in (Ardizzone et al. 2019; 202T; Chung al., [023] Hagemann et al.] [022] Mirza
& Osindero 2014). However, this approach cannot directly applied to the gradient flow setting, where the generator is not trained end-to-end.
Contribations. We simulate conditional MMD particle flows for posterior sampling in Bayesian inverse problems, To this end, we provide three kinds contributions. The first two address theoretical questions while the last one validates our findings numerically.
Conditional generative models approximate the joint distribution by learning a mapping T such that Px,y R Pr(z.y)).. but in fact we are interested the posterior distributions Px Y=y In this paper, prove error bounds between posterior and joint distributions within the MMD metrie in expectation. The proofs of these results are based on relations between measure spaces and RKHS as well as Lipschitz stability results under pushforwards.
We represent the considered particle flows as Wasserstein gradient flows of modified MMD func- tional. As side effect of this representation, we can provide a theoretical justification for the empirical method presented by Du et al. (2023), where the authors obtain convincing results by neglecting the velocity in the y-component in sliced Wasserstein gradient flows. Based on locally isometric embeddings of the N into the Wasserstein space, we can show that the result is again a Wasserstein gradient flow with respect to modified functional.
We approximate our particle flows by conditional generative neural networks and apply the aris- ing generative model in various settings. On the one hand, this includes standard test sets like conditional image generation and inpainting on MNIST, FashionMNIST and CIFARIO and su- perresolution on CelebA. On the other hand, we consider very high-dimensional imaging inverse problems like superresolution of materials' microstructures as well as limited-angle and low-dose computed tomography.
Related work. Many generative models like GANs, VAEs and normalizing flows can be used for posterior sampling by adding conditioning parameter as an additional input, see (Ardizzone et al. 2019 [02T Batzolis et al. 2021 Hagemann et al. [022 Mirza & Osindero) 2014). The loss function of these methods is based on the Kullback-Leeibe divergence. In this case, stability results were proven by AAteerügge al. (2023b) based on local Lipschitz regularity results from (prungkk (002)).
this paper, we are interested in generative models which are based on gradient flows (Fan et al. 2022 Kolouri et al.) 2019b; Hertrich et al. 2024 2023b Nguyen & Ho 2022, Mokrov et al. 2021). In this case, the above approach is not directly applicable since the network is not trained end-to-end. For the sliced Wasserstein gradient flows, (Du et al. (2023) proposed to approximate the joint distribution while neglecting the velocity in one component. They achieved very promising results, but evaluated their model empirically without giving theoretical justification.
Here, we consider gradient flows with respect to MMD with negative distance kernel, which is also known energy distance or Cramer distance (Sejdinovie et al. 2013] Székely 2002). The theoretical analysis
2

---

Published as conference paper ICLR 2024
GT
FBP
Reco 1
Reco 2
Reco 3
Mcan
Std
Figure 1: Generated posterior samples, mean image and pixel-wise standard deviation for limited angle computed tomography using conditional MMD flows.
of such gradient flows in the Wasserstein space turns out to be challenging, since discrete measures might become absolutely continuous and vice versa, see (Hertrich et al., 2023b). As a remedy. many papers consider particle flows
space discretization, see (Carrillo et al. 2020 Daneshmand & Bach [023]
Daneshmand et al. 2023 Hertrich et al. 2024). The question whether the mean-field limit of these particle flows corresponds the continuous Wasserstein gradient flow is far only partially answered and still active area of research, see (Carrillo et af.] 2020; Daneshmand & Bach 2023; Daneshmand et al.) (003))
Further, there is plenty literature covering the statistical properties of MMD general (Sriperumbudur ct al., 2011] Gretton et al. 2012 Modeste & Dombry. (023) as well as its applications to causal inference
(Kremer et 2022 2023 via conditional moments.
Outline of the paper. Section 2 briefly recalls MMD with negative distance kernel and corresponding discrete Wasserstein gradient flows. Based on this, we introduce conditional generative MMD flows Section 3 by the following path: i) we establish relations between joint and posterior distributions, ii) we present an interpretation of the conditional particle flow as Wasserstein gradient flow of an appropriate functional, and iii) suggest generative variant of the conditional MMD flow. Numerical experiments are contained in Section and conclusions are drawn in Section The appendix contains all the proofs, implementation details and additional experimental results.
2 MMD AND DISCRETE GRADIENT FLOWS
Let P, (R*),, E (0, oc) denote the space of probability measures with finite p-th moments. We are interested
in the MMD D (2) x P2 (IRd) + R of the negative distance kernel K(x,, y) == -|| - yil defined by
(1, = 2 R- [ K(r,y) du((x) dp(y) - -J | K(x, y) du(x)) du(v)
(1)
1
{ [ K dv(x)) dy(y).
+ 2
It is metric on (RD) 3 P2 (RR)), see e.g. Sejdinovic et al. 2013 : Székely & Rizzo 2013). In particular,
we have that DR (m,, v) = 0 and only if μ
For a fixed measure E (RRd), we consider Wasserstein gradient flows of the functional F, (2(
-
(-o0, oo] defined by
Fr(1) Di (,. const,
(2)
3

---

Published as a conference paper ICLR 2024
where the constant is just the third summand in (T.. For definition of Wasserstein gradient flows, we refer to Appendix A.3) Since the minumum of this functional is we can use the flows to sample from this target distribution. While the general analysis of these flows theoretically challenging. in particular for the above non-smooth and non-À-convex negative distance kernel, see Carrillo et al. (2020;; Herrtich et al. (2023b), we focus on the efficient numerical simulation via particle flows as proposed in Hertrich et al. (2024). To this end, let PN == < Σi= 62, b2 E RR D2 * Iz. + j} denote the set of empirical measures on R with N pairwise different anchor points. Given M independent samples = (Pi....,AA E (RR) of the measure v, we deal with its empirical version P/ M 2i=1 Sp. and consider the Euclidean gradient flow of the discrete MMD functional Fp: (RRJ)N
IR defined for (21,,,,, E (R4) by
N
N
1 MN
Fœ(x) = Fuss(uNN) 2N
Σ
2 xyll
b0 Pyll-
(3)
Then, a curve = (WI, uN): 00) + (R3) solves the ODE
ù ---FNPP(u)) w(0) (u (0)
(0)
N
if and only if the curve : flows with respect to the functional Jume: P2( + R [00} given by
: (0, 0o) - P2 defined by YN (t) = N E Euz(0) is Wasserstein gradient
Tun (p) = {{ Fuse (1), if PN. otherwise.
00,
In Hertrich et al. (2024), this simulation of MMD flows was used to derive generative model.
3 CONDITIONAL MMD FLOWS FOR POSTERIOR SAMPLING
In this section, we propose conditional flows for posterior sampling. We consider two random variables X E Rdd and Y E R.. Then, we aim to sample from the posterior distribution Px One the most important applications for this are Bayesian inverse problems. Here X and Y are related by Y noisy(f(X)), where f: R: + R" is some ill pooeed forward operator and "noisy" denotes some noise process. Throughout this paper, we assume that we are only given samples from the joint distribution Рx,.. In order to sample from the posterior distribution Pxjy=y- we will use conditional generative models. More precisely, we aim to find mapping T: Rdd R Rd such that
T(.,) = Px|N-
(4)
where Pz is an easy-to-sample latent distribution and ward of measure. The following proposition summarizes the main principle of conditional generative modelling and provides a sufficient criterion for 4 To make the paper self-contained, we add the proof in Appendix A.I
defines the pushfor-
Proposition 1. Let X, Z E Rd be independent random variables and Y E R be another random vari- able. Assume that T: IRd x R" - Rdd fulfills Pr(y.2), = -x,x. Then, it holds Py-almost surely that
T(.8) Pz = P
In Subsection 3.1
extend this result and show that (4) still holds true approximately, whenever the distributions P and Рx, are close to each other. Then, in Subsection 3.2 we construct such map- ping T based on Wasserstein gradient flows with respect to conditioned version of the functional Frx Finally, we propose an approximation of this Wasserstein gradient flow by generative neural networks in Subsection
4

---

Published as a conference paper ICLR 2024
3.1 POSTERIOR VERSUS LEARNED JOINT DISTRIBUTION
In Proposition 1 we assume that T is perfectly learned, which is rarely the case in practice. Usually, we can only approximate the joint distribution such that Dx (Prrz.y).; P becomes small. Fortunately, we prove under moderate additional assumptions that then the expectation with respect to y of the distance of the posteriors DK (T(., y) Pz. Px(r=w) becomes small too. A similar statement was shown by Kim jt al. (2023) Proposition 3); note that their RHS is not equal to the MMD of joint distribution, but modited version. Such a statement involving MMDs can not generally hold true, see Example
Theorem 2. Let Sn C Rddh and Sa C Rdd be compact sets. Further, let X, X E Sa and Y E S. be absolutely continuous random variables, and assume that Px, and Px,x have densities fulfitiss
P -x 1 Cs., 11 - 32l1
(5)
PXY= - -PXIY= 1 Cs., 11 3221
(6)
a.e. on for all y1, V2 E Sn. Then it holds
Ey-py
(CD Px.) P saanisys
(7)
The proof, which uses relations between measure spaces and reproducing kernel Hilbert spaces (RKHS), given in Appendix A.2 Similar relations hold true for other ,distances" the Kullback- Leibler divergence or the Wasserstein distance, see Remark [ For the latter one as well as for the MMD, it is necessary that the random variables are compactly supported. Example shows that ( is in general not correct if this assumption is neglected.
Based Theorem it was done for Wasserstein distances by Altekrüger et al. (2023b). For this, we require the existence of C, C > 0 such that
can prove pointwise convergence of sequence of mappings in similar way
(x P2. (1,, 92) Pz) - Clly y/
(stability under pushforwards),
(8)
DK Dx (PXY== PXIY= = - Čllynn 3/
(stability of posteriors)
(9)
for all 91, 3/ E Sn with py (Wi), Py (V2) 0. In Appendix A.2 we show that both stability estimates hold true under certain conditions, see Lemma 9 and 10, Furthermore, we prove following theorem.
Theorem 3. Let Sn C R and Sa C Rdd be compact sets and X, Z Sa and Y E Sn absolutely continuous random variables. Assume that lpy (y)) py
C'1ly: y2ll _ for all y1, 1/ E Sn and fixed C' O.
Let Px, have density satisfying (5. Moreover, let {T I2d x R" + Rd} be family of measurable
mappings with Dx (PP Px.) S 8, which fulfill (8) and ( with uniform C, C > 0. Further,
assume that Px,y with X = =T;(ZZ,)) have densities satisfying (6) with uniform Cs., > Then, for all
y E Sn with py (y) > 0, it holds
PxIv=)) +0 as E -
3.2 CONDITIONAL MMD FLOWS
In the following, we consider particle flows approximate the joint distribution Рx,.. Together with the results from the previous subsection, this imposes an approximation of the posterior distributions Px Let N pairwise distinet samples (pi, 4)) E Rd Rn from the joint distribution Рx, be given, and set (p, q) = ((ps, (:)))))))) )) Let Pz be a d-dimensional latent distribution, where we can easily sample from. We draw sample 2, for each and consider the particle flow t >> (u(t), q) starting ((2,, q: 9:)) where the
5

---

Published as a conference paper ICLR 2024
second component remains fixed and the first component (u1 uN): [0, 00) + Rd follows the gradient flow
i(t)= - (p.a) ((u, q)), a(O) (21 (21,...,2N)
(10)
with the function F(p-.) in (3) and the gradient V with respect to the first component. Since the MMD a metric, the function r + (p-n) ((x, q)) admits the global minimizer
(p1 (....... PN). In numerical examples, we observe that the gradient flow (IO) approaches this global minimizer as t + co. Finally, we approximate the motion of the particles by a mapping T: Rd R Red which describes how the initial particle (zi, (:) moves to the particle (4(4max) at some fixed time maxx Moreover, by the convergence of the gradient flow IOP, we have that
N ST(Zen)e
PriZY)y
Σ Emg(tmax)) 3.
16 6pea R Px,.
particular, the arising mapping T fulfills the assumptions from the previous subsection such that we obtain
T(:9) P2 = Px]Y= for Py y E...
The following theorem states that the solutions of (TP correspond to Wasserstein gradient flows with respect conditioned MMD functional. To this end, set PN..- == 1 Ein day,, (Iz; 9:) E I2 R" (zi, 4:) + (Iy, 9j), j}-
Theorem 4. For given (Pi: 9:)
1, 1,..... N, UN (u1, .... uN): [0, 00) + (Bd)) be a solution of (O),, and assume (t,(t), + (uy(t), for + j and all
# E Sprer Let
N
O. Then the curve YN.g: 00) (2 defined by
N
(t)
1=1
is Wasserstein gradient flow with respect to the functional Juna: P2(RR - RU (00) given by
if PN.g. otherwise,
Jux.9 = 1000
00,
where FuN. is the functional defined in 22
The definition of Wasserstein gradient flows and the proof are included in Appendix A.3 will see that there is no flow the second component.
In particular, we
In order to approximate the joint distribution Px,v starting in Pz,,, Du et al. (023 obtained convincing results by considering (discretized) Wasserstein gradient flows with respect sliced Wasserstein distance
SW3 v) = BeEsa-: (Ps (P PS p)) Pę(x) = (E, x).
They observed in their experiments that there is nearly no flow in the second component, but acknowledged that they "are unable to provide a rigorous theoretical justification for the time being." " Our proof of Theorem 4 more precisely Corollary 13 Appendix A.3 delivers the theoretical justification of their empirical result.
3.3 CONDITIONAL GENERATIVE MMD FLOWS
this subsection we want to learn the mapping T: Rd x R + Rd describing the evolution, how the initial particles (2,, (:) move to uu (fm()))), where (u,, ...; uN): [0, oo) - RdN solves the ODE (TO. To this
6

---

Published as a conference paper ICLR 2024
end, we adopt the generative MMD flows from Hertrich al. (2024) and simulate the ODE (TO) using an (le)
explicit Euler scheme. More precisely, we compute iteratively u (A) [ui . ..... "N by
(k+1) (k) - TNV V. (P.g) .q),,
u 24-
(11)
In order evaluate the gradient efficiently and to speed up the computations, we use the sliced computation
of V, F(no)) ((u(*),g) ,9) and the momentum form of (T1) proposed in Hertrich et al. (0024). Now, we train
neural networks P1...., 4y TRd xR= - IRd taking (us; (i) as an input, such that each network approximates a fixed number Ti of explicit Fuler steps from (11). The precise training procedure is outlined in Algorithm
the appendix. Once the networks 4 are trained, the approximating mapping T is given by 7(:
(x) 0....00 2 (- y). In particular, we can generate samples from the approximated posterior distribution Px|y=y by drawing ~ Pz and computing T(z, y), even for samples y of which are not contained in the training set.
4 EXPERIMENTS
We apply our conditional generative MMD flows to generate images for given conditions y in two settings, namely class-conditional image generation, and ii) reconstruction from posterior distributions Px|y= inverse problems. The chosen networks
( are UNets (Ronneberger et al. 2015), where we adopted the
implementation from Huang et al. (202T) based on Ho et al. (2020). Further details are given in Appendix
4.1 CLASS-CONDITIONAL IMAGE GENERATION
We choose the condition Y to be the one-hot vectors of the class labels in order to generate samples of MNIST (LeCun et al. 1998). FashionMNIST (Xiao et al. (017) and CIFARIO (Krizhevsky) (009) for given class labels. Figure illustrates the generated samples and reports the average class conditional FID values. That is, compute the FID between the generated samples from specific class with all test samples from the same class. Note that class conditional FID values are not comparable with unconditional FIDs. We compare the results with f-SWF of Du et al. 2023). J We observe that the conditional MMD flow generates samples of good visual quality and outperforms Du et al. (023). Further examples are given in Appendix D
4.2 INVERSE PROBLEMS
Inpainting. For the inpainting task with the mask operator let y be the partially observed images. In- painted images of MNIST, FashionMNIST and CIFARIO are shown in Figure 3 The observed images are the leftmost ones, while the unknown ground truth images are given in the rightmost column. The various generated samples in the middle column have very good reconstruction quality and are in particular consis- tent with the observed part. Their pixelwise standard deviation (std) is given in the second last column.
Superresolution. For image superresolution, we consider the superresolution operator f and low-resolution images y. Reconstructed high-resolution images of CelebA (Liu et al. 2015) are illustrated Figure For CelebA, we centercrop the images to 140 x 140 and then bicubicely downsample them to 64 x 64. Again, the observations are the leftmost ones, the unknown ground truth the rightmost ones and the reconstructions in the middle column are of good quality and high variety. Another superresolution example on high- dimensional and real-world images materials' microstructures is presented in Figure 4 We benchmark conditional MMD flow against a conditional normalizing flow "SRFlow" (Lugmayr et al. 2020) and
WPPFlow (Altkrügee & Hertrich, (023) in terms of PSNR and SSIM (Wang et al. 2004). A more detailed description given in Appendix
Computed Tomography. The forward operator f is the discretized linear Radon transform and the noise process is given by a scaled negative log-Poisson noise, for details see, e.g., /Alterüüee et al. 2023a);
7

---

Published as conference paper ICLR 2024
0
C
O
0
0
O
D
0
2
1
/
/
27
2
2
2
2
2
Class conditional
E-SWF
Cond. MMD Flow
3
3
3
3
3
3 4
(ours)
9
H 4 5
4 5
H
MNIST
18.9
14.6
5
5-
5
5
RishionMNIST
25.9
27.5
CIFARIO
118.1
101.8
6
6
6
7
7
7 7
7
8
8
8
9
9
9
4 9
MNIST
FaShionMNIST
CIFARIO
Figure 2: Class-conditional samples of MNIST, FNsHionMNIST and CIFAR10 and average class conditional FIDs. Note that these FID values are not comparable to unconditional FID values. more detailed version is given in Table T
7
4
4
7
7
7
3
4
4
/
4
4
4
4
4
1
7
4
4
4
4
0O0
0
b
D
0
0
0
C
O
6
6
D
O
O
O
6
0
0
8
O
8
Image inpainting (from top bottom) for MNIST, FashionMNIST, CIFARIO.
(b) Image superresolution for CelebA with magnification factor
Figure 3: Image inpainting and superresolution for different data sets
PSNR
SSIM
WPPP
25.89
0.657
AAltkkügger & Hertrich 2023
Lugmayr et al 2020
25.11
0.637
MMD
27.21
0.774
(ours)
HR image
LR image
Prediction
Prediction 2
Std
Figure 4: Two different posterior samples and pixel-wise standard deviation for superresolution using ditional MMD flows.
8

---

Published as conference paper ICLR 2024
PSNR
SSIM
WPPF
31.16
0.774
(Altekrüger Hertrich 2023
Cond. Norma lizine Flow
35.07
0.831
Denker al. 2021 Cond. MMD Flow
35.37
0.835
(ours)
GT
FBP
Mean
Error
Std
Figure 5: Generated mean image, error towards ground truth and pixel-wise standard deviation for low dose computed tomography using conditional MMD flows.
Leuschner et al. (202I). The data taken from the LoDoPaB dataset of Leuschner et al. (2021) for low-dose CT imaging. The results are shown in Figure [5] We illustrate the mean image of 100 reconstructions, its error towards the ground truth and the standard deviation of the reconstructions. A quantitative comparison with (Denker et al. 202I) for this example is given in Fig [ Here provide PSNR and SSIM of the mean images for the whole testset containing 3553 images. More examples towards limited angle CT and low-dose CT are given in Figures and 10 in Appendix D
5 CONCLUSION
We introduced conditional MMD flows with negative distance kernel and applied them for posterior sam- pling in inverse problems by approximating the joint distribution. To prove stability of our model, we bounded the expected approximation error of the posterior distribution by the error of the joint distribution. We represented conditional MMD flows as Wasserstein gradient flows, which also provides additional insights for the recent paper by Du et al. (2023). Finally, we applied our algorithm to conditional image gen- eration, inpainting, superresolution and CT. From theoretical viewpoint it would be interesting to check whether the dimension scaling in Theorem 2 is optimal. In this paper we focused on the negative distance kernel because it can be computed efficiently via slicing and sorting. It would be interesting if similar results hold for other kernels. Moreover, so far we only considered discrete gradient flows with a fixed number N of particles. The convergence properties of these flows in the mean-field limit N oo is only partially answered in the literature, see c.g. (Carrilo et al., 2020) for the one-dimensional setting or (Arbel et al.) 2019) for result with smooth kernels. From a numerical perspective, it could be beneficial consider more general slicing procedures, see e.g. (Kolouri et al. 2019a Nguyen & Ho 2023), or other kernels, see Hertrich| (024).
Limitations. It is the aim of our numerical examples to demonstrate that our method can be used for highly ill-posed and high- sith--miensionll imaging inverse problems. particular, do not claim that our computed tomography experiments are realistic for clinical applications. In practice, the availability and potential biases of high-quality datasets are critical bottlenecks in medical imaging. Moreover, even slight changes the forward operator noise model require that the whole model is retrained. which is computational costly and demands a corresponding dataset. Since the particles are interacting, an important next step would be to enable batching train model. Finally, for its scalability to nigh-dimensional and complicated inverse problems. The precise evaluation of posterior sampling algorithms in high dimensions is very hard due the lack of meaningful quality metries.
work mainly theoretical but also provide evidence
9

---

Published as a conference paper ICLR 2024
ACKNOWLEDGMENTS
P.H. acknowledges funding by the German Research Foundation (DFG) within the project SPP 2298 "Theo- retical Foundations of Deep Learning". J.H. within the project STE 571/16-1 and by the EPSRC programme grant "The Mathematics of Deep Learning" with reference EP/V026259/L, F.A. by the DFG under Ger- many's ' Excellence Strategy - The Berlin Mathematics Research Center MATH+ (project AA5-6), R.B. and J.C. within the BMBF project 'VI-Screen' (13N15754) and G.S. acknowledges support by the BMBF Project "Sale" (011S20053B).
The material data from Section been acquired in the frame of the EU Horizon 2020 Marie Sklodowska- Curie Actions Innovative Training Network MUMMERING (MUltiscale, Multimodal and Multidimensional imaging for EngineeRING, Grant Number 765604) at the beamline TOMCAT of the SLS by A Saadaldin, Bernard, and Marone Welford. We acknowledge the Paul Scherrer Institut, Villigen, Switzerland for provision of synchrotron radiation beamtime at the TOMCAT beamline XO2DA of the SLS.
We would like to thank Chao Du for providing us class-conditional samples of his method for a quantitative comparison in Figure
REFERENCES
F. Altekrüger and J. Hertrich. WPPNets and WPPFlows: the power of Wasserstein patch priors for super- resolution. SIAM Journal on Imaging Sciences, 16(3):1033-1067, 2023.
F. Altekrüger, A. Denker, P. Hagemann, J. Hertrich, P. Maass, and G. Steidl. PatchNR: learning from very few images by patch normalizing flow regularization. Inverse Problems, 39(6), 2023a.
F. Altckrüger, P. Hagemann, and G. Steidl. Conditional generative models are provably robust: pointwise guarantees for Bayesian inverse problems. Transactions on Machine Learning Research, 2023b.
F. Altekrügger, J. Hertrich, and G. Steidl. Neural Wasserstein gradient flows for maximum mean discrepancies with Riesz kernels. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (eds.), Proceedings of the ICML '23, pp. 664-690. PMLR, 2023c.
D. Alvarez-Melis, Schiff, and Y. Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022.
L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows. Birkhäuser, Basel, 2005.
M. Arbel, A. Korba, A. Salim, and A. Gretton. Maximum mean discrepancy gradient flow. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, 2019.
L. Ardizzone, J. Kruse, C. Rother, and U. Köthe. Analyzing inverse problems with invertible neural net- works. In Proceedings of the ICLR 19, 2019.
L. Ardizzone, J. Kruse, C. Lüth, N. Bracher, C. Rother, and U. Köthe. Conditional invertible neural networks for diverse image-to-image translation. In Akata, A. Geiger, and T. Sattler (eds.), Puttern Recognition. Proceedings of the DAGM GCPR '20, pp. 373-387, Cham, 2021. Springer.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In D. Precup and Y. W. Teh (eds.), Proceedings of the ICML '17, pp. 214 214-221223332 PMLR, 2017.
R. Baptista, B. Hosseini, N. B. Kovachki, Y. Marzouk, and A. Sagiv. An approximation theory framework for transpoet sampling algorithms. arXiv preprint 1302./3963. 2023.
10

---

Published as a conference paper ICLR 2024
G. Batzolis, Stanczuk, C.-B. Schönlieb, and C. Etmann. Conditional image generation with score-based diffusion models. arXiv preprint 2/11 2021.
M. G. Bellemare, I. Danihelka, W. Dabney, S. Mohamed, B. Lakshninarayanan, S. Hoyer, and R. Munos. The cramer distance as a solution to biased wasserstein gradients. arXiv preprint arXiv: 10743, 2017.
N. Bonnotte. Unidimensional and Evolution Methods for Optimal Transportation. PhD thesis, Université Paris Sud-Paris XI, France, 2013.
J. A. Carrillo, M. Di Francesco, A. Esposito, S. Fagioli, and M. Schmidtchen. Measure solutions to a system of continuity equations driven by Newtonian nonlocal interactions. Discrete and Continuous Dynamical Systems, 40(2): 191-1231, 2020.
H. Chung, J. Kim, M. T. Mccann, M. L. Klasky, and J. C. Ye. Diffusion posterior sampling for general noisy inverse problems. In Proceedings of the ICLR '23, 2023.
T. M. Cover and J. A. Thomas. Elements of Information Theory. Wiley, Hoboken, 2nd edition, 2006.
H. Daneshmand and F. Bach. Polynomial-tine sparse measure recovery: from mean field theory to algorithm design. arXiv preprint 2204.07879, 2023.
H. Daneshmand, J. D. Lee, and C. Jin. Efficient displacement convex optimization with particle gradient descent. arXiv preprint 3302.0475,, 2023.
V. De Bortoli, J. Thornton, J. Heng, and A. Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. S. Liang, and J. Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17695-17709. Curran Associates, 2021.
A. Denker, M. Schmidt, J. Leuschner, and Maass. Conditional invertible neural networks for medical imaging. Journal of Imaging, (111):243, 2021.
C. Du, T. Li, T. Pang, S. Yan, and M. Lin. Nonparametric generative modeling with conditional sliced- Wasserstein flows. In Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and Scarlett (eds.), Proceedings of the ICML '23, pp. 3565-8584, PMLR, 2023.
J. Fan, Q. Zhang, A. Taghvaei, and Y. Chen. Variational Wasserstein gradient flow. In K. Chaudhuri, Jegelka, L. Song, C. Szepesvari, Niu, and S. Sabato (eds.), Proceedings of the ICML '22, pp. 6185- 6215. PMLR, 2022.
A. Garbuno-Inigo, T. Helin, F. Hoffmann,, and B. Hosseini. Bayesian posterior perturbation analysis with integral probability metrics. arXiv priprint 2303. 2023.
I. Goodfellow, Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates, 2014.
A. Gretton, K. Borgwardt, M. Rasch, B. Schölkopf, and A. Smola. A kemnel method for the two-sample- problem. In B. Schölkopf, J. Platt, and T. Hoffman (eds.), Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006.
A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012.
P. Hagemann, J. Hertrich, and G. Steidl. Stochastic normalizing flows for inverse problems: a Markov chains viewpoint. SIAM/ASA Journal on Uncertainty Quantification, 10(3):1 162 1190, 2022.
11

---

Published as a conference paper ICLR 2024
J. Hertrich. Fast kernel summation in high dimensions via slicing and Fourier transforms. arXiv preprint 2xiV:2401..88660, 2024.
J. Hertrich, R. Beinert, M. Gräf, and G. Steidl. Wasserstein gradient flows of the discrepancy with distance kernel on the line. In Calatroni, M. Donatelli, S. Morigi. M. Prato, and M. Santacesaria (eds.), Scale Space and Variational Methods in Computer Vision. Proceedings of the SSVM '23, pp. 431 4433, Cham, 2023a. Springer.
J. Hertrich, M. Gräf, R. Beinert, and G. Steidl. Wasserstein steepest descent flows of discrepancies with Riesz kernels. arXiv preprini 22/1.02800 1.01804, 2023b.
J. Hertrich, C. Wald, F. Altekrüger, and P. Hagemann. Generative sliced MMD flows with Riesz kernels. In The Twelfth International Conference on Learning Representations, 2024.
J. Ho, A. Jain, and Abbeel. Denoising diffusson probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, vol- ume 33, pp. 6840-6851. Curran Associates, 2020.
C.-W. Huang, J. Lim, and A. Courville. A variational perspective on diffusion-based generative models and score matching. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
R. Jordan, D. Kinderlehrer, and F. Otto. The variational formulation of the Fokker-Plaan equation. SIAM Journal on Mathematical Analysis, 29(1):1-17, 1998.
Y. Kim, K. Lee, and M. C. Paik. Conditional Wasserstein generator. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(6):7208-722 7219,, 2023.
D. P. Kingma and Ba. Adam: a method for stochastic optimization. In Proceedings of the ICLR '15, 2015.
D. P. Kingma and M. Welling. Auto-encoding variational Bayes. In Proceedings of the ICLR '14, 2014.
Kolouri, K. Nadjahi, U. Simsekli, R. Badeau, and G. Rohde. Generalized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019a.
S. Kolouri, P. E. Pope, C. E. Martin, and G. K. Rohde. Sliced Wasserstein auto-encoders. In Proceedings of the ICLR '19, 2019b.
H. Kremer, J.-J. Zhu, K. Muandet, and B. Schölkopf. Functional generalized empirical likelihood estimation for conditional moment restrictions. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato (eds.), Proceedings of the ICML. '22, volume 162, pp. 11665-1 1682. PMLR, 2022.
H. Kremer, Y. Nemmour, B. Schölkopf, and J.-J. Zhu. Estimation beyond data reweighting: kernel method of moments. arXiv preprint 2305. 10898, 2023.
A. Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University Toronto, Canada, 2009.
LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 (6(11):2278-233 2324, 1998.
J. Leuschner, M. Schmidt, D. O. Baguer, and P. Maass. LoDoPaB-OC BCB, benchmark dataset for low-dose computed tomography reconstruction. Scientific Data, 8(109), 2021.
Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. Proceedings of the ICCV '15 pp. 1738, 2015.
12

---

Published as a conference paper ICLR 2024
A. Liutkus, U. Simsekli, S. Majewski, A. Durmus, and F.-R. Stöter. generative modeling via optimal transport and diffusions. In K. Chaudhuri and R. Salakhutdinov (eds.), Proceedings of the ICML '19, pp. 4104 4113. PMLR, 2019.
flows: nonparametric
Lugmayr, M. Danelljan, L. Van Gool, and R. Timofte. SRFlow: Learning the super-resolution space with normalizing flow. In ECCV, 2020.
M. Mirza and S. Osindero. Conditional generative adversarial nets. arXiv preprint 1411.1784, 2014.
T. Modeste and C. Dombry. Characterization of translation invariant MMD on IR and connections with Wasserstein distances. HAL preprint 03855093v2, 2023.
P. Mokrov, A. Korotin, L.. Li, A. Genevay, J. M. Solomon, and E. Burnaev. Large-scale Wasserstein gradient flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15243-15256, 2021.
K. Nguyen and N. Ho. Revisiting sliced Wasserstein on images: from vectorization to convolution. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (eds.), Advances in Neural Information Processing Systems, 2022.
K. Nguyen and N. Ho. Energy-based sliced wasserstein distance. arXiv preprint arXiv:2304. 13586, 2023.
E. Novak and H. Wozniakowski. Tractability of Multivariatt Problems. Volume II: Standard Information for Functionals. EMS Press, Zürich, 2010.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L.. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. Vito,o M. Raison, Tejani. S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTTrrc: an imperative style, high-performance deep learning library. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances Neural 2019.
Information Processing Systems, volume 32, pp.
J. Rabin, G. Peyré, J. Delon, and M. Bernot. Wasserstein barycenter and its application to texture mixing. In A. M. Bruckstein, B. M Haar Romeny, A. M. Bronstein, and M. M. Bronstein (eds.), Scale Space and Variational Methods in Computer Vision, Proceedings of the SSVM '11, pp. 435-446, Berlin, 2012. Springer.
O. Ronneberger, P. Fischer, and T. Brox. U-Net: convolutional networks for biomedical image segmenta- tion. In N. Navab, J. Homegger, W. M. Wells, and A. F. Frangi (eds.). Medical Image Computing and Computer-Assisted Intervention. Proceedings of the MICCAI '15 pp. 234 241, Cham, 2015. Springer.
D. Sejdinovic, B. Sriperumbudur, A. Gretton, and K. Fukumizu. Equivalence of distance and RKHS- based statistics in hypothesis testing. Annals of Statistics, 2165):2233-2211 2013.
Song and S. Ermon. Improved techniques for training score-based generative models. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12438- 12448. Curran Associates, 2020.
Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In Proceedings of the ICLR '21, 2021.
B. Sprungk. On the local Lipschitz stability of Bayesian inverse problems. Inverse Problems, 36(5):055015, 2020.
B. K. Sriperumbudur, K. Fukumizu, and G. R. G. Lanckriet. Universality, characteristic kernels and RKHS embedding of measures. Journal of Machine Learning Research, 12(7), 201
13

---

Published as conference paper ICLR 2024
I. Steinwart and A. Christmann. Support Vector Machines. Springer, New York, 2008.
I. Steinwart and J. Fasciati-Ziegel. Strietly proper kernel scores and characteristic kernels on compact spaces. Applied and Computational Harmonic Analysis, 51:510-542, 2021.
G. Székely. E-statistics: The energy of statistical samples. Technical Report 02-16, Bowling Green State University, 2002.
G. J. Székely and M. L. Rizzo. A new test for multivariate normality. Journal of Multivariate Analysis, 93 (1):58-80, 2005.
G. J. Székely and M. L. Rizzo. Brownian distance covariance. Annals of Applied Statistics, 236) 2009.
G. J. Székely and M. L. Rizzo. Energy statistics: class of statistics based on distances. Journal of Statistical Planning and Inference, 14398):1249-1272, 2013.
C. Villani. Topics in Optimal Transportation. AMS, Providence, 2003.
Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli. Image quality assessment: From error visibility structural similarity. IEEE Transactions Image Processing, 13(4):600-61) 612, 2004.
M. Welling and Teh. Bayesian learning via stochastic gradient Langevin dynamics. In L. Getoor and T. Scheffer (eds.), Proceedings of the ICML '11, pp. 681-6881 1888,88, Madison, 2011. Omnipress.
Holger Wendland. Scattered data approximation, volume 17. Cambridge university press, 2004.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint 1708.07747, 2017.
14

---

Published as a conference paper ICLR 2024
A SUPPLEMENT TO SECTION 3
A.1 PROOF OF PROPOSITION
By the definition of the conditional probability as the disintegration of the joint measure, we have to prove that Py k Py,, with k(y.)) = T (((.))+ # Pz. To this end, let A E B(R" x Rd)) be Borel measurable.
Then holds
(Py x R)(() = f J [ d)
Since y and Z are independent by assumption, this is equal to
JEN / La(T(z,))) dPr.
-Sp [ La(x, y) (V, = Pr.x(A),
where the first equality follows by the transformation formula and the last equality follows by assumption. Since this holds for all A E B(RR" x R"), this completes the proof.
A.2 SUPPLEMENT TO SUBSECTION
The proof of Theorem requires some preliminaries on RKHS which can found more detailed, e.g. in Steinwart & Christmann 2008); Wendland (2004). The negative distance kernel K ((x;, v) || y is
conditionally positive definite function meaning that for every E and a; with Σi= a; = 0 and
I; E Rdd, 1,..., N, it holds
Σ
ajaz K (1i1 xz)) 2 0,
(12)
f,=1
where equality is only possible if a, = for all = 1, 1....,,,, N. The associated kernel
K (zr y) == -1|a y/l + |>||| + ||v|l
is positive definite, i.e., K fulfills 12 without the sum constraint the az For any 14, Pi(RR), we have
DR v) = Dx (p,
A Hilbert space of functions H(T)) mapping from R.d to IR with inner product (s, () associated norm and with the property that point evaluations f + f(x) are continuous for all f E H(R") is called reproducing kernel Hilbert space (RKHS). For symmetric, positive definite functions and in particular for K, there exists a unique RKHS such that the reproducing kernel property
f(x) =
for all f E H(R))
holds true. We denote this RKHS by H. (IR)). Further, there is an injective mapping from Pa (IRd) H R (Rd), called kernel mean embedding defined by
((x) = (K(x, ),m)) = K(x,0) du(y).
15

---

Published as conference paper ICLR 2024
(2023). Note that this has nothing do with the characteristic function of Since we do not address the later one in this paper, there is no notation mismatch. The kernel mean embedding not surjective, see Steinwart & Fasciati -iegel (2021). We have for H (R.) and j Pj (RR)) the representation
sec, c.g.. Modeste & Dombry
(h, 14) = S h(1)) du(a)) = [ (h, R( 1)96 du(x)
=(h,, J K(x,-) du(x)) Hx
= (h, Mk
(13)
This implies together with the dual representation of the MMD (Novak Wozniakowski 2010 for μt, P1 (IE") that
(R(UU))
sup (h, μ 1)) sup (h, D) = 111 01| || $1
(14)
[|ll $1
Equipped with the Wasserstein-p distance defined by
Wp(u, v) = { (infren(ae))) fRx xR4 || y y|l dr (xx w)) in((p))) S || y/l dx(x,, y)
if E [1, x)), if E (0,
(15)
where 11(, v/) denotes the set of measures with marginals and v. the space (P())) becomes a complete and separable metric space, see Modeste & Dombry (2023): Villani 2003). We will only need E ( 1,2). Note that by Jensen's 's inequality,
(H, S Wi(H, v) (2(,,))
Lee. convergence in Pp stronger for larger values p. Finally, we need the dual representation of the Wasserstein- J distances
Wy(n, v) sup (h IN c S1
(16)
where Ct (Id) denotes the space of 1 H-Höle continuous functions together with the seminorm l//c2 supzzy (f(x) - f(y))/ll= yll, see Modeste & Dombry 22023); Villani (2003). For compactly supported measures, relation between the MMD with the negative distance kernel and the Wasserstein distance was proven in Hertrich et al. (2024).
Lemma 5. Let K(x -1|x yll. For U, Pi
(1R)) it holds
( S Wi(
n and are additionally supported the ball (R( then there exists a constant C. such that
Wi(uv)) < CaR 2477
28+1
Dx DK (4,
Now we can prove Theorem
Proof Theorem
For any E RR, we consider the difference of the kernel mean embedding
functions f(12)) = PxIY= - Px By 14 and 13). we obtain
D (P (P Pxjr= = (f(..), P x P = (f(:.),, Pxiy=y - PxIr=g)
16

---

Published as a conference paper ICLR 2024
and further
Ey-Pr (P PxIx=
(f(:.); PxIY=y (x)x=) Px(m=e) dPy(y)
-Je J f(2, ( px) (v) dæ dy.
Applying the definition of the posterior density Px|Y=y = PX.. this can be rewritten as py
[ J fla,v)(xxy (x,)) dr dy.
Ey
(17)
2. Next, show that the function f(x, y) is 1 -Hööder continuous with respect to both arguments. First, we conclude by assumption on the H-Hölder continuity of the posteriors with respect to y that
(x,,t) dl d (z) (x)- / K(x,t) dP (x)
f(t f(t,
gd
JRd
k(x,t) dP d (z) + S ((x,t) d (x)
Ja
(2x,))(P1)== -PX =xx dx
gd
+ [ ((x,t) PxIY (x) dr
2Cs / ((x, 1) drllyi vall 22Cs, Csyllvi - 32ll3
where Cs, maxtes, Ss. kxx 1) dr. Interchanging the role of yn and y2 yields
[f(t, 91) (t. 22)l S Cslly: v2ll=, Cs 20s, Os.
for t E Sa and all yı, V2 E Sa- Now the triangle inequality and the relation ((x))>(y)]== = khh, R(t, (v,,))) H) S | [x(x))) - that
for all h E Hg. implies for 22 E Sa and yı, 322 E S,
[f(11, - (x22 < [f( v1) f(m 32)1 + (s(1,, 92) f((r2.
S Csally - =2ll* $ + [I 22)l (| - (Z22
By the reproducing kernel property and definition of A we have
[ Rx2, :1 (I ((22, 2k 221,, 222) 21|201 xall,
that
1f(2),8)) f(z2, v2)1 S Cs, - 22ll (2)) x2
< 2
max (Os., (211((;)222l22) [((1181) (2 12
Finally, we obtain
[l (f(-12)ll DD (Px (Px|X== P
sup
h d(Px Y=
Px(r=
|| &
31 Rd
17

---

Published as conference paper ICLR 2024
(LLa
sup g St
dP x
|+| + ddxxy===)) 22
sup ||/e|oo
2Csa
where |/A|l00 ma [xxx) and the last inequality follows by [h(z)] 10h, K(, m)ugl S | (xxx)))1z = 22llll |a|| I Csul/blls for all E Sa. In summary, we proved for all
3i, 42 E Sa ad yı, 1/2 E S. that
1f(11.81) f(22, v2)l & Ò' [[ 91) - (02, v2)1
where Č == max (C..: 2v2Cs.) Hence, o 2 1
-öölder continuous with lit 1.
24
3. Using Jensen's s inequality well (16 and (17 we conclude
By
[xx(Px Px S Eyyoy
PxIV==)))
C
< (ex,
Finally, apply Lemma 5 to get
Ey [Dx Pex(Pxxve (P Pxir=y)] DDD(PPxy Px, ाaकnei
with appropriate constant C' This finishes the proof.
The assumption that the random vectors X, Xx and Y map onto compact sets cannot be neglected as the following counterexample shows.
Example 6 (Theorem 2 fails for non-compactly supported measures). For fixed m E N and < < 6-1, we consider the random variables X, X, and on R with joint distributions Px.y = 1 M (0.0) + (m..1)))) and Px, = 2/2 (m.0) + 22220000 (0.1), where U, denotes the uniform distribution on the indicated set and Qe(I, u) the e-ball around (z, y) with respect to the co-norm. The assumptions on the conditional densities in Theorem 2 are fulfilled with Cs, (1 2e)-1/2, Using Lemma we obtain Dr (Px,y) Px,)) s 2 for all N. Furthermore, have
Dx(PPxy
Py [Dx
2 Dx (q) Dx (m). vm be.
2
Consequently, there cannot exist constant C such that 7 holds true for random vectors with arbitrary non compact range.
Estimates similar 7 can also be established for other divergences.
Remark 7 (Relation between joint and conditioned distributions for different distances). i) For the Kullback- Leibler (KL) divergence, the chain rule (Cover & Thomas 2006 Thm 2.5.3) implies
Ey [KL(Px Px|y=y)] KL Px,y).
ii) Using Lemma and Jensen's inequality, the estimate 7 can be transferred to the Wasserstein-1 distance. More precisely, under the assumptions in Theorem 2 we have
Sy>PyWW((P) (X Px]Y= S C' Cy [Dx Px|r=
CC (x ( Px w PainTLRaTE & CWW((xxy Px,y) BaiiniiaaT
with constants C, C"
18

---

Published as conference paper ICLR 2024
Next, we show that the relations (8) and ( hold true under certain assumptions. We need an auxiliary lemma which was proven for more general kernels by |Baptista et al. (2023 Thm 3.2). Lemma 8. Consider K(r,, y) =-1|= yll. and let Sa C Rdd be compact set, and Z Sg be random
variable. For F. G.E. E 1/ (Sa, RR)) it holds
DK (F Pz P2) V2 -G)))=
G(z))I|+.
Now we can prove the Hölder estimates.
Lemma (Stability under Pushforward), Consider K(r,y) == || yIl and Sn C R" and Sa C B be compuct sets, and Z E Sa be random variable. Further, let T B& x Sn + Re be measurable. If the derivatives are uniform bounded by supyes, [T,yTzz,)) S
for all z E Sa then
Dx (T (+ Pz. T (2,12,) + Pz) s v2C [|l1 - -2ll for all y1, Sn.
Proof. For all 92 E Sn, Lemma 8 yields
DR (+.:)) Pz, T (. (2) Pz) S 2 Eg [1] (11 T(z, 32)l
Further, the second fundamental theorem of calculus implies
|T (z. v1) T (z, (2)32)l = / V (z, 21 + t (32 - 11)) (en y2) de
< ! v + (12 yn )1l di Il s2ll S C
[lyi
y/
for all z E Sy. Applying this estimate in the above inequality yields the assertion.
For the stability with respect the posteriors, there exist sophisticated strategies obtain optimal bounds, see Garbuno-Inigo et al. (222). However, for our setting, we can just apply result from Altekrüger et al. (2023b) Sprungk (2020) on the Wasserstein-1 I distance.
Lemma 10 (Stability of Posteriors). Consider K(x = -1I yll- Let Sn C R and Ss C Rdd be compact sets, and X Sa and Y Sn be random variables. Assume that there exists a constant M such that all E S, and for all x E SS it holds
llog Py|X=z (y1) - logpv (92) S M I1l1 - 32ll
Then, there exists a constant C >> such that for all y1, 322 Sn we have
D (PxIY=01 PxIV= S Clly1 - yell
Proof. By Lemma know that D S (W, v)) Now can apply (Altekrüger al. 2023b] Lem. 3) together with the fact that local Lipschitz continuity on compact sets implies just Lipschitz continuity.
Proof of Theorem 3 For 8 == Ey (T Pz, Px]Y= Theorem 2 implies å < C'e (lata))))) Adapting the lines the proof of (Altekrüger et al. 2023b Thm. 5) with respect to MMD instead of the Wasserstein-1 distance, we obtain
(x(T" (.v) Pz, Px]y= S D D & D aaa Do
for arbitrary y Snn, where the constants D.D > 0 depend on the dimension n, the value (y) the
diameter of Sm, the constants from (5). 6 and 9 as well as on C" from the assumptions. Finally, taking - yields the assertion.
19

---

Published as a conference paper ICLR 2024
A.3 SUPPLEMENT TO SUBSECTION 3.2
A curve y: + (2R)) on the interval / S IR is called absolutely continuous if there exists a Borel velocity field Uy: IR + IR with Jr || E dt < oo such that the continuity equation
(n(() V (w))(t)
(18)
is fulfilled Rd in weak sense, see (Amrrosoo al., 2005
Thm. 8.3.1). A locally absolutely
continuous curve (0, oc) + (22( with velocity field 2/ 7(0) Pe P2(1e (RR) is a Wasserstein gradient flow with respect to functional F: P2 (2( - (-o0, 00)
2t E -FF(((t) for a.a. a.r..e.
t>
where 22 (RR denotes the regular tangent space at μ E P2(R"). see (Ambrosio et al. 2005 Def. 8.4.1), and (F) the reduced Fréchet subdiffential at μ. E P2 (R( consisting of all (2 (R.. : satisfying
F(v) F(p)
inf TEDop"(p,)))
(E(x), - 2) dr(x, v) o(2((,,)) for all E P2(R
Rd xRd
see (Ambrosio et al. 2005 Eq. (10.3.13). Here the infimum is taken over the set popt (n, v) consisting of all optimal Wasserstein-2 transport plans between μ and i.e., the minimizer of 15 for p = 2.
order to prove Theorem 4 we first prove that an absolutely continuous curve 7 within the metric space (P2(R4 x ] Wa) does not transport mass within the second component, whenever the second marginal is a constant empirical measure. Note that the ith marginal can be written as (P:) using the projections (I x and P2(x, v) y.
Theorem 11. ye I + P2(Rd R") be an absolutely continuous curve with associate vector field i (1,2); Rd x Ih Rdd x RD. If (P2) a constant empirical measure 1 EN 5g. independent of then 012 vanishes 7(t)) for almost every E I.
N
Proof. Denote by oith E P2((R RR") x (Rd x )))) an arbitrary optimal Wasserstein-22 plan between r(1) and 7tt + h), and by P:: (R2 R=y² (IR x IR") the projections to the first and second component,
ie. (2((11 y1), (x2, 12)) (z,,8) for E {1.2). For almost every e the associate vector field v, of
r satisfies
lim (1,, h (P2 P म.
(Id, v1) vn().,
(19)
2005 Prop 8.4.6). For n E N, let where the left-hand side converges narrowly, see (Ambrosio et al. : [0, o0) - [0. 1] be a continuous function with (0) 0. (t)
for (0,m) and fa 0 for
t 2 n. Consider the integral
[ (aaa (n d(Pi, + (P2 Pi)) 4hh t ((x1, (12,
F. =
| (Red x fa fa(H vill) dml+h (aetthh(,,,00 ((11, v1), (x2,
(20)
Since (Pa) = N Σm 8q. is constant empirical measure, every plan m. 4+h supported on
U x (a:)) x (R" (ai)).
i-f=
20

---

Published as conference paper ICLR 2024
Hence, on the support of mt" ++h the norm ||y2 - vill (20 becomes either zero is bounded by 1//e v1 2 min ( bll : + j} = S.
Thus, for & S/n, the integral Fn vanishes and the narrow convergence in 19 implies
(1,2(11311) (-n, 0) U (0,
for -tee (z) E I2 R" Since is arbitrary, we obtain the assertion.
Interestingly, Theorem 11 is in general not true if the second marginal is not an empirical measure as Exam- plc 12 below shows. We can now adapt (Altekrüger et al. 2023c Prop. D.1) for proving Theorem
Proof of Theorem Let ş == (5i, .. Ew) E RáN satisfy (Si, gi) * (Ex 4j for all + j. Then, there exists an E > such that the optimal transport plan between & Σ )(s.)))) and N. Ex= ((.g)) given
by Ei= (( for all 1 E RdN with |l5 7l S c. In particular, it follows
N Σ
W
8 (EL N
1 Σ (in
1 N
(I,,, 9:) (mi, 9:) a) Σ
I&: 1 (21)
By analogous argumentation as in Alteküüger et al. (2023c), YN,g is locally absolutely continuous curve; and by Theorem I1 the second component of the associate velocity field 1 = =(v,,,, 02,2) vanishes, i.e. 02 0 for almost every & (0, o0). Exploiting (Ambrosio et al. 2005 Prop. 8.4.6), we obtain
W2(7NN (t h)), (Id + (e;, (NN.
lim
h-0
1112
W L Eis s(ul((h)))))) N
Em (eo.
lim h-40
|112
Σ h-
lim
1
N
(t + h)
m4(1) ())))))) ve (t)
1 Σ (d:(()) - (t,) ())l
N
for a.e. (0, o0), where the first equality in the last line follows from 21 In particular, this implies (t) = (t) qı) a.. such that NV, (()o)(t))t = (11,1(u
(ux (ux(()))))))))))(( (N)).
For fixed 1, we now consider an e-ball around (w. where the Wasserstein-2 optimal transport between (N..(() and measure from μ PN.g becomes unique discussed in the beginning of the proof. More precisely, the unique plan between (N, and μ== 인 Σ ((n-9.)) with W7(H, (N)(()) S e is then given by π N Em ((())(()..w)).)n)..))) Since is solution of (TO), obtain
0 < P(p 4)) (( 4)) (P ((((((()))) g)), u(I) + -(I))
N
= (1) - (1N. +
(t, 9:), (ne q:) - (u.(t), 9)) ((
N
= -J),, Jun. (1N.gg))))))))))
( (v (1)) (22, v2) (x1, vn)) dr ((21, v1), (12.
+ +oWW2(H (t)).
Since क is the unique plan in (7Ng(),,)))) 14), and since Juw. (1) = oo for & PN.g- we have i E -TTxx (T(()) showing the assertion.
Finally, we construct an explicit example showing that the restriction to empirical second marginals in The- orem is inevitable.
21

---

Published as conference paper ICLR 2024
Example 12 (Theorem 11 fails for arbitrary marginals). Let
[0, 1] be a continuously differentiable
function on the interval C IR, and consider the curve T/: + (2([R) given by
given
(r(() (1 f(t) 60 - f(t) 11
Figuratively, f controls how the mass on the interval [-1, 0] flows into or out of the point measure located zero. In order show that such curves are absolutely continuous, we exploit the associate quantile functions. More generally, for E P2((R), the quantile function defined as
(s)s = min( E R p((-oc, 2 s}.
s S (0, 1).
For specific curve, the quantile functions piecewise linear and given by
(e) (s) min min ((((t))* 1,0),
see for instance [Hertrich al. (2023a Prop 1). Due the relation between the quantile function and the Wasserstein distance (Villani 2003 Thm 2.18), we obtain
(itr [. (Qr)())(s) - ())(t)(())
f (((0(())) - (-((tt))) ds -(((()))'
If the derivative of f/f // is bounded by M I, the mean value theorem yields
M Wa(rs(t). vs
f t2l.
such that Ty is indeed absolutely continuous. Based on the curves 7 on the line, we now consider the curve 7: (1/4, 3/4] + P2( (TR: IK) given by
1 (() 2 (1 - (0.00 {(00)x(11,0) + 662.0) - )(a)x(-.a))
where Aa denotes the uniform measure on the set A C IR Restricting the transport between 7(11) and ((t2) along the lines segments {0} x [-1, and {2] [-,,0], and using the above considerations, where the derivatives bounded by .M = 4, we obtain
(2(7((1), ((2)) s 3 14 dz/. 4,42 [1/4,3/4].
Thus 7 is absolutely continuous too. Furthermore, both marginals (1)) = (ō0 +2) and T()) 2 (60 + )(-1.0)) are independent of f. If Theorem associate vector field R² + _ has to be the zero everywhere implying that is constant. Since this is not the case, we would obtain a contradiction.
[1 would hold true for non-empirical marginals, the
Theoretical justification of numerical observation in Du et al. 2023 By Ambrosio al. (2005 Thm. 8.3.1), absolutely continuous curves y: 1 + (2("")) ) in P2 (R") correspond to weak solutions of the continuity equation (18). For the sliced Wasserstein gradient flows with target measure v, analytic representation of U/ was derived Bonnotte (2013) as
(x( Egege- (Pc())))) Pε(x) = ().x),
where thx.g the Kantorovic potential between P (t) and Pegu In the case that 7(0)
EiE 8z,
and 3 E M Liml 6p.. this implies that 7(t) 1/4 ΣN Sur(e), where u (ui m.y uN) solution of
i(l) v.((()) ((,())))GG,(((()) Gp(z) sw
8 M
8ps
22

---

Published as a conference paper ICLR 2024
In the context of posterior sampling, Du al. (202)) considered sliced Wasserstein gradient flows starting 1/ 2j=1 ((p..))) such that VMra Px,y and 7(0) 2 Pz Py for continuous random variables X and Y and latent variable Z. Then, they observed numerically that the second part 11 of (t,,), IRd R - R x I is "almost zero". In order to
at 7(0) - & Ei= )(x,,.)) with target measure VM, M
apply Proposition I they set the component 12 in their simulations artificially to zero, which corresponds solving the ODE
(22) Du et al. (023)) write by themselves that they are unable to provide a rigorous theoretical justification of the functionality of their algorithm. Using analogous proof as for Theorem 4 we can now provide this justification as summarized in the following corollary. particular, the simulations from Du et al.
i(t) - V.
(2023) still Wasserstein gradient flows, but with respect different functional which has the minimizer UM5. R Px.Y-
Corollary 13. Let u = (ui..... uN):
o0) + ( R be solution of (22) and assume that (t, q.) + (1)(t), (,) for + j and all > Then, the curve YN.q: : (0, 00) - P2 (R. defined by
N
(N,4(t) N 6m.()).
is a Wasserstein gradient flow with respect to the functional
(SW((,
if PNá otherwise.
00,
ALGORITHM SUMMARY
Here summarize the training of our conditional MMD Flows, Algorithm
C IMPLEMENTATION DETAILS
The code is written in PyTorch Paszke al. 2019) and is available online"
We use UNets { which are trained using Adam (Kingma & Ba, 2015) with a learning rate of 0.0005. Since the differences between the particles (*++)) and (k) are large when starting simulating (II) and smaller for larger step k. increase the number of simulation steps Ti up to predefined maximal number of iteration steps Tmaz = 30000.
For experiments, make of several improvements, which are explained in the following.
Sliced MMD equals MMD Instead of computing the derivative of the MMD functional in directly, we use the sliced version of MMD shown in Hertrich al. (2024). More precisely, let = (z (#1..... ... ZN)
(Rd) == (PI...., PM) E (Rd)M and let (z) == (x(r) be the discrete MMD functional, where we explicitly note the dependence the dimension d. Then we rewrite the gradient of the MMD V., Fx())))) with the negative distance kernel as
V F.(2) (x) = 1. ((5, (5. x. ()))
http://hwha..ooo.oaaaaaaaaaaoaaa t ///modess/unet.pp
modified from
23

---

Published as a conference paper ICLR 2024
Algorithm I Training of conditional MMD flows
Input: Joint samples -1, PN, M, .. QN from Px,y, initial samples ut. ... "g momentum parameters my E [0,] 1) for | 1,....
Initialize (U1. ...,.... uN) = O.
for 1.... L do
Set [ū] ..... ū N
(0) = (u (1-1) ..... N
(111)
Simulate Ti steps of the (momentum) MMD flow:
for / = 1 1,.., Ti do
Update v by
(U,,-.., UN) + " ())
(8)
+ mi(oi.... uN)
Update the flow samples:
x(++1)_ = (((k) (U
end for
Train 2 such that (T.)) se u!
(ū: q:) by minimizing the loss
N
[(6))
(:) - (4) - u 112 ( (0)
N
tm1
Set (uf) ...., uN (1)
(ui (1-1)
(1-1)) ( 11-1) Q1) ... Di(uN (1-1) .qN)).
..... uN
end for
where P = ((EP))..... , (E, PM)))
and ca is a constant given by
VTI( d+1)
Cd
r(2)
Thus it suffices to compute the gradient w.r.t Pig" et al. (2024) Section 3).
which can be done in a very efficient manner, see Hertrich
Pyramidal schedules In order to obtain fast convergence of the particle flow (T] even in high dimensions, we make use of a pyramidal schedule. The key idea to simulate the particle flow on different resolutions of the image, from low to high sequentially. Given the target images p, E Rd, where d = C H W with height H. width W and C channels for i I.....N, we downsample the image by factor S. Then we
start simulating the ODE 111) with initial particles x (0) E R $2 i.e., we simulate the ODE in a substantially smaller dimension. After predefined number of steps I, we upsample the current iterate (1()) to the higher resolution and add noise onto in order increase the intrinsic dimension of the images. Then we repeat the procedure until the highest resolution is attained.
Locaely-comnected projections Motivated by Du et al. (023); Nguyen & Ho (2022), the uniformly sam- pled projections & 3d-1 can be interpreted Instead, for image tasks the use of locally-connected projections greatly improves the efficiency of the cor- responding particle flow. More concretely, for given patch size & we sample local projections & uniformly from -1 Then, we randomly choose pair (h, w) and extract a patch E(,,,n) (p) of our given image
fully-conneee layer applied to the vectorized image.
p Rdd where (h, w) is the upper left corner of the patch and E(hh,w): Rd R is the correspond-
24

---

Published as a conference paper ICLR 2024
ing patch extractor. Using this procedure, we simulate the ODE (TT) for (i (2,) and target Eihh (pi).
i 1,.., N, where the location of the patch is randomly chosen in each iteration.
order to apply the locally-connectee different scales, similar to Du et (2023) and, depending on the condition q:, locally-conneeted projections are also used here. Note that here we introduced an inductive bias, since we do not sample uniformly from Sd-1 anymore, but it empirically improves the performance of the proposed scheme. A more detailed discussion can be found in (Du et al. 2023; Nguyen & Ho 2022).
projections on different resolutions, we upsample the projections to
C.I CLASS-CONDITIONAL IMAGE GENERATION
where the conditions qi are the one-hot vectors of the class labels. We use P = 500 projections for each scale of the locally- connected projections, for the observation part we use fully-connected projections. For MNIST, we the patch size
For MNIST and FaSH use N 20000 target pairs (pt. i =,
5 and apply the projections on resolutions 5, 10, 15, 20 and 25. For MNSSi, the patch size s is used on resolutions and 27. In both cases, the networks are trained for 5000 iterations with a batch size of 100.
For CIFARIO, we N 40000 target pairs and make use of the pyramidal schedule, where we first downsample by a factor resolution 4 4 and then upsample the iterates x²(() by factor of 2 after every 700000 iterations. In the first two resolutions we use P = 500 projections and in the third resolution we use P 778 projections. On the highest resolution of 32 32, make use of locally-connected projections, where we choose the patch size 7 on resolutions 7, 14, 21 and 28. Here we use P = 400 projections and train the networks for 4000 iterations with batch size of 100.
C.2 INPAINTING
For MNIST and FNSHT we use N = 20000 target pairs i where the conditions % are the observed parts of the image. We use again P = 500 projections for each scale the locally- connected projections, for the observation part we use fully-connected projections. For MNIST, we use the patch size 8 and apply the projections on resolutions 5, 10, 15 and 20. For FashionMNIST, the patch size = 7 is used on resolutions 7, 14 and 21. In both cases, the networks with a batch size of 100. For CIFAR1O, we use = 30000 target pairs and make use of the same pyramidal schedule in the class-conditional part, but increase the resolution after every 600000 iterations.
=
trained for 4000 iterations
C.3 SUPERRESOLUTION
We use N = 20000 target pairs of CelebA, where the low-resolution images are bicubicely downsampled resolution 16 16. Similarly the pyramidal approach for CIFARIO, we downsample the particles by a factor of 8 and increase the resolution after every 600000 iterations by a factor of 2. While for the first
2 resolutions we use fully-connected projections with P 500 and p 768 projections, for resolutions 32 x 32 and 64 x 64 we make use of lecally-connected projections with P = 500 and s = 7 on resolutions 7 and for 32 x and 7, 21 and for 64 x 64. For the observations we use fully-connected projections for all resolutions. The networks are trained for 5000 iterations and batch size of 10.
C.4 COMPUTED TOMOGRAPHY
We use the first N = 400 target pairs of the LoDoPaB training set of size 362 x 362, where the observations
are the FBP reconstructions of the observed sinograms. We use P = 500 projections for each scale of the locally-connecte projections with patch size s = 15 and use all resolutions between 15 and 135 as well resolution 270.
25

---

Published as conference paper ICLR 2024
Table FID scores of the class-conditional samples for MNIST, FashionMNIST and CIFARIO. Separated for each class.
MNIST
FashionMNIST
CIFARIO
FID
ESWE
Cond. MMD Flow
FID
A-SFF
Cond. MMD Flow
FID
E-WWF
Cond. MMD Flow
(ours)
et
(ours)
Class 0
10.3
13.2
Class
27.9
30.1
Class 0
127.0
98.8
Class
16.0
23.3
Class
16.0
21.9
Class
126.9
100.9
Class
26.2
117.9
Class 3
13.8
12.8
Class
28.3
27.3
Class
141.0
99.4
Class 4
13.7
Class
25.5
Class
103.9
Class
34,0
13.7
Class
27.7
31.1
Class
126.2
112.2
Class 6
18.9
13.2
Class
28.5
29.1
Class
109.5
87.1
Class
14.6
14.3
Class
21.3
Class
112.5
109.3
Class
19.4
13.0
Class
37.2
Class
10066
94.0
Class 9
17.5
12.0
Class
20.7
20.0
Class 9
107.9
99.3
Averago
189
14.6
Average
25.9
27.5
Average
118.1
101.8
WPPFlow (mean img)
SRFlow (mean img)
Sliced MMD Flow (mean img)
PSNR
25.88 (26.92)
25.11 (27.36)
27.21 (27.93)
SSIM Wang al. 2004
0.657 (0.766)
0.637 (0.771)
0.774 (0.790)
Table 2: Comparison of superresolution results of different reconstruction schemes. The values are meaned of 100 reconstructions for each observation. The value of the resulting mean image in brackets. The best value is marked in bold.
D FURTHER NUMERICAL EXAMPLES
Here we provide additional generated samples for the experiments from Section Moreover, the FID scores for the class-conditonal image generation given Table each class separately. Note that the arising values are not comparable with unconditional FID values. Obviously, we outperform the L-SWF of et al. (203), which is the conceptually closest method.
E SUPERRESOLUTION ON MATERIAL DATA
addition to the superresolution example in Section we provide an example for superresolution on real- world material data. The forward operator f consists of blur operator with 16 x 16 Gaussian blur kernel with standard deviation and a subsampling with stride 4. The low-resolution images used for reconstruction generated by artificially downsampling and adding Gaussian noise with standard deviation 0.01. We train the conditional MMD flow with 1000 pairs of high- and low-resolution images of size 100 x 100 and 25 x 25, respectively.
In Figure II we consider an unknown ground truth of size 600 x 600 and an observation of size 150 150. We compare the conditional MMD flow with WPPFlow (Altekrüger & Hertrich (023) and SRFlow (Lugmayr ct al. 2020). Note that while the conditional MMD flow and the SRFlow are trained under the same setting, the WPPFlow just requires the 1000 observations and the exact knowledge noise model for training. We can observe that the conditional MMD flow is able to generate sharper and much more realistic reconstructions for the given low-resolution observation. More examples are given in Figure 12 A quantitative comparison of the methods is given in Table Here we consider 100 unseen test observations reconstruct 100 samples from the posterior for each observation.
the forward operator and the
26

---

Published as conference paper ICLR 2024
0
0
O
0
0
0
0
O
O
O
0
6
O
0
0
1
/
/
/
1
/
/
1
/
/
2
2
2
2
2
2
2
2
2
2
2
R
2
2
2
2
2
2
3
2
2
Z
2
2
7
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
/
4
/
4
4
4
4
4
7
4
4
4
4
4
4
4
4
4
4
5
5
5
5
5
5
5
5
5
5
5
5
5
S
5
5
5
5
5
6
5
5
5
5
5
5
6
6
6
6
6
6
6
6
6
O
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
2
?
7
7
7
7
7
9
7
7
7
7
7
7
7
7
7
7 8
7
8
8
8
8
8
8
8
8
8
6
8
8
8
8
8
8
8
8
7
8
8
8
8
9
9
7
9
8
6
9
9
9
9
9
9
8
8
8
9
9
7
9
(a) MNIST
(b) FaSH
(c) CIFARI0
Figure 6: Additional class-conditional samples of MNIST, FashionMNIST and CIFAR10.
27

---

Published as a conference paper ICLR 2024
7
9
7
5
7
4
5
7
5
9
9
4
6
7
8
4
7
9
7
9
7
3
5
5
4
7
5
9
7
7
9
7
9
9
7
9
9
9
7
7
7
5
8
7
7
7
5
7
T
7
5
3
7
7
7
7
9
/
/
/
/
/
/
i
/
4
4
/
/
J
/
/
/
/
/
/
4
4
4
2
2
2
2
2
2
2
2
Z
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
&
2
2
2
2
2
/
/
/
/
/
/
/
/
/
4
1
[
4
4
4
1
4
Y
8
1
8
8
8
8
8
8
8
7
7
8
8
8
8
8
8
8
8
8
8
5
8
8
8
8 9
8
8
9
9
7
7
7
9
9
7
9
7
7
9
6
6
6
6
6
6
6
6
6
6
6
6 6
6
6
6
6
6
6
6
6
6
5
O
(a) MNIST
(b) MNIST
(c) CIFARIO
Figure 7: Additional inpainted samples of MNIST, FashionMNIST and CIFAR10.
28

---

Published as a conference paper ICLR 2024
(a) CelebA superresolution
Figure 8: Additional superresoluted images CelebA
29

---

Published as conference paper ICLR 2024
GT
FBP
Reco 1
Reco 2
Reco
Mean
Std
GT
FBP
Reco
Reco
Reco
Mcan
Std
Figure 9: Additional generated posterior samples, mean image and pixel-wise standard deviation for low- dose computed tomography using conditional MMD flows.
30

---

Published as conference paper ICLR 2024
GT
FBP
Reco 1
Reco 2
Reco
Mean
Std
GT
FBP
Reco
Reco
Reco
Mean
Std
Figure 10: Additional generated posterior samples, mean image and pixel-wise standard deviation for limited angle computed tomography using conditional MMD flows.
31

---

Published as conference paper ICLR 2024
HR image
LR image
Prediction
Prediction
Prediction
Mean image
Std
HR image
LR image
Prediction
Prediction 2
Prediction
Mean image
Std
HR image
LR image
Prediction
Prediction
Prediction
Mean image
Std
Figure 11: Different WPPFlow (top), SRFlow (middle) and conditional MMD flow (bottom) reconstructions of the HR image and their mean and pixelwise standard deviation (normalized). The zoomed-in part marked with a black box.
32

---

Published as conference paper ICLR 2024
HR image
image
Prediction
Prediction
Prediction
Mcan image
Std
HR image
LR image
Prediction 1
Prediction
Prediction
Mean image
Std
Figure 12: Additional conditional MMD flow reconstructions of the HR image and their mean and pixelwise standard deviation (normalized). The zoomed-in part is marked with black box.
33


Reviews:
[{"review_id": "0cT1m6M8rx", "reviewer": "Reviewer_PLVu", "summary": "This paper proposes a conditional flow of the MMD with the negative distance kernel, which can be further implemented by conditional generative neural networks with application in image generation, inpainting, and super-resolution. The authors derive the convergence of the posterior under some certain stability conditions, and relate it to a Wasserstain gradient flow. Those results extend previous investigation for sliced Wasserstein flow. The work is relatively theoretical and lacks a thorough comparison with other generative models.", "questions": "1. In Equation 4, $T$ is defined, however $T_\\sharp$ is not defined.\n\n\n2. Is it possible to validate the error bound via numerical experiments somehow?\n\n\n3. Could the author comment on the difference between the proposed analysis and sliced Wasserstein flow, as the implementation is still based on the sliced version of it?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper presents some interesting theories, and extends the analysis on sliced Wasserstein gradient flow.", "weaknesses": "1. It would be better to elaborate on the pros and cons of using a negative distance kernel (efficiency, sample complexity, etc).\n\n2. The contribution is not entirely clear. Could the author comment on the effectiveness/efficiency/novelty/difficulty of the proposed method?\n\n3. A highlight of the proof techniques used by the authors to address gradient flows with respect to MMD with negative distance kernel without mean-field approximation would help to improve the importance of this work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jtEYG2YqZS", "reviewer": "Reviewer_eQGc", "summary": "In this paper, conditional MMD flow with negative distance kernel is introduced.\nThe model's stability is proven by bounding the expected approximation error of the posterior distribution.\n\nThrough theoretical justification, the authors obtain convincing results by neglecting the velocity in the y-component in sliced Wasserstein gradient flows.\nThen, the power of the method is also demonstrated by numerical examples including conditional image generation and inverse problems.", "questions": "1. Except computed tomography experiment, only visulization results of other experiments are given in the paper, however, it is difficult to quantitatively evaluate the result and to compare with other method. Hence, evaluation metrics need to be introduced or self-defined.\n\n2. The related work: Neural Wasserstein gradient flows for maximum mean discrepancies with Riesz kernels, proposed similar method, what is the strength and advantage over it? and what about the performance difference?\n\n3. Why chosing UNet? Is there a significant difference in the effect of choosing other models such ResNet and transformer.\n\n4. As Fig.7c shows, inpainting results of CIFAR are not good enough, the generated images differ from each other greatly at the unobserved part, what is the reason? and are there any solutions to improve it.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The theoretical justification of the proposed method is clear and detailed.\n2. Several experiments are conducted to prove the power of the method.\n3. Introducing negative distance kernel to MMD is a good idea and contributions are well-described.", "weaknesses": "As mentioned by the authors, the proposed approach has some limitations:\n\n1. The model is sensitive to forward operator and noise type.\n2. Lack of meaningful quality metrics to evaluate the results.\n3. Realism of the computed tomography experiment results can not be guaranteed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JdySIeobRZ", "reviewer": "Reviewer_mNmx", "summary": "This paper proposes conditional MMD flows with the negative distance kernel for posterior sampling and conditional generative modelling. By controlling the MMD of the conditional distribution using the MMD of the joint distribution, the paper provides a pointwise convergence result. In addition, the paper shows that the proposed particle flow is a Wasserstein gradient flow of a modified MMD functional, and hence provides some theoretical guarantee for [1]. Finally, the paper experiments on several image generation problems and compares with other conditional flow methods.\n\n[1] C. Du, T. Li, T. Pang, S. Yan, and M. Lin. Nonparametric generative modeling with conditional slicedWasserstein flows. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (eds.), Proceedings of the ICML ’23, pp. 8565–8584. PMLR, 2023.", "questions": "1. The paper states that MMD combining with the negative distance kernel results in many additional desirable properties, however it lacks convergence rate or discretization error analysis because “the general analysis of these flows is theoretically challenging”. Regarding this problem, what is the advantage of MMD over other discrepancies like Kullback–Leibler divergence or the Wasserstein distance especially for conditional generative modelling problems?\n2. Is it possible to provide a discretization error analysis between discrete MMD flow and the original continuous MMD flow?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and clearly-organized.\n2. The paper proves that the proposed particle flow is a Wasserstein gradient flow of an appropriate functional, thus providing a theoretical justification for the empirical method presented by [1].\n3. Abundant generated image samples are shown in the experiments.\n\n[1] C. Du, T. Li, T. Pang, S. Yan, and M. Lin. Nonparametric generative modeling with conditional slicedWasserstein flows. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (eds.), Proceedings of the ICML ’23, pp. 8565–8584. PMLR, 2023.", "weaknesses": "1. The novelty of the proposed method appears to be limited, since it is mainly the Generative Sliced MMD Flow [1] method applied to conditional generative modelling problems. Additionally, the proof of Theorem 3 partially follows [2].\n2. The theoretical comparison with different kernels (Gaussian, Inverse Multiquadric and Laplacian [1]) and discrepancies (KL divergence, W_1 [2] and W_2 [3] distance) in Theorem 2 is insufficient.\n3. The numerical results of image generation lack comparison with other methods like Generative Sliced MMD Flow in [1]. It would be better to compare the FID scores for different datasets and various methods like [1], since the proposed method adopts the computational scheme of Generative Sliced MMD Flow. It would be beneficial to compare with Conditional Normalizing Flow in the superresolution experiment and with WPPFlow, SRFlow in the computed tomography experiment.\n\n\n[1] J. Hertrich, C. Wald, F. Altekrüger, and P. Hagemann. Generative sliced MMD flows with Riesz kernels. arXiv preprint 2305.11463, 2023c\n\n[2] F. Altekrüger, P. Hagemann, and G. Steidl. Conditional generative models are provably robust: pointwise guarantees for Bayesian inverse problems. Transactions on Machine Learning Research, 2023b.\n\n[3] F. Altekrüger and J. Hertrich. WPPNets and WPPFlows: the power of Wasserstein patch priors for superresolution. SIAM Journal on Imaging Sciences, 16(3):1033–1067, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5cvMkVhyua", "reviewer": "Reviewer_MCWB", "summary": "The paper introduces conditional flows of the Maximum Mean Discrepancy (MMD) with the negative distance kernel for posterior sampling and conditional generative modeling. The joint distribution of the ground truth and the observations is approximated using discrete Wasserstein gradient flows, and an error bound for the posterior distributions is established. it is proven in the paper that the particle flow within our method indeed functions as a Wasserstein gradient flow of an appropriate functional. The paper's efficacy is demonstrated through various numerical examples, encompassing applications such as conditional image generation and the resolution of inverse problems, including superresolution, inpainting, and computed tomography in low-dose and limited-angle scenarios.", "questions": "* Standard Sliced Wasserstein is not optimal, there are other variants e.g., [3],[4]. Is standard SW preferred in this setting?\n* Can the proposed MMD flows be seen as a debiased version of Sliced Wasserstein gradient flow in the setting of discrete flows?\n\n[3] Generalized Sliced Wasserstein Distances\n[4] Energy-Based Sliced Wasserstein Distance", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposal of MMD flows with a \"generalized\" kernel kernel which is also known as energy distance or Cramer distance is new. \n* The paper can prove that the particle flow with the generalized MMD is indeed a Wasserstein gradient flow of an appropriate function.\n* The paper uses the MMD flows in the setting of sampling from the posterior which is interesting and new.\n* Experiments are conducted on class-conditional image-generation (MNIST, FashionMNIST, and CIFAR10) and inverse problems with medical images.", "weaknesses": "* There is no quantitative comparison in class-conditional image-generation with previous works e.g., score-based generative modeling (without using labels). Similarly, score-based generative models can also be used in medical image inverse-problem [1].\n* There is no comparison with Sliced Wasserstein Gradient flows e.g., with JKO scheme. [2]\n* Considering discrete flows is quite restricted. \n\n[1] Solving Inverse Problems in Medical Imaging with Score-Based Generative Models.\n[2] Efficient Gradient Flows in Sliced-Wasserstein Space", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)