Source-linked AI summary
A Mean Field View of the Landscape of Two-Layers Neural Networks
Song Mei, Andrea Montanari, Phan-Minh Nguyen
TL;DR
In large-scale applications, each example may be visited only a few times, and theoretical guarantees suggest limited advantage from multiple passes. The paper uses a propagation-of-chaos scaling analysis to approximate two-layer-network SGD by a Wasserstein gradient-flow PDE over parameter distributions. The PDE analysis proves convergence to near-global optima for several examples, while noisy SGD converges generically to a near-global minimum of regularized risk.
Problem
In large-scale applications, each example may be visited only a few times, and theoretical guarantees suggest limited advantage from multiple passes.
Source: §1 · p. 3Method
The paper uses a propagation-of-chaos scaling analysis to approximate two-layer-network SGD by a Wasserstein gradient-flow PDE over parameter distributions.
Source: §1 · p. 4, §1.1 · p. 5, §4 · p. 15Results
The PDE analysis proves convergence to near-global optima for several examples, while noisy SGD converges generically to a near-global minimum of regularized risk.
Source: §4 · p. 15Takeaways & Limitations
When N ≫D, SGD behavior is effectively independent of hidden-unit count, and distributional dynamics offers a mechanism for avoiding some non-convexity pitfalls.
Source: §4 · p. 15The convergence time can depend on dimension D and data distribution P, with the analysis suggesting generally T = eO(D).
Source: §3.1 · p. 14
Abstract
from arXiv · showhide
Multi-layer neural networks are among the most powerful models in machine learning, yet the fundamental reasons for this success defy mathematical understanding. Learning a neural network requires to optimize a non-convex high-dimensional objective (risk function), a problem which is usually attacked using stochastic gradient descent (SGD). Does SGD converge to a global optimum of the risk or only to a local optimum? In the first case, does this happen because local minima are absent, or because SGD somehow avoids them? In the second, why do local minima reached by SGD have good generalization properties? In this paper we consider a simple case, namely two-layers neural networks, and prove that -in a suitable scaling limit- SGD dynamics is captured by a certain non-linear partial differential equation (PDE) that we call distributional dynamics (DD). We then consider several specific examples, and show how DD can be used to prove convergence of SGD to networks with nearly ideal generalization error. This description allows to 'average-out' some of the complexities of the landscape of neural networks, and can be used to prove a general convergence result for noisy SGD.
1 Introduction
The paper studies supervised learning in two-layer neural networks, where SGD optimizes a non-convex population risk. It develops a scaling-limit description by distributional dynamics, a Wasserstein gradient flow that enables landscape analysis and convergence results.
1 Introduction
- Two-layer networks model labels from feature vectors using hidden-unit parameters chosen to minimize population risk, with square loss as the main setting.The training data are i.i.d. from an unknown distribution P, and the network has N hidden units with parameters θ_i.Source: §1 · p. 3
- SGD is analyzed under a one-pass assumption in which each training example is an independent draw and is never revisited.This assumption is motivated by large-scale applications where examples are visited at most a few times and by theoretical evidence that repeated passes offer limited advantage.Source: §1 · p. 3
- The paper’s central contribution is a nonlinear PDE showing that, in a suitable scaling limit, SGD becomes distributional dynamics over parameter probability measures.The PDE is a gradient flow of the asymptotic population risk R(ρ) in the Wasserstein space (P(R^D), W2), supporting landscape analysis and convergence results for noisy SGD.Source: §1 · p. 4
1.1 An informal overview
- In the infinite-width formulation, risk depends on the empirical parameter distribution, becomes convex under a positive-semidefinite kernel, yet dynamics still preserve local mass and need not reach a unique global optimum.The finite-width and measure-valued risks differ by O(1/N) under mild assumptions.Source: §1.1 · p. 4
- Distributional dynamics approximates empirical SGD when ε ≪1/D and N ≫D, while factoring out neuron permutations and enabling symmetry reductions inherited from the data distribution.The continuum description enforces local mass conservation because parameter mass moves continuously rather than jumping between disconnected regions.Source: §1.1 · p. 5
- If distributional dynamics reaches an optimum in time t*(D)=O(D), the required sample count is O(D) and the resulting population risk is independent of N; noisy SGD generically converges globally to an entropy-regularized minimizer.The noisy-SGD limit adds diffusion and converges to the minimizer of Fβ,λ(ρ) for finite β, with a number of steps independent of N.Source: §1.1 · p. 6
2 Examples
The examples illustrate how distributional dynamics analyzes SGD in simplified Gaussian classification problems, proving near-optimal convergence under suitable conditions and predicting failures when assumptions are violated. They also show close agreement between reduced PDE predictions and empirical SGD behavior.
2.1 Centered isotropic Gaussians
- For centered isotropic Gaussians, SGD succeeds despite the classification problem requiring multilayer features, and its empirical behavior closely matches a one-dimensional reduced PDE.Symmetry keeps the weight distribution spherically symmetric, reducing dynamics to the norm distribution; the example is deliberately non-realistic and uses simple data and activations.Source: §2 · p. 6, §2.1 · p. 7
- For sufficiently large dimension, neuron count, and suitable initialization and step size, isotropic-Gaussian SGD reaches near-optimal risk with high probability.With ε = 1/(C0d), the required number of SGD steps is linear in d and independent of the number of hidden units N.Source: §2.1 · p. 9
2.2 Centered anisotropic Gaussians
- For centered anisotropic Gaussians, reduced-symmetry dynamics learns the relevant subspace and converges to nearly optimal risk using k = O(d) samples, independent of N.The reduced distribution tracks the norms of weight components inside and outside the informative subspace.Source: §2.2 · p. 9
2.3 A better activation function
- With standard ReLU units that fit output weights and offsets, SGD reaches very small risk, although convergence has a nontrivial structure with long flat regions.The empirical trajectories are well captured by a four-dimensional reduced PDE for output weight, offset, and relevant and irrelevant weight norms.Source: §2.3 · p. 10
2.4 Predicting failure
- When the activation is non-monotone and violates the isotropic theorem’s assumptions, initialization determines whether SGD converges to a low-risk or high-risk limit.A one-dimensional PDE tracks both outcomes, demonstrating that SGD need not always approach a global optimum.Source: §2.4 · p. 11
3 General results
The section establishes that distributional dynamics approximates large-width SGD, while characterizing risk descent and fixed points. Diffusion dynamics has broad convergence guarantees, whereas noiseless dynamics requires stability conditions and can fail to converge to certain fixed points.
3 General results
- The asymptotic risk minimum approximates the finite-N risk minimum, and the noiseless PDE makes risk non-increasing while its fixed points generally include more than global optimizers.Global risk optimizers are fixed points, but fixed points need not be global minimizers.Source: §3 · pp. 12–13
- As N→∞ and ε→0 under explicit scaling conditions, SGD converges almost surely to the distributional-dynamics PDE, with analogous bounds for noisy SGD.The approximation error depends benignly on N and D but grows exponentially with the time horizon, limiting applicability when dynamics converge slowly or are unstable.Source: §3 · pp. 12–13
- Under A1–A3, diffusion dynamics preserves absolute continuity and has at most one finite-free-energy fixed point, characterized by its stationary density.The convergence analysis relies on regularity assumptions and an initial finite free energy.Source: §3 · p. 13
3.1 Convergence: noisy SGD
- With A1–A4 and 1/K0 ≤ λ ≤ K0, the free energy has a unique minimizer, and diffusion dynamics converges to it from any initialization with finite free energy.The proof uses free-energy dissipation, whose limiting stationary condition identifies the global minimizer by convexity.Source: §3.1 · p. 14
- Regularized noisy SGD reaches arbitrarily small error with high probability after finite time when β is sufficiently large and N and 1/ε are exponentially large in D.The convergence time is independent of N but may depend on D and the data distribution, with generally suggested scaling T = eO(D).Source: §3.1 · p. 14
3.2 Convergence: noiseless SGD
- For noiseless dynamics, a point mass δθ∗ is fixed exactly when ∇V(θ∗) + ∇1U(θ∗,θ∗) = 0; positive-definite Hessian yields locally exponentially fast convergence.Convergence holds when the initial distribution is supported in a sufficiently small neighborhood of θ∗.Source: §3.2 · p. 15
- Conversely, under the stated level-set and Hessian conditions, a fixed point containing a point mass with a negative Hessian eigenvalue cannot attract an initially bounded-density distribution.This supplies a stability-based obstruction for characterizing possible long-time limits of noiseless dynamics.Source: §3.2 · p. 15
4 Discussion and future directions
The paper shows that, when hidden units greatly exceed input dimension, SGD and noisy SGD are captured by Wasserstein gradient-flow PDEs whose dynamics converge toward near-global optima. It also identifies extensions beyond square loss and future directions involving optimal transport and glass theory.
Main contributions
- When N ≫ D, SGD and noisy SGD are approximated by PDEs (7) and (12), respectively, both corresponding to Wasserstein gradient flows.This clarifies that SGD behavior, like empirical risk minimization, becomes independent of hidden-unit count once overparameterization is sufficiently large.Source: §4 · p. 15
- The PDE (7) converges to a near-global optimum from bounded-density initialization, avoiding the need to show unique or uniformly good finite-N local minima.It can eliminate exceptional stationary points and merge multiple finite-N stationary points into similar distributions ρ.Source: §4 · p. 15
- Noisy SGD converges generically to a near-global minimum of the regularized risk in time independent of the number of hidden units.Source: §4 · p. 15
Scope and limitations
- The approach, developed for square loss, should generalize to other loss functions.The paper refers to the supplementary information for this extension.Source: §4 · p. 15
Future directions
- Optimal-transport tools may prove convergence for the Wasserstein gradient flow, while glass-theory ideas may clarify how multiple finite-N minima correspond to one limiting minimizer.Source: §4 · pp. 15–16
Related work
- Several groups later obtained closely related results to Theorem 3 [RVE18] [SS18] [CB18].Source: §4 · p. 16
Supplementary information
The supplementary information provides proofs and additional technical details for the mathematical results, together with further information about the numerical experiments.
- It contains proofs and technical details for the mathematical results, plus additional information about the numerical experiments.Source: Supplementary information · p. 19
5 Notations
This section establishes notation for vectors, matrices, scalars, probability measures, geometric sets, integrals, derivatives, Lipschitz constants, probability-measure distances, generic constants, and natural numbers.
- Vectors, matrices, and scalars are written in lowercase bold, uppercase bold, and lowercase plain type, respectively.Source: §5 · p. 19
- The notation defines probability measures P(Ω), Euclidean balls Bd(x; r), measure-function integrals ⟨f, µ⟩, derivatives, and Lipschitz constants.Source: §5 · p. 19
- The bounded Lipschitz distance dBL and Wasserstein distance Wp denote distances between probability measures.Source: §5 · p. 19
- K denotes a generic constant depending on K0, K1, K2, and K3, whose values are specified contextually, while N denotes the natural numbers.Source: §5 · p. 19
6 General results: Statics
The section establishes convexity and finiteness properties of the continuum population risk, relates finite-network and continuum minimization, and characterizes its minimizers. Under analyticity, minimizers have highly concentrated support, with a sharper discrete description in one dimension.
6 General results: Statics
- The continuum risk R(ρ) is convex over probability measures, with R(ρ)=∞ outside the class where its defining expectations are finite.The required expectations are finite under a finite second moment of x when the feature is quadratically bounded.Source: §6 · p. 20
6.1 Proof of Proposition 1
- Minimizing the population risk R_N(θ) yields results comparable to minimizing the continuum risk R(ρ), while a stated condition characterizes continuum-risk minimizers.The sufficiency argument gives R(ρ) ≥ R(ρ∗) + ⟨U, (ρ − ρ∗)⊗2⟩ ≥ R(ρ∗).Source: §6.1 · pp. 20, 22
6.2 Some additional results
- Under the stated analyticity assumptions, every minimizer ρ∗ either has support of zero Lebesgue measure or satisfies the alternative conclusion of Corollary 1.The support result follows because supp(ρ∗) lies in the zero set of a non-trivial analytic function.Source: §6.2 · p. 22
- For D = 1, the zero-measure support conclusion strengthens to a countable convex combination of point masses without accumulation, finite when Ψ(θ; ρ∗) diverges as |θ| grows.The no-accumulation property follows from the zero-set structure of analytic functions in one dimension.Source: §6.2 · p. 22
7 General results: Dynamics
This section develops the distributional dynamics PDE, establishing its well-posedness, regularity, convergence from finite-width and noisy SGD, and key properties of its solutions. It also analyzes risk, stability, measure preservation, generalization, and trajectory behavior near unstable critical points.
7 General results: Dynamics
- Under assumptions A1 and A3, the limiting PDE and its equivalent nonlinear trajectory dynamics have unique solutions, with ρ_t Lipschitz in W2 Wasserstein distance.The PDE is interpreted as an evolution of probability measures in weak sense, and the trajectory law satisfies the same equation.Source: §7 · p. 23
7.1 Proof of Theorem 3: Convergence to the PDE
- Theorem 3 is proved by a propagation-of-chaos argument that controls the discrepancy between SGD and nonlinear dynamics with high-probability estimates under A1–A3.The proof combines Lipschitz bounds, martingale concentration via Azuma–Hoeffding, union bounds, and Gronwall’s inequality; convergence holds when N →∞ and ε_N →0 with N/log(N/ε_N) →∞ and ε_N log(N/ε_N) →0.Source: §7.1 · pp. 24, 25, 27
7.2 Proof of Theorem 3: Generalization to β < ∞
- For finite temperature 1 ≤ β < ∞, the limiting PDE remains well posed in weak sense and is equivalently represented by nonlinear stochastic dynamics driven by independent Brownian motions.The assumptions on U, V, λ, and ξ ensure a unique continuous stochastic-dynamics solution, while existence and uniqueness of the PDE follow from a modified argument of [Szn91, Theorem 1.1].Source: §7.2 · p. 28
- The noisy-SGD extension controls the difference between discrete noisy SGD and the nonlinear stochastic dynamics with a high-probability bound under A1–A3 and finite-temperature assumptions.The argument follows the noiseless proof while adding Gaussian-noise estimates and using Gaussian concentration and Doob’s inequality.Source: §7.2 · pp. 30–31
- The β < ∞ generalization proof controls the approximation terms using bounded-Lipschitz estimates and concludes via Gronwall’s inequality.The resulting argument extends Theorem 3 to finite β under the cited assumptions.Source: §7.2 · p. 32
7.3 Proof of Proposition 2: Monotonicity of the risk
- The PDE risk R(ρ_t) is non-increasing, and a fixed-point distribution is characterized by support on parameters where ∇Ψ(θ; ρ)=0.Conversely, a distribution supported on this zero-gradient set remains fixed under the dynamics.Source: §7.3 · p. 33
7.4 A general continuity result
- Solutions of the two PDEs remain close with distance bounded by K exp(Kt) times the initial bounded-Lipschitz distance plus the perturbation ε0.The continuity estimate assumes bounded-Lipschitz gradients for both pairs of potentials, with K independent of D.Source: §7.4 · pp. 33–34
7.5 Some properties of the solution of the PDE (7.1)
- The PDE transports ρ0 through the ODE flow, preserves mass in invariant sets and the positive orthant, has a connected compact limit set, and preserves density with bound M_t = K M_0 exp{KDt}.The positive-orthant preservation requires the additional condition stated in Lemma 7.9, while the density result assumes twice continuously differentiable U and V.Source: §7.5 · pp. 35–36
7.6 Proof of Theorems 6: Stability conditions
- Under the stated stability conditions, a distribution initially concentrated near θ∗ remains controlled for all time and satisfies exponential decay of a localized moment.The proof shows the relevant stopping time is infinite and derives a differential inequality yielding exponential decay.Source: §7.6 · p. 39
7.7 Proof of Theorem 7: Instability conditions
- The instability proof begins with a coupling lemma for concentrated absolutely continuous measures, constructing paired samples with controlled directional separation.For one dimension, the construction uses a shifted inverse distribution function; the supplied passage establishes the one-dimensional case.Source: §7.7 · p. 40
- Separation grows at least as (1/Z)e^{λ1(t−t0)/2}, forcing both coupled trajectories to leave B(θ∗; r1) by TUB.The growth follows from the negative Hessian eigenvalue −λ1 and a coupling whose relevant event has probability at least 9/10.Source: §7.7 · pp. 42–43
- Trajectories starting in B(θ∗; r0) exit B(θ∗; r1) within TUB and never return to B(θ∗; r0).This holds for sufficiently small r0 and ε0 while dBL(ρt, ρ∗) and the mass discrepancy near θ∗ remain bounded by ε0.Source: §7.7 · pp. 44, 46
- After exit, the energy remains below the level required to enter the smaller neighborhood, preventing re-entry under the PDE closeness assumptions.The argument uses the Hessian spectrum and Lemma 7.14 to maintain an energy gap after exit.Source: §7.7 · p. 46
- Lemma 7.17 bounds the probability of returning to B(θ∗; r0) after time t0 by an arbitrarily prescribed η under Theorem 7’s conditions.The proof places most mass in a lower sublevel set and uses Lemma 7.14 to keep trajectories there.Source: §7.7 · p. 46
- Combining the exit, no-return, and probability bounds yields a contradiction with maintaining at least p∗−ε0 mass near θ∗, proving Theorem 7.The parameters are chosen with η, ε0 ≤ p∗/10 and sufficiently small neighborhood radii.Source: §7.7 · p. 47
8 Centered isotropic Gaussians
This section analyzes centered isotropic Gaussians by reducing the dynamics to radial probability measures and characterizing their minimizers and fixed points. It establishes convergence results through infinite- and finite-dimensional analyses, approximation and regularity estimates, and verification of the assumptions for a running activation example.
8 Centered isotropic Gaussians
- The analysis assumes a bounded, non-decreasing, Lipschitz activation σ, defines q(r) ≡ E{σ(rG)} for G ∼ N(0, 1), imposes S2–S4, and sets ξ(t) = 1/2 as a time reparametrization.Source: §8 · pp. 47–48
8.1 Statics
- Rotational symmetrization cannot increase the convex risk, so minimization may be restricted to rotation-invariant distributions represented by radial measures.For any distribution ρ, its rotational symmetrization ρs satisfies R(ρs) ≤ R(ρ).Source: §8.1 · pp. 48–49
- For ∆ < ∆∞, the global minimizer is a unique point mass, whereas for ∆ ≥ ∆∞ a zero-risk global minimizer exists with compact support bounded away from zero.Source: §8.1 · pp. 50–51
8.2 Dynamics: Fixed points
- The radial PDE fixed points fall into three types: zero-risk measures, point masses with nonzero risk, and mixtures supported at 0, +∞, and one finite radius.For ∆ < ∆∞, the type-(b) fixed point is unique; at ∆ = ∆∞ there is a unique type-(a) fixed point and no type-(b) fixed points; for ∆ > ∆∞, type-(b) fixed points may be multiple.Source: §8.2 · p. 52, §9.2 · p. 67
8.3 Dynamics: Convergence to global minimum for d = ∞
- The d = ∞ dynamics converges to a global minimizer from any initialization in Pgood.Source: §8.3 · p. 53
Step 1. Existence and uniqueness of solution. Mass
- The proof establishes existence and uniqueness of the PDE solution and shows that limiting sets are fixed points, with nontrivial limits restricted to mixtures of point masses at 0, +∞, and possibly one finite radius.Connectedness and analyticity reduce the limiting set to a singleton of the form a0δ0 + a∞δ∞ + aδr∗ in the relevant case.Source: Step 1. Existence and uniqueness of… · pp. 54, 55, 57
8.3 Dynamics: Convergence to global minimum for d = ∞
- For Δ<Δ∞, the PDE converges to the unique point-mass fixed point that is also the global risk minimizer, whereas for Δ≥Δ∞ it converges to a limiting set with risk 0.These conclusions apply to the d=∞ dynamics described by Theorem 10.Source: Step 1. Existence and uniqueness of… · p. 55
8.4 Proof of Theorem 1
- The finite-dimensional proof reduces convergence to showing that large finite d approximates the d=∞ dynamics, using kernel estimates and Gaussian regularity bounds that verify assumptions A2 and A3.The argument controls the relevant Gaussian quantities uniformly and handles second-derivative bounds even when σ′ may not be differentiable.Source: §8.4 · pp. 57, 59
- Theorem 1 follows after uniformly bounding the derivative terms needed for assumption A3 and controlling the PDE and SGD approximations over the stated finite time horizon.The proof uses Stein’s formula and Gaussian tail bounds, then applies the finite-dimensional and SGD approximation results under conditions relating N, ε, and d.Source: §8.4 · pp. 60–62
8.5 Checking conditions S0–S4 for the running example
- The activation satisfies S0–S4 when s1 < s2, s1 < −1, (s1 + s2)/2 > 1, (3s1 + s2)/4 ∈(−1, 1), and 0 < t1 < t2.The verification establishes analyticity and bounded derivatives for q, positivity and endpoint behavior for q′, and the remaining positivity condition through F2 and F3.Source: §8.5 · pp. 62–64
- A concrete parameter choice satisfying all conditions is s1 = −2.5, s2 = 7.5, t1 = 0.5, and t2 = 1.5.The paper notes that the requirements of the running-example lemma are not restrictive.Source: §8.5 · p. 62
9 Centered anisotropic Gaussians
This section analyzes centered anisotropic Gaussians through radial distributions over relevant and irrelevant subspaces, characterizing global minimizers and PDE fixed points. It establishes mass-support properties, zero-risk behavior in the infinite-dimensional limit, and convergence of finite-dimensional dynamics and SGD to the limiting PDE under stated scaling conditions.
9 Centered anisotropic Gaussians
- The model assumes jointly diagonalizable covariances differing on an s0-dimensional subspace, with s0 = γ · d and both subspace dimensions diverging as d →∞.The anisotropy is parameterized by 0 < ∆<1, with τ+ = 1 + ∆ and τ−=1−∆.Source: §9 · p. 64
9.1 Statics
- By convexity, risk minimization may be restricted to distributions invariant under separate rotations of the relevant and irrelevant subspaces, represented by the pair of norms (∥w1∥2, ∥w2∥2).The induced probability space is P(E2), where E2 = [0, +∞)2 ∪ {∞}.Source: §9.1 · pp. 65–66
- The static analysis identifies global minimizers including a point mass at (r∗, 0), while some global minimizers have finite support and risk zero.The threshold ∆∞∈(0, 1) separates the minimizer regimes described in the theorem.Source: §9.1 · p. 66
9.2 Dynamics: Fixed points
- The reduced PDE’s fixed points fall into three types: zero-risk fixed points, nonzero-risk point masses at (r∗, 0), and mixtures of these types, with their existence depending on ∆ relative to ∆∞.For ∆<∆∞ there is a unique type-(b) fixed point; for ∆=∆∞ there is a unique type-(a) delta fixed point and no type-(b) fixed points; for ∆>∆∞ there are type-(b) fixed points and some type-(a) fixed points.Source: §9.2 · pp. 67–68
Step 1. Existence and uniqueness of solution. Mass
- The convergence proof establishes existence and uniqueness of the PDE solution, preserves all mass in (0, ∞)2 at finite times, and constrains limiting fixed-point sets through connectedness and sign conditions.The limiting-set analysis yields the point-mass global minimizer for ∆<∆∞ and zero-risk limiting sets for ∆≥∆∞.Source: Step 1. Existence and uniqueness of… · pp. 55, 69, 70
9.3 Dynamics: Convergence to global minimum for d = ∞
- For d = ∞, the PDE converges from any initialization to the unique global minimizer when ∆<∆∞, and to a limiting set with risk 0 when ∆≥∆∞.For ∆<∆∞, the limiting minimizer is a point mass at (r∗1, 0); for ∆≥∆∞, every point in the limiting set has risk 0.Source: §9.3 · p. 68, Step 1. Existence and uniqueness of… · p. 70
Step 1. Existence and uniqueness of solution. Mass
- The mass argument shows that trajectories starting in Γk remain there, preventing finite-time escape through the coordinate boundary or infinity.This yields nondecreasing mass in Γk and supports the exclusion of fixed-point configurations contained in P2.Source: Step 1. Existence and uniqueness of… · pp. 70–72
9.4 Dynamics: Proof of Theorem 2
- As d →∞ with k = γ · d, the interaction kernel ud and its gradient converge uniformly to u∞ and its gradient.The proof uses concentration of the Gaussian-angle variables as both s0 and d − s0 diverge.Source: §9.4 · pp. 73–74
- For sufficiently large d, the finite-dimension PDE remains close to the d = ∞ PDE over the bounded time interval t ≤ 10T.The argument combines kernel convergence, initialization convergence in bounded-Lipschitz distance, and the stability estimate from Lemma 7.7.Source: §9.4 · p. 75
- Under the stated sample-size, step-size, dimension, and probability conditions, SGD trajectories approximate the PDE solution with probability at least 1 − δ.The SGD-to-PDE bound is controlled by Ke10KT · errN,d(z), with constants depending on the theorem assumptions.Source: §9.4 · p. 76
10 Finite temperature
The finite-temperature analysis establishes well-posedness and convergence for the free-energy minimizer and associated PDE, including uniqueness, regularity, and global minimization under the stated conditions. It also derives finite-temperature risk and finite-time high-probability guarantees, alongside a heuristic temperature scaling.
10.1 Statics
- Under bounded-Lipschitz U and V, Fβ,λ has a unique minimizer in K satisfying the Boltzmann fixed point condition, which itself has a unique solution.The minimizer is obtained through lower semicontinuity and strong convexity, and its density is positive almost everywhere.Source: §10.1 · pp. 77–79
- Under condition A3 and β ≥ 1, the minimizer’s regularized risk is within K(1 + λ)[d log(2 + 1/λ)]/β of inf ρ∈P(Rd) Rλ(ρ).The bound is obtained by Gaussian smoothing and comparing the finite-temperature free energy with the regularized risk.Source: §10.1 · pp. 79–80
10.2 Dynamics
- With conditions A1–A3 and finite initial free energy, the finite-temperature PDE has a unique weak solution whose densities remain in K, are absolutely continuous, and have uniformly bounded Ent and M.Existence is established using a Wasserstein JKO scheme, followed by compactness and standard uniqueness arguments.Source: §10.2 · pp. 80, 82
- Under the stronger conditions A1–A4, the PDE solution satisfies ρt(θ) ∈ C1,2((0, ∞) × Rd), giving continuous time derivatives and second-order continuous spatial derivatives.The regularity proof uses heat-kernel estimates, Duhamel’s principle, and convolution bounds.Source: §10.2 · pp. 82, 84
- Under A1–A4, ρ∈C1,2((0,∞)×R^d), remains positive, and converges weakly to the unique Boltzmann fixed point and global minimizer of Fβ,λ.Free energy is differentiable and non-increasing along the PDE trajectory.Source: §10.2 · pp. 86–88
10.3 Proof of Proposition 3, Theorem 4, and Theorem 5
- For sufficiently large β and finite T, the PDE is within η of ρ* thereafter, and the finite-width stochastic approximation achieves the target guarantee with probability at least 1−δ under stated scaling conditions.The scaling requires N,1/ε≥C0 exp{C0T}D and ε≥1/N^10.Source: §10.3 · pp. 89–90
10.4 Dependence of convergence time on D and η
- Near ρ*, convergence is heuristically governed by the Fokker–Planck log-Sobolev constant, which may be exponentially small in D, suggesting T=e^O(D)log(1/η).Theorem 5 itself does not estimate dependence on D or η.Source: §10.4 · p. 90
11 Numerical Experiments
Numerical experiments validate the distributional-dynamics PDE against SGD across isotropic and anisotropic Gaussian settings, while examining risk landscapes, optimization predictions, and theoretical conditions. They also illustrate how SGD behavior depends on initialization and κ, including convergence to either poorly generalizing local minima or well-generalizing global minima.
11.1 Isotropic Gaussians
- The isotropic-Gaussian simulations use a piecewise-linear activation with t1 = 0.5, t2 = 1.5, s1 = −2.5, and s2 = 7.5 across the numerical examples.This activation is also used for the centered-Gaussian examples corresponding to Figures 1 and 2.Source: §11.1 · p. 91
11.1.1 Empirical validation of distributional dynamics
- The isotropic-Gaussian PDE is simulated with a multiple-deltas ansatz, reducing the dynamics to numerical evolution of particle locations and using larger J for greater approximation accuracy.For the main Figure 1 setting, J = 400 is used; some cases permit J = 1 because the reduced distribution is well concentrated.Source: §11.1.1 · pp. 91–93
- In isotropic Gaussians, SGD and PDE risk predictions match well, even with J = 1 for ∆ = 0.2, where the reduced distribution remains concentrated.For ∆ = 0.2, the risk transitions rapidly from high to low; for ∆ = 0.8, the transition is slower and the distribution does not concentrate at large time.Source: §11.1.1 · p. 93
Optimizing
- The optimization approximation discretizes the distribution over K pre-chosen points, producing a quadratic programming problem; in d = ∞, the required kernel computation is simpler.The reported implementation uses K = 100 equally spaced points on [0.01, 10].Source: Optimizing · p. 95
11.1.3 Checking the condition of Lemma 1 in the main text
- The static analysis checks the minimizer r* and tests Lemma 1’s inequality over grids of r and ∆, finding the condition satisfied over an interval of ∆ values for the examined dimensions.The grid uses r from 0.1 to 10 in increments of 0.1 and ∆ from 0.01 to 0.99 in increments of 0.01.Source: §11.1.3 · p. 96
11.2 Centered anisotropic Gaussians with ReLU Activation
- SGD and the PDE show a good match in the anisotropic Gaussian ReLU experiment up to iteration 10^6, with smaller discrepancies for larger s0, d, and N.The PDE assumes infinite s0, d, and N, whereas the SGD simulations use finite values.Source: §11.2 · p. 99
11.3 Isotropic Gaussians: Predictable Failure of SGD
- The failure experiment is designed so that δ0 is a poorly generalizing local minimum, and theory predicts convergence there when SGD starts sufficiently close to it.The initialization comparison uses κ = 0.1 and κ = 0.4, with d = 320, N = 800, and ∆ = 0.5.Source: §11.3 · p. 101
- The experiment computes the PDE trajectory from t = 0 to t = 10^7ε and matches SGD iteration k to r(ε^4/3k).The PDE uses J = 400 and a single run, with 10^5 logarithmically spaced time points.Source: §11.3 · p. 102
- The κ = 0.1 initialization converges to a local minimum with nearly zero ∥w_i∥2 and trivial misclassification rate 0.5, whereas κ = 0.4 reaches a well-generalizing global minimum.The error rate is estimated by Monte Carlo averaging over 10^4 samples.Source: §11.3 · p. 102
- Figure 11.9 reports the empirical error rates for the two initializations, validating that SGD can reach either a poorly generalizing local minimum or a well-generalizing global minimum.The error rate is defined as P{sign(ŷ(x; θ)) ≠ y}.Source: §11.3 · p. 102, Figure 11.9 · p. 102
A Concentration inequalities
This section states an Azuma–Hoeffding concentration lemma for an R^d-valued martingale and outlines its proof using subgaussian martingale differences, Markov's inequality, and a stopping-time argument.
A Concentration inequalities
- Lemma A.1 considers an R^d-valued martingale initialized at zero under an almost-sure increment condition.The martingale is adapted to a filtration and takes values in R^d.Source: A Concentration inequalities · p. 102
- The proof writes the martingale increments as Z_k = X_k − X_{k−1} and invokes the subgaussian condition (A.1).Source: A Concentration inequalities · p. 103
- The argument applies Markov's inequality with ξ = 1/(2nL^2) to obtain a bound for all t ≥ 0.Source: A Concentration inequalities · p. 103
- To control max_{k≤n} ||X_k||_2, the proof introduces a stopping time at the threshold 2L√(d + t) and applies the preceding inequality.Source: A Concentration inequalities · p. 103
B On the generalization to other loss functions
This section formally extends the paper’s framework from quadratic loss to general loss functions, deriving the corresponding mean field risk and distributional dynamics. The extension is heuristic, with rigorous analysis deferred to future work.
- The framework is formally generalized to loss functions ℓ: R × R → R, but all arguments remain heuristic and lack a rigorous treatment.A rigorous study is deferred to future work.Source: B On the generalization to other loss… · p. 103
- For a general loss, the population risk induces a mean field risk R: P(R^D) → R.Source: B On the generalization to other loss… · p. 103
- The corresponding distributional dynamics retains the quadratic-loss form, with Ψ(θ; ρ) modified through the derivative of ℓ with respect to its second argument.Source: B On the generalization to other loss… · p. 103
- Setting ℓ(y, ˆy) = (y − ˆy)^2 recovers the expressions used throughout the paper.Source: B On the generalization to other loss… · p. 103
Related papers
Cited by
- Learning to Play: Reinforcement Learning and Games2026
- Global Exponential Convergence of Two-Layer Linear Network Training2026
- Machine learning and the physical sciences2019
- Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks2019
- Gradient Descent Provably Optimizes Over-parameterized Neural Networks2018