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 $\|f_{\rho^{MF}} -f_{\hat{\rho}}\|$ 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 $\|f_{\rho^{MF}} -f_{\hat{\rho}}\|^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 $\epsilon$ with only $\text{poly}(d/\epsilon)$ neurons, training samples, and GD steps.