Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond
Authors: Taiji Suzuki
Denny Wu
Kazusato Oko
Atsushi Nitanda
Venue: NeurIPS 2023 poster (2023)
Keywords: mean-field regime
feature learning
Neural network optimization
sparse parity function
classification
sample complexity

Abstract: Neural network in the mean-field regime is known to be capable of \textit{feature learning}, unlike the kernel (NTK) counterpart. Recent works have shown that mean-field neural networks can be globally optimized by a noisy gradient descent update termed the \textit{mean-field Langevin dynamics} (MFLD). However, all existing guarantees for MFLD only considered the \textit{optimization} efficiency, and it is unclear if this algorithm leads to improved \textit{generalization} performance and sample complexity due to the presence of feature learning. To fill this gap, in this work we study the statistical and computational complexity of MFLD in learning a class of binary classification problems. Unlike existing margin bounds for neural networks, we avoid the typical norm control by utilizing the perspective that MFLD optimizes the \textit{distribution} of parameters rather than the parameter itself; this leads to an improved analysis of the sample complexity and convergence rate. 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 achieves a sample complexity where the degree $k$ is ``decoupled'' from the exponent in the dimension dependence.

Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond
Taiji Suzuki! Denny Wu
Kazusato Oko 1., Atsushi Nitanda 2.5
University of Tokyo, RIKEN AIP, New York University, *Flatiron Institute, 'Kyushu Institute of Technology
taiji@mist. wtt.:/miiit..tttoooooooo jc. dennywu@nyu.c edu, oko-kazusatoßg. ecc. u-tokyo. jp. nitandačai. .wwpaik.tttttcccc....o.cc jc.
Abstract
Neural network in the mean mean-field regime is known to be capable of feature learning, unlike the kernel (NTK) counterpart. Recent works have shown that mean-field neural networks can be globally optimized by noisy gradient descent update termed the mean-field Langevin dynamics (MFL.D. However, all existing guaran- tees for MFLD only considered the optimization efficiency, and it is unclear if this algorithm leads to improved generalization performance and sample complexity due to the presence of feature learning. To fill this important gap, in this work we study the sample complexity of MFLD learning a class of binary classification problems. Unlike existing margin bounds for neural networks, avoid the typical norm control by utilizing the perspective that MFLD optimizes the distribution parameters rather than the parameter itself; this leads an improved analysis of the sample complexity and convergence rate. We apply our general framework to the learning of k-sparse parity functions, where prove that unlike kernel methods, two-layer neural networks optimized by MFLD achieves a sample complexity where the degree is "decoupled" from the exponent in the dimension dependence.
Introduction
Mcan-field Langevin dynamics. The optimization dynamics of two-layer neural networks in the mean-field regime can be described by nonlinear partial differential equation of the distribution of parameters (Nitanda and Suzuki, 2017; Chizat and Bach, 2018; Mei et al., 2018; Rotskoff 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 (it) the parameters are allowed 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-field neural networks, the mean-field Langevin dynamics (MFLD) (Mei et al., 2018; Hu et al., 2019) is particularly attractive due to the recently established quantitative optimization guarantees. MFLD arises from noisy gradient descent update on the parameters, where Gaussian noise is injected to the gradient to encourage "exploration" It has been shown that MFLD globally optimizes entropy-regularized convex functional in the space of measures, and for the infinite-width and continuous-tin time dynamics, the convergence rate is exponential under suitable isoperimetric conditions (Nitanda et al., 2022; Chizat, 2022). Furthermore. uniform-in-timg estimates of the particle discretization error have also been established (Suzuki et al., 2023a; Chen al., 2022), meaning that optimization guarantees for the mnfiiie-dimension problem can be effectively translated to finite-widie neural network.
However, existing analyses of MFLD only considered the optimization of neural networks; this alone does not demonstrate the benefit of mean-field regime nor the presence of feature learning. Therefore,
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

an important problem is to characterize the generalization of the learned models, and prove efficient 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) efficiently learns an interesting class of functions, with 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 & coordinates of the input (on hypercube); as a special case, setting k- = 2 recovers the classical XOR problem. When kt < d, this target function
now-dimensional, and hence we expect feature learning to beneficial in that can "zoom in" to relevant subspace. In contrast, for kernel methods (including neural networks in the lazy regime) which cannot adapt to such structure, has been shown that sample complexity of [(dd 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-field regime can achieve sample complexity of n O(d/e) (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 + oo) with no iteration complexity. Moreover, Wei et al. (2019); Chizat and Bach (2020) directly analyzed the infinite- wiith limit, and while Telgarsky (2023) provided low-rotation regime, and a very large number of particles N Odd)) is required. Lastly, these analyses are specialized to XOR, and do not directly generalize to the k-parity setting.
finitewwidt characterization, the dynamics is restricted to the
1.1 Our Contributions
In this work, we bridge the aforementioned gap by presenting simple and general framework to establish sample complexity of MFLD in learning binary classification 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 specifically, our contributions can be summarized as follows.
We present general framework to analyze MFLD in the learning of binary classification tasks. Our framework has two main ingredients: (i) an annealing procedure that applies to common classification losses that removes the exponential dependence on regularization parameters in the logarithmic Sobolev inequality, and (tt) a novel local Rademacher complexity analysis for the distribution of parameters optimized by MFLD. As result, we can obtain generalization guarantee for the learned neural network in discrete- time and width settings.
We apply our general framework to the k-sparse parity problem, and derived learning guarantees with improved rate of convergence and dimension dependence, as shown in Table 1. Specially, in the n x d² regime we obtain exponentially converging classification error, whereas in the d regime we achieve linear dimension dependence, Note that this improves upon the NTK analysis (which gives a sample complexity of (dd))) in that "decouple" the degree i from the exponent in the dimension dependence. Our theoretical results are supported by empirical findings.
Authors
regime/method
k-parity
class crror
d'/n width d"
iiteratinns
Ji and Telgarsky (2019)
NTK/SGD
No
d²/c
Telgarsky (2023)
NTK/SGD
No
d2/n
d²
d2/e
Barak et al. (2022)
Two phase SGD
Yes
d 1)/2 //n
0(1)
d/c
Wei et al. (2019)
mean-filddGF
No
d/n
00
0o
Telgarsky (2023)
mean-field/GF
No exp( exp(-O(y/n)) (v(rd))
d/n
dd
0c
Ours
mean-field/MFLD
Yes
O(d)
Ours
mcan-fielW/MFLD
Yes
d/n
Old)
e O(d)
Statistical and computational complexity (omitting poly-log terms) for the k-sparse (or 2-sparse) parity problem. Column "k-parity" indicates applicahility to the general k-spaase parity setting, and for references that does not handle k-parity we state the complexity for 2-parity (XOR). is the input dimensionality. and is the sample size. For SGD, n is the total sample size given by product mini-batch size and iterations.
Table 1:
1.2 Additional Related Works
In addition to the mean-field analysis, parity-like functions can be learned via other feature learning procedures such as the one gradient step analysis in Danicly and Malach (2020); Ba et al. (2022);

---

Damian et al. (2022); Barak et al. (2022). Such analysis requires nontrivial gradient concentration at initialization, which translates to a sample complexity that scales as n ((d") (Barak et al., 2022). For "narrow" neural networks, Abbe et al. (2023) showed that modified projected (online) gradient descent algorithm can learn the k-parity problem with a sample complexity of O e (for some unknown c). Also, Refinetti et al. (2021); Ben Arous et al. (2022) showed that two-layer ReLU network with more than 4 neurons can learn the Gaussian XOR problem with gradient descent.
Problem Setting
Throughout this paper, consider classification problem given by the following model:
Y =1a(2) (2) {+1}
where Z = (Z...... Zd) is the input random variable on R and 1 is the indicator function corresponding to a measurable set A E (B(R)), i.e., (Z) = 1 if Z E and 1 (2) 0 Z & A. Let Pz be the distribution of Z. We are given input-output pairs D, (zi, y/ >==1 independently identically distributed from this model as training data. Then, we construct a binary classifier that 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-field regime via the mean-field Langevin dynamics.
One important problem setting for analysis is the k-sparse parity problem defined follows. Example (k-sparse parity problem). Pz is the uniform distribution on the grid /d)/ and A {S ......., Ca) E {#1/Vd]" [1 -. S
03'
As a special case, k = 2 (XOR) has been extensively studied (Wei et al., 2019; Telgarsky, 2023).
Mean-field two-layer network. Given input 2, let (z(2) be one neuron in two-layer neural network with parameter J = (x1, 22, (2) v2a) E IRd+++1 defined as
(=) = R|tanh(=" T x1 + x2) + 2 taahhrz)]/33
where R E R is hyper-parameter determining the scale of the network2. We place an extra tanh activation for the bias term 23 E IR because the boundedness h2 is required in the convergence analysis. Let P be the set of probability measures on (Rdd ((Rd)) where d d + and ((R)) the Borel o-algebra on Rd and Pp be the subset P such that its p-th moment is bounded: [ < oc (μ P). The mean-field neural network is defined as integral over neurons h.,
fa(-) J ha (-))(d)))
for μ E P. To evaluate the performance of fu. we define the empirical risk and the population risk as L(μ) Ei (z.), L(p) E (2))),
respectively, where L I -> RR0 is a convex loss function. In particular, we consider the logistic loss {(f. v) log(1 + exp(-yf) for [+1} and f E I. To avoid overfitting, we consider regularized empirical risk F(4)= == L(x) + XEx~x [Ai]]X [J2], where X, A₁ 2 are regularization parameters, One advantage of this mean-field definition is that fy is linear with respect to and hence the functional L(H) becomes a convex functional.
Mean-field Langevin dynamics. We optimize the training objective via MFLD, which is given by the following stochastic differential equation:
dX2 = 8F(m)) (X,)dt+ 2XdW,, = Law(X)),
where Xo Ma. Law( (X) denotes the distribution of the random variable X and (Wi)zzo is the d-dimensional standard Brownian motion. Readers may refer to Theorem 3.3 of Huang et al. (2021) for the existence and uniqueness of the solution. Here, (F(H)) is the first variation of F.
Definition 1. For a functional G P +R R. the firstvvriationn &G (1) at E P is a continuous functional P x IRd + IR satisfying lim (lev+()-eeu))
8y &C (n)(x)d(v for any v E P.
6
We present the axis ali-nad setting for conciseness, but the same result holds under orthogonal transforms. We may also deal with more general activation function that is smooth and bounded,
3

---

In our setting, we have (F(H) (x) = Ela + A(A, ||.133). It is known that the
Fokker-Planche equation of the SDE (1) is given by²
/ s - V.
[ (Alog(me) -
(F())) 8v
Then, we can verify that this is equivalent to the Wasserstein gradient flow to optimize the following entropy regularized risk (Mei et al., 2018; Hu et al., 2019):
L(p) F(R) XEnt(() T(p) AKI((U, (const.)
where KL(U, (4) = f log(n//)dd is the KL divergence between and μs, and is the Gaussian distribution with mean 0 and variance 1/(2)),, i..,, // = ((0,1/(2))))
For practical algorithm, that is, we approximate the solution Pr by an empirical measure J. 전 Ei ōx, corresponding to set of finite particles 2 = ( C Rdd Let 2, (x:) C IRd be N particles at the r-th update (T {0, 1, ... }), and define fA, u.z. finite particle approximation of the population counterpart. Then, the discretized MFLD is defined as follows: X ~ Po. and 2, is updated as
need to consider space- and time-discretized version of the MFLD,
(F())) (x+) + 2Am
X = x n/
where 77 > 0 the step size, and E i.s.d. (0, I). This is the Euler- Marle--m approximation of the MFLD with a discretized measure; 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 classification error, we assume that the loss function satisfies the following condition.
Assumption 1. The convex loss function i R - RRz satisfies the following conditions:
is first order differentiable, its derivative is Lipschitz continuous and its derivative is bounded by I: ( (x'))) S C| x' and sup, (x() S.
f monotonically decreasing, and is classification calibrated: (() < (Bartlett et al., 2006).
=(u))² ((0) (e(u) vl" for any 0.
This standard assumption satisfied by several loss functions such as the logistic loss. We remark that the first assumption is used to show the well-definedness of the mean-field Langevin dynamics and derive its discretization error, and also to obtain a uniform generalization error bound through the classical contraction argument (Boucheron et al., 2013; Ledoux and Talagrand, 1991). The second and third assumptions are used to show the convergence of classification error of our estimator.
Logarithmic Sobolev inequality. Nitanda al. (2022); Chizat (2022) showed that the convergence of MFLD crucially relies on properties of the proximal Gibbs distribution whose density is given by
(x( x exp (-X 8p (x) (x)).
1 (-18F(W)
where j P. By the smoothness of the loss function (Assumption 1) and the tanh activation, show that the objective L has unique solution μ which is also proximal Gibbs measure itself.
Proposition (Proposition 2.5 of Hu et al. (2019)). The functional C has a unique minimizer in P2 that is absolutely continuous with respect to the Lebesgue measure. Moreover, μ² P2 is the optimal solution if and only if μ is absolutely continuous and its density function given by Par.
The next question is how fast the solution be converges the optimal solution μ², As will see, the convergence of MFLD heavily depends on logarithmic Sobolev inequality (LSI) on Pp-
"This should be interpreted in weak sense, that is, for any continuously differentiable function with compact support, f odj4 ødpe - VΦ (V log(14) V (F(()))
8v dd

---

Definition 2 (Logarithmic Sobolev inequality). Let be probability measure on (RR ((R))) satisfies the LSI with constant o > 0 for any smooth function Φ R2d lK with E [c0)] < 0o.
2 E V 6
E [ log ( - 1 log (E
This is equivalent to the condition that the KL divergence from μ is bounded by the Fisher divergence: I [og(de/dy)) /dp////d/d///////d//// 2 [|V (oggdd/dd))]] for any P which is absolutely continuous with respect to f.. The LSI of proximal Gibbs measure can be established via standard perturbation criteria. For T(x)) with bounded first-varition,, we may apply the classical Bakry-Fmery and Holley-Stroock arguments (Bakry and Emery, 1985; Holley and Stroock, 1987) (Corollary 5.7.2 and 5.1.7 of Bakry al. (2014)): If]
(L(a) op
oo
B is satisfied for any μ E P2. then μ and p satisfy the LSI with
a 2 A1 exp (-4B/A).
(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.
Proposition 2 (Entropy sandwich (Nitanda et al., 2022: Chizat, 2022)). Suppose that yi satisfies C(PO) < oo and the proximal Gibbs measure Pp. corresponding to the solution fr has the LSI constant a for all 2 O, then the solution μr, of MFLD satisfies
XKL(E" ,p) S C(ux) - f(x²)= S (xp((-ax))(CLoo)) - L(u")),
where μ²¹ argmin uEP L(u) (the existence and uniqueness of μ² is guaranteed by Proposition 1),
Hence we know that time horizon T = O Xa log(1/E) is sufficient to achieve ? > 0 accuracy.
Convergence of the discretized algorithm. While Proposition 2 only established the convergence rate the continuous dynamics, similar guarantee can be shown for the discretized setting. Let
L (x())) = NE
NN(N) ~ [F(uz))] [
where (M( is distribution of N particles 2 = (X*) i=1 RR.. Let j(() be the distribution of the particles D-- (X:) at the 7-th iteration. Suzuki et al. (2023b) showed that, if Xan 1/4
and 17 S 1/4, then for B² = E||l x xA, ( XXi R 1²+Ad +
O(d - x-) and d.p =
C, + An), where L = 2R + AA1 = O(1) and C₁ = 8(722 + aA B +d) -(d+1-1))
1 N BBC ( Sexp
ELC 11 c(r)) + Xa
40
Xanr 2
I.c (21) + +7) + Xa.N'
where C2 a constant depending on X. In particular, for given ē > O, the right hand side can be bounded by ē
4C2 after T O iPC TvCC 1 log(1/2) iterations with the step size Aa
A
V'aē
0(( LC + LVC ) Furthermore, the convergence of the loss function can be connected to Ani the convergence of the function value of the neural network as follows.
E M ==)upp(P2)
sup
(Sux, fnr (2))²
2
4F2 <
Xa
(cN (M
-(1²)) + 2TE
2x x JoN
sup xEsupp(P2) Pz)
1 N
E
lxx (2) - (hr
via Lemma 2, if ||=|| < 1 for any z E supp( as in the k-sparse parity problem. Hence, by taking the number of particles as N = e =>[[Aa)= + (AAI)-² + d/Ai] and letting ē Xaee with the choice of T and 77 as described above, we have suppesupp(Pz) 1fuar (z) - fx- (x)] = Op(e).
Here the second term in the right hand side can be bounded by
32 N
1 +
[xx,)=
2
+
d X,
5

---

Assumptions on the model specification. We restrict ourselves to the situation where perfect classifier with margin co is included in our model class, stated in the following assumption.
Assumption 2. There exists cu 0 and R > O such that the following conditions satisfied:
For some R. there exists μ E such that KL(V, ) S R and L(p") S E(0) co-
For any X < co/ R, the regularized expected risk minimizer MM == argmin L(m) + AKI satisfies Y funa (X) 2 co almost surely.
Importantly, we can apply the same general analysis for different classification problems, as long as Assumption verified. The advantage of this generality that we do not need to tailor our convergence proof for individual learning problems. Note that the convergence rate of MFLD is strongly affected by the values of R and It; therefore, it crucial to establish this condition using the smallest possible values of ft and R. in order to obtain a tight bound of the classification error. As an illustrative example, we now show that the k-spare parity estimation satisfies the above assumption.
Example: k-sparse parity estimation. In the k-sparse parity setting (Example I), Assumption 2 is satisfied with constants specified in the following propositions. The proofs are given in Appendix A. Proposition 3 (k-sparse parity). Under Assumption I and for R = k, there exists p P such that KL(V, ) S erk log(k)) (=
c2. Proposition 4. Under Assumption and the settings of R and given in Proposition 3, if X < c2/(2R), then Mix sutisfies
and L(p") S ((0) C2. where C1. > 0 are absolute constants.
and fuu is a perfect classifier with margin C22 ie., f (X) 2 62 2 (ax()) L(pD)) S E((0) - AR < {()) 2
In other words, Assumption 2 is achieved with R = O log(k) (d),, R = k and co = c2/2. By substituting these values of R and R to our general results presented below, we can easily derive bound for the classification сπтоr of the MFLD estimator.
4
Main Result: Annealing Procedure and Classification Error Bound
4.1
Annealing Procedure
The convergence rate of MFLD is heavily dependent the LSI constant which may impose large computational cost. We alleviate this dependency by employing novel annealing scheme where we gradually decrease the regularization parameter X. In particular, at the x-th round, we run the MFL.D
until (near) convergence with regularization parameter (()) = -- x(); po p(n-1),
dx,- = - [
(L(p)) by
(X,)dt 22( 21 X,] + 22x
(6)
which corresponds to minimizing (())() == L(n) + X (L(
μ). Then, we obtain a near optimal
solution ((()) as ((*) (f(*)) S minper (1) t for a given e > 0. We terminate the procedure
after K rounds and obtain ju as the output.
Suppose that there exists n² such that KL(P, R and L(p*) < 8 Then, as long ( 2 8" and e* < 8*, we have that
(((())) + ( (()) S L(₁²) + XL* (() + e" s 28 + X R (R+ (M)
Since ((H) (px) is monotonically decreasing during the optimization, we always have L(H+) < ((x) (1(())))) < (((x))(((n)) < (R + 1)) Now we utilize the following structure on common classification losses to ensure that [8T.(1)/5pll is small when the loss I(H) is small.
Assumption 3. There exists CL > 0, such that, for any E P, it holds that || (L(x) Ilx S cx RL()).
For example, the logistic loss satisfies this assumption:
|0, log(1 + exp( (x)( = 1-exp(-u S <log(1 + exp(-u))
cxp(-a)

---

Hence, the condition holds for ct = 1 because ||
1l P (1)) v
1 -i Re(y, f=(z)))= RL(x). Under this assumption, the Holley- -Stryy-ttrook argument yields that the log-Sobolev constant during the optimization can be bounded
A1 exp
Aci, ((R))))))
AI exp (-8
Xim)
which is independent of (()) and the final accuracy 8*, Hence, after K = log, (R((((())*))
round, we achieve that L(u(K)) S 8" + 2e" where each round takes Tx = log(1/6 (or((1/")//44 (()())))/((2)) = Ollog(L/e") (xp(8erR(RR + 2))/ for the continuous time setting.
For the diserete setting, Tx 0 ((
ac" VXimaae log(L/e") exp(Sc, 2)) ()1/(1)
iterations each round sufficient (where C1 is given for A(0)), where the step size is 7 o (CT'ae" ACT1/V/X()ae" VX(m)a+)). As long as E" =((Am) for all k. the total itera- tion number can be simplified as O (Iog(1/c*) exp(16e1 N (number of particles) can be taken as N = O(((((a))) -2 (1(1))=
and the width (x00)=2= + d)/(6)))
O ([exp(16c_R(FF (R- (A()() + dl/(c))²).
Remark 1. We make the following remarks on the annealing procedure.
The main advantage of this annealing approach is that the exponential factor induced by the LSI constant a is not dependent on the choice of the regularization parameter (1) (note that the LSI is solely determined by the intermediate solution pt). In contrast, the naive Holley-Stroock argument of Eq. (5) imposes exponential dependency on the regularization parameter.
Our annealed algorithm differs from the procedure considered in Chizat (2022, Section 4); importantly, we make use of the structure of classification loss functions to obtain refined computational complexity analysis.
4.2 Generalization Error Analysis
We utilize the local Rademacher complexity (Mendelson, 2002; Bartlett et al., 2005; Koltchinskii, 2006; Giné and Koltchinskii, 2006) to obtain faster generalization error rate. For the function class of mean-field neural networks, we introduce F ;= {fa E P), and the KL-constrained model class Fx (x²" == {fa μ E P, KL(K", , S M) for nº and M >> 0. The Rademacher complexity of function class JF defined as
Rad(F) = Ee,.. sup Ei ff(z))
where (2) are i.idd observations from Pz and (s.) is an i.i.d Rademacher sequence (P(e, = 1) P(e -1) = 1/2)). We have the following bound on the Rademacher complexity of the function class (u²).
Lemma (Local Rademacher complexity of FM (M²), Chen et al. (2020) adapted). For any fixed μº E and > 0, holds that Rad(F (x²)) s 2R
The proof is given in Appendix B.2 in the supplementary material. Combining this local Rademacher complexity bound with the peeling device argument (van de Geer, 2000), we can roughly obtain the following estaimte (note that this is informal derivation):
L(p) L(x²) (u 1²) 8L(H²) ōp
KL(H4
1 S
AKL(KC,
+ AKL(U", nd
(7)
with high probability, where à argmin L(H), μ² argmin L(x) XKL μ), 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 different types of classification error bounds (Type I and Type II in the following subsection, respectively). Note that there appears 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/ bound. In other words, localization essential to obtain our fast convergence rate.
We remark that a local Rademacher complexity technique is also utilized by Telgarsky (2023) to derive (11/n) rate. They adapted a technique developed for smooth loss function by Srebro et al. (2010), which requires the training loss (()) to be sufficiently small, that is, of order ( O(1/n). In our setting, to achieve such a small training loss, we need take large R such as R = ([)og(n)). Unfortunately, such a large R induces exponentially small log-Sobolev constant like a S exp( (ed = n In contrast, our analysis focuses on the local Rademacher complexity around J and hence do not require the training loss to be close O; instead, it suffices to have training loss that is close to or smaller than that of J=".
4.2.1 Type I: Perfect Classification with Exponentially Decaying Error
prove that the MFLD estimator attains a perfect classi- fication with an exponentially converging probability by evaluating the term (I). From Eq. (7) we can establish KL( . < (p(1 )); this KL divergence bound can be used to control the L - between fa and fu- Indeed, we can show that llfe - fp- 1I& I 2R" .i) (see the proof of Theorem 1). Then, under the margin assumption of f.- (Assumption 2), we have that fp also yields Bayes optimal classifier. More precisely, we have the following theorem.
In the regime of ())))
Theorem 1. Suppose Assumptions I and hold. Let Mo (e" + + + 1))/X(K) Moreover,
suppose that ((K) < co/ R and
4R²
A(K)
Q c
(KK) 4R+
+ +8R (4 + log log2 (8n² (8n2 M)) + nx( c² > 0, 32R2
nA(K)2
then Sp yields perfect classification, i.e., P PY [Z) > 0) = 1, with probability 1 exp(
2A(K10 32R
The proof is given in Appendix C. in the supplementary material. From Proposition we see that the requirement x 1/2 implies that in the n x d² regime, Theorem gives a perfect classification guarantee with a failure probability decaying exponentially fast.
4.2.2 Type II: Polynomial Order Classification Error
Next we evaluate the classification error bound from term (II) in Eq. (7). In this case, we do not require an - bound as in the Type I analysis above; this results in a milder dependency on X(K) and hence a better sample complexity.
Theorem 2. Suppose Assumptions I and hold. Let ( < co/ R and Mo = (E" 2 1))/A(K) exp(-tt) the classification error of fuixy is bounded
Then, with probability
8R.2
1 + π
(K)
P(Y fg uuy (2) s 0) < (co)
nX(K) (4 (4+ log log, MoR))
[RR
32R2n
The proof is given in Appendix C.2 in the supplementary material. We notice that the right hand side scales with O(1 (11/(nX(4 (1l((n(())) which is better than ((1/(nX(X)2) in Theorem 1; this implies that a sample size linear the dimensionality is sufficient to achieve small classification error. The reason for such improvement in the \ deeenndenece is that the stronger - -nrm 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 classification.
4.2.3 Computational Complexity of MFLD
From the general result in Section 4.1, can evaluate the computational complexity to achieve the statistical bounds derived above. In both cases (Theorems I and 2), we may set the optimization error e* = O(11(nx(K)) Then, the total number of iteration can be
Σ K Tx S (dd+\(k)-)nn (1)) exp(16c_ + 2)) log (og(n)(k)))
The width N (the number of particles) can be taken N Oi(e* ((()2((())) =O(n²a 2) O (n² exp(16c), R(R + 2))).
8

---

Corollary 1 (k-sparse parity setting). In the k-sparse parity setting, we may take R = O log(k)² Pd), R = k and A(K) = O() = ((1/(k (og//(ooog())) (og( d)). Therefore, the classification error is bounded by
P 0) 0 k2 log log(1/0) log loggog())))
n
with probability 1 1- 6. Moreover, nn (2(kk log(k)(dP)) then ((Y 0) = with probability
1 1 exx log(k) (gg(k))+//)) (dP)).
As for the computational complexity, we require (A: (og(k)) () log explO log( itera- tions, and the number of particles is O(exp(O(k² (
Comparison with prior results. In the 2-sparse parity setting, neural network in the kernel (NTK) regime achieves only (dd² ((d convergence the classification error, as shown in and Telgarsky (2019) and Telgarsky (2023, Theorem 2.3), whereas we demonstrate that mean-field neural network can improve the rate to (d/nn) via feature learning. Telgarsky (2023) also analyzed the learning of 2-sparse parity problem beyond the kernel regime, and showed that 2-layer ReLU neural network can achieve the best known classification error O(d/n). However, their analysis considered a low-rotation dynamics and assumed convergence at ( - whereas our framework 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 still require exponential width N O(n (xp(O)d))), such a condition is an improvement over N O(d)) in Telgarsky (2023).
Barak et al. (2022) considered 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 fit the data. For the k-sparse parity problem, they derived classification error bound of O(d(++1)/2 /n)) In contrast, our analysis yields much better statistical complexity of O dd/( A exp(- log(k)) /dd))) which "decouples" the degree k in the exponent of the dimension dependence.
5 Numerical Experiment
We validate our theoretical results by numerical experiment on synthetic data. Specifically, we consider the classification of 2-sparse parity with varying dimensionality d and sample size n.
Recall that the samples {(z,, ) are independently generated so that 2; follows the uniform distribution on {+1/v and = di. E {+1} (zi (G..1...., Sr.a)) A finite-width approximation of the mean-field neural network Eド= hr, (2) is employed with the width N 2, 000. For each neuron of the network, all parameters are initialized to independently follow the standard normal distribution (meaning that the network is rotation invariant at initialization) and the scaling parameter R is set 15. We set d take values 5, 10, ..... 150, and n 50, 100,- 2000.. We trained the network using noisy gradient descent with 77 = 0.2, A = 0.1, and X 0.1 /d (fixed during the whole training) until T 10, 000. The logistic loss used for the training objective.
Figure shows the average test accuracy over five trials. We make the following observations.
1.0
n n=d
25
n = d
0.9
The red line corresponds to the sample size
n = (dd²), above which we observe that almost- perfettt-eeetttee classification is achieved. Ac- cording Section 4.2 (Type D) the classifi- cation error gets exponentially small with re- spect to n/d", which predicts very small clas- sification error above the line of n = (d?)) which matches experimental result.
dimension
-0.8
75-
0.7
100
125
150
250
500
750 1000
1250 1500 1750 2000 sample size
The boundary of test accuracy above 50% is almost linear, as indicated by the blue line. This matches the theoretical conditions (Type II) obtain the polynomial order classifica- tion error in Section 4.2.
Figure 1: Test accuracy of two-layer neural network optimized by the MFLD to learn a d- dimensional 2-sparse parity (XOR) problem.

---

Conclusion and Discussion
We provided general framework to evaluate the classification error of a two-layer neural network trained by the mean-field Langevin dynamics. Thanks to the generality of our framework, bound for specific settings can be derived by directly specifying the parameters in Assumption 2 such as R, R. and co. We also proposed an annealing procedure 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 than existing works.
limitation of our approach is that the required width and number of iterations are exponential with respect the dimensionality d; in contrast, certain tailored algorithms for learning low-dimensional target functions such as Abbe et al. (2023); Chen and Meka (2020) do not require such exponential computation. An important open problem is whether we can reduce the width and number of iterations to poly(d) for the vanilla noisy gradient descent algorithm. Another interesting direction is investigate the interplay between structured data and the efficiency of feature learning, done in Ghorbani et al. (2020); Refinetti et al. 1) Ba et al. (2023); Cousavi-Hosseinn et al. (2023).
Acknowledgements
The authors thank Matus Telgarsky for diseussions and feedback. TS was partially supported by ISPS KAKENHI (20H00576) and JST CREST (JPMJCR2115, JPMJCR2015). KO was partially supported by JST, ACT-X Grant Number JPMIAX23C4, JAPAN. AN was partially supported by ISPS KAKENHI (22H03650).
References
E. Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-stairease property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. SGD learning neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302. 1055, 1055, 2023.
Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. digh--immnsional asymptotics of feature learning: 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 httts://openreet
fort/forum? hidimmmidddddddee d=akddwRG6EG..
J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, and D. Wu. Learning in the presence of low-dimensional structure: a spiked random matrix perspective. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
D. Bakry and M. Émery. Diffusions hypercontractives. In Azéma and M. Yor, editors, Séminaire de Probabilités XIX 1983/84, pages 177-206, Berlin, Heidelberg, 1985. Springer Berlin Heidelberg. ISBN 978-3- 778-3-540-3939779,
D. Bakry, Gentil, M. Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
B. Barak, B. L. Edelman, S. Goel, S. M. Kakade, eran malach, and C. Zhang. Hidden progress in deep learning: 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: https://oppnrrviexx wttpeevieeeeeeeettreee
P. Bartlett, O. Bousquet, and S. Mendelson. Local Rademacher complexities. The Annals of Statistics, 33:1487-1537, 2005.
I. Bartlett, M. Jordan, and J. McAuliffe. Convexity, classification, and risk bounds. Journal of the American Statistical Association, 101(473):138-15) 2006.
10

---

G. Ben Arous, R. Gheissari, and A. Jagannath. High- nigh-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349- 25362, 2022.
S. Boucheron, G. Lugosi, and P. Massart. Concentration Inqquliiiies: A Nonasymptotic Theory of Independence. OUP Oxford, 2013. ISBN 9780199535255. URL hwwph:ttt//ooooooo www.olo.oo.oooo.ooooooooooooo...
F. Chen, Z. Ren, and S. Wang. Uniform-in-ime propagation of chaos for mean field langevin dynamies. arXiv preprint arXiv:2212.03050, 2022.
S. Chen and Meka. Learning polynomials in few relevant dimensions. In 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:/ /ttpc//ee/occceoeeooo mlr, ppwss/vie2 /125// cwwneha. html.
Z. Chen, Y. Cao, Q. Gu, and T. Zhang. A generalized neural tangent kernel analysis for two-layer Larochelle, M. Ranzato, Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 13363-13373. Curran As- sociates, Inc., 2020. URL https: s/proced//srosssiinngg neurips ccsssccc/ooo #Lle//aff4441555566666666622222
neural networks.
wttpcc/ppppppeee... papppapppr.... feaape_fileeee fttp//a/paapp////20///////// Paper pap.r. pdf.
L. Chizat. Mean-field langevin dynamics Exponential convergence and annealing. Transactions on Machine Learning Research, 2022. URL https: http://pppnneeveee.eeeerrrrure net/forum witoorrmmdidBiiie FiddBDdzLHigin
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 4040-3050, 2018
L. Chizat and F. Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. arXiv preprint arXiv:2002.0444 2020.
A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 5413-5452. PMLR, 02-05 Jul 2022. URL htttps://proeeeeee
mlr. wwwpppresssseeeee damian22a, html.
A. Daniely and E. Malach. Learning parities with neural networks. arXiv preprint arXiv:2002.07400, 2020.
M. D. Donsker and S. S. Varadhan. Asymptotic evaluation of certain markov process expectations for large time. IV. Communications on pure and applied mathematics, 36(2): 1983.
B. Ghorbani, S. Mei, Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint 1904::19944.119 2019.
B. Ghorbani, S. Mci, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? arXiv preprint arXiv:2006. 13409, 2020.
E. Giné and Koltchinskii. Concentration inequalities and asymptotic results for ratio type empirical processes. The Annals of Probabilityy 14(3):1143-3) 1216, 2006.
R. Holley and D. Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics, 46(5-6): 159--199,, 1987.
D. Hsu. Dimension lower bounds for linear approaches to function approximation.
K. Hu, Ren, D. Siska, and L. Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint 1rXiv:1905.07769, 2019.
X. Huang, P. Ren, and F--Y. Wang. Distribution dependent stochastic differential equations. Frontiers of Mathematics in China, 16(2):257-301, 2021.
A. Jacot, Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances in Neural Information Processing Systems 31, pages 8580-8589. 2018.
11

---

Z. Ji and M. Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint (xXv:1909.11292. 2019.
Koltchinskii. Local Rademacher complexities and oracle inequalities in risk minimization. The Annals of Statistics, 34:2593-2656, 2006.
M. Ledoux and M. Talagrand. Probability in Banach Spaces. Isoperimetry and Processes. Springer, New York, 1991. 1R1102015.
S. Mei, A. Montanari, and P.-M. Nguyen. A mean field view the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, L15(3))E7765-E776) 2018
Mendelson, Improving the sample complexity using global data. JEEE Transactions on Information Theory, 48:1977-1991, 2002.
M. Mohri, A. Rostamizadeh, and A. Talwalkar. Foundations of Machine Learning. The MIT Press, 2012.
A. Mousavi-Hosseln D. Wu, T. Suzuki, and M. A. Erdogdu. Gradient-based feature learning under structured data. In Thirty-sevntth Conference Neural Information Processing Systems (NeurIPS 2023), 2023. 2023.
A. Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles, arXiv preprint 1eXiv:1712.054333 2017.
A. Nitanda, D. Wu, and Suzuki. Convex analysis of the mean field langevin dynamics. In G. Camps- Valls, F. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 9741-97577 PMLR, 28-30 Mar 2022.
M. Refinetti, S. Goldt, F. Krzakala, and L. Zdeborova. Classifying high- dinensiinnnii gaussian mixtures: Where kernel methods fail and neural networks succeed. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8936-8947. PMLR, 18-24 Jul 2021. URL https://prroceeddngg mlr wwt.s/rss..rreeeeeeeeeeeeeeeeeee
inett i21b html.
G. M. Rotskoff and F. Vanden-Fijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.00913. 2018.
J. Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3):1800-1822 2020.
N. Srebro, K. Sridharan, and A. Tewari, Smoothness, low noise and fast rates. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010. URL https://proceeedings wewpeeeeddiinnnnenno..o.e.eeee.e htcc/appr_filleeppprr////o1///1 httpppppee/////1//// 0oof999331ee2300aaa11tt22e72e2e3 Paperrrr..ppapp.....eeeeee Paper.ddf.
T. Suzuki, A. Nitanda, and D. Wu. Uniform-in-time propagation of chaos for the mean field gradient langevin dynamics. In Submitted The Eleventh International Conference on Learning Representations, 2023a. URL / btctps://openreeee wwtrriiwweee//eeee net. http///ormmidddd_ee met:/foorum?isss_ScUUSSSSTT. JScUkk9TUU
T. Suzuki, D. Wu, and A. Nitanda. Convergence of mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction. arXiv preprint, 2023b.
M. Telgarsky. Feature selection and low test error in shallow low-rotation ReLU networks. In
The Eleventh International Conference Learning Representations, 2023. URL https:// openreview www.:nnetffooriii d=wwEski emm99
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 neural nets vs their induced kernel. In Advances in Neural Information Processing Systems, pages 9712-9724, 2019.
12

---

Appendix
A Proofs of Propositions 3 and 4
Proof of Proposition 3. Remember that
(2) R|tanh(2 aa x2) + 2 tanh( tanh(2s))/3.
..... Let b; = 2 for O,... k. let Ç be the positive real such that EE-~N (0. (aanh) tanh((+uu) = 1 (note that, this also yields EN 1) taaaaaan + u)] = -1 by the symmetric property of tanh and the Gaussian distribution), and let
Vd. Va, Vd.
0..... ETRRd d Ie dimmensioo
& d-dimension
Let
1/(x))
O 1/(2x) 0
o) E R(GHI+DD((HHH))
Σ = 0
0
and be constant which will be adjusted later on. Then, for E24 == [log(pk)s - (bb 1), 1T E Rd and $2141 : -llogg(pk)E log(ok) + 1] E Rd for o,..., k, we define
iuz = (f2i,, E), i2-- = ((2++11, E).
Then, we can see that, for 2 E (=1/v it holds that
Exnpa (2)] RE RE N( ((.1/)) (tan/////
(b. - 1)) + u] + 13/3
because we have (aI, +2 = log(pk)((E, (b. 1))+ ΣØ Uzz + ud+1 for ~ N(ET (b;- 1)] (1((2))) where U ~ (0, 1/ (i.i.d) and Σ و 2 + udtI obeys the Gaussian distribution with mean 0 and variance Σラ 2 य xi In the same vein, also have
Bxwiect [2)(2)] - NO0./ (log((())]))))))))) (pk) 2) (b, + 1)) + u] + 1}/3.
Here, define |=| = lii (1,..., k} 2; 0}| for E (+1/Vdj which is the number of positive elements of z in the index set (1, (1..... k). For fixed number E (O...., k), we
fı(2; u) (tanh[log(pk)((E: 2) -(b. 1)) u] + 1}/3, f2(2;u) = (tanh]log(pk)((5222 (b. + 1)) + + 13/3,
then we can see that
(L1((pp))
(|2) < i), ([x) 2 1),
[= [1 001
and
(2(=; [O(1/(pk))
(|=l <<++))). 1), - -((1/pk))
[1=]2+1] 1);
because (5,z) b, = EE (iggn(z)))) 2|=| - k - by = 2[[ i). Hence, we have that
f(z; f1(2; u) fz(z;u) = [12(1)
(12l = i), (O(1/(pk)) (otherwise).
Then, since tanh(u) +1 e-
Ite if |2| = and |o| $ 1/Ar.
ou
f(z;u) i(1).
and if |z| A i and |u| < log(pk)/2.
f(z;u) ((1(p)))
13

---

Therefore, when |2| =
1//
[ ((z:)g(uudu > 2(1).
where is the density function of ((0, 1/xr), and when |=| #
log(pk)/2
(22 f
loglok)/2
lu/zlogppk)/2
eep((->1 log(pk))/2) log(pk)
S 0(1((pk)) + +
O(k/pp)))
where we used the upper-tail incquality of the Gaussian distribution in the second inequality. Hence, holds that
(|2 = i), B(()) Bz-pai + B (+)] ((1/p) (otherwise),
because R= = Therefore, by taking p > sufficiently large, we also have
(1(1) (|&) is even),
f(=)
Σ 1=0
(11)(((=)) = (-1 (|2l is odd),
2(k +1) 1)
where the constant hidden in 12(-) is uniform over any |z]. Hence, there exists de > 0 such that Yf(ZZ > 22 almost surely. Then, if we let ((aa) (A) μ(aAA for E R, probability measure and a measurable set A, then we can see that f represented as
f()) = Exzay h (-)],
where
n¹ 2(k 1)
(#22,((-19) + (214),((119)))
2=0
Then, by letting C2 E(0) - (((b), have
L(p") S E(O) C2-
Next, we bound the KT.-divergence between v and μ=" Notice that the convexity of KL-divergence yields that
(L(E,)() 2(k + 1) Σ (KL(V, (2s) + (L (2g+1))
x1 log(pk) ( [||E|]] + (max [b,l + 1)3] log(1/ (1
Ollog(A ( + )ス)) ((log())2d)) (= (=R),
because k: < d, which gives the assertion.
Proof of Proposition The first assertion can be seen by L(P()) S L((() + AKL((]lu < ((0) C2 XR < E(0) C2 + c2/2 E(O) c2/2. by the assumption of such that À < (2/() The same argument is also applied to [1)4 that is, [(pp)) & ((0) 02/2.
We show that the optimal solution fuAi should satisfy | (2)1 1 (z') for all x, (+1/V d by the symmetric property of the distribution. For that purpose, we construct a group action on (-1/Vd, 1/ For notational simplicity, consider z = f-1, +1]d instead. Zd can be equipped with a group structure where the binary operation between 2, E Zd is given by the element wise product
(2=()² Let T; Z - Zd be group action that flips the show
j-th element T
(21,..., -2j.... 2a). It is obvious that T, is bijective. Then,
14

---

that there exists a sequence of indices by the chain of the group actions (o) .... o (2) o To(1)² for k = 1, 2 elements of 2d for any initial state 2 E zd and the last state turns back to the initial state, that is, Ta(24)² = 2. This can be shown by induction. d 1. we just need to take x(1) = ( and T2( To(1)=
E fl.....d such that the "orbit" generated
covers the entire
(indeed, T=(
which satisfies the condition). Suppose that the d' - 1, then we show that it is true for d = d'. Indeed, if we write be the corresponding sequence for d = d' 1, then let
argument is true for = 1.2..... .
o' d1 (1) o' (2"
(i) = (1)
for 1, ,...,, 2d 1, then the corresponding orbit covers all the elements with the last element is fixed except the initial state. Then, we flip the last coordinate at the -d' 1-th step as
a(2" -1 d,
and then we brrcce-back" the orbit as
a22 + 1) o' 0'(20'-2 i),
for = 1, becomes
2d' -1'-1_ 1. By the construction, we notice that, after 2 -1 2d' steps, the state
T To(2"-) T To(2"" 2) T=(1)= [2 -1, -1
Therefore, by flipping the last coordinate again, the state can come back to the initial state, that is, by setting a (2d)) d, we have Ta(2d) o... Ta(1)= = z while the orbit covers the entire elements of za"
We can define the action T, to [#1/Vd}d in the same manner. We note that the distribution of Z is invariant against the action of Tj. that is, Pz + P2 where T is the push-forward induced by T3- Here, let jı == J[) and f fp. We also define an "adjoint operator" Tý P - P for j E {k + 1,...,d d} such that fT((T²) fzz for all 2 E (=1/Va]a We can easily check that there exists Tÿ yu that satisfies this condition. Indeed, we may take Tý such that Tí u(A B B C) u[T; A B C)) where A (T} E A} for any measurable set A BRRd), ) B E B(R) and C' C (RR). Since the probability on product sets uniquely determines the probability on the Borel sigma algebra ((R)) x B(R) (B) = IB (Ro+2). Tī p is uniquely determined.
Now, we consider the sequence of the group action To(.).... (2(24=)) constructed above acting on the last 1 coordinates. Since the label is independent of the last d coordinates, we have L(T)( ..... () L(î) for all i = 1,..., 2d= Moreover, the symmetricity of the Gaussian distribution yields that KL(v, n) KL(V, p). Hence, if take
N
2d-k Σ To . A
then by the convexity of the objective L. we have
24-k
c(P')
Σ
o...o.. L(ā).
2d-k
Then, by the strong convexity of the objective L (due the Kl.-divergence), we have f Í.. Then, notice that
1 2d-k
fa(2)
Σ 1=1
fa(To() o.... oo())))(())))))
2d-k
for any 2 E {+1/Vd}", which means that fa(2) is constant against any change of the last d coordinates.
Next, we consider to change the first coordinates. We also construct sequence of the group action (o()....,TT[( acting on the first & coordinates. In this case, the action of Tý z also flips the
15

---

corresponding label. Hence, if 9(2) is the label of 2, then the label for T5z is -y()) ((Tzz)- By
defining the adjoint operator T as
T -
(-B) (-C)),
for j E (l....,k), k), then by the symmetric shape of lanh, we have
fa ((=)fTz((2))
for any 2 {#1/Vd]" We also have that KL ()) KL(v, T2 by the symmetry the Gaussian
distribution. Therefore, onece again by the invariance of the Pz against the action of T and the convexity of C. we have that
(()') s L()),
where
1 2
To (11) . To Σ
(8)
2
1=1
By the strong convexity of due to the KL-ddivergence term, we have uul j.. Hence,
2 M(T (o o.... ..... Ti²),
y=fa(z)
(9)
for any = E (#1/Vd]d, which means that (=]((z) is constant.
Finally, combining Eq. (9) with the fact L(A) S {(0) c2/2, we have that P/Yff))) > 0) and
Y fz( 2 c almost surely because || S 1.
B Proofs of auxiliary lemmas
B.1 Difference between finite particle network and its integral form for the optimal solution
Lemma 2. Under Assumption 1, if]]]] S for any = E supp Pz), then it holds that
27
16R2 1+2
2R
}
d
E [x-((roox =Esupp(P2)
sup
N Σ hx
(2) (x)
<
N
(AA1)2
+
Xx
Proof. Note that the left hand side can be rewritten as
N
{.)
Σ hxx((2) - / hhale)dy' (z)
Ex Exirur
sup sEaupp(Pe) Px)
N
1=1
From the standard argument of concentration inequality corresponding to the Rademacher complexity, holds that
P sup EEuupp(Pz)
hx: (2)
h2
N
21R2
2 2
sup PE)
+
N
N
t=1
< exp(-t),
for any 0, where = ( is i.i.d. sequence the Rademacher variable, 2,, (X:) an i.i.d. sequence generated from μ and the probability is taken with respect to the realization of Vr Since tanh is Lipschitz continuous with tanh' < the contraction inequality of the Rademacher complexity yields that
Er, %.
sup =Esupp(P2)
exxx (=)
N
16

---

N
2
N
52
sup zEsupp( Pz)
c T x + xx2
+
E %.
Σ
c tanh(Xi.s)
3
1=1
R
N
2
Be, 2.
sup =Esupp(P2)
(12112 +1
Σ G x x,
X 1,2
3
N
1=1
2R
1 N
27
+
Ee,
Σ N
(anh(XX;
x,3)
3
RV2
2R
3 V N BI|X! 12 + (x
3vN
To bound the right hand side, we evaluate the moment of μ²" Since μ¹ is the stationary distribution of the SDE (I). Hence, if Xo n², the process (Xe) obeying the MFLD (1) satisfies X, ~ μ²'
for any & 2 O. Then, the infinitesimal generator of the MFLD gives that
d ||| 1 R. (2x.)) (- (L(()) dt
(x,)
-AAx + 2Ad
4 (Tāl(u)) (Xx)
]
2
[I| + IE
(2X:)
+ 2Ad
12
1
2A
[|| + IE
AA1l|Xel12
+
(Xx)
+ 2Xd,
2d1
84
which yields that
[| < (2f2)2
2d
E
+
(AA1)²
A1
Therefore,
N
2R
2R²
+ d
=: D
Be,
sup sEaupp(Pz)
Σ
Ghx; (2)
2 (Ad1)2
+1 +1
N
3/N
A1
vN
t=1
Then, for x 2 0, holds that
N
)
hxg (z) - he hx())du
P
sup sza)pp(P2)
2x
N
im1
.V
8D2
}}
< min
1, exp
4173
N
This yields that
hxx(z) (
hes / (2)dd)
E
sup sEsupp(P2)
2 Σ (x) - /hale)an ( )
1
P
sup =Esupp(P2)
2 x dr
Pz)
N i=1
8
N
N
Jo
exp
4R22
dx
8D2
4R
16R
2R² (A22)²
+ d
X1 }
<
1
N
N
This gives the assertion.
17

---

B.2
Proof of Lemma
The proof follows the line of Lemma 5.5 of Chen et al. (2020). Let the empirical Rademacher complexity of a function class F be
Rad(F) LEg
sup LuEF n
Σ fz((z))
We that Rad(.) E
Rad(()]. Then, it holds that
Rad(Far(1
1
sup (a)
Σe
M=(())dd(())
1
1=1
M + E. log
Sexp
2
du du((x)
i=1
)]
M + log
E. exp
2
Σ
ec (zi)
du((z))
where the first inequality is by the Sonsker-Varadii duality formula of the KL-divergence (Donsker and Varadhan, 1983). The second term in the right hand side be bounded
Ege
exp
7 n
Σ ((l 1=1
)
< E.
exp
n
E he (z:)
)
S exp
7
R2
n
Therefore, we have that
M + 2 R²
Rad (FM()))) S 9
The right hand side can be minimized by y = vnM R, which yields
M
(ad(FM((²) - 2R
This gives the assersion.
C Proofs of Theorems 1 and 2
To show the theorems, first prepare the following uniform bound, which is crucial obtain the fast learning rate.
Lemma 3. It holds that
[pp) I(n) +(y") I(p")
2 Bad(FFKr.wwaryy (H²)) (211ff fo- V vn log logo (8n2MR)
n
uniformly over μ E with KL( S M with probability exp( f) for any >
Proof. Let Bg = M/2-* k and Cg = 2R/2-k' for 0, and FA. == {p E P KL(H, (() I, - fa-llo < Cz). Then, the standard concentration inequality regarding to the Rademacher complexity (Theorem 3.1 Mohri et al. (2012)) with the contraction inequality (Theorem 11.6 of Boucheron al. (2013) or Theorem 4.12 of Ledoux and Talagrand (1991)) yields that
P sup L(n) L(u) L(x²) L(x²) < Rad(Fn, (n*)) + Cay 2 exp(-I), MEFN,
18

---

for each k. Then, since 11f fu S 2R for any u. E P. by taking uniform bound for all
k = O,..., log2(n²) and k' = =0.... log2 (2Rr), we have that
Rad(F (1²)) + (21 - Su- 1lo V 1 2
P I(u) I(N) + I(x²) I(m²) =
n
of E Falee) 2 21- (2+ -1- logo (nM) + logg(2Rn) exp(-t)
By resetting + t + log (og2 M RR, we obtain the assertion.
C.1 Proof of Theorem
By Assumption 2, have
max(L((())O))) ((PA)C)))))))LLX S ((0) Co:
and
P ( (M (2) =1.
Let jt = (k) and n² == 11x01]- By the optimality condition of jt, we have
I(A) + ((KOT(( S T(14)100) + X (U, + c- <1/2- 28" = =cp-
Moreover,
L(p)+ + (K)
+ ("" $ ((()) ((0 - (I(n") + AU0 ,)) (T(P) I(A) + I(0²) T.(x²))
< + (T((p) I(A)) + I(y") 1 I(y")).
By the optimality of u²," the right hand side can be lower bounded
T((2)) + (T) ( (I((p²) X (T, L(p)) - I(p²) (fi 1²)) st(("))
XLO ,.).
Then, we have KL(U", Ai) S (E* + + + 1))/>((), Hence, the uniform bound on the Rademacher complexity (Lemma 3) with Mo = (6* + + 1))/A(k), it holds that
T(p) T(ia) + T(("") T(k²))
1 + log logg (8n2 Molt)
2 2 Rad Mod(FaIInnrray ; (F²)) + (2l) fp- 11o vn
n
with probability 1 exp(-t) for any t 0. Moreover, Lemma gives
2KL(μ" P)) v-1
(od(FFKL((, (1()) s v
Therefore, by setting s/n with 1, we obtain that
I(A)- I(n²) (i- 11*) 6L(1C) )
+ ( (L)
81
2 2KL (p*, + (21].a fu. 1o0 vn
s/n + log log, (8n2 Mo
2
32 (k(
4R
(L ,
+
+ xx) Ilfa - -fu 12
(K)
812 SyT + log (Sng Mo X(k
+
8/22
32R2n2
n1
19

---

with probability - (a v 6)2 I a² b² (a, b 2 0) are used in the last inequality, Here, we know that the Hellinger distance between μ³" and ji can be bounded as
where a arithmeti--eommetric mean relation and an inequality
2 / (V dz 2da(u", .R) S KL(H", A).
Hence,
I1f fu- 1occ
11 (a())(())))))))))))))))))))))))) (
/
M(x) Vur(z))
+ / - / /(Vi(o) - vu((x))
8
8 2 / (
+ ( (x))dx
(Vi(a)- (x))
s dnl ()) S KLL
By summarizing the argument above and noticing
L(i)) -(ux)(===-)))))()) - -p--
2 (
- -1²) sLia")
holds that
8y.
L(A) L(w") (a- --)) s1(1²))
1 (K)
8p
2
$ A .p)
32R2
4R
((K)
SR
syn + log log, (8n2MMoP)
Mo
nA(Kj
32R
A(K
with probability
exp(-xxpp nV) Then, letting (K),
then we have
Ilfa fp-11 11:0
32R²
n2R + 8R (K) / log 1ogg(8822 R. nX(K)
4R
< 4R2
nA(K)2
(2(
4R
(K)2
x nA(K)2
32R
4RX (1) +
+ &
nX + log log (8n² M R))
32R2n
with probability 1 - exp(-r> (K) /n) for any > 0. Hence, if we take n sufficiently large and sufficiently small so that the right hand side is smaller than oj then have
( > 0) =1
More precisely, if we let
(K)2
4R
(KK
+ (4 + log log, (8n Mj R1) }
8n = 32/21
NA(K)2
(1K)
4R +
32
and 8, O. then P (aYSa(Z) 0) 1 with probability 1 exp(-8n).
C.2 Proof of Theorem 2
Let A := (()) and μ² : Hixtoo. From the proof of Theorem 1, with probability 1 - exp(-t)), holds that
32R² 4R (1K)
8
n\(K) (1 + logg 2 R)R))
nA(K) n 32R2n²
> L(p)- I(p²) (p. -1r (L(H²) 1
öp
2
1 L(A) I(x²) - (fs - 11)) 8L(m²)
+ (
(IK) R0)
2 2
20

---

2 L(A) I(x²) (i-- - 11") 8L(u²)
8yc
Since Jn- (X) 2 co almost surely, the Markov's inequality yields that
PaY < 0)
8L(He)
[[(i)) -(yx') (E) 11))
$ E(O) (-coe((co) + E(col)
8fc
[(R) I(m²) (A - -11)) 8L(1²)
= -20(00) 2
ōf
This yields the assertion.
D Additional Experiments
In addition to the 2-sparse parity problem, here give additional experiments for the k-sparse parity problems with & = 3 and k: 4. The setting is the same with the main text, except that the width, step size, the number of iterations
modified to N 2, 000, 7) 0.2, and T 10, 000, and take different grids of n and d. We ran the experiment 10 times and plotted the mean for each n and d. The results are provided in Figure
For h: = which is shown Figure 2 (a), we can see that there appear phase-transition lines of (dd) and = (d)) which correspond to the test accuracy of around 90% and the test accuracy of around 70%, respectively. The former almost perfect classification means the Type convergence in 4.2, where the classification error gets exponentially small with respect n/dd. The latter means the Type polynomial convergence.
For l: = 4, presented in Figure 2 (b), only the linear thass-trannstion, that corresponds to the test accuracy of around 70%, observed.
Together with Section 5, we conclude that the our theoretical results are experimentally verified for the 2-sparse parity problem well higher-order sparse parity problems. The reason why the quadratic line was not observed in k = 4 would be partially because the hidden coefficients and the required number of steps are exponentially large with respect
1.0
d2
0.9
05
20
10
aCC
0.8
0.7
15
07
1953
1001
2006
5000
(a) Result with k = 3.
(b) Result with = 4.
Figure 2: Test accuracy two-layer neural network trained by MFLD to learn a d-dimensional k-sparse parity problems.
21


Reviews:
[{"review_id": "dgY0yDOybd", "reviewer": "Reviewer_afRb", "summary": "This paper proves optimization and generalization guarantees for MFLD for the $k$-sparse parity problem. It proves that if the network is sufficiently overparameterized (width > $e^{\\Omega(d)}$) then $n = d$ samples suffice, which is independent of $k$. Furthermore it proves an exponential convergence rate when $n \\ge d^2$.", "questions": "- How difficult would it be to adapt the techniques in this paper to the case of label noise (i.e. flip each label with probability $p < 1/2$)?", "limitations": "The authors have adequately addressed the limitations of this paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is able to prove concrete Rademacher-based generalization guarantees for MFLD, including explicit $\\epsilon$ dependencies in different regimes\n\nAlso, I think the comparisons with Chen and Meka (2020) are unnecessary as that paper focuses on the case when the covariates are Gaussian which is a very different problem.", "weaknesses": "- This paper does not consider label noise which would likely break the exponential convergence rate in the $n \\ge d^2$ regime. It is not clear to me whether the techniques would extend to this setting.\n- Like many mean field papers, discretizing the dynamics requires exponential width (at least $e^{\\Omega(d)}$). The tradeoff is therefore width $O(1)$ and sample complexity $O(d^k)$ vs width $e^{O(d)}$ and sample complexity $d$. However, this problem is not unique to this paper and appears fundamental to a lot of mean field analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive comments. We address the technical comments below.\n\n**Q:** *How difficult would it be to adapt the techniques in this paper to the case of label noise (i.e. flip each label with probability $p < 1/2$)?*   \n**A:** We believe that extending our result to a situation with label flipping is not difficult and the almost same results would hold (i.e., the exponential convergence of the classification error for $d^2 < n$ and $d/n$ classification error for $d < n$). \nHowever, we need to take care of non-monotonicity of the conditional expectation of the loss $f \\mapsto \\mathbb{E}_Y[\\ell(Yf)|X]$ when there is a label noise. This requires additional technical difficulty, but we think it can be resolved by using a small $\\bar{R}$. \n\nIndeed, the exponential convergence of the classification error is shown by [R1] for $p < 1/2$ label flipping in kernel classification. The essential point of the proof is to show the $L^\\infty$-convergence of the classifier, which is what we have shown in our analysis. Hence, we believe that our analysis would be a good starting point to explore more difficult and realistic settings. \n\n[R1] Koltchinskii, V., Beznosova, O. (2005). Exponential Convergence Rates in Classification. In: Auer, P., Meir, R. (eds) Learning Theory. COLT 2005. Lecture Notes in Computer Science, vol 3559, pp.295--307, 2005. \n\n\n**Q:** *Exponentially large width.*  \n**A:** Indeed this is one main drawback of the current mean-field analysis. This being said, as shown in Table 1, our result already represents a noticeable improvement over existing mean-field analyses, in terms of both the width (prior results required $N=\\mathcal{O}(d^d)$) and the iteration complexity (prior works assumed $t\\to\\infty$), even though our results can handle the more general $k$-parity setting.  \nTherefore, we believe that our analysis serves as an important step toward more efficient (quantitative) learning guarantees for mean-field neural networks.  \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n"}, {"review_id": "6KOpb9TCB8", "reviewer": "Reviewer_1qsa", "summary": "The paper conducts a theoretical study of the generalization error and sample complexity of two layer neural network training with Mean Field Langevin Dynamics (MFLD) or (informally a version of )noisy gradient descent. Specifically the paper specializes the generalization error for subset parity problem to derive bounds. Empirical results are provided to illustrate the agreement between theory and empirical observations with synthetic data", "questions": "Please check the weaknesses section\n\n- Please consider updating Section 4 to make new work clear by perhaps moving some preliminaries to the supplement\n- Please update code to ensure the results included in the paper can be reproduced without any change(s)", "limitations": "The authors do a very nice job of noting weaknesses in their discussion. No more changes necessary ", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(Note that I am not very familiar with the details of the subject matter and also related work)\n\n- Analytical generalization bounds may be of interest to the theoretical community especially given the recent interest seen with using subset parity as a dataset in analysis\n- The paper does a nice job of comparing existing and their classification error bounds for k-sparse (subset) parity problem to highlight their contribution. The fact that their bound decouples error from k is notable to this reader\n- This reviewer appreciates the experimental results in the paper. This figure summarizes the analytical findings included in Section 4\n", "weaknesses": "(Note that I am not very familiar with the details of the subject matter and also related work)\n\n- The paper appears to be spend a lot of real estate with preliminaries at the cost of sacrificing clarity of presentation of the main result. In section 4.1, I do not quite follow how the choice of regularization parameter reduces complexity\n- While analysis does decouple the dependence on k in the k-sparse parity task, this appears to be at the cost of the network width and the number of optimization steps which are both exponential compared to the work by Barak et al. (NeurIPS 2022). I am left wondering about the usefulness of this result to the community \n- (minor) I glanced at the code for the experiments and noticed that m = 2000. I assume this is the width value used in the paper. The paper states a value of 10_000 so I want to make sure that the shared code has a typo. It would be nice if the code is updated to reflect the exact value used in the paper for clarity", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for carefully reading our paper and giving insightful comments.   \n\n**Q:** *The paper appears to be spend a lot of real estate with preliminaries.*   \n**A:** Thank you for the suggestion. Our goal is to make the main text as self-contained as possible; hence we introduced the basics of the mean-field Langevin dynamics (distribution dependent SDE, log-Sobolev inequality, and so on) before presenting the main result. In the revision, we will improve the presentation and reorganize the main text. \n\n**Q:** *I do not quite follow how the choice of regularization parameter reduces complexity.*  \n**A:** As seen in Proposition2, the LSI constant $\\alpha$ controls the convergence rate. And this rate could exponentially deteriorate to achieve the perfect classification since the coefficient $\\lambda$ appeared in Eq. (5) should be properly small depending on the number of examples $n$. However, noting the exponent of Eq. (5) is $-B/\\lambda$ which is essentially the ratio of the loss and $\\lambda$, such exponential deterioration could be resolved by the annealing for controlling the term $-B/\\lambda$. \n\n**Q:** *Comparison to Barak et al. (NeurIPS 2022).*   \n**A:** First, we note that our bound gives much better sample complexity than Barak et al. (2022) -- their bound on the classification error is $O(d^{(k+1)/2}/\\sqrt{n})$ while ours is $O(d/n)$. While we achieve such an improved sample complexity using large width, this finding suggests an interesting tradeoff between computational cost and statistical complexity.  \nMoreover, although the large width is one main drawback of the current mean-field analysis, as shown in Table 1, our result already represents a noticeable improvement over existing mean-field analyses, in terms of both the width (prior results required $N=\\mathcal{O}(d^d)$) and the iteration complexity (prior works assumed $t\\to\\infty$). \nTherefore, we believe that our analysis serves as an important step toward more efficient (quantitative) learning guarantees for mean-field neural networks.\n\n\n**Q (minor):** *I glanced at the code for the experiments and noticed that $m = 2,000$.*   \n**A:** Thank you for the close reading. This is because we also conducted experiments for $k=3,4$ parity problems in the Appendix in which we employed $m=2,000$ and we uploaded the code for that setting. Although only a few lines differ for each experiment, we will include all three experiments in the new supplementary file. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n"}, {"review_id": "iFLVslTRrv", "reviewer": "Reviewer_uexs", "summary": "This paper considers the problem of learning the k-sparse parity problem with a two-layer network in mean-field regime. The main results are in two folds: 1. the authors proposed an annealing method to obtain a better rate of convergence. 2. the authors compute the classification error via computing the local Rademacher complexity. ", "questions": "1. Regarding the annealing methods: If I understand the annealing methods correctly, the intuition is that under assumption 3, the log Sobolev constant is actually controlled by the loss, and since the loss is decreasing along the trajectory, the log Sobolev constant will be larger than the trivial bounds in ( Nitanda et al., 2022; Chizat, 2022). Thus it seems to me that annealing is added for the purpose of controlling the loss properly, rather than fundamentally speeding up the dynamics (in fact, annealing does not improve the convergence rate of the MFLD in general, see e.g. Chizat, 2022 in your reference ). Also, I'm wondering if it is possible to obtain a faster convergence through different annealing methods.?\n\n\n\n2. In Proposition 4, you characterize the margin of the stationary distribution, with small enough regularization. I'm wondering if this properties still hold in the limit of $\\lambda \\rightarrow 0$?\n\n3. Generality of the results:  I'm wondering if it's possible to generalize the result to any bounded smooth activation functions, especially propositions 3 and 4, and Theorem 1 and 2?\n\n\nMinor points and typos:\n\n1. Line 169: $\\log(0) \\rightarrow \\ell(0)$\n\n2. Line 203: $\\lambda^{(k)} \\rightarrow \\lambda^{(\\kappa)}$\n\n3. In the statement of Theorem 1, you have an assumption $Q: = ... > 0$. Given a specific setting (e.g. for logistic loss and taking the parameters in propositions 3 and 4), is this assumption verified?", "limitations": "No limitations and potential negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem of learning a k-sparse parity function in a mean-field regime seems to be interesting.\n\n2. Proposition 4, which gives a characterization of the margin at the stationary distribution seems to be novel and interesting to me.\n\n3. The annealing methods proposed in this paper seem to be interesting, which gives a way to obtain a better convergence rate via bounding log-Sobolev constant properly. \n\n4. This paper is in general well-written, the main theoretical results are stated clearly, and I feel it is rather easy to understand the main results of this paper.", "weaknesses": "I haven't observed an obvious weakness in this paper, but I do have a few questions and comments, please refer to the \"Questions\" part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful comments. We address the technical points below.   \n\n**Q:** *It seems to me that annealing is added for the purpose of controlling the loss properly, rather than fundamentally speeding up the dynamics (in fact, annealing does not improve the convergence rate of the MFLD in general, see e.g. Chizat, 2022 in your reference). Also, I'm wondering if it is possible to obtain a faster convergence through different annealing methods?*   \n**A:** Recall that the log-Sobolev constant is strongly dependent on the value of the empirical loss, and we expect an exponential dependency on the sample size $n$ in the LSI constant without proper annealing.  We note that the regularization coefficient $\\lambda$ (a.k.a. temparature) should be properly small to achieve the perfect classification, which exponentially deteriorates the convergence rate without annealing.\nOur main insight is that by controlling the temperature parameter carefully, we can avoid such as exponential dependency. \nThis property is special to the logistic loss (its derivative is bounded by its absolute value) and does not hold in general settings.   \nIn contrast, Chizat (2022) did not take this point into account, and consequently the computational complexity is worse than ours -- we will include a more detailed discussion in the revision. \nThere is definitely a possibility that better convergence rate can be obtained by a different annealing method, which we leave as future work.   \n\n**Q:** *In Proposition 4, you characterize the margin of the stationary distribution, with small enough regularization. I'm wondering if this properties still hold in the limit of $\\lambda \\to 0$?*   \n**A:** This is a good point. We believe that it holds also for $\\lambda=0$. \nThis requires minor modification of our proof, because the current analysis relies on the uniqueness of the optimal solution $\\mu^*$ which does not hold for $\\lambda = 0$. However, this point can be overcome by showing uniqueness of the ``function value'' $f_{\\mu^*}(x)$ for optimal solutions instead of the uniqueness of the optimal measure $\\mu^*$, which could be shown by the strong convexity of the logistic loss.    \n\n\n**Q:** *Generality of the results: I'm wondering if it's possible to generalize the result to any bounded smooth activation functions, especially propositions 3 and 4, and Theorem 1 and 2?*  \n**A:** These statements may not be generally true because the activation function should have sufficiently large expressive power. \nThis being said, we expect similar results for wide range of activation functions as long as they have symmetricity such as $h_x = -h_{-x}$, \neven though our proof utilizes the explicit form of the $\\tanh$ activation. \n\n**Q:** *Minor points and typos*   \n**A:** Thank you so much for pointing them out. We will correct them in the revision.    \n\n**Q:** *Condition on $Q$*    \n**A:** Yes, this can be verified easily. Essentially, this condition asserts that ``the sample size $n$ is sufficiently large'' because the parameters $c_0$ and $\\bar{R}$ can be seen as constants. \nMore specifically, if $n \\gg (\\lambda^{(K)})^{-2}$, then the condition holds. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n"}, {"review_id": "QfgcIgTCYI", "reviewer": "Reviewer_zm9j", "summary": "This paper characterizes the generalization of neural networks, and prove efficient sample complexity guarantees in the mean-field region with the presence of feature learning. It presents a general framework to establish sample complexity of MFLD for binary classification problem. Such framework can be used to obtain generalization guarantee for the learned neural network in discrete-time and finite-width settings. When applied to the k-sparse parity problem, the proposed framework yields an improved convergence rate and dimension dependence.\n", "questions": "1. What is the difference of the last two rows in Table 1? Why they have the identical region/method, width, and number of iterations, but different class error?\n\n2. The work done by Telgarsky  (2023) requires width of $O(2^d)$ not $O(d^d)$ right?\n\n\n", "limitations": "I believe the authors have a sufficient discussion on the limitations.\n\nA couple of suggestions:\n\n  - $d$ has been used in the introduction several times but never been properly defined in this section (it has been defined in the table 1 though).\n\n  - In Assumption 2, last inequality $L(\\mu^*) \\leq log(0)-c_0$, I believe $log(0)$ here is a typo. The same typo also appeared in Proposition 3.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "\nThe technique used in this analysis is quite interesting and unique. Unlike existing margin bounds for neural networks that rely on norm control, this study takes a different approach by considering that MFLD optimizes the distribution of parameters rather than the parameters themselves. This perspective allows for an improved analysis of the sample complexity and convergence rate. ", "weaknesses": "1. This study is focus on binary-classification problem. Definitely it could be a good starting point, but at the same time, it is hard to argue that the results for binary-classification problem is very significant.\n\n2. The authors focus on k-sparse parity problem (2-sparse in particular) as an illustrative example throughout the paper. While I understand the 2-sparse parity problem is a well studied problem with a rich body of literature, I still do not quite understand the significance or the motivation of studying this particular problem.\n\n3. The width and number of iterations required in this analysis is of $O(e^d)$ which I think is far from reasonable in practice.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful comments.   \n\n**Q:** *While I understand the 2-sparse parity problem is a well studied problem with a rich body of literature, I still do not quite understand the significance or the motivation of studying this particular problem.*   \n**A:** The parity setting is a canonical example of learning a low-dimensional signal (since $k\\ll d$) from high-dimensional data. The low-dimensionality of target function may be exploited by *adaptive* learning procedures (as opposed to fixed bases models such as kernel methods), and thus this function class has been studied in recent works to demonstrate the advantage of representation learning in neural networks (e.g., see Wei et al. 2019).   \nIn addition, the $k$-sparse parity function corresponds to the Fourier bases on the discrete lattice. Hence, analyzing this problem can give some insight to the understanding of function estimation in the Euclidean space with orthogonal basis expansions such as Hermite polynomial and Fourier basis. \n\n**Q:** *The width and number of iterations required in this analysis is of $O(e^d)$*    \n**A:** Indeed this is one main drawback of the current mean-field analysis. This being said, as shown in Table 1, our result already represents a noticeable improvement over existing mean-field analyses, in terms of both the width (prior results required $N=\\mathcal{O}(d^d)$) and the iteration complexity (prior works assumed $t\\to\\infty$), even though our results can handle the more general $k$-parity setting.  \nTherefore, we believe that our analysis serves as an important step toward more efficient (quantitative) learning guarantees for mean-field neural networks.  \n\n**Q:** *What is the difference of the last two rows in Table 1?*    \n**A:** The two rows correspond to different sample size regimes $n \\gg d^2$ and $n \\gg d$; in other words, they can be put together into one row where the classification error becomes $\\min\\\\{d/n,\\exp(-O(\\sqrt{n}/d))\\\\}$.    \n\n**Q:** *The work done by Telgarsky (2023) requires width of $O(2^d) $ not $O(d^d)$ right?*   \n**A:** The required width in the table is coming from Table 1 of the arXiv version of Telgarsky (2023). \nHowever, we noticed that it has been modified to $O(d^{d/2})$ in its ICLR camera ready version. \nAccordingly to the paper, it should be $O(d^{d/2})$ instead of $O(2^d)$.   \n\n**Q:** *Typos and suggestions.*   \n**A:** Thank you very much for your suggestions. We will modify our manuscript according to your suggestions.   \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n\n\n"}]

Author rebuttal:

Meta-review: The reviewers unanimously find the paper interesting and result worth publication. The area chair agrees with the evaluation after reading the discussion and the manuscript.
Decision: Accept (poster)