# Uniform-in-time concentration in two-layer neural networks via transportation inequalities

Arnaud Guillin, Boris Nectoux, Paul Stos

<figure><img src="image_1.png" /></figure>

## To cite this version:

Arnaud Guillin, Boris Nectoux, Paul Stos. Uniform-in-time concentration in two-layer neural networks via transportation inequalities. 2026. hal-05526186

HAL Id: hal-05526186

https://hal.science/hal-05526186v1

Preprint submitted on 25 Feb 2026

HAL is a multi-disciplinary open access archive for the deposit and dissemination of scientific research documents, whether they are published or not. The documents may come from teaching and research institutions in France or abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est destinée au dépôt et à la diffusion de documents scientifiques de niveau recherche, publiés ou non, émanant des établissements d'enseignement et de recherche français ou étrangers, des laboratoires publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
# Uniform-in-time concentration in two-layer neural networks via transportation inequalities

Arnaud Guillin*¹, Boris Nectoux†¹, and Paul Stos‡¹

¹UNIVERSITÉ CLERMONT-AUVERGNE, CNRS UMR 6620, LMBP, Clermont-Ferrand, France

## Abstract

We quantify, uniformly over time and with high probability, the discrepancy between the predictions of a two-layer neural network trained by stochastic gradient descent (SGD) and their mean-field limit, for quadratic loss and ridge regularization. As a key ingredient, we establish $T_p$ transportation inequalities ($p \in \{1, 2\}$) for the law of the SGD parameters, with explicit constants independent of the iteration index. We then prove uniform-in-time concentration of the empirical parameter measure around its mean-field limit in the Wasserstein distance $W_1$, and we translate these bounds into prediction-error estimates against a fixed test function $\varphi$. We also derive analogous concentration bounds in the sliced-Wasserstein distance $SW_1$, leading to dimension-free rates.

## 1 Introduction

Neural networks have become a standard tool for a variety of high-dimensional prediction problems [21, 24, 39], and their empirical success has fueled substantial interest in understanding the theoretical mechanisms underpinning their behavior [44, 40, 1].

In this paper we focus on the classical framework of supervised learning for a regression task. We observe a stream of independent data points $(x_k, y_k) \in \mathbb{R}^d \times \mathbb{R}$ drawn from a common (unknown) distribution $\pi$. The goal is to learn a predictor $\hat{y}$ of the label $y$ that generalizes to previously unseen features $x$. In a one-hidden-layer neural network of width $N$, this predictor is parametrized by $\theta = (\theta^1, \dots, \theta^N) \in \mathcal{E} = (\mathbb{R}^D)^N$ and takes the form

$$ \hat{y}_{\theta}(x) = \frac{1}{N} \sum_{i=1}^{N} \sigma_{*}(\theta^{i}, x), $$

where $\sigma_* : \mathbb{R}^D \times \mathbb{R}^d \to \mathbb{R}$ is an activation function. Typically, $\theta^i = (a_i, b_i, w_i) \in \mathbb{R}^{d+2}$ and $\sigma_*(\theta^i, x) = a_i \sigma(w_i \cdot x + b_i)$ for some $\sigma : \mathbb{R} \to \mathbb{R}$ (e.g., sigmoid, tanh, ReLU). The factor $1/N$ is the standard mean-field normalization of the network output [15]; it is made for convenience and can be absorbed into the activation by simple redefinition.

In principle, one would like to choose $\theta$ so as to minimize the population risk

$$ \mathcal{R}(\theta) = \mathbb{E}_{\pi}[\ell(y, \hat{y}_{\theta}(x))], $$

where $\ell : \mathbb{R} \times \mathbb{R} \to \mathbb{R}$ is a loss function. For regression, we use the square loss $\ell(y, \hat{y}) = \frac{1}{2}(y - \hat{y})^2$ out of simplicity, though other losses could be considered.

In practice, since we do not have direct access to the data distribution $\pi$, a standard approach consists in minimizing an empirical proxy $\hat{\mathcal{R}}$ built from the observed samples via an optimization algorithm such as stochastic gradient descent (SGD) or its variants [7, 8].

<footnote>*arnaud.guillin@uca.fr</footnote>

<footnote>†boris.nectoux@uca.fr</footnote>

<footnote>‡paul.stos@uca.fr</footnote>
In the present setting, parameters are initialized i.i.d. from $\mu_0 \in \mathcal{P}(\mathbb{R}^D)$ and updated according to online SGD with fixed learning rate $\alpha > 0$: for every $i \in \{1, \dots, N\}$,

$$ \theta_{k+1}^i = (1 - \gamma\lambda)\theta_k^i + \gamma(y_{k+1} - \hat{y}_{\theta_k}(x_{k+1}))\nabla_{\theta^i} \sigma_*(\theta_k^i, x_{k+1}), \quad (1) $$

where $\gamma = \alpha/N$. The term $-\gamma\lambda$ corresponds to a ridge regularization of the empirical risk

$$ \hat{\mathcal{R}}_{\lambda}(\theta, x, y) = \frac{1}{2}(y - \hat{y}_{\theta}(x))^2 + \frac{\lambda}{2}\|\theta\|_2^2 $$

which will be useful for our analysis below.

In the online setting, each data point is used only once and never revisited. This idealization is standard in the learning literature and is relevant whenever data are abundant or streaming, so that repeated reuse of a fixed finite dataset is not the leading phenomenon [7, 37].

Providing quantitative guarantees for the training dynamics of neural networks remains challenging, even in the basic two-layer architecture trained by stochastic gradient descent. This difficulty stems not only from high dimensionality, but also from the feedback structure of learning: at each iteration, each parameter is updated through a gradient that depends on the current network output, while this output aggregates the contributions of all $N$ neurons and is evaluated on random data samples. As a result, SGD produces a strongly coupled stochastic evolution in parameter space $\mathcal{E}$.

For wide two-layer networks, a particularly fruitful approach has been to model training as a large interacting particle system and to track it via the empirical distribution of the neurons' parameters [36], or its rescaled-in-time version [15, 16], defined for all $t \ge 0$ by

$$ \mu_t^N = \frac{1}{N} \sum_{i=1}^N \delta_{\theta^i_{[Nt]}} \in \mathcal{P}(\mathbb{R}^D). $$

As the number of neurons $N \to \infty$, this empirical process converges in probability (in an appropriate topology) to a deterministic evolution $(\bar{\mu}_t)_{t \ge 0}$ characterized as the solution to a nonlinear measure-valued equation [27, 42, 16] (see (13)). This asymptotic description yields a more tractable macroscopic model for the learning dynamics, and has proved useful for studying fluctuations [41, 16] and for analyzing stability and long-time behavior [27, 10].

Existing quantitative finite-width bounds are typically derived on a fixed horizon $[0, T]$, with constants deteriorating (often exponentially) as $T$ grows; see, e.g., [27, Theorem 3]. This is unsatisfactory in practice, since learning algorithms are routinely run for many iterations.

In this work, we address this gap by quantifying uniformly over $t \ge 0$ the deviation of the empirical parameter measure $\mu_t^N$ from its mean-field limit $\bar{\mu}_t$, under an explicit contractivity condition induced by ridge regularization. More precisely, we study three complementary levels of discrepancies:

* $\Delta_t^{N,\varphi} = |\langle\varphi, \mu_t^N\rangle - \langle\varphi, \bar{\mu}_t\rangle|$ for a fixed 1-Lipschitz test function $\varphi: \mathbb{R}^D \to \mathbb{R}$. In particular, assuming Lipschitz activation (assumption (A1) below), for any fixed input $x \in \mathbb{R}^d$, taking $\varphi = \sigma_*(\cdot, x)$ translates any bound on $\Delta_t^{N,\varphi}$ into a bound on the network output.

* $\Delta_t^N = W_1(\mu_t^N, \bar{\mu}_t)$, where $W_1$ is the classical 1-Wasserstein distance on $\mathcal{P}_1(\mathbb{R}^D)$ (the set of Borel probability measures on $\mathbb{R}^D$ with finite first moment). Recall the Kantorovich-Rubinstein duality formula for $W_1$ (see, e.g., [46, Remark 6.5]):

$$ W_1(\mu, \nu) = \sup_{\|\psi\|_{\text{Lip}} \le 1} |\langle\psi, \mu\rangle - \langle\psi, \nu\rangle|. \quad (2) $$

Hence $\Delta_t^N$ simultaneously controls $\Delta_t^{N,\varphi}$ for all 1-Lipschitz $\varphi$.

* $\Delta_t^{N,s} = SW_1(\mu_t^N, \bar{\mu}_t)$, where $SW_1$ is the 1-sliced-Wasserstein distance, defined by

$$ SW_1(\mu, \nu) = \int_{\mathbb{S}^{D-1}} W_1(P_{u\#\mu}, P_{u\#\nu}) \zeta(du), \quad P_u(w) = u \cdot w, $$
with $\zeta$ the uniform probability measure on the unit sphere $S^{D-1}$. The sliced analog $SW_1$ is widely used in high dimension and is computationally tractable [6, 34, 29] as it reduces comparisons to one-dimensional projections, while still probing a rich family of observables (all linear projections). This makes $\Delta_t^{N,s}$ a useful compromise between the very targeted discrepancy $\Delta_t^{N,\phi}$ and the fully global metric $\Delta_t^N$.

When the particular choice is immaterial, we write

$$ \Delta_t^{N,\circ} \in \{\Delta_t^{N,\phi}, \Delta_t^N, \Delta_t^{N,s}\}. $$

We refer to [38, 46, 34] for background on Wasserstein distances and optimal transport (see also [33] for a more machine-learning-oriented introduction), and to [6, 28, 29] for sliced-Wasserstein and related sliced divergences.

Our analysis hinges on two steps. First, under regularity conditions on the initialization law $\mu_0$ and the data distribution $\pi$, we establish a $T_p$ transportation inequality [22] ($p \in \{1, 2\}$) for the law $\nu_k = \mathcal{L}(\theta_k)$ of the parameter vector at iteration $k$, with explicit constants independent of $k$. This is the content of Proposition 1. For both $p = 1$ and $p = 2$, this implies sub-Gaussian tails (7) for centered Lipschitz observables of the parameters (hence of $\mu_t^N$), with constants that do not deteriorate with time. Instantiated with suitable Lipschitz functionals of $\mu_t^N$, this first step yields high-probability deviation bounds around the mean for our three discrepancies of interest (see Corollary 1). We then bound the remaining bias term $\mathbb{E}[\Delta_t^{N,\circ}]$ uniformly over $t \ge 0$ by comparing the interacting SGD particles to an i.i.d. mean-field system through a synchronous coupling. This yields a decomposition into a *dynamic error* along trajectories, handled via a standard *propagation-of-chaos* argument [45] (Proposition 3 and Corollary 2), and a *static* i.i.d. sampling error for empirical measures, controlled by classical empirical rates in Wasserstein distance [20] (Propositions 4 and 5). Combining the two steps, we obtain a *nonasymptotic*, *uniform-in-time* concentration bound around the mean-field limit. This constitutes our main result, presented as Theorem 1. These results are established under standard boundedness and Lipschitz assumptions on the activation $\sigma_*$. The final section explains how to relax boundedness by allowing at most linear growth in the parameter variable on the support of the data, hence covering more general ReLU-type activations, via localization.

**Related works.** Mean-field approximations for the online training of wide two-layer neural networks were introduced in [27], where the limiting dynamics is derived together with finite-width bounds on a fixed time horizon $[0, T]$. A trajectory law of large numbers and central limit theorem for the empirical parameter measure on any fixed time interval were proved in [42, 41], and later extended at the level of the full trajectory (under mini-batching and noisy regimes) in [16]. For complementary analyses of the mean-field dynamics, including stability and long-time behavior under various structural assumptions, we refer to [36, 10, 27]. We also mention the quantitative propagation-of-chaos analysis for (continuous-time) SGD dynamics developed in [11].

Beyond fixed-horizon bounds, obtaining *time-uniform* control typically requires some form of *dissipativity* (e.g., contractivity in a suitable metric) to prevent error accumulation; see, e.g., [18, 23, 14] in the setting of weakly interacting diffusions. Uniform-in-time exponential bounds for Wasserstein fluctuations have been established for related particle systems (including heterogeneous mean-field interactions) in [2], extending transport-based concentration estimates for empirical measures [5]; see also the general mean-field particle-model framework of [13].

In a different direction, several works extend mean-field ideas to deep architectures and other infinite-width regimes. Rigorous mean-field formalisms for multilayer networks are developed in [43, 30, 31, 19]. Mean-field viewpoints have also been proposed for transformer architectures, where token dynamics can be interpreted as interacting particle systems [35, 9].

From a more geometric perspective, the classical i.i.d. Wasserstein sampling-error bounds of Fournier and Guillin [20] have recently been sharpened in [12], where one obtains similar optimal rates with the ambient dimension replaced by the covering dimension of the support. This
suggests a potential direction for future work in our setting: if the mean-field limit concentrates on a low-dimensional subset of parameter space, then identifying an effective support and estimating (or upper-bounding) its covering dimension could yield strictly improved rates.

During the preparation of this manuscript, we became aware of [47], which establishes quantitative propagation-of-chaos bounds for general McKean-type nonlinear Markov chains and discusses long-time behavior under additional structural assumptions. While the settings and techniques differ, both works support the viewpoint that uniform-in-time propagation of chaos is fundamentally tied to an underlying contractive structure.

The rest of this paper is organized as follows. Our standing assumptions and main results are stated in Section 2. Section 3 contains the proofs. Finally, Section 4 presents an extension of the previous results to linearly growing activations.

**Notation.** All random variables are defined on a common probability space $(\Omega, \mathcal{F}, \mathbb{P})$. For a random variable $Z$, $\mathcal{L}(Z)$ denotes the law of $Z$. We use $|\cdot|_{\infty}$ both for the absolute value on $\mathbb{R}$ and the Euclidean norm on $\mathbb{R}^d$ or $\mathbb{R}^D$. The scalar product on $\mathbb{R}^D$ is denoted by $v \cdot w$. We write $|\cdot|_{\infty}$ for the sup norm on $\mathcal{Z} = \mathbb{R}^d \times \mathbb{R}$, and $\|\cdot\|_{p}$ ($p \in \{1, 2\}$) for the usual $\ell^p$ product norm on $\mathcal{E} = (\mathbb{R}^D)^N$:

$$ \|\theta\|_p = \left( \sum_{i=1}^{N} |\theta^i|^p \right)^{1/p}, \quad \theta = (\theta^1, \dots, \theta^N) \in \mathcal{E}. $$

Given a metric space $(\mathcal{X}, \mathbf{d})$ and a function $f: \mathcal{X} \to \mathbb{R}$, we define (whenever finite)

$$ \|f\|_{\infty} = \sup_{x \in \mathcal{X}} |f(x)|, \quad \|f\|_{\text{Lip}} = \sup_{x \neq y} \frac{|f(x) - f(y)|}{\mathbf{d}(x, y)}. $$

For measurable $\varphi: \mathbb{R}^D \to \mathbb{R}$ and $\mu \in \mathcal{P}(\mathbb{R}^D)$, we use the bracket notation

$$ \langle \varphi, \mu \rangle = \int_{\mathbb{R}^D} \varphi \, d\mu. $$

In particular, for any input $x \in \mathbb{R}^d$ and $t \ge 0$, we have

$$ \hat{y}_{\theta_{\lfloor Nt \rfloor}}(x) = \langle \sigma_*(\cdot, x), \mu_t^N \rangle. $$

Throughout this work, $C > 0$ denotes a generic constant (independent of $N$, $t$ and $p$) whose value may change from line to line.

# 2 Main results

## 2.1 Assumptions

We define the following set of assumptions.

(A1) The activation function $\sigma_* \in \mathcal{C}^1(\mathbb{R}^D \times \mathbb{R}^d)$; moreover, $\sigma_*$ is uniformly bounded with bounded mixed gradient, namely

$$ \|\sigma_*\|_{\infty} \le B, \quad \|\nabla_{(\theta^i, x)} \sigma_*\|_{\infty} \le M. $$

(A2) Labels are almost surely bounded: $|y| \le A \pi$-a.s.

(A3) For each $\theta^i \in \mathbb{R}^D$, the map $x \in \mathbb{R}^d \mapsto \nabla_{\theta^i} \sigma_*(\theta^i, x)$ is $L_x$-Lipschitz uniformly in $\theta^i$; for each $x \in \mathbb{R}^d$, the map $\theta^i \in \mathbb{R}^D \mapsto \nabla_{\theta^i} \sigma_*(\theta^i, x)$ is $L_\theta$-Lipschitz uniformly in $x$.
Assumptions (A1)–(A3) are standard in the learning literature: for instance, [27] assumes bounded activations and labels together with bounded-Lipschitz drift terms, and uses a sub-Gaussian control on $\nabla_{\theta^i} \sigma_*(\theta^i, \cdot)$ to handle high-dimensional inputs. Our conditions are slightly more direct (uniform bounds and Lipschitz constants), which makes the stability constants explicit and keeps the subsequent concentration arguments transparent. The main limitation is that (A1) excludes non-smooth or unbounded activations such as ReLU. This condition will be relaxed in the final Section 4.

Under these assumptions, we collect the following constants, used repeatedly in what follows.

$$ K = \alpha((A+B)L_x + M^2 + M), \quad \lambda_* = (A+B)L_\theta + M^2, \quad L_N = |1 - \gamma\lambda| + \gamma\lambda_*, \\ C_* = 8M^4 + 4(\lambda + (A+B)L_\theta)^2, \quad N_* = \left[ \frac{4\alpha C_*}{\lambda - \lambda_*} \right]. \tag{3} $$

## 2.2 Transportation inequalities for the SGD dynamics

Transportation inequalities relate relative entropy to optimal transport and offer a direct route to proving Gaussian-type concentration for Lipschitz observables. Here, they form the first ingredient of our analysis: we show that the law $\nu_k = \mathcal{L}(\theta_k)$ of the parameter vector at iteration $k$ satisfies a $T_p$ transportation inequality (for $p \in \{1, 2\}$) with constants independent of $k$ (Proposition 1), which in turn yields uniform sub-Gaussian concentration along the SGD Markov chain (Corollary 1). We briefly recall below the definition of $T_p$ and the few properties we will use; for general background about this type of inequalities and their concentration of measure consequences, we refer for instance to [25, 26, 22] and the references therein.

Given a Polish metric space $(\mathcal{X}, \mathbf{d})$ and $p \ge 1$, we say that a probability measure $\mu \in \mathcal{P}(\mathcal{X})$ satisfies an $L^p$-transportation inequality on $(\mathcal{X}, \mathbf{d})$ with constant $C > 0$, and we write $\mu \in T_p(C)$, if for all $\nu \in \mathcal{P}(\mathcal{X})$,

$$ W_p(\nu, \mu) \le \sqrt{2C H(\nu | \mu)}, \tag{4} $$

where $W_p$ denotes the usual $p$-Wasserstein distance on $\mathcal{P}_p(\mathcal{X})$ (the set of Borel probability measures with finite $p$-th moment), and

$$ H(\nu | \mu) = \begin{cases} \int_{\mathcal{X}} \log \frac{d\nu}{d\mu} d\nu & \text{if } \nu \ll \mu \\ +\infty & \text{otherwise} \end{cases} $$

is the relative entropy of $\nu$ with respect to $\mu$.

Among all values of $p \ge 1$, the cases $p = 1$ and $p = 2$ have attracted the most attention. For $T_1$, Bobkov-Götze's theorem [4] gives an equivalence between (4) and Gaussian concentration: $\mu \in T_1(C)$ on $(\mathcal{X}, \mathbf{d})$ if and only if, for every Lipschitz $f: (\mathcal{X}, \mathbf{d}) \to \mathbb{R}$, $f$ is $\mu$-integrable and

$$ \forall \xi \in \mathbb{R}, \quad \mathbb{E}_{\mu}[e^{\xi(f - \mathbb{E}_{\mu}[f])}] \le \exp\left(\frac{C\xi^2}{2}\|f\|_{\text{Lip}}^2\right). \tag{5} $$

In particular, if $\mu$ is compactly supported with, say $\text{supp } \mu \subset B(x_0, R)$ for some $x_0 \in \mathcal{X}$ and $R > 0$, then (5) immediately follows from Hoeffding's lemma, implying $\mu \in T_1(R^2/4)$. More generally, $T_1$ is also equivalent to an exponential integrability condition on $\mathbf{d}(x, x_0)^2$, see [17, Theorem 2.3]. By Chernoff's method, (5) in turn yields the sub-Gaussian tail bound

$$ \forall r > 0, \quad \mathbb{P}(|f(X) - \mathbb{E}[f(X)]| \ge r) \le 2 \exp\left(-\frac{r^2}{2C \|f\|_{\text{Lip}}^2}\right), \quad X \sim \mu. \tag{6} $$

Note that $T_2(C)$ is stronger than $T_1(C)$ (since $W_1 \le W_2$), and therefore also implies Gaussian concentration for Lipschitz functionals. In contrast with $T_1$, compact support alone does not guarantee $T_2$ in general; proving $T_2$ typically requires additional structure, such as a logarithmic Sobolev inequality [32].
We will repeatedly use the stability of $T_p$ under Lipschitz pushforward [17, Lemma 2.1]: if $\mu \in T_p(C)$ on $(\mathcal{X}, \mathbf{d})$ and $\Psi : (\mathcal{X}, \mathbf{d}) \to (\mathcal{Y}, \mathbf{d}')$ is $L$-Lipschitz, then $\Psi_{\#}\mu \in T_p(L^2C)$ on $(\mathcal{Y}, \mathbf{d}')$.

We will also rely on the following i.i.d. tensorization identities. For $n \ge 1$, equip $\mathcal{X}^n$ with the product metrics

$$ \mathbf{d}_1(x, y) = \sum_{i=1}^{n} \mathbf{d}(x_i, y_i), \quad \mathbf{d}_2(x, y) = \left( \sum_{i=1}^{n} \mathbf{d}(x_i, y_i)^2 \right)^{1/2}. $$

If $\mu \in T_1(C)$ on $(\mathcal{X}, \mathbf{d})$, then $\mu^{\otimes n} \in T_1(nC)$ on $(\mathcal{X}^n, \mathbf{d}_1)$ and if $\mu \in T_2(C)$ on $(\mathcal{X}, \mathbf{d})$, then $\mu^{\otimes n} \in T_2(C)$ on $(\mathcal{X}^n, \mathbf{d}_2)$; see [26, Proposition 2.3] and the remark after [26, Proposition 1.3].

All of our statements below hold for either choice of $p \in \{1, 2\}$. To propagate a $T_p$ inequality along the SGD Markov chain, it is natural to assume a $T_p$ inequality at initialization and for the data distribution:

(A4$_p$) The initialization law $\mu_0$ satisfies $\mu_0 \in T_p(C_0^{(p)})$ on $(\mathbb{R}^D, |\cdot|)$, and the data distribution $\pi$ satisfies $\pi \in T_p(C_\pi^{(p)})$ on $(\mathbb{R}^d \times \mathbb{R}, |\cdot|_\infty)$.

Under assumption (A4$_p$) and the one-step contractivity condition $L_N < 1$ (see (3)), we can propagate $T_p$ inequalities uniformly over the SGD iterates, in the spirit of [17, Theorem 2.5] and [3, Theorem 1.2].

**Proposition 1.** Fix $p \in \{1, 2\}$. Assume (A1)-(A3), (A4$_p$), and $L_N < 1$. Then for all $k \in \mathbb{N}$, $\nu_k \in T_p(C_N^{(p)})$ on $(\mathcal{E}, \|\cdot\|_p)$, with the explicit constants

$$ C_N^{(1)} = N C_0^{(1)} + \frac{K^2}{(1 - L_N^2)} C_\pi^{(1)}, \quad C_N^{(2)} = \frac{C_0^{(2)} \vee K^2 C_\pi^{(2)} N^{-1}}{(1 - L_N)^2}. $$

*Remark.* Under our scaling $\gamma = \alpha/N$, one typically has $\gamma\lambda_* < 1$ for $N$ large enough, in which case the condition $L_N < 1$ amounts to the admissible interval $\lambda_* < \lambda < 2\gamma^{-1} - \lambda_*$ for the regularization coefficient. The lower bound $\lambda_* < \lambda$ ensures the regularization dominates the interaction term, making the SGD update *contractive* in the parameter variable (see Lemma 1 (ii) below). The upper bound $\lambda < 2\gamma^{-1} - \lambda_*$ is a discretization artifact of (1), and is essentially nonbinding for wide networks.

As a direct consequence of Proposition 1 (see indeed (6)), for any Lipschitz $f : (\mathcal{E}, \|\cdot\|_p) \to \mathbb{R}$ and $p \in \{1, 2\}$, it holds for every $k \in \mathbb{N}$,

$$ \forall r > 0, \quad \mathbb{P}(|f(\theta_k) - \mathbb{E}[f(\theta_k)]| \ge r) \le 2 \exp \left( - \frac{r^2}{2C_N^{(p)} \|f\|_{\text{Lip},p}^2} \right) \quad (7) $$

where

$$ \|f\|_{\text{Lip},p} = \sup_{\theta \neq \theta'} \frac{|f(\theta) - f(\theta')|}{\|\theta - \theta'\|_p}. $$

For $t \ge 0$ and $\theta = (\theta^1, \dots, \theta^N) \in \mathcal{E}$, let

$$ \mu_{\theta} = \frac{1}{N} \sum_{i=1}^{N} \delta_{\theta^i} \in \mathcal{P}(\mathbb{R}^D), $$

and define for any 1-Lipschitz $\varphi : \mathbb{R}^D \to \mathbb{R}$,

$$ f_t^{\varphi}(\theta) = |\langle \varphi, \mu_{\theta} \rangle - \langle \varphi, \bar{\mu}_t \rangle|, \quad f_t(\theta) = W_1(\mu_{\theta}, \bar{\mu}_t), \quad f_t^{\mathbf{s}}(\theta) = SW_1(\mu_{\theta}, \bar{\mu}_t). $$

By construction, evaluating along the SGD iterates at iteration $k = \lfloor Nt \rfloor$, we thus recover

$$ f_t^{\varphi}(\theta_{\lfloor Nt \rfloor}) = \Delta_t^{N, \varphi}, \quad f_t(\theta_{\lfloor Nt \rfloor}) = \Delta_t^{N}, \quad f_t^{\mathbf{s}}(\theta_{\lfloor Nt \rfloor}) = \Delta_t^{N, \mathbf{s}}. \quad (8) $$
Using a canonical coupling (pairing the i-th atoms), for any $\theta, \theta' \in \mathcal{E}$,

$$W_1(\mu_\theta, \mu_{\theta'}) \le \frac{1}{N} \sum_{i=1}^{N} |\theta^i - \theta'^i| = \frac{1}{N} \|\theta - \theta'\|_1 \le \frac{1}{\sqrt{N}} \|\theta - \theta'\|_2.$$

By Kantorovich-Rubinstein duality (2), any 1-Lipschitz $\varphi$ satisfies $|\langle\varphi, \mu\rangle - \langle\varphi, \nu\rangle| \le W_1(\mu, \nu)$, and since each projection $P_u$ is 1-Lipschitz, we also have $SW_1(\mu, \nu) \le W_1(\mu, \nu)$. Hence by triangular inequality, $f_t^\varphi, f_t$, and $f_t^s$ are all $N^{-1/p}$-Lipschitz on $(\mathcal{E}, \|\cdot\|_p)$, for $p \in \{1, 2\}$.

Although the constant $C_N^{(1)}$ in Proposition 1 may scale like $N$, notice that this does not affect (7) since for both $p=1$ and $p=2$, we have $C_N^{(p)} N^{-2/p} = O(N^{-1})$. Hence, combining (7) with (8) yields the following uniform deviation inequality.

**Corollary 1.** Fix $p \in \{1, 2\}$ and assume (A1)-(A3), (A4$_p$) and $L_N < 1$. Then there exists $C > 0$ (independent of $N$ and $t$) such that for any $t \ge 0$ and $\delta \in (0, 1)$, with probability $1-\delta$,

$$|\Delta_t^{N,\circ} - \mathbb{E}[\Delta_t^{N,\circ}]| \le C \sqrt{\frac{\log(2/\delta)}{N}}.$$

## 2.3 Uniform bias decay and concentration around mean-field

Corollary 1 above provides a concentration inequality for $\Delta_t^N$ (and similarly for $\Delta_t^{N,\varphi}$ and $\Delta_t^{N,s}$) around its mean. It thus remains to control this bias uniformly in time in order to recover a concentration inequality for $\mu_t^N$ around $\bar{\mu}_t$.

Assumption (A4$_p$) introduced above is only needed to derive Proposition 1 and Corollary 1. For the bias estimates below, it is enough to assume a finite $(2+\varepsilon)$-moment at initialization:

(A5) There exists $q > 2$ such that $\mathbb{E}[|\theta_0^1|^q] < \infty$.

In particular, for either choice of $p \in \{1, 2\}$, assumption (A4$_p$) implies (A5) (since both $T_1$ and $T_2$ imply exponential integrability and therefore moments of all orders).

**Proposition 2.** Assume (A1)-(A3), (A5) and $L_N < 1$, $N \ge N_\star$ (see (3)). Then there exists $C > 0$ (independent of $N$ and $t$) such that

$$\sup_{t \ge 0} \mathbb{E}[\Delta_t^{N,\circ}] \le C \kappa_N,$$

where

$$\kappa_N = \begin{cases} N^{-1/2} & \text{for } \Delta_t^{N,\varphi} \text{ and } \Delta_t^{N,s}, \\ N^{-1/(1\vee D)} + N^{-1/2} & \text{for } \Delta_t^N. \end{cases}$$

*Remark.* The only dimension-dependent term in $\kappa_N$ appears for the Wasserstein discrepancy $\Delta_t^N = W_1(\mu_t^N, \bar{\mu}_t)$, through the rate $N^{-1/(1\vee D)}$. This term originates from the optimal i.i.d. sampling error of empirical measures in $W_1$ on $\mathbb{R}^D$ [20, Theorem 1], and reflects the intrinsic complexity of optimal matching in high dimension, also known as the *curse of dimensionality*. By contrast, both $\Delta_t^{N,\varphi}$ and $\Delta_t^{N,s}$ enjoy dimension-free $N^{-1/2}$ rates, since they reduce to one-dimensional (or scalar) averages. Sharper dimension-specific rates and possible logarithmic corrections are available in low dimension (see the discussion after Theorem 1 in [20]), but the above bound is sufficient for our purposes.

Finally, combining Corollary 1 with Proposition 2, we obtain the following uniform-in-time concentration inequality between the empirical measure of SGD parameters and its mean-field limit.

**Theorem 1.** Assume (A1)-(A3), (A4$_p$) (for either $p=1$ or $p=2$), and $L_N < 1$, $N \ge N_\star$ (see (3)). Then there exists $C > 0$ (independent of $N$ and $t$) such that for any $t \ge 0$ and $\delta \in (0, 1)$, with probability $1-\delta$,

$$\Delta_t^{N,\circ} \le C \kappa_N + C \sqrt{\frac{\log(2/\delta)}{N}}.$$

with $\kappa_N \to 0$ as $N \to \infty$ given in Proposition 2.
**Application to the network output.** Fix $x \in \mathbb{R}^d$ and let $\varphi_x = \sigma_*(\cdot, x)$. Under (A1), $\varphi_x$ is M-Lipschitz uniformly in $x$. Applying Theorem 1 to $\Delta_t^{N,\varphi}$ for $t \ge 0$, with $\varphi = \varphi_x/M$ thus yields that for any $\delta > 0$, with probability $1 - \delta$,

$$|\hat{y}_{\theta_{\lfloor Nt \rfloor}}(x) - \langle \sigma_*(\cdot, x), \bar{\mu}_t \rangle| \le C \sqrt{\frac{1 + \log(2/\delta)}{N}}.$$

# 3 Proofs

## 3.1 SGD dynamics

We start by rewriting the SGD recursion in a form that makes its Markov structure explicit, since Proposition 1 is ultimately a statement about transport inequalities along the resulting Markov chain. The proof combines a Lipschitz control of the one-step update (Lemma 1) with stability results for $T_p$ inequalities under contracting Markov kernels. For $p = 1$, propagation of $T_1$ is obtained by an elementary induction based on the Laplace transform characterization (5). For $p = 2$, we instead rely on a dependent tensorization result for contracting Markov chains [3, Theorem 1.2] applied at the level of the full trajectory, and then project onto the last coordinate to recover a $T_2$ inequality uniform in the iteration index.

For $\theta = (\theta^1, \dots, \theta^N) \in \mathcal{E} = (\mathbb{R}^D)^N$ and $z = (x, y) \in \mathcal{Z} = \mathbb{R}^d \times \mathbb{R}$, define the one-step drift

$$ 
\begin{aligned}
F_{\lambda}^{i}(\theta, z) &= (y - \hat{y}_{\theta}(x)) \nabla_{\theta^{i}} \sigma_{*}(\theta^{i}, x) - \lambda \theta^{i}, \quad i \in \{1, \dots, N\}, \\
F_{\lambda}(\theta, z) &= (F_{\lambda}^{1}(\theta, z), \dots, F_{\lambda}^{N}(\theta, z)) \in \mathcal{E},
\end{aligned}
\qquad (9) 
$$

and the associated update map $\Phi(\theta, z) = \theta + \gamma F_{\lambda}(\theta, z)$. Then (1) is equivalently

$$ \theta_{k+1} = \theta_k + \gamma F_{\lambda}(\theta_k, z_{k+1}) = \Phi(\theta_k, z_{k+1}) \qquad (10) $$

with $z_{k+1} = (x_{k+1}, y_{k+1}) \in \mathcal{Z}$. Since $(z_k)_{k \ge 1}$ are i.i.d., the parameter sequence $(\theta_k)_{k \in \mathbb{N}}$ thus forms a time-homogeneous Markov chain on $\mathcal{E}$ with transition kernel

$$ \forall \theta \in \mathcal{E}, \quad P(\theta, d\theta') = [\Phi(\theta, \cdot)_{\#\pi}](d\theta'). \qquad (11) $$

**Lemma 1.** Assume (A1)-(A3). Then, for any $\theta, \theta' \in \mathcal{E}$ and $z, z' \in \text{supp } \pi$,

$$ (i) \quad \|\Phi(\theta, z) - \Phi(\theta, z')\|_2 \le \frac{K}{\sqrt{N}} |z - z'|_{\infty}, \quad \text{and} \quad \|\Phi(\theta, z) - \Phi(\theta, z')\|_1 \le K |z - z'|_{\infty}, $$

$$ (ii) \quad \|\Phi(\theta, z) - \Phi(\theta', z)\|_p \le L_N \|\theta - \theta'\|_p \quad \text{for any } p \in \{1, 2\}, $$

with constants $K, L_N$ defined in (3).

*Proof.* (i) Let $\theta, \theta' \in \mathcal{E}$ and $z = (x, y)$, $z' = (x', y') \in \text{supp } \pi \subset \mathbb{R}^d \times [-A, A]$ (by (A2)). For every $i \in \{1, \dots, N\}$, by triangular inequality,

$$ 
\begin{aligned}
|\Phi^i(\theta, z) - \Phi^i(\theta, z')| &= \gamma|(y - \hat{y}_\theta(x))\nabla_{\theta^i}\sigma_*(\theta^i, x) - (y' - \hat{y}_\theta(x'))\nabla_{\theta^i}\sigma_*(\theta^i, x')| \\
&\quad \pm (y - \hat{y}_\theta(x))\nabla_{\theta^i}\sigma_*(\theta^i, x')| \\
&\le \gamma|y - y'| |\nabla_{\theta^i}\sigma_*(\theta^i, x')| + \gamma|y - \hat{y}_\theta(x)||\nabla_{\theta^i}\sigma_*(\theta^i, x) - \nabla_{\theta^i}\sigma_*(\theta^i, x')| \\
&\quad + \gamma|\hat{y}_\theta(x') - \hat{y}_\theta(x)||\nabla_{\theta^i}\sigma_*(\theta^i, x')|.
\end{aligned}
$$

By (A1), the map $x \in \mathbb{R}^d \mapsto \sigma_*(\theta^i, x)$ is M-Lipschitz (uniformly in $i$), hence $|\hat{y}_\theta(x') - \hat{y}_\theta(x)| \le M|x - x'|$. Using also (A3), we get

$$ |\Phi^i(\theta, z) - \Phi^i(\theta, z')| \le \gamma M |y - y'| + \gamma((A+B)L_x + M^2)|x - x'|. $$

It follows that

$$ \|\Phi(\theta, z) - \Phi(\theta, z')\|_2 \le \gamma \sqrt{N} (M |y - y'| + ((A+B)L_x + M^2) |x - x'|) \le \frac{K}{\sqrt{N}} |z - z'|_{\infty}, $$
and consequently

$$ \|\Phi(\theta, z) - \Phi(\theta, z')\|_1 \le \sqrt{N}\|\Phi(\theta, z) - \Phi(\theta, z')\|_2 \le K |z - z'|_{\infty}. $$

(ii) Again, by triangular inequality, for every $i \in \{1, \dots, N\}$,

$$ \begin{aligned} |\Phi^i(\theta, z) - \Phi^i(\theta', z)| &= |(1 - \gamma\lambda)(\theta^i - \theta'^i) + \gamma(y - \hat{y}_{\theta}(x))\nabla_{\theta^i}\sigma_*(\theta^i, x) \\ &\quad - \gamma(y - \hat{y}_{\theta'}(x))\nabla_{\theta^i}\sigma_*(\theta'^i, x) \pm \gamma(y - \hat{y}_{\theta}(x))\nabla_{\theta^i}\sigma_*(\theta'^i, x)| \\ &\le |1 - \gamma\lambda||\theta^i - \theta'^i| + \gamma|y - \hat{y}_{\theta}(x)||\nabla_{\theta^i}\sigma_*(\theta^i, x) - \nabla_{\theta^i}\sigma_*(\theta'^i, x)| \\ &\quad + \gamma|\hat{y}_{\theta'}(x) - \hat{y}_{\theta}(x)||\nabla_{\theta^i}\sigma_*(\theta'^i, x)|. \end{aligned} $$

By (A1), the map $\theta^i \in \mathbb{R}^D \mapsto \sigma_*(\theta^i, x)$ is M-Lipschitz (uniformly in $x \in \mathbb{R}^D$), hence

$$ |\hat{y}_{\theta'}(x) - \hat{y}_{\theta}(x)| \le \frac{1}{N} \sum_{i=1}^{N} |\sigma_*(\theta'^i, x) - \sigma_*(\theta^i, x)| \le \frac{M}{N^{1/p}} \|\theta - \theta'\|_p. $$

Using also (A2) and (A3), we get

$$ |\Phi^i(\theta, z) - \Phi^i(\theta', z)| \le (|1 - \gamma\lambda| + \gamma(A + B)L_{\theta}) |\theta^i - \theta'^i| + \gamma \frac{M^2}{N^{1/p}} \|\theta - \theta'\|_p. $$

It follows by Minkowski inequality that $\|\Phi(\theta, \cdot) - \Phi(\theta', \cdot)\|_p \le L_N \|\theta - \theta'\|_p$ for all $p \in \{1, 2\}$. $\square$

With Lemma 1 in hand, we are now in position to pass from one-step stability to uniform-in-time concentration along the SGD iterates.

*Proof of Proposition 1.* Case $p = 1$. The proof goes by induction on $k \in \mathbb{N}$. Since $\mu_0 \in T_1(C_0^{(1)})$, by tensorization we have $\nu_0 = \mu_0^{\otimes N} \in T_1(NC_0^{(1)})$ on $(\mathcal{E}, \|\cdot\|_1)$. Now, assume that $\nu_k \in T_1(c_k)$ for some $c_k > 0$ and $k \in \mathbb{N}$. Let us show that $\nu_{k+1} \in T_1(c_{k+1})$ with

$$ c_{k+1} = L_N^2 c_k + K^2 C_{\pi}^{(1)}, \qquad (12) $$

where $K$ and $L_N$ are the Lipschitz constants of Lemma 1. Here, we use the equivalence between $T_1$ and Gaussian concentration (5). Let $f : (\mathcal{E}, \|\cdot\|_1) \to \mathbb{R}$ be 1-Lipschitz. Take $Z \sim \pi$, and let for $\theta \in \mathcal{E}$, $g(\theta) = \mathbb{E}[f(\Phi(\theta, Z))]$. By Lemma 1, it follows that $g$ is $L_N$-Lipschitz, and for each fixed $\theta \in \mathcal{E}$, the map $z \mapsto f(\Phi(\theta, z))$ is $K$-Lipschitz. Since $\pi \in T_1(C_{\pi}^{(1)})$, the concentration bound (5) gives, for all $\theta \in \mathcal{E}$,

$$ \forall \xi \in \mathbb{R}, \quad \mathbb{E}[e^{\xi(f(\Phi(\theta, Z)) - g(\theta))}] \le \exp\left(\frac{C_{\pi}^{(1)}\xi^2}{2} K^2\right). $$

Now, let $\theta = \theta_k \sim \nu_k$ be random and independent of $Z$, so that $\Phi(\theta_k, Z)$ is equal in distribution to $\theta_{k+1}$ (see (10)). Conditioning on $\theta_k$, we get for any $\xi \in \mathbb{R}$,

$$ \begin{aligned} \mathbb{E}[e^{\xi(f(\theta_{k+1}) - \mathbb{E}[f(\theta_{k+1})])}] &= \mathbb{E}[e^{\xi(g(\theta_k) - \mathbb{E}[g(\theta_k)])}] \mathbb{E}[e^{\xi(f(\theta_{k+1}) - g(\theta_k))} | \theta_k] \\ &\le \exp\left(\frac{C_{\pi}^{(1)}\xi^2}{2} K^2\right) \mathbb{E}[e^{\xi(g(\theta_k) - \mathbb{E}[g(\theta_k)])}]. \end{aligned} $$

Since $\nu_k \in T_1(c_k)$ and $g$ is $L_N$-Lipschitz, we have again by (5)

$$ \forall \xi \in \mathbb{R}, \quad \mathbb{E}[e^{\xi(g(\theta_k) - \mathbb{E}[g(\theta_k)])}] \le \exp\left(\frac{c_k \xi^2}{2} L_N^2\right). $$

Combining the last two inequalities thus yields

$$ \forall \xi \in \mathbb{R}, \quad \mathbb{E}[e^{\xi(f(\theta_{k+1}) - \mathbb{E}[f(\theta_{k+1})])}] \le \exp\left(\frac{\xi^2}{2}(L_N^2 c_k + K^2 C_{\pi}^{(1)})\right) $$
which is in turn equivalent to $\nu_{k+1} \in T_1(c_{k+1})$ with $c_{k+1}$ as in (12). Under $L_N < 1$, it follows that for all $k \ge 1$,

$$c_k = L_N^{2k} N C_0^{(1)} + K^2 C_\pi^{(1)} \sum_{j=0}^{k-1} L_N^{2j} \le N C_0^{(1)} + \frac{K^2}{(1-L_N^2)} C_\pi^{(1)} = C_N^{(1)},$$

which is the claimed bound for $p = 1$.

Case $p = 2$. Here, we adopt a different approach: we first establish a $T_2$ inequality for the joint law of the whole trajectory $(\theta_0, \dots, \theta_k) \in \mathcal{E}^{k+1}$ using a dependent tensorization theorem for contracting Markov chains, and then project onto the last coordinate to recover a $T_2$ inequality for $\nu_k$, with a constant independent of $k$.

Fix $k \in \mathbb{N}$ and denote by $\nu^{(k)}$ the law of the trajectory $(\theta_0, \dots, \theta_k)$ on $\mathcal{E}^{k+1}$, equipped with the $\ell^2$ product distance

$$d_2((\theta_0, \dots, \theta_k), (\theta'_0, \dots, \theta'_k)) = \left( \sum_{j=0}^{k} \|\theta_j - \theta'_j\|_2^2 \right)^{1/2}.$$

We apply [3, Theorem 1.2] with $s = 2$ to the measure $\nu^{(k)}$ on $(\mathcal{E}^{k+1}, d_2)$, and verify that assumptions (i) and (ii) there are satisfied.

First, since $\mu_0 \in T_2(C_0^{(2)})$, by tensorization we have $\nu_0 = \mu_0^{\otimes N} \in T_2(C_0^{(2)})$. Moreover, using stability under Lipschitz pushforward and $\pi \in T_2(C_\pi^{(2)})$, we get by Lemma 1 (i) and (11) that $P(\theta, \cdot) \in T_2(C_P)$ with $C_P = C_\pi^{(2)} K^2/N$ independent of $\theta$. Therefore, condition (i) in [3, Theorem 1.2] holds with constant $C_0^{(2)} \lor C_P$.

Second, let $Z \sim \pi$ and consider the coupling $(\Phi(\theta, Z), \Phi(\theta', Z))$ of $P(\theta, \cdot)$ and $P(\theta', \cdot)$. The contraction condition (ii) in [3, Theorem 1.2] is then a direct consequence of Lemma 1 (ii) since

$$W_2^2(P(\theta, \cdot), P(\theta', \cdot)) \le \mathbb{E}[\|\Phi(\theta, Z) - \Phi(\theta', Z)\|_2^2] \le L_N^2 \|\theta - \theta'\|_2^2.$$

Under $L_N < 1$, [3, Theorem 1.2] gives $\nu^{(k)} \in T_2(C_N^{(2)})$ on $(\mathcal{E}^{k+1}, d_2)$ with constant $C_N^{(2)} = (1-L_N)^{-2}(C_0^{(2)} \lor C_P)$ independent of $k$. Finally, the projection $(\theta_0, \dots, \theta_k) \mapsto \theta_k$ is 1-Lipschitz from $(\mathcal{E}^{k+1}, d_2)$ to $(\mathcal{E}, \|\cdot\|_2)$, therefore $\nu_k \in T_2(C_N^{(2)})$, which is the claimed bound for $p = 2$. $\square$

## 3.2 Uniform bias decay

We now derive a uniform-in-time bound on the bias term $\mathbb{E}[\Delta_t^{N,0}]$ appearing in Corollary 1. To this end, we compare the SGD iterates to an auxiliary i.i.d. particle system driven by the mean-field dynamics, defined below.

Let us first formalize the mean-field limit by introducing the evolution equation governing the law of a typical particle in the infinite-width regime. We define the mean-field limit $(\bar{\mu}_t)_{t \ge 0}$ as the unique (deterministic) solution in $\mathcal{C}(\mathbb{R}_+, P_1(\mathbb{R}^D))$ to the continuity equation

$$\partial_t \bar{\mu}_t + \alpha \nabla \cdot (G_\lambda(\cdot, \bar{\mu}_t) \bar{\mu}_t) = 0, \quad \bar{\mu}_0 = \mu_0, \qquad (13)$$

where $G_\lambda$ represents the mean-field drift: for any $(\theta^i, \mu) \in \mathbb{R}^D \times \mathcal{P}(\mathbb{R}^D)$,

$$G_\lambda(\theta^i, \mu) = \mathbb{E}_\pi \left[ (y - \langle \sigma_*(\cdot, x), \mu \rangle) \nabla_{\theta^i} \sigma_*(\theta^i, x) \right] - \lambda \theta^i. \qquad (14)$$

Equation (13) is to be understood in the weak sense (see, e.g., [38, Definition 4.1]), namely for any smooth compactly supported test function $\varphi \in C_c^1(\mathbb{R}^D)$,

$$\frac{d}{dt} \langle \varphi, \bar{\mu}_t \rangle = \alpha \langle \nabla \varphi \cdot G_\lambda(\cdot, \bar{\mu}_t), \bar{\mu}_t \rangle.$$
Under (A1)-(A3), existence and uniqueness for (13) in the penalized case follows from a straightforward adaptation of the arguments in [16, §2.3.1] (the regularization term adds a linear dissipative drift and the Lipschitz estimates in $W_1$ used there remain unchanged).

A useful probabilistic tool for the analysis of (13) is given by the associated nonlinear dynamics: following the arguments of [45, Theorem 1.1], there exists a unique (in law) trajectory solution $(\bar{\theta}_t)_{t \ge 0}$ to

$$ \dot{\bar{\theta}}_t = \alpha G_\lambda(\bar{\theta}_t, \rho_t), \quad \bar{\theta}_0 \sim \mu_0, \quad \rho_t = \mathcal{L}(\bar{\theta}_t), \qquad (15) $$

and $(\rho_t)_{t \ge 0}$ satisfies (13). In addition, under (A1), (A2) and (A5), it is straightforward to show that $(\rho_t)_{t \ge 0} \in \mathcal{C}(\mathbb{R}_+, \mathcal{P}_1(\mathbb{R}^D))$ and that $t \ge 0 \mapsto \rho_t$ is in fact Lipschitz-continuous in $W_1$. It follows by uniqueness of the solution to the mean-field PDE (13) that $\rho_t = \bar{\mu}_t$ for all $t \ge 0$. Therefore, we introduce, as it is customary, the i.i.d. particle system $(\bar{\theta}_t^i)_{t \ge 0}$, $i \in \{1, \dots, N\}$, starting from the same points as (1) and evolving independently for $t \ge 0$ according to the nonlinear dynamics (15):

$$ \forall t \ge 0, \quad \bar{\theta}_t^i = \theta_0^i + \alpha \int_0^t G_\lambda(\bar{\theta}_s^i, \bar{\mu}_s) ds, \qquad (16) $$

so that $\mathcal{L}(\bar{\theta}_t^i) = \bar{\mu}_t$ for all $t \ge 0$. The associated empirical measure is denoted by

$$ \bar{\mu}_t^N = \frac{1}{N} \sum_{i=1}^N \delta_{\bar{\theta}_t^i}. $$

By triangular inequality, we decompose the bias into

$$ \mathbb{E}[\Delta_t^N] \le \mathbb{E}[W_1(\mu_t^N, \bar{\mu}_t^N)] + \mathbb{E}[W_1(\bar{\mu}_t^N, \bar{\mu}_t)], \qquad (17) $$

and similarly for $\mathbb{E}[\Delta_t^{N,\varphi}]$ and $\mathbb{E}[\Delta_t^{N,s}]$. The first term in (17) captures a *dynamic error* between the SGD particles and their mean-field counterparts. It is controlled uniformly over $t \ge 0$ by a standard *propagation-of-chaos* argument [45], based on a synchronous coupling.

**Proposition 3.** Assume (A1)-(A3), (A5), and $\lambda > \lambda_*$, $N \ge N_*$ (see (3)). Then there exists $C > 0$ (independent of $N$ and $t$) such that

$$ \sup_{t \ge 0} \mathbb{E}[W_1(\mu_t^N, \bar{\mu}_t^N)] \le \frac{C}{\sqrt{N}}. $$

Recalling $|\langle\varphi, \mu\rangle - \langle\varphi, \nu\rangle| \le W_1(\mu, \nu)$ for any 1-Lipschitz $\varphi$ and $SW_1 \le W_1$, Proposition 3 directly implies Corollary 2 below.

**Corollary 2.** Under the assumptions of Proposition 3, there exists $C > 0$ (independent of $N$ and $t$) such that

$$ \sup_{t \ge 0} \mathbb{E}[|\langle \varphi, \mu_t^N \rangle - \langle \varphi, \bar{\mu}_t^N \rangle|] \le \frac{C}{\sqrt{N}}, \quad \sup_{t \ge 0} \mathbb{E}[SW_1(\mu_t^N, \bar{\mu}_t^N)] \le \frac{C}{\sqrt{N}}. $$

The second term in (17) is a purely *static* i.i.d. sampling error of empirical measures, and is controlled by the optimal $W_1$ rate of [20, Theorem 1], which we recall in the following proposition.

**Proposition 4.** Assume (A1), (A2) and (A5). Then there exists $C > 0$ (independent of $N$ and $t$) such that

$$ \sup_{t \ge 0} \mathbb{E}[W_1(\bar{\mu}_t^N, \bar{\mu}_t)] \le CN^{-1/(1 \vee D)}. $$

We also require static i.i.d. sampling rates for the empirical measure $\bar{\mu}_t^N$ when tested against a test function $\varphi$ and through $SW_1$. Unlike for $W_1$, we expect these two quantities to admit a uniform $N^{-1/2}$ rate independent of the parameter dimension $D$, since the first is simply a scalar sample average and the second reduces to one-dimensional Wasserstein distances through projections. This intuition is formalized in the next proposition.
**Proposition 5.** Assume (A5). Then there exists $C > 0$ (independent of $N$ and $t$) such that

$$ \sup_{t \ge 0} \mathbb{E}[|\langle \varphi, \bar{\mu}_t^N \rangle - \langle \varphi, \bar{\mu}_t \rangle|] \le \frac{C}{\sqrt{N}}, \quad \sup_{t \ge 0} \mathbb{E}[\mathrm{SW}_1(\bar{\mu}_t^N, \bar{\mu}_t)] \le \frac{C}{\sqrt{N}}. $$

Combining (17), Propositions 3–5 and Corollary 2 yields Proposition 2 (note that condition $L_N < 1$ in Proposition 2 implies $\lambda > \lambda_*$ in Proposition 3). The proofs of Propositions 3, 4 and 5 are carried out in the next two sections.

### 3.2.1 Proof of Propositions 4-5

We begin by establishing a uniform-in-time moment bound for the i.i.d. mean-field particles $\bar{\theta}_t^i$. It will be invoked repeatedly throughout the proofs of Propositions 3, 4, and 5.

**Lemma 2.** Assume (A1)–(A2), and $\mathbb{E}[|\theta_0^1|^q] < \infty$ for some $q > 1$. Then there exists $C_q > 0$ (independent of $N$ and $t$) such that

$$ \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|^q] \le C_q < \infty. $$

*Proof.* Fix $q > 1$ such that $\mathbb{E}[|\theta_0^1|^q] < \infty$, and let $f_q(t) = \mathbb{E}[|\bar{\theta}_t^1|^q]$, $t \ge 0$. For $q > 1$, the map $N_q: w \mapsto |w|^q$ is continuously differentiable on $\mathbb{R}^D$, with $\nabla N_q(w) = q|w|^{q-2}w$ for $w \ne 0$ and $\nabla N_q(0) = 0$. Therefore, using (15) and the chain rule, we have

$$ \begin{aligned} f'_q(t) &= \mathbb{E}[\nabla N_q(\bar{\theta}_t^1) \cdot \dot{\bar{\theta}}_t^1] \\ &= q\alpha \mathbb{E}[|\bar{\theta}_t^1|^{q-2}\bar{\theta}_t^1 \cdot G_\lambda(\bar{\theta}_t^1, \bar{\mu}_t)] \\ &\le -q\alpha\lambda \mathbb{E}[|\bar{\theta}_t^1|^q] + q\alpha C \mathbb{E}[|\bar{\theta}_t^1|^{q-1}] \end{aligned} $$

with $C = (A+B)M$ in the last inequality by (A1)–(A2). By Young’s inequality with exponents $q$ and $q' = q/(q-1)$, for any $\varepsilon > 0$,

$$ C|\bar{\theta}_t^1|^{q-1} \le \frac{\varepsilon^{q'}}{q'}|\bar{\theta}_t^1|^q + \frac{C^q}{q\varepsilon^q}. $$

Choosing $\varepsilon^{q'} = \lambda$ and injecting into the previous differential inequality, we get

$$ f'_q(t) \le -\alpha\lambda f_q(t) + \alpha \frac{C^q}{\lambda^{q-1}}. $$

A Grönwall estimate then gives, for all $t \ge 0$,

$$ f_q(t) \le e^{-\alpha\lambda t} f_q(0) + \frac{C^q}{\lambda^q} (1 - e^{-\alpha\lambda t}) \le \mathbb{E}[|\theta_0^1|^q] + \frac{C^q}{\lambda^q} = C_q < \infty, $$

which concludes the proof. $\square$

*Proof of Proposition 4.* Let $q > 2$ be given by (A5). By Lemma 2, there exists $C_q > 0$ (independent of $t$ and $N$) such that

$$ \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|^q] \le C_q. $$

Therefore, applying [20, Theorem 1] with $p = 1$ and exponent $q > 2$ to the i.i.d. sample $(\bar{\theta}_t^1, \dots, \bar{\theta}_t^N)$ of law $\bar{\mu}_t$ yields, for all $t \ge 0$,

$$ \mathbb{E}[\mathrm{W}_1(\bar{\mu}_t^N, \bar{\mu}_t)] \le C (\mathbb{E}[|\bar{\theta}_t^1|^q])^{1/q} N^{-1/(1\vee D)}, $$

where $C$ depends only on $D$ and $q$. The result follows by taking the supremum over $t \ge 0$. $\square$
Proof of Proposition 5. Let $t \ge 0$ and set $U_t^i = \varphi(\bar{\theta}_t^i)$, which are i.i.d. with $\mathbb{E}[U_t^1] = \langle \varphi, \bar{\mu}_t \rangle$. By the Cauchy-Schwarz inequality,

$$ \mathbb{E}[|\langle \varphi, \bar{\mu}_t^N \rangle - \langle \varphi, \bar{\mu}_t \rangle|] = \mathbb{E}\left[\left|\frac{1}{N}\sum_{i=1}^N (U_t^i - \mathbb{E}[U_t^1])\right|\right] \le \frac{\sqrt{\mathrm{Var}(U_t^1)}}{\sqrt{N}}. $$

Since $\varphi$ is 1-Lipschitz, $\mathrm{Var}(U_t^1) \le \mathbb{E}[|U_t^1 - \varphi(0)|^2] \le \mathbb{E}[|\bar{\theta}_t^1|^2]$. By (A5), $\mathbb{E}[|\theta_0^1|^2] < \infty$. We may thus apply the uniform moment bound of Lemma 2 with $q = 2$ and take the supremum over $t \ge 0$ in the previous display to obtain

$$ \sup_{t \ge 0} \mathbb{E}[|\langle \varphi, \bar{\mu}_t^N \rangle - \langle \varphi, \bar{\mu}_t \rangle|] \le \frac{C}{\sqrt{N}}. $$

Let us now turn to the analogous bound for the sliced-Wasserstein distance. By Fubini's theorem,

$$ \mathbb{E}[\mathrm{SW}_1(\bar{\mu}_t^N, \bar{\mu}_t)] = \int_{\mathbb{S}^{D-1}} \mathbb{E}[\mathrm{W}_1(P_{u\#\bar{\mu}_t^N}, P_{u\#\bar{\mu}_t})] \zeta(du). \quad (18) $$

For each $u \in \mathbb{S}^{D-1}$, $P_{u\#\bar{\mu}_t^N}$ is the empirical measure of $N$ i.i.d. samples with common law $P_{u\#\bar{\mu}_t}$ on $\mathbb{R}$. By (A5) and Lemma 2, there exists $q > 2$ and $C_q > 0$ (independent of $D$, $N$ and $t$) such that

$$ \sup_{t \ge 0} \sup_{u \in \mathbb{S}^{D-1}} \int_{\mathbb{R}} |w|^q P_{u\#\bar{\mu}_t}(dw) = \sup_{t \ge 0} \sup_{u \in \mathbb{S}^{D-1}} \mathbb{E}[|\langle u, \bar{\theta}_t^1 \rangle|^q] \le \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|^q] \le C_q. $$

Therefore, applying [20, Theorem 1] in dimension $D = 1$ with $p = 1$ gives a constant $C > 0$ (depending only on $q$) such that for all $t \ge 0$ and for all $u \in \mathbb{S}^{D-1}$,

$$ \mathbb{E}[\mathrm{W}_1(P_{u\#\bar{\mu}_t^N}, P_{u\#\bar{\mu}_t})] \le \frac{C}{\sqrt{N}}. $$

Integrating over $u \sim \zeta$ and taking the supremum over $t \ge 0$ thus yields the desired bound. $\square$

### 3.2.2 Proof of Proposition 3

Although the network's parameters are initialized i.i.d., this independence is generally lost under the SGD dynamics. However, it is easy to see that exchangeability of the system is preserved at all times, i.e. the law of the parameters is left invariant under permutations of their coordinates. Indeed, this is a direct consequence of the permutation-equivariance of the update map $\Phi$: for any permutation $\sigma$ of $\{1, \dots, N\}$,

$$ \forall \theta \in \mathcal{E}, \quad \Phi(\theta^{\sigma}, \cdot) = \Phi(\theta, \cdot)^{\sigma} \quad (19) $$

where $\theta^{\sigma} = (\theta^{\sigma(1)}, \dots, \theta^{\sigma(N)})$. In addition, even though the SGD and mean-field systems are coupled through the same initialization (so they are a priori not independent), the resulting family of paired particles is still exchangeable. This is made precise in the following lemma.

**Lemma 3.** For every $t \ge 0$, the law of $(\theta_{\lfloor Nt \rfloor}, \bar{\theta}_t)$ is exchangeable.

*Proof.* Fix a permutation $\sigma$ of $\{1, \dots, N\}$, and let $t \ge 0$ and $k = \lfloor Nt \rfloor$. By definition (16), there exists a deterministic map $\Psi_t: \mathbb{R}^D \to \mathbb{R}^D$ (independent of $i$) such that for every $i \in \{1, \dots, N\}$, $\bar{\theta}_t^i = \Psi_t(\theta_0^i)$. In particular,

$$ \bar{\theta}_t^\sigma = (\bar{\theta}_t^{\sigma(1)}, \dots, \bar{\theta}_t^{\sigma(N)}) = (\Psi_t(\theta_0^{\sigma(1)}), \dots, \Psi_t(\theta_0^{\sigma(N)})) = \Psi_t(\theta_0^\sigma), $$

where by abuse of notation we extend $\Psi_t: \mathcal{E} \to \mathcal{E}$ componentwise.

On the other hand, by iterating (10) we may write $\theta_k = \Phi^{(k)}(\theta_0, z_1, \dots, z_k)$, where $\Phi^{(k)}: \mathcal{E} \times \mathbb{Z}^k \to \mathcal{E}$ is defined recursively as $\Phi^{(0)}(\theta) = \theta$ and for $j \in \mathbb{N}$,

$$ \Phi^{(j+1)}(\theta, z_1, \dots, z_{j+1}) = \Phi(\Phi^{(j)}(\theta, z_1, \dots, z_j), z_{j+1}). $$
Therefore, by iterated applications of (19), it follows that for all $\theta \in \mathcal{E}$ and $(z_1, \dots, z_k) \in \mathbb{Z}^k$,

$$ \Phi^{(k)}(\theta^{\sigma}, z_1, \dots, z_k) = \Phi^{(k)}(\theta, z_1, \dots, z_k)^{\sigma}. $$

Hence the maps $\Psi_t$ and $\Phi^{(k)}$ are also permutation-equivariant. Using exchangeability of the initial parameter vector $\theta_0$ and independence between $\theta_0$ and the datapoints $(z_1, \dots, z_k)$, we thus have

$$ (\theta_k, \bar{\theta}_t)^\sigma = (\Phi^{(k)}(\theta_0^\sigma, z_1, \dots, z_k), \Psi_t(\theta_0^\sigma)) \stackrel{d}{=} (\Phi^{(k)}(\theta_0, z_1, \dots, z_k), \Psi_t(\theta_0)) = (\theta_k, \bar{\theta}_t) $$

which proves exchangeability of the synchronously coupled particles. $\square$

Using exchangeability and the same synchronous coupling as in Lemma 3, we have

$$ \mathbb{E}[\mathbb{W}_1(\mu_t^N, \bar{\mu}_t^N)] \le \frac{1}{N} \sum_{i=1}^N \mathbb{E}[|\theta_{[Nt]}^i - \bar{\theta}_t^i|] = \mathbb{E}[|\theta_{[Nt]}^1 - \bar{\theta}_t^1|]. \quad (20) $$

Thus exchangeability reduces the $\mathbb{W}_1$ discrepancy to the expected deviation of a single tagged particle.

*Proof of Proposition 3.* By the above argument, it suffices to bound $\mathbb{E}[|\theta_{[Nt]}^1 - \bar{\theta}_t^1|]$ uniformly in $t \ge 0$. For $t \ge 0$, set $k = \lfloor Nt \rfloor$ and $t_k = k/N$ (so that $t \in [t_k, t_{k+1})$). Define for $i \in \{1, \dots, N\}$,

$$ \delta_k^i = \theta_k^i - \bar{\theta}_{t_k}^i, \quad \delta_k = \delta_k^1, \quad D_t = \bar{\theta}_{t_k}^1 - \bar{\theta}_t^1, $$

so that

$$ |\theta_{[Nt]}^1 - \bar{\theta}_t^1| \le |\delta_k| + |D_t|. \quad (21) $$

By (16), (A1) and (A2), for all $t \ge 0$,

$$ \mathbb{E}[|D_t|] \le \alpha \int_{t_k}^{t} \mathbb{E}[|G_{\lambda}(\bar{\theta}_s^1, \bar{\mu}_s)|] ds \le \gamma((A+B)M + \lambda \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|]) $$

where we have used in the last inequality $t - t_k \le 1/N$ and $\gamma = \alpha/N$. By Lemma 2, under the assumption (A5), we have $\sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|] \le 1 + \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|^2] < \infty$. Hence, the discretization error $D_t$ is uniformly bounded as

$$ \sup_{t \ge 0} \mathbb{E}[|D_t|] \le \frac{C}{N}. \quad (22) $$

It remains to control $\mathbb{E}[|\delta_k|]$ uniformly over $k \in \mathbb{N}$. By (16), we have for all $k \in \mathbb{N}$,

$$ \bar{\theta}_{t_{k+1}}^1 = \bar{\theta}_{t_k}^1 + \alpha \int_{t_k}^{t_{k+1}} G_{\lambda}(\bar{\theta}_s^1, \bar{\mu}_s) ds, $$

hence, adding and subtracting $\gamma G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k})$,

$$ \delta_{k+1} = \delta_k + \gamma (F_{\lambda}^{1}(\theta_k, z_{k+1}) - G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k})) + E_{k+1} \quad (23) $$

where

$$ E_{k+1} = \gamma G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k}) - \alpha \int_{t_k}^{t_{k+1}} G_{\lambda}(\bar{\theta}_s^1, \bar{\mu}_s) ds. $$

We decompose as in [27, Lemma 7.2]

$$ \begin{aligned} F_{\lambda}^{1}(\theta_k, z_{k+1}) - G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k}) &= F_{\lambda}^{1}(\theta_k, z_{k+1}) - G_{\lambda}(\theta_k^1, \mu_{t_k}^N) \\ &\quad + G_{\lambda}(\theta_k^1, \mu_{t_k}^N) - G_{\lambda}(\theta_k^1, \bar{\mu}_{t_k}) \\ &\quad + G_{\lambda}(\theta_k^1, \bar{\mu}_{t_k}) - G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k}) \\ &\equiv A_k + B_k + C_k. \end{aligned} $$
Set $m_k = \mathbb{E}[|\delta_k|^2]$. Note that $m_0 = 0$ since $\delta_0 = 0$ under the synchronous coupling. Using the recursion for $\delta_{k+1}$ (23) and the above decomposition, expanding $|\delta_{k+1}|^2$ and taking expectations yields

$$m_{k+1} \le m_k + \gamma^2 \mathbb{E}[|A_k + B_k + C_k|^2] + 2\gamma \mathbb{E}[\delta_k \cdot (A_k + B_k + C_k)] \\ + \mathbb{E}[|E_{k+1}|^2] + 2\mathbb{E}[\delta_k \cdot E_{k+1}] + 2\gamma \mathbb{E}[(A_k + B_k + C_k) \cdot E_{k+1}].$$

We study each term on the right-hand side in the following technical lemma.

**Lemma 4.** Assume (A1)-(A3) and (A5). Then the following hold for all $k \in \mathbb{N}$:

(i) Let $(\mathcal{F}_k)_{k \in \mathbb{N}}$ be the filtration given by $\mathcal{F}_k = \sigma(\theta_0, z_1, \dots, z_k)$. Then $\mathbb{E}[A_k | \mathcal{F}_k] = 0$ and $\mathbb{E}[|A_k|^2] \le C$.

(ii) $\mathbb{E}[\delta_k \cdot B_k] \le M^2 m_k + CN^{-1}$, and $\mathbb{E}[|B_k|^2] \le 2M^4 m_k + 2M^2 B^2 N^{-1}$.

(iii) $\mathbb{E}[\delta_k \cdot C_k] \le -(\lambda - (A+B)L_\theta)m_k$, and $\mathbb{E}[|C_k|^2] \le (\lambda + (A+B)L_\theta)^2 m_k$.

(iv) $\mathbb{E}[|E_{k+1}|] \le CN^{-2}$, $\mathbb{E}[|E_{k+1}|^2] \le CN^{-4}$, and $\mathbb{E}[\delta_k \cdot E_{k+1}] \le \frac{1}{4}\gamma(\lambda - \lambda_\star)m_k + C\gamma^2$.

The proof of Lemma 4 is postponed to the end of this section. We now conclude.

Since $\delta_k, B_k, C_k, E_{k+1}$ are $\mathcal{F}_k$-measurable and $\mathbb{E}[A_k | \mathcal{F}_k] = 0$ by Lemma 4 (i), all cross-terms involving $A_k$ vanish:

$$\mathbb{E}[\delta_k \cdot A_k] = \mathbb{E}[A_k \cdot B_k] = \mathbb{E}[A_k \cdot C_k] = \mathbb{E}[A_k \cdot E_{k+1}] = 0.$$

Hence, using $2\gamma(B_k + C_k) \cdot E_{k+1} \le \gamma^2|B_k + C_k|^2 + |E_{k+1}|^2$ and $|B_k + C_k|^2 \le 2|B_k|^2 + 2|C_k|^2$, we get

$$m_{k+1} \le m_k + \gamma^2 \mathbb{E}[|A_k|^2] + 4\gamma^2 \mathbb{E}[|B_k|^2 + |C_k|^2] + 2\gamma \mathbb{E}[\delta_k \cdot B_k] + 2\gamma \mathbb{E}[\delta_k \cdot C_k] \\ + 2\mathbb{E}[|E_{k+1}|^2] + 2\mathbb{E}[\delta_k \cdot E_{k+1}].$$

Using the bounds of Lemma 4, we obtain after rearranging terms

$$m_{k+1} \le m_k - \frac{3}{2}\gamma(\lambda - \lambda_\star)m_k + C_\star\gamma^2 m_k + C\gamma^2 + \frac{C}{N^3}$$

with $C_\star = 8M^4 + 4(\lambda + (A+B)L_\theta)^2$ (see (3)). Since $N \ge N_\star$, we have $C_\star\gamma^2 \le \frac{1}{4}\gamma(\lambda - \lambda_\star)$, and hence

$$m_{k+1} \le (1 - \gamma(\lambda - \lambda_\star))m_k + C\gamma^2 + \frac{C}{N^3}.$$

Since $\lambda > \lambda_\star$, a discrete Grönwall argument and $m_0 = 0$ give

$$\sup_{k \in \mathbb{N}} m_k \le \frac{C\gamma^2}{\gamma(\lambda - \lambda_\star)} \le \frac{C}{N} \qquad (24)$$

Finally, combining (22) with (24) in (21) yields

$$\sup_{t \ge 0} \mathbb{E}[|\theta^1_{[Nt]} - \bar{\theta}^1_t|] \le \sup_{k \in \mathbb{N}} \sqrt{m_k} + \frac{C}{N} \le \frac{C}{\sqrt{N}}.$$

Recalling (20), this concludes the proof of Proposition 3. $\square$

**Proof of Lemma 4.** Fix $k \in \mathbb{N}$ and recall $t_k = k/N$, so that $\mu_{t_k}^N$ is the empirical measure of the SGD parameters at step $k$.

(i) By definitions (9) and (14), since $z_{k+1}$ is independent of $\mathcal{F}_k$ and $\theta_k$ is $\mathcal{F}_k$-measurable,

$$\mathbb{E}[F_\lambda^1(\theta_k, z_{k+1}) | \mathcal{F}_k] = G_\lambda(\theta_k^1, \mu_{t_k}^N),$$
hence $\mathbb{E}[A_k | \mathcal{F}_k] = 0$. Moreover, by (A1)-(A2),

$$|A_k| = |(y_{k+1} - \hat{y}_{\theta_k}(x_{k+1}))\nabla_{\theta^1}\sigma_*(\theta_k^1, x_{k+1}) - \mathbb{E}_\pi[(y - \hat{y}_{\theta_k}(x))\nabla_{\theta^1}\sigma_*(\theta_k^1, x)]| \le 2(A + B)M,$$
so that $\mathbb{E}[|A_k|^2] \le 4(A + B)^2 M^2$.

(ii) Write $B_{1,k} = B_{1,k} + B_{2,k}$ with

$$B_{1,k} = G_{\lambda}(\theta_k^1, \mu_{t_k}^N) - G_{\lambda}(\theta_k^1, \bar{\mu}_{t_k}^N) \quad B_{2,k} = G_{\lambda}(\theta_k^1, \bar{\mu}_{t_k}^N) - G_{\lambda}(\theta_k^1, \bar{\mu}_{t_k}).$$

Using $\|\nabla_{(\theta^i, x)}\sigma_*\|_\infty \le M$ and the M-Lipschitzness of the map $\theta^i \in \mathbb{R}^D \mapsto \sigma_*(\theta^i, x)$ from (A1), we have

$$|B_{1,k}| \le M \mathbb{E}_\pi[|\hat{y}_{\theta_k}(x) - \langle \sigma_*(\cdot, x), \bar{\mu}_{t_k}^N \rangle|] \le \frac{M}{N} \sum_{i=1}^N \mathbb{E}_\pi[|\sigma_*(\theta_k^i, x) - \sigma_*(\bar{\theta}_{t_k}^i, x)|] \le \frac{M^2}{N} \sum_{i=1}^N |\delta_k^i|.$$

Hence, by exchangeability (see Lemma 3),

$$\mathbb{E}[|B_{1,k}|^2] \le M^4 \mathbb{E}\left[\left(\frac{1}{N}\sum_{i=1}^N |\delta_k^i|\right)^2\right] \le \frac{M^4}{N}\sum_{i=1}^N \mathbb{E}[|\delta_k^i|^2] = M^4 m_k.$$

Using again $\|\nabla_{(\theta^i, x)}\sigma_*\|_\infty \le M$ and Jensen's inequality, we have

$$|B_{2,k}|^2 \le M^2 \mathbb{E}_\pi\left[\left|\frac{1}{N}\sum_{i=1}^N \sigma_*(\bar{\theta}_{t_k}^i, x) - \mathbb{E}[\sigma_*(\bar{\theta}_{t_k}^1, x)]\right|^2\right] = \frac{M^2}{N}\mathbb{E}_\pi\left[\mathrm{Var}(\sigma_*(\bar{\theta}_{t_k}^1, x))\right] \le \frac{M^2 B^2}{N}.$$

Combining the two bounds for $B_{1,k}$ and $B_{2,k}$, it follows that

$$\mathbb{E}[|B_k|^2] \le 2\mathbb{E}[|B_{1,k}|^2] + 2\mathbb{E}[|B_{2,k}|^2] \le 2M^4 m_k + \frac{2M^2 B^2}{N}.$$

Hence, by the Cauchy-Schwarz inequality,

$$\mathbb{E}[\delta_k \cdot B_k] \le \sqrt{m_k} \sqrt{\mathbb{E}[|B_k|^2]} \le \sqrt{C} m_k + \sqrt{\frac{C}{N}} \sqrt{m_k}.$$

Finally, applying Young's inequality

$$\forall a, b \in \mathbb{R}, \forall \varepsilon > 0, \quad ab \le \frac{\varepsilon}{2}a^2 + \frac{1}{2\varepsilon}b^2 \qquad (25)$$

with $a = \sqrt{m_k}$, $b = \sqrt{C/N}$ and $\varepsilon = 2M^2$, we obtain

$$\mathbb{E}[\delta_k \cdot B_k] \le M^2 m_k + \frac{C}{N}.$$

(iii) By definition (14),

$$C_k = \mathbb{E}_\pi\left[(y - \langle \sigma_*(\cdot, x), \bar{\mu}_{t_k}^N \rangle)(\nabla_{\theta^1}\sigma_*(\theta_k^1, x) - \nabla_{\theta^1}\sigma_*(\bar{\theta}_{t_k}^1, x))\right] - \lambda\delta_k.$$

Using the Cauchy-Schwarz inequality and (A3), we have

$$|\mathbb{E}_\pi[(y - \langle \sigma_*(\cdot, x), \bar{\mu}_{t_k}^N \rangle)(\nabla_{\theta^1}\sigma_*(\theta_k^1, x) - \nabla_{\theta^1}\sigma_*(\bar{\theta}_{t_k}^1, x))]| \le (A + B)L_\theta |\delta_k|.$$

Hence

$$\mathbb{E}[\delta_k \cdot C_k] \le -(\lambda - (A + B)L_\theta) m_k,$$

$$\mathbb{E}[|C_k|^2] \le (\lambda + (A + B)L_\theta)^2 m_k.$$
(iv) Rewrite $E_{k+1}$ as

$$E_{k+1} = \alpha \int_{t_k}^{t_{k+1}} (G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k}) - G_{\lambda}(\bar{\theta}_s^1, \bar{\mu}_s)) ds.$$

Using the Lipschitz bounds

$$\begin{aligned} |G_{\lambda}(\theta, \mu) - G_{\lambda}(\theta', \mu)| &\le (\lambda + (A + B)L_{\theta}) |\theta - \theta'|, \\ |G_{\lambda}(\theta, \mu) - G_{\lambda}(\theta, \nu)| &\le M^2 W_1(\mu, \nu), \end{aligned}$$

for all $\theta, \theta' \in \mathcal{E}$ and $\mu, \nu \in \mathcal{P}(\mathbb{R}^D)$, together with the fact that $s \mapsto \bar{\theta}_s^1$ is Lipschitz in $L^1$ and $s \mapsto \bar{\mu}_s$ is Lipschitz in $W_1$ (with constants controlled by $(A + B)M + \lambda \sup_{t \ge 0} \mathbb{E}[|\bar{\theta}_t^1|] < \infty$ under (A5)), we obtain

$$\mathbb{E}[|G_{\lambda}(\bar{\theta}_{t_k}^1, \bar{\mu}_{t_k}) - G_{\lambda}(\bar{\theta}_s^1, \bar{\mu}_s)|] \le C |s - t_k|.$$

Therefore,

$$\mathbb{E}[|E_{k+1}|] \le \alpha \int_{t_k}^{t_{k+1}} C(s - t_k) ds \le \frac{C}{N^2},$$

and similarly,

$$\mathbb{E}[|E_{k+1}|^2] \le \frac{C}{N^4}.$$

Finally, by Young's inequality (25) with $a = |\delta_k|$, $b = |E_{k+1}|$ and $\varepsilon = \frac{1}{2}\gamma(\lambda - \lambda_*)$, it holds

$$\mathbb{E}[\delta_k \cdot E_{k+1}] \le \frac{1}{4}\gamma(\lambda - \lambda_*) m_k + \frac{C}{\gamma(\lambda - \lambda_*)N^4} \le \frac{1}{4}\gamma(\lambda - \lambda_*) m_k + C\gamma^2.$$

This concludes the proof. □

# 4 Extension to unbounded activations

In this section, we extend our previous results to activations that are no longer uniformly bounded. To this end, we drop assumption (A1) and replace it with the following localization assumptions (B1)–(B3). We show that under these conditions (and a suitable choice of $\lambda$), the SGD iterates remain uniformly bounded for all times and thus evolve within a deterministic compact set. Once localization is established, the analysis reduces to the bounded case on that compact set.

## 4.1 Localization assumptions

In what follows, work under the following three assumptions in place of (A1).

(B1) There exist a compact set $K_x \subset \mathbb{R}^d$ and radius $R_0 > 0$ such that

$$\mathrm{supp} \pi \subset K_x \times [-A, A], \quad \mathrm{supp} \mu_0 \subset \bar{B}(0, R_0) \subset \mathbb{R}^D.$$

(B2) There exist $b, c \ge 0$ such that for all $(\theta^i, x) \in \mathbb{R}^D \times K_x$, $|\sigma_*(\theta^i, x)| \le b + c|\theta^i|$.

(B3) The activation function $\sigma_* \in C^1(\mathbb{R}^D \times \mathbb{R}^d)$ is continuously differentiable, with uniformly bounded gradient on $\mathbb{R}^D \times K_x$:

$$\forall (\theta^i, x) \in \mathbb{R}^D \times K_x, \quad |\nabla_{(\theta^i, x)} \sigma_*(\theta^i, x)| \le M.$$

Note that (B1) implies (A2). Moreover, by Hoeffding's inequality, (B1) guarantees sub-Gaussian Laplace bounds for Lipschitz observables, which is equivalent to a $T_1$ inequality. However, this condition alone is not sufficient to yield a $T_2$ inequality in general. The smoothness condition in (B3) is purely technical. In particular, it excludes non-smooth activations such as ReLU; one may alternatively work with smoothed approximations (e.g. softplus) or piecewise-smooth variants (e.g. leaky-ReLU) without affecting the localization argument below.
## 4.2 A uniform localization bound

We will use the following constants:

$$a_{\infty} = R_0 \lor \frac{M(A+b)}{\lambda - Mc}, \quad R_{\infty} = R_0 \lor \frac{M(A+b+ca_{\infty})}{\lambda}.$$

**Lemma 5.** Assume (B1)–(B3), and $\lambda > Mc$, $\gamma\lambda \le 1$. Then, almost surely,

$$\sup_{k \in \mathbb{N}} \max_{1 \le i \le N} |\theta_k^i| \le R_{\infty}.$$

*Proof.* By (1) and (B1)–(B3), for all $k \in \mathbb{N}$ and $i \in \{1, \dots, N\}$, we have

$$|\theta_{k+1}^i| \le (1 - \gamma\lambda)|\theta_k^i| + \gamma M |y_{k+1} - \hat{y}_{\theta_k}(x_{k+1})| \le (1 - \gamma\lambda)|\theta_k^i| + \gamma M (A + b + ca_k)$$

where $a_k = \frac{1}{N} \sum_{i=1}^{N} |\theta_k^i|$. Averaging over $i$, we get

$$a_{k+1} \le (1 - \gamma(\lambda - Mc)) a_k + \gamma M(A + b). \quad (26)$$

Under $\gamma\lambda \le 1$ and $\lambda > Mc$, we have $0 \le 1 - \gamma\lambda \le 1$ and $\gamma(\lambda - Mc) \le \gamma\lambda \le 1$, so that $1 - \gamma(\lambda - Mc) \in [0, 1)$. It follows by iteration of (26) that

$$\sup_{k \in \mathbb{N}} a_k \le a_0 \lor \frac{M(A+b)}{\lambda - Mc} \le a_{\infty}$$

where we used $a_0 \le \max_{1 \le i \le N} |\theta_0^i| \le R_0$ from (B1). Plugging this bound back into (26) yields

$$|\theta_{k+1}^i| \le (1 - \gamma\lambda)|\theta_k^i| + \gamma M(A + b + ca_{\infty}).$$

Finally, iterating again this inequality and using (B1), we obtain for all $k \in \mathbb{N}$ and for every $i \in \{1, \dots, N\}$

$$\begin{aligned} |\theta_k^i| & \le (1 - \gamma\lambda)^k |\theta_0^i| + \frac{M(A+b+ca_\infty)}{\lambda} (1 - (1 - \gamma\lambda)^k) \\ & \le R_0 \lor \frac{M(A+b+ca_\infty)}{\lambda} = R_\infty, \end{aligned}$$

which is the claimed bound. $\square$

*Remark.* Here, the condition $\lambda > Mc$ plays the same role as the lower bound $\lambda > \lambda_*$ in the bounded case: it enforces that the ridge regularization dominates the effective interaction induced by the activation growth. The additional requirement $\gamma\lambda \le 1$ is a convenient discrete-time stability condition ensuring $1 - \gamma\lambda \ge 0$, and is automatically satisfied for large $N$ under the scaling $\gamma = \alpha/N$.

## 4.3 Reduction to the bounded case

Lemma 5 shows that, under (B1)–(B3) and sufficiently strong regularization, the SGD iterates remain within the closed ball $\bar{B}(0, R_\infty)$ for all times. Hence, we may work on the compact set $\bar{B}(0, R_\infty) \times K_x$, on which the activation and its gradients are uniformly bounded. To apply the arguments of Section 3, we additionally require a localized Lipschitz property for $\nabla_{\theta^i} \sigma_*$ in place of (A3):

(B4) For every $R > 0$, there exist constants $L_x(R), L_\theta(R) > 0$ such that for all $x, x' \in K_x$ and $\theta^i, \theta'^i \in \bar{B}(0, R) \subset \mathbb{R}^D$,

$$|\nabla_{\theta^i} \sigma_*(\theta^i, x) - \nabla_{\theta^i} \sigma_*(\theta^i, x')| \le L_x(R) |x - x'|,$$

$$|\nabla_{\theta^i} \sigma_*(\theta^i, x) - \nabla_{\theta^i} \sigma_*(\theta'^i, x)| \le L_\theta(R) |\theta^i - \theta'^i|.$$
Assumption (B4) is mild in the present setting: for instance, it holds whenever $\sigma_*$ is $C^2$ with bounded mixed second derivatives on compact sets, which includes standard smooth approximations of non-smooth activations used in practice. The arguments of Section 3 then carry over with only notational changes, replacing the global constants ($B, M, L_x, L_\theta$) by their localized counterparts on $\bar{B}(0, R_\infty) \times K_x$.

## Acknowledgements

A.G is supported by the ANR-23-CE-40003, Conviviality, and has benefited from a government grant managed by the Agence Nationale de la Recherche under the France 2030 investment plan ANR-23-EXMA-0001. P.S is supported by the Projet I-SITE Clermont CAP 2025.

## References

[1] Francis Bach. *Learning Theory from First Principles*. Adaptive computation and machine learning. MIT Press, 2024.

[2] Erhan Bayraktar and Ruoyu Wu. "Graphon particle system: uniform-in-time concentration bounds." *Stochastic Processes and their Applications* 156 (Feb. 2023), pp. 196–225. doi: 10.1016/j.spa.2022.11.008.

[3] Gordon Blower and François Bolley. "Concentration of measure on product spaces with applications to Markov processes." *Studia Mathematica* 175 (2006), pp. 47–72. doi: 10.4064/sm175-1-3.

[4] Sergey G. Bobkov and Friedrich Götze. "Exponential integrability and transportation cost related to logarithmic Sobolev inequalities." *Journal of Functional Analysis* 163.1 (1999), pp. 1–28. doi: 10.1006/jfan.1998.3326.

[5] François Bolley, Arnaud Guillin, and Cédric Villani. "Quantitative concentration inequalities for empirical measures on non-compact spaces." *Probability Theory and Related Fields* 137.3–4 (2007), pp. 541–593. doi: 10.1007/s00440-006-0004-7.

[6] Nicolas Bonneel et al. "Sliced and Radon Wasserstein Barycenters of Measures." *Journal of Mathematical Imaging and Vision* 51.1 (2015), pp. 22–45. doi: 10.1007/s10851-014-0506-3.

[7] Léon Bottou. "Large-Scale Machine Learning with Stochastic Gradient Descent." *Proceedings of COMPSTAT'2010*. Physica-Verlag, 2010, pp. 177–186. doi: 10.1007/978-3-7908-2604-3_16.

[8] Léon Bottou, Frank E. Curtis, and Jorge Nocedal. "Optimization Methods for Large-Scale Machine Learning." *SIAM Review* 60.2 (2018), pp. 223–311. doi: 10.1137/16M1080173.

[9] Shi Chen et al. *Quantitative Clustering in Mean-Field Transformer Models*. 2025. arXiv: 2504.14697 [cs.LG].

[10] Lénaïc Chizat and Francis Bach. "On the global convergence of gradient descent for over-parameterized models using optimal transport." *Proceedings of the 32nd International Conference on Neural Information Processing Systems* (Montréal, Canada). Curran Associates Inc., 2018, pp. 3040–3050.

[11] Valentin De Bortoli et al. "Quantitative Propagation of Chaos for SGD in Wide Neural Networks." *Advances in Neural Information Processing Systems*. Vol. 33. Curran Associates, Inc., 2020, pp. 278–288.

[12] Jérôme Dedecker, Aurélie Fischer, and Bertrand Michel. *Concentration of the empirical measure in Wasserstein distance: bounds involving the covering dimension*. 2026. arXiv: 2601.11133 [math.PR].

[13] Pierre Del Moral and Emmanuel Rio. "Concentration inequalities for mean field particle models." *The Annals of Applied Probability* 21.3 (2011), pp. 1017–1052. doi: 10.1214/10-AAP716.

[14] François Delarue and Alvin Tse. "Uniform in time weak propagation of chaos on the torus." *Annales de l'Institut Henri Poincaré, Probabilités et Statistiques* 61.2 (May 2025), pp. 1021–1074. doi: 10.1214/23-AIHP1451.

[15] Arnaud Descours. "Towards a mathematical understanding of deep neural networks using a mean field analysis." PhD thesis. Université Clermont Auvergne, 2023. url: https://theses.hal.science/tel-04528800.
[16] Arnaud Descours et al. "Law of large numbers and central limit theorem for wide two-layer neural networks: the mini-batch and noisy case." *Journal of Machine Learning Research* 25 (2024), pp. 1-76.

[17] Hacène Djellout, Arnaud Guillin, and Liming Wu. "Transportation cost-information inequalities and applications to random dynamical systems and diffusions." *Annals of Probability* 32.3B (2004), pp. 2702-2732. doi: 10.1214/009117904000000531.

[18] Alain Durmus et al. "An Elementary Approach To Uniform In Time Propagation Of Chaos." *Proceedings of the American Mathematical Society* (2018).

[19] Cong Fang et al. "Modeling from Features: a Mean-field Framework for Over-parameterized Deep Neural Networks." *Proceedings of Thirty Fourth Conference on Learning Theory (COLT)*. Vol. 134. Proceedings of Machine Learning Research. PMLR, 2021, pp. 1887-1936.

[20] Nicolas Fournier and Arnaud Guillin. "On the rate of convergence in Wasserstein distance of the empirical measure." *Probability Theory and Related Fields* 162 (2014), pp. 707-738. doi: 10.1007/s00440-014-0583-7.

[21] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. *Deep Learning*. MIT Press, 2016.

[22] Nathael Gozlan and Christian Léonard. "Transport inequalities. A survey." *Markov Process and Related Fields* 16 (4 2010), pp. 635-736.

[23] Daniel Lacker and Luc Le Flem. "Sharp uniform-in-time propagation of chaos." *Probability Theory and Related Fields* 187.1-2 (2023), pp. 443-480. doi: 10.1007/s00440-023-01192-x.

[24] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. "Deep learning." *Nature* 521.7553 (2015), pp. 436-444. doi: 10.1038/nature14539.

[25] Michel Ledoux. *The Concentration of Measure Phenomenon*. Vol. 89. Mathematical Surveys and Monographs. Providence, RI: AMS, 2001.

[26] Michel Ledoux. *Measure concentration, transportation cost, and functional inequalities*. Summer School on Singular Phenomena and Scaling in Mathematical Models. Bonn, Germany, 2003. url: https://lsp.ups-tlse.fr/Ledoux.

[27] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. "A mean-field view of the landscape of two-layer neural networks." *Proceedings of the National Academy of Sciences* 115.33 (2018), E7665-E7671. doi: 10.1073/pnas.1806579115.

[28] Kimia Nadjahi et al. "Statistical and Topological Properties of Sliced Probability Divergences." *Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, December 6-12, 2020, virtual*. Ed. by Hugo Larochelle et al. 2020.

[29] Kimia Nadjahi et al. "Fast Approximation of the Sliced-Wasserstein Distance Using Concentration of Random Projections." *Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, December 6-14, 2021, virtual*. Ed. by Marc'Aurelio Ranzato et al. 2021, pp. 12411-12424.

[30] Phan-Minh Nguyen. *Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks*. 2019. arXiv: 1902.02880 [stat.ML].

[31] Phan-Minh Nguyen and Huy Tuan Pham. "A rigorous framework for the mean field limit of multilayer neural networks." *Mathematical Statistics and Learning* 6 (2023), pp. 201-357. doi: 10.4171/MSL/42.

[32] Felix Otto and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality." *Journal of Functional Analysis* 173.2 (2000), pp. 361-400. doi: 10.1006/jfan.1999.3557.

[33] Gabriel Peyré. *Optimal Transport for Machine Learners*. 2025. arXiv: 2505.06589 [stat.ML].

[34] Gabriel Peyré and Marco Cuturi. "Computational Optimal Transport." *Foundations and Trends in Machine Learning* 11.5-6 (2019), pp. 355-607. doi: 10.1561/2200000073.

[35] Philippe Rigollet. *The Mean-Field Dynamics of Transformers*. 2025. arXiv: 2512.01868 [cs.LG].

[36] Grant M. Rotskoff and Eric Vanden-Eijnden. "Trainability and Accuracy of Neural Networks: An Interacting Particle System Approach." *Communications on Pure and Applied Mathematics* 75.9 (2022), pp. 1889-1935. doi: 10.1002/cpa.22074.

[37] David Saad and Sara A. Solla. "Dynamics of On-Line Gradient Descent Learning for Multilayer Neural Networks." *Advances in Neural Information Processing Systems*. 1995.
[38] Filippo Santambrogio. *Optimal Transport for Applied Mathematicians. Calculus of Variations, PDEs, and Modeling*. Vol. 87. Progress in Nonlinear Differential Equations and Their Applications. Cham: Birkhäuser, 2015. doi: 10.1007/978-3-319-20828-2.

[39] Jürgen Schmidhuber. "Deep learning in neural networks: An overview." *Neural Networks* 61 (2015), pp. 85–117. doi: 10.1016/j.neunet.2014.09.003.

[40] Shai Shalev-Shwartz and Shai Ben-David. *Understanding Machine Learning: From Theory to Algorithms*. Cambridge University Press, 2014.

[41] Justin Sirignano and Konstantinos Spiliopoulos. "Mean Field Analysis of Neural Networks: A Central Limit Theorem." *Stochastic Processes and their Applications* 130.3 (2020), pp. 1820–1852. doi: 10.1016/j.spa.2019.06.003.

[42] Justin Sirignano and Konstantinos Spiliopoulos. "Mean Field Analysis of Neural Networks: A Law of Large Numbers." *SIAM Journal on Applied Mathematics* 80.2 (2020), pp. 725–752. doi: 10.1137/18M1192184.

[43] Justin Sirignano and Konstantinos Spiliopoulos. "Mean Field Analysis of Deep Neural Networks." *Mathematics of Operations Research* 47.1 (Feb. 2022), pp. 120–152. doi: 10.1287/moor.2020.1118.

[44] Namjoon Suh and Guang Cheng. "A Survey on Statistical Theory of Deep Learning: Approximation, Training Dynamics, and Generative Models." *Annual Review of Statistics and Its Application* 12 (2025), pp. 177–207.

[45] Alain-Sol Sznitman. "Topics in propagation of chaos." *Ecole d'Été de Probabilités de Saint-Flour XIX – 1989*. Lecture Notes in Mathematics. Berlin, Heidelberg: Springer, 1991, pp. 165–251.

[46] Cédric Villani. *Optimal Transport. Old and New*. Grundlehren der mathematischen Wissenschaften. Berlin, Heidelberg: Springer, 2008. doi: 10.1007/978-3-540-71050-9.

[47] James Vuckovic. *Propagation of Chaos for Nonlinear Markov Chains*. 2026. arXiv: 2602.07537 [math.PR].