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

Abstract: We consider one-hidden layer neural networks trained in the feature-learning regime using gradient descent, and relate the output of the finite-width network $f_{\hatρt^m}$ to its infinite-width counterpart $f{ρt^{MF}}$, which evolves in the mean-field dynamics. While constant-time horizon bounds for $|f{ρt^{MF}} - f{\hatρt^m}|$ 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 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 excess MSE 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(\log d)$, we obtain the uniform in time bound $|f{ρt^{MF}}- f{\hatρ_t^m}|^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, and extends seamlessly to other forms of discretization, including finite number of samples and time discretization. A key takeaway of our result is that whenever the convergence rate of the mean-field, population-loss dynamics is faster than $t^{-2}$, we can attain a loss of $ε$ with only $\text{poly}(d/ε)$ neurons, training samples, and GD steps.

• The paper establishes that shallow networks under gradient descent can uniformly approximate their infinite-width mean-field limit with weak propagation-of-chaos bounds.
• It leverages the decay rates of the mean-field loss and a refined coupling ODE to control finite-size fluctuations without requiring strong convexity or noise.
• The study implies that finite-width networks need only polynomial resources to achieve target accuracy on smooth regression tasks, marking a significant theoretical advance.

Uniform-in-Time Weak Propagation-of-Chaos in Shallow Neural Networks: An Expert Analysis

Overview and Problem Formulation

The paper "Uniform-in-Time Weak Propagation-of-Chaos in Shallow Neural Networks" (2605.22010) analyzes the longstanding question of how finite-width, feature-learning two-layer neural networks (1HL NNs) track their infinite-width mean-field (MF) limits under gradient descent (GD). Unlike the kernel regime—where the parameters evolve little and learning is tightly controlled—the feature-learning regime presents pronounced non-convexity and coupling between neurons, complicating uniform-in-time approximation to the mean-field trajectory.

Previous analyses, reliant on Grönwall-type estimates for coupled particles, yield short-horizon ($O(\exp(L t))$) error control, becoming ineffective for long training times due to exponential divergence. Conversely, uniform-in-time results, typical in mean-field Langevin settings, depend on strong convexity or an injected noise mechanism enabling contraction (e.g., via Logarithmic Sobolev Inequalities), which does not directly extend to deterministic GD in neural networks.

This work focuses on "weak" Propagation of Chaos (PoC): the convergence of the empirical network output $f_{\hat\rho^m}$ to its mean-field analog $f_\rho$ in the $L^2$ loss, over arbitrary time horizons, for the noiseless gradient flow trained on feature-learner networks.

Theoretical Contributions and Main Results

The core technical advance is the derivation of non-asymptotic, uniform-in-time PoC bounds for the observable function discrepancy $\|f_{\hat\rho^m} - f_\rho\|^2$ incurred by empirical (finite-width $m$) GD against MF training. Critically, the result dispenses with strong convexity or explicit regularization; instead, it leverages decay rates of the MF loss $\mathcal{L}_t$ to control finite-size fluctuations in the absence of stochastic contraction mechanisms.

Letting $L_t$ denote the MF excess MSE at time $t$, the key theorem asserts:

• If $\int_0^\infty \sqrt{L_t}\,dt < O(\log d)$ and $f_{\hat\rho^m}$0 decays as $f_{\hat\rho^m}$1 with $f_{\hat\rho^m}$2, the following uniform-in-time error bound holds:

$f_{\hat\rho^m}$3

whenever $f_{\hat\rho^m}$4. This scaling generalizes: for $f_{\hat\rho^m}$5, the optimal $f_{\hat\rho^m}$6 rate is recovered, up to dimension factors and regularity constants.

• Crucially, no assumptions are made on local strong convexity or geometry near the optimum, nor is noise required. As a corollary, achieving population loss $f_{\hat\rho^m}$7 requires only $f_{\hat\rho^m}$8 resources (neurons, samples, GD steps): a strong claim compared to prior bounds with exponential-in-time dependencies.

  Figure 1: Approximate loss $f_{\hat\rho^m}$9 and integrated root-loss $f_\rho$0 for $f_\rho$1, demonstrating sufficiently fast decay rates and convergence of the characteristic integral as assumed by the main theorem.

Fine-Grained Error Analysis

The analysis introduces a refined coupling ODE for the fluctuation vector $f_\rho$2: tracking the differences between coupled particles in the finite and mean-field systems. A new decomposition of the fluctuation dynamics distinguishes a "constant" (near convergence) error source from higher-order vanishing terms, allowing for sharp estimates. When the MF loss decays rapidly, the local Hessians $f_\rho$3 contract fluctuations, and the function error remains uniformly bounded.

A key lemma upper-bounds the operator norm of the local fluctuation amplifier $f_\rho$4 by $f_\rho$5, meaning the dissipativity of the finite-particle system tightens in tandem with mean-field convergence.

Figure 2: Examples of target densities (top), trajectories of $f_\rho$6 (middle), and the corresponding integral $f_\rho$7 (bottom), illustrating settings where the main theorem’s assumptions are respected (i.e., $f_\rho$8 converges).

Empirical Illustration

Theoretical requirements are tested on synthetic single-index models with parameterized Sobolev smoothness. For various $f_\rho$9 (regularity exponent), networks are trained on targets of increasing smoothness, and loss curves are tracked.

• In low smoothness settings ($L^2$0), the integral $L^2$1 may diverge, violating the theorem's hypothesis and leading to insufficient control over long-term fluctuations.
• For higher $L^2$2, empirical loss decay is fast enough for the integral to converge, reflecting the practical range where the stated uniform-in-time control applies.

  Figure 3: Plot of $L^2$3 for various values of $L^2$4, with increased smoothness yielding improved convergence properties and ultimately better propagation-of-chaos behavior.

Implications and Future Directions

The practical upshot is that for many natural regression tasks with smooth target functions and favorable data geometry, finite-width networks accurately track their infinite-width MF limit at all times, provided MF convergence is fast enough. This provides algorithmic guarantees for resource scaling in overparameterized regimes without requiring noise or strong convexity, and extends to discretizations involving finite samples and learning rates.

Theoretically, the result indicates that contraction induced by loss decay can replace explicit stochastic contraction for certain observables—suggesting a broader universality for mean-field descriptions in deep learning dynamics.

Possible future research avenues include:

• Extending the uniform-in-time PoC control to other functionals (beyond output loss).
• Weakening the required decay rate (e.g., for cases where $L^2$5 with $L^2$6).
• Generalizing to models with deeper architectures (e.g., ResNets, Transformers), where MF theory has shown promise but remains analytically challenging.
• Investigating settings where MF convergence exhibits dimension-dependent "burn-in" phases, as these may limit the extent of uniform-in-time controllability.

Conclusion

This work provides a non-asymptotic, uniform-in-time weak propagation-of-chaos result for shallow neural networks in the feature-learning regime, grounded in the decay of the mean-field loss and bypassing commonly assumed convexity or regularization conditions. The findings bolster the relevance of mean-field tools for analyzing non-kernel, high-dimensional neural learning, while opening new directions for deterministic propagation-of-chaos analysis and scalable non-convex learning guarantees.