(De)-regularized Maximum Mean Discrepancy Gradient Flow

TL;DR

This paper introduces DrMMD, a novel de-regularized MMD gradient flow that guarantees near-global convergence, is computationally feasible with only samples, and is applicable to large-scale problems.

Contribution

The paper proposes DrMMD, a new MMD gradient flow that ensures convergence and can be implemented in closed form using only samples, improving over existing methods.

Findings

DrMMD guarantees near-global convergence for broad target classes.

The method is implementable in closed form with only sample data.

Numerical experiments demonstrate effectiveness in large-scale settings.

Abstract

We introduce a (de)-regularization of the Maximum Mean Discrepancy (DrMMD) and its Wasserstein gradient flow. Existing gradient flows that transport samples from source distribution to target distribution with only target samples, either lack tractable numerical implementation ($f$-divergence flows) or require strong assumptions, and modifications such as noise injection, to ensure convergence (Maximum Mean Discrepancy flows). In contrast, DrMMD flow can simultaneously (i) guarantee near-global convergence for a broad class of targets in both continuous and discrete time, and (ii) be implemented in closed form using only samples. The former is achieved by leveraging the connection between the DrMMD and the $\chi^2$-divergence, while the latter comes by treating DrMMD as MMD with a de-regularized kernel. Our numerical scheme uses an adaptive de-regularization schedule throughout the flow to optimally trade off between discretization errors and deviations from the $\chi^2$ regime. The potential application of the DrMMD flow is demonstrated across several numerical experiments, including a large-scale setting of training student/teacher networks.