Quantitative Convergence of Wasserstein Gradient Flows of Kernel Mean Discrepancies

2026-03-02

AI Summary

This work investigates the quantitative convergence of the kernel mean discrepancy (KMD) functional under Wasserstein gradient flows, with applications to shallow neural network training and interacting particle systems. By integrating Yudovich’s theory, Sobolev space analysis, and the Wasserstein–Fisher–Rao framework, the study establishes, for the first time, global exponential convergence in the case $s = 1$, and provides explicit local polynomial convergence rates for $s > 1$ that depend on regularity—thereby quantifying previously non-quantitative results. Compact convergence estimates are derived in both energy spaces and higher-order regularity classes. In particular, concrete local convergence rates are obtained for ReLU shallow networks, and theoretical predictions are validated through one-dimensional numerical experiments.

Technology Category

• Machine Learning: Kernel Methods
• Search and Optimization: Non-convex Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Models for Web evolution
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

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}^ν_{s}(μ)= \frac{1}{2}\lVert μ-ν\rVert_{\dot H^{-s}}^{2}$ for any $s\geq 1 $ and $ν$ 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 $μ$ and $ν$. These rates hold both at the energy level and in higher regularity classes and are tight for $ν$ 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.

Problem

• Wasserstein gradient flow
• Kernel Mean Discrepancy
• quantitative convergence
• shallow neural networks
• interacting particle systems

Innovation

• Wasserstein gradient flow
• Kernel Mean Discrepancy
• quantitative convergence
• shallow neural networks
• Sobolev energy