1

Convergence of mean field
Langevin dynamics:
Duality viewpoint and neuralnetwork optimization
Taiji Suzuki
The University of Tokyo / AIP-RIKEN(Deep learning theory team)

29 th /July/2023
Duality Principles for Modern Machine Learning
ICML2023 workshop@Hawaii

2Atsushi Nitanda(Kyusyu Institute of Technology/RIKEN-AIP)

Denny Wu

(University of Toronto/Vector Institute)

Noisy gradient descent
Optimization of neural network is basically non-convex.
➢ Noisy gradient descent (e.g, SGD) is effective for nonconvex optimization.
NoiseGradient descent

Noisy perturbation is helpful to escape a local minimum.
➢ Likely converges to a flat global minimum.

3Gradient Langevin Dynamics (GLD)

4(Non-convex)
𝛽: inverse temperature(Gradient Langevin dynamics)

Discretization[Gelfand and Mitter (1991); Borkar and Mitter(1999); Welling and Teh (2011)](Euler-Maruyama scheme)

Stationary distribution：

Can stay around theglobal minimum of 𝐿(𝑥).

GLD as a Wasserstein gradient flow
: Distribution of 𝑋𝑡 (we can assume it has a density)
PDE that describes 𝜇𝑡 ’s dynamics [Fokker-Planck equation]:

This is the Wasserstein gradient flow to minimize the following objective:[linear w.r.t. 𝝁]= Stationary distributionc.f., Donsker-Varadan duality formula5

Objective of mean field NN
Vanilla GLD

LinearConvex objective

Nonlinear extension!

convex + strongly convex = strongly convex

Application:➢ Optimization of 2-layer neural network in mean field regime➢ Variational inference

6

Example of loss function𝑀

Non-linear with respect to the parameters 𝑟𝑗 , 𝑤𝑗 𝑗=1 .

7…

★Mean field limit:
𝑀→∞

Linear with respect to 𝜇.
[Nitanda&Suzuki, 2017][Chizat&Bach, 2018][Mei, Montanari&Nguyen, 2018][Rotskoff&Vanden-Eijnden, 2018]

Loss function (empirical risk + regularization):

Convex w.r.t. 𝜇 if the loss ℓ𝑖 is convex (e.g., squared / logistic loss).

8Particles move to fit themodel to the data.
(distribution)

The parameter of one neuron

General form of mean field LD

9Mean field Langevin dynamics:
convex➢ SDE the Fokker-Planck equation of which corresponds to the Wasserstein GF:

𝐹

Gradient

Distribution dependent SDE
GLD:,Definition (first variation)
𝛿𝐹
The first variation 𝛿𝜇: 𝒫 × ℝ𝑑 → ℝ is defined as a continuous functional such asProximal Gibbs measure

10

Linearized objective at 𝝁:
𝐹Gradient

Minimizer

Proximal Gibbs measure➢ The proximal Gibbs measure is a kind of “tentative” target.
➢ It plays important role in the convergence analysis.

Dual objective (informal)

11[Nitanda, Oko, Wu, Suzuki (ICML2023); Nitanda, Wu, Suzuki (AISTATS2022); Oko, Suzuki, Nitanda,Wu (ICLR2022)]

Primal=(Fenchel’s duality theorem)

Dual

Primal-Dual variable correspondence:
(P)

(D)

(P)

Duality gap and divergence:𝐹••(optimality condition)

GradientEntropy sandwich
Proximal Gibbs measure:

Theorem (Entropy sandwich)[Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (2022)]

12=

Convergence rate

13

Proximal Gibbs measure:

Assumption (Log-Sobolev inequality)c.f., Polyak-Lojasiewicz condition

There exists 𝛼 > 0 such that for any probability measure 𝜈 (abs. cont. w.r.t. 𝑝𝜇 ),KL-div

Theorem (Linear convergence)

Fisher-div[Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (2022)]

If 𝑝𝜇𝑡 satisfies the LSI condition for any 𝑡 ≥ 0, thenThis is a non-linear extension of well known GLD convergence analysis.

Example14

L2-regularized loss function for mean field 2-layer NN:
whereProximal Gibbs:Bounded (≤ 𝐵)

Strongly convex

If sup ℓ′𝑖 𝑓𝜇 (⋅) ℎ𝑥 (⋅) ≤ 𝐵, the proximal Gibbs measure 𝑝𝜇 satisfies the LSI𝑧

with a constant 𝛼 with

∵ Bakry-Emery criterion (1085) and Holley-Strook small perturbation lemma (1987)

Proof outline of convergence• MF-LD obeys the following nonlinear FokkerPlanck equation:

Then,[Continuity equation]=: −𝑣𝑡(∵continuityequation)

15

Mass: 𝜇𝑡 (𝑥)Vector field: 𝑏(𝑥, 𝜇𝑡 )(Definition of 𝑝𝜇𝑡 )LSI & Entropy sandwich𝑡 ) nonlinearly depends on 𝜇 , we say “nonlinear Fokker-Planck”.
※ Since 𝛿𝐹(𝜇𝑡𝛿𝜇Other applications

16

Mean field Langevin dynamics can be applied to several
problems where a distribution is optimized.
• Nonparametric density estimation via MMD minimization

𝑘: positive definite kernel
where 𝑘𝜇 = ∫ 𝑘 𝑥,⋅ 𝜇(d𝑥) (kernel embedding).
➢➢: Empirical distribution (training data)(see also Chizat (2022,TMLR))• Variational inference to approximate Bayesian posterior(KSD: Kernel Stein Discrepancy from a posterior distribution)

17

Finite particle &discrete time algorithm

18

We have obtained a convergence ofinfinite width and continuous time dynamics.
Question:
Can we obtain a finite particle & discrete
time (i.e., implementable) algorithm?
(distribution of 𝑋𝑡 )Neuron(Finite particle approximation)𝑥(vector field)

Difficulty

19• SDE of interacting particles (McKean, Kac,…, 60’)
Propagation of chaos [Sznitman, 1991; Lacker, 2021]:
The particles behave as if they are independent
as the number of particles increases to infinity.
Finite particle approximation error can be amplified through time.
→ It is difficult to bound the perturbation uniformly over time.

• A naïve evaluation
gives exponential
growth on time:
exp 𝑡 /𝑁
[Mei et al. (2018, Theorem 3)]

𝑡=1

𝑡=2

𝑡=3

𝑡=4

➢ Weak interaction/Strong
regularization in existingworkOutline of research
Infinite particles / Continuous time
Linear convergence ofmean field Langevin:[Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (TMLR2022)]

20

Finite particle / Discrete time
Double loop method:• PDA [Nitanda, Wu, Suzuki: NeurIPS2021]• P-SDCA [Oko, Suzuki, Wu, Nitanda:
ICLR2022]• Infinite-dim extension [Nishikawa, Suzuki,Nitanda: NeurIPS2022]

Difficult：
Propagation of chaos (McKean, Kac,…, 60’s)
Finite particle / Continuous time
Uniform-in-time propagation of chaos:• Super log-Sobolev ineq.

[Suzuki, Nitanda, Wu (ICLR2023)]• Leave-one-out typeevaluation/Uniform-log-Sobolev[Chen, Ren, Wang (arXiv2022)]

Finite particle / Discrete timeSingle loop method:
Time-space discretization,stochastic gradient[Suzuki, Wu, Nitanda(arXiv:2306.07221)](1) Double loop algorithm

21

Apply convex optimization techniques developed in finite dimensional settings.

Approximate 𝑞𝑡
by 𝑀 particles

Sampling

Update the target distribution(Gradient descent)Update bygradient descentDouble loop algorithms
PDA:

P-SDCA:

Particle Stochastic Dual Coordinate AscentParticle Dual Averaging[Nitanda, Wu, Suzuki: NeurIPS2021][Oko, Suzuki, Wu, Nitanda: ICLR2022]

PrimalApprox

ഥ (𝒕) is determined by dual averaging method.
𝒈=Approximate this by a linear functional of 𝑞.
(linear approx)

DualSolution:

where→ We can sample particles from this distribution byusing the following GLD.
Time discretization

Computational complexity：
1. Inner loop:
2. Outer loop:

by Fenchel duality theorem•

Randomly choose a coordinate of thedual variable and optimize theselected coordinate.
→ stochastic coordinate ascentComputational complexity：# of outer loops to obtain the duality gap 𝝐𝑷 ：(by GLD)⇒Total: 𝑶(𝝐−𝟑 ) times gradient update➢ The first polynomial time method➢➢

Exponential order convergenceRelax the dependency on sample
size

22

Outline of research
Infinite particles / Continuous time
Linear convergence ofmean field Langevin:[Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (TMLR2022)]

23

Finite particle / Discrete time
Double loop method:• PDA [Nitanda, Wu, Suzuki: NeurIPS2021]• P-SDCA [Oko, Suzuki, Wu, Nitanda:
ICLR2022]• Infinite-dim extension [Nishikawa, Suzuki,Nitanda: NeurIPS2022]

Difficult：
Propagation of chaos (McKean, Kac,…, 60’s)
Finite particle / Continuous time
Uniform-in-time propagation of chaos:• Super log-Sobolev ineq.

[Suzuki, Nitanda, Wu (ICLR2023)]• Leave-one-out typeevaluation/Uniform-log-Sobolev[Chen, Ren, Wang (arXiv2022)]

Finite particle / Discrete timeSingle loop method:
Time-space discretization,
stochastic gradient[Suzuki, Wu, Nitanda(arXiv:2306.07221)][Suzuki, Wu, Nitanda: Convergence of mean-field Langevin dynamics: Time and space
discretization, stochastic gradient, and variance reduction. arXiv:2306.07221]

(2) Single loop method

𝑁 particles(time discretization)

𝑖𝑋𝑘24𝑁𝑖=1

andwhere(stochastic gradient)(space discretization)➢ Noisy gradient descent on 2-layer NN with finite width.

• Time discretization: 𝑡 → 𝑘𝜂 (𝜂: step size, 𝑘: # of steps)• Space discretization: 𝜇𝑡 is approximatd by 𝑁 particles
𝜇𝑡 → 𝜇ො𝑘 = 𝑁1 ∑𝛿𝑋 (𝑖)• Stochastic gradient: ∇ 𝛿𝜇 → 𝑣𝑘𝑖
𝛿𝐹 𝜇𝑘Numerical experiment𝑁→∞

Training error with 𝑟 𝑥 = 𝑥 4

Test error with 𝑟 𝑥 = 𝑥 225

Convergence analysis

26: proximal Gibbs measure
Theorem (One-step update) [Suzuki, Wu, Nitanda (2023)]

Suppose that 𝑝𝜇 satisfies log-Sobolev inequality with a constant 𝛼.
Under smoothness and boundedness of the loss function, it holds thatNaïve bound:

Timediscr.

Spacediscr.

Stochasticapprox.

Assumption:
1. 𝐹: 𝒫 → ℝ is convex and has a form of 𝑭 𝝁 = 𝑳 𝝁 + 𝝀𝟏 𝔼𝝁 𝒙 𝟐 .
2. (smoothness)(boundedness)𝛿𝐿 𝜇∇ 𝛿𝜇𝛿𝐿 𝜇∇ 𝛿𝜇𝛿𝐿 𝜈𝑥 −∇ 𝛿𝜇 𝑦𝑥≤ 𝑅.

≤ 𝐶(𝑊2 𝜇, 𝜈 + 𝑥 − 𝑦 ) and(+ second order differentiability)[Suzuki, Wu, Nitanda: Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. arXiv:2306.07221]

27

Uniform log-Sobolev inequality

28(2)

𝑋𝑘(𝑁)
𝑋𝑘(1)
𝑋𝑘

𝒳𝑘 =𝑁𝑁(𝑵)on ℝ𝒅×𝑵 :𝑖
𝑋𝑘

Potential of the joint distribution 𝝁𝒌∼ 𝜇𝑘 : Joint distribution
𝑖=1of 𝑁 particles.

where➢ The finite particle dynamics is the Wasserstein gradient flow that minimizes.

(Approximate) Uniform log-Sobolev inequality [Chen et al. 2022]
For any 𝑵,

whereRecall(Fisher divergence)[Chen, Ren, Wang. Uniform-in-time propagation of chaos
for mean field langevin dynamics. arXiv:2212.03050, 2022.]

Log Sobolev for Lipschitz cont obj

29Proximal Gibbs measure:

Assumption:
𝜇 satisfies the LSI if there exits 𝛼 > 0 such that for any 𝜙 s.t. 𝜇 𝜙 2 = 1, it holds that
1. Holley—Strook argument:[Bakry & Emery, 1985; Holley & Stroock, 1987]⇒(New)
2. Lipschitz perturbation argument + Miclo’s trick:[Cattiaux & Guillin, 2022; Bardet et al., 2018](Lipschitz continuous)⇒Computational complexity
SG-MFLD(finite sum),(stochastic gradient)(Mini-batch size = 𝐵)

Iteration complexity:
By setting
the iteration complexity becomes

Timediscr.

Spacediscr.

Stochasticapprox.

,

to achieve 𝜖 + 𝑂(1/(𝜆2 𝛼𝑁)) accuracy.
➢𝐵=1/(𝜆2 𝛼𝜖) is the optimal mini-batch size. → 𝑘 = 𝑂 log 𝜖 −1 Τ𝜖 .

30

Variance reduction

31

SVRG-MFLD:(finite sum),Variance reduction(𝑋ሶ is updated onceat every 𝑚 steps)
Tighter than the analysis
in linear GLD [Kinoshita,Suzuki: NeurIPS2022]
Time
discr.

Space
discr.

Stochasticapprox.

(𝑚 = 𝐵 = 𝑛)

# of update：where 𝐵 = 𝑚 = 𝑛1/3 .

Total complexity：

𝑛 by Kinoshita&Suzuki(2022)Generalization analysis Reference• ℓ𝑖 : logistic loss
• ℎ𝑧 𝑥 = 𝑅ത ⋅ tanh 𝑥1 , 𝑧 + 𝑥2 /2• Learning XOR function on high dimensional data.

➢ 𝑋 ∼ Unif( −1,1 𝑑 ) (up to freedom of rotation)➢ 𝑌 = 𝑋𝑘 𝑋𝑙 for 𝑘, 𝑙 ∈ 𝑑 with 𝑘 ≠ 𝑙.
Q: Can we learn XOR function with GD?
How large is the computational cost?

Table 1 of [Telgarsky: Feature selection and low test
error in shallow low-rotation ReLu networks, ICLR2023].

32

33

• Setting 1: 𝑛 > 𝑑2

Reference➢ Comp complexity: exp(O(𝑑))➢ Test error (classification error) = 𝐎(exp(− 𝒏/𝒅))• Setting 2: 𝑛 > 𝑑➢ Comp complexity: exp(O(𝑑))➢ Test error (classification error) = 𝐎(𝒅Τ𝒏)Conclusion

34

Mean field Langevin dynamics• Entropy sandwich➢ Connecting duality gap with KL-div between thecurrent solution and its proximal Gibbs measure.
➢ Exponential convergence• Finite particle approximation➢ Uniform-in-time propagation of chaos
Many other interesting topics:• Entropic fictious play [Chen, Ren, Wang (2022); Nitanda et al. (ICML2023)]• Learning theory, better sample complexity than NTK [Suzuki et al. (2023)]• Application to Reinforcement Learning: Policy-Gradient [Yamamoto et al. (2023)]• Infinite dimensional mean field Langevin [Nishikawa, Suzuki, Nitanda (NeurIPS2022)]