Contribution
Provides a concise, self-contained convergence rate analysis of mean field Langevin dynamics (MFLD) for KL-regularized objectives in both continuous and discrete time, using a novel “proximal Gibbs distribution” p_q that makes the analysis parallel classical finite-dimensional convex optimization. Establishes linear convergence in continuous time for any regularization strength, and a discrete-time rate for noisy gradient descent on mean field neural networks. Additionally shows p_q exactly characterizes the duality gap in empirical risk minimization.
Key insight
The proximal Gibbs distribution p_q(θ) ∝ exp(-(1/λ) δF/δq (q)(θ)) plays the role of a “gradient” and “proximal point” in the space of measures: KL(q || p_q) simultaneously serves as a Polyak-Łojasiewicz-style upper bound on the optimality gap L(q) - L(q*) and, in ERM settings, equals exactly the duality gap. Combined with a log-Sobolev inequality on p_q, this turns convergence analysis of nonlinear Fokker-Planck dynamics into a near-mechanical translation of classical convex optimization arguments.
Assumptions
- F is convex on P with smooth functional derivative δF/δq(q)(θ) = O(1 + ||θ||
- Proximal Gibbs distribution p_q satisfies log-Sobolev inequality uniformly in q with constant α (LSI)
- For MF neural nets: bounded loss derivative |∂_z ℓ| ≤ C_1, bounded activation |h_θ(x)| ≤ C_5, L2 regularization r(θ) = ||θ||^2 (gives α = 2λ’/(λ exp(4 C_1 C_5/λ)))
- Gradient Lipschitz conditions on δF/δq for discrete-time analysis
- Population/expected risk or IID empirical samples (IID assumed; not a distributed/federated setting)
Key results
- Theorem 1 (continuous time). L(q_t) - L(q*) ≤ exp(-2αλ t) (L(q_0) -
L(q*)): linear convergence of MFLD in continuous time
Holds when: Requires uniform LSI constant α > 0 on p_q; holds for any λ
0
- Corollary 1. inf_{s∈[0,t]} KL(q_s || p_{q_s}) ≤ exp(-2αλ(t-1))/(2αλ^2) · (L(q_0) - L(q*)) Holds when: Same LSI assumption
- Theorem 2 (discrete time). L(q^(k)) - L(q*) ≤ δ_η/(2αλ) + exp(-αλη k)(L(q^(0)) - L(q*)) Holds when: δ_η is discretization error O(η) for MF neural nets under bounded gradients
- Corollary 2 (iteration complexity). O((1/(ε α^2 λ^2)) log(1/ε)) iterations to reach ε-accurate solution Holds when: Exponential dependence on 1/λ through α
- Theorem 3 (duality gap). L(q) - D(g_q) = λ KL(q || p_q): exact duality gap in ERM Holds when: L2 regularization, smooth convex per-sample losses
Claims
| id | claim | strength | support |
|---|---|---|---|
| C1 | MFLD converges linearly to the global optimum of the KL-regularized nonlinear objective in continuous time, for any regularization strength, given a uniform log-Sobolev inequality. | strong | Theorem 1 with direct Lyapunov argument |
| C2 | Noisy gradient descent (discretized MFLD) achieves O((1/(ε α^2 λ^2)) log(1/ε)) iteration complexity for ε-accurate optimization of mean field two-layer networks. | moderate | Theorem 2 + Corollary 2; constants have exponential dependence on 1/λ |
| C3 | The proximal Gibbs distribution p_q exactly equals the primal-dual duality gap (λ KL(q / / p_q) = L(q) - D(g_q)) in ERM, enabling convergence to be monitored without knowing L(q*). | strong | Theorem 3 + empirical illustration on student-teacher regression |
| C4 | Finite-particle approximation error grows exponentially with time horizon and is not negligible unless the linear convergence rate is sufficiently large. | weak | Informal argument citing Mei et al. 2018 Theorem 3 |
Method
Mean Field Langevin Dynamics (MFLD) / Noisy Gradient Descent.
Continuous time SDE: dθ_t = -∇(δF/δq)(q_t)(θ_t) dt + √(2λ) dW_t, whose density q_t follows a nonlinear Fokker-Planck equation. Discrete time: θ^(k+1) = θ^(k) - η ∇(δF/δq)(q^(k))(θ^(k)) + √(2λη) ξ^(k) with Gaussian noise ξ^(k). For analysis, introduce p_q ∝ exp(-(1/λ) δF/δq(q)), show L(q) - L(q*) ≤ λ KL(q || p_q), use LSI on p_q to get a PL-type inequality, and apply the Vempala-Wibisono one-step interpolation trick for discrete-time bounds.
- Proximal Gibbs distribution p_q as a surrogate/gradient in measure space
- Log-Sobolev inequality (uniform over q) to get PL-like condition
- One-step interpolation argument from Vempala & Wibisono (2019)
- Primal-dual formulation linking p_q to Fenchel duality gap
Concepts
- Proximal Gibbs distribution p_q — p_q(θ) ∝ exp(-(1/λ) δF/δq(q)(θ)); minimizer of linearization of L at q plus λ KL(·||q); fixed point equation q* = p_{q*} characterizes the optimum.
- Mean field Langevin dynamics — SDE whose density obeys a nonlinear Fokker-Planck equation arising from KL-regularized risk minimization over parameter distributions of a mean field neural net.
- Log-Sobolev inequality (LSI) — Ent_{p}(g^2) ≤ (2/α) E_p[||∇g||^2]; implies KL(q||p) ≤ (1/(2α)) Fisher information; gives linear convergence.
- Duality gap characterization — In ERM with L2 regularization, L(q) - D(g_q) = λ KL(q || p_q), with g_q the vector of per-sample loss derivatives.
Connections
Builds on.
- Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices — Adapts the Vempala-Wibisono one-step interpolation argument from linear Fokker-Planck / LMC to the nonlinear MFLD setting.
- On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport (LIT-298) — Part of the mean field neural network convergence program that this paper sharpens with quantitative rates.
- Particle Dual Averaging (Nitanda, Wu, Suzuki 2021) — Prior work by same authors using LSI-based analysis of MFLD in a dual averaging inner loop; this paper removes the dual averaging outer loop.
- Mean-field Langevin dynamics: Exponential convergence and annealing (Chizat 2022) — Concurrent independent work using the same proximal Gibbs distribution observation; focuses on continuous-time and annealing rather than discrete-time.
Related.
- Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices — Key technical ingredient extended from linear to nonlinear Fokker-Planck.
Recommendations
- R1 — When analyzing noisy SGD / Langevin-type training in measure space, check whether a log-Sobolev inequality holds (uniformly) on the proximal Gibbs distribution — it yields PL-type linear convergence in continuous time and dissolves many technical barriers. [not filed as a practice: advice on proof technique rather than on training] Topic: convergence analysis tooling · Strength: strong · When: KL- regularized objective, smooth bounded functional derivative, bounded activation/loss derivative.
- R2 — Monitor optimization via the primal-dual duality gap λ KL(q || p_q) rather than requiring knowledge of L(q*); this can be estimated by sampling from p_q with an inner Langevin loop. Topic: training diagnostics · Strength: moderate · When: ERM settings where a primal-dual formulation exists (L2-regularized convex per-sample loss).
- R3 — Expect exponential dependence on 1/λ in convergence constants for MFLD under generic LSI assumptions; plan annealing schedules or structural assumptions if small λ is required. Topic: regularization tuning · Strength: moderate · When: MF neural network training with small KL regularization.
Bearing on the record
Mean-field Langevin dynamics as convex optimization in measure space. Its R2 — monitor the duality gap rather than needing to know the optimum — is the one part that is a diagnostic rather than a proof technique.
Limitations
- Uniform log-Sobolev constant α has exponential dependence on 1/λ in the general bounded-perturbation analysis, degrading rates for weak regularization.