Mean-field limits of two-layer networks

Convergence theorems. Along the noiseless PDE $R(\rho_t)$ is nonincreasing, and $\rho$ is a fixed point iff $\mathrm{supp}(\rho)\subseteq\{\nabla_\theta\Psi(\cdot;\rho)=0\}$, a larger set than the minimizers. Theorem 4: with noise, under $C^4$ regularity and $\lambda$ bounded above and below, $F_{\beta,\lambda}$ has a unique minimizer, the Gibbs fixed point $\rho_*\propto\exp\{-\beta\Psi_\lambda(\theta;\rho_*)\}$; the flow converges to it from any $\rho_0$ of finite free energy, and $R(\rho_*^{\beta,\lambda})\le\inf_\rho R_\lambda(\rho)+K(1+\lambda)D\log(2+1/\lambda)/\beta$ with $R_\lambda=R+\lambda\int\|\theta\|^2d\rho$. Theorem 5 combines this with Theorem 3: for $\beta\ge KD$ there is $T\lt\infty$, independent of $N$, such that noisy SGD with $N,1/\varepsilon\ge C_0e^{C_0T}D$ has $R_N(\theta^k)\le\inf_\rho R_\lambda(\rho)+\eta$ for $k\in[T/\varepsilon,10T/\varepsilon]$ with probability $1-\delta$. Theorems 6 and 7 (noiseless stability): $\delta_{\theta_*}$ is a fixed point iff $\nabla V(\theta_*)+\nabla_1U(\theta_*,\theta_*)=0$; if $H_0=\nabla^2_\theta\Psi(\theta_*;\delta_{\theta_*})\succ0$, any $\rho_{t_0}$ supported in a small ball around $\theta_*$ converges to it exponentially. Conversely, if $\rho_*=p_*\delta_{\theta_*}+(1-p_*)\tilde\rho_*$ is a fixed point at which $\nabla^2\Psi(\theta_*;\rho_*)$ is nonsingular with a negative eigenvalue and $\tilde\rho_*$ sits strictly below the level $\Psi(\theta_*;\rho_*)$ (plus a compactness condition on level sets), no $\rho_0$ with bounded density converges to $\rho_*$.

The static problem is convex, and optimality is a condition on the potential. $U\succeq0$ makes $R$ convex in $\rho$, $R_N$ is $R$ at the empirical measure, and sampling neurons i.i.d. from $\rho_*$ loses only the diagonal of the double sum, $K/N$. For optimality expand $$R(\rho)-R(\rho_*)=2\langle\Psi(\cdot;\rho_*),\rho-\rho_*\rangle+\langle U,(\rho-\rho_*)^{\otimes2}\rangle.$$ Moving mass $t$ from near a support point where $\Psi$ is higher by $\Delta$ to $\arg\min\Psi$ changes the risk by $-2\Delta t+O(t^2)\lt0$; conversely the support condition makes the linear term nonnegative, and the quadratic term is nonnegative by positive semidefiniteness. The gap between stationarity ($\nabla\Psi=0$ on the support) and optimality ($\Psi$ minimal on the support) is the whole difficulty of the dynamics.

SGD stays close to i.i.d. copies of a McKean-Vlasov process (propagation of chaos). Couple each neuron to the nonlinear dynamics started at the same point, $$\bar\theta_i^t=\theta_i^0-2\int_0^t\xi(s)\nabla\Psi(\bar\theta_i^s;\rho_s)\,ds,\qquad \rho_s=\mathrm{Law}(\bar\theta_i^s),$$ which exists and is unique by Sznitman's theorem and makes the $\bar\theta_i$ i.i.d. with law solving the PDE. The gap comes from time discretization (order $\varepsilon$), a martingale (order $\sqrt\varepsilon(\sqrt D+\sqrt{\log N}+z)$, by Azuma-Hoeffding) and the empirical measure in the drift, which costs the average coupling distance over the other neurons plus an i.i.d. fluctuation of order $\sqrt{(D+\log N)/N}$, since $\nabla_1U$ is bounded. Grönwall on the maximal deviation closes the argument and is where $e^{CT}$ comes from.

The PDE descends the risk, and its stationary points carry no current. By the weak form of the PDE, $$\frac{d}{dt}R(\rho_t)=-4\xi(t)\int\|\nabla_\theta\Psi(\theta;\rho_t)\|^2\rho_t(d\theta),$$ so the risk is monotone and a fixed point is a measure supported on critical points of its own potential. If $\Psi(\cdot;\rho)$ has a deeper valley where $\rho$ has no mass, transport cannot reach it, and $\rho$ is stationary without being optimal.

With noise the free energy is strictly convex, so the flow has a unique limit. Entropy makes $F_{\beta,\lambda}$ strictly convex, with the Gibbs equation as first-order condition, and the diffusion PDE dissipates it, $$\frac{d}{dt}F_{\beta,\lambda}(\rho_t)=-2\xi(t)\int\big\|\nabla_\theta\big(\Psi_\lambda(\theta;\rho_t)+\beta^{-1}\log\rho_t(\theta)\big)\big\|^2\rho_t(d\theta),$$ so the dissipation vanishes along a sequence, tightness gives a limit point with a density, and there $\beta\Psi_\lambda+\log\rho$ is constant, which makes it the unique Gibbs measure. The temperature cost comes from smoothing: convolving a near-optimal $\rho$ with $N(0,\tau^2I)$ raises $R$ by $O(\tau^2D)$ while its entropy is at least $\frac D2\log(2\pi e\tau^2)$, and at $\tau=1/\beta$ the $\log\beta$ terms cancel against the entropy upper bound, leaving $O(D\log(1/\lambda)/\beta)$. No rate in $t$ comes out; the relevant log-Sobolev constant can be exponentially small in $D$, suggesting $T=e^{O(D)}\log(1/\eta)$.

Without noise, a stable atom attracts, and a saddle cannot be reached from a bounded density. For stability, test against $\varphi_2=\frac12\|\theta-\theta_*\|^2$: for $\rho$ supported in a small ball the positive Hessian gives $\int\langle\theta-\theta_*,\nabla\Psi(\theta;\rho)\rangle d\rho\ge\frac\lambda2\int\|\theta-\theta_*\|^2d\rho$, so $\langle\varphi_2,\rho_t\rangle$ decays exponentially, and a second test function keeps the support inside the ball. At a nondegenerate saddle the characteristics expand along the negative eigendirection, so mass converging to $\theta_*$ must start near a codimension-one stable set, and a bounded density puts too little mass in a thin slab around it to deliver the atom $p_*$. This is why the examples assume bounded densities: they exclude initializations sitting exactly on saddles.

In the Gaussian example symmetry reduces the PDE to one dimension, and at $d=\infty$ the risk depends on two numbers. The PDE preserves rotation invariance (which no finite set of neurons has), so $\rho_t$ is determined by the law of $r=\|w\|$. As $d\to\infty$ independent uniform directions are orthogonal, so each class output concentrates at $\langle q_\pm,\rho\rangle$, $q_\pm(r)=\mathbb E\sigma((1\pm\Delta)rG)$, and the risk is a sum of squares of the class residuals $\lambda_\pm$. The radial force factors, $$\partial_r\psi_\infty(r;\rho)=\lambda_+q_+'(r)+\lambda_-q_-'(r)=\lambda_-q_+'(r)\Big[\frac{\lambda_+}{\lambda_-}+Z(r)\Big],\qquad Z=q_-'/q_+'\ \text{increasing},$$ so it changes sign at most once, and the signs of $(\lambda_+,\lambda_-)$ determine the outcome: equal signs leave only stationary points with risk at least 1, never reached from below 1; one mixed sign traps mass at $\delta_{r_*}$, the global minimizer when $\Delta\lt\Delta_\infty$; the other puts $\delta_{r_*}$ at a maximum of the potential, which Step 5 rules out. For $\Delta\ge\Delta_\infty$ the limit has zero risk, and finite $d$ follows by continuity of the PDE in $(V,U)$ and Theorem 3.

Homogeneity plus spread initialization turns stationary limits into global minima. A limit $\nu$ has $F'(\nu)=0$ on its support; optimality needs $F'(\nu)\ge0$ everywhere. If $F'(\nu)$ is negative somewhere, take a negative regular value $-\eta$ of its restriction to the sphere: near $\nu$ the cone over that sublevel set is invariant and, by 2-homogeneity, the radial velocity there is at least $\eta r$, so $|u_t|\ge|u_0|e^{\eta t}$ and any mass in the cone grows exponentially and pushes the flow out of every neighborhood of $\nu$. The support condition keeps mass in every cone: by a topological degree argument, if $\mathrm{supp}(\mu_0)$ separates two spheres then so does $\mathrm{supp}(\mu_t)$ (in the 1-homogeneous case the sets separated are $w=\pm r_0$, and the mass in the cone grows at least linearly instead of exponentially). So a convergent flow cannot converge to a non-optimal point; the result gives no rate but needs no noise. In the experiments particle gradient descent needs only slightly more particles than the sparsest minimizer (for a 4-neuron ReLU teacher in $d=2$, 5 fail and 10 succeed).