Quantitative Convergence of Wasserstein Gradient Flows of Kernel Mean Discrepancies

TL;DR

This paper analyzes the convergence rates of Wasserstein gradient flows of Kernel Mean Discrepancy functionals, with applications to neural network training dynamics and particle systems, providing explicit rates and numerical validation.

Contribution

It establishes the first quantitative convergence rates for Wasserstein gradient flows of KMD functionals, including neural network training and particle systems, with explicit polynomial and exponential rates.

Findings

Exponential convergence for s=1 under minimal assumptions

Polynomial convergence rates for s>1 depending on regularity

Numerical experiments validating theoretical rates

Abstract

We study the quantitative convergence of Wasserstein gradient flows of Kernel Mean Discrepancy (KMD) (also known as Maximum Mean Discrepancy (MMD)) functionals. Our setting covers in particular the training dynamics of shallow neural networks in the infinite-width and continuous time limit, as well as interacting particle systems with pairwise Riesz kernel interaction in the mean-field and overdamped limit. Our main analysis concerns the model case of KMD functionals given by the squared Sobolev distance $ \mathscr{E}^{\nu}_{s}(\mu)= \frac{1}{2}\lVert \mu-\nu\rVert_{\dot H^{-s}}^{2}$ for any $s\geq 1 $ and $\nu$ a fixed probability measure on the $d$-dimensional torus. First, inspired by Yudovich theory for the $2d$-Euler equation, we establish existence and uniqueness in natural weak regularity classes. Next, we show that for $s=1$ the flow converges globally at an exponential rate under minimal assumptions, while for $s>1$ we prove local convergence at polynomial rates that depend explicitly on $s$ and on the Sobolev regularity of $\mu$ and $\nu$. These rates hold both at the energy level and in higher regularity classes and are tight for $\nu$ uniform. We then consider the gradient flow of the population loss for shallow neural networks with ReLU activation, which can be cast as a Wasserstein--Fisher--Rao gradient flow on the space of nonnegative measures on the sphere $\mathbb{S}^d$. Exploiting a correspondence with the Sobolev energy case with $s=(d+3)/2$, we derive an explicit polynomial local convergence rate for this dynamics. Except for the special case $s=1$, even non-quantitative convergence was previously open in all these settings. We also include numerical experiments in dimension $d=1$ using both PDE and particle methods which illustrate our analysis.