Generative Sliced MMD Flows with Riesz Kernels

Thank you very much for your review. We extended the experimental part by including the CelebA dataset which is of higher dimension than the other datasets.

In order to circumvent the dimensionality effect, we use on CelebA two further techniques.

• We use the multiscale approach to generate the images. More precisely, we first generate a downsampled version of the CelebA dataset, use an upsampling step and approximate the higher-dimensional version. From a theoretical side this correspond to change the latent distribution.

• Following the references (Du et al. 2023) and (Nguyen and Ho, 2022), we included local projections.

We describe both techniques in more detail within the implementation details setting in Appendix H.

To explain the dependence on the dimensions theoretically, we have already included Remark 5 in the main paper.

References:

(Du et al, 2023) Nonparametric Generative Modeling with Conditional Sliced-Wasserstein Flows

(Nguyen and Ho, 2022) Revisiting sliced Wasserstein on images: From vectorization to convolution

Thank you very much for your comments. We want to emphasize some advantages of our method which explain why it works that well.

• The MMD and the energy distance (= MMD with Riesz kernel and $r=1$) have shown to be very effective two-sample tests to compare high-dimensional probability measures, see (Szekely, 2004). Therefore, minimizing these metrics via gradient flows is a natural approach. However, previous works on MMD flows (e.g. Arbel et al, 2019) were only applicable for small problems due to the quadratic computational cost in the number of samples. Only with our reduction of the computational cost it is possible to compute such MMD flows on datasets with 50000 images or more in a reasonable time.

• Accurate gradient evaluations: We have proven convergence guarantees stating that the sliced MMD gradients converge to the original MMD gradients. Moreover, we do not require to solve an optimization problem within the steps which is necessary for other gradient flow methods which are often based on implicit schemes (like JKO).

• The sequential learning approach is computationally efficient. The whole generative model $\Phi_L\circ\cdots\circ\Phi_1$ has a large number of parameters and is (as usual for generative models) a highly irregular mapping. However, during the training we only have to consider one of the networks $\Phi_l$ at a time. This leads to a sequence of easier learnable problems instead of one hard one.

We added a paragraph "Discussion" in the conclusions section which summarizes these points.

For the architecture of the $\Phi_l$ we used a standard UNet with the implementation of (Huang et al. 2021). Moreover, in our experiments changes in the architecture did not have a large effect on the results as long as it is expressive enough.

References:

(Szekely et al, 2004) Testing for equal distributions in high dimension

(Arbel et al, 2019) Maximum Mean Discrepancy Gradient Flow

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

Many thanks for your thorough and detailed review. We answer to all of your comments separately.

• Concerning Theorem 2: the statement of Theorem 2 does not follow from a corresponding statement about the Coulomb kernel as we explain below. The Coulomb kernel in Chafaï et al. is defined by $K_C(x,y)=\\|x-y\\|^{2-d}$ for $d\ge3$ without the minus. Therefore, we cannot bound the Riesz kernel from below by the Coulomb kernel. Moreover, Thm 1.1 from Chafaï et al. assumes that $\int\int K_C(x,y)d \mu(x)d \mu(y)<\infty$ and $\int\int K_C(x,y)d\nu(x)d \nu(y)<\infty$. This assumption is violated whenever $\mu(\\{x\\})>0$ or $\nu(\\{x\\})>0$ for any $x$. In particular, this assumption excludes the case of empirical measures which is the most important one for our applications. Nevertheless, since this is an interesting idea and since the bound looks similar at the first glance, we added a comment about that at the end of the proof of Theorem 2 in Appendix B. Even though we do not utilize this bound, we found the relations to the Wasserstein-1 distance quite interesting and would like to keep it in the paper.

• Concerning the sharpness of the bounds in Theorem 2: We do not think that a sharper version of Theorem 2 without the dimension scaling holds true for the following reason: It is well-known that MMD has sample complexity $O(n^{-1/2})$ (Gretton et al. 2012) and Wasserstein-1 has a sample complexity of $O(n^{-1/d})$ (see e.g., Peyré and Cuturi, 2020). If it would hold true that $W_1(\mu,\nu)\leq C\mathcal D_K(\mu,\nu)$ for all $\mu,\nu$ this would imply that also Wasserstein-1 would have a sample complexity of $O(n^{-1/2})$ which is a contradiction. We included a comment about that in the paper. For the Coulomb kernel this argumentation does not apply, since it is not bounded on compact sets (which is an assumption for the sample complexity $O(n^{-1/2})$ for MMD) and since the result of Chafai et al. explicitly excludes empirical measures.

• Regarding Theorem 3: Note that the potential energy $V$ is convex. Consequently, we can use subgradients instead of gradients. We added a comment on this in the paper after Theorem 3.

• Regarding the dimension dependency: We thank the reviewer for raising this point. The fact that $c_{d,r}\in O(d^{r/2})$ is exactly the reason why the $\sqrt{d}=d^{1/2}=d^{r/2}$ appears within the error bound for the gradients in Theorem 4, where $r=1$. We already discussed the consequences in Remark 5. In particular, this causes that the computational complexity for the estimation of gradients of the MMD up to a fixed error $\epsilon$ via slicing depends quadratically on $d$ and not linearly. In the updated version, we inserted comments after Theorem 1 and in Remark 5 in order to make this point more clear. Moreover, we inserted a plot with the approximation error on the right of Figure 1.

• About training a sequence of neural networks (Section 4.2): We train the networks for two reasons. First, the neural networks speed up the generation process and allow us to draw a single sample (for the particle flow, we always have to compute a batch of samples). The second and more important reason is that the particle flow itself converges exactly to the training images, see Figure 9 and the corresponding explanation in the Appendix I of the paper for an illustration. In some sense, we could say that the particle flow works "too well" since it approximates the discrete target measure exactly. The resulting FID of the flow particles would be 1.37 for MNIST, 1.90 for FashionMNIST and 3.90 for CIFAR10. If we use the sequence of neural networks, we can interpolate between the training points in order to generalize the dataset. Finally, note that we do not train a network per step of the particle flow, but one network for 30000 steps of the flow after a warm start procedure. A detailed explanation towards this is given in Appendix H. Thus the training of the networks is not very costly compared to the simulation of the particle flow.

• Regarding related work: We agree that several references from the related work section were relevant for providing a broader context of the topic, but were not directly related to our paper. We moved these references to the general introduction and reorganized the related work section. We hope that this makes it more clear now.

About (Chafaï et al., 2016): you are absolutely correct and I missed that point, this is a good remark. For the generalization/memorization discussion I thank you for the answers. I think it would be good to describe this in the paper: especially the formalism with the push-forward operator and the interpolation path (which were enlightening; at least for me).

All my concerns are addressed and I change my score accordingly.

Thank you very much for the review.

• We extended our experimental part by the CelebA dataset which admits more structure than (Fashion)MNIST. Here, our method admits very good results. In particular, we outperform WGAN significantly (32.1 vs 41.3 in the FID). Moreover, on CIFAR10 our method is at least competitive with other gradient-flow-based methods though we agree that there is a gap between gradient-flow-based and score-based methods like NCSN (in our paper as well as in the literature). Compared with WGAN, we obtain with 54.8 vs 55.2 a similar (even slightly better) FID. Note that previous works on MMD flows (e.g. Arbel et al, 2019) were not able to scale to image generation because of the quadratic computational cost within the number of samples. As mentioned by reviewer HuoH our fast gradient computation "opens up the possibility of employing MMD flows in generative modeling tasks that were previously considered impractical for these types of functions".

• Regarding the theoretical part, we would like to emphasize that our results provably lower the computational complexity of (gradient) evaluations of the MMD/energy distance. All in all the proofs of our theorems are non-trivial and require a wide range of mathematical tools. They are directly related to the applications part and result in a highly efficient algorithm for MMD (gradient) computations.

We hope that the reviewer will reconsider our paper with these explanations.

References:

(Arbel et al, 2019) Maximum Mean Discrepancy Gradient Flow