Paper: Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel (ICLR-2024)

Reviewer comment (Experiments):
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].

Author response:
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.

Reviewer comment (Experiments):
There is no comparison with Sliced Wasserstein Gradient flows e.g., with JKO scheme. [2]

Author response:
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.

Reviewer comment (Theory):
Can the proposed MMD flows be seen as a debiased version of Sliced Wasserstein gradient flow in the setting of discrete flows?

Author response:
- 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

Reviewer comment (Novelty):
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].

Author response:
## 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.

Reviewer comment (Theory):
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.

Author response:
- 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.

Reviewer comment (Experiments):
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.

Author response:
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.

- 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.

Reviewer comment (Theory):
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?

Author response:
- 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.

Reviewer comment (Theory):
Is it possible to provide a discretization error analysis between discrete MMD flow and the original continuous MMD flow?

Author response:
- 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.

Reviewer comment (Experiments):
The model is sensitive to forward operator and noise type.

Author response:
- 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.

Reviewer comment (Evaluation):
Lack of meaningful quality metrics to evaluate the results.

Author response:
Thank you very much for your comments. We added quantitative evaluations of our method as outlined in the general answer.

Reviewer comment (Evaluation):
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.

Author response:
Thank you very much for your comments. We added quantitative evaluations of our method as outlined in the general answer.

Reviewer comment (Novelty):
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?

Author response:
- 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.

Reviewer comment (Experiments):
Why chosing UNet? Is there a significant difference in the effect of choosing other models such ResNet and transformer.

Author response:
- 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.

Reviewer comment (Evaluation):
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.

Author response:
- 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.

Reviewer comment (Theory):
It would be better to elaborate on the pros and cons of using a negative distance kernel (efficiency, sample complexity, etc).

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

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Writing):
In Equation 4, $T$ is defined, however $T_\sharp$ is not defined.

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

Reviewer comment (Experiments):
Is it possible to validate the error bound via numerical experiments somehow?

Author response:
- 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.

Reviewer comment (Novelty):
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?

Author response:
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.