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

Thank you very much for the review. To address the first two points from the "weaknesses" part, we added a comparison for class-conditional image generation with the previous method of (Du et al. 2023) which is based on sliced Wasserstein gradient flows. For details, we refer to the general answer. We address your other comments below.

• Regarding other slicing variants: Consider other slicing procedures could be an interesting line of future research. In this paper, we focused on the standard mean-slicing procedure, because of its theoretical tractability. We added a sentence to the future-work section that we want to generalize our results for other slicing variants and transfer these results to (sliced) MMD functionals. Note that the focus of the paper is on conditional Wasserstein gradient flows of the MMD functional. Clearly, we have for our special negative distance kernel that "MMD = sliced MMD" which can be used in the numerical part, but plays no role in the rest of the paper.

• Regarding gradient biases: (Bellemare et al. 2017) show that Wasserstein distances have biased sample gradients. Since in 1D Wasserstein and sliced Wasserstein coincide, this probably also holds true for the sliced Wasserstein distance, even though we are not aware of a formal proof of this statement. In Appendix A and Proposition 3 of the same paper (Bellemare et al. 2017) it is shown that the energy distance (= MMD with negative distance kernel) has unbiased sample gradients. In 1D, the energy distance coincides with the so-called Cramer distance, which is presented in Bellemare et al. 2017) as "debiased Wasserstein distance". In this sense, one could indeed view our MMD flows as "debiased sliced Wasserstein gradient flows". However, since we do not focus on gradient biases, we do not want to use this term within the paper.

(Bellemare et al, 2017) The Cramer Distance as a Solution to Biased Wasserstein Gradients

Many thanks for your comments. We added several comparisons in the numerical part. In particular, we added a comparison with WPPFLows on CT and a FID comparison for class-conditional image generation.

Regarding the novelty

Note this is not only "a generative Sliced MMD Flow [1] method applied to conditional generative modelling problems". Our paper is about conditional Wasserstein gradient flows of the MMD functional. There is just a small hint that "MMD = sliced MMD" for the negative distance kernel, see [1], which is irrelevant for the theory part.

The consideration of the conditional flows brings indeed a lot of new theoretical challenges. Two of them are addressed in our paper:

i) Conditional generative models approximate the joint distribution by learning a mapping $T$ such that $P_{X,Y} \approx P_{T(Z,Y)),Y}$, but in fact we are interested in the posterior distributions $P_{X|Y=y}$. In this paper, we prove error bounds between posterior and joint distributions within the MMD metric. To this end, we use relations between measure spaces and RKHS as well as Lipschitz stability results under pushforward measures which are quite involved. These results are new and highly non-trivial to prove. Clearly they are based on preliminary work in particular on [2], but we just cited the necessary results and use them in completely new proofs.

ii) We want to characterize a conditional Wasserstein gradient flow of a functional $F$ as usual Wasserstein gradient flow of a modified functional to better understand the behaviour of the first one. To establish a relation to Wasserstein gradient flows, we locally lift the $\mathbb R^{N d}$ into the Wasserstein space using local isometries. As a byproduct we give a theoretical explanation for the results of (Du et al. 2023).

We have rewritten the contributions paragraph in the introduction to highlight the points i) and ii).

In the applications part, the conditional MMD flows for large-scale imaging inverse problems like computed tomography is new. Of course we use [1] for reducing the computational cost and we agree that this is crucial to obtain a tractable algorithm, but this is not the point of our contribution.

Other comments

• Thanks for the suggestion to compare with a conditional normalizing flow for superresolution and with SRFlow for CT. We would like to point out that SRFlow is the same as a conditional normalizing flow, where the authors adapted the architecture for superresolution. We added a comment about that in the paper.

• Different kernels: Since the negative distance kernel is the only kernel, where we can compute the gradient of MMD in $O(N\log(N))$ instead of $O(N^2)$, it is the only kernel, where our method is applicable. From a theoretical viewpoint, generalizing Theorem 2 to a larger class of kernels would be nice and is a current line of our research. For the KL divergence, we have to restrict to absolutely continuous measures, but get a stronger version of Theorem 2 with an equality statement. We refer to Remark 7 in Appendix A.2 for a detailed discussion of Theorem 2 with KL or Wasserstein-1.

• Advantage over KL and Wasserstein: Neither the KL divergence nor the Wasserstein distance admit a closed-form two-sample formulation. For the KL divergence, we have to estimate the density (or the score) of one of the measures, for Wasserstein we have to solve an optimization problem. In contrast, we can just compute the MMD and its derivative directly on two sets of samples. In order to lower the complexity of the computations, one can apply the slicing procedure from [1] which admits a provable error bound.

• Regarding error bounds between discrete and continuous MMD flows: In the one-dimensional case, (Carillo et al. 2020) proved that the discrete MMD flows with negative distance kernel converge in the mean-field limit to the continuous ones and we cited this in our paper. In higher dimensions, the MMD functional with negative distance kernel is not $\lambda$-convex along generalized geodesics such that most of the theory of Wasserstein gradient flows is no longer applicable. In (Hertrich et al. 2022), the authors derive analytic formulas for MMD flows with negative distance kernel in higher dimensions. Despite these examples, we are not aware of any theoretical results regarding existence, uniqueness or convergence results for such flows for $d\geq 2$. Therefore, we feel that there is a long line of theoretical work which has to be done before a proof of such error bounds would be reachable. Therefore they are beyond the scope of our paper.

References:

[1], [2] from the review

(Carillo et al. 2020) Measure solutions to a system of continuity equations driven by newtonian nonlocal interactions.

(Hertrich et al. 2022) Wasserstein steepest descent flows of discrepancies with Riesz kernels.

(Hertrich et al. 2023) Generative sliced MMD flows with Riesz kernels.

Thank you very much for your comments. We added quantitative evaluations of our method as outlined in the general answer. We address your other comments below.

• Regarding the sensitivity to the forward operator/noise: We want to clarify that this issue is not specific for our method. Instead, any supervised learning method for inverse problems is sensitive to mismatches between the training and testing models. We emphasize this issue in the limitations section, because it is particularly relevant for medical imaging applications, where such mismatches are unavoidable. However, in our paper we want to demonstrate the usability of our method for highly ill-posed and high-dimensional imaging inverse problems. In particular, calibrating it for clinical applications is not within the scope of our paper. This is common-sense in the machine learning community and unfortunately most authors do not even mention that.

• The paper (Altekrüger et al, 2023) does not provide any "proper" generative model. Instead the authors consider forward and backward schemes for non-discretized Wasserstein gradient flows. Consequently, they have to train a generative model in each step of the gradient flow. This allows an extensive analysis of such flows, but is computational costly and therefore not realizable for large problems. Indeed, their MNIST example does not consider the whole dataset but only a subset of a few hundred images.

• The U-Net is a standard choice for many imaging applications as well as in diffusion models, see (Huang et al, 2021). We did also some brief experiments with a lightweight transformer network and achieved slightly worse results. We think that there is not a very large difference as long as the considered network is expressive enough.

• Regarding the inpainting example: CIFAR10 is a highly diverse dataset such that a large variation of the generated samples for inpainting is natural. Consequently, the generated images have to differ from each other greatly. Note that for less diverse datasets like (Fashion)MNIST or CelebA, we obtain much less variation in the reconstructed images.

References

(Du et al, 2023) Nonparametric Generative Modeling with Conditional Sliced-Wasserstein Flows, Du, Li, Pang, Shuicheng, Lin, ICML 2023

(Altekrüger, 2023) Neural Wasserstein Gradient Flows for Discrepancies with Riesz Kernels, Altekrüger, Hertrich, Steidl, ICML 2023

(Huang et al, 2021): A variational perspective on diffusion-based generative models and score matching, Huang, Lim and Courville, NeurIPS 2021

Thank you very much for the review. We added comparisons with other generative models as outlined in the general answer. We address your other comments below.

Main proof techniques and contributions

We have two main theoretical contributions in addition to numerical applications:

i) Conditional generative models approximate the joint distribution by learning a mapping such that $P_{X,Y} \approx P_{T(Z,Y)),Y}$, but in fact we are interested in the posterior distributions $P_{X|Y=y}$. In this paper, we prove error bounds between posterior and joint distributions within the MMD metric. To this end, we use relations between measure spaces and RKHS as well as Lipschitz stability results under pushforward measures which are quite involved.

ii) We want to characterize a conditional Wasserstein gradient flow of a functional $F$ as usual Wasserstein gradient flow of a modified functional to better understand the behaviour of the first one. To establish a relation to Wasserstein gradient flows, we locally embed the $\mathbb R^{N d}$ into the Wasserstein space using local isometries. As a byproduct we give a theoretical explanation for the results of (Du et al. 2023).

iii) Numerically, we approximate our conditional MMD flows by conditional generative neural networks and apply them in various settings like class conditional image generation, image restoration and CT.

We have rewritten the contributions paragraph in the introduction to highlight the points i), ii) and iii).

Difference to sliced Wassestein flows

We want to emphasize that our proposed method does not use sliced distances, also not the sliced Wasserstein distance. We consider conditional gradient flows for the functional $F$=MMD within the $W_2$ geometry. In contrast, conditional sliced Wasserstein gradient flows (Du et al. 2023) consider gradient flows for the functional $F$=sliced Wasserstein within the $W_2$ geometry. Consequently, both methods act in the $W_2$ geometry, but minimize completely different functionals. Our paper contains just a hint that for the negative distance kernel (and only for this kernel) it holds ``MMD = sliced MMD''. This is an exceptional nice property proven in (Hertrich et al. 2023) which we clearly exploit in the numerical part for speeding up the computations.

Other comments:

• Regarding the numerical verification of bounds we have to note that this is a worst case bound. Therefore it is hard to do numerical verification (as obtaining the ground truth posterior is hard/impossible in high dimensions) and it is likely that real world examples actually attain better rates.

• We added some nice properties of the negative distance MMD (efficient calculation and good sample complexity) in the introduction.

• $T$#$\mu$: This denotes the pushforward measure of $\mu$. We included the definition in the paper.

References

(Hertrich et al. 2023) J. Hertrich, C. Wald, F. Altekrüger, and P. Hagemann. Generative sliced MMD flows with Riesz kernels. arXiv preprint 2305.11463