Feature learning via mean-eld neural networks andanisotropic features
Taiji Suzuki

1,2

3

2,4

Joint work with Denny Wu , Atsushi Nitanda

1

Graduate School of Information Science and Technology, the University of Tokyo,

2

RIKEN Center for Advanced Intelligence Project,

3
4

1,2

, Kazusato Oko

Center for Data Science, New York University,

Faculty of Computer Science and Systems Engineering, Kyushu Institute of Technology.

Abstract
Neural network in the mean-eld regime is known to be capable ofthe kernel (NTK) counterpart.

feature learning, unlike

Recent works have shown that mean-eld neural networks

mean-eld Langevindynamics (MFLD). However, all existing guarantees for MFLD only considered the optimization eciency, and it is unclear if this algorithm leads to improved generalization and sample
can be globally optimized by a noisy gradient descent update termed thecomplexity due to the presence of feature learning. To ll this gap, in this work we study thestatistical and computational complexity of MFLD in learning a class of binary classication
problems. Unlike existing margin bounds for neural networks, we avoid the typical norm control via the perspective that MFLD optimizes the

distribution of parameters rather than theparameter itself; this leads to an improved analysis of the sample complexity and convergencerate. We apply our general framework to the learning of k -sparse parity functions, where we
prove that unlike kernel methods, two-layer neural networks optimized by MFLD achievesa sample complexity where the degree k is decoupled from the exponent in the dimension
dependence.
In addition to that, we consider a setting where the input feature has anisotropic property,
that is, it has a covariance that is not the identity matrix.

Then, we show that such ananisotropic property helps reduce generalization error and computational complexity for the
MFLD. This implies that the anisotropic property of the real world dataset would help thetraining of neural network in terms of both statistical and computational performances.

1 Introduction
Mean-eld Langevin dynamics.

The optimization dynamics of two-layer neural networks

in the mean-eld regime can be described by a nonlinear partial dierential equation of the
distribution of parameters (Nitanda and Suzuki, 2017; Chizat and Bach, 2018; Mei et al., 2018;
Rotsko and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020).

Such a description has

multiple advantages: (i) global convergence guarantees can be obtained by exploiting convexity of
the loss function, and (ii) the parameters are allowed to evolve away from initialization and learn
informative features, in contrast to the neural tangent kernel (lazy) regime (Jacot et al., 2018).
Among the gradient-based optimization algorithms for mean-eld neural networks, the mean-eld Langevin dynamics (MFLD) (Mei et al., 2018; Hu et al., 2019) is particularly attractive due tothe recently established quantitative optimization guarantees. MFLD arises from a noisy gradientdescent update on the parameters, where Gaussian noise is injected to the gradient to encourageexploration.

It has been shown that MFLD globally optimizes an entropy-regularized convex

functional in the space of measures, and for the innite-width and continuous-time dynamics,
the convergence rate is exponential under suitable isoperimetric conditions (Nitanda et al., 2022;
Chizat, 2022). Furthermore, uniform-in-time estimates of the particle discretization error have also
been established (Suzuki et al., 2023a; Chen et al., 2022), meaning that optimization guarantees
for the innite-dimensional problem can be eectively translated to a nite-width neural network.

1However, existing analyses of MFLD only considered the optimization of neural networks; thisalone does not demonstrate the benet of mean-eld regime nor the presence of feature learning.
Therefore, an important problem is to characterize the generalization of the learned models, andprove ecient sample complexity guarantees.

The goal of this work is to address the following

question.

Can we show that neural network + simple noisy gradient descent (MFLD) eciently learns aninteresting class of functions, with a better rate of convergence compared to the lazy regime?

Learning sparse parity functions.

One particularly relevant learning task is the

k -sparse

parity problem, where the response y is given by the sign of the product of k coordinates of the
input (on hypercube); as a special case, setting k = 2 recovers the classical XOR problem. Whenk ≪ d, this target function is low-dimensional, and hence we expect feature learning to be benecialin that it can zoom in to relevant subspace. In contrast, for kernel methods (including neuralnetworks in the lazy regime) which cannot adapt to such structure, it has been shown that a samplekcomplexity of n = Ω(d ) is unavoidable (Ghorbani et al., 2019; Hsu; Abbe et al., 2022).
For the XOR case (k = 2), recent works have shown that neural networks in the mean-eldregime can achieve a sample complexity of n = O(d/ϵ) (Wei et al., 2019; Chizat and Bach, 2020;
Telgarsky, 2023), which indeed improves upon the NTK complexity (Ji and Telgarsky, 2019). However, all these results directly assumed convergence of the dynamics (t → ∞) with no iteration
complexity.

Moreover, Wei et al. (2019); Chizat and Bach (2020) directly analyzed the innite-width limit, and while Telgarsky (2023) provided a nite-width characterization, the dynamics isdrestricted to the low-rotation regime, and a very large number of particles N = O(d ) is required.
Lastly, these analyses are specialized to XOR, and do not directly generalize to the k -parity setting.

1.1

Our Contributions

In this work, we bridge the aforementioned gap by presenting a simple and general framework toestablish sample complexity of MFLD in learning binary classication problems. We then apply
this framework to the sparse k -parity problem, and obtain improved rate of convergence for the
fully time- and space-discretized algorithm. More specically, our contributions can be summarized
as follows.

• We present a general framework to analyze MFLD in the learning of binary classication tasks.
Our framework has two main ingredients: (i) an annealing procedure that applies to common
classication losses that removes the exponential dependence on regularization parameters in
the logarithmic Sobolev inequality, and (ii) a novel local Rademacher complexity analysis for the
distribution of parameters optimized by MFLD. As a result, we can obtain learning guarantees
for the mean-eld neural network in discrete-time and nite-width settings.

• We apply our general framework to the k -sparse parity problem, and derived learning guarantees
2with improved rate of convergence and dimension dependence. Specially, in the n ≍ d regime
we obtain exponentially converging classication error, whereas in the n ≍ d regime we achieve
linear dimension dependence. Note that this improves upon the NTK analysis (which gives aksample complexity of n = Ω(d )) in that it decouples the degree k from the exponent in thedimension dependence. Our theoretical results are supported by empirical ndings.

2 Problem Setting
Throughout this paper, we consider a classication problem given by the following model:

Y = 1A (Z) − 1Ac (Z) ∈ {±1}Z = (Z1 , . . . , Zd ) is the input random variable on Rd and 1A is the indicator functiondcorresponding to a measurable set A ∈ B(R ), i.e., 1A (Z) = 1 if Z ∈ A and 1A (Z) = 0 if Z ̸∈ A.
n
Let PZ be the distribution of Z . We are given input-output pairs Dn = (zi , yi )i=1 independently
whereidentically distributed from this model as training data.

Then, we construct a binary classierthat predicts the label for the test input data as accurate as possible. To achieve this, we learn a
two-layer neural network model in the mean-eld regime via the mean-eld Langevin dynamics.

2

One important problem setting for our analysis is the k -sparse parity problem dened as follows.

√d
Example 1 (k-sparse parity √problem). PZ is the uniform distribution on the grid {±1/ d} andA = {ζ = (ζ1 , . . . , ζd ) ∈ {±1/ d}d | ζ1 · · · ζk > 0}1 .

As a special case, k = 2 (XOR) has been extensively studied (Wei et al., 2019; Telgarsky, 2023).

Mean-eld two-layer network.

Given input z , let hx (z) be one neuron in a two-layer neural
d+1+1dened as

network with parameter x = (x1 , x2 , x3 ) ∈ R

hx (z) = R̄[tanh(z ⊤ x1 + x2 ) + 2 tanh(x3 )]/3,
where R̄ ∈ R is a hyper-parameter determining the scale of the network. We place an extra tanh

activation for the bias term x3 ∈ R because the boundedness of hx is required in the convergenced¯d¯¯ = d + 2 and B(Rd¯)
analysis. Let P be the set of probability measures on (R , B(R )) where d¯
dis the Borel σ -algebra on R and Pp be the subset of P such that its p-th moment is bounded:pEµ [∥X∥ ] < ∞ (µ ∈ P). The mean-eld neural network is dened as an integral over neurons hx ,fµ (·) =∫hx (·)µ(dx),for µ ∈ P . To evaluate the performance of fµ , we dene the empirical risk and the population riskasL(µ) := n1∑ni=1 ℓ(yi fµ (zi )),L̄(µ) := E[ℓ(Y fµ (Z))],

respectively, where ℓ : R → R≥0 is a convex loss function. In particular, we consider the logistic
loss ℓ(f, y) = log(1 + exp(−yf )) for y ∈ {±1} and f∈ R. To avoid overtting, we consider a:= L(µ) + λEX∼µ [λ1 ∥X∥2 ], where λ, λ1 ≥ 0 are regularization
parameters. One advantage of this mean-eld denition is that fµ is a linear with respect to µ,
and hence the functional L(µ) becomes a convex functional.

regularized empirical risk F (µ)

Mean-eld Langevin dynamics.

We optimize the training objective via MFLD, which is givenby the following stochastic dierential equation:

dXt = −∇√δF (µt )(Xt )dt + 2λdWt ,δµµt = Law(Xt ),(1)

where X0 ∼ µ0 , Law(X) denotes the distribution of the random variable X and (Wt )t≥0 is the ddimensional standard Brownian motion. Readers may refer to Theorem 3.3 of Huang et al. (2021)
δF (µt )for the existence and uniqueness of the solution. Here,is the rst variation of F .
δµDenition 1. For a functional G : P → R, the rst-variation δGδµ (µ) at µ ∈ P is a continuousfunctional P × Rd→ R satisfying lim G(ϵν+(1−ϵ)µ)
=
ϵ

∫ δG

δµ (µ)(x)d(ν − µ) for any ν ∈ P .

ϵ→0In our setting, we have∑n ′δF (µ)
12
i=1 ℓ (yi fµ (zi ))yi hx (zi ) + λ(λ1 ∥x∥ ). It is known thatδµ (x) = n
2

the Fokker-Planck equation of the SDE (1) is given by

]
[
(
)]
[
(µt )
(µt )
= ∇ · µt ∇ λ log(µt ) + δFδν.
∂t µt = λ∆µt + ∇ · µt ∇ δFδν

(2)

Then, we can verify that this is equivalent to the Wasserstein gradient ow to optimize the following
entropy regularized risk (Mei et al., 2018; Hu et al., 2019):

L(µ) = F (µ) + λEnt(µ) = L(µ) + λKL(ν, µ) + (const.)
where KL(ν, µ) =∫(3)log(µ/ν)dµ is the KL divergence between ν and µ, and ν is the Gaussian

distribution with mean 0 and variance I/(2λ1 ), i.e., ν = N (0, I/(2λ1 )).

1 We present the axis-aligned setting for conciseness, but the same result holds under orthogonal transforms.
2 This should be interpreted in a weak sense, that is, for any continuously dierentiable function ϕ with a compact

∫

∫

∫ ∫

(µt )
)dµτ dτ .
support, ϕdµt − ϕdµs = − st ∇ϕ · (∇ log(µt ) − ∇ δFδν3

For a practical algorithm, we need to consider a space- and time-discretized version of the∑Nµt by an empirical measure µX = N1 i=1 δXi
i N
d¯
i N
d¯
corresponding to a set of nite particles X = (X )i=1 ⊂ R . Let Xτ = (Xτ )i=1 ⊂ R be
MFLD, that is, we approximate the solution

N particles at the τ -th update (τ ∈ {0, 1, 2, . . . }), and dene µτ = µXτ as a nite particle

approximation of the population counterpart. Then, the discretized MFLD is dened as follows:
X0i ∼ µ0 , and Xτ is updated as

Xτi +1 = Xτi − η∇iwhere η > 0 is the step size, and ξτ√δF (µτ ) i(Xτ ) + 2ληξτi ,δµ(4)∼i.i.d. N (0, I). This is the Euler-Maruyama approximationof the MFLD with a discretized measure; we present the discretization error bounds in the next
section.

3 Main Assumptions and Theoretical Tools
In this section, we introduce the basic assumptions and technical tools for our analysis.

Condition on the loss function

To derive the convergence of the classication error, we assume

that the loss function satises the following condition.

Assumption 1. The convex loss function ℓ : R → R≥0 satises the following conditions:• ℓ is rst order dierentiable, its derivative is Lipschitz continuous and its derivative is bounded
′
′ ′
′
′
by 1: |ℓ (x) − ℓ (x )| ≤ C|x − x | and supx |ℓ (x)| ≤ 1.
• ℓ is monotonically decreasing, and is classication calibrated: ℓ′ (0) < 0 (Bartlett et al., 2006).
• ψ(u)−1 := ℓ(0) − (ℓ(u) − uℓ′ (u)) > 0 for any u > 0.
This standard assumption is satised by several loss functions such as the logistic loss.

We

remark that the rst assumption is used to show the well-denedness of the mean-eld Langevindynamics and derive its discretization error, and also to obtain a uniform generalization error boundthrough the classical contraction argument (Boucheron et al., 2013; Ledoux and Talagrand, 1991).
The second and third assumptions are used to show the convergence of classication error of ourestimator.

Logarithmic Sobolev inequality.

Nitanda et al. (2022); Chizat (2022) showed that the con-vergence of MFLD crucially relies on properties of the proximal Gibbs distribution whose density()
1 δF (µ)
is given by pµ (X) ∝ exp −(X), where µ ∈ P . By the smoothness of the loss function
λ δµ(Assumption 1) and the tanh activation, we can show that the objective L has a unique solutionµ∗ which is also a proximal Gibbs measure of itself (Proposition 2.5 of Hu et al. (2019)). The∗
next question is how fast the solution µt converges to the optimal solution µ . As we will see, theconvergence of MFLD heavily depends on a logarithmic Sobolev inequality (LSI) on pµ .

Denition 2 (Logarithmic Sobolev inequality). Let µ be a probability measure on (Rd , B(Rd )). µd2satises the LSI with constant α > 0 if for any smooth function ϕ : R → R with Eµ [ϕ ] < ∞,Eµ [ϕ2 log(ϕ2 )] − Eµ [ϕ2 ] log(Eµ [ϕ2 ]) ≤ α2 Eµ [∥∇ϕ∥22 ].
This is equivalent to the condition that the KL divergence from µ is bounded by the Fisher di∫∫log(dν/dµ)dν ≤ α2 ∥∇ log(dν/dµ)∥2 dµ, for any ν ∈ P which is absolutely continuouswith respect to µ. The LSI of proximal Gibbs measure can be established via standard perturbavergence:tion criteria. For L(µ) with bounded rst-variation, we may apply the classical Bakry-Emery andHolley-Stroock arguments (Bakry and Émery, 1985; Holley and Stroock, 1987) (Corollary 5.7.2
δL(µ)∗
and 5.1.7 of Bakry et al. (2014)): If ∥δµ ∥∞ ≤ B is satised for any µ ∈ P2 , then µ and pX
satisfy the LSI withα ≥ λ1 exp (−4B/λ) .

(5)

With the LSI condition on the proximal Gibbs distribution, it is known that the MFLD converges
to the optimal solution in an exponential order by using a so-called Entropy sandwich technique.

4

Proposition 1 (Entropy sandwich (Nitanda et al., 2022; Chizat, 2022)). Suppose that µ0 satises
L(µ0 ) < ∞ and the proximal Gibbs measure pµt corresponding to the solution µt has the LSI
constant α for all t ≥ 0, then the solution µt of MFLD satises
λKL(µ∗ , µt ) ≤ L(µt ) − L(µ∗ ) ≤ exp(−2αλt)(L(µ0 ) − L(µ∗ )),where µ∗= argminµ∈P L(µ).

Hence we know that time horizon T = O(1
λα log(1/ϵ̃)) is sucient to achieve ϵ̃ > 0 accuracy.

Convergence of the discretized algorithm.

While Proposition 1 only established the conver-gence rate of the continuous dynamics, similar guarantee can be shown for the discretized setting.
Let

LN (µ(N ) ) = N EX ∼µ(N ) [F (µX )] + λEnt(µ(N ) ),
(N )
i N
d¯
is a distribution of N particles X = (X )i=1 ⊂ R . Let µτ
be the distribution of thei Nparticles Xτ = (Xτ )i=1 at the τ -th iteration. Suzuki et al. (2023b) showed that, if λαη ≤ 1/4 and[()]η ≤ 1/4, then for B̄ 2 := E[∥X0i ∥2 ]+ λλ1 1 14 + λλ1 1 R̄2 +λd = O(d+λ−2 ) and δη := C1 L̄2 (η 2 +λη),
where µ(N )+ λλ1 B̄ 2 + d) = O(d + λ−1 ),)() ( N (N )
)
1
4 2 (
4Cλ
E[L (µ0 )]
λαητ
)
∗
∗
E[LN (µ(N)]−L(µ
)
≤
exp
−
−L(µ
)
+
L̄ C1 λη + η 2 +
,
τ
2
N
N
λα
λαN

where L̄ = 2R̄ + λλ1 = O(1) and C1 = 8(R̄2where Cλ is a constant depending on λ.

( In2 particular,) for a given ϵ̃ > 0, the right hand side
√
L̄
C1
4Cλ
L̄ C1
1
√
can be bounded by ϵ̃ +λαN after T = Oαϵ̃ +λα log(1/ϵ̃) iterations with the step sizeλαϵ̃(( L̄2 C1 L̄√C1 )−1 )√η=O. Furthermore, the convergence of the loss function can be connected to
αϵ̃ +
λαϵ̃
the convergence of the function value of the neural network as follows,[EXτ ∼µ(N )
k

L̄2
≤ 4λα

]
sup
z∈supp(PZ )( N(fµXτ (z) − fµ∗ (z))))
L (µ(N)τ− L(µ∗ )N

2[+ 2EX∗ ∼(µ∗ )⊗N(supz∈supp(PZ )

1
N

∑N

i=1 hX∗i (z) −Here the second term in the right hand side can be bounded by

32R̄2
N

∫

∗

hx (z)dµ (z)

)2

]
.

[
(
)]
2R̄2
d¯
1 + 2 (λλ
+
, if
2
λ1
1)

∥z∥ ≤ 1 for any z ∈ supp(PZ ) as in the k -sparse parity problem. Hence, by taking the number of−2[(λα)−2 + (λλ1 )−2 + d/λ1 ] and letting ϵ̃ = λαϵ2 with the choice of T and η as
described above, we have supz∈supp(P ) |fµX (z) − fµ∗ (z)| = Op (ϵ).
ZTparticles as N = ϵ

Assumptions on the model specication.

We restrict ourselves to the situation where aperfect classier with margin c0 is included in our model class, stated in the following assumption.

Assumption 2. There exists c0 > 0 and R > 0 such that the following conditions are satised:• For some R̄, there exists µ∗ ∈ P such that KL(ν, µ∗ ) ≤ R and L(µ∗ ) ≤ ℓ(0) − c0 .

• For any λ < c0 /R, the regularized expected risk minimizer µ[λ] := argmin L(µ) + λKL(ν, µ)satises Y fµ[λ] (X) ≥ c0 almost surely.
Importantly, we can apply the same general analysis for dierent classication problems, aslong as Assumption 2 is veried. The advantage of this generality is that we do not need to tailor
our convergence proof for individual learning problems. Note that the convergence rate of MFLD
is strongly aected by the values of R̄ and R; therefore, it is crucial to establish this condition
using the smallest possible values of R̄ and R, in order to obtain a tight bound of the classication
error. As an illustrative example, we now show that the k -sparse parity estimation satises theabove assumption.

5

Example: k-sparse parity estimation.

In the k -sparse parity setting (Example 1), Assump-tion 2 is satised with constants specied in the following propositions.

Proposition 2 (k-sparse parity). Under Assumption 1 and for R̄ = k, there exists µ∗ ∈ P suchthat

KL(ν, µ∗ ) ≤ c1 k log(k)2 d (= R),∗
and L(µ ) ≤ ℓ(0) − c2 , where c1 , c2 > 0 are absolute constants.

Proposition 3. Under Assumption 1 and the settings of R and R̄ given in Proposition 2, if
λ < c2 /(2R), then µ[λ] satisesmax{L̄(µ[λ] ), L(µ[λ] )} ≤ ℓ(0) − c2 + λR < ℓ(0) − c22 ,cand fµ[λ] is a perfect classier with margin c2 , i.e., Y fµ[λ] (X) ≥ 22 .
2

In other words, Assumption 2 is achieved with R = O(k log(k) d), R̄ = k and c0 = c2 /2. By

substituting these values of R and R̄ to our general results presented below, we can easily derivea bound for the classication error of the MFLD estimator.

4 Main Result
4.1

Generalization Error Analysis

We utilize the local Rademacher complexity (Mendelson, 2002; Bartlett et al., 2005; Koltchinskii,
2006; Giné and Koltchinskii, 2006) to obtain a faster generalization error rate. For the function
class of mean-eld neural networks, we introduce F := {fµ | µ ∈ P}, and the KL-constrained◦◦◦model class FM (µ ) := {fµ | µ ∈ P, KL(µ , µ) ≤ M } for µ ∈ P and M > 0. The Rademacher
complexity of a function class F̃ is dened as[]∑nRad(F̃) := Eεi ,zi supf ∈F̃ n1 i=1 ϵi f (zi ) ,nnwhere (zi )i=1 are i.i.d. observations from PZ and (εi )i=1 is an i.i.d. Rademacher sequence (P (εi =1) = P (εi = −1) = 1/2)). We have the following bound on the Rademacher complexity of the◦function class FM (µ ).

Lemma 1 (Local Rademacher complexity of FM (µ◦ )√, Chen et al. (2020) adapted). For any xedµ◦ ∈ P and M > 0, it holds that Rad(FM (µ◦ )) ≤ 2R̄
Combining

this

local

Rademacher

complexity

M
n .

bound

with

the

peeling device argument(van de Geer, 2000), we can roughly obtain the following estaimte (note that this is an informal
derivation):δ L̄(µ∗ )L̄(µ̂) − L̄(µ ) − (µ̂ − µ )+ λKL(µ∗ , µ̂) ≲
δµ
|
{z
} |
{z
}
∗

∗

(II)

√

KL(µ∗ , µ̂)
1≲+ λKL(µ∗ , µ̂),
n
nλ

(6)

(I)

with high probability, where µ̂ = argminµ∈P L(µ), µ∗= argminµ∈P L̄(µ) + λKL(ν, µ), R and R̄
are regarded as constants, and the last inequality is by the AM-GM relation. Observe that on
the left hand side, we have two non-negative terms (I) and (II). Corresponding to each term, we
obtain dierent types of classication error bounds (Type I and Type II in the following subsection,
respectively). Note that there appears a O(1/n) factor in the right hand side, which cannot be√

obtained by a vanilla Rademacher complexity evaluation because it only yields an O(1/n) bound.

In other words, localization is essential to obtain our fast convergence rate.
We remark that a local Rademacher complexity technique is also utilized by Telgarsky (2023)
to derive aO(1/n) rate.

They adapted a technique developed for a smooth loss function by

Srebro et al. (2010), which requires the training loss L(µ̂) to be suciently small, that is, of orderL(µ̂) = O(1/n). In our setting, to achieve such a small training loss, we need to take large R̄ suchas R̄ = Ω(log(n)). Unfortunately, such a large R̄ induces exponentially small log-Sobolev constant−cd
like α ≲ exp(−cd log(n)) = n. In contrast, our analysis focuses on the local Rademacher∗
complexity around µ , and hence we do not require the training loss to be close to 0; instead, it∗suces to have a training loss that is close to or smaller than that of µ .

6Type I: Perfect Classication with Exponentially Decaying Error

In the regime ofn = Ω(1/λ(K)2 ), we can prove that the MFLD estimator attains a perfect classication withan exponentially converging probability by evaluating the term (I). From Eq. (6) we can estab∗(K)2
lish KL(µ , µ̂) ≤ Op (1/(nλ)); this KL divergence bound can be used to control the L∞ -norm
2
2∗
between fµ̂ and fµ∗ . Indeed, we can show that ∥fµ̂ − fµ∗ ∥∞ ≤ 2R̄ KL(µ , µ̂) (see the proof of
Theorem 1). Then, under the margin assumption of fµ∗ (Assumption 2), we have that fµ̂ also
yields a Bayes optimal classier. More precisely, we have the following theorem.

Theorem 1. Suppose Assumptions 1 and 2 hold. Let M0 = (ϵ∗ + 2(R̄ + 1))/λ(K) . Moreover,suppose that λ(K)

Q := c20 −< c0 /R and

)
]
[
(
4R̄2
λ(K)
2
2
(K) ∗
(K)
+
8
R̄
(4
+
log
log
(8n
M
R̄))
+
nλ
ϵ
> 0,
λ
4
R̄
+
0
2
nλ(K)232R̄2 nthen fµ̂ yields perfect classication, i.e., P (Y fµ̂ (Z) > 0) = 1, with probability 1 − exp(−

nλ(K)2Q).
32R̄4Type II: Polynomial Order Classication Error

Next we evaluate the classication error∞
In this case, we do not require an L -norm bound as in the(K)
Type I analysis above; this results in a milder dependency on λand hence a better sample
bound from term (II) in Eq. (6).

complexity.

Theorem 2. Suppose Assumptions 1 and 2 hold. Let λ(K) < c0 /R and M0 = (ϵ∗ +2(R̄+1))/λ(K) .
Then, with probability 1 − exp(−t), the classication error of fµ(K) is bounded as[
P (Y fµ(K) (Z) ≤ 0) ≤ 2ψ(c0 )) 1
8R̄2 (
4 + t + log log2 (8n2 M0 R̄) +
(K)
n
nλ

We notice that the right hand side scales with O(1/(nλ

(K)

(
4R̄ +

λ(K)
32R̄2 n)]+ ϵ∗ .

)), which is better than O(1/(nλ(K)2 ))in Theorem 1; this implies that a sample size linear in the dimensionality is sucient to achieveK
small classication error. The reason for such improvement in the λ -dependence is that the∞stronger L -norm convergence is not used in the proof; instead, only the convergence of the loss
is utilized. On the other hand, this analysis does not guarantee a perfect classication.

4.1.1 Computational Complexity of MFLD
From the general result in Section ??, we can evaluate the computational complexity to achieve thestatistical bounds derived above. In both cases (Theorems 1 and 2), we may set the optimization∗(K)
error ϵ = O(1/(nλ)). Then, the total number of iteration can be

∑K
κ=1

(
)
Tκ ≤ O (d + λ(K)−1 )n exp(16cL R̄(R + 2)) log(nλ(K) ) .

The width N (the number of particles) can be taken as N()
O n2 exp(16cL R̄(R + 2)) .

= O((ϵ∗ λ(K) α)−2 ) = O(n2 α−2 ) =Corollary 1 (k-sparse parity setting). In the k-sparse parity setting, we may take R =O(k log(k)2 d), R̄ = k and λ(K) = O(1/R) = O(1/(k log(k)2 d)). Therefore, the classication
error is bounded by

P (Y fµ(K) < 0) ≤ O)( 2
k log(k)2 d(log(1/δ) + log log(n)) ,nwith probability 1 − δ . Moreover, if n = Ω(k6log(k)4 d2 ), then P (Y fµ(K) > 0) = 1 with probability

1 − exp(−Ω(nk 6 log(k)4 /d2 )).
2
2
2
As for the computational complexity, we require O(k log(k) dn log(nd) exp[O(k log(k) d)]) itera2
2
tions, and the number of particles is O(exp(O(k log(k) d)))).

7

Comparison with prior results.
nel (NTK) regime achieves only

In the 2-sparse parity setting, neural network in the ker-O(d2 /n) convergence in the classication error, as shown in

Ji and Telgarsky (2019) and Telgarsky (2023, Theorem 2.3), whereas we demonstrate that meaneld neural network can improve the rate to O(d/n) via feature learning. Telgarsky (2023) also
analyzed the learning of 2-sparse parity problem beyond the kernel regime, and showed that 2layer ReLU neural network can achieve the best known classication error O(d/n). However, theiranalysis considered a low-rotation dynamics and assumed convergence at t

→ ∞, whereas ourframework also provides a concrete estimate of the computational complexity. Indeed, the number of iterations can be bounded as O(dn log(nd) exp[O(d)]). In addition, while we still require2dexponential width N = O(n exp(O(d))), such a condition is an improvement over N = O(d ) in
Telgarsky (2023).
Barak et al. (2022) considered a learning method in which one-step gradient descent is performed for the purpose of feature learning, and then a network with randomly re-initialized bias
units is used to t the data. For the k -sparse parity problem, they derived a classication error√(k+1)/2
bound of O(d/ n). In contrast, our analysis yields a much better statistical complexity ofO(k 2 log(k)2 d/n ∧ exp(−Ω(nk 6 log(k)4 /d2 ))), which decouples the degree k in the exponent of
the dimension dependence.

5 Conclusion and Discussion
We provided a general framework to evaluate the classication error of a two-layer neural network
trained by the mean-eld Langevin dynamics. Thanks to the generality of our framework, an error
bound for specic settings can be derived by directly specifying the parameters in Assumption 2
such asR, R̄, and c0 .

We also proposed an annealing procedure to alleviate the exponential

dependencies in the LSI constant.

As a special (but important) example, we investigated the

k -sparse parity problem, for which we obtained more general and better sample complexity thanexisting works.

References
E. Abbe, E. B. Adsera, and T. Misiakiewicz.

The merged-staircase property:

a necessary and nearlysucient condition for sgd learning of sparse functions on two-layer neural networks. In

Learning Theory, pages 47824887. PMLR, 2022.

Conference on

E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. SGD learning on neural networks: leap complexity andsaddle-to-saddle dynamics.

arXiv preprint arXiv:2302.11055, 2023.

J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of featurelearning: How one gradient step improves the representation. In A. H. Oh, A. Agarwal, D. Belgrave, and
K. Cho, editors,

Advances in Neural Information Processing Systems, 2022. URL https://openreview.

net/forum?id=akddwRG6EGi.

J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, and D. Wu.
structure: a spiked random matrix perspective. In

Processing Systems (NeurIPS 2023), 2023.

Learning in the presence of low-dimensional

Thirty-seventh Conference on Neural Information

Séminaire de
Probabilités XIX 1983/84, pages 177206, Berlin, Heidelberg, 1985. Springer Berlin Heidelberg. ISBND. Bakry and M. Émery. Diusions hypercontractives. In J. Azéma and M. Yor, editors,
978-3-540-39397-9.
D. Bakry, I. Gentil, M. Ledoux, et al.

Analysis and geometry of Markov diusion operators, volume 103.

Springer, 2014.
B. Barak, B. L. Edelman, S. Goel, S. M. Kakade, eran malach, and C. Zhang. Hidden progress in deeplearning: SGD learns parities near the computational limit. In A. H. Oh, A. Agarwal, D. Belgrave, and
K. Cho, editors,

Advances in Neural Information Processing Systems, 2022. URL https://openreview.

net/forum?id=8XWP2ewX-im.

P. Bartlett, O. Bousquet, and S. Mendelson. Local Rademacher complexities.
33:14871537, 2005.

8

The Annals of Statistics,

P. L. Bartlett, M. I. Jordan, and J. D. McAulie. Convexity, classication, and risk bounds.

the American Statistical Association, 101(473):138156, 2006.

Journal of

G. Ben Arous, R. Gheissari, and A. Jagannath. High-dimensional limit theorems for sgd: Eective dynamics and critical scaling.

Advances in Neural Information Processing Systems, 35:2534925362, 2022.

Concentration Inequalities: A Nonasymptotic Theory of Independence. OUP Oxford, 2013. ISBN 9780199535255. URL https://books.google.co.jp/books?

S. Boucheron, G. Lugosi, and P. Massart.

id=koNqWRluhP0C.

F. Chen, Z. Ren, and S. Wang. Uniform-in-time propagation of chaos for mean eld langevin dynamics.

arXiv preprint arXiv:2212.03050, 2022.

S. Chen and R. Meka. Learning polynomials in few relevant dimensions. In J. Abernethy and S. Agarwal,
editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of
Machine Learning Research, pages 11611227. PMLR, 0912 Jul 2020. URL https://proceedings.

mlr.press/v125/chen20a.html.

Z. Chen,Y. Cao,Q. Gu,

layer neural networks.
itors,and T. Zhang.

A generalized neural tangent kernel analysis for two-In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, ed-

Advances in Neural Information Processing Systems, volume 33, pages 1336313373. Cur-URL https://proceedings.neurips.cc/paper_files/paper/2020/file/
9afe487de556e59e6db6c862adfe25a4-Paper.pdf.
ran Associates, Inc., 2020.

Mean-eld langevin dynamics : Exponential convergence and annealing. Transactions on
Machine Learning Research, 2022. URL https://openreview.net/forum?id=BDqzLH1gEm.

L. Chizat.

L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using
optimal transport. In

Advances in Neural Information Processing Systems 31, pages 30403050, 2018.

L. Chizat and F. Bach. Implicit bias of gradient descent for wide two-layer neural networks trained withthe logistic loss.

arXiv preprint arXiv:2002.04486, 2020.

A. Damian, J. Lee, and M. Soltanolkotabi.

Neural networks can learn representations with gradient

Proceedings of Thirty Fifth Conference on Learning
Theory, volume 178 of Proceedings of Machine Learning Research, pages 54135452. PMLR, 0205 Jul
descent.

In P.-L. Loh and M. Raginsky, editors,

2022. URL https://proceedings.mlr.press/v178/damian22a.html.
A. Daniely and E. Malach. Learning parities with neural networks.

arXiv preprint arXiv:2002.07400, 2020.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in highdimension.

arXiv preprint arXiv:1904.12191, 2019.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel
methods?

arXiv preprint arXiv:2006.13409, 2020.

E. Giné and V. Koltchinskii. Concentration inequalities and asymptotic results for ratio type empirical
processes.

The Annals of Probability, 34(3):11431216, 2006.

R. Holley and D. Stroock.

Logarithmic sobolev inequalities and stochastic ising models.

statistical physics, 46(5-6):11591194, 1987.

Journal of

D. Hsu. Dimension lower bounds for linear approaches to function approximation.
K. Hu, Z. Ren, D. Siska, and L. Szpruch. Mean-eld langevin dynamics and energy landscape of neuralnetworks.

arXiv preprint arXiv:1905.07769, 2019.

X. Huang, P. Ren, and F.-Y. Wang. Distribution dependent stochastic dierential equations.

Mathematics in China, 16(2):257301, 2021.

Frontiers of

A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neuralnetworks. In

Advances in Neural Information Processing Systems 31, pages 85808589, 2018.

Z. Ji and M. Telgarsky. Polylogarithmic width suces for gradient descent to achieve arbitrarily small test
error with shallow relu networks.

arXiv preprint arXiv:1909.12292, 2019.

9

V. Koltchinskii. Local Rademacher complexities and oracle inequalities in risk minimization.

of Statistics, 34:25932656, 2006.

M. Ledoux and M. Talagrand.

The Annals

Probability in Banach Spaces. Isoperimetry and Processes. Springer, New

York, 1991. MR1102015.
S. Mei, A. Montanari, and P.-M. Nguyen. A mean eld view of the landscape of two-layer neural networks.

Proceedings of the National Academy of Sciences, 115(33):E7665E7671, 2018.

S. Mendelson. Improving the sample complexity using global data.

Theory, 48:19771991, 2002.

A. Mousavi-Hosseini, D. Wu, T. Suzuki, and M. A. Erdogdu.
structured data.

2023), 2023.

In

IEEE Transactions on Information

Gradient-based feature learning underThirty-seventh Conference on Neural Information Processing Systems (NeurIPSA. Nitanda and T. Suzuki.

arXiv:1712.05438, 2017.

Stochastic particle gradient descent for innite ensembles.

arXiv preprint

A. Nitanda, D. Wu, and T. Suzuki. Convex analysis of the mean eld langevin dynamics. In G. Camps-

Proceedings of The 25th International Conference on ArticialIntelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 97419757.
Valls, F. J. R. Ruiz, and I. Valera, editors,

PMLR, 2830 Mar 2022.
M. Renetti, S. Goldt, F. Krzakala, and L. Zdeborova. Classifying high-dimensional gaussian mixtures:

Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine
Learning Research, pages 89368947. PMLR, 1824 Jul 2021. URL https://proceedings.mlr.press/
Where kernel methods fail and neural networks succeed. In M. Meila and T. Zhang, editors,

v139/refinetti21b.html.

G. M. Rotsko and E. Vanden-Eijnden.
particle system approach.

Trainability and accuracy of neural networks: An interacting

arXiv preprint arXiv:1805.00915, 2018.

J. Sirignano and K. Spiliopoulos. Mean eld analysis of neural networks: A central limit theorem.

Processes and their Applications, 130(3):18201852, 2020.

Stochastic

N. Srebro, K. Sridharan, and A. Tewari. Smoothness, low noise and fast rates. In J. Laerty, C. Williams,

Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010. URL https://proceedings.neurips.cc/paper_files/
J. Shawe-Taylor, R. Zemel, and A. Culotta, editors,

paper/2010/file/76cf99d3614e23eabab16fb27e944bf9-Paper.pdf.
T. Suzuki, A. Nitanda, and D. Wu.
langevin dynamics. In

Uniform-in-time propagation of chaos for the mean eld gradient

Submitted to The Eleventh International Conference on Learning Representations,2023a. URL https://openreview.net/forum?id=_JScUk9TBUn.

T. Suzuki, D. Wu, and A. Nitanda. Convergence of mean-eld langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction.

arXiv preprint, 2023b.

The EleventhInternational Conference on Learning Representations, 2023. URL https://openreview.net/forum?

M. Telgarsky. Feature selection and low test error in shallow low-rotation ReLU networks. In

id=swEskiem99.

S. van de Geer.

Empirical Processes in M-Estimation. Cambridge University Press, 2000.

C. Wei, J. D. Lee, Q. Liu, and T. Ma. Regularization matters: Generalization and optimization of neuralnets vs their induced kernel. In

Advances in Neural Information Processing Systems, pages 97129724,

2019.

10