Generative Sliced MMD Flows with Riesz Kernels

TL;DR

This paper introduces an efficient method for computing MMD flows with Riesz kernels, enabling scalable training of generative models for images by leveraging sliced MMD and Riesz kernel properties.

Contribution

The paper demonstrates that MMD with Riesz kernels can be computed efficiently using sliced versions, reducing computational complexity and enabling scalable generative modeling.

Findings

Efficient gradient computation for MMD flows with Riesz kernels.

Reduction of complexity from quadratic to near-linear for certain cases.

Successful application to image generation on MNIST, FashionMNIST, and CIFAR10.

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