Asymptotic Analysis of Two-Layer Neural Networks after One Gradient Step under Gaussian Mixtures Data with Structure
Authors: Samet Demir
Zafer Dogan
Venue: ICLR 2025 Poster (2025)
Keywords: deep learning theory
random features
Gaussian equivalence
universality
high-dimensional asymptotics
Primary area: learning theory
TL;DR: We study the impacts of Gaussian mixtures data assumption to feature learning in neural networks trained with one gradient step in order to bridge the gap between isotropic data assumption and real datasets.

Abstract: In this work, we study the training and generalization performance of two-layer neural networks (NNs) after one gradient descent step under structured data modeled by Gaussian mixtures. While previous research has extensively analyzed this model under isotropic data assumption, such simplifications overlook the complexities inherent in real-world datasets. Our work addresses this limitation by analyzing two-layer NNs under Gaussian mixture data assumption in the asymptotically proportional limit, where the input dimension, number of hidden neurons, and sample size grow with finite ratios. We characterize the training and generalization errors by leveraging recent advancements in Gaussian universality. Specifically, we prove that a high-order polynomial model performs equivalent to the non-linear neural networks under certain conditions. The degree of the equivalent model is intricately linked to both the "data spread" and the learning rate employed during one gradient step. Through extensive simulations, we demonstrate the equivalence between the original model and its polynomial counterpart across various regression and classification tasks. Additionally, we explore how different properties of Gaussian mixtures affect learning outcomes. Finally, we illustrate experimental results on Fashion-MNIST classification, indicating that our findings can translate to realistic data.

Published as a conference paper at ICLR 2025
ASYMPTOTIC ANALYSIS OF TWO-LAYER NEURAL NETWORKS AFTER ONE GRADIENT STEP UNDER GAUSSIAN MIXTURES DATA WITH STRUCTURE
Samet Demir', Zafer Dogan
'MLIP Research Group. KUIS Al Center, Kog University *Department of EEE, Koç University {sdemir20, zdogan) kuu edu.tr tr
ABSTRACT
In this work, study the training and generalization performance of two-layer neural networks (NNs) after one gradient descent step under structured data mod- eled by Gaussian mixtures. While previous research has extensively analyzed this model under isotropic data assumption, such simplifications overlook the com- plexities inherent in real-world datasets. Our work addresses this limitation by analyzing two-layer NNs under Gaussian mixture data assumption in the asymp- totically proportional limit, where the input dimension, number of hidden neu- rons, and sample size grow with finite ratios. We characterize the training and generalization errors by leveraging recent advancements Gaussian universality. Specifically, prove that high-order polynomial model performs equivalent the nonlinear neural networks under certain conditions. The degree of the equiv- alent model is intricately linked to both the "data spread" and the Icarning rate employed during one gradient step. Through extensive simulations, we demon- strate the equivalence between the original model and its polynomial counterpart across various regression and classification tasks. Additionally, we explore how different properties of Gaussian mixtures affect learning outcomes. Finally, il- lustrate experimental results on Fashion-MNIST classification, indicating that our findings can translate realistic data.
INTRODUCTION
Understanding how neural networks learn from data and generalize to unseen examples is a funda- mental problem in deep learning theory (Goodfellow et al. 2016). While significant progress has been made in analyzing two-layer neural networks (Pennington Worah, 2017 Ba et al. 2022) under simplified data models, such isotropic Gaussian inputs, these data models fail capture the intricate structures present in practical datasets (Ba al. 2023). Therefore, our research focuses on characterizing the training and generalization performance of two-layer neural networks under more realistic data conditions.
The theoretical analysis of two-layer neural networks has largely focused on the so-called lazy- training regime (Ghorbani et al.) 2019). where features experience minimal or no training. One example of this regime called random feature model (Rahimi & Recht 2007), where the first layer is randomly initialized and fixed while the second layer is trained. Indeed, the random features have been crucial for exploring phenomena such as "double descent" (Mei Montanari (022) observed in over-parane parameteriize models. However, the random feature model lacks feature learning. Recent studies have begun to fill this gap by analyzing two-layer networks where the first layer trained with a single gradient descent step (Ba et al.) 2022 Damian et al., (022). This shift significantly enhances our understanding how feature learning influences generalization performance.
Despite recent advancements, much of the existing work with feature learning mainly relies on overly simplistic data distributions, such as isotropic Gaussian or spherical inputs (Moniri et al, 2024 Dandi et al. 2023a; Cui et al. 2024), and occasionally spiked covariance models (Ba et al. 2023 Mousavi-Hosseini et al.] (202)). While these assumptions simplify analysis and provide in- sights, they overlook the complex nature of real-world data. In practical learning tasks, data is often

---

Published as a conference paper at ICLR 2025
better represented mixture of distributions (Seddik et al. 2020 Dandi al. 2023b). Moreover, real-world datasets are typically high-dimensional but often exhibit low intrinsic dimensionality (Facco et al. 2017| Spigler et al. (020). This gap between theoretical data assumptions and actual data distributions highlights the need for more sophisticated analytical frameworks that capture the mixture nature and low-dimensional structure of real-world datasets.
In this work, we study the performance of a two-layer neural network trained with a single gradient descent step under Gaussian mixture data with covariances including low-dimensional structure. Our data model captures both the mixture nature and intrinsic low-dimensionality in wolld datasets. By leveraging recent advancements in Gaussian universality, we provide comprehensive characterization of training and generalization errors in the asymp- totically proportional limit, where the input dimension, number of hidden neurons, and sample size grow proportionally. Our analysis shows that, under specific conditions, finite-degree polynomial model- to as the "Hermite model"- achieve equiva- lent training and generalization performance to that of nonlinear neural network. We find that the degree of the equivalent polynomial model is closely linked to both "data spread" and the learning rate used during the training of the first layer. We release all our code at ht /ttp:://githie com/ KU-MLIP/
Our contributions can be summarized as follows:
We establish theoretical framework for characterizing the training and generalization errors of two-layer neural networks under Gaussian mixtures data with covariances featuring additional low-dimensional structures.
We demonstrate that a finite-degree polynomial model serves as an equivalent performance model, simplifying the analysis of neural networks under Gaussian mixtures data assumption.
Through extensive simulations, including Fashion-MNIST classification, we validate our findings and highlight the significant impact of the structure of data on learning outcomes.
Notations We adopt the standard notations established by Goodfellow et al. (2016) throughout this work unless specified otherwise. The spectral norm of matrix F is denoted as [| We use the notation f(-) x 9(-) to indicate that the functions ff and are of the same order with respect to the parameters k. n, and m. The notation (-)) represents the big-oh notation in relation to these parameters, and we also define (()(-) as shorthand for ((((-) polylog k). effectively allowing us to omit polylogarithmic factors for clarity. Element-wise multiplication is indicated by the symbol o. Additionally, we denote the conditional expectation of a random variable X given a condition C as E[X C). while Xlo is used to refer to the conditional random variable X C. We write 3/4- to denote 3/4 for some small 0.
2 SETTING
We consider a supervised learning setup through two-layer neural network (NN) defined by
1
w o (Far)
To
Vk
where æ E RR" represents the input vector and R + R the nonlinear activation function. In this framework, F E R2xx and w E IR* denote the weights (parameters) of the first and second layers of the NN, respectively. In this study, we focus on the training and generalization performance of this model under simplified training procedure (one gradient descent step on the first layer) under structured Gaussian mixtures data assumption. Our analysis will be in the proportional asymptotic limit where the number of training data, the input dimension, and the number of features jointly diverge with finite ratios. This regime intuitively represents a scenario in which the width of the network is proportional to the size of the dataset, aligning with common practices in model scaling. We consider data samples drawn from a Gaussian mixture model:
Data model
T ~ Σ CpNN(Hg, and == a. (E" x c),
where C E = is the number of mixture components and e is a random variable denoting the com- ponent assignment for the sample (ær, with IP(c = j) Pj for j {I,...; C). Furthermore,
2

---

Published as a conference paper at ICLR 2025
My E IR" and Σ, R denote the mean and covariance of j-th component, respectively. We fur- ther assume that 3 exhibits a certain low dimensional structure that can be described as finite-rank plus identity (see assumption (A.4) in Section 4). We then define Σ == Cov(æ) the covariance matrix of the input vector æ, and note that its spectral norm Σ will be the measure of data spread in our context. Finally, v, TR2 R is unknown label generation function that be nonlinear function of E" x for regression problems or the component index e for classification problems. Note that label y depends only on a single direction (a.k.a single-index target functioon.. This is moti- vated by the fact that the NN (trained with one gradient descent step) can only learn one direction target functions is left future work.
about the labels (Lemma]) I]. Extension
Training procedure We restrict ourselves to a simplified two-stage training procedure introduced in (Ba et al. | Damian et al. 2022), where we first learn features by taking one gradient descent step on the first-layer parameters and then estimate the second-layer parameters separately. This procedure is summarized as follows:
i) Gradient descent on the first layer: Given { D:) 0:)] a set of training samples drawn from 2 first fix the initialization and perform single gradient descent step on F with respect to squared loss. The gradient update with learning rate 1) > 0 is given as
F == F + GG
(3)
where gradient matrix G is defined
1
1
www "oo'o(FFX (FXX) od(F)) x
wÿ
(4)
Vk
with X
[ 22
Zm [T and = [81. 12. Dm
ii) Ridge regression for the second layer: With the trained first layer F, then train the second layer weight vector using a new set of training samples {< 3:) drawn from 22. as follows:
2
arg min
Vi
w o (F
2 |
E
2m
V/
where À 2 0 is the regularization constant.
Note that the ridge regression is performed on the same data, then after one gradient step, F will no longer be independent of X, which would significantly complicate the analysis. Instead, new set of training data is used to circumvent this difficulty, following the prior work by(Ba et al. (2022).
Performance metries To evaluate the performance of the two-layer neural network defined in equation ① and trained on the data model outlined in equation (2). we establish key metrics that quantify both training and generalization errors. After training both layers, define the training error as:
y1
1 uo oo Vk
(Fxx))
2 + 11001
(6)
2m
which captures the discrepancy between the predicted outputs and the actual labels for the training dataset, incorporating a regularization term controlled by the parameter À to prevent overfitting. In addition, we assess the generalization error, denoted as:
(1-
1 T o (F 21
G E (m.
which reflects the expected prediction error on unseen data. This metric provides insight into how well the model is likely to perform in practice, beyond the training set.
Scalings of data spread and learning rate Our analysis indicates that the combined scaling of the learning rate 17 and data spread | is more critical to our theoretical framework than their individual scalings. To capture this relationship, we introduce a "strength parameter" E [0, 1] governing the scaling 17| D x n². We also define a "weighting parameter" E [0 that controls individual (((1o)) and 17 no This parameter interpolates between two extremes: one where 1 = n² and |Σ|| x 1, and another where y x 1 and |Σ| x n This setting covers a wide range of scenarios for generalization performance, which we further illustrate in Appendix D.
scalings: Σ
3

---

Published as a conference paper at ICLR 2025
3 RELATED WORK
The random feature model (RFM) was initially proposed as computation- ally efficient approximation kernel methods (Rahimi & Recht 2007). RFMs are closely related to the Neural Tangent Kernel (NTK) (Jacot et a 2018 since both of them provide linear approxi- mations of two-layer neural networks (Ghorbani al. 2020, 202I). Despite their simplicity. RFMs have proven instrumental in understanding various facets of machine learning, including general- ization (Mei & Montanari (022), transfer learning (Tripuraneni et al. (021). out-of-distribution performance (Lee et al. (023), uncertainty quantification (Clarté et al. (023), and robustness (Has-) sani & Javanmard 2024). Recently, the RFM has garnered renewed interest as a means to investigate the behavior of two-layer neural networks, particularly within the lazy training regime, where the pa- rameters of the network experience minimal changes during training (Pennington & Worah (017).
Random features -
Comprehensive asymptotic analyses have been conducted for RFMs (Dhifallah & Luj Lu] [220 Goldt iet al. 2020 Mei & Montanari] 2022 Goldt et al. 2022 Hu & Lu. 2023) and their deep counterparts
(Schröder al.
2023) Zavatone- Veth & Pehlevan [0023] Bosch al. 2023). further elucidating
their theoretical underpinnings and practical implications.
Universality - One approach to the asymptotic analysis of random feature models (RFMs) in- volves utilizing equivalent models. Under the assumption of isotropic data, specifically when æ ~ (O, I,), it has been demonstrated that the random feature model w" (o"'((F)) is equivalent to the following linear model (Goldt et al. 2022 Hu & Lu 2023):
w (hol + by Fx + hīz),
(8)
where 2 ~ (O, I&) and ho, h3 > are constants. This equivalence is framed within the con- cept of "universality" from random matrix theory (Couillet & Liao 2022). as the random features or (Fx) can be replaced with Gaussian features that share equivalent mean and covariance proper- ties (Hu & Lu 2023). The universality of random features has since been extended to empirical risk minimization, allowing for broader analyses beyond RFMs (Montanari & Saeed 2022). While the results by Montanari & Saced (2022) focused on covariate inputs. Dandi et al. (2023b) recently broadened this framework to encompass inputs distributed as mixtures. Furthermore, Demir & Do- gan (2024) extended the universality of random features to Gaussian inputs with spiked covariance, highlighting the significance of structured data in RFM applications.
Feature learning- - While the existing literature on random features and universality provides valuable insights. it often overlooks the crucial aspect of feature learning inherent in neural net- works. Several studies have explored the dynamics of two-layer networks, particularly in the mean- field regime, which examines training behavior with small learning rates (Mei et al.. 2018 Bordelon & Pehlevan 2024). In this context, we focus on two-layer neural networks where the first layer trained with a single gradient descent step (Ba al. (2022), addressing the feature learning de- ficiencies found in random feature models. Notably, Ba et al. (022) established the importance of the learning rate in surpassing the performance of the linear model represented in equation (8). Subsequent works by Dandi et al.] (2023a). MMniri et al. 2024). and Cui et al.] (2024) have fur- ther analyzed neural networks after one gradient step through equivalent models, specifically for learning rates v) x ke with studies, our approach considers Gaussian mixture inputs as defined in equation (2), allowing us to investigate the intriguing effects of data distribution on feature learning. While Ba et al. (2023) and Moooussvi-Hosseinn et al. (202)) examined Gaussian inputs with spiked covariance, their findings lack equivalent models for precise performance characterization and also lack mixture aspect of our data model 2 highlighting the novelty and significance our work in this area.
[0, and isotropic Gaussian inputs. In sharp contrast to these
Gaussian mixtures- - Most of the works discussed in this section, with the exception of Dandi et al. (20236), have assumed Gaussian or spherical inputs, which do not adequately capture the mixture nature of class-based problems. Additionally, many of these studies have relied on isotropic covariance, limiting their applicability to simplified scenarios. Recently, there has been a growing interest in analyzing the asymptotic performance of various machine learning problems under the assumption that data is generated from Gaussian mixture model (Mai & Liao 2019; Mignacco jet al., 2020; Loureiro et al.] 2021 Kini & Thrampoulidis 202T] Refinetti et al. 2021). Notably,
Refinetti et al. (002)) compared a two-layer neural network trained in the mean-field regime with
4

---

Published as a conference paper at ICLR 2025
random feature model using Gaussian mixture inputs in toy example setting. Furthermore, Loureiro et al. ( provided an asymptotic performance characterization for generalized linear models under the Gaussian mixture assumption. In contrast to these studies, work includes feature learning and introduces straightforward equivalent models for analyzing two-layer neural networks, thereby enhancing the understanding of how data distribution impacts learning dynamics.
ASSUMPTIONS
(A.1) The number of training samples m, imput dimension n, and number of hidden neurons k jointly diverge with finite ratios, which means m, n, + oo while //m, m/k E R
(A.2)
Σ
x n ((1 (1-)) and r x n for a, B E [0, 1]. Thus, 7/ Σ x n
(A.3) The data is generated according to 2 Furthermore, we let μ.c = 0 and Tr(E.) = Tr(Ea) for all c, E {1,.... C. to simplify analysis.
(A.4) Furthermore, E. admits the following decomposition for c (I,...,C),
de
Σ. I, +
0 Yos
where de E zt, oer > O for all i E {l...., do}, and (eab) dd is set of orthonormal
vectors in R", Note that max.,, Oz,i x n [o) by (A.2)
(A.5) Let ||_|| 1/21 for some C' > 0 so that E[((T x) = Õ(1)
(A.6) Let F= [fi, f2. ..... fa] with f: N (0. I,, (r()), and N(0, Iz/k).
(A.7) The target function a. R22 - R is a Lipschitz function.
(A.8) The activation a R + is a function with bounded derivatives and it satisfies E (b))) oo, which allows the following Hermite expansion
a<)) Σ 1
(10)
where H2 + R denotes j-th probabilist's Hermite polynomial (O OO'Donnell 2014
Chapter 11 11.2), hy = (.-.0)..)(44(()0(()))) and b = vn/Tr(E) e IR by (A.4)
this work, we adopt several key assumptions to facilitate our anal- ysis, starting with (A.1) which defines the proportional asymptotic limit (or linear scaling regime). This assumption, commonly used in the literature (Hu & Lu 2023), allows us to interchangcably parameters n, m, and & in derivations. While recognize the potential for extending to polynomial scaling regime (Hu et al., 2024). we leave that exploration for future work. Assumptions
Discussion of Assumptions
(A.2) through (pecfficcy address Gaussian mixture data model in equation (2). Assump-
tion (A.2) outlines the necessary range for the strength parameter 3, with potential for valuable
insights from an extended range. Assumption (A.3) is used streamline the derivation of Theorem
4 Yet, the zero-mean assumption f 0 for the mixture components can be relaxed as discussed in
Appendix F Additionally, (A.4)]extend) the spiked covariance model (Johnstone) |001 Baik al.,
2005 Ba et al., 2023) by positing a finite-rank plus identity covariance model, inspired by the low intrinsic dimensions of real-world data (Facco et al. 2017]] Spigler al., 2020). Assumption (A.5) is included to prevent diverging labels. Furthermore, relates to standard initialization practices for neural network parameters, ensuring [E( x²
and ||[lw]]
Assumption (A.7)
reflects typical expectations regarding labeling functions, while assumption (A.8) pertains to the ac- tivation function used in our proofs. Although some functions like polynomials may not meet the bounded derivatives criterion directly, results remain valid long the derivatives are bounded with high probability for inputs of the form bz where z ~ NV, 1). Thus, the equivalent polynomial activation in Theorem 4 is encompassed by our activation function assumption.
5

---

Published as a conference paper at ICLR 2025
5 MAIN RESULTS
In this section, we present our main results that enhance our understanding of two-layer neural net- works under structured data, We start by analyzing the gradient G' defined in equation (4), deriving decomposition in Lemmal 1 Then, decompose Fx, revealing its "structure" and "bulk" compo- nents (Lemma 2 This decomposition helps identify conditional feature map equivalent to a (Fæz) in terms of trainnn and generalization (Theorem 3) To simplify the conditional feature map, we approximate it using polynomial function, again leveraging the structure-bulk composition of Fx. Specifically, in Theorem O we show that the neural network is equivalent to a polynomial model with regards to the training and generalization errors. The equivalent polynomial model facilitates the analysis of the nonlinear activation function through a reduced set of coefficients.
First of all, we consider the gradient matrix G and its decomposition into spike and bulk components in the following lemma, which characterizes the decomposition.
Lemma 1 (Spike+bulk decomposition of the gradient). Consider the gradient G defined in 4 admits the following decomposition
Ir
G= uv + A.
(11)
where u == h w and v == X // (m//k) where h1 == (snoooo.))o''s))) Also, ||u]| = O(1)
||l|| = 0 (k~ and ||||
O (k- with high probability, where := 1 B(1 a) 2 0.
Proof. Appendix
This lemma is pivotal it provides decomposition of the gradient G' into a dominant rank-one term uvd and negligible residual term A. In this context, h2 w represents the scaled second- m V X captures the covariance between the inputs and labels. The residual term A exhibits a magnitude of O(k-)), indicating its diminishing contribution relative to the dominant spike term uoT, which has norm of O(k Consequently, the spike term uo? emerges as the primary component driving the updates to the first-layer weights. This allows us to express the updated feature matrix as F F + 7/ + ryue"" With this formulation hand, we can proceed to analyze Fx in the subsequent lemma, leveraging the insights gained from the gradient decomposition to further extend our understanding of how F and æ interact.
layer weights, while e
Lemma (Structure+bulk decomposition of Fa:). Suppose that 2
conditioned on e-th Gaussian Eic ~ ~O, E.). Thus, Fz can be equivalently written as FE!! Ø1c ~ (0, I.). Then, we use the orthogonal decomposition: z = To + z²
component of the mixture
for = = =>e -1/
where r. [v, Ye,,, Ye.2, -e.d,) for a set of vectors (Te,,) defined in assumption (A.4)
Re I -ITT and z= = (I, TT... - This leads to
FX!/22 = FF + A))= FE!/DE
(12)
Fizi (Buik)
llmc. (Structure)
where F- == + A and alp. := /
Proof. The result directly follows from the definitions, assumption (A.4) and the orthogonality.
Intuitively, Lemma 2 reveals that Fx behaves like noise with mean for given pair (c, k.). This insight enables us to interpret a (Fx) as random features associated with the specific conditions of (c. k..). By leveraging this observation in conjunction with established universality results for random features (Hu & Lu| 2023. Montanari & Saced 2022 Dandi et al. (2033)) we derive the following theorem. This theorem presents a conditional feature map that is equivalent to o (Fx) in terms of both training and generalization performance. This equivalence not only underscores the relevance of structured data in feature learning but also enriches our theoretical framework by connecting the behavior of neural networks to the well-studied properties of random feature mod- els. Through this approach, we enhance our understanding of how the underlying data distribution influences learning dynamics, paving the way for more nuanced analyses in subsequent sections.

---

Published as a conference paper at ICLR 2025
Theorem 3 (Conditional Gaussian equivalence). Under the assumptions (A.1 (A.8) consider fea- ture map ((x)) of (ax)) with the definitions in Lemma 2 feature map (conditioned on c, Kic)
Then, define the following conditional
o(z;c, = =v(c,Ke) + 2e (c,,no)) + +((c, (c)
(13)
vlc, B c, ciio] (icc == B (o(F
where
c, C:Ni
(14)
((c,ke)) == Cov (o(Fx)) c, (Ree)) (c, RT))
=Co
g ~ (0, In).
(15)
Consider replacing the feature map 6(x) with the conditional feature map ô(σ; c, ke). Then,
(i) the training error T with the feature map @(x), and that with the conditional feature map b(a; c, Ko), both converge in probability to the same value,
(ii) the corresponding generalization errors G also converge in probability to the same value if additional assumption [A.9) provided in Appendix B hold.
Proof. Appendix
Recall that e denotes the index of the Gaussian component in the input mixture, while Ke represents
the aligmment with the subspace of the structure defined in Lemma 2 Theorem D establishes that, after conditioning on (c, Ke). the feature map <(xx) a (Fx:) can effectively substituted with the conditional feature map p(x; c, Ke) without impacting training and generalization errors. This substitution streamlines our analysis by allowing the feature map to be expressed in terms of con- ditional expectations and covariances. Theorem 3 sets itself apart from prior results on Gaussian equivalence through its unique application of conditioning. While
findings build on the work of Hu & Lu (2023) and Dandi et al. (2023b), their results do not include similar conditioning due the absence of the structure described in Lemma 2 Moreover, the conditional Gaussian equiva- lence from Dandi et al. (2023a) is limited to conditioning on spike in the gradient under isotropic Gaussian data. contrast, our result incorporates conditioning on both the mixture component and highlighting a more nuanced interplay between data characteristics and feature learning dynamics. While Theorem is compelling on its own, we can further approximate the conditional feature map using polynomial feature map. To do so, we consider the i-th ele- ment of Fx with the decomposition given in Lemma 2 which leads aijk, + (fft)²²z 2 where
the structure in Lemma 2
asjmc = f: 1/22T Totsc and f == f: + nA, For Osjsz Lemma asserts that if { < B 4
then [a:]x.|' | = Ō(1/k1-9) for some e > with high probability. The vanishing nature of [ajx. enables us to approximate o [d:]k, + ( with ōy + +(--)²z++) where ō: R + IR is a polynomial activation function defined in equation (16). By utilizing this approximation, we can derive the following theorem, which further elucidates the relationship between conditional feature maps and polynomial approximations within our framework.
Theorem 4. Under the assumptions (A.1) (A.8) et be activation function. Suppose that there exist E + such that 1-2
[ < 3 < Define another activation function
J (111(1/6)) + thiz with ~~ N(0,1,
6((z)
(16)
and hi = v (o - with b = /n/Tr(E), Consider replacing the activation (x) with the polynomial activation 61(2) after the training of the first layer F as in (3. Then,
where hg = (H,
(i) the training error T with activation o, and that with the polynomial activation o both converge in probability to the same value,
(ii) the corresponding generalization errors G also converge in probability to the same value if an additional assumption [..9].ooovidee in Appendix B hold.
Proof. Appendix
C
7

---

Published as a conference paper at ICLR 2025
0,55
Hemile Medel
Nara Network
tnn
ReL
0.150
0.15
gonmstration stree
Signai
0.125
gereralicttion emnr
0.03
0.075
0.0
0.000
0.05
0.1
0.8
1.0
0.0
0.2
0.6
0.8
1.0
(a) Impact of k/m
(b) Impact of
(c) Impact of
(3 3/4 and a = 1/2)
(B 3/4 and k/m 1)
1/2 and k/m 1)
Figure 1: Generalization error comparison between neural network and the Hermite model. We set both the input dimension and the number of samples to n = m = 1000, with two Gaussian components (C 2) and covariance matrix ranks of d1 d2 1. The mixture ratio for both components is set to f1 = P2 = 0.5, and a regularization constant of À = 1e is applied. For the labels, utilize y ReLU ()) and we limit the maximum degree of the Hermite polynomial to I == 5 for numerical stability. The figure presents averages from 20 Monte Carlo simulations.
Theorem is pivotal as it establishes that, under specific conditions, the activation function a (x) in the neural network can be effectively substituted with a polynomial activation (z(x). which is
constructed from Hermite polynomials (O Donnell 2014, Chapter 11.2) up to degree [ (-1 1, without
compromising training and generalization errors. The strength parameter / log(nil EII plays
(o)(n)
critical role in determining the necessary degree of the polynomial activation. Notably, Theorem
4 extends the equivalence results presented by Moniri et al. (2024) to encompass more general data scenarios. However, similar to the limitations identified in Moniri et al. (2024) regarding the maximal scale for the learning rate, Theorem 4 does not address the maximal value for the strength parameter (33 = 1), in contrast to Theorem 3 While / + 1 implies / - oo, we observe that a finite 1 value enough achieve the equivalence of generalization errors in numerical simulations for 3 x 1. Furthermore, the choice of Hermite polynomials is particularly useful due to their orthogonality properties (Lemma in Appendix when applied to Gaussian inputs. Consequently, we refer to the resulting model as the equivalent "Hermite model"::
w (Fx),
(17)
where on is defined finite sum of Hermite polynomials scaled by coefficients hj. supplemented by a Gaussian noise term that accounts for residuals. This equivalence significantly simplifies our analysis by transforming the nonlinear activation into a polynomial form, thereby enhancing the tractability of the model while maintaining its performance characteristics. The polynomial repre- sentation offers two key advantages: first, has superior theoretical properties, such as easier per- formance characterization compared to general nonlinear forms; second, it defines an equivalence class of activation functions based on their polynomial coefficients. Activation functions with the same coefficients will yield identical performance outcomes, opening intriguing possibilities for fu- ture research. This equivalence class could be crucial in the search for optimal activation functions, enabling more targeted exploration of their effects model performance. Overall, by leveraging this framework, we can more readily explore the implications of feature learning and structured data on neural network behavior.
6 SIMULATION RESULTS AND DISCUSSION
In this section, we present our numerical results and provide detailed discussion of their implica- tions. Each result illustrates the generalization errors corresponding to the parameter of interest for each specific plot. To showcase the applicability of our theoretical findings, we evaluate three widely used activation functions simultaneously: ReLU (rectified linear unit), tanh (hyperbolic tangent), and Sigmoid (logistic function). This comparison allows to assess how different nonlinearities impact generalization behavior. validating our theoretical predictions and highlighting the practical relevance of our results in guiding activation function selection in neural network architectures.
8

---

Published as a conference paper at ICLR 2025
1.0
Hermit Model
22 2sterallastinn c000
0.030
gaeeiilzatoon irme
0.8
goncalaatio
0.6
0.025
0.06
0.4
ReL.U
0.04
0.2
unh
0.015
Sigmood
0.010
0.02
0.2
0.66
-1.0
-05
1.0
muxttre ratio (with Pt Ph)
mittre aligmmet 7a
covariance makk (with di
(a) Impact of mixture ratio for classification: V// 2c 3 (d1 = d2= I, no-alignment)
(b) Impact of mixture alignment for
(c) Impact of covariance rank for
regression: = (U Φ)
regression: y LU x)
(d) du = 1,
0.5)
P2 0.5, no-alignment)
Figure 2: Impacts of properties of the Gaussian mixture data model on generalization performance. Here, we the number of Gaussian components to C 2, with equal input dimensions and sample sizes of m = n = k = 1000. The parameters are configured with 3 3/4_.a =1/2, / = 4, and regularization constant of À 1e 4. For (a) and (b), the eigenvalues of the covariance matrix ( for each Gaussian component are fixed at 01,1 = 22. n², while in (c), the eigenvalues (0c, )
are sampled uniformly from the interval (0, n²). The results displayed are averages from 20 Monte Carlo simulations, with data resampled for each run.
Effeet of model complexity - Figure demonsttttttts that the generalization errors of both the neural network and the equivalent Hermite model closely align across all values of kt/ m, reinforcing our theoretical findings. Supporting this, Figure (given in Appendix (D3) reveals that the training errors for both models also match closely. However, since generalization performance is of greater interest than training performance, we will concentrate on generalization errors in the subsequent plots. It is worth noting that our remaining simulation results are presented for the case of kz /m 1, although similar outcomes can be observed for other k/m ratios, indicating the robustness of our findings across different settings.
High data spread instead of high learning rate for better performance In Figure lb we investigate the impact of a on generalization error while keeping 3 3/4 constant. Here, influ- ences the ratio of the learning rate 17 to the data spread, i.e., the norm of the input covariance matrix ||]], characterized by ||E|| x (22==)) and 1) Sa The results reveal that as a increases, the gen- eralization error also rises, indicating that the strength of the data's 's structure is more beneficial for generalization performance than the strength of the learning rate in the first layer. This observation underscores the importance of structured data in enhancing model performance.
Larger strength parameter |9 for improved generalization - In Figure Ic we examine the effect of 3 on generalization error while keeping a 1/2 constant. The parameter governs the product of the learning rate and the norm of the input covariance matrix E|l (a.k.a. data spread) through the relationship 7/ Σ|| x n The results indicate that as increases, generalization errors decrease, reflecting the benefits of more complex data distributions (larger ||E|| and higher learning rates associated with larger 3 values. Note that while higher 3 leads to better generalization in setting, value shapes the curve of the generalization error with respect to B, as illustrated in D.2
Appendix D.2
Properties of Gaussian mixture- - Next, we investigate how the properties of the Gaussian mix- ture data model (2) influence generalization performance. Figure 2 illustrates the generalization errors associated with various data characteristics. First, we explore the effects of imbalanced data through different mixture ratios in Figure 2a While both ReL.U and Sigmoid activations achieve zero error for completely imbalanced cases, the mixture ratio has minimal impact on generalization error with the tanh activation function due to its symmetry and the zero-mean input, which prevents bias toward any class. Next, we examine the significance of mixture diignment-speciffee ment between the spiked directions two Gaussian components on generalization performance in regression. As shown in Figure 2b increasing alignment InT,, 72.1 towards 1 decreases generaliza-
the align-
tion error since the Gaussian mixture reduces single Gaussian distribution when hn 72.1 1.

---

Published as a conference paper at ICLR 2025
1.0
Hermee Model
tach
0.6
Sigmoid
geseraliazton ernr
0.5
0.6
0.4
0.2
0.2
4,70
4//00
(a) Class O vs. Class
(b) Class vs. Class 9
(T-shirt/top Trouser)
(Sandal vs. Ankle boot)
binary classification for |Σ| = n and 1) = 1. The data is generated from conditional GAN trained on Fashion- dataset and pre-processed. For the pre-processing, the inputs from each class are demeaned, re- re-scaled and added noise such that assumptions (A.2)-(A.4) are satisfied. m 500, À le 4 and 5. Details for the simulations and examples of input images after the pre-processing are provided in Appendix
Figure 3: Simulation results on
This suggests that mixture data presents greater challenges than single Gaussian data. Finally, Figure relreveals that higher effective rank (dz) of the covariance matrix leads to decreased generalization error. This occurs because a higher rank, on average, leads to a larger maximum eigenvalue due to random sampling of eigenvalues for the covariance matrices.
Collectively, the results in Figures 1 and confirm that the generalization performances of neural networks align closely with those of the equivalent Hermite model, underscoring the strength of our theoretical findings across diverse scenarios.
Towards theoretical results on real data In this work, we provide theoretical insights on realistic data, focusing on Gaussian mixture models, since previous studies (Seddik et al. 2020 Dandi et al. 2023b) have shown that data generated by generative adversarial networks (GANs) resemble Gaussian mixtures. We present simulation results from a conditional GAN (cGAN) (Mirza & Osindero 2014) trained on the Fashion-MNIST dataset (Xiao et al., (017), enabling samples conditioned on specific classes. This allows us to create two binary classification tasks: Class (T-shirt/top) vs. Class 1 (Trouser) and Class 5 (Sandal) vs. Class 9 (Ankle boot). For this setting, our simulation results in Figure 3 reveal that finite-degree Hermite model (I = 5) achieves nearly the same generalization error as the neural network. Therefore, our method allows for a direct examination of the neural network's activation function a via its equivalent Hermite model's 's activation function ôj Additionally, we discuss the effects of learning rate 1j Appendix D.4 suggesting potential improvements in generalization error when Σ x n and x n. Thus, extending our work to cover (1, 2] would an intriguing direction for future research.
generate
7 CONCLUSION
In this work, have explored the behavior of two-layer neural networks after single gradient step within the framework of the asymptotic proportional limit, specifically under the assumption of Gaussian mixture data. Our analysis provides comprehensive understanding of how structured data and feature learning jointly influence the generalization performance of these networks. We have established that conditional Gaussian model is equivalent to the two-layer neural network in terms of both training and generalization performance. Furthermore, we demonstrated that finite-degree polynomial model can effectively approximate this conditional Gaussian model, thereby showing that polynomial models can also perform equivalently to neural networks. We also highlighted po- tential avenues for future research, particularly the extension of the range of 3, which governs the joint scaling of feature learning and data spread. Our simulation results illustrate the impact of various properties within the Gaussian mixture data setting, reinforeing our theoretical findings. Im- portantly, applied our theoretical results to practical classification problem using the Fashion- MNIST dataset, where input images were generated through a cGAN, This work not only enhances understanding of neural networks in structured data contexts but also sets the stage for further investigations into their performance on realistic datasets.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We acknowledge that this work supported in part by the TÜBITAK-ARDEB 1001 program under project 124E063. We extend our gratitude to TÜBITAK for its support. S.D. is supported by an AI Fellowship provided by Koç University & iş Bank Artificial Intelligence (KUIS AI) Research Center and a PhD Scholarship (BİDEB 221 1) from TÜBITAK.
REFERENCES
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. In Advances in Neural Information Processing Systems, 2022.
Jimmy Ba. Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the pres- ence of diw-dimesioonal structure: A spiked random matrix perspective. Advances in Neural Information Processing Systems, 2023.
Jinho Baik, Gérard Ben Arous, and Sandrine Péché. Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. The Annals of Probability, 33(5):1643 1697, 2005.
Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. Advances in Neural Information Processing Systems, 2024.
David Bosch, Ashkan Panahi, and Babak Hassibi. Precise asymptotic analysis of deep random feature models. In The Thirty Sixth Annual Conference on Learning Theory, pp. 4132-4179. PMLR, 2023.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Info- gan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in Neural Information Processing Systems, 2016
Lucas Clarté, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. On double-descent in un- certainty quantification in overparametrized models. In International Conference on Artificial Intelligence and Statistics, pp. 7089-7125. PMLR, 2023.
Romain Couillet and Zhenyu Liao. Random matrix methods for machine learning. Cambridge
University Press, 2022.
Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue Lu, Lenka Zdeborova, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. In International Conference Machine Learning, pp. 9662-9695. PMLR, 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pp. 14522. PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, (giant) step at time. arXiv preprint arXiv: 18270, 2023a.
Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová. Univer- sality laws for gaussian mixtures in generalized linear models. In Advances in Neural Information Processing Systems, 2023b.
Samet Demir and Zafer Dogan. Random features outperform linear models: Effect of strong input- label correlation in spiked covariance data. arXiv preprini wrXii:2449.2025:1 2024.
Oussama Dhifallah and Yue M Lu. A precise performance analysis of learning with random features. arXiv preprint ar arXiv:2008.1 /904, 2020.
Elena Facco, Maria d' Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimen- sion of datasets by a minimal neighborhood information. Scientific reports, 721): 2017.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. Advances in Neural Information Processing Systems, 2019.

---

Published as a conference paper at ICLR 2025
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 2020.
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029) 1054, 2021.
Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure on learning in neural networks: The hidden manifold model. Physical Review X, 10(4):041044, 2020.
Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mézard, and Lenka Zde- borová. The gaussian equivalence of generative models for learning with shallow neural networks. In Mathematical and Scientific Machine Learning, pp. 426-471. PMLR, 2022.
lan Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http: / //www. deepl. ea rningbook.org org
Hamed Hassani and Adel Javanmard. The curse of overparametrization in adversarial training: Precise analysis of robust generalization for random features regression. The Annals of Statistics, 52(2):441-45 2024.
Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Trans. Inf. Theory, 693))1132-1194, Mar. 2023.
Hong Hu, Yue M Lu, and Theodor Misiakiewicz. Asymptotics of random feature regression beyond the linear scaling regime. arXiv preprint arXiv:2403.08160, 2024.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen-
eralization in neural networks. Advances in Neural Information Processing Systems, 2018.
lain M. Johnstone. On the distribution of the largest eigenvalue in principal components analysis.
The Annals of Statistics, 29(2):295 327, 2001.
Ganesh Ramachandra Kini and Christos Thrampoulidis. Phase transitions for one-vs-one and one- vs-all linear separability in multiclass gaussian mixtures. In ICASSP 2021-202 IEEE Interna- tional Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 2021.
IEEE,
Donghwan Lee, Behrad Moniri, Xinmeng Huang, Edgar Dobriban, and Hamed Hassani. Demys- tifying disagreemeent-on-the-iiie in high dimensions. In International Conference on Machine Learning. pp. 19053-19093. PMLR, 2023.
Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová. Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions. Advances Neural Information Processing Systems, 2021.
Xiaoyi Mai and Zhenyu Liao. High dimensional classification via regularized and unregularized empirical risk minimization: Precise error and optimal loss. arXiv preprint arXiv:1905.1377 2019.
Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Commun. Pure Appl. Math., 75(4))667-766. 2022.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two- layer neural networks. Proceedings of the National Academy of Sciences, 115633))76665-7777, 2018.
Francesca Mignacco, Florent Krzakala, Yue Lu, Pierfrancesco Urbani, and Lenka Zdeborova. The role of regularization in classification of high-dimensional noisy gaussian mixture. In Interna- tional conference on machine learning, pp. 6874-6883. PMLR, 2020.
Mehdi Mirza and Simon Osindero. DoXiv:1410 17844 2014.
Conditional generative adversarial nets. arXiv preprint
12

---

Published as a conference paper at ICLR 2025
Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non-linear feature learning with one gradient step in two-layer neural networks. In International Conference Machine Learning, pp. 36106-36159. PMLR, 2024.
Andrea Montanari and Basil N Saeed. Universality of empirical risk minimization. In Conference on Learning Theory, pp. 4310-4312. PMLR, 2022.
Alireza Mousavi-Hosseing Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient-based feature learning under structured data, Advances Neural Information Processing Systems, 2023.
Ryan O'Donnell Analysis of Boolean Functions. Cambridge University Press, 2014.
Jeffrey Pennington and Pratik Worah. Nonlinear random matrix theory for deep learning. In Ad-
vances in Neural Information Processing Systems, pp. 2637-2646, 2017.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in Neural Information Processing Systems, 2007.
Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference Machine Learning, pp. 3936-8947. PMLR, 2021.
Dominik Schröder, Hugo Cui, Daniil Dmitriev, and Bruno Loureiro. Deterministic equivalent and eтror universality of deep random features learning. In International Conference on Machine Learning, pp. 30285-30320. PMLR, 2023.
Mohamed Amine Seddik, Cosme Louart, Mohamed Tamaazousti, and Romain Couillet. Random matrix theory proves that deep learning representations of gan-data behave as gaussian mixtures. In International Conference on Machine Learning, pp. 8573-8582, PMLR, 2020.
Stefano Spigler, Mario Geiger, and Matthicu Wyart. Asymptotic learning curves of kernel methods: empirical data versus teacher-student paradigm. Journal of Statistical Mechanics: Theory and Experiment, 2020(12):12400) 2020.
Nilesh Tripuraneni, Ben Adlam, and Jeffrey Pennington. Covariate shift in high- random feature regression. arXiv preprint arXiv:2111.08234, 2021.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027, 2010.
Roman Vershynin. digh-dimensional probability: An introduction with applications in data science, volume 47. Cambridge University Press, 2018.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint 17Xiv:100.0777 2017.
Jacob Zavatone- Veth and Cengiz Pehlevan. Learning curves for deep structured gaussian feature models. Advances in Neural Information Processing Systems, 2023.
13

---

Published as a conference paper at ICLR 2025
NOTE TO READERS OF THE PROOFS: HIGH-PROBABILITY EVENTS
In the following sections, we leverage high-probability bounds to establish our results, focusing on the core components of the proofs while omitting detailed tail bounds for brevity. We provide appropriate references for readers seeking comprehensive derivations. We are confident that the proofs, as presented, are both rigorous and accessible, adhering to the standards of conference paper format. Furthermore, we plan to release an extended version of this work for future journal publication, which will include more detailed proofs and additional insights, thereby enriching the understanding of our findings and their implications in the field.
A SPIKE+BULK DECOMPOSITION OF THE GRADIENT
Following Ba et al.
2022).
we study the gradient G as follows:
G m
1
wÿ"
www (FXT))
o od((F))) X,
(18)
Vk
VI
hi
WyJ @ ( x x- 1
(ww"" (FF x,
x
myk
myk
mk
(19)
where we use the orthogonal decomposition: d'(z) hı + o'(''' and hi == 10 (z)). Here, the expectation is over = ~ NN 1) since for all E {I,...,kk k). J x has Gaussian distribu- tion (conditioned on component c) for fixed f; and we have Ex [f." 2 and E æ [f }) = 1
by (A.3) and (A.6) Furthermore, hı = (1) since (.) has bounded derivatives by (A.8)
Then, we have = we- for u := h1 w and v == X (T//mm/) while A == = C, which gives us G uod + A. Here, A matrix represents the spike structure of the gradient while A is called the bulk. Next, we show
A = Ô (k=*) with high probability by studying the norms of B and C,
where / : 1 3(1 a) 2 O. We start with bounding [|]
[B| <
@ Oo' III XII
(20)
mvk
diag w) jo (FX diag g) [X|,
(21)
mvk
w -00
و |
|
(22)
myk
for which control the norms the last line one-by-one follows. First, we get
|||| = Ō(k 1/2).
(23)
with high probability, due to hub-Gaussian tail bound (Vershynin 2018 Proposition 2.5.2) and (A.6) Next, we have
9 ~ = Ö
(24)
with high probability, due to Gaussian concentration of Lipschitz functions (Vershynin 2018 The- orem 5.2.2) and since we can equivalently write ū; = o. (2;) for 2; NoO0, C.) with C. (1) by
(A.3) A.5) while a. is a Lipschitz function by (A.7) Next, we get bounds on [[] and |X] as
F| (1), and
X 11 /2
(11/2))
(25)
with high probability, due to concentration of norms of aub-Gaussian matrices (Vershynin] 2018
Theorem 4,4.5). The bound for ||F||| is due to (A.6) and Tr(E) x k by (A.4) The effect of the mixture (2) is handled by considering X as concatenation of Gaussian matrices. Note that k. m can used interchangeably in the bounds due to (A.1)) Using (25). we get
| (FX ||| Ô =(()))))00))))1111-0 (x1)2+(1-1)/2)
(26)
14

---

Published as a conference paper at ICLR 2025
with high probability. Here, the rows of o' (FXT)" are independent sub- Gaussian vectors, which
follows from Gaussian concentration of Lipschitz functions Vershynim 2018
Theorem 5.2.2) and
the boundedness of derivatives d' by (A.8) Thus, we get lo (FX IIl S || polylog
X
with high probability using the concentration of the norm of matrix with independent sub- Gaussian
rows Vershynin 2010 Theorem 5.39 and Equation 5.26). Furthermore.
O(k ( by (A.2)
which makes
X
O (111/2(11-)/) with high probability due 25 and allow reach 26
Putting everything together, we get B = Ö (k: () with high probability.
Similarly, we focus on ||C|],
||C|| < mk
ww o oFF O (FX*) )| |*|||
(27)
diag(w))) (FX )daag( w or (FX
||
(28)
mk
//wloooll' of I ||||||||
(29)
mkk
for which find high probability bounds for the norms the last line one-by-one as follows. Note that bounds for |/// and || are found before in (23) and (25 respectively, Thus, we study the
remaining two terms starting with [o'(FXX ) follows:
|Io' 4 || (k:)
(30)
with high probability, which follows from [|o'(FX') X S 11o' (FX*) 1 where 1 denotes all ones matrix of dimensions k m. Here, h1 (1) as discussed at the beginning of the section and bound for ||o (FX")) is given in 26 For the last term, we have
I (FX = δ (k00-0)7 11/2)
(31)
with high probability, due to the Gaussian concentration of Lipschitz functions (Vershynin) 2018 Theorem 5.2.2). Note that [w] of Gaussian vector (Vershynin 2018 Theorem 3. and (A.6) Combining these bounds, we get
Ö(1) with high probability due to the concentration of the norm
||C|| = Ô (k ') with high probability.
Using the bounds on |B] and ||C|], we reach
A = B
C|I
<
B
+
||||
= Ô (K-9.
(32)
with high probability. Similarly, we also get
uul
|h wo
O(),,
and
||||||
X =ō (K-1/2)
(33)
mvF
with high probability using the found norm bounds, which completes proof.
B PROOF OF CONDITIONAL GAUSSIAN EQUIVALENCE (THEOREM 3
Here, we prove the conditional Gaussian equivalence under Gaussian mixtures data setting following the proof technique in Dandi et al. (2023a). We first provide the following lemma describing a conditional central limit theorem (CLT) in our setting.
Lemma 5 (Conditional CLT). For any Lipschitz function db VR, E Rdd++1
R2 - RR Vo E (1,...,C) and
lim 200
sup
E d (w ((x)){TT))
[c,tcc]
c, C 0, v T (())) m) - E
WES.. 1//
(34)
where Sk = {10 E R
| & C₁.
|| (2)k==)) for some Cı, Cz 0.
15

---

Published as a conference paper at ICLR 2025
Proof. Due to the decomposition of Fx (Lemma 2 we can equivalently study oaa + +F x instead of o (F To proceed, we define := F- (I, - r in order to focus
on the following equivalent features: o (a + Fž) for ă NOO, Then, we can define the following neuron-wise activation functions
Oalm (u) (ailke + b.u) Ba--((a.1) [or (asjm + b,ūi]
(35)
where b; = 1illl for i E (I,...,k). kj.
Now, we define / = [f.... fx1" where f: f./b, Thus, we focus on the feature mapping of the form Ji
(ff 2) which equiva- lent to random features mapping (Rahimi & Recht 2007 Hu & Lu 2023) with activation functions differing across neurons. The one-dimensional CLT for random features (Goldt et al., 2022 Hu & Lu 2023 Montanari & Saced (202)) holds even when the activation functions differ across neurons as observed by Dandi et al. (2023b a). Therefore, we check the assumptions used in showing the one-dimensional CLT for random features in Hu & Lu (203)) Here, the following events hold with high probability:
Jr dúg
-(1/11/2))
supsyett.....
1 (1),
which follow from the bounds for [|] and ||||| provided in Appendix A and the fact that I|,, - - < 1 T.T
O(1) by definition. Furthermore, Ews. (w) = Therefore, we can utilize the Lemma 2 in Hu & Lu (2023). Note that the odd activation function assumption, which is used to match covariances in Hu & Lu 20033) can be dropped here since (xx) and d(xx have exactly the same mean and covariance. Finally, to cover the effect of the second parameter of the test function 1/, we set fo 21/2 // Theorem 2 in Hu Lu (023)
D 1/2 g
and use
Lemma 5 useful for showing performance-wine equivalence of the two feature maps since it states that the conditional expectation of any test function for the original feature map o(x) is equal to that for the equivalent conditional feature map ò(æ; c, K%)) in the limit. To apply Lemma 5 our results about training and generalization error, we need to assume in VV Te in 55) to satisfy 2 w E S& with high probability, where Sy is defined in the lemma. Note that Sş is used in the proof of Lemma 5 when we utilize the CLT results from Hu & Lu (003).. One can show that v VT E S₁ holds with high probability similar to Lemma 17-18 and Lemma in Hu & Lu (003),, which omitted here. Alternatiell,, one can include w v S₁ as a constraint into the optimization objective of in 5). Either way, we continue assuming this condition holds.
Recall the original feature map and the equivalent conditional feature map at this point:
d(x) == o( (Fxx)
(36)
0(z; R.) == v(c, ne) L (c)K)))+ + #b K.)g.
(37)
where == =-= 1/2 0/c- T.ke + z-, and
v(c, : F. [f()x))))))))))()))
c, Kic
B ((
c, Kic
(38)
(c, = Cov (o(Fx)
c, 150
v (c, R.) g N0, In).
(39)
In the rest of the proof. we use the universality result for Gaussian mixtures in generalized linear models by Dandi et al. (2023b), which is an extension of the universality result by Montanari & Saeed (202)) to mixture settings. There are couple of points to be modified in order apply their results in our case. First, while Dandi et al. (2023b) assumes a conditional CLT (conditioned on component index c) in their Assumption 4, replace it with Lemma 5 which is another conditional CLT (conditioned on c and K. together). Furthermore, Dandi et al. (2023b) supposes bounded mean in their Assumption 2, which translates to ||v s..)
C. for some C 0. In our case, this
16

---

Published as a conference paper at ICLR 2025
assumption does not hold, which is the main challenge. To overcome this challenge, we study the conditional mean of the features v/C, and the demeaned features 6(x) - v(c, Ke) separately. About the demeaned features d(x) M(c, Re),, we have the following lemma similar to Lemma 20-21 in Dandi et al. (2023a).
Lemma 6. Under our assumptions, for a given Re;; C,
(i) the random vector ((x²) v(c, ke) is sub-Gaussian with sub-Gaussian norm independent of Re, e and dimensions k.
(ii) the matrix + each row of which is sample of <(x) vlc, Kie). satisfies P 1 2 < 2 2exp(-C2kk some C2,C2> 0.
Proof. Here, (i) follows from the Gaussian concentration of Lipschitz functions (Vershynin 2018 Theorem 5.2.2) and the boundedness of derivatives of a.
(ii) is due to (i and the concentration of the norms of matrices with independent sub-Gaussian rows Vershynin 2010 Theorem 5.39 and Equation 5.26).
Next, focus on vlc, Kc). To relax the assumption its norm, we have the following lemma similar to Lemma 22 in Dandi et al. (2023)).
Lemma 7. Suppose our assumptions. Let in as defined in 5). Define v(x)) == v(c, KG) where e and Ke are depending on as defined in Lemma 2 Then, the Jollowing holds with high probubility
m
1
w vo(i) = =(11)
(40)
m
Proof. First, recall that
i
arg min
V/i
1
w" o(xz))
+
+ 1
(41)
2m
Vk
w
Then, we have
1/
1
o0 c((::)
2
+ 2
A |1
S
1
Σ
ve Ö(1),
(42)
2m
2m
where the first step is due i being the optimal solution of |ya| = O1) with high probability as mentioned in Appendix A
41 while the second step is due to
This leads to
T
o(0.))
<
20 p(xx))
Yk:
(43)
2m
1=1
m
T o(Ix))
2
(8)
1
t o(z:))
1 m
(44)
2m
i=1
1
1
2 (16)
1
2T o(xx))
2
yf
(45)
Vi
m
m
1
:t o(xx))
2
(11)
Õ(1).
(46)
Vi
m
? -o v(c, v(c, Nio) ((
(bm) Õ(1),
(47)
m
where (i) is due to the Cauchy- Seawhrr inequality; (ii) follows from basic algebraic manipulation: (iii) is due to 1y:l O(1) with high probability; finally, (iv) follows from Lemma
17

---

Published as a conference paper at ICLR 2025
Now, utilize the universality of training enror for Gaussian mixtures (Dandi al. 2023b Theorem and Corollary 2) with the following modifications:
(i) Their Assumption 1 (Loss and regularization) holds as is in setting.
(ii) Their Assumption 2 (Boundedness and concentration) requires boundedness for [e, ko)ll- Thus, we need relax it mentioned before. The assumption used for free energy approximation which directly follows (unchanged) from MMntanari & Saeed (2022). The corresponding assumption Montanari & Saced (2022) is Assumption 5, which is used for their Lemma 5 and 6. Our Lemma (6) and ensures that Lemma and in Montanari & Saeed (2022) hold for the case of unbounded and variable means v(c, Ke)) across samples.
(iii) We have a modification to handle their Assumption 3 (Labels)). The assumption does not directly allow labels y depend considering ( æz] E g as the input of the generalized linear model in Dandi et al. (2023b) and constraining the last n parameters of the model to be O. Note that Lemma (conditional CLT) includes ç x as the second parameter of the test function, which makes the CLT valid for such an input. Thus, Lemma cooovrrss the dependence of labels to æ.
However, such dependence
be incorporated by
(iv) Their Assumption 4 (CLT) holds in our case due our Lemma (conditional CLT) and the law total expectation.
For the universality of the generalization error, we utilize Theorem 4 in Dandi et al. (023b), which requires the following additional assumption:
(A.9) Define perturbed optimization objective
Qm (s) == min WESp 2m i=1
Vi
1 Vk
uT (Fa:))
A | sG(w),
(48)
where s E and ((w) is the generalization error defined in when the second layer weight is w. Then, there exist s* O such that, for [-s",s" s"). the function 4m (s)
converges pointwise a function ((s) that is differentiable s 0.
Note that the additional assumption (A.9) corresponds to Assumption 5 in Dandi et al. (2023b),
argument. Furthermore, assumptions similar to (A.9) can be found in Hu & Lu (2023) (see their Assumption 9) and Montanari & Saeed| (2022) (see their Theorem 3). Finally, we observe that Assumption in
which is used prove the equivalence generalization errors using a convexity-
Dandi et al. (203b)) trivially holds in our case since we have single minimizer v in 5 and it can be simply found by ridge regression.
C PROOF OF THEOREM
In this section, we prove the equivalence of the two-layer neural network after training with one gradient step to the Hermite model. First of all, recall the bulk+structure decomposition of Fir for the case of the input e on e-th Gaussian,
FE!/2 (F (F+nA)=*
(49)
Fizi (Balk)
(Structure)
where z = 1/2 1/ ~ NO0, In) and we use the orthogonal decomposition: 2 Tote + x1, with
To = [v, Yo,1, Yo.2.... Yo.del, Ke TT"
and z- = (I, - -.- r.) -1r?
Then, instead of a F (o)e)) we focus on the following equivalent features: a aaa,. F2) where
(0 In) and F == FI P - -I?)). This can be equivalently written as follows for each neuron i:
(aije + bru),
(50)
where b; = |llll 0 and u, (f+ /b; N(O, 1) for all i E {1....,k) Furthermore, (wi, uz) is jointly Gaussian with [B(muy) f: f//bbbb)) for all i
18

---

Published as a conference paper at ICLR 2025
For the rest of this proof, our task is to show that the polynomial activation 0z (1(ajjn. mn + b,ug), defined in (16), has conditional mean and covariance that are equivalent to those of a (( + bzuz).
First, we start with the following lemma, allowing us to decompose (a + b,u,). Note that I 1/2 reduces the equivalent activation oy to the noisy linear activation (model) that was studied in prior work (Hu & Lu
2023 Moniri et al. 2024), so we focus on the case of 1/2 in our proof,
which corresponds to 2, i.e. beyond the linear case.
Lemma 8. If $ 3 < 3 < 7 then |a]]kal O PL and laajell' ((1/k+t9 hold with high probability, where = 1(1 - > O.
Proof. Recall that ds/me = f EC
Then, we have the following with high probability
lasjel1 s ft 2//2 In) + muzo") Tetic
+ (fftt)"Ta
(51)
ST EX// In)) || | ||To
+ ((S (Tt
(52)
(f) + +uuv)" (I!!² -In) + || |l) ||Tttell + |ff: ((f::)(Tat TTrell
(53)
(ftt)² (I!!² - -Tn)l + "ooo" //! - In
+ n [/luul>)) ||a ell + ((ftt)
Tcke
(54)
where we apply the triangle inequality and Cauchy- Schwarz inequality repeatedly to reach the last line. Then, we study each of the terms in the last line by using the following facts:
= ~~~N0. In) by definition 49
[Ta(TTT)) I r ((1) by assumption (A.4) and its definition (49)
f. ~ I. by assumption (A.6)
Tr(E) O(k) and Tr(E. - In) ((k((1-)) by assumptions (A.2) and (A.4)
k and n can used interchangeably since k/n E R by (A.1)
First, we observe that |T = Ir, (Tr.) follows from Hanson-Wrightt inequality (Vershynin 2018 Theorem 6.2.1).
Õ(1) holds with high probability, which
Then, we have o/l Ö(n--1) and nllu,eT E2!? II O(K--1) holding with high
probability as explained in the following. Here, |us| = w/rwal ((k-1/2) with high probability since wy (0, 1/k) by assumption (A.6) Furthermore, |||| Ö(((6(L-)--1))) with high probability by Lemma while (k and |/// I..Il (1.((1)) by (A.2)
Next, we get ((
In)Il Ö( with high probability. Here, first recall that f Si + A, then we have 7/l | S 1/| A and 7/|| | is vanishing by Lemma 1 Note that tighter bounds can be found about rows of A similar Lemma 12 in Dandi et al. (2023a), which is omitted
here. Therefore, we focus on !| ( IAI = ff (E!/ I,,) f. and utilize Hanson- Haighn-Wright inequality.
For the last term, we reach the high probability event (xx Ö(L--1) follows. Similar to the previous case, 17, does not affect our bound here. Therefore, we continue with T T (.) Ir z|. Considering that the entries of fi and the entries of z are sub-Gaussian, the product of any entry of f: with any entry of z is sub-exponential with the sub-exponential norm of term can be considered
Ck-1/2 for some C > 0. Then, (T.T..)
sum of zero-mean sub-exponential random variables. Therefore, we reach (T -Ir'' 2| Ō(k 112) with high probability by Bernstein's 's inequality (Vershynin 2018 Theorem 2.8.2).
Combining these together in (54). we reach [a:]=.,| ((k Using
1(1 - -) - 1 > 0, we
also get |am]xal' Ô (k=(1+)) with high probability.
19

---

Published as a conference paper at ICLR 2025
Since |az|e=! (() with high probability (Lemma we decompose a (asicc + b,u4) as
a(aine bb44)
(0)
(beu))a
R₂
(55)
No+s H, (,(0///))un) agme + Rr
(56)
2
ho+ aim
Hj((b,/b)ux) Ris
(57)
j!
(a.).. (Hb(b./b) Ro,
(58)
where R; = ((1/k++) is the remainder since |0/
-O(1/k++) with c > O. Here, we apply
Taylor's expansion in the first step, Hermite expansion (T)) in the second step, and basic algebraic manipulations for the rest of the steps.
Before continuing, it is beneficial to state the orthogonality of Hermite polynomials (Lemma 9 at this point, which is utilized in the following derivations of conditional mean and covariances of (.).
Lemma 9 (Orthogonality of Hermite polynomials). Let H, denote i-th probabilist's 's Hermite poly- nomial. If (a, b) is jointly Gaussian with zero mean, then the following holds
En, (a) = =(((E,,lab))' 8.g-
(59)
Proof. See O' O'Donnell 2014
Chapter 11.2) for the unit variance case EL
= E[b²] = 1). For an
extension that allows B [[a=] 7 1 while [E[2 1, see Demir & Dogan 2024 The same technique
can be used for further extension to allow E[a2] * 1 and #1.
Now, we are in the position to study conditional mean and covariance of (aan. + b,u.). Note that u; is the source of randomness in the rest of the derivations while aijs, and 6; are fixed since conditioned on (c, k.). First, we study the conditional mean as follows:
velc, K.) Eug (a(as)n + beu)) c, kc).
(60)
1 ba M., [1, (bb./b).) + Ri,
(61)
ho(asjo Bru, (Ho((br/b)us)) + R.
(62)
1 o!
Maagin + R
(63)
Eu, (aqik. + bruz) c, ke] R.
(64)
where R, = (11/+++)) is the remainder which may vary line to line. Here, the first step is due the decomposition in (58) while the second step follows from the orthogonality of Hermite polynomials (Lemma 9p. The third step is due E,,, (Hb function has the same first Hermite coefficients hy as those of the function. Using (64). we then reach the following norm bound on the difference of the conditional means
= Finally, the last step is because the a
Σ R = Ö (AIZZAA)
(o(c,K)) Blor -Blo,
Niol S
(65)
Next, we study the conditional cross-covariance between aagjo + beus) and 2//- Here, recall the relation between iN and - by definition: u, = (f+t)T=+ 2t /b.. Thus, (us, 2t) is jointly Gaussian
20

---

Published as a conference paper at ICLR 2025
with zero mean and B|u,2]] Bi,) /9//,, which leads to the following derivation
V (cg(c, (aijx + bzue)az c, mel,
(66)
he (asjm, ((lH, (bb//))u))=)) + Ris.
(67)
h (agh J (b., (b./b)) + REss
(68)
Mo (b. (o (b//b) Ray.
(69)
hot 1 ii (( + Ras
(70)
Mot1aiie (1t E(= + R
(71)
(a + bess)a c, Ke] + Ri.j-
(72)
where Ri, = (11kk+t) is the remainder which can change line to line. Here, we start by applying
the same steps we use for the derivation of (64). Then, we utilize E|u;z]] 1 (1a, }
/b; to case is
reach (70). For (71), we utilize the following figh-probabilitty event so that the o moved into the remainder R²
(73)
which follows from Lemma 8 E[(,, }]] Õ(1) by defimition, b E R$ by (A.8) and the following
gigh-probability event: lf(tt O(1/k'=)). Here, lff. S l + |n
اار Furthermore, lf =
((L/k2/2) since fig is sub-Gaussian with sub-Gaussian norm of C/ v (T(E) for some C' > 0 by
(A.6) and Tr(E) & by (A.4) Finally, |nA.,l S / (1/kk² by Lemma | which
completes our explanation about 773
To reach the last line (72) in our derivation of the cross-covariance, we use the fact that a and o have the same first / Hermite coefficients hy. Using 722 we get
L kc) ((a((Fx) c, c.ll < L ( k.)
Elo (Fx( I c. (c.
(74)
1 kc
R$ Ô
(75)
where we use the assumption that n/k E TR. (A.I)
Before continuing with the covariance between a (agje, + byu:) and (agjke + byuz), we have the following lemma about ff fú, which is helpful for the upcoming derivations on the covariance. Lemma 10. 1/2 < B < 1, then the following holds with high probability
max Esesc isujsk
ft J, bb8,y Ō(1/k'=8)
(76)
where b. = / //) Furthermore, 1- B = e1) by definition in Lemma g
Proof. Recall that f: = (In - -T (T) = - -T..T. (f: + n A.) Then,
fl f, b36d, = -I(fi nA,) 3d (In To T. + + - b26by].
(77)
polylog A: レー
(I, IT. ff, b 136
(78)
polylog k k1-9
1f - 66 + I...T. 'T f,
(79)
polylog k k1-9
(80)
21

---

Published as a conference paper at ICLR 2025
where the first step is due to triangle inequality, n ا < = )() by Lemma and
I| - TT T & 1++ [ T {()) by definition and assumption (A.4)
The second step follows from the triangle inequality. For the last step, we use the following two high-
probability events: (i) If f;- bb = Ô (k-1/2) and (ii) T.
fy 0 1/2).
We prove (i) as follows. For + j. we have fi., and f,.- as Sub-Gaussian random variables with sub- Gaussian norm bounded by v [T) for some C > 0 due to Vershynin 2018 Example 2.5.8) and
(A.6) Therefore. fi. is sub sub-exponential random variable with sub exponentplt norm bounded by TREy due to Vershynin 2018 Lemma 2.7.7). Since ff 5; Σ Ji,ndJ,s, we can apply Bernstein's 's inequality Vershynm 2018 Theorem 2.8.2). The case of = j can be derived similarly.
Note that Tr(E) x by (A.4)
Similarly, we show (ii) in the following. The key point is that f. is sampled independent of the fixed matrix T. by (A.6) For i J. we utilize Hanson- Hangon-Wright inequality (Vershynin 2018 Tr(T. -TT))TTT((()) [ T (1) by definition and (A.4) The case of + j follows similarly.
Theorem 6.2.1) with " T.(
O(1/k) and
Then, we focus on the conditional covariance between o(aj,, + b,u,) and aajjn, + b2,5). For i * (the off-diagonal entries of the covariance), we derive the covariance follows:
Ela(a. + b,u))o(ajin + byيي) I e, Kel.
(81)
Σ
1
R. (.. (.,((./b))u) R
Σ ss
8
1 ( + R₁)
(82)
s
(st)² he (ailk (Ha
+ Rag.
(83)
fr160) + Righ
(84)
(F71118) + Rag.
hatoaain
(85)
o!
Notoogine (FF + Rass
hoto
Σ
(86)
= =Elōi buu))r((ajn + bbuz)) kcl Rig:
(87)
where RA, R..j - Ō(1/k++) are the remainder terms that may vary from line to line. For the first step, we use the decomposition in (58) while the second and third steps are due to the orthogonality of Hermite polynomials (Lemma 9 and the fact that (u, uy) is jointly Gaussian with Elu,uy] = ft' (b,bb, To reach (85). we use the tigh-probability event: If. fz/b²/ = O(1/k**s), which
follows from Lemma 10 To arrive at (86), we utilize the following gigh-probability event
F.. f. = O(L/kttt), for o 0+s21,
(88)
which again follows from Lemma g Lemma and b R+ by (A.8) Finally, the last line (87) is due to the fact that oG function has the same first / Hermite coefficients hy as those of a function.
For i = j (the diagonal entries of the covariance), we apply Taylor's expansion to (a + b,u)) around bu, as follows:
(aqp. + beu,)² = albu,)² + Riçi
(89)
where Ra,; = 20 (1) (M)a)) +(b=b)uu) some in-between bu: and 0gike +b2u. Here, we see
that [t] = Ō(1) with high probability by definition, which makes lor (tto (t) Õ(1) with high prob- ability since a has bounded derivatives by Assumption (A.8) Then, since |a.je. Ö(1 ((1/k(1+6)/t)
22

---

Published as a conference paper at ICLR 2025
by Lemma 8 and 1b, bj = O(1/R((+)))) by Lemma 10 and u, ~ NOO,, 1), we can conclude that RA, O(1 (11/k(+)))() holds with high probability. Using this result, we reach
Blo(a + b,uu)² rc] = (bo, I c, Ke] + R
(90)
= El + b,uu)² Ki] + R4,55
(91)
where R2, = O(1/k(+)))) is again the remainder that can differ from line to line. To get the final line, we use [bo = (lor(bu,; (h7)² Σド /z/ /Gt by the definition of o (16).
Therefore, we reach the following norm bound on the difference of the covariances for and ōJ
(Fx))))))(Zx) c, re] - Blo Kol|
S max Ra, +
R
(92)
= - -11/mmnn(00000))))
(93)
where we first separate diagonal and off- parts in the first line while the last line follows from (II) and (87) This completes our task of showing the conditional mean and covariances of on( [ ) are equivalent to those of (z)) when conditioned on (c, ke).
Conditional Gaussian Equivalence under Equivalent Means and Covariances From the previous derivations, we have
E (o((a)))
c, kic
E (F(()ar) E (F 4
c. Kic
o(1),
(94)
B of (x 1T{ e, ci
c, ne
o(1),
(95)
Cov (a(Fa)) (, Cov (oi(F
c, Ke)
o(1),
(96)
from 665 75 and
93 respectively, under the conditions of Theorem
In the rest of this section, we describe why (94)- are enough to utilize the conditional Gaussian equivalence (Theorem 3 for the proof of Theorem 4 First of all, note that we only the exact mean and exact covariances in the proof of the conditional CLT (Lemma 5) while exact equality
conditions can be directly replaced with (44 (96) for the rest of our proof for Theorem 3
First, we observe that (44 and (95) do not cause any significant change in comparison to the exact equality since ab in Lemma Lipschitz function. To show this, we define the following remainder
r(c, Ke) = (D(c, Ke) (.)) + ( (c, 1 (c, xz
(97)
where (e, Ke) and S (c, k.)) are defined as
(c, k.)) = E (Fx) c, Kin
( k.) = E [z (z C, C,
(98)
Then, we have the following about the test function ub in Lemma 5
| (woT (d(x) + r(c, (x),,,.o)))))) - <-x² - vb (w" ((x), E x)
< L :ū r(c, ke)],
(99)
for some L > O. which follows from the fact that ab is a Lipschitz function. Since vo E S₁ in Lemma 5 we have || (1). Thus, we just need to show that [|r ke)]l is vanishing, which would then imply that replacing d(x) with d(x) + r(c, Ke) has a vanishing effect on the test function 99). To bound the norm of []r(c, sc))l., we first apply the triangle inequality follows:
[Ir xa)l < [[0(c k.) x2)ll + |1 (c, k.)) (e,, xz)))= ll,
(100)
where the first term is vanishing by (94) while the second term deserves further elaboration. Here, the second term has vanishing bound with high probability by Hanson-Wright Wright inequality (Ver- shynin 20018, Theorem 6) since x is a random vector with independent mean zero sub- Gaussinn coordinates and have 95 This completes explanation about the impact of (94) and (5
23

---

Published as a conference paper at ICLR 2025
Now, focus the covariance (96). At this point, recall that our proof of Lemma reduces to
Theorem 2 and Lemma in Hu & Lu (2023). To cover the effect of (96). we replace Lemma 5 in
Hu & Lu (023) with our bound (96) for the covariance. Note that the vanishing bound in (96) is enough to replace the bound in Hu Lu (2023) for our case since we just need a vanishing bound in their Lemma 2 and Theorem 2. For the sake of completeness, we would like mention the
location where the bound for the covariance is utilized in Hu & Lu (2023): it is applied just above equation (79) in the proof of their Lemma which then resulted in equation (68) in their Lemma and equation (67) in their Theorem 2 after some derivation.
We showed that the conditional Gaussian equivalence Theorem applicable with (94 96 Therefore, we utilize Theorem with (94)-(96), which completes our proof for Theorem
D ADDITIONAL SIMULATION RESULTS
Here, we provide additional simulation results about (i) landscape of generalization error with re- spect and (Figure 4). (ii) impact of 3 for different a values (Figure (5). (iii) training errors (Figure () and (iv) the effect learning rate for TaNSST classification (Figure 7). These are beneficial for understanding various aspects of our setting that we could not cover in the main text due to page limitations.
D.1
HEAT MAPS GENERALIZATION ERRORS WITH RESPECT TO AND
100
1.0
0.150
1.0
0.8
0.125
Senerall Paliiaattoo 00000
50
difference percentage
0.6
0.100
25
0.075
0.4
0.050
0.2
0.025
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
(a) Generalization error of the neural network. The data spread scales as |Σ|| x (110), while the learning rate scales as 7 x aa
(b) Percentage difference between the generaliza- tion errors of the neural network and the equiva- lent linear model (
Figure 4: Generalization performance of the neural network with respect to and B as heat maps. We set n = 400, m 500, and used the ReLU activation function (a = a. = ReLU). The number of classes is C = 2 with dimensions d1 d2 1. The parameters 01, 02.1 8(811a) and = 1 employed control the model's 's behavior. & (22,1+721)//(0 + (/2 is used to ensure high alignment between and the data covariance. The results presented are averages from 20 Monte Carlo simulations.
Figure 4a Ilustrates the overall landscape of the generalization error with respect to a and (8, high- lighting the significance of the strength parameter 8, which governs the scale of the learning rate y combined with the data spread ||>||| Our findings indicate that as we increase 3 and/or decrease generalization performance improves, underscoring the importance of structured data relative to first-layer learning. Additionally, Figure 4b demonstrates that surpassing linear models' perfor- mance important benchmark in the literature (Cui et al. 2024 requires increasing / while sometimes only reducing a suffices within certain ranges (3 E 0.44, 0.7]). Overall, these observa- tions highlight the profound impact of the structure of data learning outcomes.
24

---

Published as a conference paper at ICLR 2025
D.2
IMPACT OF 3 FOR DIFFERENT VALUES
gnerairation stree
0.15
goveraitation error
0.15
Sigmoid
inth
generalizsion erver
0.16
0.10
0.10
Nouni Netwolk Hermite Model
0.12
0.05
6.05
0.00
0.08
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.5
01.5
1,0
0.0
0.2
0.4
0.6
0.8
1.0
(a) 0
(b) 1/2
(c) =1
Figure 5: Impact of 3 for various a values in the setting of Figure 1
In Figure 5 we observe how the generalization curve with respect to 3 gets affected by a. Overall, see that increase in 3 is beneficial for improving the generalization enror, while the value shapes the generalization curve with respect to B3.
D.3 TRAINING FRRORS
Nesral Network
ReUU
Hemmie Model
0.025
Signaid
uabling oIN
8.000
0.015
A/m0
Figure 6: Equivalence of the training errors with respect to k/m in the setting of Figure la
In the main text, we focus on the generalization error for our simulation results since it is of more interest in comparison to the training error. However, it is important note that our results hold for the training error as well, which may be of particular interest on its own. Therefore, here, we provide simulation result about the training errors Figure We observe that the training error decreases as k/ //m (governing the number of hidden neurons) increases, which is expected. More significantly, see that the training errors of two-layer neural networks match those of the equivalent Hermite model for all activation functions in the simulation.
D.4 EFFECT OF LEARNING RATE BEYOND 3 E [0, 1] FOR FASHION-MNIST CLASSIFICATION
1.0
Neural Network
Rel.I
1.0
geteraliaation ennor
0.8
Sigmold
0.6
0.6
0.4
8.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
4//m
log(n)/kog(s)
(a)
Σ and 1)
Impact of for Σ>|
(8 1)
(B E (1, 2] and k/m 1)
Figure 7: Impact of the learning rate for Fashion classification in the setting of Figure Ra.
25

---

Published as a conference paper at ICLR 2025
While result in Figure 3 for Fashion-MNIS, classification is interesting its question the impact of the learning rate 7) in the setting of the figure. Here, we answer the aforemen- simulation result in Figure 7 Note that Figure 77 is the same Figure
tioned question with 3/. which is given again for the sake of side-to-side comparison with 7b. In Figure (7), we can see n which aligns with
that the generalization performance gets better as approach towards the predictions for the case of isotropic data. However, we lose the equivalence since & > 1. Thus, the corresponding Hermite model is dropped from the plot. To study the result (b), one needs to extend our main results such that they cover / E (1, 2], which is left to future work.
(a) Class O (T-shir/top)
(b) Class (Trouser)
(c) Class 5 (Sandal)
(d) Class 9 (Ankle boot)
Figure 8: Examples of input images after the pre-processing for the result in Figure
E DETAILS FOR THE FASHION-MNIST SIMULATIONS
In this section, we detail our numerical simulation for Fashion-MNIST classification. We begin by training a conditional Generative Adversarial Network (cGAN) using the Fashion-MNIST dataset, following the implementation of InfoGAN (Chen et al. 2016) outlined by Dandi et al. (2023b). The cGAN is trained for 50 epochs with a batch size of 64, utilizing the Adam optimizer with learning rate of 0.0002 and 3 parameters set to (0.5, 0.999). The model is optimized using binary cross-entropy (BCE) loss for both the generator and discriminator. This structured approach allows effectively generate class-conditioned samples, facilitating robust classification tasks while adhering to established training protocols.
Next, preprocess the input images generated by the trained cGAN model to align with sumptions outlined in (A.2) (A.4) Each sample is flattened into a vector in IR", where n = 28 28 784. Before preprocessing,
calculate the mean and covariance for each class using Monte Carlo estimation with million samples per class. We then compute the ratio f² == Tr(E,) 1 Σg) to adjust the scale of the second class samples, ensuring that Tr(E,) Tr as required (A.3) The preprocessing involves demeaning samples from both classes and scaling the second class by 7, resulting μ1 = P22 and Tr(E,) Tr(E2). To further satisfy assumption
(A.4) we introduce noise to the samples using the formula æ := / /n/Tr(Ei] + where ~ (N, and & de- notes the preprocessed sample. The multiplier (n/TT(EE) controls the signal-to-noiss ratio (SNR), consistent with our data assumptions. Labels are assigned as +1 for the first class and for the second. This completes dure are shown in Figure g illustrating clear distinguishability between classes. Additionally, we use b = 1 for calculating the Hermite coefficients in our Fashion-MNIST classification simulations.
preprocessing steps, and sample images generated through this proce-
F EXTENSION TO NON-ZERO MEANS FOR MIXTURE COMPONENTS
We assume zero mean for the mixture components (μt 0) in Assumption A.33 to simplify our proofs. However, it is important to recognize that the mean pee of each Gaussian component functions similarly to the spikes Ye,i in the covariance structure described by equation 9 To extend our analysis include non-zero means, we assume the existence constant C such that || < C| This allows us to incorporate non-zero means p. * by adding the term Fj.. into the structure outlined in Lemma 2
In our proofs, we would need to bound additional terms arising from these non- omit these details here for brevity. Instead, present numerical evidence demonstrating that
non-zero zero means; however,
26

---

Published as a conference paper at ICLR 2025
results remain valid even when accounting for non-zero means. Figure 9 illlssrates numerical simulations for Fashion-MNIST classification with non-zero means (ft 7 0). The results indicate that the Hermite model performs equivalently to the neural networks, for the inputs of each class. Notably, the generalization errors in the case with non-zero means (as shown in Figure are lower compared to those observed with means (Figure(3). suggesting that classification becomes easier when different means are introduced. This finding further underscores the robustness of our framework and its applicability more complex data scenarios.
with non-zero means
0.20
Nearal
Neural Network
ReLU
Tlemii Medel
Signsid
lan
generiiaationn crree
0,3
Hennile Model
Sigmeid
gonestiration
0.19
0.2
0.05
0.00
0.0
(a) Class 0 vs. Class
(b) Class Class 9
(T-shirt/top vs. Trouser)
(Sandal vs. Ankle boot)
Figure 9: Impact of non-zero means Fashion-MNIST classification for |Σ|
and v 1. m 500, = le - 5, and = 5. This is the same setting as Figure 3 however inputs for each class demeaned separately; instead, all data is demeaned together to ensure E[x] while maintaining He * Comparison with results from Figure 3 highlights the effects non-zero means on generalization error.
27


Reviews:
[{"review_id": "gMwkhnQWR6", "reviewer": "Reviewer_uU4P", "summary": "This paper proposes a theoretical framework to study the training and generalization errors of a two-layer neural network under a more realistic assumption of Gaussian mixture data with additional low-dimensional structures and shows that a finite-degree polynomial model serves as an equivalent performance model, which can simplify the analysis of neural networks. Lastly, the paper compares the model performance using the “Hermite model” with Relu, Tanh, and Sigmoind functions in simulation data and real data generated with generation models and studies the impact of data spreading.", "questions": "No", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Theoretical framework of a two-layer neural network for more realistic data assumption;\n2. A finite-degree polynomial model serves as an equivalent performance model.", "weaknesses": "Finding the weakness of the paper is beyond my current ability.", "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": "v5QhSjQm3W", "reviewer": "Reviewer_SJBV", "summary": "1. This paper establishes a theoretical framework for characterizing the training and generalization errors of two-layer neural networks under Gaussian mixture data, while previous work mainly focuses on Gaussian data. \n\n2. This paper demonstrates that a finite-degree polynomial model serves as an equivalent performance model, simplifying the analysis of neural networks under the Gaussian mixture data assumption. \n\n3. Through extensive simulations, including classification tasks on Fashion-MNIST, this paper validates the findings and highlights the significant impact of data structure on learning outcomes.", "questions": "What prevents the analysis from covering cases with non-zero mean Gaussian components?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The framework provides a novel lens for understanding neural networks trained on Gaussian mixture data with structured covariances, marking a clear advancement over previous works that lacked these structural considerations.\n\n2. The paper includes a compelling mix of theoretical insights and numerical experiments, with applications to Fashion-MNIST classification. These experiments effectively validate the theoretical claims and underscore the significant impact of data structure on neural network learning outcomes.", "weaknesses": "1. The primary weakness lies in Assumption (A.3), which requires all Gaussian components in the Gaussian mixture model to be zero-centered. In my view, such a Gaussian mixture model does not introduce significant challenges, thus providing limited theoretical advancement in understanding neural network training. Specifically, $ E_{x\\sim P}f = E_{x\\sim N(0,\\sigma_1^2)}f + E_{x\\sim N(0,\\sigma_2^2)}f + \\cdots $; since $ x $ is always multiplied by the trainable weights $w$, a linear transformation of a zero-mean Gaussian can be reduced to the isometric case. This observation leads me to believe this analysis could be a straightforward extension of existing works that consider the isometric case.\n\n\n2. The theoretical insights intended to guide the practical training of neural networks are unclear. Specifically, the connection between the proposed data model and real-world applications needs to be clarified. The motivation for extending beyond existing Gaussian distribution assumptions to study this particular type of data model also remains ambiguous.\n\n3. The numerical experiments would be more compelling if implemented on larger-scale datasets, e.g., ImageNet.", "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": "D3W9aFEQgs", "reviewer": "Reviewer_rEU8", "summary": "This manuscript investigates the training and generalization performance of two-layer neural networks after a single gradient descent step, under a structured covariate model given by a Gaussian mixture with finite-rank + identity covariance. The analysis is conducted in the proportional high-dimensional limit where the covariate dimension, hidden-layer width, and sample size diverge at constant rate. Leveraging recent progress in the Gaussian universality literature, the authors show that:\n\n- After a single-gradient step on the first layer weights, the gradient admits a decomposition in terms of a rank-1 spike and a correction which is asymptotically small in operator norm.\n- A conditional Gaussian equivalence theorem showing that the training and generalization errors after a single-gradient step are asymptotically equivalent to conditional Gaussian covariate model. This equivalence can be further written in terms of a Hermite decomposition of the non-linear activation.\n\nFinally, the authors provide empirical validation via simulations on Fashion-MNIST classification tasks, supporting the theoretical predictions and highlighting how data structure influences learning outcomes.", "questions": "- **[Q1]**: In L199-200:\n> *\"Notably, (Ba et al., 2022) established that a learning rate $\\eta$ exceeding $\\tilde{O}(\\sqrt{k}) is necessary to surpass the performance of the linear model represented in equation (8).\"*\n\nIn my understanding, (Ba et al., 2022) showed that the maximal learning rate is rather *sufficient* to achieve to surpass the performance of the best linear model. Indeed, (Moniri et al., 2023) showed that an intermediate scaling leads to partially learning non-linear components of the activation.\n\n- **[Q2]**: You assume that all the means of the GMM are zero *\"to simplify the analysis\"* (Assumption A.3). While on a high-level I can see how the means can play a similar role to the spikes in the covariance (as you mention in the discussion that follows), this is an important simplification which should be better discussed. In particular, do you expect to see any effect of multi-modality which are not captured by the spiked covariance? If yes, it would be good to add a discussion on this. If not, please provide quantitive evidence of that.\n\n- **[Q3]**: From the discussion of the conditional Gaussian equivalence (cGET, Theorem 3), it is not very clear how this result differs from the cGET proved by Dandi et al., (2023a,b). Only by going to Appendix B, one sees a more detailed discussion of how Thm 3 builds on previous results from Hu and Lu (2023) and Dandi et al. (2023b), and hence what are the key differences. I believe that it would be beneficial to the reader to have this discussion summarized in the main - in particular what results from (Hu and Lu 2023; Dandi et al., 2023b) are used and clearly stressing how Thm 3 extends / differs from them.\n\n- **[Q4]**: Another point which I missed is a closer discussion on the relationship between the regime studied here and the ones from (Ba et al., 2022; Dandi et al., 2023a; Moniri et al., 2023; Cui et al., 2024). In particular, these works can be split in three different regimes:\n\n    1. A regime where Gaussian universality holds and there is effectively no feature learning (performance is lower-bounded by a linear method) (Ba et al., 2022)\n    2. An intermediate regime where there is partial feature learning (Moniri et al., 2023).\n    3. Maximal update / critical step size: (Dandi et al., 2023a, Cui et al., 2024)\n\nMy understanding is that your Lemma/Theorem 1-3 holds covers all the regimes, while Theorem 4 is for an intermediate regime similar to Moniri et al., 2023. Is this correct? Would be good to add a more explicit discussion on that point. \n\n- **[Q4]**: Appendix A, L724-725 is confusing. Do you mean $f_{i}^{T}x$ is Gaussian distributed conditioned on the class?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This work is inscribed in the recent literature seeking to quantify feature after a single but aggressive gradient step in two-layer neural networks (Damien et al., 2022; Ba et al., 2022; Dandi et al., 2023a; Moniri et al., 2023; Cui et al., 2024). The main difference with respect to these works is in the data model, which includes a multi-modal model with structured covariance. The manuscript is well-written.", "weaknesses": "The main weakness is the lack of a close comparison with the related literature on the technical part of the work. Indeed, the authors do a good job in situating their contribution to the related literature in the general discussion (introduction, related works), but this falls short in the following. This is particularly problematic since the technical results heavily builds on previous work (in particular, the proofs in the appendix are based on reductions to previous results of Ba et al., 2022 and Dandi et al., 2023a,b). In the *Questions* below, I give a few pointers where I missed a discussion.", "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": "JbsEZYaZSh", "reviewer": "Reviewer_Veqf", "summary": "The authors consider the problem of characterizing the hidden-layer features of a two layer neural network, after one gradient step on the first layer weights, in the asymptotic limit of comparably large dimension, width and number of samples. They consider (mixtures of) Gaussians with spiked covariance structure. They establish a conditional Gaussian equivalent map for the features, and further establish that the activation function can be replaced by a polynomial function with asymptotically equivalent test and train errors. Different aspects of the model (learning rate -to- data spread ratio, alignment between the spikes, cluster weights) are varied and their effect on the test error discussed. Many numerical experiments are performed to support or illustrate the discussion.", "questions": "-  In (17), does $\\hat{F}$ still refer to (12), namely the first layer weights after a gradient step on the _original_ network with activation $\\sigma$, and not $\\hat{\\sigma}_l$? In other words, does the $\\sigma\\leftrightarrow \\hat{\\sigma}_l$ equivalence hold only when training the readout, or does it hold for the full training procedure, i.e. including when performing the gradient step on the first layer? This is not clear from the discussion in the main text.\n\n- Can the authors add a discussion of why considering zero means (assumption A.3) is necessary, or if it can be relaxed to accommodate non-zero means ?  \n\n- l. 410 : \"This equivalence significantly simplifies our analysis by transforming the nonlinear activation into a polynomial form, thereby enhancing the tractability of the model while maintaining its performance characteristics.\" Can the authors be more specific on which point of the analysis this equivalence aids? This point, if any, is not sufficiently discussed in the current state of the manuscript.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The paper is well written; results are sufficiently discussed and put in context. Section 6 provides ample intuition and discussion on the effect of the different parameters of the problem, allowing to illustrate and support the experimental results. Although I did not carefully check the proof, technical results seem sound.\n\nTechnically, the analysis of conditional Gaussian universality to structured covariances is interesting and relevant to extend the stream of recent related works to more realistic data distributions. In particular, the analysis of the joint scaling of $\\eta, \\lVert \\Sigma\\lVert$ is interesting, and provides a insightful picture of the tradeoff. I appreciated section 6, and in particular the discussion of the effect of the alignment between different covariance spikes.", "weaknesses": "I do have a number of clarifying questions, which I list in the section below. Overall, I believe the work is sound and will be of interest to the community working on those topics. I do not give a higher score, because the paper leaves open rather natural questions which are naturally attached to its setting, beyond its current scope (namely establish Gaussian universality for the considered distribution). In particular :\n\n- Can a description of the features spectrum be given in the case of a single-index target function, in the likeliness of Moniri et al. (2023)? Do the spikes persist, or are they altered by the presence of the structure in the covariances ?\n\n- Is it possible to reach a tight asymptotic expression for the test and train errors?\n\nIn regards to this, I believe the title would gain to be slightly toned down slightly to be better representative of the scope of this paper, which the abstract accurately on the other hand reflects. I still stress however that the contribution is interesting and sound.", "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)