1

Deep Learning Theory

Taiji Suzuki
The University of Tokyo
Deep Learning Theory Team/AIP-RIKEN

12 nd /Mar/2024
MLSS2024@OIST

Success of deep learning

AlphaGo/Zero

Large language model
Learning efficiency
of few shot learning

Deep learning has shown
great performances in
the AI research field.
→ Why?

2[Alammar: How GPT3 Works - Visualizations andAnimations,
https://jalammar.github.io/how-gpt3-worksvisualizations-animations/]

Performance of few-shotlearning against model size

[Brown et al. “Language Models are Few-Shot Learners”, NeurIPS2020]

Generative models (diffusion models)[Silver et al. (Google Deep Mind): Mastering the game of Go withdeep neural networks and tree search, Nature, 529, 484—489, 2016][Ho, Jain, Abbeel: Denoising Diffusion Probabilistic Models. 2020]

Image recognition[He, Gkioxari, Dollár, Girshick: Mask R-CNN, ICCV2017]

Stable diffusion,
2022.

Jason Allen "Théâtre D'opéra
Spatial“ generated by Midjourney.
Colorado State Fair’s fine art competition,
Generated by NovelAI
1st prize in digital art category

What we need to solve?

3

Why does deep learning work well?
• Several theoretical work has been conducted.
• There are still many things that should be explored.
• Clarification of principle of deep learning• What is essential to realize a “good” learning system?
→We may find a new method beyond DL.
Issue in academic conference

Issue in industries• We don’t want to use black-box
system.
• Accountabilities of companies.

Ali Rahimi’s talk at NIPS2017 (test of time award).
“Random features for large-scale kernel methods.”

Criticism that DL is “alchemy.”

Deep learning theory

Role of mathematics

4

Physics• Theory of Relative➢ Riemannian geometry• Quantum mechanics➢ Functional analysis

Mathematics

physical phenomenon

Several mathematicians/physicists
join the ML community.

Machine learning••••

Deep learning•
Prob. theory•
Functional anal.
Wasserstein geom. •Diffusion equation

Statistics
Optimization
Numerical
analysis

Mathematics

Layers of deep learning theory

5

Application

Interpretability：
Accountability, visualization,easier maintenance

Understanding
behaviors of DL➢ Accountability
Analysis of several techniques： ➢ Clarifying possibility
and limitations
Analysis of architecture, loss function➢ Guideline fordesign, analysis of optimization techniquedesigning a learning
methodPrinciple of deep learning：
Representation ability, generalization
ability, convergence analysis of opt.

Essence of learning：
Characterization of “good” learning
methods, unified theory，beyond DL
Foundation

Today’s topic

3 issues of deep learning theory
Representation ability
What kind of functions can
DNN approximate?

Generalization ability
How well can DL generalize
from finite observations?
Optimization ability
How fast can we find the
optimal parameter?

67Outline
1. Representation ability + Generalization
ability➢Universal approximator➢Depth separation➢Adaptivity of deep learning➢ Inhomogeneity of smoothness➢ Curse of dimensionality➢Foundation models➢ Diffusion model➢ Transformer

2. Optimization ability➢Noisy gradient descent➢Mean field Langevin➢CSQ lowerbound8

9

Representation abilityof neural networks
Representation
ability

Generalization
ability

Optimization

Representation ability
Universal approximator
Neural networks can approximate“any function” with “any precision”.

2-layer NNs can approximate any function,by increasing the number of neurons.
[Hecht-Nielsen,1987][Cybenko,1989]
Year

Basis function

Space

1987

Hecht-Nielsen

Depending on the target

𝐶(𝑅 𝑑 )

1988

Gallant & White

Cos

𝐿2 (𝐾)

Irie & Miyake

integrable

𝐿2 (𝑅 𝑑 )Carroll & Dickinson

Continuous sigmoidal

𝐿2 (𝐾)

CybenkoContinuous sigmoidal

𝐶(𝐾)Funahashi

Monotone & bounded

𝐶(𝐾)

1993

Mhaskar + Micchelli

Polynomial growth

𝐶(𝐾)

2015

Sonoda + Murata

Unbounded, admissible

1989

is any compact set.

𝐿1 𝑅 𝑑 , 𝐿2 𝑅 𝑑

10

What are we missing?
• [Theory] Kernel method is also auniversal approximator．
• [Practice] DL performs better.
(in some case)→ Why?
Kernel method(linear model)

11

Deep learning

Deep model

Trainable

Trainable

Fixed

Trainable→ We compare “accuracy” of estimation/approximation.

Feature learning

12

Adaptivebasis learning

Nonlinear• Linear model• Kernel model

Fixed

Fixed

Trainable• Neural network

Trainable

Trainable Trainable13

Kernel model with fixed basis

𝜑1RadiusHuge radius
𝜑2

Target function

14

Kernel model with fixed basis

NN with 𝜑(𝑥; 𝜃1 )𝜑1Radius𝜑2

Target function

NN with 𝜑(𝑥; 𝜃2 )

Deep network is exponentially powerful

15

Width vs Depth

𝐿

=

𝐿’

Exponentially large width is required.

[Arora, Basu, Mianjy, Mukherjee: Understanding Deep Neural Networks with Rectified Linear Units.
ICLR2018.]

Curse of dimensionality/Barron class16• Mean field model defined by 𝜋

whereApprox. error• Neural network model with 𝑴 neurons
The first layer can be tuned：set of NNs• Random feature model with 𝑴 neurons

Curse of dimensionality

Generated randomly(fixed independent of the target)(To obtain 𝜖 accuracy, 𝑀 = 𝜖 −Ω(𝑑) is required)[E, Ma, Wu: A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics. Science China
Mathematics volume 63, 1235–1258 (2020)][E, Ma, Wu: A priori estimates of the population risk for two-layer neural networks. COMMUN. MATH. SCI. 17(5), 1407–1425 (2019)]

Outline
1. Representation ability + Generalization
ability➢Universal approximator➢Depth separation➢Adaptivity of deep learning➢ Inhomogeneity of smoothness➢ Curse of dimensionality➢Foundation models➢ Diffusion model➢ Transformer

2. Optimization ability➢Noisy gradient descent➢Mean field Langevin➢CSQ lowerbound

17

18

Analysis in nonparametric regression-Superiority of deep learning-

Representation
ability

Generalization
ability

Optimization

Non-parametric regression
Non-parametric regression

where 𝜉𝑖 ∼ 𝑁(0, 𝜎 2 ) and 𝑥𝑖 ∈ 0,1 𝑑 ∼ 𝑃(𝑋) (i.i.d.).

We estimate 𝑓 o from 𝑥𝑖 , 𝑦𝑖 𝑛𝑖=1 .

Estimation error:

A similar argument can be applied to classification.

19

Bias-Variance decomposition
Model ℱ: 𝑑-dimensional parameter𝑛: data size

Model’s degrees of freedom𝑀𝑛Predictive error =(mean squared error)

Variance+ [Approx. error]
Model Bias

• Bias-variance trade-off
Large model is not necessarily good
• 𝑑 can be replaced by an “intrinsic”
dimensionality.
• Several learning theory is reduced toevaluate the bias and variance terms.
[Gine&Nickl: Mathematical Foundations of InfiniteDimensional Statistical Models. 2015][Schmidt-Hieber, 2017;Hayakawa&Suzuki, 2020]ℱ

20

Bias-Variance decomposition
Model ℱ: 𝑑-dimensional parameter𝑛: data size

(covering number)(mean squared error)

Variance

Model Bias

• Bias-variance trade-off
Large model is not necessarily good
• 𝑑 can be replaced by an “intrinsic”
dimensionality.
• Several learning theory is reduced toevaluate the bias and variance terms.
[Gine&Nickl: Mathematical Foundations of InfiniteDimensional Statistical Models. 2015][Schmidt-Hieber, 2017;Hayakawa&Suzuki, 2020]ℱ

21

Covering number

22

The smallest number of balls with radius 𝜖 measured
by the norm ‖ ⋅ ‖∞ to cover the function class ℱ.

• The covering number measures how large the space ℱ is.
• In other words, it represents “complexity” of the model.

Scaling law[Kaplan et al.: Scaling Laws for Neural Language Models, 2020]
Fixed model size

（Data size）

Reducible loss

[Henighan et al.: Scaling Laws for Autoregressive Generative Modeling, 2020]

log(Predictive error)=−𝛼 log 𝑛 + log(𝐶)[Brown et al.: Language Models are Few-Shot Learners, 2020] ←Analysis of GPT-3

23

Analysis of kernel model
Model(ONS in L2)

24

Assumption

Observation:

Training method
𝑴 dimensional model

Predictive error

Least squares estimator

Variance

Bias

Optimal rate𝑀𝑎log Pred. error = − 1+𝑎log 𝑛 + log(𝐶)𝑀−𝑎Variance=(dim of model)/n
Bias=L2-norm of residual

25

Predictive error with the optimal model size𝑎log(Predictive error) = − 1+𝑎log 𝑛 + log(𝐶)log(Pred. error)
Predictive error for model size M

Worse
prediction

Betterprediction

Smaller data

Larger datalog(𝒏)(data size)

26

Statistical learning theory of kernel methods• Caponnetto and De Vito. Optimal Rates for the Regularized LeastSquares Algorithm. Foundations of Computational Mathematics,
volume 7, pp.331–368 (2007).
• Steinwart and Christmann. Support Vector Machines. 2008.

Related recent papers• Mei, Misiakiewicz, Montanari. Generalization error of random featuresand kernel methods: hypercontractivity and kernel matrix
concentration. arXiv:2101.10588.
• Bordelon, Canatar, Pehlevan. Spectrum Dependent Learning Curves in
Kernel Regression and Wide Neural Networks. ICML, 1024-1034, 2020.
• Canatar, Bordelon, Pehlevan. Spectral Bias and Task-Model Alignment
Explain Generalization in Kernel Regression and Infinitely Wide Neural
Networks. Nature Communications, volume 12, Article number: 2914
(2021).

Predictive error of deep neural network

27

To bound the predictive error of deep learning,
we evaluate➢ Variance of a DL estimator (Sample variance)➢ Bias of the model (Approximation error)
Typical function class• Barron class• Hölder class• Sobolev class• Besov class
Approximation theory:• Mhaskar: Neural networks for localized approximation of real functions. In Neural Networks for
Processing III, Proceedings of the 1993 IEEE-SP Workshop, 190–196, 1993.
• Pinkus: Approximation theory of the mlp model in neural networks. Acta Numerica, 8:143–195, 1999.
• Mhaskar: Neural networks for optimal approximation of smooth and analytic functions. NeuralComputation, 8(1):164–177, 1996.

Why deep?

28• There are many theories...
Reduced rank
regression

Piece-wisesmooth function

If there is low
dimensional
representation,
deep is better.

Deep is betterto estimate a
discontinuousfunction.

[Imaizumi&Fukumizu, 2019]

Besov space[Suzuki, 2019]To estimate
functions with
non-uniform
smoothness,
DL is better.

Low dimensional
data structure[Schmidt-Hieber, 2019] [Nakada&Imaizumi,
2019][Chen et al., 2019][Suzuki&Nitanda, 2019]

If data are
distributed on low
dim-manifold, DL
is better.

Est. error

Deep Kernel

Typical situation29

Non-smooth

Smooth•Approximation by Gaussian kernel•Approximation by Gaussian kernel
Same Gaussian width→Less efficient•Approximation by NN

Different resolution depending on location→More efficient

Kernel method cannot specify
informative directions.
→Curse of dimensionality•Approximation by NN
Informative directions is extracted ininternal layers.
→Avoids curse of dimensionality

Hölder, Sobolev, Besov space30

Spatial inhomogeneity0

Smoothness

Hölder, Sobolev, Besov space

Intuition

31

Smoothness

Uniformity of smoothness
Spatial inhomogeneity0

Smoothness

Deep learning has adaptivity

32• DL constructs basis function “adaptively”.
→Efficient learning• Shallow learning should use “redundant” model.
→Inefficient learning (affected by redundant noise)𝑠𝑓 ∘ ∈ 𝐵𝑝,𝑞
( 0,1 𝑑 ): “Besov space”[Suzuki: Adaptivity of deep ReLU network for learning in Besov and mixedsmooth Besov spaces: optimal rate and curse of dimensionality, ICLR2019]

Linear estimator (Shallow method)

Deep learninge.g., kernel ridge regression:≫
Suboptimal

Optimal(𝑛: sample size，𝑝: uniformity of smoothness，𝑠: smoothness)
Uniform
resolution
Convergence rate of estimation error (mean squared error E 𝑓መ − 𝑓 ∗- The rate of error decrease as the sample size 𝑛 → ∞.

Adaptive
resolution2)

Linear estimator“Shallow” learning methods
Kernel ridge regression:

33

Trainable

Fixed(feature map)

Fixed

Gram matrix (kernel function)

Nadaraya-Watson estimator:(see also [Imaizumi&Fukumizu, 2019])

Linear estimator: linear to the observation 𝑌 = 𝑦𝑖 𝑛𝑖=1 .

Examplelinear••••Kernel ridge estimator
Sieve estimator
Nadaraya-Watson estimator
k-NN estimatorRelation to sparse estimation

34

Wavelet basis

Resolution
𝑘=0
𝑘=1

Low
frequency

𝑘=2

𝛼0,1
𝑗=1

𝛼1,1
𝑗=1
𝛼2,1 𝛼2,2
𝑗=1

𝑗=2

𝛼1,2
𝑗=2
𝛼2,3 𝛼2,4
𝑗=3

𝑗=4

𝑘=3
High
frequency

Sparse linear combination of wavelet
basis functions is effective to capture
input-dependent smoothness(0 < 𝑝)(informal)

Wavelet basis decomposition

Small 𝑝 = Sparse coefficient

Non-uniform smoothness
over the space

Relation to sparse estimation

35

Wavelet basis

Resolution

𝛼0,1

𝑘=0
𝑘=1

Low
frequency

𝑘=2

𝑗=1
𝛼1,1
𝑗=1
𝛼2,1 𝛼2,2
𝑗=1

𝑗=2

𝛼1,2
𝑗=2
𝛼2,3 𝛼2,4
𝑗=3

𝑗=4

𝑘=3
High
frequency

Sparse linear combination of wavelet
basis functions is effective to capture
input-dependent smoothness(0 < 𝑝)(informal)

Wavelet basis decomposition

Small 𝑝 = Sparse coefficient

Wavelet-shrinkage

Non-uniform smoothness
over the space

Proof strategy

36• Step 1 : Find basis function expansion.

⋅∵ Adaptive approximation by
B-Spline [DeVore & Popov, 1988; Dung, 2011]
−𝑠/𝑑
2
≤
𝑁
𝐿

• Step 2 : Approximate basis functions.
: Approximation by DNN.
: linear combination.

• Step 3: Combine the bounds of Step 1 and 2.

≤ O(𝑒 −𝐿 )≤ 𝑁 −𝑠/𝑑

Bias variance decomposition[Local Rademacher complexity bound]Classic analysis on nonparametric regression.
See [Schmidt-Hieber, 2019; Hayakawa&Suzuki,2020] for the analysis of DNN.

Depth

Width

Sparseness(# of non-zero parameters)

Upper bound of absolute valueof each parameter

Choose the network size 𝑁 to balance bias and variance:෨ 2𝑑Τ 2𝑠+𝑑 ).
𝑁 = 𝑂(𝑛37

Hardness: Convex hull argument
A function with a property that is destroyed by convex combinationis hard to estimate by linear estimators.
e.g., “Spatial inhomogeneity of smoothness”

Theorem(This can be extended to Q-hull for the fixed design setting)[Donoho & Johnstone, 1994] [Satoshi Hayakawa and Taiji Suzuki: 2020]

Ex. Piecewise constant function with 3 jumps.

0.5x

Deep:1/𝑛, Kernel: 1/ 𝑛

0.5x

• Linear estimator should prepare redundant basis functions.

=+
3 jumps

3 jumps

6 jumps38

General theory
Piece-wise smooth

39

Besov space

Reduced rank
regression

Low dim. data
Variable smooth
Besov space

Non-convexity
sparsity

Curse of dimensionality
Estimation error bound：→ Curse of dimensionality

MNIST: 784 dim/13.4 intrinsic-dim[Facco et al. 2017]40

Dimensionality: Manifold regression

41

-dim

-dim

MNIST:
784 dim/13.4 intrinsic-dim[Facco et al. 2017]• Classic nonparametric method: Bickel & Li (2007); Yang & Tokdar(2015); Yang & Dunson (2016).
• Deep learning: Nakada & Imaizumi (2019); Schmidt-Hieber (2019);
Chen et al. (2019).

More realistic setting

Invariantdirection

42

Variable
direction

𝑠1 , 𝑠2 , 𝑠3 : smoothness(non-smooth) 𝑠1 , 𝑠2 ≪ 𝑠3 (smooth)

Data are hardly distributed exactly on low-dim manifold.

• Smoothness could depend on directions.
• Local coordinate.

Anisotropic Besov space

43

Def. (Anisotropic Besov space)(finite difference)(modulus of smoothness)(𝑠2 times differentiable)(𝑠1 times differentiable)

𝐿𝑝 -norm of 𝑠𝑖 -times derivative.
• 𝑠𝑖 : smoothness to the -th coordinate.
• 𝑝: Uniformity of smoothness overthe input space.

Anisotropic Besov space

44

Def. (Anisotropic Besov space)(finite difference)(modulus of smoothness)(𝑠2 times differentiable)(𝑠1 times differentiable)

𝐿𝑝 -norm of 𝑠𝑖 -times derivative.
• 𝑠𝑖 : smoothness to the -th coordinate.
• 𝑝: Uniformity of smoothness overthe input space.

45

𝑦

𝑥2
𝑆2 : large

𝑥1
𝑆1 : small

Estimation error bound46(least squares estimator)※ Here, we do not discuss optimization ability.

TheoremLet,

The rate of convergence is determined by smoothness parameters.

When 𝑯 = 𝟏,

Comparison to linear estimator：
DNN

Linear（Affected by dimension）

If 𝑠𝑖 s are small (non-smooth) toward
small numbers of directions andlarge toward other directions, DNN
can avoid the curse of dimensionality.
If 𝒔𝟐 = ⋯ = 𝒔𝒅 = ∞, then 𝒏−

𝟐𝒔𝟏
𝟐𝒔𝟏 +𝟏.

• Linear estimator cannot find important directions.
Then, the rate of convergence is strongly affectedby the most non-smooth (𝑠1 ) parameter.
• We used a “convex hull argument” to show therate of convergence.

Comparison to linear estimator𝑓 ∘ depends only 𝐷-dimensional subspace.

Deep

Linear estimator

≪
(𝑛

2𝑠
−
2𝑠+1 when 𝐷 = 1)

𝑐 = 1 if 𝐷 < 𝑑/2, 𝑐 = 0 if 𝐷 ≥ 𝑑/2.
2𝑠+𝑑

(𝑛

−2𝑠+2𝑑Suboptimal

OptimalDeep can ease curse of dim.,but linear estimators directly
suffers from curse of dim.

when 𝐷 = 1 and 𝑝 = 1)

𝑊

𝑊

47

Outline
1. Representation ability + Generalization
ability➢Universal approximator➢Depth separation➢Adaptivity of deep learning➢ Inhomogeneity of smoothness➢ Curse of dimensionality➢Foundation models➢ Diffusion model➢ Transformer

2. Optimization ability➢Noisy gradient descent➢Mean field Langevin➢CSQ lowerbound

48

49

Analysis of diffusion model• Kazusato Oko, Shunta Akiyama, Taiji Suzuki: Diffusion Models are
Minimax Optimal Distribution Estimators. ICML2023.

Diffusion model「An astronaut riding ahorse in aphotorealistic style」

50

DALL·E: [Aditya Ramesh, Mikhail
Pavlov, Gabriel Goh, Scott Gray,
Chelsea Voss, Alec Radford, MarkChen, Ilya Sutskever: Zero-Shot
Text-to-Image Generation.
ICML2021.]
DALL·E2:[Aditya Ramesh, Prafulla
Dhariwal, Alex Nichol, Casey Chu,
Mark Chen: Hierarchical TextConditional Image Generation with
CLIP Latents. arXiv:2204.06125]

Stable diffusion, 2022.
Jason Allen "Théâtre D'opéra Spatial“ generated by Midjourney.
Colorado State Fair’s fine art competition, 1st prize in digital art
category

Generated by NovelAI

Movie generation

51• SORA (OpenAI, 2024)➢ Diffusion Transformer

A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears
a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red
lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of thecolorful lights. Many pedestrians walk about.

Diffusion model[Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Song et al., 2020; Ho et al., 2020; Vahdat et al., 2021]

Forward process：Convert the target distribution to a noise distribution (e.g., Gaussian)(𝑌𝑡 ∼ 𝑋𝑇−𝑡 )Reverse process：Convert the noise distribution to the target distribution[Vahdat, Kreis, Kautz: Score-based Generative Modeling in Latent Space. arXiv:2106.05931]

52Forward process
Forward process:

OU process

where 𝜇𝑡 = exp −𝑡 , 𝜎𝑡2 = 1 − exp −2𝑡 .

The forward process converges to the noise distribution (standard
normal) exponentially:

Standard normal
OU-process[Vahdat, Kreis, Kautz: Score-based
Generative Modeling in Latent Space.
arXiv:2106.05931]

53

Reverse process
Reverse process:𝑌0 ∼ 𝑝𝑇

54[Haussmann & Pardoux, 1986]ത
(𝑡 ∈ [0, 𝑇])

Fact：𝑌𝑡 ’s distribution=𝑋 𝑇−𝑡
ത ’s distribution
That is, 𝑌𝑡 ∼ 𝑝𝑇−𝑡
By following the forward process in reverse, noise that follows a (nearly) normal
distribution can be gradually modified to reproduce the original distribution of theimages.

Goal

Start

Benefit of diffusion model
We can sample from multimodal distribution
efficiently.
➢“Easy distribution” →”Difficult distribution”
If we try to sample directly fromthe original distribution, then it
could not get over the “gap”.

• Even though the score of the original distribution is complex, thedistribution of the diffused 𝑋𝑡 is smooth → easy to estimate →easy to generalize.
• Learning is more stable because it uses information of theintermediate distribution 𝑝𝑡 instead of directly learning end-toend mapping from the noise to the source distribution.

55

Toy data[https://github.com/Kei18/tiny-tiny-diffusion]
This is not generating an image of a dinosaur but the shape of thedensity function looks like a dinosaur. Each point corresponds toeach image (𝑋𝑡 ).

56

Score estimation

57

Reverse process:(unknown)(unknown)[Haussmann & Pardoux, 1986]⇒ 𝑌𝑡 ∼ 𝑝𝑇−𝑡

𝒀𝑻ഥ ∼ 𝒑𝟎

Approximated process (generative model):(𝑁(0, 𝐼) is close to 𝑝𝑇 )Theorem (Girsanov’s theorem)⇒ It suffices to estimate the score function 𝛻log(𝑝𝑡 ) as accurate
as possible.

Score matching

58

Reference

Unknown. We cannot calculate.
(𝑋 𝑇−𝑡ത と𝑌𝑡 は同じ分布)

Score matching

59

ReferenceScore matching

Observation (𝑛 data points 𝐷𝑛 = 𝑥𝑖 𝑛𝑖=1 ):Empirical score matching loss:

Can be sampled via OU process

Explicit form is available

60Error analysis of diffusion models• Reverse SDE characterization: Song et al. (2021)[Approximation error analysis]• KL-divergence bound via Girsanov’s theorem: Chen et al. (2022)• Error bound with LSI: Lee et al. (2022a)➢ With smoothness: Chen et al. (2022) and Lee et al. (2022b)• Error propagation with manifold assumption: Pidstrigach (2022)[Generalization analysis]• Wasserstein distance bound: O(𝑛Bortoli (2022)−

1𝑑) with manifold assumption: DeQ:
1. How accurately can we estimate the score functions?
2. How strongly does the estimation error of score functions affectthe final result?

61

Problem setting

62

Assumption 1
The true distribution 𝑝0 is supported on −1,1 𝑑 andwith 𝑠 > 1Τ𝑝 − 1Τ2 + as a density function on −1,1 𝑑 .
Assumption2

is sufficiently smooth on the edge of the
support𝑠 (Ω))
Besov space (𝐵𝑝,𝑞Very smooth
Besov space
Smoothness

Spatial inhomogeneity

Problem setting

63

Assumption 1
The true distribution 𝑝0 is supported on −1,1 𝑑 andwith 𝑠 > 1Τ𝑝 − 1Τ2 + as a density function on −1,1 𝑑 .
Assumption2

is sufficiently smooth on the edge of the
support𝑠 (Ω))
Besov space (𝐵𝑝,𝑞Smoothness

Intuition

Very smooth
Uniformity of smoothness

Besov space
Smoothness

Spatial inhomogeneity

Convergence rate result

𝑇

64

𝑇

Theorem (Estimation error in TV-distance)

Let 𝑇 = 𝑛−𝑂(1) , 𝑇 = 𝑂(log(𝑛)). Then, the empirical risk minimizer 𝑠Ƹ
in DNN satisfies

This is minimax optimal, that is, it holds

Although 𝑠(𝑥,Ƹ 𝑡) is a function with 𝑑 + 1-dimensional input, there appears “𝑑” in thebound instead of 𝑑 + 1. This is because Gaussian convolution induces smoothness.

B-spline basis decomposition

65

Reference
Approximate each term by DNNs

• B-spline decomposition of a Besov function 𝑝0

Cardinal B-spline of order:→ Piece-wise polynomial of order m.
Tensor product B-spline:

Cardinal B-spline interpolation (DeVore & Popov, 1988)

66

Reference• Atomic decomposition:

Wavelet/multi-resolution expansion
j=1

Scale

k=0

𝑠 can be decomposed into𝑓 ∈ 𝐵𝑝,𝑞

𝛼2,1
)

k=2

𝛼1,2
j=2

𝛼1,1
j=1

k=1
such that(where

𝛼0,1

j=1

𝛼2,2
j=2

𝛼2,3

j=3

𝛼2,4
j=4

k=3

(Norm equivalence)

DNN can approximate eachB-spline basis efficiently.

𝑁 terms (should be appropriately chosen depending on 𝑓)(see also Bolcskei, Grohs, Kutyniok, Petersen: Optimal Approximation with Sparsely Connected Deep Neural Networks. 2018)

Proof outline (1)67Approximate each term by DNNs

• B-spline decomposition of a Besov function 𝑝0
Approximation error• Diffused B-spline basis expansion of 𝑝𝑡

Decompose=: 𝐸𝑎𝑗 ,𝑏𝑗 (𝑥, 𝑡)

Diffused B-spline➢ We approximate Diffused B-splines by DNNs.

Approximation error of Diffused B-spline

68

Lemma (Approximation error of diffused B-spline)෠ ℝ𝑑 × ℝ+ → ℝ𝑑 such that
There exists a deep neural network 𝜙:with depth 𝐿 = 𝑂 log 4 𝜖 −1 , width 𝑊𝑖 = 𝑂(log 6 (𝜖 −1 )), sparsity (#of non-zero parameters) 𝑆 = 𝑂(log(𝜖 −1 )), and ℓ∞ -norm bound 𝐵 =𝑂(exp(𝑂(log 2 𝜖 −1 ))) on parameters.

: Deep neural network# of non-zero parameters:≤ O(𝑒 −𝐿 )≤ 𝑁 −𝑠/𝑑Error bound of score
Difficult𝑇

Non-smooth

Easy

𝑡∗
Smooth69

Very smooth𝑇ത𝑡• Bound by diffused B-spline approximation➢ Similar argument is applied to 𝛻𝑝𝑡 :(𝜎𝑡2 = 1 − exp −2𝑡 )• A tighter bound on the smooth part (𝑡 > 𝑡∗ )- Useful for W1 bound.
- Smoothness around theedge (A2) is not requires.

(take 𝑘 = 𝑠 + 1)

Error decomposition

70(Pinsker’s inequality)

Score matching loss

Truncation loss
at 𝑇.

Bias

Truncation loss
at 𝑇.

Variance

𝑇𝑇Low dimensional structure
ℝ𝑑

′
𝑑
ℝ

The support of the target distribution isin a low dimensional subspace.

The estimated distribution is never absolutely
continuous to the target distribution.
→ Wasserstein distance71

𝑾𝟏 -distance convergence rate
Theorem (Estimation error in W1-distance)

For any fixed 𝛿 > 0, by slightly changing the estimator, theempirical risk minimizer 𝑠Ƹ in DNN satisfies

This is also known as minimax optimal (up to 𝛿) [Niles-Weed & Berthet(2022)].
• 𝑑′ appears instead of 𝑑: Diffusion model can avoid curse ofdimensionality.
• The minimax rate of Wasserstein distance is faster than that of TV
distance, which makes it difficult to establish the bound.
➢ We need more precise estimate of the score around 𝑡 = 0.
(TV)

(W1)

72

73

Lemma (tighter bound on W1 distance error)• For large 𝑡, we can estimate the score more accurately.
• For small 𝑡, the error does not propagate so much dueto the term 𝑡.
→ Better rate.

Bound for W1 distance

𝑡∗

𝑇
𝑖
𝑌𝑇−𝑡

𝛻log(𝑝𝑡 )

𝑡

2𝑡 ∗
𝑡𝑖

4𝑡 ∗

(= 2𝑖 𝑡 ∗ )𝑠Ƹ,(negligible)(exp(−𝑇))𝑡∗𝐾ത
𝑇 = 2 𝑡∗74

Transformer

75[Shokichi Takakura, Taiji Suzuki: Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions
with Infinite Dimensional Input. ICML2023]

Properties of Transformer• It can output a value from wide
rage of tokens.
→ Curse of dimensionality?
• It can choose important tokens
depending on input.
→ Can avoid curse of dim!

⋯

𝑌−1

𝑌0

⋮

⋮

𝑌2
⋮

⋯

Self-attention

FNN

We showed minimax optimality to estimate
a sequence-to-sequence function.

Theorem (estimation error)

𝑌1
⋮

⋯ 𝑥−1

𝑥0

𝑥1

𝑥2

𝑌

これはペンです．

⋯

𝐹∘(almost minimax optimal)➢ It achieves polynomial order convergenceeven though input is infinite-dimensional.

𝑋

This is a pen.

Remark on the rate of convergence

76

Remark： Even if the rate is better, the method does
not necessarily achieve better prediction.
Predictive error
Deep learning

Shallow learning

Method blue is betterat this sample size.

Method red is betterat this sample size.

Sample size

Outline
1. Representation ability + Generalization
ability➢Universal approximator➢Depth separation➢Adaptivity of deep learning➢ Inhomogeneity of smoothness➢ Curse of dimensionality➢Foundation models➢ Diffusion model➢ Transformer

2. Optimization ability➢Noisy gradient descent➢Mean field Langevin➢CSQ lowerbound

77

78

Optimization of NN

Representation
ability

Generalization
ability

Optimization

Optimization of NN

79

We should “optimize” the parameters.

𝑊: parameter

Loss function for the 𝑖-th data
point.

Loss function：degree of fit to the data

Loss function optimization(W could be billions dimensional)
Usually, stochastic
gradient descent is used.

OptimalLocal optimality
Objective function of deep learning is non-convex.
Critical point
Local optimal

Global optimal• For linear deep NN, every local optimal is global optimal：
Kawaguchi, 2016; Lu&Kawaguchi, 2017.
※True only for linear NN.
→ Sufficient conditions that a critical point is a global optimal
was also derived by Yun, Sra&Jadbabaie (2018).

• Low rank matrix completion has no spurious local
minimum：Ge, Lee&Ma, 2016; Bhojanapalli, Neyshabur&Srebro, 2016.

80

Loss landscape81

Wide neural network does not have any isolated local minimum.
(a local optimal solution is connected to global optimal solutions)※This does not indicate GD can reach the global optimal.

TheoremSuppose that we are given 𝑛-training data 𝑥𝑖 , 𝑦𝑖 𝑛𝑖=1 , and the loss
function ℓ is convex.
⊤
For two layer NN model 𝑓 𝑎,𝑊 (𝑥) = σ𝑀
𝑚=1 𝑎𝑚 𝜂(𝑤𝑚 𝑥) with
continuous activation function, if the width is not smaller than thedata size (𝑀 ≥ 𝑛), every arcwise connected component of a levelset of the empirical loss 𝐿෠ 𝑎, 𝑊 = 𝑛1 σ𝑛𝑖=1 ℓ(𝑦𝑖 , 𝑓 𝑎,𝑊 (𝑥𝑖 )) contains
the global optimal solution.
[Venturi, Bandeira, Bruna: Spurious Valleys in One-hidden-layer Neural Network Optimization Landscapes.
JMLR, 20:1-34, 2019.]

This does not happen(not connected)

This happens

Overparameterization
Wide neural network does not have spurious
local minima.
Narrow network…

0
Wide network

Since the model complexity is
increased, the initial solution isalready close to the global optimal.

0• Two types of analysis➢ Neural Tangent Kernel (NTK)➢ Mean-field analysis

82

Two regimes

83• Neural Tangent Kernel regime (lazy learning )[Jacot+ 2018][Du+ 2019][Arora+ 2019]➢ 𝑎𝑗 = O(1/ 𝑀)(Xavier initialization/He initialization)• Mean field regime➢ 𝑎𝑗 = Ο(1/𝑀)[Nitanda & Suzuki (2017), Chizat & Bach(2018), Mei, Montanari, & Nguyen (2018)]

Different scaling of initial solution yields
different behavior.

84

NTK: Large scale initialization → features are (almost)freezed.
Optimization trajectoryof first layer parameters
in a 2-layer NN:

True function:[Ba et al., 2022]

Mean field: Small scale initialization → features need to movesignificantly.

ABC-parameterization

85• ABC-parameterization [Yang&Hu, 2021]
𝑛: width(1) Parameterization(𝑤 𝑙 is the actual trainable parameter)

A(2) Initialization(3) Learning rate

B

C

(Appropriate scaling)[Yang&Hu: Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks. ICML2021.]

86

The optimal hyper-parameter in a small size network can be transferred to hugemodel.
It is used to train GPT-3.5. Billions of dollars cost could be saved.

[Yang et al.:Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. arXiv:2203.03466]

https://github.com/microsoft/mutransformers (𝜇P for Transformers)

Implicit regularization

87

Binary classification with exp-loss:
where

Optimization in the space of signed measures.
If we start from small initialization, only neurons that arenecessary for classification “grow up.”
[Chizat&Bach: Implicit Bias of
Gradient Descent for Wide Two-layer
Neural Networks Trained with the
Logistic Loss. COLT2020.]

Optimization dynamics implicitly
regularize the solution.

→ Sparse solution：implicit regularization
The solution converges to the max-margin solution under L1-norm constraint(if the sequence converges to a “global minimizer” direction):

Gradient descent and implicit regularization 88• Dynamics starting from a small initialization converges
to the minimum norm solution.
→ Implicit regularization

Set of solutions
Optimization with gradient
Initialization（around origin）

The “simplest” solution
Deeper network leads to “sparser” solution.

[Gunasekar et al.: Implicit Regularization in Matrix Factorization, NIPS2017][Soudry et al.: The implicit bias of gradient descent on separable data. JMLR2018][Gunasekar et al.: Implicit Bias of Gradient Descent on Linear Convolutional Networks, NIPS2018][Moroshko et al.: Implicit Bias in Deep Linear Classification: Initialization Scale vs Training
Accuracy, arXiv:2007.06738]

Implicit regularization in each regime
Regime89

Implicit regularization

NTK, kernel methodwith early stopping

L2-regularization

Mean-fieldL1-regularization• Deep learning uses several “explicit regularization”.
→ Batch normalization, Dropout，Weight decay，MixUp，...
• On the other hand, the “implicit regularization” induced by
the deep structure and optimization dynamics is also very
important.
→ Benign overfitting, Grokking, Flat-minimum, ...

Grokking/Benign-overfitting

90

Grokking: [Power et al.: Grokking: Generalization beyond overfitting on small
algorithmic datasets. arXiv:2201.02177][Xu et al.: Benign Overfitting and Grokking in ReLU Networks for XOR Cluster Data.
arXiv2310.02541]

It is also called “hidden progress” [Barak et al. 2022].
See also [Meng et al.: Benign Overfitting in Two-Layer ReLU Convolutional
Neural Networks for XOR Data. arXiv2310.01975]

Outline
1. Representation ability➢Universal approximator➢Adaptivity of deep learning➢ Inhomogeneity of smoothness➢ Curse of dimensionality

2. Generalization ability➢Double descent, Benign overfitting foroverparameterized model➢Generalization gap analysis➢ Norm based bound➢ Compression based bound

3. Optimization ability➢Neural Tangent Kernel➢Dynamics in a feature learning regime➢Mean field analysis

91

92

Noisy gradient descent andits global optimality

Noisy gradient descent
The model is not linearly approximated.
We need to solve “non-convex” optimization.

Noise
Gradient descent

SGD is a noisy gradient descent.
Noisy perturbation is helpful to escape local minimum.

93

Sharp minima vs flat minima
It is said that SGD likely stay in “flat local minimum”→ Gives good generalization error.

94

This is not theory,
Just explanation.

Keskar, Mudigere, Nocedal, Smelyanskiy, Tang (2017):
On large-batch training for deep learning: generalization gap and sharp minima.

≅Normal distribution→ Random walk is likely to be
captured in a flat region.

• (criticism) The concept of “flat”
depends on the choice of coordinate
system. (Dinh et al., 2017)•PAC-Bayesian analysis (Dziugaite, Roy,2017)

Smoothing by noisy gradientsmoothing[Kleinberg, Li, and Yuan, ICML2018]

Stochastic gradient ⇒ Noise is added ⇒ Objective is smoothed

SGD optimizes a “smoothed” objective:95

GLD/SGLD

96• Stochastic Gradient Langevin Dynamics (SGLD)(Non-convex)(Gradient Langevin dynamics)
Stationary distribution：
Discretization[Gelfand and Mitter (1991); Borkar andMitter (1999); Welling and Teh (2011)]

GLD:(Euler-Maruyama scheme)

SGLD:

Stochastic gradientGLD as a Wasserstein gradient flow97

: 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 formula

Continuity equation
Continuity equation:
The meaning of this equation(∀𝑓: compact support，𝐶 ∞ -class)d𝑇• Let 𝑇𝑡 be a map generated by the vector field 𝑣𝑡 : 𝑡 𝑥 = 𝑣𝑡 𝑇𝑡 𝑥 .
d𝑡𝑑𝑑• 𝜇𝑡 is the push-forward of 𝜇0 by a map 𝑇𝑡 : R → R : 𝜇𝑡 = 𝑇𝑡# 𝜇0 .
That is, 𝝁𝒕 is the distribution of 𝑻𝒕 (𝒙) where 𝒙 ∼ 𝝁𝟎 .

[continuity equation]

98

Stationary distribution

Stationary distribution of the continuous time dynamics:

The stationary distribution
concentrates around the
optimal solution.

99

Wasserstein gradient flow100

We neglect this term below

By the continuity equation 𝜇𝑡 = −𝛻 ⋅ [𝑣𝑡 𝜇𝑡 ], it holds that=0

Wasserstein勾配流101

In particular, if(GLD)then this is the steepest gradient descent direction
such that=: −𝑣𝑡Fisher divergence:

GLD is the Wasserstein gradient flow to minimize the KL-div from 𝝁∗ .

Log-Sobolev inequality

102

Stationary distribution:

Def (log-Sobolev inequality)
There exists a constant 𝛼 > 0 such thatfor any probability measure 𝜈 (absolutely-continuous w.r.t. 𝜇∗)

KL-div

Geometric ergodicity

E.g.：• Quadratic+Bounded• Weak Morse function

Fisher-div

𝜇𝑡 : distribution of 𝑋𝑡(by log-Sobolev)

Linear convergence w.r.t. KL-div[Bakry, Gentil, and Ledoux: Analysis and Geometry of Markov Diffusion Operators. Springer, 2014. Th. 5.2.1]

103
Sufficient condition for log-Sobolev inequality

Strongly convex (Bakry-Emery criterion):𝛻𝛻 ⊤ 𝐿 𝑥 ≽ 𝜇𝐼
Ex.: OU-process．𝐿 𝑥⇒[Bakry and Émery, 1985]𝑥2= ⇒𝜇=1
2

Bounded perturbation lemma (Hollley-Stroock):
Suppose thatand 𝜇 satisfies 𝛼′-LSI, then𝜇∗ satisfies 𝛼-LSI with

[R. Holley and D. Stroock. Logarithmic sobolev inequalities and stochastic Ising models. Journal of statistical physics, 46(56):1159–1194, 1987.]

Ex.: 𝐿 𝑥 = ℓ 𝑥 + 𝜆1 𝑥 2 and ℓ 𝑥≤ 𝐵, then 𝜇∗ satisfies LSI with 𝛼 = 2𝜆𝜆1exp(−4𝐵/𝜆).

References• Finite dimensional Langevin dynamics:➢ Convergence in low (convex case): Dalalyan and Tsybakov, 2012;
Dalalyan, 2016; Durmus and Moulines, 2015, ..
➢ Non-convex Optimization: Raginsky et al., 2017; Xu et al., 2018;
Erdogdu, Mackey and Shamir, 2018
➢ Log-Sobolev inequality: Vempala and Wibisono, 2019.

• Infinite dimensional Langevin dynamics:➢ Continuous time:• Existence & Uniqueness of invariant measure: Da Prato and Zabczyk,1992;
Maslowski, 1989; Sowers, 1992.
• Geometric ergodicity: Jacquot and Royer, 1995; Shardlow, 1999; Hairer, 2002,
Its explicit rate: Goldys and Maslowski, 2006.

➢ Discrete time:• Weak approximation rate of discretized scheme: Hausenblas, 2003;
Debussche, 2011; Bréhier, 2014; Bréhier and Kopec 2016.

Other topics (MCMC in Hilbert space):• preconditioned Crank–Nicolson (pCN): Hairer et al., 2014; Eberle, 2014;
Vollmer, 2015; Rudolf and Sprungk, 2018.
• Metropolis-Adjusted Langevin Algorithm (MALA): Durmus and Moulines,
2015; Beskos et al., 2017.

104

Related work: Graduated optimization

105• Graduated non-convexity
Blake and Zisserman: Visual reconstruction, volume 2. MIT press Cambridge, 1987.

• Convolution with Gaussian kernel
Z. Wu. The effective energy transformation scheme as a special
continuation approach to global optimization with application to molecularconformation. SIAM Journal on Optimization, 6(3):748-768, 1996.

• Graduated optimization
Hazan, Levy, and Shalev-Shwartz: On graduated optimization forstochastic non-convex problems. International conference on machinelearning, pp. 1833-1841, 2016.

෠𝛿 𝑥 = E➢ 𝜎-nice property: 𝐿𝑢∼𝑈(B(R𝑑 )) [𝐿(𝑥 + 𝛿𝑢)]
➢ Polynomial time convergence.
Survey:Mobahi and Fisher III. On the link between gaussian homotopy
continuation and convex envelopes. Energy Minimization Methods
in Computer Vision and Pattern Recognition, pp. 43-56, 2015.

106

Optimization theory
in mean field regime

2-layer NN in mean-field scaling• 2-layer neural network:

Non-linear with respect to parameters 𝑟𝑗 , 𝑤𝑗

𝑀
𝑗=1

.

where 𝑋 (𝑗) = 𝑟𝑗 , 𝑤𝑗 andLoss function:loss

Non-convex

L2 regularization

107

Application of GLD

108

Noisy gradient descent update:⇔

Does it converge?
Naïve application of existing theory in gradient Langevin dynamics yields

iteration complexity to achieve 𝜖 error.
→ Cannot be applied to wide neural network.

[Raginsky, Rakhlin and Telgarsky, 2017; Xu, Chen, Zou, and Gu, 2018; Erdogdu, Mackey and Shamir, 2018; Vempala and WibisonoMean field limit

109

𝑀

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

★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).

…

General form of mean field LD

110convex + strictly convex = strictly convex

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

𝐹

Gradient

GLD:,

Definition (first variation)
𝛿𝐹
The first variation 𝛿𝜇: 𝒫 × ℝ𝑑 → ℝ is defined as a continuous functional such as

MF-LD to optimize mean field NN

111

Loss function:(distribution of 𝑋𝑘 )Neuronℎ𝑥 (⋅)

Discrete time MFLD:𝑥

Proximal Gibbs measure

112

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.

Entropy sandwich
Proximal Gibbs measure:

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

Optimality condition113

Duality (informal)

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

Primal=(Fenchel’s duality theorem)[Rockafellar (1967)]

Dual

Primal-Dual variable correspondence:
(P)

(D)

(P)

Duality gap and divergence:𝐹••

Gradient(optimality condition)

Convergence rate

115

Proximal Gibbs measure:

Assumption (Log-Sobolev inequality)c.f., Polyak-Lojasiewicz condition𝑓 𝑥 − 𝑓 𝑥 ∗ ≤ 𝐶 𝛻𝑓 𝑥 2

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

Theorem (Linear convergence) [Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (2022)]

If 𝑝𝜇𝑡 satisfies the LSI condition for any 𝑡 ≥ 0, thenrate ofextension
convergenceis characterized
by LSIThis isThea non-linear
of well known
GLD convergence
analysis.

Proof outline of convergence116• MF-LD obeys the following nonlinear Fokker-Planck
equation:
Mass: 𝜇𝑡 (𝑥)Then,[Continuity equation]=: −𝑣𝑡(∵continuity
equation)

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

117

L2-regularized loss function for mean field 2-layer NN:
where

Proximal Gibbs:

Bounded perturbation Gaussian

Bounded (≤ 𝐵)

Strongly convex

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

LSI with a constant 𝛼 with

∵ Bakry-Emery criterion (1985) and Holley-Strook bounded perturbation lemma (1987)

Other applications

118

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)

119

Finite particles &discrete time algorithm

120

We have obtained a convergence of
infinite width and continuous time dynamics.
Question:
Can we evaluate a finite particles & discrete time
approximation errors?
(distribution of 𝑋𝑡 )Neuron(Finite particle approximation)𝑥(vector field)

Difficulty121• 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Practical algorithm

𝑀 particles(time discretization)

𝑖
𝑋𝑘

𝑀
𝑖=1

and

where

(space discretization)(stochastic gradient)➢ 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:

𝛿𝐹 𝜇
𝛻
𝛿𝜇

→ 𝑣𝑘𝑖

𝑘

122

Convergence analysis

123: 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 that
𝐎(𝟏/𝑴)

Naïve bound:Time
discr.

Space
discr.

Stochastic
approx.

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.
NeurIPS2023]

Uniform log-Sobolev inequality

124

Reference
(2)

𝑋𝑘
(𝑁)
𝑋𝑘

(1)
𝑋𝑘

𝒳𝑘 =

𝑖
𝑋𝑘

𝑀

𝑀

∼ 𝜇𝑘 : Joint distribution
𝑖=1
of 𝑀 particles.
(𝑴)

Potential of the joint distribution 𝝁𝒌

on ℝ𝒅×𝑴 :where➢ The finite particle dynamics is the Wasserstein gradient flow that minimizes.

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

where
Recall[Chen, Ren, Wang. Uniform-in-time propagation of chaos
for mean field Langevin dynamics. arXiv:2212.03050, 2022.]

Computational complexity

125

SG-MFLD(finite sum),(stochastic gradient)(Mini-batch size = 𝐵)

Iteration complexity:Time
discr.

Space
discr.

Stochastic
approx.

➢ By
Approximation
errors are uniform in time.
setting,complexity
becomes on 𝑴 (number of➢ the
No iteration
exponential
dependency
neurons).
to achieve 𝜖 + 𝑂(1/(𝜆2 𝛼𝑁)) accuracy.
➢𝐵=

1/(𝜆2 𝛼𝜖) is the optimal mini-batch size. → 𝑘 = 𝑂 log 𝜖 −1 Τ𝜖 .

Numerical experiment
𝑀→∞

Test error v.s. Number of steps(regularization term: 𝑟 𝑥 = 𝑥 2 )

126

127

Generalization error analysis

Generalization error analysis
So far, we have obtained convergence of
MFLD.
⇒ How effective is the feature learning of
MFLD in terms of generalization error?

• Benefit of feature learning?

Neural network vs Kernel method(NTK vs mean field)

128

Classification task

129

Problem setting (classification):

(-1)
(+1)

Loss function and model:➢Logistic loss: ℓ 𝑦𝑓 = log(1 + exp(−𝑦𝑓))➢Tanh activation: ℎ𝑥 𝑧 = 𝑅ത ⋅ [tanh 𝑥1 , 𝑧 + 𝑥2 + 2 ⋅ tanh 𝑥3 ]/3Assumptions
Objective of MFLD:

where 𝜈 = 𝑁(0, 𝜆/(2𝜆1 )).

Assumption

KL-regularization

There exists 𝜇 ∗ such that
1. KL 𝜈, 𝜇 ∗ ≤ 𝑅,
2. 𝑌𝑓𝜇∗ 𝑍 ≥ 𝑐0 (a. s. ),for some constants 𝑅, 𝑐0 > 0.
(+ classification calibration condition)

The Bayes classifier is attained by 𝝁∗ with a bounded KL-div from 𝝂.
𝑓𝜇∗ (𝑧)𝑐0

𝑌 = −1supp(𝑃𝑍 )𝑌=1
supp(𝑃𝑍 )𝑧130

Main theorem

131

Theorem 1
Suppose that 𝜆 = Θ(1Τ𝑅), then it holds that

Class. error𝑅ത 2 𝑅O𝑛with probability 1 − exp −𝑡 .

Existing bound: Chen et al. (2020); Nitanda, Wu, Suzuki (2021)

Class. Error ≤ O

1.
𝑛(Rademacher complexity bound)• Our bound provides fast learning rate (faster than 1/ 𝑛).
O 𝑅Τ𝑛 ≪ O 1Τ 𝑛
• 𝜇∗ : KL 𝜈, 𝜇∗ ≤ 𝑅, 𝑌𝑓𝜇∗ 𝑍 ≥ 𝑐0
• ℎ𝑥 𝑧 = 𝑅ത ⋅ [tanh 𝑥1 , 𝑧 + 𝑥2 + 2 ⋅ tanh 𝑥3 ]/3Main theorem 2

132

Theorem 2
Suppose that 𝜆 = Θ(1/𝑅) andthen it holds thatwith probability

𝑅

Theorem
𝔼[Class. Error] ≤ O.
𝑛1:Theorem
𝔼[Class. Error] ≤ O exp(−O(𝑛/𝑅2 )) if 𝑛 ≥ 𝑅2 .
2:
If we have sufficiently large training data,
we have exponential convergence of test error.
We only need to evaluate 𝑅 to obtain a test error bound.

Proof sketch
Expected loss

Local Rademacher complexity yields the following bound:

Adaptive to the KL-divergence

By setting 𝜆 = O(1/𝑅),• The term (I) gives the first bound: O 𝑅Τ𝑛 .
• The term (II) gives the second bound: O exp(−O(𝑛/𝑅2 )) .

133

Example: k-sparse parity problem

134• 𝑘-sparse parity problem on high dimensional data

➢ 𝑍 ∼ Unif( −1,1 𝑑 )➢ 𝑌 = ς𝑘𝑗=1 𝑍𝑗(up to freedom of rotation)※ Assume we don’t know which coordinate corresponds to 𝑍𝑗 .

Q: Can we learn sparse 𝒌-parity with GD?
Is there any benefit of neural network?

Complexity to learn XOR function (𝑘 = 2)

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

𝑘 = 2: XOR problem
𝑑 = 3, 𝑘 = 2

Generalization bound

135

ReminderSuppose that there exists 𝜇 ∗ such that𝜇∗ : KL 𝜈, 𝜇∗ ≤ 𝑅, 𝑌𝑓𝜇∗ 𝑍 ≥ 𝑐0 (perfect classifier with margin 𝑐0 )Then,
𝑅
Theorem 1:
𝔼[Class. Error] ≤ O.
Theorem 2:𝑛𝔼[Class. Error] ≤ O exp(−O(𝑛/𝑅2 )) if 𝑛 ≥ 𝑅2 .

We can evaluate 𝑅 required for the 𝑘-sparse parity problem:

Lemma
For the 𝑘-parity problem, we may take

Generalization
error
bound
Our analysis
provides136

Corollary (Test accuracy of MFLD)•

better sample complexity• discrete-time/finite-width analysis• Setting 1: 𝑛 > 𝑑 • 𝑑 and 𝑘 are “decoupled.”➢ Test error (classification error) = 𝐎(𝒅Τ𝒏)• Setting 2: 𝑛 > 𝑑2
➢ Test error (classification error) = 𝐎(exp(−𝒏/𝒅𝟐 ))(Computational complexity is exp O 𝑑(But, can be relaxed to O(1) if X is anisotropi

These are better than NTK (kernel method);
Sample complexity of NTK 𝒏 = 𝛀 𝒅𝒌 vs NN 𝒏 = 𝐎(𝒅)Trade-off between computational complexity and sample complexity.

Discussion

137• The CSQ lower bound states that O 𝑑𝑘−1
sample complexity is optimal for methods withpolynomial order computational complexity.
[Abbe et al. (2023); Refinetti et al. (2021); Ben Arous etal. (2022); Damian et al. (2022)]• On the other hand, our analysis is about full-batch GD.
Minibatchsize# of iterations

Samplecomplexity

Our analysis

𝒏

𝒆𝒅

𝒅

SGD
(CSQ-lower bound)

1

𝑑𝑘−1

𝑑𝑘−1

We obtain a better sample complexity than O(𝑑 𝑘−1 ) with higher
computational complexity.
→ We can obtain a polynomial order method with MFLD for
anisotropic input.

CSQ algorithm138

Def (Correlational Statistical Query (CSQ) algorithm)[Ben-David, Itai, Kushilevitz, 1995; Kearns, 1998; Bshouty, Feldman, 2002]

A CSQ algorithm can access the data only via queries 𝜙: ℝ𝑑 → [−1,1] andreturns 𝑔 ∈ ℝ𝑑 with tolerance 𝜏 such thatEx.

Online SGD for a squared loss:(CSQ)• Boolean case:𝒌-parity:

where 𝑆 = 𝑘.

• Gaussian case:
Single index model:
where 𝑔: ℝ → ℝ and 𝑤 ∈ ℝ𝑑 .

CSQ lower bound139

For the Gaussian single index model, the information exponent plays
an important role.
Hermite polynomial expansion of the link function:

Def (information exponent [Ben Arous, Gheissari, Jagannath, 2021])

The computational complexity of a CSQ algorithm is lower bounded as:
Theorem (CSQ lower bound [Abbe, Boix-Adser`, Misiakiewicz, 2023])
A CSQ algorithm with error tolerance 𝜏 requires at least 𝑁 queries2∘መመ
to obtain an estimator 𝑓 s.t. 𝔼 𝑓 − 𝑓≤ 0.1 where(we suppose 𝑘 ∗ > 2)

Note that the gradientcomputation at each
iteration consumes 𝑂(𝑑)queries. Thus, 𝑑𝑘−1
iterations are enough.

Recent progress

140• SGD with smoothing operation achieves the Gaussian
optimal rate:
Damien et al.: Smoothing the Landscape Boosts the Signal for SGD Optimal
Sample Complexity for Learning Single Index Models. NeurIPS2023.

• Near optimal complexity of SGD to learn XOR problem:
Glasgow: SGD Finds then Tunes Features in Two-Layer Neural Networks
with near-Optimal Sample Complexity: A Case Study in the XOR problem.
ICML2024.

• Optimal SQ sample complexity to learn Gaussian
single index model with the “generative” information
exponent:
Damian, Pillaud-Vivien, Lee, Bruna: The Computational Complexity of
Learning Gaussian Single-Index Models. arXiv:2403.05529.

141

The setting of 𝑘 ∗ =1.

Feature learningwith one-step gradient descent[Ba, Erdogdu, Suzuki, Wang, Wu, Yang: High-dimensional Asymptotics of
Feature Learning: How One Gradient Step Improves the Representation.
NeurIPS2022]

Gradient descent and kernel alignment

142

Question：Can we obtain “good” features from data by updating thefirst layer parameter 𝑊 by gradient descent?
Result：GD with large step size can extract the leading term of the true
function. Especially, for the single index model (𝑓 ∗ 𝑥 = 𝜎 ∗ (〈𝑥, 𝑤 ∗ 〉)), thepredictive risk provably outperforms random feature methods.
→ Kernel alignment, feature learning.
We consider the proportional limit (𝒏, 𝒅, 𝑵 → ∞), and evaluate predictive
risk of one-step GD.

➢ 𝜼 = 𝚯( 𝑵) can outperform random feature models.
➢ 𝜼 = 𝚯(𝟏) can outperform the initial setting of 𝑊.
➢ 𝜼 = 𝒐(𝟏) does not improve the performance.
Gaussian equivalence property + Random matrix theory→ Exact risk evaluation.

Dot: Simulation
Solid line: Theory

Related work

143

The first few step of GD with large learning rate can extract
informative features.

• Staircase function[Abbe et al., NeurIPS2021; Abbe et al., arXiv2202.08658]

Small number of gradient descent can extract
nonlinear features to estimate “staircase” function.
The trained features for GD can outperform random
feature model.

• Benign overfitting with feature learning[Cao et al., arXiv:2202.06526; Frei et al., arXiv:2202.05928]

Gradient descent in two-layer NN can yield benign
overfitting and achieves almost the Bayes error in binary

Problem setting

144

Observation model:

where 𝑥𝑖 ∼ 𝑁 0, 𝐼 , 𝜖𝑖 ∼ 𝑁(0,1), and 𝑥𝑖 ∈ 𝐑𝑑 .
➢ We fit 2-layer NN of mean field scaling:(∵ 𝑎𝑖 = 𝑂𝑝 (1/ 𝑁))
Mean field regime 𝑂(1/𝑁)

where 𝑎𝑖 ∼ 𝑁(0,1/𝑁) and 𝑊𝑖𝑗 ∼ 𝑁(0,1/𝑑).
var

Empirical risk:

varPredictive risk:

Question: Can we provably improve the predictive risk by gradient descen
We analyze the risk especially for the single index model:

145
Feature learning with optimization guarantee

We consider the proportional limit (𝒏, 𝒅, 𝑵 → ∞ with 𝑛Τ𝑑 → 𝜓1 , 𝑁Τ𝑑 → 𝜓2 ).
It allows to derive precise risk.

We evaluate predictive risk of one-step GD.

Take home message:
GD with Large step-size can outperform any
random feature model by only one-step update.
[Outline of our result]➢ 𝜼 = 𝚯( 𝑵) can get out of NTK regime andoutperform random feature models.
➢ 𝜼 = 𝚯(𝟏) can outperform the initial setting of 𝑊.
➢ 𝜼 = 𝒐(𝟏) does not improve the performance.

Dot: Simulation
Solid line: Theory

Ridge regression with RF

146

Feature learning vs Random feature
Random features (without feature learning):• Conjugate kernel at initialization:
Precise asymptotics has• NTK (Neural tangent kernel):been extensively studied.
(e.g.,[Louart, Liao, and Couillet,
2018; Mei and Montanari,2019])RF ∈ {CK,NTK}

Trained feature:

Limitation of RF

147(1) Random feature models and(2) GD updates with small learning rate
can learn only linear functions in the proportional
limit.
[El Karoui (2010); Ghorbani et al. (2019), Hu and Lu (2020), ...]
Theorem (Lower bound of predictive risk for RF)
If the step size is not large 𝜂 = Θ(1), then for any finite number steps 𝑡, we
haveNonlinear part cannot be trained!

where 𝑃≤1 is the projection operator in 𝐿2 (𝑃𝑋 ) to the subspaceconsisting of linear functions and constants.
Remark: The same is true for “rotational invariant kernel” [El Karoui (2010)].

This is because in high dimensional setting, a central limit theorem yields(linear function;
Gaussian equivalence)Effect of large step-size update

148

Theorem: almost rank 1
Spectral distribution of 𝑊1

Spectum of 𝑾𝟎(remains unchanged)

Estimated signal
Alignment with the linear
component of 𝑓 ∗ .
(feature learning)Improvement over the Initial CK
Large step size GD update

149

True function

Best linear model(random feature, small step size GD)

Initial kernelSpace of linear functions• 𝜂 = Θ( 𝑁) (large learning rate):
Known as maximal update parameterization (𝜇P) [Yang and Hu,
2020].

(measure for model misspecification)is assumed.

••𝜏 ∗ = 0 if 𝜎 = 𝜎 ∗ = erf.
𝜏 ∗ ≪1 if 𝜎 = 𝜎 ∗ = tanh.

𝑛Τ → 𝜓 , 𝑁Τ → 𝜓12𝑑𝑑Large learning rate yields feature learning and can bebetter than the small step size regime if 𝜏 ∗ ≪ 𝑃>1 𝑓 ∗ 2 .

Implications
Corollary150

If 𝜎 = 𝜎 ∗ = erf, then 𝜏 ∗ = 0.
In particular, we have 𝑅𝑊1 𝜆 = Θ 𝜓1−1 = Θ 𝑑/𝑛 .

Small step size

Lower bound for RF models

Θ(𝑑/𝑛)

Outperform any RF models

Predictive risk of ridge regression on CK obtained by one step GD (empiricalsimulation, 𝑑 = 1024): brighter color represents larger step size scaled as 𝜂 = 𝑁 𝛼
for 𝛼 ∈ [0,1/2]. We chose 𝜎 = 𝜎 ∗ = erf, 𝜓2 = 2, 𝜆 = 10−3 , and 𝜎𝜖 = 0.1.

Summary151• Representation/Generalization ability➢Depth separation➢Adaptivity of deep learning: separation between linear(shallow) and deep methods• Optimization ability➢Overparameterization➢Noisy gradient descent: a near global optimum✓ Estimation error separation between kernel and deeplearning➢Mean field Langevin➢CSQ lower boun

Deep learning theory that makes DL white box thatcan be controllable.
It would reveal the essence of “good learning system”
which would be useful to develop methods beyond DL.