Improved statistical and computational complexity of the mean-field Langevin dynamics under structured data
Authors: Atsushi Nitanda
Kazusato Oko
Taiji Suzuki
Denny Wu
Venue: ICLR 2024 poster (2024)
Keywords: neural network optimization
feature learning
mean-field Langevin dynamics
Primary area: learning theory

Abstract: Recent works have shown that neural networks optimized by gradient-based methods can adapt to sparse or low-dimensional target functions through feature learning; an often studied target is the sparse parity function on the unit hypercube. However, such isotropic data setting does not capture the anisotropy and low intrinsic dimensionality exhibited in realistic datasets. In this work, we address this shortcoming by studying how gradient-based feature learning interacts with structured (anisotropic) input data: we consider the classification of $k$-sparse parity on high-dimensional orthotope where the feature coordinates have varying magnitudes, and analyze the learning complexity of the mean-field Langevin dynamics (MFLD), which describes the noisy gradient descent update on two-layer neural network. We show that the statistical complexity (i.e. sample size) and computational complexity (i.e. network width) of MFLD can both be improved when prominent directions of the anisotropic input data align with the support of the target function. Moreover, by employing a coordinate transform determined by the gradient covariance, the width can be made independent of the target degree $k$. Lastly, we demonstrate the benefit of feature learning by establishing a kernel lower bound on the classification error, which applies to neural networks in the lazy regime.

Published as a conference paper at ICLR 2024
IMPROVED STATISTICAL AND COMPUTATIONAL COM- PLEXITY OF THE MEAN-FIELD LANGEVIN DYNAMICS UNDER STRUCTURED DATA
Atsushi Nitanda" 1,2, Kazusato Oko 3.4, Taiji Suzuki Denny Wu
Alphabetical ordering. 'IPPC, Agency for Science, Technology and Research, Singapore, 2 CFAR, Agency for Science, Technology and Research, Singapore,
$ DDepartment of Mathematical Informatics, the University of Tokyo,
*Center for Advanced Intelligence Project, RIKEN, "Center for Data Science, New York University,
"Center for Computational Mathematics, Flatiron Institute
-ssaars----... edu. ed. sg, tai wei.immmss....... wutpoyy.....ttoooooooooo jp, dennywu@nyu. edu
oko kazusatoog ecc.u wotkoooooaocooo ac. jp.
ABSTRACT
Recent works have shown that neural networks optimized by gradient- based meth- ods can adapt sparse or low-dimensional target functions through feature learn- ing; an often studied target is the sparse parity function on the unit hypercube. However, such isotropic data setting does not capture the anisotropy and low in- trinsic dimensionality exhibited in realistic datasets. In this work, we address this shortcoming by studying how gradient- based feature learning interacts with struc- tured (anisotropic) input data: we consider the classification of k-sparse parity on high-dimensional orthotope where the feature coordinates have varying magnitudes, and analyze the learning complexity of the mean-field Langevin dynamics (MFLD), which describes the noisy gradient descent update on two-layer neural network. We show that the statistical complexity (i.e. sample size) and computational complexity (i.e. network width) of MFLD can both be improved when prominent directions the anisotropic input data align with the support of the target function. Moreover, by employing a coordinate transform determined by the gradient covariance, the width can be made independent of the target degree k. Lastly, we demonstrate the benefit of feature learning by establishing kernel lower bound on the classification error, which applies neural networks the lazy regime.
INTRODUCTION
We consider the learning of two-layer nonlinear neural network (NN) with N neurons:
f(=) = ΣE (=), = ERd, 11,00 (z) Rd + R,
where hage (z) is one neuron with trainable parameter .:((,, e.g., we may (+) Rd and hg (2) 0(22, (l)) with some nonlincarity IR- - R. One crucial benefit of the model (1) is the ability to learn representation that adapts to the learning problem, such as sparsity and low-dimensional structures. Indeed, it has been shown that this feature learning ability enables NNs trained with gradient-based algorithms to avoid the curse of dimensionality and outperform non-adaptive methods such kernel models in learning various low-dimensional target functions (Abbe al., 2022; Ba et al., 2022; Damian et al., 2022; Bietti et al., 2022; Hoossvi-Hosseinn et al., 2022; Abbe et al., 2023).
noticeable example of low-dimensional problem the classification of k-sparse parity, where the target label is defined as the sign of the product of k < coordinates: fz,), = sign II 22). where denotes the i-th coordinate of vector The classical XOR problem corresponds to the case where the layt- parameters of the NN to identify the relevant k-dimensional subspace, which can be achieved via gradient- based feature learning (Danicly and Malach, 2020; Refinetti et al., 2021; Frei al., 2022; Barak et al., 2022; Ben Arous et al., 2022).
and input on the unit hypercube. Efficiently learning this target function requires

---

Published as a conference paper at ICLR 2024
One particularly relevant feature learning paradigm for the parity problem is the mean-field analysis, which lifts the optimization problem into the space of measures (Nitanda and Suzuki, 2017; Chizat and Bach, 2018; Mei et al., 2018; Rotskoff and Vanden- Eijndenn 2018). For isotropic input (2₃ e {-1,+1 +1)), mean-field NN can learn the parity function with linear sample complexity. Specifically, Wei et al. (2019); Chizat and Bach (2020); Telgarsky (2023) proved a (dd/n) classification error for 2-parity (XOR), when the NN is optimized by (modified) gradient flow. Very recently, Suzuki et al. (2023b) considered noisy variant of gradient descent termed the mean-field Langevin dynamics (MFLD), and showed that the O(d/n) rate remains valid for the isotropic k-parity problem with dimension- k. While the computational complexity is demanding due to the exponential width required in the mean-field analysis, one remarkable feature is the statistical complexity decouples the degree & from the exponent in the dimension dependence; this contrasts the NTK analysis where a sample size of (Ghorbani et al., 2019; Mei et al., 2022), and thus demonstrates the benefit of feature learning.
(2dd)) ) is typically needed to learn degree-k polynomial on isotropic input data
Feature learning under structured data. Noticeably, most existing analyses on the parity problem are restricted to the isotropic setting, where the input features do not provide any information of the support of the target function. On the other hand, realistic datasets are often structured, and different feature directions may have different magnitudes that guide the algorithm towards efficient learning. For example, real-world data often has low intrinsic dimensionality (Fodor, 2002; Pope et al., 2021), 1, and the observation that input directions with larger variation tend to have good predictive power has motivated various data preprocessing procedures such as PCA (Hastie al., 2009).
Recent works have indeed illustrated that in certain regression settings with low-dimensional target, anisotropic input data can improve the performance of both kernel methods and NNs. However, these results either did not take into account the optimization dynamics of NN (Suzuki and Nitanda, 2019; Ghorbani et al., 2020), or characterized the feature learning dynamics in "narrow-width" setting (Ba et al., 2023; Mousavi-Hosseini et al., 2023) which differs from the mean-field regime. Moreover, classification and regression problems have fundamentally different structures, and thus existing regression analyses do not directly translate to the k-parity classification problem. Therefore, our goal is to investigate the interplay between structured data and feature learning in the problem setting of classifying k-sparse parity function on anisotropic input data with mean-field NN.
1.1 OUR CONTRIBUTIONS
We study the statistical and computational complexity of the mean-field Langevin dynamics in learning a k-sparse parity target function on anisotropic input data. Specifically, we consider the following generating process of the data-label pairs (z,y),
2 z= AE, y sign [Te. Z),, where 2n aid Uni(((11/VQ.11V))))
for some matrix A e gaxd which controls the amisotropy. Extending the convergence analysis of MFLD in Suzuki et al. (2023a;b), we prove discrete-time and finite-width learning guarantees for two-layer neural network optimized with noisy gradient descent for this general data setting. We then specialize our general learnability result to specific examples where MFLD benefits from structured data through an improvement of the constant in the logarithmic Sobolev inequality of a proximal Gibbs measure associated with the training dynamics, and demonstrate that both the statistical and computational complexity improves upon that in the isotropic setting. In particular, we show that
When the feature directions of with large magnitude align with the support of the target function Tk. then MFLD can achieve better statistical complexity (required sample size) and computational complexity (required network width) compared to the isotropic setting in Suzuki et al. (2023b).
If we apply coordinate transformation based on the gradient covariance matrix, then the required width for MFLD can be made independent of the degree k. This is equivalent to an anisotropic & regularization, and prove that the weighting matrix can be estimated from the first gradient step.
We also provide a classification error lower bound for kernel methods in the anisotropic parity setting, which highlights the advantage of gradient- based feature learning via MFLD.
In Table we summarize and compare our results against prior works on learning sparse parity functions. To clearly illustrate the improved dimension dependence, we state our rates for a simple spiked covariance model analogous the setting in Ghorbani al. (2020); Ba et al. (2023) (see
2

---

Published as a conference paper at ICLR 2024
data
result type
regime/method
sample size
width
iterations
authors
NTK/SGD
d'
d
d'fe
and Telgarsky (2019)
two-phase SGD
d+t1/2
O(1)
d/e²
Barak et al. (2022)
upper bound
mean-field/GF
d/c
Wei et al. (2019)
Isotropic
mean-field/Gi
d/c
da/2
Telgarsky (2023)
MFLD
d/c
exp(d)
exp(d)
Suzuki et al. (2023b)
lower bound
random features
Barak et al. (2022)
MFLD
dan/c
expldaa))
expldar
Theorem 1
upper bound
Anisotropic
MFLD (transformd)
Cad 1/c
d
O(1)
Theorem 2
lower bound
kernel
der
Theorem
Table 1: Statistical and computational complexity for the k-sparse parity classification, omitting polylogarithmic terms. d is the dimensionality, and v. is the sample size. We set y sign IT &), A diag(s1, #2, ....... 8d), where
d² for 2 0, and $x+1 the effective dimension den d'=a < when (2023) do not cover the general k-parity setting, so we state the complexity for 2-parity (XOR). For the lower bound, we restate (Barak et al., 2022, Theorem 5) for bounded norm random features predictor. Finally, for the kernel lower bound in Theorem we only track the dimension dependence.
sæ 1: following Ghorbani et al. (2020) we define O. We note that results from Wei et al. (2019): Telgarsky
(10) in Section 4.1): the data-label pairs (z, y) are generated as sign [I 2i) for k = Oa(L), where the informative directions are given as the uninformative directions 2, E 1/2 (i k +, d). In this example, larger (hence smaller den = d'-a) corresponds to stronger anisotropy, which facilities feature learning due to the alignment between the oow-dimesiional structure and the target function. As shown in Table 1, this benefit is evident in both the original MFLD and the coordinate-transffrnna version.
(1-/2-1)//) (i =, L..... k) for [0, 1]. and
2 PROBLEM SETTING
k-sparse parity classification. We consider the following binary classification problem. Definition 1 (k-sparse parity problem under linear transformation). Given invertible matrix A, the impul random variable and the corresponding label Y are generated as
z= A?, Y =igi IIle 7,)
where Ż is distributed from the uniform distribution on +11/ Vdy" and ||||| < 1 almost surely.
This definition includes the well-studied XOR (Wei et al., 2019; Telgarsky, 2023) as a special case. Example (Isotropic XOR). take = Ia and = sign(2, (k: 2).
Similarly, also cover k-parity on isotropic data (Barak et al., 2022; Suzuki et al., 2023b).
As for anisotropic data, an example that we will consider in the subsequent sections is the following axis-aligned setting, where the coordinates are independent but may have different magnitudes.
Example 2 (Axis-aligned anisotropic parity). There exist positive reals > 0 (i 1.....d) such that the support of Pz (the distribution of Z) is given by S == { x {=52) x x { where Ej= 8 = 1, and (z+) are mutually independent and P( = =;) = P(z: ; = -s;) = 1/2. The k-sparse parity label corresponds to the sign of the product of k--idices J C (1.....,,
Mean-field two-layer network. Let h-))) : Idd + IR be one neuron associated with parameter =x1 (21, F2) F:3) dd+111 in two-layer neural network: given an input Rd
(z) Rt]Lanh(=T Tn + x2) + 2 tanh(ra))/33
where Ri E IR is an output scale of the network and an extra tanh activation for the bias term E R placed to make the function bounded following Suzuki al. (2023b). Let P be the set of Borel probability measures on d where d d + and P, be the subset of with finite p-th moment: |X < ec (μ E P). The mean-fiee field neural network is defined by integrating infinitely many neurons hr over ddd with the distribution μ P: f.(() hr (d)())))
Let((.)) IR XR + Rx0 be a smooth and convex loss function for the binary classification. Typically, consider the logistic loss function {(f.)) y) log(1 + exp(-y where E TR, y E (+1). We
3

---

Published as a conference paper at ICLR 2024
also denote E(yf) E(f. v) Then, the empirical risk and the population risk of f/ are defined as
L(n) -- - ビ ((nl((2(1)), L(n) [[
To avoid overfitting, we consider regularized empirical risk F(μ) == L(p) + AE xE
with the regularization parameters A, A 2 0. In addition, we introduce the entropy regularized risk: ((1) = F(x) + XEnt(()..
(3)
We can immediately see that L. is equivalent to L(x) AKT(UV, jr) up constant, where KT, jo) log(u/v)dp) is the KL divergence between // and μ., and v is the Gaussian distribution with mean O and variance 1 12(2)) A remarkable advantage of meean-field parameterization is that the above objectives become convex functional with respect to the distribution μe, since j linearly acts on fu.
3 MEAN-FIELD LANGEVIN DYNAMICS
The mean-field Langevin dynamics is defined by the following stochastic differential equation: dx,- = -V FF(m) dy4 (X,) + 22dWW, 1t La,(X,),
starting at Xo ~ no. where (20) is the d-dimensional standard Brownian motion, and (F(())) du
is the (F(M) (z) = E ((v.fu(=(())))hhaaa
first variation of F, which, in our setting, is written as \(A1l1z(13). The Bolkeer-Planck equation of SDE (4) is given by'
dep = ^Δيu + V
P4 V ( dv
V v (Alog(me) + (F(m))
(5)
Several studies (Mei et al., 2018; Hu al., 2019) showed that MFLD (4) globally optimizes the objective (3), that is, when t + do we have L(w) - ((UN)), where Mdd == argmin C(u).
For practical algorithm,
need consider a space- and time-discretized version of the MFLD, that is, we approximate the solution ft by an empirical measure 2. 1 EN 5x, corresponding to set of finite particles 2- = (x*) C R2. Let 2, (X+) C R be N particles at the r-th update (r (0. 1, }), and define fl, J12 as finite particle approximation of the population counterpart. Then, the discretized MFLD defined as follows: X ~ μo, and Ly, is updated as
N
X++1 X v/
6 AF(H (Xt) + 2
(6)
of
where 1) > 0 is the step size, J i...d. NO, 1). Note that in the context of mean-field neural network (1), the discretized update (6) simply corresponds to the noisy gradient descent algorithm, where Gaussian perturbation is added at each gradient step. We write fx == Spx for simplicity of notation.
3.1 LOGARITHMIC SOBOLEV INEQUALITY
Nitanda et al. (2022); Chizat (2022) have established the exponential convergence of MFLD by exploiting the proximal Gibbs distribution Pμ associated with μ E P. The density of pn is given by
(x(X) x exp (- 5F(M) (x)) مز
The smoothness of the loss function and the tanh activation guarantee the existence of the unique minimizer μt of C, which also solves the equation: μ = Pn (see Proposition 2.5 of Hu et al. (2019)). The key in their proofs is to show a logarithmic Sobolev inequality (LSI) the Gibbs measure Pp- Definition (Logarithmic Sobolev inequality). Let μ be a Borel probability measure on Rd. We say j satisfies the LSI with a constant a > 0 for any smooth function Rd -> R with E, [6]] co,
Ep 2 log( ())) Epp [d²] log log(E,((d)) < E V 11
a
We can apply the classical Bakry-Emery and Holley-Stroock arguments (Bakry and É'ery, 1985; Holley and Stroock, 1987) to derive the LSI constant on the Gibbs distribution whose potential is the sum of the strongly convex function and bounded function. In particular, if / |loo < B, we can establish the LSI for the proximal Gibbs distribution with a 2 X exp (-4B/>) In our case, since the logistic loss is employed and each neuron ht is bounded by R, we have B = R and therefore A exp (-AR/A)-
op
'This should be interpreted in weak sense, that is, for any continuously differentiable function with compact support, f (ddu -f ødp,. S Vø. (V (Vlog(m) V (f(e)) dd,,.
4

---

Published as a conference paper at ICLR 2024
3.2 QUANTITATIVE ANALYSIS OF MFLD
Convergence guarantee. As shown in Chen et al. (2022): Suzuki et al. (2022), the LSI constant determines not only the convergence rate of the continuous dynamics, but also the number of particles (i.e.. width of the neural network) to approximate the mean-field -iind-fiidddd limit. Let us consider the linear )
functional of distribution ((N) of N particles X = (X*) C IRd defined by
L (n(()) = NE NE ~p(N) (F()2) + XEnt(u()))
(N) be the distribution of particles 2, = (x; at the 7-th iteration, and define A, = Let MTT
1 L (1f ((xx)) Suzuki al. (2023a) established the convergence rate of MFLD follows. N
Proposition 1. Let == IEIIX
क (++ + and b, = C,L2(17 + Xn), where L = 2R- + Λd1 and C₁ = 8(R² + X2, +d)
+++ =
=d== = =(d+ Then, if aam S 1/4 and nj S 1/4,
then the neural network trained by MFLD converges to the optimal network f(x! as
E
SuPsesupp(Pa) (sx, (z) Jum (=))²]
47. Aor
A, + 2 R
where D, is further bounded by A, cxp --aann//3 /2) t 2 I (An + n/² + 4C,
Xa
xaN
In particular, for a given c² > 0, the right hand side can be bounded by e* + 2R"
after T
N
O 1 Xon log(1/c")) iterations with the step size 17 O'aa² /C1 + \e In terms of generalization error (Proposition 2). the optimization error can be set c" ((1/(n))?). Then, the required total number of iteration T and the number of particles N can be bounded by
(8) From this evaluation, we see that it is crucial select the regularization strength À so that the loss is sufficiently small. In the following section, we investigate how structured data affects the choice of X.
T O ((d + >) exp(16R) (og((A))) N s ((((((A))-))))))) = 0 (n² exp(8R/A))
Generalization error bound. Now we state the classification error bound of the neural network optimized by MFLD. For this purpose, we introduce the following assumption which will be verified later on for the anisotropic parity setting.
Assumption 1. There exists on > 0 and R > O such that the following conditions are satisfied:
There exists p* P such that KT.(vilp") S R and T(p²) E(O)) CD-
For any < col R, the risk minimizer x of L(p) satisfies Y f (X) 2 Co almost surely.
Here ep characterizes the margin of a solution μ²' and R controls "difficulty" of the problem. Indeed, if larger R is required, the Bayes optimal solution should be far away from the prior v. Hence, expect that obtaining a good classifier is more difficult. Let fa be an approximately optimal solution of L with e* accuracy: L(A) minjer L(n) + e"; we have the following generalization error bounds.
Proposition 2 (Suzuki al. (2023b)). Let Mo = (e* + 2(R + 1))/> and suppose that X < co/ R.
(i) the sample size n satisfies
(R+ A + Re + log (n2 + mnee] =:
with an absolute constant C, then fa satisfies P (Y fo (Z) < 0) 0 (the Bayes optimal classifier) with probability 1 exp(
n²²
(c SSnn))
32
(i)) When the sample size does not satisfy the condition n. > S, we still have an alternative error bound: there exists an absolute constant C such that
P < 0) < CB(co) f
(1 + t log logg (2 + (R Tp +c]
with probability 1 exp(p (), where B(co) == 1/(((0) (E(co) co (co)))
5

---

Published as a conference paper at ICLR 2024
This result states that if we take the regularization parameter À sufficiently small < O(1/R), then for sufficiently large sample size such that n > S (21/2²), we have an exponential convergence of the expected classification emror p- [Y (Z) 0)] exp(=S((n>)) otherwise, sill have a linear decay p= PPY (2( I 0) O(1 (m))() Hence, the classification error and its convergence rate is almost completely characterized by R through the choice of O R): for problem with large R, we need to pay greater sample complexity.
It is also worth noting that the value of affects not only the statistical complexity but also the compu- tational complexity. Remember that the number of iterations T and the network width N also depend on through Eq. (8). Indeed. by taking N O (expl (RR/co)), co)), which has exponential dependence on R.
co/ R, we arrive at T O ((exp(16RR/c) log(n)) and
Therefore, the goal of the subsequent sections is answer the following question in the affirmative:
Can we utilize the anisotropy of input data reduce the value of R. hence improving the statistical and computational complexity of MFLD?
4 MAIN RESULT: LEARNING UNDER STRUCTURED DATA
4.1 STATISTICAL AND COMPUTATIONAL COMPLEXITY FOR ANISOTROPIC DATA
Now we analyze how the anisotropic property of the input affects the generalization error and the computational complexity through the aforementioned measure of problem difficulty R. We first
present a framework for the general problem setting in Definition 1. Let ē (ői..... oa) E R as
o. = { V Te);;
(9)
(x).
Then, we have the following proposition that controls R in terms of the transformation matrix A. Proposition 3. Define e A where ò is defined by Eq. (9). For R there exists p E P and R such that
(L(vllue) S = ci (1 k²) log(k))²
and L(p²) S E(0) Cz. where C₁, C2 > 0 are absolute constants.
Under the conditions in this proposition, we can show that the minimizer of the MFLD objective achieves the Bayes optimal classifier with a positive margin as follows.
Proposition 4. Assume that there exists μ*' P such that the conditions in Proposition is satisfied with R and R in the statement. Then, if we choose the regulaization parameter X < c2/(2R) then the minimizer ( of the MFLDD objective satisfies
max(T L(xx)) {()) 2/
and fuy is a perfect classifier with margin cz i.e. f (Z) 2 c2 almost surely.
The proofs of both propositions can be found in Appendix in the appendix. These general results state that Assumption is satisfied for the general problem setting in Definition 1. Now we consider special cases where concrete sample complexity and computational complexity be derived. For example, we have the following evaluation for the k-sparse parity with anisotropic covariance.
Example: Anisotropic k-sparse parity. In the k-parity setting (Example 2), Assumption is satisfied with constants specified in the following propositions, which follow from Proposition 4.
Corollary 1 (Anisotropic k-sparse parity). Suppose that (Z, is generated from the anisotropic k parity problem (Example 2). Then, for k. there exists μ= E P satisfying KL(yllp²) S R where
R =
Σ 5-2 ΣS²
lkg(
NE1+
and T((²) < E(0) C2 where C1 C2 > O absolute constants.

---

Published as a conference paper at ICLR 2024
This result highlights the benefit of structured data. Observe that isotropic covariance corresponds to Vd (i ...... d), where R needs to be (kdd), which then leads to exponential dimension dependency in the computational complexity, and also simension-deeenders sample complexity, as shown in Suzuki et al. (2023b). On the other hand, if the input covariance is anisotropic so that s) > ((1/k) for j E I (i.e., the input 2, is large for the informative coordinates j E In and other coordinates are small), then the value of R becomes dimension-free R = O log(k)))-
Substituting the values of R and R. to the generalization error and computational complexity bounds, we obtain the following corollary.
Theorem (k-sparse parity setting). Define S7 == Eje,, 8/² Under the same setup as Corollary
1, we may take R = O log(k)²), = k and X = O(1 O(1/(5,, log(k)²)) so that the classification evvor is bounded by
K: log( (log(1/8) +og log log(n))
PPY < 0) 0
n
with probability 1 5. Moreover, if = S2( log(k))), then P(Y Sun S 0) = with probability
1 exppp_[2(mm/ (k* (k log(k))))]
For the computational cost, suffices to take the number of iterations and network width N as
T OS}, log(k) n log(nd) exp[O(kS], log(k))))), N = On" exp(O(kS)) log(k)?)))), respectively, to achieve the same statistical complexity as described above.
As mentioned above, for sufficiently anisotropic data such that S7, = k2, the computational complex- ity becomes completely polynomial order with respect to the dimension d; this is in stark contrast to the isotropic setting, where the complexity has exponential order with respect to d.
Now we provide two examples of the k-parity problem in Example 2 (i.e., 4s = {I...., covariance structure allows us to smoothly interpolate between the isotropic and anisotropic setting.
.a....,k) k)) where
Power-law decay. We set
of = cai 0 where a E [0, 1).
We have that S7. (dd-a) leading to R = (1 log(k)²). This interpolates between the isotropic and the completely anisotropic setting 57, = h² by adjusting between (0, 1).
Spiked covariance. Similar to Ghorbani et al. (2020); Ba et al. (2023), we set
(d for E T, 8, = O(d otherwise, where a E [0, 1].
(10)
In this case we have R = kd log(k)² which becomes dimension-free when a approaches 1. Remark For the spiked covariance setting above, the computational and statistical complexity of MFLD is governed by the effective dimensions daff = d''= defined in Ghorbani et al. (2020). As the input becomes more anisotropic, deg decreases and hence the learning problem becomes easier.
4.2 ENHANCING ANISOTROPY VIA COORDINATE TRANSFORM
From the previous analysis, computational complexity. This being said, it is worth noting that unless the problem is sufficiently anisotropic such that R becomes cost, the computational cost would still be super-polynomial in terms dimension dependence. The goal of this section is to show that the computational complexity can be further improved by exploiting the anisotropy of the learning problem. Specifically, we utilize the gradient covariance matrix to estimate the informative subspace, similar to the one-step gradient feature learning procedure studied Ba et al. (2022); Damian al. (2022); Barak al. (2022).
that anisotropic data indeed improve both the statistical and
Let aww 2) = (z) for (21, = (w, by be) for fixed by and bz. We initialize the particles Do {(wr, b,, b2 N/2 U ((-wr, -b,, --,,,,,)) ),)) N/2 by generating wy from the uniform distribution u(B..,) on the ball with sufficiently small radius > O. The gradient for each neuron is given as
j=
g(w))= = Σ (yifa ((nffx,<) (x) (w: 2()).
7

---

Published as a conference paper at ICLR 2024
Note that we have f2r (Z) 0 almost surely. We then calculate the covariance
G G= E ((wr)g(w))"
estimate the informative subspace. Define the "regularized covariance" G = G + 1oT. For this choice of G, we apply the following coordinate transform to the input Z:
2 + 1
where CA is a scaling parameter so that || S 1 almost surely. We denote by é; = accordingly. After the transformation, we train the neural network via MFLD; that is, we optimize the objective + 1 Ein1 (2:))) + (AA [|]+
Intuitively, this coordinate transform amplifies the informative coordinates (j I&) and suppress the non-informative coordinates (j & Ig). Hence the covariance of the input features becomes more well-specifier to the target signal Y leading better LSI constant. We remark that such coordinate transformation is equivalent employing an anisotropic weight decay regularization on the weight parameters r(x) ||0:||
Taken into account the sample complexity to estimate the gradient covariance, we obtain the following evaluation of the KL divergence between the prior distribution and Bayes optimal solution μ².
Theorem 2. Assume that d maxj'r s). O(). Suppose that eo is taken sufficiently small such that Σn=1 w < almost surely for w ~ u(c..) and B[]wy]] = O(1)², and the regularization parameter An is set to be so = =II sg. max 8z.- We assume that the sample size n and the number of particles N satisfies
n C2 kdR 1og(2/8))
Rdd log (d/o)
N C
(11)
II of
maxy $Ic Sj.
for given 5 (0, 1), where Cj a constant depending on k. Then, for R and sufficiently small Ck. there exists μ P such that L(μ²) I E(0) - and KL(v|I < R where
maaxyel4 83 minyel 5/' +11 log(k)}
R= ckk
for a constant 01 independent of the dimensionality d, with probability 1- 8. Here, the probability is with respect to the randomness of training data and generating the initial parameters (wi)
We make the following remarks on the thcorem.
This theorem implies significant improvement on the L.SI constant since R independent of d as max + (1), which is satisfied even for the isotropic setting. The dimension-free R
long as minyreh, " then implies that no exponential dependence is present in the computational complexity, Moreover, the runtime and the network width both "decouples" k from the exponent in dimension dependence.
In order to accurately estimate the gradient matrix, there is an additional cost in the statistical complexity. For the isotropic setting, (11) implies a sample complexity of n = ((d*+1)), which matches the total sample size in the stochastic gradient descent procedure as in Barak et al. (2022).
the input anisotropic so that IIye,, 83, > d=* then the sample complexity to estimate the informative direction is also improved. For instance, the spiked setting (10), the sample complexity is improved to n x d(1-a)+d = d, and in the most extreme case, when the signal well-specified by the principle components (i.e., denominator is ((1)), the complexity becomes linear d. This observation also demonstrates the benefit of structured data in feature learning.
Tradcoff between statistical and computational complexity. By comparing the complexity derived in Theorem and Theorem 2, we observe tradeoff between the statistical and computational complexity: estimating the gradient covariance matrix requires additional samples, but consequently the required width and iterations of MFLD significantly decrease. An interesting question is whether such tradeoff naturally occurs in more general data settings and feature learning procedures.
This condition implicitly assumes si O(1/d) (Vj): observe that Σ, wÿs} 3 < max; 55 E m3
max, ddd which can be hounded by 0(1) only when max, 83 O(1/d).
8

---

Published as a conference paper at ICLR 2024
5 KERNEL LOWER BOUND FOR ANISOTROPIC PARITY
To complement our feature learning results, in this section we prove classification lower bound for kernel methods on the (axis-aligned) anisotropic k-parity problem in Example 2 in a spiked covariance setting. We assume y = ((2) = sign == 21), and for k- I k. the input features satisfy
| = d for 1, k".
|2:| 1 for k- + d.
(12)
This is to say, the first h- coordinates of the k-parity target function are aligned with the prominent directions of the input, but the target can also depend on additional k k- coordinates that are not amplified. clear that the following two settings are special cases of the above definition:
(i) a = 0: isotropic data.
(i) = k" the (well-specified) spiked covariance model in (10).
We emphasize that most existing kernel lower bounds are tailored for regression with the squared error (Ghorbani et al., 2020; Hsu, 2021; Abbe al., 2022). Even for the simplest isotropic setting (i). to our knowledge the only classification lower bound for the parity problem is given in Wei et al. (2019) which only handles the k 2 case (XOR).
Our lower bound applies to rotationally invariant kernels: K(2, 2') = [[[|=]]: |2/|], ||, (2, z²)) as consid- ered in El Karoui (2010); Donhauser et al. (2021). In the hypercube setting, |= fixed, and hence we may restrict ourselves to a positive semidefinite inner-product kernel which can be written as
K ((,,²)= = Ei= (2.22)
{m)) non-negative and bounded.
Note that this covers the wide range of NTK of fully-connected NNs (Liang et al., 2020; Donhauser al., 2021). Given n i..d. samples, we construct the kernel estimator fø(z) with R" chosen arbitrarily: fa(=) Σ 3,K(2, 20)) We have the following classification error lower bound.
Theorem Consider the spiked covariance setting in (12). Given any fixed 0, if the sample size n S -
then for sufficiently large d, with probability at least 0.99 over the samples, for all choices of B E TR.. = 1 (2)())) 2(0) will misclassify the sign of y at 22(1) fraction of the time, that is,
PI-Pz (J<(2) < 0] = 2(1).
The proof can be found Appendix C. First, we lower bound the failure probability by the probability that [/a(2)] is large, by extending Wei et al. (2019) based finer evaluation the correlation yk(= 2)). Then reduce the problem controlling the lowest eigenvalue of some kernel matrix.
We make the following remarks on the kernel lower bound.
Recall that E (0, 1) controls the anisotropy of input features. When - 0. the input becomes isotropic, and we obtain a n = d²=A lower bound on the sample complexity for the classification error, which matches the regression lower bound in Ghorbani et al. (2019).
In the "well-specified" setting where A: = k* as in (10), the kernel sample complexity simplifies to n x d1--a)k = di which is strictly worse than that of MFLD given in Theorem for l: >> 1. On the other hand, the required sample size d times better than the covariance estimation procedure in Theorem 2 (although we believe the factor d stems from a technical artifact of the proof); however, in terms of computational complexity test time, the kernel estimator needs to store n training points which scales with d² whereas for the NN we only need to store poly(d, k)) neurons, which decouples the degree & in the exponent the dimension dependence.
6 EXPERIMENT
We validate our theoretical analysis by numerical experiments. We consider the anisotropic 3-sparse parity problem: 21 22 $2 $3 = k/ /d, and ... ... = 1/vd. We fixed d = 300 and varied n and & to train the neural network (2) with the logistic loss. More details are Appendix D.

---

Published as a conference paper at ICLR 2024
In Figure I we plot the test accuracy func-
10125
tion of the sample size n and R, which con- trols the level of anisotropy. As clearly seen, increasing i enables smaller the model to learn the problem with smaller sample complexity n, which demonstrates how anisotropy helps learn- ing. Moreover, let focus on the "phase transi- tion" boundary between yellow and blue regions. According Theorem 1, the classification error is bounded by Zj 8, 8///n /n x 2d/n up to a constant, which predicts that there would be a boundary around x² (n)) as indicated by the red line in the figure. We therefore conclude that the empirical findings match the theoretical result in Theorem
a²
10
0.9
0.8
10925
0.7
1095
-0.6
10925
0.5
10a
10a
sample size n
Figure 1: Test accuracy of NN trained by MFLD to learn an anisotropic d-dimensional 3-parity problem.
7 CONCLUSION
We studied the interplay between structured data (in the form of input anisotropy) and the efficiency of feature learning in the context of classifying kisparse parity using two-layer neural network optimized by noisy gradient descent (mean-field Langevin dynamics). We showed that anisotropy can improve both the statistical and computational complexity of MFLD, leading to a separation from kernel methods (including neural networks in the NTK regime). Interesting future directions include (t) extending this observation to more general class of target functions, (iz) improving the sample complexity to obtain the covariance estimator Theorem 2, and (iii) deriving a more precise description on the tradeoff between statistical and computational complexity in NN training.
ACKNOWLEDGEMENT
TS partially supported by JSPS KAKENHI (20H00576) and JST CREST (JPMJCR2115, JP- MJCR2015). KO was partially supported by JST, ACT-X Grant Number JPMJAX23C4, JAPAN. AN was partially supported by JSPS KAKENHI (22H03650).
REFERENCES
E. Abbe, E. B. Adsera, and Misiakiewicz. The -ttird-staircsse property: a 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 on neural networks: leap complexity and saddle-to tiee-tto-dddlll dynamics. arXiv preprint arXiv:2302.11055,
2023. 2023.
J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. nigh-dimensional 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 /tttps://cperreee
net/ / /e//fforue forum? d= akddwwGGGG..
Ba, M. A. Erdogdu, Suzuki, Z. Wang, and D. Wu. Learning in the presence of low-dimensional structure: spiked random matrix perspective. In Thirry-evvntth Conference Neural Information Processing Systems (NeurIPS 2023). 2023.
D. Bakry and M. fimery. Diffusions hypercontractives. In J. 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- 978-3-540-3939779
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,
10

---

Published as a conference paper at ICLR 2024
D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL nttps://ogg hatp:s///ppeeevieww net 2tt.://wwwwe/eeueeetmmem "oorummdd=8X00 2ewX- im.
G. Ben Arous, R. Gheissari, and A. Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349- 25362, 2022.
A. Bietti, J. Bruna, C. Sanford, and M. J. Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, 2022.
F. Chen, Z. Ren, and S. Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050, 2022.
L. Chizat. Mean-ficlu langevin dynamics Exponential convergence and annealing. Transac- tions on Machine Learning Research, 2022. URL pettps::/openrevieee net fott/fooummddd
EDqzIH1e gEm.
L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. 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 with the logistic loss. arXiv preprint arXiv:2002.04486, 2020.
A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In P.-LL Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 5413-5452. PMLR, 02-05 Jul 2022. URL https: http/poocceeeeee ngs ml wal.r.pres//ve /visss//vvvv /am.////mmmmmmmmmmmm 222 hhmml. html.
A. Daniely and E. Malach. Learning parities with neural networks. arXiv preprint arXiv:2002.07400, 2020.
K. Donhauser, M. Wu, and F. Yang. How rotational invariance of common kernels prevents general- ization in high dimensions. In International Conference on Machine Learning, pages 1804-28144 PMLR, 2021.
N. El Karoui. The spectrum of kernel random matrices. The Annals of Statistics, 38(1):1-50, 2010.
I. K. Fodor. A survey of dimension reduction techniques. Technical report, Lawrence Livermore National Lab., CA (US), 2002.
S. Frei, N. S. Chatterji, and P. L. Bartlett. Random feature amplification: Feature learning and generalization in neural networks. arXiv preprint arXiv:2202.07626, 2022.
B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. 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:20066113409 2020.
T. Hastie, R. Tibshirani, J. H. Friedman, and J. H. Friedman. The elements of statistical learning: data mining, inference, and prediction, volume 2. Springer, 2009.
R. Holley and Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics, 46(5-6): 1159-1194, 1987.
D. Hsu. Dimension lower bounds for linear approaches to function approximation. Daniel Hsu homepage, 2021.
K. Hu, Z, Ren, D. Siska, and I.. Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint / 2019.
Z. Ji and M. Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint arXiv: 1909, /2292, 2019.
11

---

Published as a conference paper at ICLR 2024
T. Liang, A. Rakhlin, and X. Zhai. On the multiple descent of minimum- norm interpolants and restricted lower isometry of kernels. In Conference on Learning Theory, pages 2683-2711. PMLR, 2020.
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, 115(3))E7665-E7771, 2018.
S. Mei, T. Misiakiewicz, and A. Montanari. Generalization error of random feature and kernel meth- ods: Hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis, 59:3-84, 2022.
Misiakiewiez. Spectrum of inner- mener-prouuct kernel matrices in the polynomial regime and multiple descent phenomenon in kernel ridge regression. arXiv preprint arXiv:2204. /0425, 2022.
A. Mousavi-Hosseini, S. Park, M. Girotti, Mitliagkas, and M. A. Erdogdu. Neural networks effi- ciently learn low-dimensional representations with sgd. In The Eleventh International Conference on Learning Representations, 2022.
A. Mossevi- D. Wu, T. Suzuki, and M. A. Erdogdu. Gradient-based feature learning under structured data. In Teirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
A. Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv:1712.054335 2017.
A. Nitanda, D. Wu, and T. Suzuki. Convex analysis of the mean field langevin dynamics. In G. Camps- Valls, F. J. 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-9757. PMLR, 28-30 Mar 2022.
R. O' O'Donnell. Analysis of boolean functions. Cambridge University Press, 2014.
Pope, C. Zhu, A. Abdelkader, M. Goldblum, and T. Goldstein. The intrinsic dimension of images and its impact on learning. arXiv preprint 204.:v::244.000 08894,, 2021.
M. Refinetti, S. Goldt, F. Krzakala, and L. Zdeborova. Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In M. Meila and T. Zhang, editors, 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 press/ vi re refneetiill html.
ht https: https://ppcceeedngg
G. M. Rotskoff and E. Eijnen- Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:/805.00913 2018.
Suzuki and A. Nitanda. Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic besov space. arXiv preprint 1rXiv:1910.12799 2019.
T. Suzuki, A. Nitanda, and D. Wu. Uniform-in-time propagation chaos for the mean-field gradient langevin dynamics. In The Eleventh International Conference on Learning Representations, 2022.
T. Suzuki, D. Wu, and A. Nitanda. Mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023a.
T. Suzuki, D. Wu, K. Oko, and A. Nitanda. Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023b.
M. Telgarsky. Feature selection and low test emror in shallow low-rotation ReLU networks. In The Fleventh International Conference on Learning Representations, 2023. URL https://
openreview nnt/ forrmm/rrumm//
R. Vershynin. High-dimensional probability. University of California, Irvine, 2020.
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

---

Published as a conference paper at ICLR 2024
Appendix
A PROOFS OF PROPOSITIONS 3 AND 4 AND COROLLARY
Proof of Proposition 3. We follow the proof strategy from Suzuki et al. (2023b). Remember that
(z) R|tanh(z x1 [Tx1
x2) + 2 tanh (anh(z))]/3.
Let b. = 2i for i .,...,k k let >> 0 be the positive real such that (o..N( (2 (anh(C +{]]] = 1 (note that, this also yields (o, tanh(-S + u)] = -1 by the symmetric property of tanh and the Gaussian distribution). Let
(2
O 1/(2))) 0
o)
E R R(CHHTID((CHHH)
Σ =
0
and be constant which will be adjusted later on. Then, for E23 =- [og(pk) log(pk)(b)-
1), CT u2 and E2j+1 == -[log(pk)o" log(pk) + 1).51 T E Rd for j O....,k, k, we define
Poj =- NE2j, E), A2j41:= := N( E).
By noticing that for E supp( P2) there exists e (=1/Vaj such that 2 = Aă, we can see that
Exmpas [!x(2)] = RE (tw.. (by - -1) + u| + 11/3
because we have
[0)] (b; - -(0 + (I1, + +x2= (() -
Σ
U + ud+1
1=1
= ((- AE) - (bj - 1)) + Σ
1,
U Md+1,
for x N(ØT (b; - 1))T I/(2x)) where u4 ~ NO0 1/2x1))2) (i.i.d) and Σ Di=1 4: + ud++
obeys the Gaussian distribution with mean 0 and variance 1
2\, || -
1 (1 || S
for all 2 supp(Pz), where we used the assumption on A. In the same vein, also have
E Epppay+i [2)] = N(0 (=11)//2,) (tanh[log(pb)) hAt.://NOOC.H/////////aaaaaalaO 2) - (b, + 1)) + u] 13/3.
Here, define | = lii I 2, > 0}| for E supp(Pz) which is the number of positive elements ......,,, k), we let
of 2 in the informative index set Ig. For a fixed number j E (O.....
(b, 1)) m| 1}/3, f2(2;u) = (tanh(log(p+)((0..) - (b, + 1)) + w| + 13/3.
(=) (tanh
then we see that
f= = [O(1/(pk)) (lžlo < j),
(1-0(1/ppk)
(lã 2 j),
and
[O(1/(pk))
(13l < + 1),
f2(z:0) {1 (1-((1((pkb)
1),
because (6,2)-b, b, -b) E3 sign( b, = 2/0/u --
b, = 2(1² - 5). Hence, we have that f(z: u) fi(z; u) fe(z:u) [2(11
(1 j) (otherwise),
(0(1(pp))
and f(z; u) > 0 for |// = j. Then, since tanh(u) + 1 2"t4
c--e""
1+e if 17/0 = j and
|u| S 1/Ar,
f(z;u) 2(1),
13

---

Published as a conference paper at ICLR 2024
and if [ē|u # and |u| log(pk)/2.
f(z;u) & ((1/ppk)
Therefore, when -ž1o = j.
1/21
Lo=N Pu..)o.. 2 -1//A, ((=:))g(u)du > 22(1).
where is the density function of (O, (1 |=11) and when iž/0 +
log(pk)/2
((::))
f(2; u)g(u)d=
log(pk)/2
1ul21og(ah)/2
(exp(->) log(pk)9/2) log(pk)
- {((11(pk) + +0
=(LL(ph))
where we used the upper-tail inequality of the Gaussian distribution in the second inequality. Hence, holds that
(2())) (L(pp)
(1/ j),
f(22:= 1 (2)] + Bx-aaz4 [h,
(otherwise),
because R = k. Therefore, by taking p sufficiently large, also have
(2(1))
(lēlo is even),
f(z) 2(k 11)
Σ (-1) (
(11(1) (lēlo is odd),
t=
where the constant hidden in 52(-) is uniform over any 12/0- Hence, there exists dg 0 such that Y fZ) cg almost surely. Now if we let ((a) (B) μ(aB) for R, probability measure and measurable set B, we can see that is represented as
f(+) = E Exry (/2 (-)l
where
&
Σ (A2x + (211),
/² 2(k + 1)
2=0
Then, by letting C2 ((0) ((b)), we have
L(p") S {(O) C2.
Next, we bound the KL divergence between v and μ². The convexity of KL divergence yields that
KL(P,p") 2(k: + 1)
(KL + KL(V, P2)+1))
S A1 log( ) + (max |byl 1) + log(1/(2X1)) + +1 (?)
0 2 (116112 k²))
This gives the assertion.
Next, we prove Proposition
Proof of Proposition The proof of this statement resembles Proposition of Suzuki al. (2023b).
The key step in their proof is to show that the optimal solution satisfies
(=)x (=')
for any z, E supp(Pz). We prove that this still holds in our general setting, Let TA : Rd - R be
TA
(Am1,
14

---

Published as a conference paper at ICLR 2024
where (21, 22, 135) for 21 E R 22 R and 23 R. Then,
that
fx(2) = (TA)
for μa E P and Ta is the push-forward with respect to Ta. and z = Až. Based on this coordinate transform, we can reduce the problem the standard parity setting where the input obeys the uniform distribution on {#1/vd] According this coordinate transform, the prior distribution is transformed to := TA v, which is again normal distribution with mean 0 and variance AAT (2X1). We also let T, be the map which flips the sign of the i-th coordinate. Then, the key argument the proof of Suzuki al. (2023b) show that
(L(vAlle) K (vallTj*H))
for measure Y E P (which is supposed to be TA ja for population risk minimizer jt). This equality true because the normal distribution point symmetric. Indeed, we have
KL(vAllu) K A (4)) = KL(vallT;**),
where the first equality by the invariance of the KL divergence against any bijective coordinate transform and the second equality is by the point symmetry of the normal distribution. Then, following the same argument to Suzuki et al. (2023b), we obtain the assertion.
Finally, Proposition can be obtained as a corollary of Proposition where we set diag syv............ suvd). For this setting, we can easily see that
|
Σ
$
jEL
Combining with this evaluation and the fact that
k=
Σ
Σ s,5,
8
Σ
8
IE
iEh
EE
IE
iE
we obtain the assertion.
B ESTIMATING THE INFORMATION MATRIX
Without loss of generality, we may take 47 (L....,kk) (w, bi, b2) for fixed by and by. Then,
k}. Let ( 2) kx(2)) for (I1,
( T 2) = Σ
1=1 El (0)( 2)
Note that the gradient of the loss with respect to wy be written
( e' (v.lno (2:)) (wTzz) 2x).
Suppose that fno (2)) = O, then since Y IL (87 Zz), its expectation can be expressed as
gj(w) = E
II (57) 20 (w 2) J'El
(1) If j E Tg. then we have that
9 (ww :
II
Sy
II J'Ele'
Zyo'(w Z)
j'ENN
15

---

Published as a conference paper at ICLR 2024
Then, by the Taylor expansion of σ. it holds that
(,(w) 83
II
E
II
Zzod ((0r)) 22)
2.)
el
S'EJN
t=0
J'En'
6=0
+
E
II Zy(o'((0)) 2)
k=
J'E
6=0
=
II
3''
E
II
Zy' (k - 1)! (w) Z)k
Ck
VEL
J'ENU
E
Zy
CE+1 A
(w (w Zyy Z)0
J'ETA
8j
II PELL
sj
II J'ELN
Cj
(lc 1)! JEEL
II
wy + (higher order term)
(k 1)!
C
II sy
II J'EL
wy + (higher order term).
Je
The higher order term (a) in the above expression can be evaluated as
E II
C++1 Z
(w
Z)d
!!
t=k
j'El
00
Zz'
C++1
Σ
Σ
Σ
Wy Lg'
8=k
d'El
('ElTu
22-1-11
6--60-1
Σ
Wy Zz:
Wis Zi
Σ
Wz
Wjs Z,5
1-f
Σ
W.yZ. Zi'
Win- Lioes (wT2) 2) -y---
Je Se-z
8
E
II
(55
C++1
Σ
E
Σ
wy Zy.
&=k
JEN
Giunth 1)
(PElTe))
8--1-11
1--6-11
Σ
Σ
2y
6x-1- 88-2-1
}
Σ yeleAUSn.hh.......
Wy. Zzj'
(w T 2) 19
Then, applying Hölder's inequality, we can bound the right hand side as
6-b+1
c-
Σ
II
(g)
C2+1 1
Σ
Σ
E
Σ
Wy Zz.Zzz
t= J'ENU
Cta-
('ELLN
16

---

Published as a conference paper at ICLR 2024
4- 86-111
f-
E
E
Wy
E
E
wy Zy
('E(TA((isi))
J'E(Te Midaimtt-a))
-4+1
E3
w
Now, the moment bound of sub-Gaussian random variables to bound each term wy Zz-1<-+t1] "u-fu=1+'
x-x-1 Indeed, we can see that, for any index set J C I. wj N sub-Gaussian random variable with parameter |wy O 85 where sub-Gaussian random
E1 Σ
variable with a parameter u satisfied [[]]]]]] I (eu) (Ae 2 1) with an absolute constant e (see Proposition 2.5.2 of Vershynin (2020), for example). Hence, by noticing that [|w, O s) | < || s||2
and 1 (66 4--1 1) + f--1
f- the right hand side can be bounded by
II PEEA
II J'ENAV
By
C++1 !
(k
1)1 (e 1s + 1) 1)! (cll 8|]) ( k 1)(C
e!
t=
- k+1 (L k + ()
II
II JELV
Σ
(cllWe| 2 a|))
Wy
(E ( +
JELL
t=
Then, by the Stirling's formula, the right hand side can be bounded by
(e- h (c-K+1)(-R+1))
II
II
a-illw O s|| c-1
V2T - 1 1-++1)=-+11 122 (f--h)
JEELM
J'EL
t=
6-h+1
II
2
II
Σ
c-illl N 8|| 1-k+1
=
2T
(e ( k 1)1/2
(I (e-k) k 1
J'ELN
J'ENN
t=k
Ck
II
II jELV
|wyb
2
j'ELVV
where we used the assumption
|| O s|| is sufficiently small such that cetieellu @
t-h+1
f|]) V2n
@ Therefore, we can see that
(e-8+1)173
(y( CK
II 8y'
II
Wyy + (higher order term),
JEEL
J'EN'
by (w)) 2
Ck
Syy
II J'EW'i
|uy.
2
j'El
(
<
2
Ck
8y'
II
|uyl
(13)
j'L
(2) In the same vein, also have that, for
Ig. it holds that
(w) = Ck+2
II
SY'
II J'ELUj
wy + (higher order term),
j'Elnuj
(w) I 2ck+2
II
$y'
II J'Elujj
wyl.
(14)
JEE
Next, we show the concentration of the empirical gradient 9 (w) around its expectation. Observe that
sup Y.2
( fr, (Z)) Z) < Rs,,
Here, 1 : (I.y))
17

---

Published as a conference paper at ICLR 2024
(wTZ) Z)] s sJ. Vary, (Y (2)) (w
Therefore, by Bernstein's 's inequality, we obtain that
P (lss(w)) (w)) 2
4Rs v
log log (2/0) 2
for any 5 E (0, Hence, if we let n
16k R log(2N/5)
(Coc* IIyen, $y) 2'
for sufficiently small constant Co. then we have that 1g (wi) (y(w:) S Cock II f'Ele
8y 8 Vid,
(15)
uniformly over 1 = I,..., N with probability 8.
For that purpose, we evaluate the expectations of gi (w) (w) carefully. Let H(w) 04+1 Ei= (e-k+1)) E2 [(w 2) f3]++]]] 21 C s||? + Ei=o (1+20)!
0444+20 Ez [(w 2) We evaluate
for each condition on ji and je. (a) If j2 Lx. then it holds that
Ew 1ū (W) (W)) ok II
Sy
II jEEA
W3(1+ H(W))²
JLE
Ω
II j'El
(b) j1 # j2 and J2 E Ig. then holds that
Ew (W) (W)) o Ir, of.E [L W3. W,, W ja (1 + H(W)) 0,
where we used that the distribution of is symmetric and H() satisfies H(W) H(-W).
(c) If + ja and jı 4 and ja Is then
Ew Di (W) (W) = CxCk+2
II
II W3 (1 HHW))2 Jrelh
f'EL
(d) ji @ la and j2 E Jx. then
Ewla, (W (W)gi, (W)] = II
W3 Ws, W (1 (1+ H(W))2 j'El
II
y'E
( # je),
o LO 87)
(51 j2)-
Summarizing these evaluations, we see that G (G j E Rdxdd defined by
ddd
G -Ew W))ga (W))
is diagonal matrix where G)n.j for f1 E Js has larger values than that for jn Ig. We define its
empirical average version G = (G j1.t2 i E Rxd as
ddd
N
G =
Σ
gn (w)).
N
E=
18

---

Published as a conference paper at ICLR 2024
Now, we show the concentration of G around its population version G. Note that
95 (1))))((()))
(o (wx) 9t. (wi) + 9t. (wx)) (w)) IoH (wi) + 955 (wi))
N
N
(9 (wi) D2 (wi) (wi) 9ia (wi))
N
(e (wl) - Ds. (wr))ā;=(u)
N
Σ
Oga (wi)
Isa (2x)) (w))
N
N
1
+
Σ
Ii (wr)āss
N
Therefore, if we let AG Gj Gj Gr to En Is (wr)ar, (wi) and Girzz
N Ei=log (wi) Iú. (()z(())) 9i (wx)), then for any TRdd it holds that
1
( N Σ (g(w)) - g(wr)) [(wi)) + 2u
lu AGul =2u T
Σ (1m)) (()))
22u Gu + 2u (G G)u + 2u Gu.
(16)
Here, the term G - G can be bounded by the matrix Bernstein's 's inequality as
p j Gllop K (v Q (P(L+log(d)) log(d)) (1 log(d)1Q) )
exp((tt),
N
N
where K is an absolute constant and Q d fore, taking N = (dd2x² log(d/5)/ max ælx s4,)) for sufficiently small Co yields that
Il
En, $7. because (g 0g (wr))lop S SQ()). There-
|| - Gl||| up O
Co II J'E
b max 22y
with probability 1 1- 8. In the same manner, we also have that
G c GG- -.
(17)
In the following, we let Q1 = IIve 03. and Q2 = =II 83. maxyen 8z.
Then, by modifying the objective as
L(u) A1 1 (
with regularization parameter &o Q2- This is equivalent the alternative objective
[(x) + >1 [,, [II [ where the input is transformed as Z +- AZ where A = CAVG G + AT with B diag ()) vd. ...... ... .d......vdd and constant CA = ((kQ)) maxy -3/2- such that ||4]]| S 1.
Indeed, we can take such ea because
VG+ +OIIBI
< k 20 + (G (G-G)+33) BŽ
d
Σ
Q1/(kd) + Q2 (mxxmax
max 87)) + max sy (kQ] + (d- -x)(2)
3' af
fid1=1
19

---

Published as a conference paper at ICLR 2024
S kQ max sfy
where we used the assumption of Σ 57 and the fact that dQ2 dQ1 maxyer. 83. x 1 due to the assumption d maxx EI, 83. = ((1) Then, we can see that
||A 8|| = = 2- B (G (o1)- B = c
(G )oo) )oo)
for Ca = (0T ...... Sk ; 0)T. Now, let
G G+
[1.]]
G1.22
G(2.21
Then, we can see that
(G + ~01)-1= (G11,1) G1,2 G2.1))
We know that || llop $ 11(Aoo)) llop = O5' and [| llop 2 Co koykkldk)Qiie 1 (kd) by the argument Eq. (16) and the assumption d dmxyy gr, 57 O(1). Hence, that
[2.2]
-1,11--Gip - G1.2 G2.1 2 Q1 ((C)(k(d
2 [Q1 O(C80)14,
by Eq. (17). Therefore, by taking Co sufficiently small, we have that
(Gp.a) G2 GR.n)) 3 Qi I.
Therefore, finally arrive
||-2 S 2 11s ||, G1.,3 12.2 G2.11) 1lop
maxy's 8
S kQ maxx))))
og Qi k22 j'E
min
minyel $3
C KERNEL LOWER BOUND
In this section, derive the kernel lower bound for the k-parity classification problem (Example 2) in the spiked covariance setting, Our proof is divided into two steps. First, we translate the event when prediction fails into when the value of [fa(2)] is away from zero. We combine the proof for 2-parity (Wei al., 2019) and an additional observation that K(2, 2*) have d-(1-a)k" (a111)1""-k==- -(k-k") correlation to to get the tighter bound for general higher order parities than that in Wei et al. (2019). Then, we show that the probability of that event is evaluated by the the smallest eigenvalue of some kernel matrix defined in Lemma 3. Finally, we apply the lower bound of the smallest eigenvalue using the technique of Misiakiewicz (2022).
Note that, we do not need to prove Theorem for 1 < a S 1, because L(1 ar = 0 holds for such a. Thus in the following we assume 1 1. 1, hence 1 0.
Lemma 1. For n S (11-))++((e-X)) with probability 1 - exp(-S((d) over the random draws of the training sample, we have
PI-Pz (o(z)y
< 2
[s(()| d'1-a)
Σ 136b
1/d,
where e a constant depending on and {-))-
Proof. Consider randomly drawn 2k41:d, which fix in the following. Suppose fo(=)b(2) 2 0 for all choices of zi:a and [fx(2)] Z माक umx En [&&] for some 2isk for the sake of contradiction (with high probability). Then, consider the average of K(2, 2' over the choices of 21s as follows:
TTT2)
E.st K(2, x((z)))(()) 2k+1dd = Er
Σ
a
8(2)
2k+1nd
d
t=0
20

---

Published as a conference paper at ICLR 2024
0 Σ
2 1² 1
2 1:d
(18)
1=
y'=1
Let us evaluate E [(
I1 2z syl2++1al For k
2(11 aa)k (1-
we expand
1-0
(z = (E
2,2} y obtain
k-k*
E Bxak
2'
2
2k+lid
S
(d-k)) k))
d
1
5'=1
P=
terms containing each 21, Sh. at least once
$ d (1-a)k"
(k--k")
For 1 2 [ [(1=0)]*+(h==3) + 1, we have ا S d d1-100) Mogd with probability 1
a/d over the choice of 2k+1:d. and therefore Σ -1=2k-1 Ba d |2 X d1-1aa (k- By using them for (18), we have
Exsa [(2, ()))) 1o41:dd] = EZL.
Σ
2(2) 2k+ S (1-1--))---(k6
d
1=0
for randomly drawn Fstidd, with probability more than 1 - 1/ (2)= +()=""((+++ +(=- Therefore,
E [Jo(=)((z) =x+1:d] = Beis
Σ Br (,,((z=')u(l
2k+1:d
(19)
Σ
18:1
(20)
d(-aa)k +(k-k*1
1=1
with probability more than - 1/d.
On the other hand, if f=(2))(2) 2 0 for all 2i and (fa(z) 2 200 +(6- En= 21:k, we have
18 for some
Eii. [fa(z)y(2) 24 =
3x41:a] 2
1
Σ (s())((2) Bisk
131
(21)
2 d(1-a)k- + (k-
By comparing (20) and (21), we have the contradiction for more than - 1/d probability of the choice of #k41:d by taking sufficiently large. Therefore, if [f(()]] ス (ol=ap=-ax- > |8: for
some 21:k, there exists some 21sh that yields fz(z)y < 0, for 2k-1:d that is drawn with probability more than 1 1/d, which yields the conclusion.
Now we evaluate the probability PZ-Pz [[ʃ((z)] 2 nu on la Ey El [Bl]]. Since fa(2) have very high order terms, make the following approximation. Lemma 2. Let us define g1: [-1,1]
R
(t)
n
1
Suppose n S -1--a)k-(kk-R Then,
1 Tz
PI~PE Ps
[ [n],
Die
K(,,2t)- 91
- d d ((--x))
2 21 -1-1)))))-(((1)))
l/d.
Proof. First, we note
2(1--)6"+(0-4311
(2 TT)
T xt
2 TT at
(((22))- - 91
Σ
71
d
d
1=0
21

---

Published as a conference paper at ICLR 2024
z xTx²
(22)
d
2(0 +(8
[+1
With probability 1 1/d ++1 ++ 111//2
d
(1-a)/2 Vlog d. This means that (22) is
(I
bounded by S
log -
(1--a))+((k-+ for sufficiently large d. By
taking the uniform bound over all i, we get the assertion.
Owing to this lemma, it suffices to bound P P2~P2 Ei=1 Bigi ( by ((1), because
+(1-1) Ei=1 [3:]]
PP Pa
[ 2 +11-a)b++(1--
13,1
2 Pe
Σ Begn
c+1
Σ
| - 1/d.
d d( ajk +k
1=1
For this, we lower bound the second moment, which captures the variation of fg. Lemma 3. Suppose on are all positive and define K2 E R as
2k+ldd "RtLid
(K2)as = Σ
d d-k
I=
Then, for sufficiently large d, we have
Ex
Σ
3.91
2
2 d 111 1(10-))"++(kkX3)) K28.
d
The proof requires several auxiliary lemmas follows. We defer the proofs of them after the proof of Lemma 3.
Lemma 4. For any integers p.g
Σ 31(z T2
Σ Bi(=T22l 24)0
E,
]
2 Es
Be(2k E (u(ertaaeettdd
Σ ta
20
Lemma S. Let 2t, a E {[1,111, E 11, be vector sampled uniformly from the hypercube, and let be any integer. Then, we can expand the expectation as
2d
23
E2
caiI
d
P=0
Furthermore, for sufficiently large d, ca,, 0 and especially Ca,1, (11)².
Proof of Lemma 3. Let us first expand the target:
27
12
2
E
Σ
3igl
d
2(
E,
3
7
d
1=1
1=0
22

---

Published as a conference paper at ICLR 2024
[[00=)) +(6===")
E,
71
B
d
1m0
Y4 n2
Eg
Bi
Σ
3
(23)
From Lemma and My 712 > each term is non-negative and (23) is lower bounded by
27
Σ
n B
Σ
3c
d
2001 "W"" 1.kk 101- 12
))
27
n Σ
Ik
T
20
71
E
EEE l:d
Bi
(24)
d
i=1
Let us define matrix K1 _nx= so that (24) is equal to 3 K13, For that, we define
2(000)) (200-0)=++0--0 +(0-6"))
T (
)
(Ki)as= =
H Ezy+ii
d d
- k
According to Lemma 5,
2(1-m))
Shande #b441ed 1'
(K1)as
(d "30''s-kiit
d
1=0
P==
20
20(1 wi 1e8 11111
2 2k+1dd T 2bbt d d- k
Nh -
(d- 16) 16)-1" cd-Ral'
1=0
1==
Because Ga-k,1",1 2 O and Cd-k.111 = (11)2, (d - k) 4 =
2k EE>== (d- "c- - 2 d
holds. Thus, we have (d - k) 2 d [1-aakk| -- -k-kk) for all 1 S [(1 a)k} ] + (k - k*) for
sufficiently small c, and by defining K2, K3 E maxn as
((1-a)k' (+(k
Iklldd 2
(K2)as =
Σ
d k
(11aa) J4k-
(henidd T T 1:d
(Kg)ss =
Σ
((d-k) k) - (1((11-)) -1(1-0))(1+
d k
1=
201 a) K"
(d-k) sletidd T 2 #441d
Σ
d k
t-
we have K, = cd- ((1-a)(1 J+k= K2 + K3. Moreover, K3 is positive semi-definite because Kg is
written as sum of polynomial kernels with positive coefficients. Thus, we can lower bound |3 K1 , by d' [(1-a) [o]]*][([ j-(k-k)) 323 (up to constant factor).
Proof of Lemma 4. The proof idea comes from Lemma B.9 of Wei et al. (2019). For a set S C [k], we let 2 - 1i=z 2i. and for set T S [d] \. [k], we let 2T = I 2. Expand (a (xT21) 2')0 as
d p
(2 12) =
Σ
2
) (xT(())
S.T
23

---

Published as a conference paper at ICLR 2024
where 11/7 2 0 depends only on |S], T, and due to symmetry. Also, we let
Σ
ma o (a (#k+11d)
T
Note that TT, 2 1 2 O, because Co.1T1,p considers the case where z,(i [k:] is multiplied for even number of times.
As basic fact in Boolean analysis, we have Ez 2T 2T 0 unless S = S' and T T'. Therefore,
Σ
Σ
Belz B
1=1
[
= E,
Σ
3,
ms (199(1)))))
3: IGsuM
3,7
8,
Σ Ez
Σ
B.C ACssmm())))((()))
Σ
Bcss,, Tl,0= (T())(())))())
3,T
Σ
2 11sT
Σ
8,
1=1
Σ
Co.1TI.Doo1TT
Σ Bi
(25)
Where we used CISI.TT1.p
C1S), Fl.4 2 0. On the other hand,
n Σ B((2 1e
Σ (x 1a ) - Emm.mmm
2
Esp
Σ
Bg
2
(26)
Because (/S), 2 Cr,p and 91//TT1 2 Cr.gr comparing (25) and (26) yields
n E B,(eTTo TaypTaype)) Σ B,(x
(aaTTtt)) 24j0
E;
Eig-ii Σ
n
] 20.
E
Σ
which concludes the proof.
Proof of Lemma 5. LHS is determined by how many coordinates are different between 2 and 2?, which is captured by 2t 23, Thus, LHS is the polynomial of 2 24 Moreover, its degree is at most I because the degrees of 2 and & 2d are at most in LHS. Thus the LHS can be written as Ekeop Ca.,r ( TT yr Note that, when is even, LHS is invariant the replacement z + -2,, and therefore Ca = 0 for odd l'. On the other hand, when is odd, Ca.,I O for even 1'.
Let us evaluate ca..r. By multiplying d' on both sides, we have
2
2)
Eg
CALI
2T ), LHS will converge into
By taking d -> 2o (while fixing the angle
9 x²
2) Vd
(27)
g
Vd
24

---

Published as a conference paper at ICLR 2024
here 9 follows gd -d- (vd). Now consider the Hermite expansion of Ei= Ct, H.y If
even, C1P" 2 T 1 )m!
0 for even l' and Crp for odd l'. If odd, C,P 2
for odd 1' and Cg 0 for even I'. By using these Hermite coefficients, (27) is equal to
Σ
2
G
d
M=0
Note that, as a function of the angle
E [1,1, the convergence is uniform. Therefore, we get
d canu + cG, (c)
for all and P'. When c.,- = 0, Ca.,,'' 0 for all d as we saw above. When c. > 0, there exists d such that Car for all d' > d. Therefore, for sufficiently large d, we have ca, 2 Moreover, by direct calculation, Ca. (1!)².
After obtained Lemma we would like to bound d- d-1(1-a)k) B K2B. For this, we use the lower bound the smallest eigenvalue of K2.
Let ((a) (d = 1, be sequence of inner-product kernels with ((a)(z, h Consider the case when each K(a) is associated with the same kernel function h: [-1, + R, so that M4a) h holds for all 2, 2' E [-1.1]d The following Lemma requires h is degree-k polynomial and its coefficients are positive for all degrees. Note that K₂ satisfies these conditions.
Lemma (Misiakiewicz (2022)). Assume the following conditions hold:
(a) (0) 0 for K' .....
(b) h (AA) (0) > O.
(c) h(()) is k-times differentiable.
Now fix > arbitrarily, and assume that d >> 1 and dhe-oavlood for some (aa) with ad + cold + oc). Given n ii.dd sample )= from Pz, we construct kernel matrix K E Baxn as (K(a)ss = h0= 2 Then, the kernel matrix K(d) can be decomposed into two positive semi-definite kernel K k and Kg and the spectrum of K> is bounded by
E (6)2,
+0 (d a o0).
Proof. See Section 3.2 of Misiakiewicz (2022), where we take ń = k 5.
Therefore, for any fixed > 0 and d >> 1 and n S dl(1-a)k") |+ (k= all the assumptions are satisfied for K2 with 1 [[1 - k=) + (k- k²) (if we regard K2 kernel in Rd x Rd- Note that we can take a, (log d) so that and d' - aay/ngd > 6 Then, the smallest eigenvalue of Kyt-1 is lower bounded by ?2(1) with probability at least 0.99 over the randomly drawn sample, for sufficiently large d. This immediately implies that the smallest eigenvalue of K₂ bounded by ((1) with probability at least 0.99.
Now we finalize the proof of Theorem 3.
Proof of Theorem 3. According to Lemmas 3 and 6, for all choices of B, with probability least 0.99 over the randomly drawn sample, we have
Txt d
)
27
E,
Σ
3.g1
> d (1--))))-(
82
(28)
i=1
> |la-a)k-]+(k--1
Σ |.8al
(29)
25

---

Published as a conference paper at ICLR 2024
Σ
(30)
de(1-a)k-) k-)-6
Because 9i is degree-2k: polynomial, Bonami's Lemma (e.g., Theorem 9.21 of (O' OoDDnnell, 2014)) yields
27
Eg
Σ
Bigi
1)ab Es
Σ
3.g
d
(2k
As result, the Paley Paley-Zygming Zygmund inequality (see Theorem 9.4 of (O'Donnell, 2014)) yields
27
2 T
2'
(1 (1-t3) -22
P2
Σ
3.91
2
Σ
B,91
2
2
(31)
(2k: 1)48
1=1
for all 0 1
Combining (28) and (31), with probability 0.99 over the sample, we have
8.91
2 d.(1a))k" J+ - -- 5/2
13,1.
d
with probability S2(1) over the choice of By taking sufficiently large d, -
is
larger than 2.00-)) (c is a constant from Lemma 1). Thus, using Lemma 2, we get
1tc
P PEWP2
[[a(2)] 2 d(ll-a)ee k+
3 2 1 - 1ld.
k-)
1=1
Now we apply Lemma and finally obtain
Pe-Pa 2//d,
which concludes the proof.
D DETAILS OF THE EXPERIMENT
We describe the experiment settings for Figure/ We considered an anisotropic d-dimensional 3-sparse parity problem (Example 2): = 21 81 = 82 83 a/ /d,, and 84 = 1/ V.. Here controls the alignment of the distribution to the feature, the signal-to-noise ratio. We fixed the dimension d to 300, and varied n and a. We trained the neural network (2) with Specifically, employed the width N 2000 as finite neuron approximation, and initialized neurons so that each them followed the standard normal distribution (and thus the network was rotation invariant at the initialization). By using the logistic loss, updated the network by the discretized MLFD (6) by setting = 0.25, Å = 0.1, and À = 0.la² d (fixed during the training) by following Corollary 1, until 7 10000. We ran the experiment 5 times with different seeds and plotted the mean for each n and a.
15.
26


Reviews:
[{"review_id": "YxRqCQBHYK", "reviewer": "Reviewer_obgQ", "summary": "In this work, the authors study the statistical and computational complexity of mean-field Langevin dynamics (MFLD) with anisotropic input data. In particular, they show that both complexities can be improved when prominent directions of the anisotropic input data align with the support of the target function.", "questions": "If we have real data while $A$ is generally unavailable or can only be estimated from data, what can we say about this case with the results of this paper? \n\nI also wonder if the problem can be tackled with a preconditioned version of MFLD, instead of performing coordinate transforms on the input. \n\nTypo:\n- page 6, after Proposition 2: otherwise, we “still”\n- page 6, after Proposition 3: Under this “condition”", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is technically solid and study an important problem. MFLD is a recent framework to help us understand the behavior of two-layer nonlinear NN, particularly relevant to the isotropic $k$-parity problem. Extending the isotropic setting to the anisotropic setting is definitely interesting to study, and the authors have established learning guarantees for two-layer nonlinear NN for this anisotropic setting.", "weaknesses": "It appears that the whole work assumes that the matrix $A$ is known a priori or prespecified. To me, this assumption might be too strong and make the problem of more theoretical interest than very practically relevant, although it seems to me that the motivation of studying anisotropic data arises from the case of realistic datasets as mentioned in the abstract.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SHK052LKMn", "reviewer": "Reviewer_e822", "summary": "This work investigates mean-field Langevin dynamics on structured anisotropic input data with a k-sparse parity target function. It provides results for specific 2-layer neural networks around both statistical and computational complexity. In particular, it proves discrete-time and finite-width learning guarantees thus extending on the results of previous work. These results are verified empirically on synthetic data.", "questions": "* Given the very niche nature of this paper, I'm not sure that ICLR is the best venue for this work. Given that (I would imagine) this work would be of interest to quite a small subcommunity, would it not be more suitable to submit to a more specific venue rather than a generalist conference like ICLR? This is not to speak negatively about this work, but it seems that this topic is less generally relevant in its nature and may not be best suited to this large-scale style of conference. \n\n* Could the authors provide a definition of anisotropic input data? I don't think this was clearly defined. Does it exactly refer to the setting defined mathematically at the beginning of Sec 1.1?\n\n* Why can the existing results for the regression setting not be directly applied in the binary classification setting by converting it to a regression problem?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Disclaimer: I have not checked the validity of the proofs in the appendix and, therefore, cannot comment on the correctness of the results.\n\n* This work is generally well written.\n* The related work seems to be well addressed (although I am not familiar with the literature). \n* The contributions are quite clearly laid out in the introduction. \n* Section 2 (and 3) does a reasonable job of laying out the problem setting.", "weaknesses": "* I am not well-positioned to comment on the significance of the results within this research area. However, it is not clear to me why we should be interested in this formal setting beyond the fact that some previous works have considered it. The assumptions seem highly contrived and no attempt is made to link them to any practical task in a meaningful way. Could the authors explain why progress in this research direction is worth pursuing? I would argue that if this work is being submitted to a broad conference such as ICLR, more effort should be made to broaden its appeal by explaining why its setting is relevant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BLHR2S5kXZ", "reviewer": "Reviewer_Gjnf", "summary": "The authors address how gradient-based feature learning interacts with anistropic input data. They do this by studying the sparse parity problem where feature coordinates have varying magnitudes. To this end, they analyze the learning complexity of the mean-field Langevin dynamics, which describes the noisy gradient descent update on two-layer neural networks. They show one can use the anisotropy of the data to improve the statistical complexity (i.e. sample size) and computational complexity (i.e. network width) of the mean-field Langevin dynamics. This improvement is found when the main directions of the anisotropic input data aligns with the support of the target function. They also provide a method using coordinate transformations determined by the gradient covariance to show that the computational complexity can be improved by exploiting the anisotropy.", "questions": "- One page 7, under Corollary 1, the authors state: \n  > On the other hand, if the input covariance is anisotropic so that ... then the value of $R$ becomes dimension-free: $R = O(k^2 log(k)^2 )$.\"\n  \n  Can similar results apply to real-world datasets?\n\n- How applicable is the coordinate transformation used in the paper, to conventional neural network training for real-world datasets?\n\n- For equation (2), we're using a smooth activation function (tanh). Can the results apply for nonsmooth activation functions, namely ReLU networks? i.e. can the intuition \"anisotropy helps\" also apply for ReLU networks?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The results seem novel and interesting. Namely they show how one can utilize the anisotropy of the input data to improve learning for a specific setting.\n- The result of coordinate transformations which leverages anisotropy could be of interest for practical neural network training.\n- The implication of a tradeoff between statistical and computational complexity also seems interesting. \n- The paper is clearly written and the results clearly stated.", "weaknesses": "- In my opinion, the claim that \"anisotropy helps\" seems too strong, a bit specific to the problem in the paper, and not necessarily a general statement that can be made for neural networks based on the assumptions and results of the paper. It would be great if the authors could better motivate a connection between this problem and more general neural networks with real-world datasets.\n\n- The authors state on page 2, under \"Feature learning under structured data\", that \n\n  > in certain regression settings with low-dimensional target, structured data with a spiked covariance structure can improve the performance of both kernel methods and optimized NNs (Ghorbani et al., 2020; Ba et al., 2023; Mousavi-Hosseini et al., 2023; Suzuki et al., 2023b). However, these regression analyses do not directly translate to the binary classification setting which the k-parity problem belongs to.\n\n  This implies that the intuition \"anisotropy helps\" is perhaps already observed in the literature, which may weaken the impact of the paper. If the authors could expand on why the $k$-parity problem is worth studying separately, that would really help motivate things.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)