Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics
Authors: Alireza Mousavi-Hosseini
Denny Wu
Murat A Erdogdu
Venue: ICLR 2025 Poster (2025)
Keywords: mean-field Langevin dynamics
feature learning
multi-index models
neural networks
gradient descent
Primary area: learning theory
TL;DR: We provide guarantees for learning a multi-index model using a two-layer neural network trained with the mean-field Langevin dynamics.

Abstract: We study the problem of learning multi-index models in high-dimensions using a two-layer neural network trained with the mean-field Langevin algorithm. Under mild distributional assumptions on the data, we characterize the effective dimension $d_{\mathrm{eff}}$ that controls both sample and computational complexity by utilizing the adaptivity of neural networks to latent low-dimensional structures. When the data exhibit such a structure, $d_{\mathrm{eff}}$ can be significantly smaller than the ambient dimension. We prove that the sample complexity grows almost linearly with $d_{\mathrm{eff}}$, bypassing the limitations of the information and generative exponents that appeared in recent analyses of gradient-based feature learning. On the other hand, the computational complexity may inevitably grow exponentially with $d_{\mathrm{eff}}$ in the worst-case scenario. Motivated by improving computational complexity, we take the first steps towards polynomial time convergence of the mean-field Langevin algorithm by investigating a setting where the weights are constrained to be on a compact manifold with positive Ricci curvature, such as the hypersphere. There, we study assumptions under which polynomial time convergence is achievable, whereas similar assumptions in the Euclidean setting lead to exponential time complexity.

Published as a conference paper at ICLR 2025
LEARNING MULTI-INDEX MODELS WITH NEURAL NETWORKS VIA MEAN-FIELD LANGEVIN DYNAMICS
Alireza Mousavi-Hosseini* Denny Wr Murat A. Erdogdu
University of Toronto, Veector Insitute, "New York University, "Flatiron Institute {mousavi, ardogdu)(cs toronto, edu, dennywuenyu, edu
ABSTRACT
We study the problem of learning multi-index models in high-dimensions using a two-layer neural network trained with the mean-field Langevin algorithm. Under mild distributional assumptions on the data, we characterize the effective dimen- sion detf that controls both sample and computational complexity by utilizing the adaptivity of neural networks to latent low-dimensional structures, When the data exhibit such structure, deft be significantly smaller than the ambient dimension. We prove that the sample complexity grows almost linearly with deat bypassing the limitations of the information and generative exponents that appeared in recent analyses of gradient-based feature learning. On the other hand, the compu- tational complexity may inevitably grow exponentially with der in the worst-case scenario. Motivated by improving computational complexity, we take the first steps towards polynomial time convergence of the mean-field Langevin algorithm by investigating setting where the weights are constrained to be on a compact manifold with positive Ricci curvature, such as the hypersphere. There, we study assumptions under which polynomial time convergence is achievable, whereas similar assumptions in the Euclidean setting lead to exponential time complexity.
INTRODUCTION
A key characteristic of neural networks is their adaptability to the underlying statistical model. Several works have shown that shallow neural networks trained by (variants of) gradient descent can adapt to inherent structures in the learning problem, and learn functions of fow-dimensional projections with a sample complexity that depends on properties of the nonlinear link function such as the information exponent (Ben Arous et al., 2021) or generative exponent (Damian al., 2024) for single-index models, and the leap complexity (Abbe et al., 2023) for multi-index models. Specifically, prior works typically established a sample complexity of n ~ d(())) for gradient- based learning, where s can be the information/leap exponent (Abbe et al., 2022; Bietti et al., 2022; Damian et al., 2023; Ba et al., 2023; Mousavi-Hosseini et al., 2023b; Bietti et al., 2023; Dandi et al., 2023) or the generative exponent (Dandi al., 2024; Lee et al., 2024; Arnaboldi al., 2024; Joshi et al., 2024), depending on the implementation of gradient descent. This sample complexity is also predicted by the framework of statistical query lower bounds (Damian al., 2022; Abbe et al., 2023; Damian et al., 2024).
On the other hand, neural networks can efficiently approximate arbitrary multi-index models regard- less of the generative/leap exponent (Barron, 1993; E et al., 2022); moreover, if the (polynomial) optimization budget is not taken into consideration, there exist computationally inefficient training algorithms that can achieve sample complexity independent of s (Bach, 2017: Liu al., 2024; Damian et al., 2024). Intuitively speaking, function depending on à == Oa(1) directions of the input data has kd = O(d) parameters to be estimated, and hence the information theoretically optimal algorithm on isotropic data only requires d samples. However, thus far it has been relatively unclear whether standard first-order optimization algorithms for neural networks inherit this optimality.
A promising approach to close the sample complexity gap is to consider neural networks in the mean-fiell field regime (Nitanda & Suzuki, 2017: Chizat & Bach, 2018; Mei et al., 2018; Rotskoff & Vanden -ijnden, 2018; Sirignano & Spiliopoulos, 2020), where overparameterization is utilized to lift the gradient descent dynamics into the space of measures so that global convergence can be established. While most existing results in this regime focus on optimization instead of generalizzaiion/leaaablittt

---

Published as a conference paper at ICLR 2025
recent works have shown that under restrictive data and target assumptions (such XOR), mean-fiell neural networks achieve a sample complexity that does not depend on the leap complexity (Wei et al., 2019; Chizat & Bach, 2020; Telgarsky, 2023; Suzuki et al., 2023b). Among these works, Suzuki et al. (2023b) proved quantitative convergence guarantees for learning k-parity with n x d samples, despite the target function having leap index k. Key to this result is the convergence rate analysis of the mean-field Langevin algorithm (MFLA) (Hu et al., 2019; Nitanda et al., 2022; Chizat, 2022b). However, existing learnability guarantees in the mean-field regime fall short in the following aspects:
Learning general multi-index models. Prior works established optimal sample complexity for mean-field neural networks under stringent assumptions on the data distribution (isotropic Gaussian, hypercube, etc.) as well as on the target function such as single-index models with specific link functions (Berthier al., 2023; Mahankali al., 2023), or k-sparse parity classification (Wei et al., 2019; Telgarsky. 2023; Suzuki et al., 2023b). Hence, the problem of universally learning functions of low-dimentional projections with minimal data assumptions using neural networks with standard training procedure remains largely open.
Polynomial computational complexity. To achieve optimal sample complexity, the computational complexity of the training algorithm in Telgarsky (2023); Suzuki et al. (2023b) is exponential in the ambient (input) dimension. Although such exponential dependence may be unavoidable in the most general setting, sufficient conditions under which the mean-field algorithu can achieve statistical efficiency with polynomial compute relatively under-explored, with the exception of recent work that studied the specific example of k-parity on anisotropic data (Nitanda et al., 2024).
1.1 OUR CONTRIBUTIONS
Motivated by the above discussion, in this work we address two key questions. First, we ask
Can we train two-layer neural networks using the MFLA to learn arbitrary multi-inde index models with an (information theoretically) optimal sample complexity?
We answer this in the affirmative by showing that empirical risk minimization on standard variant of a two-layer neural network can be achieved by the MFLA. This result handles arbitrary multi- index models on subGaussian data with general covariance, hence enabling us to obtain sample complexity with optimal dimension dependence up to polylogarithmic factors with standard gradient- based training. However, such universal guarantee will inevitably suffer from exponential computational complexity; thus, the second fundamental question we aim to answer is
Are there conditions under which the computational complexity of the MFLA can be improved from exponential to (quasippolynomial dimension dependence?
We provide a positive answer in two problem settings. In the Euclidean setting, we show that the complexity of MFLA is governed by the effective dimension of the learning problem, instead of its ambient dimension; this implies an improved efficiency of MFLA when the data is anisotropic. In the Riemannian setting, we outline concrete conditions on the Ricci curvature of the compact manifold defining the weight space under which MFLA converges in polynomial time.
1.2 RELATED WORKS
Mean-field Langevin dynamics. The training dynamics of neural networks in the mean-field regime is described by nonlinear partial differential equation in the space of parameter distributions (Chizat & Bach, 2018; Mei et al., 2018; Rotskoff & Sanden-Eijnden 2018). Unlike the neural tangent kernel (NTK) description (Jacot et al., 2018; Chizat et al., 2019) that freezes the parameters around random initialization, the mean-field regime allows the parameters to travel and learn useful features, leading to improved statistical efficiency. While convergence analyses for mean-field neural networks are typically qualitative, in that they do not specify the speed of convergence discrepancy, the mield Langevin algorithm that we study is a noticeable exception, for which the quantitative convergence rate (Hu et al., 2019; Nitanda et al., 2022; Chizat, 2022b) and uniform-in- time propagation of chaos (Chen et al., 2022; Suzuki et al., 2022; 2023a; Kook et al., 2024; Nitanda, 2024: Chewi et al., 2024) have been established.
width
A recent work (Takakura & Suzuki, 2024) considered a two- too-timeccale MFLD where the second layer is optimized infinitely faster than the first layer, and provided statistical guarantees for learning
2

---

Published as a conference paper at ICLR 2025
Barron spaces with bounded activation function. The concurrent work of Wang al. (2024) studied this two-timesee vvo-iimeccale approach to MFLD in a more general setting of optimization over signed measures without considering the estimation aspect and statistical guarantees. Our formulation here bypasses the need for two-timescale dynamics while learning a similarly large class of target functions.
Learning low-dimensional targets. The benefit of feature learning has also been studied in a "narrow-width" setting, where parameters of the neural network align with the low-dimensional target function during gradient-based training. Examples of low-dimensional targets include single-index models (Ben Arous et al., 2021; Ba et al., 2022; Bietti et al., 2022; Mousavi-
et al., 2023a;
Damian et al., 2023; Lee al., 2024) and multi-index models (Damian et al., 2022; Abbe et al., 2022; 2023; Dandi et al., 2023: Bietti et al., 2023; Collins- Wooddfnn et al., 2023; Vural & Erdogdu, 2024). threshold for learning such functions is n 2 d (for isotropic data) While the information-theoret Barbier et al., 2019; Damian al., 2024), the complexity of gradient- (Mondelli & Montanari, 2018; based learning is governed by properties of the link function. For instance, in the single-index setting, prior works established a sufficient sample size of n 2 (e(s) where 5 is the information exponent for one-pass SGD on the squared/correlaaion loss (Dudeja & Hsu, 2018; Ben Arous et al., 2021; Bietti et al., 2022; Damian et al., 2023), and the generative exponent (Damian et al., 2024) when the algorithm can reuse samples access a different loss (Dandi et al., 2024; Lee et al., 2024; Arnaboldi achievable
et al., 2024; Joshi et al., 2024). This presents a gap between the Coormation-thhoreticaa sample complexity and the performance of neural networks optimized by gradient descent.
Notation. We denote the Euclidean inner product with (: the Euclidean norm for vectors and the operator norm for matrices with ||-I], and the Frobenius norm with I|-I|p- Given a topological space W endowed with an underlying metric and Lebesgue measure, we use P(W), P2(W). and (W) to denote the set of (Borel) probability measures, the set of probability measures with finite second moment, and the of absolutely continuous probability measures with finite second moment, respectively. Finally, we use wwo to denote the Dirac measure at wh, i.e. J (ww) (w) = h(w)).
2 PRELIMINARIES OPTIMIZATION IN MEASURE SPACE
Statistical model. In this paper, we consider the regression setting where the input a E Rd is generated from some distribution and the response R is given by the following multi-inde model
(u))_a) Vk
(ue xe vk
+
(2.1)
Here, g : R - IR is the unknown link function, & is zero-mean ç-subGaussian noise independent from x; for simplicity, we assume 5² S 1. Without loss of generality, we assume that the unknown directions u1....; Uy are orthonormal, and define U = (u//Vk, .... u /K)] E B&xd, thus, we can use the shorthand notation y = g(Ux) x) + E. Throughout the paper, we consider the setting < and treat & as an absolute constant independent from the ambient input dimension d.
For a student model x + (x: W) with W denoting its model parameters, we consider loss functions of the form ((ū,)) y) = p(ù y) where p RR + R+ is convex. In the classical regression setting where we observe n i.i.d. samples (xx(1) ,, ) from the data distribution, the regularized population risk and the regularized empirical risk are defined respectively as
(W) = B[((ū(=; W), + R(W) and (A(W) = 1
A
{(0( (1 W), y ()) + R(W),
where R is some regularizer on the model parameters and the expectation is over the joint distribution of (xc y). In practice, we minimize the empirical risk JA as the finite sample approximation of the population risk Jy. anticipating that both minimizers are close to each other.
We use a tworr neural network coupled with l regularization to learn the statistical model (2.1), where learning constitutes recovering both unknowns U and g. Denoting the m neurons with matrix W == (w)... Wm) the student model and the fz-regularizer are given as
1 W m
Im (x: W)
Σ (x; wx and R(W)= =
Σ
Wy
(2.2)
m
m
3

---

Published as a conference paper at ICLR 2025
where V/ R W
- R is the activation function, and w W with W denoting a Riemannian manifold. In this formulation, the second layer weights are all fixed to be +1.
To minimize an objective J denoting either JA or JA, we will consider a discretization of the following set of SDEs, which essentially define interacting particle system over neurons:
dw = - mm w J ( ) + -28--1
for 1 1 S m,
(2.3)
where V is the Riemannian gradient and (B is a set of independent Brownian motions on W. We scale the learning rate by m to compensate for the fact that the gradient will be of order m -1 with respect each neuron. The case 8 oo corresoonds to the classical gradient flow over J, while the Brownian noise can help escaping from spurious local minima and saddle points.
Optimization in measure space. Notice that the neural network and the regularizer in (2.2) are both
invariant under permutations the weights (w,,..., Wm ): thus, an equivalent integral representation of these functions can be written using Dirac measures dwz, centered wj. namely
Ú(:; pw) [/ ddww and R(ww) -/11 dpd with jW
(a::
1 m
m Ewy (2.4)
Here, pa is the empirical measure supported on m atoms. Indeed, j(:: pw) fm (z: W) and R(uw)= R(W), ). and this formulation allows extension infinite-width networks by removing the condition that measures are supported on m atoms, and by expanding the feasible set of measures to μ E P2)... Thus, we rewrite the population and the empirical risks in the space of measures as Ix(ww) == (W) and fx(w) == (W),
with domain / (2V)). Let J PPW) + R be the population risk Tx the empirical risk Fx. We can equivalently state the interacting SDE system (2.3) as (see e.g. (Chizat, 2022b, Prop. 2.4))
dw'- - v [u::: )dl + V for 1 1 S m,
(2.5)
where J'[l] L2W)) denotes the first variation of (p) defined via
Av E P2(W), which is unique up to additive constants when exists (Santambrogio, 2015, Definition 7.12).
[() ()) limeto
(2.6)
As co, the empirical measure (w: weakly converges deterministic measure μ4 for all fixed i, a phenomenon known as the propagation of chaos (Sznitman, 1991). Furthermore, de can be characterized the law of the solution of the following SDE and non-linear Fokker-Planck equation
dw = - Tn. )dl + / 2BB and 0:14: = V (1 VJ (m() + 3-'AAu, (2.7)
where V. and A are the Riemannian divergence and Laplacian operators, respectively. Due to the existence of mean-field interactions, (2.7) is known as the mean-field Langevin dynamics (MFLD).
For a pair of probability measures μ < v both P(W), we define the relative entropy H(p v) and the relative Fisher information I(u p/) respectively as
M()lv)= = ///0///000 In dp du
I(plz)= = // Vin
du dy
and
du.
(2.8)
dy
It is well-known at this point that μ4, in (2.7) can be interpreted as the Wasserstein gradient flow of the entropic regularized functional Fs(11) == I Tp)) 3 HH(1)) /T) where is the uniform measure on compact W or the Lebesgue measure on Euclidean space (Jordan et al., 1998; Ambrosio et al., 2005; Villani, 2009). For this gradient flow to converge exponentially fast towards the minimizer H4 arg min, (1)),, we require gradient domination condition on 2 in the space of probability measures, given as
H(P (2)) S CL.SI I(HIp3), Vμ E P(W), 2
(2.9)
which is referred to as the log- log-Sobolev inequality (LSI). If the measure di,, (oo(-BJ'[uu)dd satisfies LSI with constant CLST for all 2 O, μ4, enjoys the following exponential convergence
FB(m)) - FF((()) S e JeLsi (Fs(Ha) Fs(Lz));
(2.
see e.g. (Chizat, 2022b, Theorem 3.2) and (Nitanda al., 2022, Theorem 1).
4

---

Published as a conference paper at ICLR 2025
3 LEARNING MULTI-INDEX MODELS IN THE EUCLIDEAN SETTING
In this section, we consider learning multi-index models in the Euclidean setting. For technical reasons, we use approximation of ReLU denoted by + (x,,(2) for some K,,. 1, which is given by (z) x= In(1 + exp(xz)) for E (-00, and extended on (t/2, 0o) such that tx,, is C2 smooth, p S &/ < 1, and [ S k. Note that Φx., recovers ReL. as K. - co. Recall that we freeze the second-layer weights at +1. Consequently, non-negative activations can only learn non-negative functions. To alleviate this, we choose W = 222d+2 and use the notation w = (w) w with w1, w2 E IRd+1 to denote the first and the second half weight coordinates, and use the activation function
(2; w) = ((t,(()))))) ((E, W2)),
(3.1)
where ā = (x, E Rd+1 for constant iz corresponding to a bias unit. The above can also be seen as a 2-layer neural network with activation Φx,, and second-layer weights frozen +1.
We use the neural network and the regularizer in (2.2) with weights W == (wi...., Wm), and minimize the resulting empirical risk (W) via the mean-field Langevin algorithm (MFLA), which is a simple time discretization of (2.3) with the stepsize 17 and the number of iterations > 0.
w
1+1 = w ma wy (W) + 22 E) i5 S m,
(3.2)
where 5 are independent standard Gaussian random vectors. When the stepsize is sufficiently small, MFLA approximately tracks the system of continuous-time SIES (2.3) as well as their equivalent formulation in the measure space (2.5). If, in addition, the network width m is sufficiently large, prop- agation of chaos will kick in and the dynamics will be an approximation to MFLD (2.7), ultimately minimizing the corresponding entropic regularized objective FB,)(H) Fx(1) + (H((TT)
We make the following assumption on the input distribution. Assumption 1. The input x has zero mean and covariance Σ. Further, || and ||Uxe|| are subGaus- sian with respective norms o., M 1/21 || and o., |, | ||r for some absolute constants on Tu-
One example for the above assumption is the Gaussian case x ~ (O, Σ) where || is subGaus- sian with norm || All for any matrix A. In settings we consider, we can replace the operator norm with the Frobenius norm to obtain a weaker assumption, since || is roughly concentrated near its mean, scaling with ||>/// A lpp. Assumption 1 can cover much broader settings than Gaussianity, e.g. it is satisfied when æ = I z and z has mean-zero i.i.d. subGaussian coordinates (Vershynin, 2018, Theorem 6.3.2.). Without loss of generality, we will consider a scaling where ||| / 1.
A key quantity in our analysis is the effective dimension which governs the algorithmic guarantees. Definition 1 (Effective dimension). Define deat /2//2 where C9 tr( r. = - ||F.
The effective dimension dea can be significantly smaller than the ambient dimension d, leading to particularly favorable results in the following when deu polylog(d). This concept has numerous applications from learning theory to statistical estimation; see e.g. Vershynin (2018); Wainwright (2019); Ghorbani et al. (2020); Ba al. (2023). In covariance estimation, for example, the effective dimension is typically defined as tr(E) / Σ (e.g. (Wainwright, 2019, Example 6.4)), which is equivalent to deff in Definition 1 provided that U lives in the top eigenspace of Σ. However, in general, deff might be larger than learning setup where the variations of x provide very little information about target directions U, making estimation more difficult. We make the following assumption on the link function in (2.1).
which is expected as one can imagine a supervised
for for 21, Z2 E Assumption 2. The link function is locally Lipschitz: (g(21) g(22) I L||z1 z2ll
IR satisfying [|>1]] || V ||
< = r1 + o. 22/ + 1))n(n) for some q > 0 and L = (11rr)).
We also assume E 2 1
Note that the above Lipschitz condition is only local, thus allowing polynomially growing link functions g. We scale the Lipschitz constant with 1/r, to make sure has a variance of order (1).
Recall from Section 2 that in order to prove convergence of the MFLD (and its time/particle dis- cretization MFLA), it is sufficient for the Gibbs potential Upw, exp(-85% (wwJ) to satisfy LSI
5

---

Published as a conference paper at ICLR 2025
uniformly along its trajectory. Here, it is straightforward to derive the first variation as
[lll(w))= = ([]]w +
A + "|lool with (oll((w) Σ d'i(aa;;))) - ( (3.3)
The following assumption introduces the uniform LSI constant for the trajectory of MFLA.
Assumption 3. Ler W (wi...., w't denote the trajectory of MFLA. assume the measure Pww X exp(-333 (ww))) satisfies the LSI (2.9) with constant for all I 2 0, and CLSI 2 B.
The above condition is stated simplify the exposition and will be verified in our results by using the boundedness of Ōx,:; JA can be considered as a bounded perturbation of strongly convex potential, thus satisfies LSI by the Holley-Stroock argument (Holley & Stroock, 1986).
Proposition 2. Suppose p is Cp-Lipschitz. Then for any μ E (2)(2?++),, the probability measure VA & exp(-83%(x)) with J2 given by (3.3) satisfies the LSI (2.9) with constant
1
CLSI S exp(4Cp1.3).
(3.4)
BA
For the squared loss, we can replace Cp above with Fo(pw) )1/2. With this, as Fo(uw) is uniformly bounded along the trajectory, convergence of the infinite-width MFLD can be established. However, for the finite-width MFLA, controlling (o(w)) is challenging there is non-trivial probability that neurons incur large loss, which is why we require Lipschitz p. Note that the right hand side of (3.4) independent of k; thus, by letting - co. the proposition implies the same LSI constant for bounded variant of ReLU. However, for MFLA (the time discretization of MFLD), we additionally require smoothness of the activation.
3.1 STATISTICAL AND COMPUTATIONAL COMPLEXITY OF MFLA
The main result of this section is stated in the following theorem.
Theorem Under Assumptions 1, and 3, consider MFLA (3.2) with parameters 2 (dott/) and 1 S Ó (iaier(((+r where iz = ||>|| Fa. Suppose x,
on (1),
0 e ) the loss satisfies lp') " & 1, and the algorithm initialized with the weights sampled b
Li.d. from some distribution w Po with [1]]]] S 1. Then, with the number of samples n, the number of neurons m. and the number of iterations that can respectively be bounded by
n = Ö(da)) m = 5
d
))
1= 1 CLSI1
(3.5)
3A
17
with probability at least On for some > O, the excess risk satisfies
Ew Ex (m; W w' [p(E)] < on (1).
(3.6)
The above theorem demonstrates that (i) the effective dimension of Definition I controls the sample complexity. and (fi) the LSI constant of Assumption 3 controls the computational complexity. To that end, we can employ the LSI estimate of Proposition
arrive at the following corollary.
Corollary 4. In the setting of Theorem 3, using the LSI estimate of Proposition 2, with the number of samples, the number of neurons, and the number of iterations, respectively bounded by
(dett)
-((de 1 = (de
(3.7)
MFLA can achieve the excess risk bound (3.6) with X-1 A-1, x polylog(n).
We observe that the above corollary demonstrates a certain adaptivity to the effective low-dimensional structure, both terms of statistical and computational complexity. Remarkably, this property of MFLA emerges without explicitly encoding any information about the covariance structure in the algorithm. In contrast, consider "fixed-gni -gid"" methods for optimization over the space of measures P ((R22+2) (see Chizat (2022a) and references therein), in which the algorithm fixes the first-layer

---

Published as a conference paper at ICLR 2025
Work
Class of Targets
Sample Complexity
Input
Covariance
dor dos-adaptivity
Telgarsky (2023)
2-parity
hypercuhe
isotropic
x
Suzuki et al. (2023b)
k-parity
d
hypercube
isotropic
x
Nitanda et al. (2024)
k-parity
tr(2)) > |
uill
parallelotope
full-rank
/
Bach (2017)
multi-index
bounded
general
x
Theorem
multi-index
tr tr(>))||>>>
subGaussian
general
Table I:
Learning guarantees of neural networks with exponential compute (we state the dimension dependence). Our Theorem improves upon prior bounds, with potentially significant gap depending on the problem setup.
of two-layer network's 's representation and only trains the second solving convex problem similar to the random features regression (Rahimi & Recht, 2007). However, fixed-grit -rii methods do not show any type of adaptivity low-dimensions, and in particular their computational complexity always scales exponentially with the ambient dimension d, unless information about the covariance structure is explicitly used when specifying the fixed representation.
Table compares recent works in various aspects. Bach (2017) requires d sample complexity for learning general k-index models, which is worse than the complexity den of Theorem 3 even in the worst case deft = d. The improvement in our bound is due to a refined control over ||Ux/]] ||: while Bach (2017) assumes this quantity scales with vd, it can be verified that for centered æ. its expectation is independent of d. Further, Bach (2017) does not provide a quantitative analysis of the optimization complexity, and it is not clear if their algorithm is adaptive to the covariance structure. Nitanda et al. (2024) studied learning s-sprse parities, a subclass multi-index models we considered. for which it is considerably simpler to construct optimal neural networks with bounded activation. While the effective dimension (and the resulting sample complexity) of Nitanda et al. (2024) not explicitly acale-invariant, we derive scale-invarian translation of their bound in Appendix C, and show that it always lower bounded by our effective dimension, especially when E is nearly deficiein.nnt.
Remark. We make the following remarks on the complexity of learning multi-index models.
Even though the complexity in Corollary scales exponentially with dett, in Section 3.2 we outline problem settings where dest = polylog(d)) (d), under which it is possible achieve quasipolynomial runtime for the MFLA. That said, the exponential dependence in detf is unavoidable in general in LSI-based analysis (Menz & Schlichting, 2014), and consequently present in the mean-field literature (Chizat, 2022b; Suzuki et al., 2023a;b).
the isotropic setting Σ It, recent works have shown that certain variants of SGD can learn single-index polynomials with almost linear sample complexity (Dandi et al., 2024; Lee et al., 2024; Arnaboldi et al., 2024), which matches our sample complexity without needing exponential compute. However, these analyses crucially relied on the polynomial link function, which has generative exponent most 2 (Damian et al., 2024) and is SQ-Iearnable with Óa(d) samples (Mondelli & Montanari, 2018; Barbier et al., 2019; Chen & Mcka, 2020). In contrast, our assumption on the link function allows for arbitrarily large generative exponent, and hence the computational lower bound in Damian et al. (2024) implies that achieving learnability in the n > d scaling requires exponential compute for statistical query learners.
3.2 UTILIZING THE EFFECTIVE DIMENSION
To better demonstrate the impact of effective dimension den, we consider two covariance models.
Spiked covariance. We consider the spiked covariance model of Mousavi- et (2023b).
Namely, given spike direction 0 E Sd-1, suppose the covariance and the target directions satisfy
Σ = (1 (1+a) a1+ (I. + aooT), a x d² | x " 72 E [0, 1].
71 E [o, 1/2]. (3.8)
Note that in high digh-diiensional settings, 71 = 1/2 corresponds to a regime where 0 is sampled uniformly over 1d11, whereas 71 0 corresponds to the case where has a strong (perfect) correlation with U. We only consider 7/2 < 1 since 7/ > corresponds to a setting where the input is effectively one-dimensional, In this setting, effective dimension depends on ?1 and 72-
7

---

Published as a conference paper at ICLR 2025
500
ttaget aligmmer)
1.50
8 (sample complexity
1.00
300
0.75
0.55
2tnd
0.25
der (effective dimension)
(eigenvalue decay)
(a)
(b)
Figure 1: (a) deft according to Corollary 6. (b) Test loss from MFLA, details in Appendix E.
Corollary 5. Under the spiked covariance model (3.8), we have der x d' (1112-291)00)
To get improvements over the isotropic effective dimension d, either the spike magnitude the spike-targe target alignment ||@@| needs to be sufficiently large so that 7/2 > 2% Recall that the effective dimension in the covariance estimation problem is (r(E/ / | = d Therefore, detf = O. i.e. has a significant
in Corollary only matches its unsupervised counterpart when correlation with the target directions U. As 72 + and 71 the effective dimension will be smaller than polylog(d), leading to computational complexity that is quasipolynomial in d.
Scaling laws under power-law spectra. Next, we consider a more general power-law decay for the eigenspectrum. Specifically, suppose E ==Σ A, the spectral decomposition Σ, and
100 10
x for 1sisd d,
(3.9)
A1
for some absolute constants a, > O. Notice that Σi_ U0,1 = | = The following corollary characterizes dent in terms of the parameters and y. Corollary 6. Under the power-law eigenspectrum for the covariance matrix (3.9), we have
(2(22-0)))
<<,,7 <1 a a<1,7> 21, a a> 1
deft dl-a
(3.10)
d(1-7)vo
where = above hides polylog(d) dependencies.
The scaling of der (hence the sample complexity) is illustrated in Figure la. We remark that the power-law assumption (3.9) is parallel to the source condition and capacity condition in the nonparametric regression literature (Cucker Smale, 2002; Caponnetto capacity condition measures the decay of feature eigenvalues, and the source condition measures the alignment between the target and feature eigenvectors.
Vito, 2007), where
Also, based on (3.10), the width and number of iterations in Corollary 4 both become quasipolynomial in when a, 7 2 1. This corresponds to the setting where Σ is approximately low-rank with most of its eigenspectrum concentrated in the first few principal components, and the corresponding eigenvectors are aligned with the row space of U.
POLYNOMIAL TIME CONVERGENCE IN THE RIEMANNIAN SETTING
The strong statistical learning guarantees in the previous section come at computational price; MFLA may need exp(dan) many iterations and neurons to converge. This complexity arises since in the worst case, the LSI constant that governs the convergence of MFLD will be exponential in the inverse temperature parameter 3 (Menz & Schlichting, 2014). In this section, provide
8

---

Published as a conference paper at ICLR 2025
first steps towards achieving polynomial-time complexity for MFLD. In particular, we show that we constrain the weight space to be a compact Riemannian manifold with a uniformly lower bounded Ricci curvature such as the hypersphere s we can establish uniform LSI constant with polynomial dimension dependence, while the same set of assumptions in the Euclidean setting results in exponential dimension dependence. Notice that due to the manifold constraint on the weights, we no longer require 6 fz-regularization, and simply consider the objective Fz(x) Fo(m) - F).
Let (W, g) be a (d - 1) )-dimensional compact Riemannian manifold with metric tensor g. We denote the Ricci curvature of W with Ricg. We recall the neural network (( μJ) f v ( w (ydp(w)) (w) where, in this case, we choose a C²-smooth activation (x;:) : W + IR defined on the manifold. We consider the following model example to demonstrate our results.
Example 7. W is the hypersphere gd-1 equipped with its canonical metric tensor, and the activation is V (æ; w) = (ww, (w,z)) for some smooth ø : RR + R. Suppose (0'), |ø"| S 1, and the distribution of æ satisfies the conditions of Assumption 1.
The following assumption plays an important role in the analysis. Assumption 4. (W. g) satisfies the curvaturedimension condition Ricg i for an absolute constant > O. Further, there exists some ji E PPW) such that Fo(p) S E and H(ū (r) S A for some constants E, Å, where is the uniform distribution on W.
For the unit sphere S we have Ricg > (d 2)g; thus, the curvature-dimension condition is satisfied for sufficiently large d. Moreover, if there exists some with To(u) S E (e.g. the minimizer of T)) for which H(( r) oc, one can construct j such that (ū) I ( and H((A [) < (d), by smoothing μ via convolution with box kernels (see (Chizat, 2022a, Theorem 4.1) and its proof). Therefore in the worst-case, we have Á ((d). However, under reasonable model assumption, we can verify Assumption 4 with A o(d), which is demonstrated in the below proposition.
Proposition 8. Let y = f (a;; )dp= for some n² P(W) such that du² el dr for f W IR. Then, Jo(x²) = 0 and H(u (T) S f(dy" - S osc(f) where osc(f) == sup f 1---inn inf f.
In the above result, the constants in Assumption 4 can be identified E = 0 and - osc(f) which is the oscillation of the log-density of μ=" Consequently, if the neurons in the teacher model are sufficiently present in all directions of the weight space, we get osc(/) extreme case μ¹ = 7 which implies f is constant. Interestingly, in the case of k-multi-index models, this condition implies that k grows with dimension, ruling out the case A O(). We include a natural example of target functions of interest in the form of Proposition 8 in Appendix D.
o(d); consider e.g. the
For MFLD to converge to minimizer of Jo. the parameter /9 needs to satisfy 33 > ((A to ensure the entropic regularization not the dominant term in the objective F3. In the Euclidean setting. this implies an LSI constant of order ex)()(A)), resulting in a computational complexity exp(O(A)) (A) as shown in Theorem 3. In what follows, we demonstrate via the Bakry- -Enery theory (Bakry & Émery, 1985) that in the Riemmanian setting, under uniform lower bound on the Ricci curvature, the LSI constant can be independent of A and d as long as we have A = o(d). We include a natural example of target functions of interest in the form of Proposition 8 in Appendix D.
Proposition 9. Suppose Assumption 4 holds and the loss p is Ci-Lipschitz. Then, for all μ E P(W) and B < ed/ C/,,C,K, the probability measure U O exp(-3356lv satisfies the LSI with constant
CLSI S (od - 3Cpk)
(4.1)
Es,, (K V V (a; and Ess, denotes the expectation under empir- ical data distribution over n samples.
where K = -
Remark. In the setting of Example with 2 ar | we have K S Σ with probability at least 1 - (n -) for some constant > O. Consequently, the LSI constant is independent of d.
We can now present the following global convergence guarantee to the minimizer of To for large d. Theorem 10. Suppose Assumption 4 holds, and let K be in Proposition 9, Let (Hx)22 denote the law of the MFLD. For any > 0, let B = A/e and d ' 2C, K A/ ge. Then, we have
A In
Fo(Ho)
To(uT) S 58+6,
whenever T
(4.2)
Egd
E
9

---

Published as a conference paper at ICLR 2025
Moreover, in the setting of Example 7 and for 1-Lipschitz loss function, we have d 2 A/ De and (1 + ( [ V (11//e"), then
A Egd
Fa(mn)
JoluT) SEtE,
whenever
T
In
(4.3)
with probability at least 1 - O(n 9) over the randomness of data, for some constant > 0.
The sample complexity is controlled by the maximum of Á and tr (r up to log factors. We remark that dependence on is not our main focus, and it may be possible to improve 1/E²* with a more refined analysis. Remarkably, the time complexity improves in high dimensions, thanks the effect of the Ricci curvature. While the above result is for the continuous time infinite-width MFLD, the uniform-in-time propagation of chaos for MFLD strongly suggests that the cost of time/width discretizations will be polynomial, see e.g. Suzuki et al. (2023a) for the Euclidean setting, and Li & Erdogdu (2023) for the dine-iicretization of the Langevin diffusion on the hypersphere under LSI.
To compare the setting of this section to that of Section 3, as explored in Appendix A, we remark that the Euclidean l and entropic regularizations can be combined into a single effective regularizer of the form 8 / [(s]] ([][)))), where = (0, (A3) 12442); therefore, in the Euclidean setting, plays the role т. Further the proof of Lemma 20, we show that in the Euclidean setting, Å x X3/r2 and & x c,/ VA3; thus, to learn with any non-trivial accuracy, we have A x Cp/
/ dett. As
discussed above, controlling the effect of entropic regularization necessitates 3 2 (()) Unlike its Riemannian counterpart, the Euclidean LSI estimate of Proposition 2 scales with exp(3), ultimately resulting in large computational gap between the settings under the same A. This leaves open the question of whether Á x deff can be achievable in the Riemannian setting for e-multi-index models with & = O(),, which is interesting direction for future exploration.
5 CONCLUSION
In this paper, we investigated the mean-field Langevin dynamics for learning multi-index models. We proved that the statistical and computational complexity of this problem can be characterized by an effective dimension which captures the low-dimensional structure in the input covariance, along with its correlation with the target directions. particular, the sample complexity scales almost linearly with the effective dimension, while without additional assumptions, the computational complexity may scale exponentially with this quantity. Through this effective dimension, we showed both statistical and computational adaptivity of the MFLD to low-dimensions when training neural networks, outperforming rotationally invariant kernels and statistical query learners in terms of statistical complexity, and fixed-grid convex optimization methods in terms of computational complexity. Further, we studied conditions under which achieving polynomial LSI in the inverse temperature, and subsequently polynomial-in-d runtime guarantee for the MFLD is possible. Specifically, we showed that under certain assumptions, which are verified for teacher models with diverse neurons, constraining the weights the hypersphere can lead to such polynomial dependence. In contrast, the same assumptions in the Fuclidean setting result in an LSI constant scaling exponentially with the inverse temperature.
Riemannian manifold with positive Ricci curvature such as
We conclude with some limitations of our work, along with future directions.
Further assumptions are required to go beyond the current exponential computational complexity of the MFLD. We leave the study of such conditions as important direction for future work.
While we focused on A O(1), the versatility of the MFLD analysis may allow us to let k grow with dimension as in Ghorbani et al. (2019); Martin et al. (2023); Oko et al. (2024), or g to exhibit a more complex hierarchical structure (Allen-Zhu & Li, 2020; Nichani et al., 2023). Learning these functions with the MFLD an interesting direction for future research.
Another important future direction is developing lower bounds for learning multi-indes models with bastienttbased methods, under more realistic assumptions (e.g., non-adversarial noise) than the statistical query setup. These lower bounds can highlight when exponential computation is inevitable for optimal sample complexity, and present rigorous information-computati tradeoffs.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors thank Lénaic Chizat, Mufan (Bill) Li, Fanghui Liu, and Taiji Suzuki for useful discussions. MAE was partially supported by the NSERC Grant [2019-06167], the CIFAR AI Chairs program, and the CIFAR Catalyst grant.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: necessary and nearly sufficient condition for sgd learning of sparse functions two-layer neural networks. In Conference on Learning Theory, 2022.
Emmanuel Abbe, Enric leap complexity and saddle-to-sades dynamics. arXiv preprint arXiv:2302.11055 2023.
and Theodor Misiakiewicz. Sgd learning on neural networks:
Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv e-prints, pp. arXiv-2001, 2020.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Luca Arnaooldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iu- vant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprins arXiv:2405. 15459, 2024.
Jimmy Ba. Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representa- tion. arXiv preprint wwxhv:2205.01445 2022.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dimensional structure: spiked random matrix perspective. Advances in Neural Information Processing Systems, 36, 2023.
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-66 2017.
Dominique Bakry and Michel Émery. Diffusions hyperoontractives. In Séminaire de Probabilités XIX 1983/84: Proceedings, pp. 177-206. Springer, 1985.
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116( (16(12):54515546/ 2019.
Andrew R Barron. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information theory, 39(3):930-94) 1993.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from ligensiimnsi inference. J. Mach. Learn. Res., 22:106-1, 2021.
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055, 2023.
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In Advances in Neural Information Processing Systems, 2022.
Alberto Bietti, Joan Bruna, and Loucas Pilien.-Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310. 19793, 2023.
Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematics, 7(3):311366, 2007.
11

---

Published as a conference paper at ICLR 2025
Fan Chen, Zhenjie Ren, and Songbo Wang. Uniform-in-tin propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050, 2022.
Sitan Chen and Raghu Meka. Learning polynomials few relevant dimensions. In Conference on Learning Theory, 2020.
Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. Advances in Neural Information Processing Systems, 33: 13363-13373, 2020.
Sinho Chewi, Atsushi Nitanda, and Matthew $ Zhang. Uniform-in-n log-sobolev inequality for the mean-field langevin dynamics with convex energy. arXiv preprint arXiv:2409.10440 2024.
Lénaic Chizat. Convergence rates of gradient methods for convex optimization in the space of measures. Open Journal of Mathematical Optimization, 3:1-19. 2022a.
Lénaïc Chizat. Mean-field langevin dynamies: Exponential convergence and annealing. arXiv preprint arXiv:2202.01000 2022b.
Lenaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. In Advances Neural Information Processing Systems, 2018.
Lénaic Chizat and Francis Bach. Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss. Conference on Learning Theory, 2020.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. In Advances in Neural Information Processing Systems, 2019.
Elizabeth -ooodin, Courtney Paquette, Elliot Paquette. and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. arXiv preprint ar arXiv:2308.0890 2023.
Felipe Cucker and Steve Smale. On the mathematical foundations of learning. Bulletin of the American mathematical society, 39(1):1-49) 2002.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2023.
Alex Damian, Loucas Pillaud- Vivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural Networks can Learn Representa- tions with Gradient Descent. In Conference on Learning Theory, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at time. arXiv preprint arXiv:2305. 18270, 2023.
Yatin Dandi, Emanuele Troiani, Luca Amnaooli, Luca Pesce, Lenka Zdeborová, and Florent Krzakala.
The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.032220 2024.
Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications on pure and applied mathematics, 3632):183-212, 1983.
Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pp. 1887-1930. PMLR, 2018.
Weinan E. Chao Ma. and Lei Wu. The barron space and the flow-induced function spaces for neural network models. Constructive Approximation,
406, 2022.
B. Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations Lazy Training of Two-layers Neural Networks. In Advances in Neural Information Processing Systems, 2019.
12

---

Published as a conference paper at ICLR 2025
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When Do Neural Networks Outperform Kernel Methods? In Advances in Neural Information Processing Systems, 2020.
Richard Holley and Daniel W Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of Statistical Physics, 46, 1986.
Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint arXiv: STf:::1905077000 2019.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernl:: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems, 2018.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Nirmit Joshi, Theodor Misiakiewicz, and Nathan Srebro. On the complexity of learning sparse functions with statistical and gradient queries. arXiv preprint arXiv:2407.05622 2024.
Yunbum Kook, Matthew S Zhang. Sinho Chewi, Murat A Erdogdu, and Mufan (Bill) Li. Sampling from the mean-field stationary distribution. arXiv preprint arXiv:2402.07355, 2024.
Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoret, limit. arXiv preprint arXiv:2406.01581, 2024.
Mufan (Bill) Li and Murat A Erdogdu. Riemannian langevin algorithm for solving semidefinite programs. Bernoulli, 29(4):3093-3113, 2023.
Fanghui Liu, Leello Dadi, and Volkan Cevher. Learning with norm constrained, over-paramettrized, two-layer neural networks. Journal of Machine Learning Research, 25(138):1 2024.
Arvind Mahankali. Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. Advances in Neural Information Processing Systems, 36, 2023.
Simon Martin, Francis Bach, and Giulio Biroli. On the impact of overparamettriization the training of a shallow neural network in high dimensions. arXiv preprint arXiv:2311.033794, 2023.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665-E777) 2018.
Georg Menz and André Schlichting. Poincaré and logarithmic Sobolev inequalities by decomposition of the energy landscape. The Annals of Probability, 42(5), 2014.
Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory, pp. 1445-1450, PMLR, 2018.
Alireza Mousavi-Hosseine Sejun Park, Manuela Girotti, Joannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In The Eleventh International Conference on Learning Representations, 2023a.
Alireza Mousavi-Hosseing Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient-based feature learning under structured data. Advances Neural Information Processing Systems, 36, 2023b.
Eshaan Nichani, Alex Damian, and Jason D Lec. Provable guarantees for nonlinear feature learning in three-layer neural networks. Advances in Neural Information Processing Systems, 36, 2023.
Atsushi Nitanda. Improved particle approximation error for mean field neural networks. arXiv preprint arXiv:2405.15767 2024.
Atsushi Nitanda and Taiji Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint ArXiv:/7/2.088 005488, 2017.
13

---

Published as a conference paper at ICLR 2025
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 9741-9757. PMLR, 2022.
Atsushi Nitanda, Kazusato Oko, Taiji Suzuki, and Denny Wu. Improved statistical and computational complexity of the mean-field langevin dynamics under structured data. In The Twelfth International Conference on Learning Representations, 2024.
Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: compu- tational hardness and efficient gradient-based training for ridge combinations. In Conference on Learning Theory. PMLR, 2024.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems, 2007.
Grant M Rotskoff and Eric Yanden-Eijnden Neural networks as Interacting Particle Systems: Asymptotic convexity of the Loss Landscape and Universal Scaling of the Approximation Error. arXiv preprint arXiv:1805.0099 00910, 2018.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-33):94,
2015.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal Applied Mathematics, 80(2):25--75, 2020.
Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Uniform-in-time propagation of chaos for the mean-field gradient langevin dynamics. In The Eleventh International Conference on Learning Representations, 2022.
Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence mean-field langevin dynamics: Time and space discretization. stochastic gradient, and variance reduction. In Thirty-seventh Conference on Neural Information Processing Systems, 2023a.
Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. Feature learning via mean-fiel field langevin dynamics: classifying sparse parities and beyond. In Thirty-sevnnth Conference Neural Information Processing Systems, 2023b.
Alain-Sol Sznitman. Topics in propagation of chaos. Lecture notes in mathematics, pp. 165-251, 1991.
Shokichi Takakura and Taiji Suzuki. Mean-field analysis on two-layer neural networks from a kernel perspective. arXiv preprint arXiv:2403.14917, 2024.
Matus Telgarsky. Feature selection and low test error in shallow low-rotation relu networks. In The Eleventh International Conference on Learning Representations, 2023.
Roman Vershynin. Righ-DDmensional Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2018.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Nuri Mert Vural and Murat A. Erdogdu. Pruning is optimal for learning sparse features in high- dimensions. arXiv preprint arXiv:2406.08658, 2024.
Martin Wainwright. Hiigh-Dieensional Statistics: A Non-Asymptoti Viewpoint. Cambridge University Press, 2019.
Guillaume Wang, Alireza Mousavi-Hossele and Lénaic Chizat. Mean-field langevin dynamics for signed measures via bilevel approach. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
14

---

Published as a conference paper at ICLR 2025
A PROOFS OF SECTION 3
Before presenting the layout of the proofs, we introduce a useful reformulation of the objective (B))(1). Recall that
Fo(u) = 2 R(u) 1 +
X H((m).
Faa(u)=
Let y x cxp ( A | can rewrite the above
)
be the centered Gaussian measure on Rdd+2 with variance 1/(AB). Then,
FB))(M) = To(u) + (((1l)) + d In
A
2/3
2m
As result, we can define
Fa,(k)= = Jo(m)) + =
(A.1)
which is non-negative and equivalent to F/ up an additive constant. Notice that
U == arg min FB, = arg min (p))
#
This reformulation, which was also used in Suzuki et al. (2023b), allows us to combine the effect of weight decay and entropic regularization into single non-negative term H(p [2). Furthermore, the simple density expression for the Gaussian measure allows us to achieve useful estimates for H(u 77). partieular, as we will show below, possible to control H(ug (7) with effective dimension rather than ambient dimension, which leads dependence on deft rather than d in our bounds.
We break down the proof of Theorem into three steps:
1. In Section A.2 we show that there exists measure μr (22 where ((:;²)) can approximate on the training set with bounds on R(x("). This construction provides upper bounds on (((y)) and H(G I7).
2. In Section A.3, given the bound on (ug (7), we perform generalization analysis via Rademacher complexity tools which leads bound on To (pz).
3. Finally, in Section A.4, we estimate the LSI constant and constants related to smooth- ness/discretization along the trajectory, which imply that F3, (n4")) converges to Fo (1))) where F3 is an adjusted objective over ((R (2T+(22))) defined (A.6). This bound implies the convergence of -- [Jo (W)] to Jo(u)), which was bounded in the previous step.
Before laying out these steps, in Section 1.1, we will introduce the required concentration results. In the following, we will use the unregularized population (o(p) = E[[(j(æ: x: μ), y)] and empirical To(u) = IEs,, [l(](x; ; μ), y)] risks, and also consider the width versions Jo (W) = Jo(uw) and (W) = Fo(ww). Additionally, we will use &o (z) = V 0 to denote the ReLU activation.
A.1 CONCENTRATION BOUNDS
We begin by specifying the definition of subGaussian and subexponential random variables in our setting.
Definition 11 (Wainwright (2019)). A random variable x is a-subGaussian if EE H[=) < x2-d/12 for all X E R. and is (v, a. subbexxxpoooti is o-suubauussian then
(f ( < x22v2/2 for all |k| S 1/a. If
P(x - E[x] >
exp
(A.2)
202
If x is (v, a) -ubexponentiall, then
2
P(x E 2 <
exp
min
(A.3)
2
15

---

Published as a conference paper at ICLR 2025
Moreover, for centered random variables, let Hy and I-,,, denote the subGaussian and subexponen- tial norm respectively (Vershynin, 2018, Definitions 2.5.6 and 2.7.5). Then is o-subGaussian if and only if x x E|z]|eg and (v, subexpxxpoott if and only if = | E(z]ler
Next, we bound several quantities that appear in various parts of our proofs. Lemma 12. Under Assumption for any q > 0 and all 1 S i S with probability at least n
U (0) - Tu (1 + ou 2(q + 1) Inn) = Pg.
(A.4)
Proof. By subGaussianity of |Ux] from Assumption and the subGaussian tail bound, with probability at least 1 - 4-1
U (i)
S E + ouro 2(g 1)
T + ouT 2 + Inn.
The statement of lemma follows from union bound over S i is n.
Lemma 13. Under Assumption we have Es,, ||0|| x 2 with probability least 1 cxp(-S(n)).
Proof. By the triangle inequality,
||ااا||
|=][|=0|| 1)) + |B[[|a]] l1 x on >²/²| 1p
+ tr(E)) s tr(2)
Recall o2 = tr( Σ). Furthermore, by (Vershynin, 2018, Lemma 2.7.6) we have
|| s
=1|==113, S c².
We arrive at similar result for the centered random variable || ||[+=||| ||= ca. We conclude the proof by the subexponential tail inequality,
||]
2 tet) s ())
Es
Lemma 14. Under Assumption we have Es,, [y]] < 1 with probability at least 1 n n-9
Proof. By the local Lipschitzness of g, on the event of Lemma 12, have
|y| 3g) 30 30(1/r)) Un 3g².
By a similar argument to Lemma we have
|x>|| |||
[[[[ + 2 2 || $ (1 of))
E||
vin
since E ||x = 12. As a result, by the subexponential tail bound,
Es
|ux|
E | S (1 + 22) S
with probability at least (xp)-S((n)). Similarly. {²] 1/ Es.
< 1s
< 2
therefore,
|s-] - E c S 1
with probability at least 1 - 1-eeppppppp- exn)- The statement of the lemma follows by a union bound.
Lemma 15. Under Assumption 1, for any > 0 and
2
(1 + 02 + ln In(dn"),
with probability at least - ((n we have |Es,, æx
||| Further, if g 2 1, then
Es,, [x::
1/2
< | 1/2
16

---

Published as a conference paper at ICLR 2025
Proof. First, note that by subGaussianity of | ||. for every fixed i, have with probability least
1 n
E|||=|| I on
| V2 1) Inn.
Since E|||=||| S Ca. via a union bound, with probability at least - 7
! Sc + oncs Ons. /2( 1) In =: Cr-
Define the clipped version of a via H. = x(1 A
[(]]
&s
Then, on the above event,
Es, ax 1 Ess, [ar
Moreover,
E
sup E (aes w)²1 ||
< sup E (ar, w)²} |w
|| æx T
]||-
Finally, by the covariance estimation bound of (Wainwright, 2019, Corollary 6.20) for centered subGaussian random vectors and the condition given in the statement of the lemma,
Es. x
||
x
with probability least 1 O( "). Consequently, have least 1 ((n
[Es,
[x
'Ill
| with probability
For the second part of the lemma, let E denote the event which the above holds. Then,
[Es, xx |] A Σ
Es, ax 11 1/21
1/27
+ E (EE) [Es, [xer T] 1/2
E ((E) Ess. xa
Σ 1/2 + ((F.O) 1/2
B Es.. (a::: []]]"* 1/2
E
Σ|
1/2 + On (ra-9/2)cc cc.
Suppose q 2 1. Then for 2 2/ E|| we have I.. || [xa:: T1
S | 1/2
which completes
the proof.
We summarize the above results into a single event.
Lemma 16. Suppose n 2 [
(1 + of 1) In(n)) In(dn"). There exists an event such that
P(E) 2 1 - O0-0OOnn))))) ?), and on E:
1. (vx((t) < iz for all sisn.
[|=]|]]] S cτ
2. Es,
3. [Es, xx ||| S ||
4. E ESS xx "1l""
5. Es, [v2] <
We recall the variational lower bound for the KI. divergence, which will be used at various stages of different proofs to relate certain expectations to the KL divergence.
Lemma 17 (Donsker-Varadhan Varadhan Variational Formula for KI. Divergence (Donsker & Varadhan, 1983)). Let a and be probability measures W. Then,
/ fdu In
(f
(f eddvr.
H(p v) sup
f-W-R
17

---

Published as a conference paper at ICLR 2025
Finally, state the following lemma which will be useful in estimating smoothness constants in the convergence analysis.
Lemma 18. Suppose (z, 0) I Idd are drawn from a probability distribution D. Then,
E [20]]
ED //D[=>]]|EB[e] T]I--
Proof. We have
|Ba[aa]]||
sup (w, Eo(22])) |ells1
sup |el|
ED ED(=(U, x)]
sup |
Ed[=2] Ep (0, 20)²
(Cauchy-Schwarde
[ED[22]
sup (u, EE (aeer Jo) |el|s11
Ep [2 / (xx T||
Notice that the distribution D can be both the empirical as well as the population distribution.
A.2 APPROXIMATING THE TARGET FUNCTION
We begin by stating the following approximation lemma which is the result of (Bach, 2017, Proposi- tion 6) adapted to our setting.
Proposition 19. Suppose g IR R L- L-Lisschitz and lg(0)] O(LFs). On the event of Lemma 16, there exists measure (22222+) with R(H) / such that
max -
- C&LT
In
14
LF
Li
for all A 2 Ck. where Ck is a constant depending only on k. provided that the hyperparameter l satisfies CL LFF
2k/(k+1)
Proof. Throughout the proof, we will use C& to denote constant that only depends on k, whose
value may change across instantiations. Let z Ux I and ž (=T, F.) E IR k+1 Recall
that on the event of Lemma 12 we have [|(()] S F. and (0) S Lf, for all
S n. Let
denote the uniform probability measure on S* By (Bach, 2017, Proposition 6), for all A 2 Cc. there exists p E (T( with (| I such that
max g(=())
/ pl
((; v,
s C LF
LF
In fact, we have stronger guarantee on p. Specifically, p(v) given by
p(v)= Σ ())().))
where r (0,1), g* + R are introduced by (Bach, 2017, Appendix D), In particular, h(v)
-f(" 11:k
Uk+1;
uk+1
with the spheircal harmonics decomposition h(v) = Ejz0 (v)) is shown in (Bach, 2017, Appendix D.2) that Xy I (k and one can prove through spherical harmonics calculations (omitted here for brevity) that [b)(w)] & Ca sup (k-))((-1)22
1)/2 CaLF (6-1)2 As result,
|p())] S Σ 2, 3 (v)) - Σ 2 (1, CaLi, Σ < (1-r)
120
J21
18

---

Published as a conference paper at ICLR 2025
Using 1 1-r= (CCLL/CC///)))) 2/(e+1)
as in (Bach, 2017, Appendix D.4) yields
S = (
A
2k/(k+11)
Li
Define p+ (v) = :ppp) V 0
and (v) = (-p(v) V 0. Then, by positive 1-homogeneity of ReLU,
1 p(o) Fz
(u, dr(v)=
[
D+ (0)
(-(()) (v, 2) dr(1)) (o,2))d(())
1
/ p-(0)) Fre (v, x) dr(w)
(
(
(p+(1)
(0, x) dd(1)) f do (P-(v) ( () dr(w))
d00
gh
Fx
get
a (11,)) - /
Jatt: Do0(11, {))d(1((1))
| Ratt do0 ((w,w)))
Sean Lo )))dp)(
where Ar == (1+(()) Fo #T and u = (p-()) #T are the corresponding pushforward measures, 14 Fo #o and μr2 = Tu where (v) = (U Us UK+1) Rd+1 for a (o] uk+1) E R++1 In other words, w ~ fn is generated by sampling u ~ βu and letting w = (U Tk ok+1) T with similar procedure for ~ p2. Furthermore,
1 |w (dd)dd)( + feu. (pz[)) JRSE1 R(v)=
= f J (f)))
+
R++1
{ p(u)2 dr(u) <
F2
The last step is to replace do with Pk... Note that for all i, and almost surely over w Pl. 2k/(++1)
have (w. a ) S (v) S CkLia
with similar bound holding for w 12. As
Lr=
result, by choosing 2 LF
A
2k/(k+1)
we have 6xn, (we, (1) &xx ( w. (o) for
Lt=
all i and almost surely over w ~ μl and w ~ N2. By the triangle inequality, we have
(o(Ux()) - -0(0000) m) s {{ On.,
(w.,
2x -0xo (1w,, x0 ((u,)())) (w0)
{ Oxx (om((1,,, ((w, 220 (ww, -2o
(ww,))(0) 2)) docs
+
((Ux) - -f 0: (w. (dui(w) duz(w))
+
2
K + Ca
In
Li
Li
which completes the proof.
Next, we control the effect of entropic regularization on the minimum of Fo.A via the following lemma. Lemma 20. Suppose p Cp Lipschitz. For every n²= E 2(R22)+2), we have
+ \ R(x)) + 2/2C, p Es. [| Fa.(m) < Jo(w") = LE
min
VTA3
Proof. We will smooth 4t' by convoliving it with y, i.e. we consider μ = 1²" * x Let u ~ ? independent of ~ jr" and denote (u, u,, with u1, u2 IRd+1, We first bound Jolp" x7).
19

---

Published as a conference paper at ICLR 2025
Using the Lipschitzness of the loss and of Px,,, we have
Jolp" x)) Jo(ie) =Es, (e()
v( w)d(( * (ww) v) -(/ - w)
s ss
V(2; +)) - ((z
Bs.
(2: w u) - [ (n) (p (20)
< Ls
Ee Eullond + u1, x)) 2n, (w ||
/ Eullon, + u2, き)) - Φn,, (Ww x [d)
+ Cs
/
scp.ss.
(Buillku, R)ll + Eu, (2, d (w)
2 2/20p PE,. [|=]]]
Next, we bound the KI. divergence via its convexity in the first argument,
( 7(- w' w')dn" "dx" - / |( w') (xx))) (w)")
Hlp * 717) 76
Furthermore,
A3 (- 1| w0'l (||m||1) 11,01))>(<cw w') 2
A3 200
H(7(-- (o)(--w')110)
2
Consequently,
A3 Hld 771) S R(1²),
which finishes the proof.
Combining above results, we have the following statement.
Corollary 21. Suppose the event of Lemma 16 holds, p is Cp Lipschitz, and X S 1. Then,
min MEPAC (Rid==2) Fa,((H) - [p(E)] S Co
To TT
In
A=A
C2
A
Colc, +
Fa
F5
K6
F
VX8
for all 2 Ck. provided that i 2 ChA2k/(k+1)) (ra////t)))))ta( (ra/Fz) (k- 1)/(k+1),
Proof. We will use Lemma 20 with 1²¹ P(R22+22 constructed in Proposition 19. Then, for all 2 Ck.
Jolmr) = Es., (p((((;; - y)l
E.s., (p(j(<r; 11)) - (U( - E)]
S Es., [p(E) + C, ||; 11)) g(Ux)]]
< Es. [p(E)) C&C, Fz r=A)
r C, In
In
rx
Px
Fzz
Furthermore, Proposition 19 guarantees R(x() - 2 / Combining these bounds with Lemma 20 completes the proof.
A.3 GENERALIZATION ANALYSIS
Let
11 = arg min Fa.((1)
arg min Fg,((1). DEPao(R2414
HEP (EP21
20

---

Published as a conference paper at ICLR 2025
Corollary 21 gives upper bound on Tu(p)). In this section, we transfer the bound to Ju(p)) via Rademacher complexity analysis. Since Corollary 21 implies bound on Hip ()), we will control the following quantity,
Jo(p) - Fo(p). (G(p))
sup H= M
To be able to provide guarantees with high probability, we will prove uniform convergence over truncated version of the risk instead, given by
(T) To (1):
sup WHHllT)EA
where
(n) = (lpx((((;; v)), (n) 1ss (px((((; -)],
and px(() pl-) x. We will later specify the choice of x.
We are now ready to present the Rademacher complexity bound.
Lemma 22 (CChen et al., 2020, Lemma 5.5), (Suzuki al., 2023b, Lemma 1)). Suppose p is either
Lipshcchi loss or the squared error loss. Let & 22x for the squared ervor loss and Cp for the L++1
Lipschitz loss. Recall (0,
Then,
A3
To(m) - (5((m) S 400
2M n
E
sup
Proof. We repeat the proof here for the reader's convenience. Let (Es) denote i.i.d Rademacher random variables. Notice that for the squared error loss, Pse is /2x Lipschitz. Then, by a standard symmetrization argument and Talagrand's contraction lemma, we have
E
sup p:: H:H 7)
To(m) Fo(m)
2E
sup Lu I)S n
(( v)
20 IE
sup [N IISM
(0)
Next, we proceed to bound the Rademacher complexity. Specifically,
Eg
sup p:H(ul M
& V((x ;((=() wddd(ur)
F.
sup
a
M
1
5 (2x0) 20) dr((o)
Ee
In
exp
n
M
(. w)) dr(cos.
In
EE
exp
where the first inequality follows from the KL divergence lower bound of Lemma 17. Additionally, by sub-Gaussianity and independence of (E,) and Lipschitzness of Φx.., we have
Eg
exp
(Zx w))
< exp
2n i=1
a²
Σ
V (1) w)
i=1
22a212
S exp
n
Plugging this back into our original bound, obtain
M +
2012
Eg
sup
{.(
p=
21

---

Published as a conference paper at ICLR 2025
Choosing = /
Ma 2
we obtain
1 Σ
2M
Bg
sup jiHl
S pL)
- 221
S 7
which completes the proof.
We can convert the above bound in expectation to a high-probabilityy bound as follows. Lemma 23. In the setting of Lemma 22, for any > 0, we have
M + x To -0t S vi n
In(1/6)
sup HEPhe 1)
with probability at least
Proof. As the truncated loss is bounded by the result is an immediate consequence of McDiarmid's inequality.
Next, we control the effect of truncation by bounding To(u) via Th" (p),,, which is achieved by the following lemma.
Lemma 24. Suppose H(H1 ()) & M. Then,
-2(()) - -o7 (N) = (000)s((((())))) (1+E (6- + -
Proof. Notice that since the loss is C,-Lipschitz and p(0) 0. we have p - -> S C, yl. Recall that we use L for the Lipschitz constant of g, and [0(;; μi1 S 2t. Then,
Jolu) (n) < (p((((())))) y) 2 x)p(ū(<; -
c)p((l((E;))) v) 2 x) 1/3 [(( c + 1b1 2 x/C,)' (1
Additionally, by local Lipschitzness of g.
P2(2 + lol 2 x/C,) - (( + lv/ [|]])] 2 iz))
- P(2r +
Ux|| 2 fx)
< P(2 + |g(0) 4 |g] +n
Furthermore, Let x/C, 2 4t + 2|)(0)] + 2Lr, + [[]]] and recall that L = 0(1rr)) Then, by subGaussian concentration bound, we have
+ | x/,2 se -9
(22 + lo(0) L
We conclude the proof by remarking that by our assumptions, ou and Cp are absolute constants.
Finally, we combine the steps above give an upper bound on To(uz), stated in the following lemma.
Lemma 25. Suppose X Xv2 and 3 dante3/ /r2 5ad
Let = ((A + for e, S G,
Suppose n 2 dan+22/122 and Z x- aia (fa/ra) Net
20+11) k+2 Then,
Jo(HG) - E[((()] SE, and 8 3 (m) S Blp(E)] + $ 1.
Proof. By Corollary 21 and a standard concentration bound on Es,, [p(E] with sufficiently large n to induce neglibile error in comaprison with the rest of the terms in the corollary, we have
Jo(m)) + 3 "HHG 17) Blp(E)] S
F
Ti
Æ+1
In
F
(cz + 1z)
1
F
VXB
K
22

---

Published as a conference paper at ICLR 2025
By choosing
A
rx
and assuming Oz 2 fz.
k
Fx
Tz
B H (E) +
In
Fi
VA,
Note that the above choice on A translates to a lower bound on & in Corollary 21, given by
By choosing À 2r2 and using the fact that A. Orx)) and B =
rs we have the simpification,
2
8- H) + ((A+)) + 1
1
and.
To(mg) B[p(S) & 0(1 + =
Note that (pz) & To(H)). Using the generalization bound of Lemma 23 with the choice of 6 4 for some constant q O, we have with probability 1 O(n 9),
3
In
-5 (Kg) E "V
dett
Inn
S V
x
(A.5)
nde²
Furthermore, by Lemma 24 we have
To(ug) Jo" (pz) S Le -M(xt)
Combining the above with (A.5) and choosing on x In, we have
delt ndg²
In²
Jo( - E VV
+
n
which holds with probability at least O(n over the randomness of Sn-
A.4 CONVERGENCE ANALYSIS
So far, our analysis has only proved properties of μğ. In this section, we relate these properties to s4t via propagation of chaos. In particular, Suzuki et al. (2023a) showed that for W ~ 1t",, j(æ; mi") converges to ((:;1²} in a suitable sense characterized shortly, as long as the objective over nl"" converges to (P)). Notice that μ²" is measure on ((R instead of 2))24+2). Thus, we need to adjust the definition of objective by defining the following
E m J.(W) + B(W) + m3 H(m²).
F (A m
(A.6)
We can use the same reformulation introduced earlier in (A.1) to define
fm (4 = Ew- [W(W]]
+ H(p 11 @m),
1
(A.7)
m,
which is equivalent to F3A up an additive constant. With these definitions, we can now control Ewg" [Jo(e]")] via To())) The following lemma is based on (Suzuki et al., 2023a, Lemma 4),
with more careful analysis to obtain sharper constants.
23

---

Published as a conference paper at ICLR 2025
Lemma 26. Let F2 == |E|| 1/2 V P.. and suppose p is Cp S 1-Lipschitz. Then,
F2w
H3 8m) + L2
Ew-ap (W( - Fo(1)))) -
(A.8)
m
In particular, combined with (Suzuki et al., 2023a, Lemma 3), the above implies
r2 P230LSI
62
Ew- [J([W] - 2 V
(F (FXX(()) Fo,((1)))
(A.9)
m
m
Proof. Notice that Ew-up [Jo(W)] = Ew [L [p(0((;; - (x;; ug) + ((:; ((a) - w)]l
s En ((((((;;)))) v)1 + C2 VEw Ex (w; (22;
S Toluz) + Cpy E. FRw (0(0; j(x: 13))8]
Suppose W (wI...., Wm) ju and W' = -ww',... w'm) M3 8m Let I' denote the optimal W2 coupling between W and W', and assume W. W r. Then.
Ew [((()) (²)=;""") Eww w [(((I: - -((:; ((z; ((x))))))
$ 2 wwE w [(0(); - 0(x: uw:) E. (((x;; j( 1b)))
Moreover, by Jensen's incquality,
1 m
Ew,w w [(((;; - ((:: (w:) )"]
v 178
Ew w [ (x; (xx; w())
2
Ew w (wi: w//1,,)) -3)²
2
E Ew w [w12 w/2 x)²
771
i=1
Hence,
2 2/1 (): ))
Eg W. [(((); - (w:
Ew w
W w w
m -
W2 (14 m
11 )m).
For the second term, notice that (<:; 143) Ew [0(; pw:)]] Ew; (x: w()] for all 5 < m.
By independence of (w)) and Jensen's inequality, we have
1
Ew [[[(2; pw:) ((:; (3))
Enw [(V wo) (<:;
m
27 [((M)):)0) v va: (d) (w))
Ew
m
2
m
Thus, the rest of this section deals with establishing convergence rates for F (n(") + (1)))))) To use the one-step decay of optimality gap provided by Suzuki et al. (2023a), we depend on the following assumption.
Assumption 5. Suppose there exist constants Cr, and R, such that
1. (Lipschitz gradients of the Gibbs potential) For all u. p¹ E P2(22d+2) and w, w' e B 2d
(Vllk)(w)) - Vbll'( s L(W2(1,)1))) u() + w|w0 - w w'll),
(A.10)
where W2 is the 2- Wasserstein distance.
24

---

Published as a conference paper at ICLR 2025
2. (Bounded gradients of the Gibbs potential) For all μ E P2 (2 and Rd++ have V [()
R. R.
3. (Bounded second variation) Denote the second variation of Iu(u) at w via To' ul w'), which is defined as the first variation of p. + (ww) (see (2.6) for the definition of first variation). Then, for all y P22 (2)(22242) and w. w' E RR2d+2
(w, w²) s T(1 + (2 + (-11oo1()))
(A.11)
We now state the one-step bound.
Theorem 27. (Suzuki et al., 2023a, Theorem 2) Suppose To satisfies Assumption 5. Assume À S 1, 3, L, R 2 1, and the initialization satisfies E | S R² for all S i S m. Then, for all 7 S 1/4,
F ( FBa((()) S exp 2,30Ls
(FTA(H")) - FF.((13)) + +mAmBad.. (A.12)
where
R² Am,B.An = c 1'(d+ ) + ()
R²
d
CL
+ 5)) (A.13)
)
mB
CLSI
X2
A3
CLSI
for some absolute constant C > 0.
We now focus on bounding the constants that appear in Assumption 5.
Lemma 28 (Lipschitzness of V.G). Suppose p is either the squared error loss or is Cp Lipschitz and has C₁ Lipschitz derivative. Assume 1. Notice that for the squared error loss, C 1. Then, for all p. v E (2 and w. w E 122d+2, we have
V Hbl((w)) Voll'((")) s kC, kkpp| Bs.
|| - w')l + C%
Es,
WS(C,, 11),
for the Lipschitz loss, and
v Vill((w) (ohlo( S v [o(n)]|Bs [x 112- 1 w w/|l+| Rs., n. T
W2(us,
for the squared error loss, where |ss.
0004
112-22 sup | Es. (ã, w) x. 1|
Proof. Recall that (w) = Es, Io' (i(x; u) - y) L (z; w)), where v( w) = @n,(((w1, x))
(,,2 ))). We start with the triangle inequality,
[Wll[( - (ww) < [Vlll((w) v Volul(wo)) / (ww
We now focus on the first term. For the Lipschitz loss,
V., (w))
Hill(())) Es, [e' (0(x;; pd) v) (1W1, #)) D'an(1u (1w0 .+))) ||
E. (6' ((w1, x)) - B'e, (ww) 2) 1/2
Es., 30 T
1/2
(wi) w/, x)) 1/2
1/2
CAnBs
Es. ~ T
Cpp Es.,
|ساا will=
25

---

Published as a conference paper at ICLR 2025
where the first inequality follows from Lemma 18, and the second inequality follows from the fact that 16%| I K. For the squared error loss, we have
Veni We, (w)) - Vw, Foll)(w' [ss, [((ax; p1) - () (1w, x3) 6n,, ((w0, , .;))=[]]
sup
Es, [(0(=; pم) - () (wW1 x)) - Hya (w/ x)) ) (v, ز)]
||
sup 1
Ess ((, - 3) 1s
- (wo', x))) (v, })}
To(m))
v. (v,Es, |Ww) - w/1: .x) n: 11)
sup 1
S xV To(m) Bs. (w) w¹1 -))² -τة 2
T
x (o(2) ||ss, | |ساا w''l ||.
Similar bounds apply to the gradient with respect to w2. which completes the bound on the first term of the triangle inequality.
We now consider the second term of the triangle inequality. Here we consider Lipschitz losses and the squared error loss at the same time since both have a Lipschitz derivative.
V Toll((w)) - VurTllul(a))
( (0(=;))) v) p'(((x; - y)) (w,,(((),)))
s Es. (o'(((z)) w)) - v) Pf (0(20; n²) - v)) 21 1/2
1/2
Es..
$ Css.s (z) - 0(x; / ()) 1/2 Bs. zw T
1/2
(A.14)
where the first inequality follows from Lemma 18. Let E P2 (2 x R22+2) be a coupling of μ and pl (i.e. the first and second marginals of y are equal to and μ² respectively). Recall that,
((x; ()) - Ú(a;; n²)= = / (0),,, ((W1; #)) th.,, ((w2, a)) dn., (W') m)) + dr., (WW2. .>) dr w²).
Therefore by the triangle inequality for the L2 norm Es. [() 1/2 and Jensen's inequality,
Es, [(((z; (4) ú( (()) <Es. / (oru((W), x)) - on (n (ww .))) 2 dy
1/2
I (on.((w2, x)) - Sn.c (wW)) x))) 2 dd
1/2
+ Es.
/ E.s.. (wi - w0 2/2 dy + / Ess. [w22 w/z, .t) 2)/g//2 dr
1/2 / ([]W) - will |l + llwa - will)dh(w,, W2)
Esn.
w w11 w²). / |/
<
2
By choosing 7 whose transport cost attains (or converges to) the optimal cost, we have
Es. [(0(x; p1) ((;;1()) < /2| Es. - T
(2(
Plugging the above result into (4. we have
Vw, F(w)[("")) V V Mi'dd) < V2C%
BBs
Rs.
j. T
We(u,
Notice that the same bound holds for gradients with respect to w2. Thus the bound of the second term in the triangle inequality and the proof is complete.
26

---

Published as a conference paper at ICLR 2025
Lemma 29 (Boundedness of V Vi).. In the same setting Lemma 28, for all μ E (2( and w E R22+2 we have
1/2
v (w)) v2C, Es.. T:
where C₂ = Cp when p is Lipschitz and Cp = v (T((H) when p is the squared error loss.
Proof. Notice that (onel 1. Therefore,
V (ol]ooll(w))) Es, [o'(0 ((:--))kaaa (1w) (,, )
y)
$ V Es. 1e' - -y))]
Es.. .
i 1/2
C Es.
where the first inequality follows from Lemma 18.
Lemma 30 (Boundedness of Ji'). In the same setting Lemma 28, for all μ P2( and w, w' E 2d2+2 we have
Jo" [1] (., w') s C2 it (1(001 (1/2++1101)))
where recall that C² for the squared error loss.
Proof. Via the definition given by (2.6), it is straightforward to show that
fo [u w') Es,. lo" (i(=;))) - y) Then, by the Cauchy-Schward inequality,
Jo (w w') S C, (2) (=);w")))" Moreover, by the Lipschitzness of On...
Es, (zx 1/2
< Es,, (w) 23
1/2 + Es, (wzz 2)"]1/²
1/2 ||w/|
I Es..
We similarly bound the expression for w' and arrive the statement of the lemma via Young's inequality,
[l][[ w') 22 Es.
C% ||0|||
|o'll
[s..
"z T
||| [11
||/0'113 112).
We collect the smoothness estimates and simplify them under the event of Lemma 16 the following Corollary.
Corollary 31. Suppose p and pl are Cp and C² Lipschitz respectively, with Cps C% S 1. Recall that
Σ = E [> On the event of Lemma 16, we have
Es,
Σ V F2, and consequently,
Tg satisfies Assumption 5 with constants L S K Σ V
F2),
R
ΣI 1/2 v V and CL =
Using the estimates above, we can present the following convergence bound F (((g)))) F FB. ΣΣ| i.. and for simpleity assume CLSI 2 B. For any $ 1, suppose
Proposition 32. Let F. the step size satisfies
S
27

---

Published as a conference paper at ICLR 2025
the width of the network satisfies,
52 1+
(5$
+
d ( +
xF
and the number of iterations satisfies
BCTstI
Fm (16) - -F
1 12
In
7
E
Then, we have (1 - (1)) S.
Proof. Throughout the proof, we will assume the event of Lemma 16 holds. Let Fe = Fo.ll1)). Notice that by iterating the bound of Theorem 27, we have
-ln
F(((+))) FB, I exp
FBAA) + 1- exp 230 (FF (x(())))
2,CLSI
S exp
(F (x)")))) FBA) + ABCLSIAAm,B..
2,CLiI
where the second inequality holds for 1 S 25CLSI since 1 1 2 x/2 for E [0, 1]. We now bound Am. so that the RHS of the above is less than O(e) by choosing a sufficiently large m and sufficiently small Recall that given constants L and R from Assumption 5,
Am,B.n x (d T
R (12 +
+
R
Cr ( 2)
d
mB
CTSTI
X
AB
CTST
From Corollary 31, LX x Vi
R x | 1/2 Viz, and C₁ = k To avoid notational clutter,
let 22 = | V 22. Then, to control the terms containing 7,, suffices to choose
A Cisin
3
for which we can simply choose
Cdstnerdd++21
Further, to control the term containing the number of particles m, need
+ + x
+
KF
KF
#
To drive the suboptimality bound below E, we also need to let the number of iterations satisfy
3CLsI In
F Fg (46)
FF
"
E
With the above conditions, we can guarantee
(HA(AT)) FR. Sc
which finishes the proof.
Further, we now present the proof of the LSI estimate given by Proposition 2.
Proof. [Proof of Proposition 2] Recall that
Kluww:((n)) = -illw: + ||
Thus we have Uywi (w) & (ww exp(-33illm Since 7 satisties the LSI with constant 1/(3X), by the Holley-Strooct perturbation argument (Holley & Stroock, 1986). upw satifies the LSI with constant
CL.I exp(8 osc ( (Fllw::))
13
28

---

Published as a conference paper at ICLR 2025
Additionally,
(llwwi)(w)) =
p (0(( -x)
- pCpl
which completes the proof.
Finally, we are ready to present the proof of Theorem.
A.5 PROOF OF THEOREM 3 AND COROLLARY
Recall that À = Ar and let | dutti,/12 and n 2 (dea (f. for some E S 1, where
=²X
det
غ (Ā722 -E- Then, long 2
from Lemma 25, we have To(uG) Elp(E)] x ă. Ar?
12
Note that while Lemma 25 only asks for A statement of Theorem 3 so that the choice of does not depend on k.
(fz/rx)
2(k-11) k42 we simplify this expression in the
On the other hand, given the step size 77, width m. and number of iterations by Proposition 32, have F (1,"" FBA (M)) I &. Therefore.
P28CLSIE 22 m
Ew ~n" [Jo(W)] - Jo(p3)
m
m
Additionally, from Lemma 25, have
8 8-HH(pg S +
+ 21.
Consequently, for m 2 t1 Crs) v L2 we have Ew [Jo(W)] ( -To((() S Therefore,
X=3
combining the bounds above, we have
Ew [Jo(W)] -Ep(()(())) 2 ((Att)) +
Consequently, Theorem 3.
take Ā (1), (1),), x-1 (1), which finishes the proof of
We finally remark that under the LSI estimate of Proposition 2 and the choice of hyperparameters in Theorem 3, the sufficient number neurons and iterations can be bounded by
P2 + Oldnt))) +
deft d
42 + @(ddu)))) - ((deet)
m S Ō
deft
r
da
c
and
1 o
0
r c (Oidal)) $ S(ddeddda))
which completes the proof of Corollary 4.
B PROOFS OF SECTION 4
We begin with the proof of Proposition
Proof. [Proof of Proposition 8] Note that 3o(2)) by definition. Moreover, the bound on H((" [T) is simple application of Jensen's 's inequality, namely,
H(u" (T) = / In Teldr dp- =
ef
/ Idu" In / d s S(dp"
-d)) dr).
Next, using the Bakry-Emery curvature- condition (Bakry Émery, 1985), we prove the following dimension-free LS bound.
29

---

Published as a conference paper at ICLR 2025
Proof. [Proof of Proposition 9] By the pevvature-dimnsiion condition (Bakry al., 2014, Section 5.7), the Gibbs measure Up exp(-3 3Tll)) satisfies the LSI with constant CLSI a- as long as Ricg + 2V To (w) ž ag,
for all w E W and some a > 0. By the bound on the Ricci curvature from Assumption 4, it suffices show
edg + 3V2 Tol) x ag.
Recall that
To(u) Es, | f V (( -w) -w] ]
Therefore,
[1l)(ww) = Es, (0(2; - (2; w)).
and
V.. (w) Ess. ((((;)))) y) V (x;
Consider the case where p is Cp Lipschitz. Then,
Amin V V Flll(ww)
= inf Es. ( (V((;; p)) - ( V [(x
1|
2 Cp sup Es, (0, V v(;
11l,,
Cpp..
Before stating the proof of Theorem 10, we adapt the generalization analysis of Appendix A.3 to the Riemannian setting of this section. Recall the truncated risk functions To (1)
E[p(ū(<: /)) y) A x] and fo" (u) Es. (p(j((;;))) v) A x]. Then, have the following uni- form convergence bound.
Lemma 33. Under the setting of Example 7, where recall [sp(0) S and [g'(()) S for all we have
TT) - (5)
M CCp
1/2
sup (EP 1 1T)
S
+
n
where E = n E x (6) (1) Combined with McDiarmid's inequality, the above bound implies
M n
- 1/2
In(1/8)
To (p) 2 Cop H(glT)MM
1 [++
sup
+ x
n
with probability at least
Proof. Based on the same argument as Lemma 22, for any > 0, we have
E sup
(M) (w)
$ 2C, E
sup #(N F)SM n
E (0 :n))
H(
where (Ei) are i.i.d. Rademacher random variables. Once again following Lemma 22, we have,
E sup H(K|T)SM
Σ cz0 <
M
1
E
In
exp
a²
V (x( (1) w02² dr(ww)
2n
im
Furthermore,
f exp
2
(0))
a 1111 1
Σ w, x0(1) )² dr(w))
Σ
1 (x ; w)2 dr(ww S exp
exp
2n² 1=1
n²
n²
(²))(
az 12?1:p' 11 1120
w. Ene) dr(w))
exp
exp
aa
(00)
||
1
exp
30

---

Published as a conference paper at ICLR 2025
Therefore,
1 Σ
M
(0))
+
1 2
Σ
gg
sup Hi(a +)SM n
(1)
+
a
a
==1
Using (+(()), 1 alo S
and optimizing over a yield
g
sup HH(ITTS
:4)
V
M
1
+
Σ
1)
1/2
Taking expectation with respect to the training set concludes the proof.
We also control the effect of truncating similar that of Lemma 24. Lemma 34. Under the setting of Example 7, for any μ E P(W), we have
Jolu) - - (M) 20 20(0)) + 2 2 12 E B
x
Proof. Similarly to the arguments in Lemma 24, by using the Cauchy-Schwartz and Markov inequali- ties, we have
Tolm) - (n) s ((1((((((;))) - y) 2 x)p(ū(z:) u))
S p(p(i((;))) - y) 2 E y) 1/2
s CCP((((E;))) y 2 x x/C,) 1/2 (Ú(2; - -) 1/2
(((0(x;1) -)
s c
x E [(i(=: [)7] + E V
< 20
Moreover, we have
E B I ((::))dp(w))
s 22p(0)² 21140111 [ (w. x) du((w)
E B[i(0;)))) s
2 21 11 |E|,
concluding the proof of the lemma.
Finally, we state the proof of the main theorem of this section.
Proof. [Proof of Theorem 10] Note that given guarantees that CLSI < 2/(d) along the trajectory. Consequently, by the convergence guarantee of (2.10), we have
and 2 2C,KK/(0E), Proposition
Fe(ur) & Fe(8)) + t EdT (Fo(Ho) -F((()))))))) < Fa( +
Further, by j of Assumption 4, we have
Fo(m)) s Fo 1
As result, Fo(P)) S F < E 2s, and similarly H(HT () S B(E + 2c).
Note that J0 (1T) Jo(uT). Using the fact that Cp, |E|| E of Lemma 33 and Lemma 34, with porbability of failure ā O(n have
w/ 2
X 1, and combining the bounds
4) for some constant q >
A(6 + + x
Inn n %
To(pr) To(HT)
+
n
31

---

Published as a conference paper at ICLR 2025
Optimizing over implies
B(E + +
In 1/4
Jo(uT) SEtET
+
71
n
Á(1 + E/E)
In
1/4
2 8
n
n
Choosing n according to the statement the theorem completes the proof.
C COMPARISONS WITH THE FORMULATION OF NITANDA ET AL. (2024)
Here, we provide a number of comparisons with results of Nitanda et al. (2024). In Section C.I, we show that the statistical model (2.1) is more general than their formulation, even for parity learning problems. In Section C.2, we provide an informal comparison of their effective dimension to our setting, exhibiting the improvement in our definition of effective dimension.
C.I GENERALITY OF THE FORMULATIONS
We begin by pointing out that the formulation of k-index model of (2.1) is strictly more general than that of Nitanda al. (2024), even for learning k-sparse parities. Recall that in their setting, they consider inputs of the type x = //2 for some positive definite E, where & ~ Unif((=1}() (their original formulation uses 2 ~ Unif([#1/v[]]]
but we rescale the input to be consistent with the notation of this paper). The labels are given by
sign [ (ū., 2)) 1=1
sign (II M 11211/2;;10 x3 )))
(C.1)
where (ūre are orthonormal vectors. Then, we can define an orthonormal set of vectors {u, D such that span(u1..... ..... ux) span(E 1/2
-
E 3-1/2 and define g such that
1/2 u1
1/2 ukk z
( (u),2) Vi
(ux, xt)
sign
1
V
V
for all z E [tt.ddd Therefore, the parity formulation of (C.1) can be seen as special case of the k-index model (2.1). Note that g is only defined on 2d points, and can extend all of R& such that g R + IR is Lipschitz continuous.
In contrast, the k-index model can represent parity problems that cannot be represented by (C.1). Starting from an orthonormal set of vectors {u:}e in Rd let
0
(ui, x²) Vk
(uk, 2r) Vi
) sign II
(u, xz) 2
(C.2)
Consider the case where ho = 2, then sign ((x)/2 u 2 1/2 uuu z)).
To be able to
reformulate this to (C.1), we need to be find orthonromal ún tz E Rdd such that
sign (<> (x (x ug
222 x))
sign (ūr z) [ūg z)), Vz (=1]d.
If has rank less than dl such that 1/2 / u2 then the above implies sign( (ū) , z) (ū2, 2)) >
0 for all z E {+1]d In particular, we must have some where sign( ((ū) z) (ūzz z)) > 0, which implies that
2d
(ū, 2:) (ūz 2i) U1
2,2 uz = 2d ( ū2) 0,
(C.3)
1=1
which is in contradiction with (ū), úنن = O. Therefore, for such E, we cannot formulate (C.2) as a special case of (C. This argument robust with respect small perturbations of Σ which make it
32

---

Published as a conference paper at ICLR 2025
full-rank. Specifically,, suppose Σ
Σ uz 1/2 21 8. Notice that can choose Σ u₁ such that (Ii/2u1,>) =) * for all {+1)4, e.g. by choosing E x es, i.e. the first standard basis vector. It straightforward construct full-rank E, ui, and u2 such that || || is arbitrarily small, in which case
sign (E (
sign (( u1
+ u1, ) (6.2) 20.
Following (C.3), once again the above would be in contradiction with (ū), i22) = O. This implies that the k-index model (2.1) is strictly more general than (C.1) when considering full-rank covariance matrices.
C.2 COMPARISON WITH THE EFFECTIVE DIMENSION OF NITANDA ET AL. (2024)
close inspection of the proofs in Nitanda et al. (2024) demonstrates that one can define their From the effective dimension in scale invariant manner as doff == (r(E) Σ Σ E-
previous section, observed that reduce their setting ours, need to choose set {u, ) of normalized vectors that spans the set of vectors { (E-1/2ū;) ( In particular, we can choose While {u, are not orthogonal, our proofs do not strictly rely the orthogonality assumption and it is only made for simplicity. Hence, we have
Σ 1/20,
equivalently write ü, E//2u, >
u, j u.
k
tr(E)
u
k Σ
Σ 1/2
u
1/2
Note that the above upper bound is sharp when 1, and is lower bounded by our definition of effective dimension stated in Definition 1. Therefore, we use the above bound in Table 1.
D EXAMPLES OF TARGET FUNCTIONS IN PROPOSITION 8
In this section, we provide a natural example of target function of the form given in Proposition A A'22
8. Suppose W = S and (x;; w) = (w, x)) Define μ² =T where T(v) some PSD matrix A to be defined later. Let
for
A
= ú( 1x)) = (x, Em--y- [ww [w)) = x, Enpr 2/2 0v
A/2/2
||A 1/2 1/2
ac
By defining B d E
A pm A1/2
|f we have y 11/2
Additionally, note that
( T) = x| Before proceeding further, we remark that the typical ||^=|| should be order vd to get (1) output, which is due to choosing the unit sphere the weight space.
Next, we construct A. Suppose distribution
such that with probability 1/2 we have x = e1.
the first standard basis vector. Let A diag( (1, 1,...,1) 1). Then,
y = ((e:;())) = E Bonr I+ 1)02
S cA1
where e is an absolute constant. On the other hand, ÿ(ei;7)
Choosing \ι 1. observe
that for sufficiently large d, with probability at least 1/2 the distance (;(æ;; ((x;)) | 2 C' for some absolute constant C' 0. This means that μ and are meaningfully different, and from an initialization of r one needs to train the network (e.g. with MFLD) to recover j²" It remains to find an estimate on H(p" (T)
Let I Rd + 1d-1 denote the normalization mapping, i.e. I(v) =
Tell Then, note that r I#N(O,, La), and μ (0,(0, A). Therefore, by the data processing inequality,
tr(A)
de) s In d
H(4 [T) H((N(O, A) (N((,, = 2 2
2
2
33

---

Published as a conference paper at ICLR 2025
where A = 2 Therefore, Á lnd o(d), and apply Theorem 10 to obtain polynomial convergence time, as desired.
E NUMERICAL SIMULATION
this section, we perform numerical simulations to verify the intuitions from Theorem 3. Specifically, train two-layer neural network with width m 50 and ReLU activation, where the first layer weights are initialized uniformly on the sphere, and fix the first half the second layer coordinates at +1/m, and the second half - -1/m. The input follows the distribution extra appended for bias), where Σ diag(o²) with of and of d for input dimension d 50. The labels generated by single-index model of the following form
N Σ) (with an
d-1
g((e), x)) = (el x)
V22
Therefore, the effective dimension from Definition is exactly equal to dett. We train the neural network using the squared loss with MFLA, with stepsize of 0.1, weight decay parameter 0.01, temperature 0.001.
Figure lb shows the test loss at the end of 200 it- erations of MFLA for different numbers of training samples n and effective dimension deft- For each value of n and deft. we average the test loss over independent runs with different realizations of data and initialization. In Figure we measure the gener-
20:200
Canioalationn Gip
505
0.4
alization gap, i.c. the average loss difference on the training set of n samples, and a test set of 100000 samples, the end of 3000 iterations of training with MFLA. For this experiment, we try = 100,
d.tt
200, and 500. As seen from both figures, detf controls the generalization gap and test loss, both of which decay with larger n.
Figure 2: Generalization gap measured by variy- ing the effective dimension.
The code to reproduce the experimental results is provided at: https:: /ttt:///gttttbb www./o/wwavi////o//oo MFLD Learnability.
34


Reviews:
[{"review_id": "xQhB9jZCSA", "reviewer": "Reviewer_hWwK", "summary": "This paper explores learning multi-index models in high dimensions using a two-layer neural network trained with the mean-field Langevin algorithm (MFLA). The authors identify an \"effective dimension\" that can significantly reduce sample complexity by adapting to low-dimensional structures in data. While worst-case computational complexity may still grow exponentially, they propose a solution to achieve polynomial-time convergence by constraining weights to a compact manifold under specific conditions.", "questions": "* The authors claim \"the Brownian noise can help escaping from spurious local minima and saddle points\" after presenting (2.3). Can you provide some reference of this statement? This is different from SGD which uses inexact oracle (stochastic gradient) leading to different randomness at each iteration. Adding Brownian motion is equivalent to adding randomness with the same magnitude at each iteration.\n\n* In the row between 230 and 231, should the weights be $W\\in\\mathbb{R}^{(2d+2)\\times m}$ instead of $W\\in\\mathbb{R}^{2d+2}$?\n\n* What is $\\mathbb{E}_{\\xi}[\\rho(\\xi)]$ in Thm 3?\n\n* Assumption 1 is restrictive to me. Could you provide some references which also consider this assumption on dataset? Or could you verify this assumption on some dataset? Is this assumption only for the proof? If so, why you need this assumption in your proof should be clearly stated in the main paper. It would be helpful to provide some proof sketch of Thm 3 discussing that.\n\n* Since you establish better sample complexity under subgaussian data assumption, it would be interesting to do more simulation even with artificial mean-zero and subgaussian data for empirical comparison with other results in Table 1.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The presentation is good, which is easy to follow.\n* Discussion on previous works is extensive and informative, and comparison with previous works is clear.\n* This paper uses MFLA to learn multi-index model and consider the effective dimension, which is novel to me.\n* This paper achieves optimal sample complexity, albeit leading to exponential dependence of $d_{eff}$ in $\\mathbb{R}^d$. But the authors well explain this phenomenon and improve this dependence in Riemannian manifold under some curvature conditions.", "weaknesses": "* Discussion on Assumption 1 is weak.\n* Some notations are not clearly stated.\n* This paper is short of some simulation.\n\nSee more details in Questions below.", "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": "2fQ2gmlbvs", "reviewer": "Reviewer_aq6U", "summary": "This paper examines the sample and computational complexity of learning general multi-index models in high-dimensional settings using two-layer neural networks and the mean-field Langevin algorithm. Under very mild assumptions, the authors successfully demonstrate that the sample complexity for such learning scales linearly with the effective input dimension, $d_{\\text{eff}}$, which is information-theoretic optimal. However, the computational complexity typically scales exponentially with $d_{\\text{eff}}$, which is also in general unavoidable. To address this, the authors explore specific conditions under which computational complexity might have a polynomial dependence on $d_{\\text{eff}}$. They identify certain settings—such as particular target functions and constraints on weights to lie on the hypersphere—that allow for the desired polynomial-time convergence.", "questions": "See the weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well-written, clear, and mathematically rigorous.\n\n2. Theoretical guarantees and insights into feature learning in deep learning are notoriously challenging, so any solid progress in this area is highly valuable. In particular, I find the paper’s second contribution quite interesting. Identifying cases where polynomial-time complexity (and thus efficient learning) can be achieved under the mean-field regime is of great importance, and this work is among the first to address this question, even if the current assumptions are somewhat restrictive.", "weaknesses": "I feel that the result in the first part is somewhat expected (though I recognize that proving it rigorously is still a quite valuable accomplishment. For the second part, I wonder if it might be possible to establish a polynomial-time convergence guarantee for a broader class of target functions (beyond the teacher neuron setup) by furthermore just considering the case where we restrict the weights just to the sphere. In general, I am convinced that a more detailed analysis/discussion would be valuable regarding the setups in which polynomial-time convergence with optimal sample complexity can be achieved, even when the input distribution is inherently high-dimensional, such as standard normal distribution with mean zero and covariance $\\Sigma = I_d$.", "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": "CfcI9F4o08", "reviewer": "Reviewer_bvhi", "summary": "In this work, the authors study the task of learning a general multi-index model using neural networks and the \nmean-field Langevin Algorithm (MFLA). Specifically, the authors introduce a notion of effective dimension in this context, \nwhich can potentially be much smaller than the ambient dimension when the input distribution is spiked the spikes \nalign with the relevant subspace of the target multi-index model. Then, the author derive bounds for the sample and \ncomputational complexity of MFLA based the effective dimension in both Euclidean and Riemannian settings.", "questions": "See the Weakness part of the review.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Overall, the paper is well-written, and the presentation is clear. \n* Proving polynomial convergence guarantees for MFLA is an important and challenging subject. This paper makes process in \n  this direction by proving a bound that is exponential in the effective dimension. They also show that for certain \n  models, the effective dimension can be as small as $O(\\mathrm{poly}\\log d)$, which leads to a quasi-polynomial \n  convergence guarantee.", "weaknesses": "* My main concern with this paper is that the effective dimension, as defined in Definition 1, is essentially about \n  how spiked the input distribution is (and whether it aligns with the target function) and does NOT capture the potential \n  low-dimensional structure of the target function. \n  In fact, if the input distribution is isotropic, the effective dimension (as defined in Definition 1) will be\n  $d$ even when the target function is a single-index model. \n  This limits the applicability of this paper's results, as it does not cover even the simplest Gaussian single-index case.\n* The comparison with the information/generative exponent (Section 3.1) is misleading. The authors suggest that one \n  advantage of the MFLA analysis over the information/generative exponent analysis is that the sample complexity is \n  always $\\tilde{O}(d)$ (in the isotropic case), while for SGD, (the bounds on) the sample complexity can be $\\tilde{\\omega}(d)$\n  when the information/generative exponent is larger than $2$. However, this is not true if we are allowed to have \n  exponentially many neurons, as assumed in the setting of Corollary 4. It has been shown in the information exponent paper [1]\n  that after the neuron attains a constant correlation with the ground-truth direction, the number of samples needed \n  in the final \"descent\" phase is $\\tilde{O}(d)$ regardless the information exponent. Meanwhile, if we have \n  exponentially many neurons, we can ensure that some neuron will have a constant correlation with the ground-truth \n  direction at initialization with at least constant probability. In addition, the number of steps needed in the \n  information exponent analysis is also $\\tilde{O}(d)$, which is better than the exponential bound given in Corollary 4. \n* In Sec. 4 (the Riemannian setting), the authors require the existence of a good distribution $\\mu$ with a small loss \n  and $H(\\mu | \\tau) = o(d)$ where $\\tau$ is the uniform distribution on the manifold. They justify the second condition \n  in Prop. 8 by showing that if the ground-truth distribution is given by $\\mu^* \\propto e^f \\mathrm{d} \\tau$, then we have \n  $H(\\mu^* | \\tau) \\le \\mathrm{osc}(f)$. However, this bound is vacuous if $\\mu^*$ is supported only on a few points as $f$ would then take $\\infty$ (or a large value if \n  we smooth $\\mu^*$). This excludes models such as the single-index model and $k$-multi-index model with $k$ being small.\n  To be fair, the authors do mention this limitation in the paper. I just think this assumption is too strong/unnatural. \n  It would be great if the authors could give a natural example where this condition holds. \n\n\n[1] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. 2021", "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": "o8LsA9Xfac", "reviewer": "Reviewer_PdGr", "summary": "This paper studies the sample and computational complexities of the mean field Langevin algorithm (MFLA). The analysis is centered aroung a newly introduced ``effective dimension'' which measures the complexity of the learning task (in terms of the data distribution and the teacher model). Under several regularity conditions and conditions on various constants, it is proven that the sample complexity and the computational complexity can both be controlled in terms of the effective dimension, which can be significantly smaller than the ambient dimension of the input space. The authors estimate (theoretically) this effective dimension on several examples. While the sample complexity bound is very satisfying and scales linearly in the effective dimension, the computational complexity may scale exponentially in the effective dimension and linearly in the ambient dimension. While this might be an issue in the general case, the authors show that this issue can be alleviate in a compact Riemannian setting. Improving the exponential dependence in the Euclidean setting remains open.", "questions": "- Can we derive a LSI like in Proposition 2 in a setting where there is no explicit regularization, but instead we make a dissipativity assumption on the model? Could it be of any interest to extend the analysis beyond regularization?\n\n - Could you explain further the meaning of the expectation $E_\\xi [\\rho(\\xi)]$ is Equation (3.6)?\n\n - Could part of the analysis be extended to more complex function classes than 2 layers neural networks?\n\n - Is it possible to construct explicit examples where the computational complexity becomes exponential in the effective dimension?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Both the statistical and the computational complexity are bounded in terms of a newly introduced effective dimension, which is a measure of the complexity of the learning task. It is interesting that the same quantity can be used for both notions of complexity.\n\n - The study of Mean-Filed Langevin Dynamics may provide important insights on the sample complexity of modern deep learning models.\n\n - The limitations of the paper are well-discussed.\n\n - The paper is well-written. Even though the subject is quite technical, good technical background and references are provided, making it sufficiently clear. The assumptions and theorems are clearly stated.", "weaknesses": "- The computational complexity may scale exponentially with the effective dimension (in the worst case). While the authors show that this might not be an issue in a compact Riemannian setting, the issues persists in the Euclidean setting. Additional discussion on how to go beyond this Riemanian setting could enhance the paper. In particular, what are (maybe intuitively) the conditions that could lead to non-exponential complexity?\n\n - As discussed in the paper, the analysis is restricted to a particular class of 2 layers neural networks. Moreover, the analysis is restricted to ``simple'' learning tasks, ie, the teacher models uses $k \\ll d$ neurons. It could be interesting to further discuss what the results become when $k$ Is large.", "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)