Generative Sliced MMD Flows with Riesz Kernels

Johannes Hertrich, Christian Wald, Fabian Altekrüger, Paul Hagemann

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: sliced maximum mean discrepancy, energy distance, gradient flows, Riesz kernels, generative modelling

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

Maximum mean discrepancy (MMD) flows suffer from high computational costs in large scale computations.
In this paper, we show that MMD flows with Riesz kernels $K(x,y) = - \|x-y\|^r$, $r \in (0,2)$
have exceptional properties which allow their efficient computation.
We prove that the MMD of Riesz kernels, which is also known as energy distance, coincides with the MMD of their sliced version.
As a consequence, the computation of gradients of MMDs can be performed in the one-dimensional setting.
Here, for $r=1$, a simple sorting algorithm can be applied to reduce the complexity
from $O(MN+N^2)$ to $O((M+N)\log(M+N))$ for two measures with $M$ and $N$ support points.
As another interesting follow-up result, the MMD of compactly supported measures
can be estimated from above and below by the Wasserstein-1 distance.
For the implementations we approximate the gradient of the sliced MMD by using only a finite number $P$ of slices.
We show that the resulting error has complexity \smash{$O(\sqrt{d/P})$}, where $d$ is the data dimension.
These results enable us to train generative models by approximating MMD gradient flows by neural networks even
for image applications. We demonstrate the efficiency of our model by image generation on MNIST, FashionMNIST and CIFAR10.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 2606