Uniform-in-Time Weak Propagation of Chaos in Shallow Neural Networks

We consider two-layer neural networks trained in the feature-learning regime using gradient descent, and relate the output of the finite-width network $f_{\hat{\rho }}$ to its infinite-width counterpart $f_{{\rho }^{MF}}$, which evolves in the mean-field dynamics. While constant-time horizon bounds for $\Vert f_{{\rho }^{MF}}-f_{\hat{\rho }}\Vert$ may be obtained via standard Grönwall estimates, the long-time behavior of the fluctuation is a more delicate matter. Uniform-in-time bounds often rely on (local) strong convexity in the landscape or Logarithmic Sobolev inequalities present in noisy gradient dynamics. % In this work, we study a noiseless setting and do not make assumptions on the geometry of the landscape near the optimum. Instead, we impose a weaker condition on the convergence rate of the mean-field deterministic Wasserstein-gradient-flow dynamics. In this work, we establish non-asymptotic weak propagation-of-chaos that holds uniformly in time, obtained by exploiting instead the convergence rate of the mean-field deterministic Wasserstein-gradient-flow dynamics. Specifically, denoting by $L_t$ the mean-field loss at time $t$ and $m$ the number of neurons, under standard regularity assumptions and the condition ${\int }_0^{\infty }{L}_t^{1/2}dt=O_d(1)$, we obtain the uniform in time bound $\Vert f_{{\rho }^{MF}}-f_{\hat{\rho }}{\Vert }^2\lesssim \text{poly}(d)m^{-\min (1,c/6)}$ whenever $L_t\lesssim t^{-c}$. Our result holds in a noiseless setting and does not make any assumptions on the geometry of the landscape near the optimum. A key implication of our result is that whenever the convergence rate of the mean-field, population-loss dynamics is faster than $1/t^2$, we can attain a loss of $ϵ$ with only $\text{poly}(d/ϵ)$ neurons, training samples, and GD steps.