Title: Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond
Venue: NeurIPS 2023
Accepted: yes
Keywords: mean-field regime

feature learning

Neural network optimization

sparse parity function

classification

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

Feature learning via mean-field Langevin dynamics: classifying sparse parities and beyond

Abstract

Neural network in the mean-field regime is known to be capable of feature learning , unlike the kernel (NTK) counterpart. Recent works have shown that mean-field neural networks can be globally optimized by a noisy gradient descent update termed the mean-field Langevin dynamics (MFLD). However, all existing guarantees for MFLD only considered the optimization efficiency, and it is unclear if this algorithm leads to improved generalization performance and sample complexity due to the presence of feature learning. To fill this important gap, in this work we study the sample complexity of MFLD in learning a class of binary classification problems. Unlike existing margin bounds for neural networks, we avoid the typical norm control by utilizing the perspective that MFLD optimizes the distribution of parameters rather than the parameter itself; this leads to an improved analysis of the sample complexity and convergence rate. We apply our general framework to the learning of k -sparse parity functions, where we prove that unlike kernel methods, two-layer neural networks optimized by MFLD achieves a sample complexity where the degree k is 'decoupled' from the exponent in the dimension dependence.

1 Introduction

Mean-field Langevin dynamics. The optimization dynamics of two-layer neural networks in the mean-field regime can be described by a nonlinear partial differential equation of the distribution of parameters (Nitanda and Suzuki, 2017; Chizat and Bach, 2018; Mei et al., 2018; Rotskoff and VandenEijnden, 2018; Sirignano and Spiliopoulos, 2020). Such a description has multiple advantages: ( i ) global convergence guarantees can be obtained by exploiting convexity of the loss function, and ( ii ) the parameters are allowed to evolve away from initialization and learn informative features, in contrast to the neural tangent kernel ('lazy') regime (Jacot et al., 2018).

Among the gradient-based optimization algorithms for mean-field neural networks, the mean-field Langevin dynamics (MFLD) (Mei et al., 2018; Hu et al., 2019) is particularly attractive due to the recently established quantitative optimization guarantees. MFLD arises from a noisy gradient descent update on the parameters, where Gaussian noise is injected to the gradient to encourage 'exploration'. It has been shown that MFLD globally optimizes an entropy-regularized convex functional in the space of measures, and for the infinite-width and continuous-time dynamics, the convergence rate is exponential under suitable isoperimetric conditions (Nitanda et al., 2022; Chizat, 2022). Furthermore, uniform-in-time estimates of the particle discretization error have also been established (Suzuki et al., 2023a; Chen et al., 2022), meaning that optimization guarantees for the infinite-dimensional problem can be effectively translated to a finite-width neural network.

However, existing analyses of MFLD only considered the optimization of neural networks; this alone does not demonstrate the benefit of mean-field regime nor the presence of feature learning. Therefore,

an important problem is to characterize the generalization of the learned models, and prove efficient sample complexity guarantees. The goal of this work is to address the following question.

Can we show that neural network + simple noisy gradient descent (MFLD) efficiently learns an interesting class of functions, with a better rate of convergence compared to the 'lazy' regime?

Learning sparse parity functions. One particularly relevant learning task is the k -sparse parity problem, where the response y is given by the sign of the product of k coordinates of the input (on hypercube); as a special case, setting k = 2 recovers the classical XOR problem. When k ≪ d , this target function is low-dimensional, and hence we expect feature learning to be beneficial in that it can 'zoom in' to relevant subspace. In contrast, for kernel methods (including neural networks in the lazy regime) which cannot adapt to such structure, it has been shown that a sample complexity of n = Ω( d k ) is unavoidable (Ghorbani et al., 2019; Hsu; Abbe et al., 2022).

For the XOR case ( k = 2 ), recent works have shown that neural networks in the mean-field regime can achieve a sample complexity of n = O ( d/ϵ ) (Wei et al., 2019; Chizat and Bach, 2020; Telgarsky, 2023), which indeed improves upon the NTK complexity (Ji and Telgarsky, 2019). However, all these results directly assumed convergence of the dynamics ( t →∞ ) with no iteration complexity. Moreover, Wei et al. (2019); Chizat and Bach (2020) directly analyzed the infinite-width limit, and while Telgarsky (2023) provided a finite-width characterization, the dynamics is restricted to the low-rotation regime, and a very large number of particles N = O ( d d ) is required. Lastly, these analyses are specialized to XOR, and do not directly generalize to the k -parity setting.

1.1 Our Contributions

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

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

We apply our general framework to the k -sparse parity problem, and derived learning guarantees with improved rate of convergence and dimension dependence, as shown in Table 1. Specially, in the n ≍ d 2 regime we obtain exponentially converging classification error, whereas in the n ≍ d regime we achieve linear dimension dependence. Note that this improves upon the NTK analysis (which gives a sample complexity of n = Ω( d k ) ) in that it 'decouple' the degree k from the exponent in the dimension dependence. Our theoretical results are supported by empirical findings.

Table 1: Statistical and computational complexity (omitting poly-log terms) for the k -sparse (or 2-sparse) parity problem. Column ' k -parity' indicates applicability to the general k -sparse parity setting, and for references that does not handle k -parity we state the complexity for 2-parity (XOR). d is the input dimensionality, and n is the sample size. For SGD, n is the total sample size given by product of mini-batch size and iterations.

1.2 Additional Related Works

In addition to the mean-field analysis, parity-like functions can be learned via other feature learning procedures such as the one gradient step analysis in Daniely and Malach (2020); Ba et al. (2022);

Damian et al. (2022); Barak et al. (2022). Such analysis requires nontrivial gradient concentration at initialization, which translates to a sample complexity that scales as n = Θ( d k ) (Barak et al., 2022). For 'narrow' neural networks, Abbe et al. (2023) showed that a modified projected (online) gradient descent algorithm can learn the k -parity problem with a sample complexity of n = O ( d k -1 + ϵ -c ) (for some unknown c ). Also, Refinetti et al. (2021); Ben Arous et al. (2022) showed that a two-layer ReLU network with more than 4 neurons can learn the Gaussian XOR problem with gradient descent.

2 Problem Setting

Throughout this paper, we consider a classification problem given by the following model:

where Z = ( Z 1 , . . . , Z d ) is the input random variable on R d and 1 A is the indicator function corresponding to a measurable set A ∈ B ( R d ) , i.e., 1 A ( Z ) = 1 if Z ∈ A and 1 A ( Z ) = 0 if Z ̸∈ A . Let P Z be the distribution of Z . We are given input-output pairs D n = ( z i , y i ) n i =1 independently identically distributed from this model as training data. Then, we construct a binary classifier that predicts the label for the test input data as accurate as possible. To achieve this, we learn a two-layer neural network model in the mean-field regime via the mean-field Langevin dynamics .

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

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

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

Mean-field two-layer network. Given input z , let h x ( z ) be one neuron in a two-layer neural network with parameter x = ( x 1 , x 2 , x 3 ) ∈ R d +1+1 defined as

where ¯ R ∈ R is a hyper-parameter determining the scale of the network 2 . We place an extra tanh activation for the bias term x 3 ∈ R because the boundedness of h x is required in the convergence analysis. Let P be the set of probability measures on ( R ¯ d , B ( R ¯ d )) where ¯ d = d + 2 and B ( R ¯ d ) is the Borel σ -algebra on R ¯ d and P p be the subset of P such that its p -th moment is bounded: E µ [ ∥ X ∥ p ] < ∞ ( µ ∈ P ) . The mean-field neural network is defined as an integral over neurons h x ,

for µ ∈ P . To evaluate the performance of f µ , we define the empirical risk and the population risk as

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

Mean-field Langevin dynamics. We optimize the training objective via MFLD, which is given by the following stochastic differential equation:

where X 0 ∼ µ 0 , Law( X ) denotes the distribution of the random variable X and ( W t ) t ≥ 0 is the d -dimensional standard Brownian motion. Readers may refer to Theorem 3.3 of Huang et al. (2021) for the existence and uniqueness of the solution. Here, δF ( µ t ) δµ is the first variation of F .

Definition 1. For a functional G : P → R , the first-variation δG δµ ( µ ) at µ ∈ P is a continuous functional P × R d → R satisfying lim ϵ → 0 G ( ϵν +(1 -ϵ ) µ ) ϵ = ∫ δG δµ ( µ )( x )d( ν -µ ) for any ν ∈ P .

In our setting, we have δF ( µ ) δµ ( x ) = 1 n ∑ n i =1 ℓ ′ ( y i f µ ( z i )) y i h x ( z i ) + λ ( λ 1 ∥ x ∥ 2 ) . It is known that the Fokker-Planck equation of the SDE (1) is given by 3

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

where KL( ν, µ ) = ∫ log( µ/ν )d µ is the KL divergence between ν and µ , and ν is the Gaussian distribution with mean 0 and variance I/ (2 λ 1 ) , i.e., ν = N (0 , I/ (2 λ 1 )) .

For a practical algorithm, we need to consider a space- and time-discretized version of the MFLD, that is, we approximate the solution µ t by an empirical measure µ X = 1 N ∑ N i =1 δ X i corresponding to a set of finite particles X = ( X i ) N i =1 ⊂ R ¯ d . Let X τ = ( X i τ ) N i =1 ⊂ R ¯ d be N particles at the τ -th update ( τ ∈ { 0 , 1 , 2 , . . . } ), and define µ τ = µ X τ as a finite particle approximation of the population counterpart. Then, the discretized MFLD is defined as follows: X i 0 ∼ µ 0 , and X τ is updated as

where η > 0 is the step size, and ξ i τ ∼ i.i.d. N (0 , I ) . This is the Euler-Maruyama approximation of the MFLD with a discretized measure; we present the discretization error bounds in the next section.

3 Main Assumptions and Theoretical Tools

In this section, we introduce the basic assumptions and technical tools for our analysis.

Condition on the loss function To derive the convergence of the classification error, we assume that the loss function satisfies the following condition.

Assumption 1. The convex loss function ℓ : R → R ≥ 0 satisfies the following conditions:

ℓ is first order differentiable, its derivative is Lipschitz continuous and its derivative is bounded by 1: | ℓ ′ ( x ) -ℓ ′ ( x ′ ) | ≤ C | x -x ′ | and sup x | ℓ ′ ( x ) | ≤ 1 .

ℓ is monotonically decreasing, and is classification calibrated: ℓ ′ (0) < 0 (Bartlett et al., 2006).

ψ ( u ) -1 := ℓ (0) -( ℓ ( u ) -uℓ ′ ( u )) > 0 for any u > 0 .

This standard assumption is satisfied by several loss functions such as the logistic loss. We remark that the first assumption is used to show the well-definedness of the mean-field Langevin dynamics and derive its discretization error, and also to obtain a uniform generalization error bound through the classical contraction argument (Boucheron et al., 2013; Ledoux and Talagrand, 1991). The second and third assumptions are used to show the convergence of classification error of our estimator.

Logarithmic Sobolev inequality. Nitanda et al. (2022); Chizat (2022) showed that the convergence of MFLD crucially relies on properties of the proximal Gibbs distribution whose density is given by

where µ ∈ P . By the smoothness of the loss function (Assumption 1) and the tanh activation, we can show that the objective L has a unique solution µ ∗ which is also a proximal Gibbs measure of itself.

Proposition 1 (Proposition 2.5 of Hu et al. (2019)) . The functional L has a unique minimizer in P 2 that is absolutely continuous with respect to the Lebesgue measure. Moreover, µ ∗ ∈ P 2 is the optimal solution if and only if µ ∗ is absolutely continuous and its density function is given by p µ ∗ .

The next question is how fast the solution µ t converges to the optimal solution µ ∗ . As we will see, the convergence of MFLD heavily depends on a logarithmic Sobolev inequality (LSI) on p µ .

Definition 2 (Logarithmic Sobolev inequality) . Let µ be a probability measure on ( R d , B ( R d )) . µ satisfies the LSI with constant α > 0 if for any smooth function ϕ : R d → R with E µ [ ϕ 2 ] < ∞ ,

This is equivalent to the condition that the KL divergence from µ is bounded by the Fisher divergence: ∫ log(d ν/ d µ )d ν ≤ 2 α ∫ ∥∇ log(d ν/ d µ ) ∥ 2 d µ, for any ν ∈ P which is absolutely continuous with respect to µ . The LSI of proximal Gibbs measure can be established via standard perturbation criteria. For L ( µ ) with bounded first-variation, we may apply the classical Bakry-Emery and Holley-Stroock arguments (Bakry and Émery, 1985; Holley and Stroock, 1987) (Corollary 5.7.2 and 5.1.7 of Bakry et al. (2014)): If ∥ δL ( µ ) δµ ∥ ∞ ≤ B is satisfied for any µ ∈ P 2 , then µ ∗ and p X satisfy the LSI with

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

Proposition 2 (Entropy sandwich (Nitanda et al., 2022; Chizat, 2022)) . Suppose that µ 0 satisfies L ( µ 0 ) < ∞ and the proximal Gibbs measure p µ t corresponding to the solution µ t has the LSI constant α for all t ≥ 0 , then the solution µ t of MFLD satisfies

where µ ∗ = argmin µ ∈P L ( µ ) (the existence and uniqueness of µ ∗ is guaranteed by Proposition 1).

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

Convergence of the discretized algorithm. While Proposition 2 only established the convergence rate of the continuous dynamics, similar guarantee can be shown for the discretized setting. Let

where µ ( N ) is a distribution of N particles X = ( X i ) N i =1 ⊂ R ¯ d . Let µ ( N ) τ be the distribution of the particles X τ = ( X i τ ) N i =1 at the τ -th iteration. Suzuki et al. (2023b) showed that, if λαη ≤ 1 / 4 and η ≤ 1 / 4 , then for ¯ B 2 := E [ ∥ X i 0 ∥ 2 ] + 1 λλ 1 [( 1 4 + 1 λλ 1 ) ¯ R 2 + λd ] = O ( d + λ -2 ) and δ η := C 1 ¯ L 2 ( η 2 + λη ) , where ¯ L = 2 ¯ R + λλ 1 = O (1) and C 1 = 8( ¯ R 2 + λλ 1 ¯ B 2 + d ) = O ( d + λ -1 ) ,

where C λ is a constant depending on λ . In particular, for a given ˜ ϵ > 0 , the right hand side can be bounded by ˜ ϵ + 4 C λ λαN after T = O ( ¯ L 2 C 1 α ˜ ϵ + ¯ L √ C 1 √ λα ˜ ϵ ) 1 λα log(1 / ˜ ϵ ) iterations with the step size η = O (( ¯ L 2 C 1 α ˜ ϵ + ¯ L √ C 1 √ λα ˜ ϵ ) -1 ) . Furthermore, the convergence of the loss function can be connected to the convergence of the function value of the neural network as follows,

Here the second term in the right hand side can be bounded by 32 ¯ R 2 N [ 1 + 2 ( 2 ¯ R 2 ( λλ 1 ) 2 + ¯ d λ 1 )] via Lemma 2, if ∥ z ∥ ≤ 1 for any z ∈ supp( P Z ) as in the k -sparse parity problem. Hence, by taking the number of particles as N = ϵ -2 [( λα ) -2 +( λλ 1 ) -2 + d/λ 1 ] and letting ˜ ϵ = λαϵ 2 with the choice of T and η as described above, we have sup z ∈ supp( P Z ) | f µ X T ( z ) -f µ ∗ ( z ) | = O p ( ϵ ) .

Assumptions on the model specification. We restrict ourselves to the situation where a perfect classifier with margin c 0 is included in our model class, stated in the following assumption.

Assumption 2. There exists c 0 > 0 and R > 0 such that the following conditions are satisfied:

For some ¯ R , there exists µ ∗ ∈ P such that KL( ν, µ ∗ ) ≤ R and L ( µ ∗ ) ≤ ℓ (0) -c 0 .

For any λ < c 0 /R , the regularized expected risk minimizer µ [ λ ] := argmin L ( µ ) + λ KL( ν, µ ) satisfies Y f µ [ λ ] ( X ) ≥ c 0 almost surely.

Importantly, we can apply the same general analysis for different classification problems, as long as Assumption 2 is verified. The advantage of this generality is that we do not need to tailor our convergence proof for individual learning problems. Note that the convergence rate of MFLD is strongly affected by the values of ¯ R and R ; therefore, it is crucial to establish this condition using the smallest possible values of ¯ R and R , in order to obtain a tight bound of the classification error. As an illustrative example, we now show that the k -sparse parity estimation satisfies the above assumption.

Example: k -sparse parity estimation. In the k -sparse parity setting (Example 1), Assumption 2 is satisfied with constants specified in the following propositions. The proofs are given in Appendix A.

Proposition 3 ( k -sparse parity) . Under Assumption 1 and for ¯ R = k , there exists µ ∗ ∈ P such that

and L ( µ ∗ ) ≤ ℓ (0) -c 2 , where c 1 , c 2 > 0 are absolute constants.

Proposition 4. Under Assumption 1 and the settings of R and ¯ R given in Proposition 3, if λ < c 2 / (2 R ) , then µ [ λ ] satisfies

and f µ [ λ ] is a perfect classifier with margin c 2 , i.e., Y f µ [ λ ] ( X ) ≥ c 2 2 .

In other words, Assumption 2 is achieved with R = O ( k log( k ) 2 d ) , ¯ R = k and c 0 = c 2 / 2 . By substituting these values of R and ¯ R to our general results presented below, we can easily derive a bound for the classification error of the MFLD estimator.

4 Main Result: Annealing Procedure and Classification Error Bound

4.1 Annealing Procedure

The convergence rate of MFLD is heavily dependent on the LSI constant which may impose a large computational cost. We alleviate this dependency by employing a novel annealing scheme where we gradually decrease the regularization parameter λ . In particular, at the κ -th round, we run the MFLD until (near) convergence with a regularization parameter λ ( κ ) = 2 -κ λ (0) : µ 0 = µ ( κ -1) ,

which corresponds to minimizing L ( κ ) ( µ ) := L ( µ ) + λ ( κ ) KL( ν, µ ) . Then, we obtain a near optimal solution µ ( κ ) as L ( κ ) ( µ ( κ ) ) ≤ min µ ∈P L ( κ ) ( µ ) + ϵ ∗ for a given ϵ ∗ > 0 . We terminate the procedure after K rounds and obtain µ ( K ) as the output.

Suppose that there exists µ ∗ such that KL( ν, µ ∗ ) ≤ R and L ( µ ∗ ) ≤ δ ∗ . Then, as long as λ ( κ ) ≥ δ ∗ and ϵ ∗ < δ ∗ , we have that

Since L ( κ ) ( µ t ) is monotonically decreasing during the optimization, we always have L ( µ t ) ≤ L ( κ ) ( µ ( κ -1) ) ≤ L ( κ -1) ( µ ( κ -1) ) ≤ ( R + 2) λ ( κ -1) . Now we utilize the following structure on common classification losses to ensure that ∥ δL ( µ ) /δµ ∥ ∞ is small when the loss L ( µ ) is small.

Assumption 3. There exists c L > 0 , such that, for any µ ∈ P , it holds that ∥ δL ( µ ) δµ ∥ ∞ ≤ c L ¯ RL ( µ ) .

For example, the logistic loss satisfies this assumption:

Hence, the condition holds for c L = 1 because ∥ δL ( µ ) δµ ∥ ∞ ≤ 1 n ∑ n i =1 | ℓ ′ ( y i , f µ ( z i )) |∥ h z i ( · ) ∥ ∞ ≤ 1 n ∑ n i =1 ¯ Rℓ ( y i , f µ ( z i )) = ¯ RL ( µ ) . Under this assumption, the Holley-Strook argument yields that the log-Sobolev constant during the optimization can be bounded as

which is independent of λ ( κ ) and the final accuracy δ ∗ . Hence, after K = log 2 ( R/ ( λ (0) ϵ ∗ )) round, we achieve that L ( µ ( K ) ) ≤ δ ∗ + 2 ϵ ∗ where each round takes T κ = log(1 /ϵ ∗ ) / (2 αλ ) = O (log(1 /ϵ ∗ ) exp(8 c L ¯ R ( R +2)) / ( λ 1 λ ( k ) )) for the continuous time setting.

For the discrete setting, T κ = O (( C 1 αϵ ∗ + √ C 1 √ λ ( κ ) αϵ ∗ ) log(1 /ϵ ∗ ) exp(8 c L ¯ R ( R +2)) / ( λ 1 λ ( κ ) ) ) iterations in each round is sufficient (where C 1 is given for λ ( K ) ), where the step size is η = O ( C -1 1 αϵ ∗ ∧ C -1 / 2 1 √ λ ( κ ) αϵ ∗ ) . As long as ϵ ∗ = O ( λ ( κ ) /α ) for all κ , the total iteration number can be simplified as O ( log(1 /ϵ ∗ ) exp(16 c L ¯ R ( R +2)) C 1 / ( λ ( K ) ϵ ∗ ) ) , and the width N (number of particles) can be taken as N = O ([( λ ( K ) α ) -2 + ( λ ( K ) ) -2 + d ] / ( ϵ ∗ ) 2 ) = O ( [exp(16 c L ¯ R ( R +2)) / ( λ ( K ) ) 2 + d ] / ( ϵ ∗ ) 2 ) .

Remark 1. We make the following remarks on the annealing procedure.

The main advantage of this annealing approach is that the exponential factor induced by the LSI constant α is not dependent on the choice of the regularization parameter λ ( k ) (note that the LSI is solely determined by the intermediate solution µ t ). In contrast, the naive Holley-Stroock argument of Eq. (5) imposes exponential dependency on the regularization parameter.

Our annealed algorithm differs from the procedure considered in Chizat (2022, Section 4); importantly, we make use of the structure of classification loss functions to obtain a refined computational complexity analysis.

4.2 Generalization Error Analysis

We utilize the local Rademacher complexity (Mendelson, 2002; Bartlett et al., 2005; Koltchinskii, 2006; Giné and Koltchinskii, 2006) to obtain a faster generalization error rate. For the function class of mean-field neural networks, we introduce F := { f µ | µ ∈ P} , and the KL-constrained model class F M ( µ ◦ ) := { f µ | µ ∈ P , KL( µ ◦ , µ ) ≤ M } for µ ◦ ∈ P and M > 0 . The Rademacher complexity of a function class ˜ F is defined as

where ( z i ) n i =1 are i.i.d. observations from P Z and ( ε i ) n i =1 is an i.i.d. Rademacher sequence ( P ( ε i = 1) = P ( ε i = -1) = 1 / 2) ). We have the following bound on the Rademacher complexity of the function class F M ( µ ◦ ) .

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

The proof is given in Appendix B.2 in the supplementary material. Combining this local Rademacher complexity bound with the peeling device argument (van de Geer, 2000), we can roughly obtain the following estaimte (note that this is an informal derivation):

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

obtained by a vanilla Rademacher complexity evaluation because it only yields an O (1 / √ n ) bound. In other words, localization is essential to obtain our fast convergence rate.

We remark that a local Rademacher complexity technique is also utilized by Telgarsky (2023) to derive a O (1 /n ) rate. They adapted a technique developed for a smooth loss function by Srebro et al. (2010), which requires the training loss L (ˆ µ ) to be sufficiently small, that is, of order L (ˆ µ ) = O (1 /n ) . In our setting, to achieve such a small training loss, we need to take large ¯ R such as ¯ R = Ω(log( n )) . Unfortunately, such a large ¯ R induces exponentially small log-Sobolev constant like α ≲ exp( -cd log( n )) = n -cd . In contrast, our analysis focuses on the local Rademacher complexity around µ ∗ , and hence we do not require the training loss to be close to 0; instead, it suffices to have a training loss that is close to or smaller than that of µ ∗ .

4.2.1 Type I: Perfect Classification with Exponentially Decaying Error

In the regime of n = Ω(1 /λ ( K )2 ) , we can prove that the MFLD estimator attains a perfect classification with an exponentially converging probability by evaluating the term (I). From Eq. (7) we can establish KL( µ ∗ , ˆ µ ) ≤ O p (1 / ( nλ ( K )2 )) ; this KL divergence bound can be used to control the L ∞ -norm between f ˆ µ and f µ ∗ . Indeed, we can show that ∥ f ˆ µ -f µ ∗ ∥ 2 ∞ ≤ 2 ¯ R 2 KL( µ ∗ , ˆ µ ) (see the proof of Theorem 1). Then, under the margin assumption of f µ ∗ (Assumption 2), we have that f ˆ µ also yields a Bayes optimal classifier. More precisely, we have the following theorem.

Theorem 1. Suppose Assumptions 1 and 2 hold. Let M 0 = ( ϵ ∗ + 2( ¯ R + 1)) /λ ( K ) . Moreover, suppose that λ ( K ) < c 0 /R and

The proof is given in Appendix C.1 in the supplementary material. From Proposition 3 we see that the requirement n ≍ 1 /λ ( K )2 implies that in the n ≍ d 2 regime, Theorem 1 gives a perfect classification guarantee with a failure probability decaying exponentially fast.

4.2.2 Type II: Polynomial Order Classification Error

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

Theorem 2. Suppose Assumptions 1 and 2 hold. Let λ ( K ) < c 0 /R and M 0 = ( ϵ ∗ +2( ¯ R +1)) /λ ( K ) . Then, with probability 1 -exp( -t ) , the classification error of f µ ( K ) is bounded as

The proof is given in Appendix C.2 in the supplementary material. We notice that the right hand side scales with O (1 / ( nλ ( K ) )) , which is better than O (1 / ( nλ ( K )2 )) in Theorem 1; this implies that a sample size linear in the dimensionality is sufficient to achieve small classification error. The reason for such improvement in the λ K -dependence is that the stronger L ∞ -norm convergence is not used in the proof; instead, only the convergence of the loss is utilized. On the other hand, this analysis does not guarantee a perfect classification.

4.2.3 Computational Complexity of MFLD

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

The width N (the number of particles) can be taken as N = O (( ϵ ∗ λ ( K ) α ) -2 ) = O ( n 2 α -2 ) = O ( n 2 exp(16 c L ¯ R ( R +2)) ) .

Corollary 1 ( k -sparse parity setting) . In the k -sparse parity setting, we may take R = O ( k log( k ) 2 d ) , ¯ R = k and λ ( K ) = O (1 /R ) = O (1 / ( k log( k ) 2 d )) . Therefore, the classification error is bounded by

with probability 1 -δ . Moreover, if n = Ω( k 6 log( k ) 4 d 2 ) , then P ( Y f µ ( K ) > 0) = 1 with probability

As for the computational complexity, we require O ( k log( k ) 2 dn log( nd ) exp[ O ( k 2 log( k ) 2 d )]) iterations, and the number of particles is O (exp( O ( k 2 log( k ) 2 d )))) .

Comparison with prior results. In the 2-sparse parity setting, neural network in the kernel (NTK) regime achieves only O ( d 2 /n ) convergence in the classification error, as shown in Ji and Telgarsky (2019) and Telgarsky (2023, Theorem 2.3), whereas we demonstrate that mean-field neural network can improve the rate to O ( d/n ) via feature learning. Telgarsky (2023) also analyzed the learning of 2-sparse parity problem beyond the kernel regime, and showed that 2-layer ReLU neural network can achieve the best known classification error O ( d/n ) . However, their analysis considered a low-rotation dynamics and assumed convergence at t →∞ , whereas our framework also provides a concrete estimate of the computational complexity. Indeed, the number of iterations can be bounded as O ( dn log( nd ) exp[ O ( d )]) . In addition, while we still require exponential width N = O ( n 2 exp( O ( d ))) , such a condition is an improvement over N = O ( d d ) in Telgarsky (2023).

Barak et al. (2022) considered a learning method in which one-step gradient descent is performed for the purpose of feature learning, and then a network with randomly re-initialized bias units is used to fit the data. For the k -sparse parity problem, they derived a classification error bound of O ( d ( k +1) / 2 / √ n ) . In contrast, our analysis yields a much better statistical complexity of O ( k 2 log( k ) 2 d/n ∧ exp( -Ω( nk 6 log( k ) 4 /d 2 ))) , which 'decouples' the degree k in the exponent of the dimension dependence.

5 Numerical Experiment

We validate our theoretical results by numerical experiment on synthetic data. Specifically, we consider the classification of 2 -sparse parity with varying dimensionality d and sample size n .

Recall that the samples { ( z i , y i ) } n i =1 are independently generated so that z i follows the uniform distribution on {± 1 / √ d } d and y i = dζ i, 1 ζ i, 2 ∈ {± 1 } ( z i = ( ζ i, 1 , . . . , ζ i,d )) . A finite-width approximation of the mean-field neural network 1 N ∑ N j =1 h x j ( z ) is employed with the width N = 2 , 000 . For each neuron of the network, all parameters are initialized to independently follow the standard normal distribution (meaning that the network is rotation invariant at initialization) and the scaling parameter ¯ R is set to 15 . We set d to take values 5 , 10 , · · · , 150 , and n = 50 , 100 , · · · , 2000 . We trained the network using noisy gradient descent with η = 0 . 2 , λ 1 = 0 . 1 , and λ = 0 . 1 /d (fixed during the whole training) until T = 10 , 000 . The logistic loss is used for the training objective.

Figure 1 shows the average test accuracy over five trials. We make the following observations.

The red line corresponds to the sample size n = Θ( d 2 ) , above which we observe that almost-perfect classification is achieved. According to Section 4.2 (Type I), the classification error gets exponentially small with respect to n/d 2 , which predicts very small classification error above the line of n = Θ( d 2 ) , which matches our experimental result.

The boundary of test accuracy above 50% is almost linear, as indicated by the blue line. This matches the theoretical conditions (Type II) to obtain the polynomial order classification error in Section 4.2.

250

500

750

1000 1250 1500 1750 2000

sample size

n

25

50

75

100

125

150

dimension

d

n

d

n

d

2

0.5

0.6

0.7

0.8

0.9

1.0

test acc

6 Conclusion and Discussion

We provided a general framework to evaluate the classification error of a two-layer neural network trained by the mean-field Langevin dynamics. Thanks to the generality of our framework, an error bound for specific settings can be derived by directly specifying the parameters in Assumption 2 such as R , ¯ R , and c 0 . We also proposed an annealing procedure to alleviate the exponential dependencies in the LSI constant. As a special (but important) example, we investigated the k -sparse parity problem, for which we obtained more general and better sample complexity than existing works.

A limitation of our approach is that the required width and number of iterations are exponential with respect to the dimensionality d ; in contrast, certain tailored algorithms for learning low-dimensional target functions such as Abbe et al. (2023); Chen and Meka (2020) do not require such exponential computation. An important open problem is whether we can reduce the width and number of iterations to poly ( d ) for the vanilla noisy gradient descent algorithm. Another interesting direction is to investigate the interplay between structured data and the efficiency of feature learning, as done in Ghorbani et al. (2020); Refinetti et al. (2021); Ba et al. (2023); Mousavi-Hosseini et al. (2023).

References

Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

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

Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=akddwRG6EGi .

Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, and D. Wu. Learning in the presence of low-dimensional structure: a spiked random matrix perspective. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023) , 2023.

Bakry and M. Émery. Diffusions hypercontractives. In J. Azéma and M. Yor, editors, Séminaire de Probabilités XIX 1983/84 , pages 177-206, Berlin, Heidelberg, 1985. Springer Berlin Heidelberg. ISBN 978-3-540-39397-9.

Bakry, I. Gentil, M. Ledoux, et al. Analysis and geometry of Markov diffusion operators , volume 103. Springer, 2014.

Barak, B. L. Edelman, S. Goel, S. M. Kakade, eran malach, and C. Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=8XWP2ewX-im .

Bartlett, O. Bousquet, and S. Mendelson. Local Rademacher complexities. The Annals of Statistics , 33:1487-1537, 2005.

L. Bartlett, M. I. Jordan, and J. D. McAuliffe. Convexity, classification, and risk bounds. Journal of the American Statistical Association , 101(473):138-156, 2006.

Ben Arous, R. Gheissari, and A. Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems , 35:2534925362, 2022.

Boucheron, G. Lugosi, and P. Massart. Concentration Inequalities: A Nonasymptotic Theory of Independence . OUP Oxford, 2013. ISBN 9780199535255. URL https://books.google.co. jp/books?id=koNqWRluhP0C .

Chen, Z. Ren, and S. Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050 , 2022.

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

Chen, Y. Cao, Q. Gu, and T. Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 13363-13373. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/ file/9afe487de556e59e6db6c862adfe25a4-Paper.pdf .

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

Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31 , pages 3040-3050, 2018.

Chizat and F. Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. arXiv preprint arXiv:2002.04486 , 2020.

Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 5413-5452. PMLR, 02-05 Jul 2022. URL https://proceedings.mlr.press/v178/damian22a.html .

Daniely and E. Malach. Learning parities with neural networks. arXiv preprint arXiv:2002.07400 , 2020.

D. Donsker and S. S. Varadhan. Asymptotic evaluation of certain markov process expectations for large time. IV. Communications on pure and applied mathematics , 36(2):183-212, 1983.

Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. arXiv preprint arXiv:1904.12191 , 2019.

Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? arXiv preprint arXiv:2006.13409 , 2020.

Giné and V. Koltchinskii. Concentration inequalities and asymptotic results for ratio type empirical processes. The Annals of Probability , 34(3):1143-1216, 2006.

Holley and D. Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics , 46(5-6):1159-1194, 1987.

Hsu. Dimension lower bounds for linear approaches to function approximation.

Hu, Z. Ren, D. Siska, and L. Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint arXiv:1905.07769 , 2019.

Huang, P. Ren, and F.-Y. Wang. Distribution dependent stochastic differential equations. Frontiers of Mathematics in China , 16(2):257-301, 2021.

Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems 31 , pages 8580-8589, 2018.

Ji and M. Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint arXiv:1909.12292 , 2019.

Koltchinskii. Local Rademacher complexities and oracle inequalities in risk minimization. The Annals of Statistics , 34:2593-2656, 2006.

Ledoux and M. Talagrand. Probability in Banach Spaces. Isoperimetry and Processes . Springer, New York, 1991. MR1102015.

Mei, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

Mendelson. Improving the sample complexity using global data. IEEE Transactions on Information Theory , 48:1977-1991, 2002.

Mohri, A. Rostamizadeh, and A. Talwalkar. Foundations of Machine Learning . The MIT Press, 2012.

Mousavi-Hosseini, D. Wu, T. Suzuki, and M. A. Erdogdu. Gradient-based feature learning under structured data. In Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS 2023) , 2023.

Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv:1712.05438 , 2017.

Nitanda, D. Wu, and T. Suzuki. Convex analysis of the mean field langevin dynamics. In G. CampsValls, F. J. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pages 9741-9757. PMLR, 28-30 Mar 2022.

Refinetti, S. Goldt, F. Krzakala, and L. Zdeborova. Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 8936-8947. PMLR, 18-24 Jul 2021. URL https://proceedings.mlr.press/v139/refinetti21b.html .

M. Rotskoff and E. Vanden-Eijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.00915 , 2018.

Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Srebro, K. Sridharan, and A. Tewari. Smoothness, low noise and fast rates. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, Advances in Neural Information Processing Systems , volume 23. Curran Associates, Inc., 2010. URL https://proceedings.neurips.cc/paper_files/paper/2010/file/ 76cf99d3614e23eabab16fb27e944bf9-Paper.pdf .

Suzuki, A. Nitanda, and D. Wu. Uniform-in-time propagation of chaos for the mean field gradient langevin dynamics. In Submitted to The Eleventh International Conference on Learning Representations , 2023a. URL https://openreview.net/forum?id=_JScUk9TBUn .

Suzuki, D. Wu, and A. Nitanda. Convergence of mean-field langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction. arXiv preprint , 2023b.

Telgarsky. Feature selection and low test error in shallow low-rotation ReLU networks. In The Eleventh International Conference on Learning Representations , 2023. URL https:// openreview.net/forum?id=swEskiem99 .

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

Wei, J. D. Lee, Q. Liu, and T. Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems , pages 9712-9724, 2019.

Let

and ρ > 1 be a constant which will be adjusted later on. Then, for ξ 2 i := [log( ρk ) ξ ⊤ , -log( ρk )( b i -1) , ζ ] ⊤ ∈ R ¯ d and ξ 2 i +1 := -[log( ρk ) ξ ⊤ , -log( ρk )( b i +1) , ζ ] ⊤ ∈ R ¯ d for i = 0 , . . . , k , we define

Then, we can see that, for z ∈ {± 1 / √ d } d , it holds that

because we have ⟨ x 1 , z ⟩ + x 2 = log( ρk )( ⟨ ξ, z ⟩-( b i -1))+ ∑ d j =1 u j z j + u d +1 for x ∼ N ([ ξ ⊤ , ( b i -1)] ⊤ , I/ (2 λ 1 )) where u j ∼ N (0 , 1 / (2 λ 1 )) (i.i.d.) and ∑ d j =1 u j z j + u d +1 obeys the Gaussian distribution with mean 0 and variance ∑ d j =1 1 2 λ 1 1 d + 1 2 λ 1 = 1 λ 1 . In the same vein, we also have

Here, define | z | := |{ i ∈ { 1 , . . . , k } | z i > 0 }| for z ∈ {± 1 / √ d } d which is the number of positive elements of z in the index set { 1 , . . . , k } . For a fixed number i ∈ { 0 , . . . , k } , we let

then we can see that and

A Proofs of Propositions 3 and 4

Proof of Proposition 3. Remember that

Let b i = 2 i -k for i = 0 , . . . , k , let ζ > 0 be the positive real such that E u ∼ N (0 , 1) [2 tanh( ζ + u )] = 1 (note that, this also yields E u ∼ N (0 , 1) [2 tanh( -ζ + u )] = -1 by the symmetric property of tanh and the Gaussian distribution), and let

̸

because ⟨ ξ, z ⟩ -b i = ∑ k j =1 sign( z j ) b i = 2 | z | -k -b i = 2( | z | -i ) . Hence, we have that

Then, since tanh( u ) + 1 = e u -e -u e u + e -u +1 = 2 1+ e -2 u , if | z | = i and | u | ≤ 1 /λ 1 ,

and if | z | = i and | u | ≤ log( ρk ) / 2 ,

Therefore, when | z | = i ,

̸

where g is the density function of N (0 , 1 /λ 1 ) , and when | z | = i ,

where we used the upper-tail inequality of the Gaussian distribution in the second inequality. Hence, it holds that

because ¯ R = k . Therefore, by taking ρ > 1 sufficiently large, we also have

where the constant hidden in Ω( · ) is uniform over any | z | . Hence, there exists c ′ 2 > 0 such that Y ˆ f ( Z ) > c ′ 2 almost surely. Then, if we let µ ⟨ a ⟩ ( A ) := µ ( aA ) for a ∈ R , a probability measure µ and a measurable set A , then we can see that ˆ f is represented as

Then, by letting c 2 = ℓ (0) -ℓ ( c ′ 2 ) , we have

Next, we bound the KL -divergence between ν and µ ∗ . Notice that the convexity of KL-divergence yields that

because k ≤ d , which gives the assertion.

Proof of Proposition 4. The first assertion can be seen by ¯ L ( µ [ λ ] ) ≤ ¯ L ( µ ∗ ) + λ KL( ν || µ ∗ ) ≤ ℓ (0) -c 2 + λR < ℓ (0) -c 2 + c 2 / 2 = ℓ (0) -c 2 / 2 , by the assumption of λ such that λ < c 2 / (2 R ) . The same argument is also applied to L ( µ [ λ ] ) , that is, L ( µ [ λ ] ) ≤ ℓ (0) -c 2 / 2 .

We show that the optimal solution f µ [ λ ] should satisfy | f µ [ λ ] ( z ) | = | f µ [ λ ] ( z ′ ) | for all z, z ′ ∈ {± 1 / √ d } d by the symmetric property of the distribution. For that purpose, we construct a group action on {-1 / √ d, 1 / √ d } d . For notational simplicity, we consider Z d := {-1 , +1 } d instead. Z d can be equipped with a group structure where the binary operation between z, z ′ ∈ Z d is given by the element wise product z · z ′ = ( z i z ′ i ) d i =1 . Let T j : Z d → Z d be a group action that flips the j -th element T j z = ( z 1 , . . . , -z j , . . . , z d ) . It is obvious that T j is bijective. Then, we can show where

that there exists a sequence of indices σ (1) , . . . , σ (2 d ) ∈ { 1 , . . . , d } such that the 'orbit' generated by the chain of the group actions T σ ( k ) ◦ · · · ◦ T σ (2) ◦ T σ (1) z for k = 1 , . . . , 2 d covers the entire elements of Z d for any initial state z ∈ Z d and the last state turns back to the initial state, that is, T σ (2 d ) z = z . This can be shown by induction. If d = 1 , we just need to take σ (1) = σ (2) = 1 (indeed, T σ (1) z = -z and T σ (2) ◦ T σ (1) z = z , which satisfies the condition). Suppose that the argument is true for d = 1 , 2 , . . . , d ′ -1 , then we show that it is true for d = d ′ . Indeed, if we write σ ′ (1) , . . . , σ ′ (2 d ′ -1 ) be the corresponding sequence for d = d ′ -1 , then let

for i = 1 , . . . , 2 d ′ -1 -1 , then the corresponding orbit covers all the elements with the last element is fixed except the initial state. Then, we flip the last coordinate at the 2 d ′ -1 +1 -th step as

and then we 'trace-back' the orbit as

for i = 1 , . . . , 2 d ′ -1 -1 . By the construction, we notice that, after 2 d ′ -1 +2 d ′ -1 -1 steps, the state becomes

Therefore, by flipping the last coordinate again, the state can come back to the initial state, that is, by setting σ (2 d ′ ) = d , we have T σ (2 d ′ ) ◦ · · · ◦ T σ (1) z = z while the orbit covers the entire elements of Z d ′ .

We can define the action T j to {± 1 / √ d } d in the same manner. We note that the distribution of Z is invariant against the action of T j , that is, P Z = T j # P Z where T j # is the push-forward induced by T j . Here, let ˆ µ := µ [ λ ] and ˆ f = f ˆ µ . We also define an 'adjoint operator' T ∗ j : P → P for j ∈ { k +1 , . . . , d } such that f µ ( T j z ) = f T ∗ j µ ( z ) for all z ∈ {± 1 / √ d } d . We can easily check that there exists T ∗ j µ that satisfies this condition. Indeed, we may take T ∗ j µ such that T ∗ j µ ( A × B × C ) = µ ( T j A × B × C ) where T j A = { T j z | z ∈ A } for any measurable set A ∈ B ( R d ) , B ∈ B ( R ) and C ∈ B ( R ) . Since the probability on product sets uniquely determines the probability on the Borel sigma algebra B ( R d ) × B ( R ) × B ( R ) = B ( R d +2 ) , T ∗ j µ is uniquely determined.

Now, we consider the sequence of the group action T σ (1) , . . . , T σ (2 d -k ) constructed above acting on the last d -k coordinates. Since the label is independent of the last d -k coordinates, we have L ( T ∗ σ (1) ◦ · · · ◦ T ∗ σ ( i ) ˆ µ ) = L (ˆ µ ) for all i = 1 , . . . , 2 d -k . Moreover, the symmetricity of the Gaussian distribution yields that KL( ν, T ∗ j µ ) = KL( ν, µ ) . Hence, if we take

then by the convexity of the objective L , we have

Then, by the strong convexity of the objective L (due to the KL-divergence), we have ˆ µ ′ = ˆ µ . Then, we notice that

for any z ∈ {± 1 / √ d } d , which means that f ˆ µ ( z ) is constant against any change of the last d -k coordinates.

Next, we consider to change the first k coordinates. We also construct a sequence of the group action T σ (1) , . . . , T σ (2 k ) acting on the first k coordinates. In this case, the action of T j on z also flips the

corresponding label. Hence, if y ( z ) is the label of z , then the label for T j z is -y ( z ) = y ( T j z ) . By defining the adjoint operator T ∗ j as

for j ∈ { 1 , . . . , k } , then by the symmetric shape of tanh , we have

for any z ∈ {± 1 / √ d } d . Wealso have that KL( ν, ˆ µ ) = KL( ν, T ∗ j ˆ µ ) by the symmetry of the Gaussian distribution. Therefore, onece again by the invariance of the P Z against the action of T j and the convexity of L , we have that where

By the strong convexity of L due to the KL-divergence term, we have µµ ′ = ˆ µ . Hence,

for any z ∈ {± 1 / √ d } d , which means that y z f ˆ µ ( z ) is constant.

Finally, combining Eq. (9) with the fact ¯ L (ˆ µ ) ≤ ℓ (0) -c 2 / 2 , we have that P ( Y f ˆ µ ( Z ) > 0) = 1 and Y f ˆ µ ( Z ) ≥ c 2 / 2 almost surely because ∥ ℓ ′ ∥ ∞ ≤ 1 .

B Proofs of auxiliary lemmas

B.1 Difference between finite particle network and its integral form for the optimal solution

Lemma 2. Under Assumption 1, if ∥ z ∥ ≤ 1 for any z ∈ supp( P Z ) , then it holds that

Proof. Note that the left hand side can be rewritten as

From the standard argument of concentration inequality corresponding to the Rademacher complexity, it holds that

for any t > 0 , where ϵ = ( ϵ i ) N i =1 is an i.i.d. sequence of the Rademacher variable, X ∗ = ( X i ∗ ) N i =1 is an i.i.d. sequence generated from µ ∗ , and the probability is taken with respect to the realization of X ∗ . Since tanh is Lipschitz continuous with | tanh ′ | ≤ 1 , the contraction inequality of the Rademacher complexity yields that

To bound the right hand side, we evaluate the moment of µ ∗ . Since µ ∗ is the stationary distribution of the SDE (1). Hence, if X 0 ∼ µ ∗ , the process ( X t ) t ≥ 0 obeying the MFLD (1) satisfies X t ∼ µ ∗ for any t ≥ 0 . Then, the infinitesimal generator of the MFLD gives that

which yields that

Therefore,

Then, for x ≥ 0 , it holds that

This yields that

This gives the assertion.

B.2 Proof of Lemma 1

The proof follows the line of Lemma 5.5 of Chen et al. (2020). Let the empirical Rademacher complexity of a function class ˜ F be

We can see that Rad( · ) = E ( z i ) n i =1 [ ̂ Rad( · )] . Then, it holds that

where the first inequality is by the Donsker-Varadha duality formula of the KL-divergence (Donsker and Varadhan, 1983). The second term in the right hand side can be bounded as

Therefore, we have that

The right hand side can be minimized by γ = √ nM/ ¯ R , which yields

This gives the assersion.

C Proofs of Theorems 1 and 2

To show the theorems, we first prepare the following uniform bound, which is crucial to obtain the fast learning rate.

Lemma 3. It holds that

uniformly over µ ∈ P with KL( µ, µ ∗ ) ≤ M with probability 1 -exp( -t ) for any t > 0 .

Proof. Let B k = M/ 2 -k and C k = 2 ¯ R/ 2 -k ′ for k = 0 , 1 , 2 , ... , and F k,k ′ := { µ ∈ P | KL( µ, µ ∗ ) ≤ , ∥ f µ -f µ ∗ ∥ ∞ ≤ C k } . Then, the standard concentration inequality regarding to the Rademacher complexity (Theorem 3.1 of Mohri et al. (2012)) with the contraction inequality (Theorem 11.6 of Boucheron et al. (2013) or Theorem 4.12 of Ledoux and Talagrand (1991)) yields that

for each k, k ′ . Then, since ∥ f µ -f µ ′ ∥ ∞ ≤ 2 ¯ R for any µ, µ ′ ∈ P , by taking uniform bound for all k = 0 , . . . , log 2 ( n 2 ) and k ′ = 0 , . . . , log 2 (2 ¯ Rn ) , we have that

By resetting t ← t +loglog 2 (8 n 2 M ¯ R ) , we obtain the assertion.

C.1 Proof of Theorem 1

Let ˆ µ := µ ( K ) and µ ∗ := µ [ λ ( K ) ] . By the optimality condition of ˆ µ , we have

Moreover,

By the optimality of µ ∗ , the right hand side can be lower bounded as

Then, we have KL( µ ∗ , ˆ µ ) ≤ ( ϵ ∗ +2( ¯ R +1)) /λ ( K ) . Hence, the uniform bound on the Rademacher complexity (Lemma 3) with M 0 = ( ϵ ∗ +2( ¯ R +1)) /λ ( K ) , it holds that

with probability 1 -exp( -t ) for any t > 0 . Moreover, Lemma 1 gives

Therefore, by setting t = s n with s < 1 , we obtain that

with probability 1 -exp( -s √ n ) , where a arithmetic-geometric mean relation and an inequality ( a ∨ b ) 2 ≤ a 2 + b 2 ( a, b ≥ 0) are used in the last inequality. Here, we know that the Hellinger distance between µ ∗ and ˆ µ can be bounded as

Hence,

By summarizing the argument above and noticing ¯ L (ˆ µ ) -¯ L ( µ ∗ ) -(ˆ µ -µ ∗ ) δ ¯ L ( µ ∗ ) δµ ≥ 0 , it holds that

with probability 1 -exp( -s √ n ) . Then, letting s = τλ ( K ) , then we have

with probability 1 -exp( -τλ ( K ) √ n ) for any τ > 0 . Hence, if we take n sufficiently large and τ sufficiently small so that the right hand side is smaller than c 2 0 , then we have

More precisely, if we let

C.2 Proof of Theorem 2

Let ˆ µ := µ ( K ) and µ ∗ := µ [ λ ( K ) ] . From the proof of Theorem 1, with probability 1 -exp( -t ) , it holds that

Since Y f µ ∗ ( X ) ≥ c 0 almost surely, the Markov's inequality yields that

This yields the assertion.

D Additional Experiments

In addition to the 2 -sparse parity problem, here we give additional experiments for the k -sparse parity problems with k = 3 and k = 4 . The setting is the same with the main text, except that the width, step size, the number of iterations are modified to N = 2 , 000 , η = 0 . 2 , and T = 10 , 000 , and we take different grids of n and d . We ran the experiment 10 times and plotted the mean for each n and d . The results are provided in Figure 2.

For k = 3 , which is shown in Figure 2 (a), we can see that there appear phase-transition lines of n = Θ( d 2 ) and n = Θ( d ) , which correspond to the test accuracy of around 90% and the test accuracy of around 70% , respectively. The former almost perfect classification means the Type I convergence in 4.2, where the classification error gets exponentially small with respect to n/d 2 . The latter means the Type II polynomial convergence.

For k = 4 , presented in Figure 2 (b), only the linear phase-transition, that corresponds to the test accuracy of around 70% , is observed.

Together with Section 5, we conclude that the our theoretical results are experimentally verified for the 2 -sparse parity problem as well as higher-order sparse parity problems. The reason why the quadratic line was not observed in k = 4 would be partially because the hidden coefficients and the required number of steps are exponentially large with respect to k .

1500

3000

4500

6000

7500

sample size

n

20

40

60

dimension

d

n

d

n

d

2

0.6

0.7

0.8

0.9

1.0

test acc

1500

3000

4500

6000

7500

sample size

n

5

10

15

dimension

d

n

d

0.6

0.7

0.8

0.9

1.0

test acc

(a) Result with

k

= 3

.

(b) Result with

k

= 4

.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper characterizes the generalization of neural networks, and prove efficient sample complexity guarantees in the mean-field region with the presence of feature learning. It presents a general framework to establish sample complexity of MFLD for binary classification problem. Such framework can be used to obtain generalization guarantee for the learned neural network in discrete-time and finite-width settings. When applied to the k-sparse parity problem, the proposed framework yields an improved convergence rate and dimension dependence.\n\nSTRENGTHS:\nThe technique used in this analysis is quite interesting and unique. Unlike existing margin bounds for neural networks that rely on norm control, this study takes a different approach by considering that MFLD optimizes the distribution of parameters rather than the parameters themselves. This perspective allows for an improved analysis of the sample complexity and convergence rate.\n\nWEAKNESSES:\n1. This study is focus on binary-classification problem. Definitely it could be a good starting point, but at the same time, it is hard to argue that the results for binary-classification problem is very significant.\n\n2. The authors focus on k-sparse parity problem (2-sparse in particular) as an illustrative example throughout the paper. While I understand the 2-sparse parity problem is a well studied problem with a rich body of literature, I still do not quite understand the significance or the motivation of studying this particular problem.\n\n3. The width and number of iterations required in this analysis is of $O(e^d)$ which I think is far from reasonable in practice.\n\nQUESTIONS:\n1. What is the difference of the last two rows in Table 1? Why they have the identical region/method, width, and number of iterations, but different class error?\n\n2. The work done by Telgarsky  (2023) requires width of $O(2^d)$ not $O(d^d)$ right?\n\nLIMITATIONS:\nI believe the authors have a sufficient discussion on the limitations.\n\nA couple of suggestions:\n\n  - $d$ has been used in the introduction several times but never been properly defined in this section (it has been defined in the table 1 though).\n\n  - In Assumption 2, last inequality $L(\\mu^*) \\leq log(0)-c_0$, I believe $log(0)$ here is a typo. The same typo also appeared in Proposition 3.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of learning the k-sparse parity problem with a two-layer network in mean-field regime. The main results are in two folds: 1. the authors proposed an annealing method to obtain a better rate of convergence. 2. the authors compute the classification error via computing the local Rademacher complexity.\n\nSTRENGTHS:\n1. The problem of learning a k-sparse parity function in a mean-field regime seems to be interesting.\n\n2. Proposition 4, which gives a characterization of the margin at the stationary distribution seems to be novel and interesting to me.\n\n3. The annealing methods proposed in this paper seem to be interesting, which gives a way to obtain a better convergence rate via bounding log-Sobolev constant properly. \n\n4. This paper is in general well-written, the main theoretical results are stated clearly, and I feel it is rather easy to understand the main results of this paper.\n\nWEAKNESSES:\nI haven't observed an obvious weakness in this paper, but I do have a few questions and comments, please refer to the \"Questions\" part.\n\nQUESTIONS:\n1. Regarding the annealing methods: If I understand the annealing methods correctly, the intuition is that under assumption 3, the log Sobolev constant is actually controlled by the loss, and since the loss is decreasing along the trajectory, the log Sobolev constant will be larger than the trivial bounds in ( Nitanda et al., 2022; Chizat, 2022). Thus it seems to me that annealing is added for the purpose of controlling the loss properly, rather than fundamentally speeding up the dynamics (in fact, annealing does not improve the convergence rate of the MFLD in general, see e.g. Chizat, 2022 in your reference ). Also, I'm wondering if it is possible to obtain a faster convergence through different annealing methods.?\n\n\n\n2. In Proposition 4, you characterize the margin of the stationary distribution, with small enough regularization. I'm wondering if this properties still hold in the limit of $\\lambda \\rightarrow 0$?\n\n3. Generality of the results:  I'm wondering if it's possible to generalize the result to any bounded smooth activation functions, especially propositions 3 and 4, and Theorem 1 and 2?\n\n\nMinor points and typos:\n\n1. Line 169: $\\log(0) \\rightarrow \\ell(0)$\n\n2. Line 203: $\\lambda^{(k)} \\rightarrow \\lambda^{(\\kappa)}$\n\n3. In the statement of Theorem 1, you have an assumption $Q: = ... > 0$. Given a specific setting (e.g. for logistic loss and taking the parameters in propositions 3 and 4), is this assumption verified?\n\nLIMITATIONS:\nNo limitations and potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper conducts a theoretical study of the generalization error and sample complexity of two layer neural network training with Mean Field Langevin Dynamics (MFLD) or (informally a version of )noisy gradient descent. Specifically the paper specializes the generalization error for subset parity problem to derive bounds. Empirical results are provided to illustrate the agreement between theory and empirical observations with synthetic data\n\nSTRENGTHS:\n(Note that I am not very familiar with the details of the subject matter and also related work)\n\n- Analytical generalization bounds may be of interest to the theoretical community especially given the recent interest seen with using subset parity as a dataset in analysis\n- The paper does a nice job of comparing existing and their classification error bounds for k-sparse (subset) parity problem to highlight their contribution. The fact that their bound decouples error from k is notable to this reader\n- This reviewer appreciates the experimental results in the paper. This figure summarizes the analytical findings included in Section 4\n\nWEAKNESSES:\n(Note that I am not very familiar with the details of the subject matter and also related work)\n\n- The paper appears to be spend a lot of real estate with preliminaries at the cost of sacrificing clarity of presentation of the main result. In section 4.1, I do not quite follow how the choice of regularization parameter reduces complexity\n- While analysis does decouple the dependence on k in the k-sparse parity task, this appears to be at the cost of the network width and the number of optimization steps which are both exponential compared to the work by Barak et al. (NeurIPS 2022). I am left wondering about the usefulness of this result to the community \n- (minor) I glanced at the code for the experiments and noticed that m = 2000. I assume this is the width value used in the paper. The paper states a value of 10_000 so I want to make sure that the shared code has a typo. It would be nice if the code is updated to reflect the exact value used in the paper for clarity\n\nQUESTIONS:\nPlease check the weaknesses section\n\n- Please consider updating Section 4 to make new work clear by perhaps moving some preliminaries to the supplement\n- Please update code to ensure the results included in the paper can be reproduced without any change(s)\n\nLIMITATIONS:\nThe authors do a very nice job of noting weaknesses in their discussion. No more changes necessary", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper proves optimization and generalization guarantees for MFLD for the $k$-sparse parity problem. It proves that if the network is sufficiently overparameterized (width > $e^{\Omega(d)}$) then $n = d$ samples suffice, which is independent of $k$. Furthermore it proves an exponential convergence rate when $n \ge d^2$.

STRENGTHS:
- The paper is able to prove concrete Rademacher-based generalization guarantees for MFLD, including explicit $\epsilon$ dependencies in different regimes

Also, I think the comparisons with Chen and Meka (2020) are unnecessary as that paper focuses on the case when the covariates are Gaussian which is a very different problem.

WEAKNESSES:
- This paper does not consider label noise which would likely break the exponential convergence rate in the $n \ge d^2$ regime. It is not clear to me whether the techniques would extend to this setting.
- Like many mean field papers, discretizing the dynamics requires exponential width (at least $e^{\Omega(d)}$). The tradeoff is therefore width $O(1)$ and sample complexity $O(d^k)$ vs width $e^{O(d)}$ and sample complexity $d$. However, this problem is not unique to this paper and appears fundamental to a lot of mean field analysis.

QUESTIONS:
- How difficult would it be to adapt the techniques in this paper to the case of label noise (i.e. flip each label with probability $p < 1/2$)?

LIMITATIONS:
The authors have adequately addressed the limitations of this paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09