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 Esν(μ)=21∥μ−ν∥H˙−s2 for any s≥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 Sd. 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.
Figures & tables
Figure 1: Regularity regimes of Riesz kernels Ks on the d -torus (bracket orientation indicates inclusion/exclusion of the endpoint) and some values of particular interest (arrows). The negative distance and ReLU neural network cases are not strictly speaking Riesz kernels on the d -torus, but their regularity on the diagonal matches that of the indicated exponent. For d=1 , the kernel is never singular for s≥1 (as then the Coulomb case coincides with the negative distance case).
Figure 2: The case s=1 and d=1 integrated with finite volume discretization with upwind scheme. (a) For s=1 , the regularity of μˉ and ν do not impact the exponential convergence rate. Here γ0 (resp. γν ) is the largest scalar such that μˉ∈H˙γ0 (resp. ν∈H˙γν ) and the densities are lower-bounded by 0.2 so our theoretical rate is O(e−0.2t) . (b) As theory predicts, the convergence rate is upper-bounded by the minimum density of ν , and independent of the minimum density of μˉ . Here m0 (resp. mν ) indicates the minimum of μˉ (resp. ν ). In this experiment, (γ0,γν)=(1,1) . (c) When ν has a positive lower bound, the zero-density areas of μˉ (in black) shrink exponentially fast, as theory predicts (here (γ0,γν)=(1,1) ); see Remark 1.3 and Lemma 3.2 .
Figure 3: The case s=2 and d=1 integrated with finite volume discretization with upwind scheme. We write γ0 (resp. γν ) for the largest scalar such that μˉ∈H˙γ0 (resp. ν∈H˙γν ) and the densities are lower-bounded by 0.2 . (a) Snapshots of the density of μt along the evolution (here (γ0,γν)=(2,2) ). We can observe the absence of a maximum principle (cf. x=0.9 ), and the appearance of high frequency components (cf. x=0.8 ). Both phenomena are absent from the case s=1 . (b) Convergence rate in H˙−2 norm. The dashed lines show approximate fits of the asymptotic rates.
Figure 4: The case of shallow ReLU Neural Network with d=1 , implemented via gradient descent on the population loss with a “student” and “teacher” neural network of width 800 that discretize the measure μt and ν respectively (in other words, this is an interacting particle system approximation of the PDE with 800 particles for each measure). We initialize μˉ with a uniform density ( γ0=+∞ ) and vary the target regularity (we indicate the largest γν such that ν∈H˙γν ). We compare the Wasserstein (W) and Wasserstein–Fisher–Rao (WFR) dynamics. Although (WFR)’s energy decay is slightly faster at initialization (thanks to the extra term in the energy dissipation formula ( 1.18 )), it does not converge faster in general. The dashed lines show approximate fits of the asymptotic rates.