Low Rank Gradients and Where to Find Them
Authors: Rishi Sonthalia
Michael Murray
Guido Montufar
Venue: NeurIPS 2025 poster ()
Keywords: One step gradient
low-rank gradients
two-layer neural networks
spiked covariance
regularization
high dimensional probability
Primary area: theory

Abstract: This paper investigates low-rank structure in the gradients of the training loss for two-layer neural networks while relaxing the usual isotropy assumptions on the training data and parameters. We consider a spiked data model in which the bulk can be anisotropic and ill-conditioned, we do not require independent data and weight matrices and we also analyze both the mean-field and neural-tangent-kernel scalings. We show that the gradient with respect to the input weights is approximately low rank and is dominated by two rank-one terms: one aligned with the bulk data–residue, and another aligned with the rank one spike in the input data. We characterize how properties of the training data, the scaling regime and the activation function govern the balance between these two components. Additionally, we also demonstrate that standard regularizers, such as weight decay, input noise and Jacobian penalties, also selectively modulate these components. Experiments on synthetic and real data corroborate our theoretical predictions.

Low Rank Gradients and Where to Find Them
Rishi Sonthalia"
Michael Murray
Department of Mathematics Boston College
Department of Mathematics
Bath University mjm253@bath. ac.
rishi. wwwtsssaaacta..eeeoo.o
Guido Montúfar
Mathematics and Statistics & Data Science
University of California, Los Angeles
montuf ar@math. ucla. edu
Abstract
This paper investigates low-rank structure in the gradients of the training loss for two-layer neural networks while relaxing the usual isotropy assumptions on the training data and parameters. We consider spiked data model in which the bulk can be anisotropic and ill-conditioned, we do not require independent data and weight matrices and also analyze both the mean-fieel and neural-tane courel-tangent- kernel scalings. We show that the gradient with respect to the input weights is approximately low rank and is dominated by two rank-one terms: one aligned with the bulk data-residue, and another aligned with the rank one spike in the input data. We characterize how properties of the training data, the scaling regime and the activation function govern the balance between these two components. Additionally, we also demonstrate that standard regularizers, such as weight decay, input noise and Jacobian penalties, also selectively modulate these components. Experiments on synthetic and real data corroborate our theoretical predictions.
1
Introduction
Feature learning is critical driver behind the success of deep learning. Despite this, a theoretical characterization of it remains elusive. In order to drive understanding, a line of research [3] [2] has emerged studying two-layer networks whose inner weights are trained or updated via one step of gradient descent. In this context, feature learning can be characterized through the emergence of a low-rank structure the network weights. In particular, when the weights of many of the neurons align in a predominant direction, the matrix of weights becomes approximately low rank. Moreover, Ba et al. 3 proved that a ridge estimator trained on such features can outperform random feature models and other kernel methods. However, these prior investigations require idcalized conditions, for example isotropic data or weights, which diverge from real-world scenarios where data typically exhibits anisotropy in this context have also been underexplored.
14
28 45)
ill-conditioned covariance. In addition, the effects of regularization
This paper addresses two questions: I) how do low-rank gradient phenomena arise and behave under more general conditions of anisotropy and ill-conditioning? and 2) what impuct do common regularizers have feuture learning in this context? Our analysis accommodates spiked data with an anisotropic ill-conditioned bulk. This allows us to explore the effect of the size of the data spike, controlled by parameter 2 0, as well spectral decay profiles of the bulk, controlled by parameter a Z O. Our central finding is that the gradient the inner-layer weights is generically well approximated by rank-two matrix. This structure arises from the interplay of two primary rank-one components: S, driven by the input bulk and target residue, and Sz, driven by the leading
"Corresponding Author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

eigenvector of the data covariance. The relative prominence of these components, and consequently the direction of feature learning, is determined by the interplay of data properties, the scale of the network parametrization, the choice of loss and activation function as well as the use of regularization. We corroborate our theoretical findings with experiments on both synthetic data (Sections 3.2
and 4) and real data (MNIST, CIFAR-10 embeddings) A summary of our key contributions is:
We provide theoretical framework (Section 3 Theorems I.I and 3.2 characterizing the low-rank structure of the gradient under significantly relaxed assumptions on data and weight matrices (anisotropy. ill-conditioning; Section 2
Generalized Theory of Low-Rank Gradients:
Identification of a Dominant Rank-Two Structure: We show (Theorems 3.1 and 3.2) that the
gradient is often better approximated by a rank-two matrix than the rank-one structures identified in prior specialized settings. We provide conditions under which each of these components dominates. Modulation by Activation Function and Regularization: We show how activation functions and common regularizers selectively modulate the components of the gradient. We reveal that ReI.U can suppress the contribution from the residue S₁ (Section (.2), while input noise and Jacobian penalty can promote the residue component and data spike component (Section 4 respectively.
Mean Field (MF) versus Neural Tangent Kernel (NTK) scaling: We demonstrate differences in dominant spike alignments, S ~ Xgy in MF vs. S X in NTK, at initialization (Section and the subsequent impact during training.
1.1 Related work
Low rank gradients in two layer networks: For a Mean Field (MF) like regime, prior work has shown that the gradient is approximately rank one [ [3), which results in alignment between the leading cigenvector of the hidden feature kernel with the target [45]. Dandi et al. [14] showed that to learn & directions, opposed Under an NTK scaling Moniri et al. (28) showed that a learning rate which grows with the sample size introduces multiple rank-one components in the hidden feature kernel. However, these results rely on well coliicondt input data and weight matrices. To ameliorate this issue a number of works have also incorporated ill-conditioning via a spiked covariance models (A I + n² q²qq")) with 4 found dominant rank-one gradients aligned with the data spike q (if aligned with the target 33_ and
single direction, we need high sample complexity (n 2(dd))
single-index targets (o,, ((3., x))) [4, [2]... Ba et al.
/2), enabling efficient learning, whereas Mousavi- et al. [29] showed that gradient flow might yield weights nearly orthogonal to even under seemingly favorable conditions (v 1, 3. = q). Our work continues this line of research. providing results for more general anisotropic and ill-condtiioned data and weight matrices as well the effects of regularization.
Understanding the spectral evolution of the network's 's weight and features matrices, particularly the 'bulk' components beyond dominant spikes, remains challenging. While significant progress has been
32 34 44] and after a single step of gradient descent (12), the dynamics over longer timescales are complex.
made in characterizing spectra at initialization [[ 6 16 30
Convergence to low-rank weights: While our analysis focuses on the gradient updates that drive learning, related studies investigate the implicit bias of gradient-based optimization towards low-rank solutions [17 21 27 33 4T]. Our findings complement this body of work by characterizing the generically dominant rank-two structure within the gradient updates themselves, providing insight into the mechanisms potentially driving this convergence.
2 Setup and Assumptions
In this section, we provide the technical details required for analysis. A summary of notation and discussion of examples of when the assumptions hold can be found in Table Appendix [B] We consider shallow networks with d input dimensions, m hidden neurons, and n training data points.
Assumption 1 (Proportional scaling). Let 11, th2 E E>0 be fixed constants. We consider m, as functions of d such that n/d + v/= < 1 and m/d dd + 1/2 as + 0c.
Data: We consider random input data J; E d for i E [m]. sampled i.i.d. These are stored row-wise in matrix X e Rndd For each I,, the corresponding label is y E IR, and labels stored as y E IR".
All code is available at the anonymous Github repository:
https://ithabb.com/sssootal///o Rank Graaiii

---

Assumption 2 (Input features distribution). Ler E E Rdxd for which there exists an a 2 0 such that the k-th eigenvalue satisfies Xx (E) = ftr- =ll....,,d d. Let q E s0-1 and define Ç =" for some z 0. We assume each input data point x; is sampled i.i.d. from a multivariate Gaussian distribation N0, Σ), where the full covariance E E R given by Σ E + 52""""
Assumption 2 models bulk component via Σ and spike component via q (magnitude () and allows general forms of ill-conditioning with {x()) - 0 if a > and Y (2) - 30 if > 0. This generalizes typical data distribution assumptions like isotropic Gaussian (Σ Ia) or uniform on
a sphere K [14 28 31 45), anisotropic data with bounded condition number [15] 18. divergent
largest eigenvalue and bounded smallest eigenvalue 4 22 24 29 40), or bounded largest eigenvalue
and decaying smallest eigenvalue [5 43
Network: We consider a two-layer neural network with input-output map f: Rd - I defined as f(x) 7ma" (W()) x)) E R.
(I) Here, W = [wi.... ...,, Wm, 1 E RRaxd is the matrix of inner (first-layer) weights, wy E Id is the weight vector for the j-th hidden neuron and a E I is the vector of outer (second-layer) weights. The activation function
IR- + I. is applied glement-wise to the preactivations Wx. The parameter 7'm E 1x is a non-trainable scaling constant that depends on the network width m. Assumption 3 (Network parameters). We assume the following for W, a and T/m
1. Outer weights: Elements aj are sampled ii.d. from ([nffor((-111)))
2. Inner weights: Rows w, of W have unit length, Wi E @d-1
3. Scaling parameter: 7/m = ((1/ (NTK scaling) or 7/m- = (11/m) (MF scaling).
The assumption is standard. The assumption W (unit-norm rows) relaxes typical literature requirements (c.g., isotropic Gaussian or uniformly spherical wy). This allows modeling anisotropic weights. possibly dependent on X, to analyze updates throughout training, not just at initialization.
m/)m) [2] 20 23 25), associated with lazy training where inner weights vary little [10 25), and MF (%m ~ 1/m), asso- ciated with feature learning 99, 26
The scaling parameter "/m defines two common regimes: NTK (fmm ~ 1/ /))
36 39]. These scalings yield different initial output variances (Var( (x)) @(1) in NTK vs. (1) in MF), impacting dynamics.
Assumption (Activation function). The activation function o : R + R satisfies:
1. Smoothness: o' and o", first and second derivatives of a, exist almost everywhere on R.
2. Lipschitznes:: and o' are L-Lipschitz for some conslunt L>> 0.
3. Non-trivial expected derivative: For [(0, Σ) (per Assumption 2 and W (per Assumption 3)
let μú = [o' [w] x)] (expectation over x for a given W.). We assume p. = (11) for all j. Let
P1 ........, Mml "T
We define ('(2z); o' o'( 2) Pdj
Common activation functions, such Sigmoid, Tanh, ELU [1], Swish B3) Softplus, satisfy the Smoothness and Lipschitzness. We note that the derivative of ReLU is not Lipschitz. The condition M4 = 2(1) (non-vanishing expected derivative) is satisfied by ELU, Swish, and Softplus generically, and for Sigmoid and Tanh as long as w/ Σwi- O(1). See Section B. for more details.
Parameter update via GD: Let é: IR R + Rzo be a function which measures the loss between label and prediction. With f defined as per Equation [ we define the loss given dataset (X, v) = (I;; Me ())e)())) with respect to the inner-layer weights W as L(W) = Σn {(f(x)) Mi) + XR(W). R denotes a regularization function (c.g., the 2-norm R(W) = I|W and À E Rz0 is the regularization parameter. We consider update to W arising from one step of GD, W + W - L y), where > denotes the step size. We define the residue vector as r - [0((f(=)) ())
(xx))) ER"
(2)
To motivate this terminology, consider that for the Mean Squared Error (MSE) loss, r corresponds to the vector of residues [f(x)) yele. More generally. for many losses r can typically be interpreted as the component of the targets not captured by the predictions of the model (see Section
Proposition 2.1 (Gradient of the loss). If Asssumptio 4 holds and R is differentiable, then
G = Vwr = =mm" XT [(raa) (oo((XW)) + XV R(W)- E Rdxmm
exists for almost every W in gmxd
3

---

07
5.5
2.4
200
<25
630
500
1000
200
420
600
850
2000
200
400
602
300
1020
(a) Tanh, BCE, isotropic W. (b) Swish, Hinge,
non- (c) Softplus. BCE, isotropic W.
isotropic W.
Figure 1: Singular value distribution of the gradient G for varying activation, loss and and weight distribution. Red, and blue lines show the singular value Si, and S2 respectively. (a) the rows of W are i.idd. uniformly random the unit sphere, we denote this W Ws. In (b) and (c) then W = Ws + n 1/"1q". where W is then normalized. The following parameters are constant across all experiments: a = 0, 7m v (NTK) by a triple index model, see Appendix both isotropie and non-isotropic W. For [0.25, 0.5), the gradient is approximately rank two.
750, d 1000, m 1250. The targets y are given C For v/ < 0.25, single residue-aligned spike is seen for
For our results hold we require the following technical assumption on the residues.
Assumption 5 (Residue concentration). Under the proportional scaling regime (Assumption 7p. with probability 1 - o(1) over the training data (X, y). the residue r satisfies
|| |||
log vn
o
We emphasize that Assumption 5s mild condition: ensures that no single component of the residue vector disproportionately dominates its overall Eg norm. Such a condition typically holds if the residues 7; are i.i.d. subgaussian random variables. See Appendix B.3 for more discussion.
Our analysis also depends on the alignment between the residue r and specific structural components of the input data X. From Assumption /
we have the following decomposition of the input features,
= XB + Xs XB 5zq" I
where XB has rows sampled i.i.d. from A E), NN(0, 7), and is a unit vector. We note that for sufficiently large Ç, is approximately the principal eigenvector of X One of the inputs to our analysis will be the degree of alignment between the residue vector r and the spike component of the input data. The projection of the residue onto the principal eigenvector of X X is a natural statistic interest and has been considered in prior works 20 38). In Appendix B. we provide estimates for 192 scenarios.
Assumption (Residue alignment). With probability 1 - 1-0(1), Vnllr = = (d-3)/2)
3 Spiked Data Leads to Low-Rank Gradient
In this section we analyze the role of spiked data in shaping the gradient with no explicit regularization, = 0. We demonstrate that for spiked data covariance the gradient either approximately rank one or rank two, depending primarily on the size of the spike. To demonstrate this we define the following three rank-one matrices:
Residue Spike: S = Y' (XTr) (aat)), n
Data Spike: S₂ = T'm q [=T(((r))) jooo))ooooo (XW"))
n
ZTr Interpolant: S12 : 7m q (aou)).
We remark that S; is studied in 3 and S2 is analogous to the gradient update in [4. The matrix S12 interpolates between the two: in particular, S12 and have the same right singular vector and S12 and S2 have the same left singular vector. Hence S₁ + S12 and S12 + S2 are both rank one.

---

1.3
020
50
200
611
400
200
600
250
180
(50
150
120
102
400
(a) ReLU
(b) Swish
ELU
(d) Softplus
Figure 2: ReLU suppresses the residue spike (S1) compared to smooth activations. Fixed parameters:
1/8, a = 5/9, = 750, 1000, and m 1250.
3.1 Smalr-to-Moderate Spike E (0,0.5)): C² Activations
The key contribution of this section is Theorem 3.1 which characterizes the approximate low rank structure of the gradient for small-to-moderatt spike sizes. We also note that Theorem 3.1 generalizes [3] Proposition 2] by covering broader range of covariance structures, loss functions and initialization scalings In the small spike setting, // E [0, 1/4), the gradient approximately rank and aligns with the residue plus interpolant S₁ + S12- By contrast, in the moderate spike setting, / E [1/4.1/2] the gradient becomes rank two. We empirically verify these our theoretical results (Figures D across range of activation and loss functions under the NTK scaling.
Theorem 3.1 (Gradient approximation). Suppose Assumptions) are independent, and a is a C2 function. Define G
3
4
6 are satisfied, X and W S2 Then, for all v, a Rz0.
5
S1
S12
S1 S12112
O (1W]lenn(
|| S1 S12 S2lle O (IWIleen** -
(4)
Vmmmm 1
/mmm///lrll |
with probability o(1) ) as d, n1, m co. Moreover, if < then with the same probability
||1||| = Ω
|S2l12
I1(= of (XX3)))))
n²
||
log n W/le/
EE
=
logn
|| |
(5)
|
n & log n W1l
Ω6n u 4)$
||122|2
s ~Onn" -(("-2t9) Z14).
Ω
(6)
|1'|
S
Observe that for < 1/4, if W log = o(n then G is approximately equal to the rank-one matrix S₁ + S12. Further, > 2v + a then the gradient dominated by S₁ and the spike is aligned with the data-residue term XTr. However, if B < 2y then the gradient term dominated by S12, which is aligned with the data spike q. addition, for E [1/4, 1/2),
||W]|2 log = o(n oini-me - and
10 (XW [[ To (XWT))
(
n" = w
log
(7)
then the gradient approximately the rank-two matrix S + S12 + S2. Note this distinct from prior works 3
4 14
45
where the gradient is only ever approximately rank one.
3.2
Small-to-Moderaat Spike: ReLU Activation
Theorem 3.1 requires the activation to C2, As detailed in the proof, this is needed to estab- lish that | (XWT) < O(|]W]|n***)) Indeed, when ||]]2 = @(1) and v < 1/ we have I|o' (XW' = o(n), which is key for S₁ to separate from the bulk spectrum. However, ReLU is not C2. To understand, the effect using ReLU we provide Proposition 3.1
a, and the row of W are i.i.d. from the unit sphere, then with probability 1 - o(1) we have that o' (XWT) sign (2g) sign(W
Proposition 3.1 (ReLU gradient). 2f
we see that for ReLU the operator norm of o' (XWT) is (n). This is significant increase compared to the o(n) scaling for C² activations and suggests that the norm of
From Proposition
3.1
Thhe result 3 Proposition 2] requires
0, isotropic data and MSE loss with MF scaling.

---

E and S2 are larger for ReLU. The increased size of E, S2 results in the relative suppression of the contribution of S₁ and an enhancement of the contribution of S2 the spectrum of the gradient. We empirically verify this phenomenon in Figure 2 where we compare ReI.U to its C² activations ELUU, Swish, and Softplus. We see that, for ReL.U the relative residue contribution (S1) significantly smaller when compared with its smooth approximations.
Remark 1 (Convolutional filters inherit the rank-two gradient). A ID valid convolution with stride I and filter R can be written as a two-layer network with a sparse, weight-tied matrix W Rxxd whose nonzeros are shifted copies of w, where m d k + 1. Treating the entries of W as independent parameters yields the gradient
ØL/oW E Rmxd x By Theorems 3.1 3.2 G admits
a decomposition G u( (u())T= + u (v())
E with at most two dominant rank-on terms and small bulk E. Weight tying maps G to the true filter gradient as follows:
aL Owe
m
Σ ==1
Gr
1= k.
Hence, letting 5.00
(1) ddii+1 for j = 1,2,
VoL u 5(1) + (2)) 7(2) (error),
the convolutional filter gradient lies in subspace of dimension most two, upto small error term.
3.3 Large Spike (v 2 0.5): Non C2 Activations and Dependence between W and X
The preceding analysis focused on v < 0.5. For large data spikes (v 2 0.5), we note that the C² smoothness of the activation function and independence between W and X are longer required.
Theorem 3.2 (Large data-spike gradient approximation). Suppose Assumptions 7
and we have
are satisfied, and define EL_ = - S12 S2. Then, with probability - o(1)
|ELll2 01), V mm |r|
| |2
n""
| E2
St ( n" [I( 0' (XW
Ω
log
logn 1lo' (XWT)]
(8)
Note this a generalization of 4 which required alignment between the targets y and the spike q. Theorem 3.2 shows that if > 3. or if Equation holds, then the gradient approximately rank one. contrast to the v < case, this rank-one gradient aligns closely with the data spike plus interpolant S12 + S2 rather than the residue S1. This is empirically verified in Figure 3 for non-C2 activations ReLU, as well dependent and independent W and X.
30/00
12600
1000
20100
5000
120
142
(a) ReLU, Isotropic
(b) Swish, Ws + /ng1" T
(c) Sigmoid, Wsl0
(d) ELU, W
Figure Singular value distributions of the gradient G under various activation functions and weight matrix initializations and structures, with a large data spike y = 3/4. Wss denotes the random matrix with rows drawn mutually i.i.d. uniformly from the unit sphere. The rows of W s are uniform on the sphere and orthogonal to g. All weight matrices are subsequently normalized to have unit norm rows. Fixed parameters: bulk decay exponent a 0, n 750, d = 1000, m 1250, NTK-like scaling (Tm /m), MSE loss, and triple-index model targets.
3.4 Impact of the Scale Parameter: MF vs. NTK Scaling
We consider the implications of our results for the two scaling regimes and highlight three important distinctions. As with prior work, we consider the large step-size regime. Specifically, we use step

---

Aver dot product
Average cot product
X Xgly f(X))
K
Xg (D0)
MF Scaling
(b) NTK Scaling
Figure 4: Empirical alignment (normalized inner product) of the top singular vector of the gradient G with XEy. X gg and for data from single-index model y = Sigmoid( (x)) + noise. We use isotropic X, ReLU activation, and MSE loss. We average over 500 samples of W, X, y. The error bars are the 25th and 75th percentile.
(a) MF scaling (notice small range)
(b) NTK scaling (oscillatory)
(c) Mean principal angle in degrees between the NTK and MF.
Figure 5: Evolution of the gradient direction and weight matrix during training under GD with Weight Normalization (WN) for the MF and NTK scalings. Fixed parameters are v = 0, = while using the Sigmoid activation function and the MSE loss. Plots (a) and (b) show the alignment (normalized inner product) between the leading left singular vector of the initial gradient Go (epoch 0) and that of G, (epoch (). Plot (c) shows the mean principal angle between the weight matrices learned under the MF and NTK scalings with identical initialization and training data.
size of ym To avoid exploding gradients deploy Weight Normalization (WN) [37]. We limit our focus the MSE loss. See Appendix Doo discussion of which assumptions hold during training.
Alignment initialization: residue versus target vector y. Recall from Theorem 3.1 that in the small spike regime the gradient is dominated by Sj. Further, for the MF scaling the residue approximately equal the target y, while for the NTK scaling the residue can be quite distinct from This implies the alignment of the gradient may differ significantly depending on whether an MF or NTK scaling is used. Suppose y = sigmoid(w" x) + E. then Figure presents the normalized inner products between the leading left singular vector of G and three candidate directions X X X and w. For the MF scaling, we see that the gradient's dominant direction aligns well with X X BV, consistent with Theorem 3.1 and 3 For the NTK scaling, consistent with Theorem 3.1 the gradient exhibits strong alignment with Xzrr. This differs notably from both 7 and the w alignment directions predicted 28 which believe be erroneous.
2) Stability of the gradient during early training. Let G, denote the gradient after iterations of GD. In F plot the alignment between the leading left singular vector of Go and subsequent leading left singular vectors of G, under both MF and NTK scalings. The following is quite striking: the dominant gradient direction under the MF scaling remains stable throughout training while for the NTK scaling it evolves significantly. This leads to a divergence in the trajectories of the weight matrix with identical initialization and training data.
Towards explaining this, suppose the conditions Theorem 3.1 hold at least approximately up to some iteration / < T' Then under an MF scaling the gradient is approximated by rank-one matrix whose left singular vector is nearly constant X x XHyy Therefore it remains stable over number of iterations. If the NTK scaling used instead, then S is proportional Xfr, & X and the gradient depends the residuals which evolve throughout training.
7

---

1.0
08
08
Resioue-X X
0.5
Resiove-K XXy
Dasa
Dasa
0.4
0.4
0.2-
0.2
0.7
0.2
20
40
63
83
100
20
65
83
300
20
40
63
03
200
Epoch
Epoch
Epoch
(a) Small V: v 0.125
(b) Medium v 0.4375
(c) Large v v 0.75
Figure 6: Evolution of the alignment of the leading left singular vector of G2 with data spike q and
residue (XTr)) during training. Fixed parameters: MF scaling, Tanh activation, MSE loss, a = 0.
3) Phase transitions both by epoch and data spike size. Figure 6 we observe the evolution the alignment of the gradient versus the data spike and the residue under the MF scaling. Moving from small to large spike sizes we observe a transition in the gradient alignment from the residue Xgr to the data spike q. We remark that this is predicted by Theorem 3.1 and Theorem 3.2 at initialization. Of particular interest is the middle spike size setting, where we witness phase transition during training of the gradient alignment from residue data spike. Interestingly, this transition is not discernible from the training loss, which smoothly decays during training. We only pause to highlight this interesting phenomenon here and leave more thorough analysis future work.
Effect of Regularization
We analyze three common regularization techniques: lg weight decay, isotropic input noise, and Jacobian (or gradient) penalization. We investigate how each technique influences the relative magnitudes of the residue-aligned spike S₁ and the data-aligned spike S2. For what follows let G(o)
denote the un-regularized gradient matrix derived in Proposition 2.
l2 weight decay. Adding the term || Il to the loss function modifies the gradient to G())
( XW. Proposition 4.1 implies that if W||W||2 o(mm))) it cannot suppress S or S2. however, if A||W|| ||2 = w( w(V Vmmmn") then it suppresses both spikes.
Proposition 4.1. Given Assumptions
and 1f||lll = O(vū), then with probability (1) have that || 2| S Ol)mmm)) S Olymymn* )), and ||2l|2 O(ymmmm'
7
Isotropic Gaussian imput noise. This regularization technique involves adding independent isotropic Gaussian noise & (O, () each input x; without changing the corresponding labels yr- 7 showed that training with input noise is equivalent under certain conditions to adding a Tikhonov regularizer the loss, often related to Σ ll |/ (x;)1l More recent work [46] connects adding isotropic noise to the data to controlling the trace of the Hessian of the loss function.
Let us define x/ = + &. This changes the input data distribution, effectively modifying the bulk covariance from Σ to Σ' = E + T² Consequently, derived quantities such the residue vector r', the alignment parameter 3¹, the gradient components S/, S2, S/2, the error term E', and the effective bulk spectral decay a' are denoted with primes.
Proposition 4.2 (Isotropic Gaussian noise). Assume the setup of Assumptions 7 2 with independent X and W. Assume a satisfies Assumption foro the noisy data X'. Additionally, suppose the modified residues satisfy rf (11) with probability 1 o(1), and Assumption 6 holds for r' with scaling parameter B'. 7 nP and [|o' (X'WT) ||2 = o(n), then with high probability:
||
| lz
|| |
2 w(1),
Onn"-8),
< -o(n"-9-4)
|
| 1|2
analyzes the effect of input noise. It indicates that the residue spike S remains prominent relative to the error term E'. Conversely, if the noise sufficiently strong, the data spike components Sy and S/2 become suppressed relative to E'. Intuitively, adding noise with variance 7ي nº increases the variance of the bulk data component. This boosts the overall scale of terms involving (X')². Simultancously, the added noise tends to make the pre- areaactivations W X"
Proposition 4.2

---

200
65
1000
186
500
85
2000
200
45
200
800
(a) 7 ReLU sup-(0) (b)
0.25. The (c) 0. Residue spike (d) 100. Data spike dominant.
presses the residue spike. residue spike re-appears. is dominant.
Figure 7: Effect of regularization. Panels (a), (b) are for isotropic Gaussian noise. Parameters: 750, d 1000, m 1250, V = 1/8, activation, MSE loss. Panels (c), (d) are for Jacobian norm penalization. As À increases, the size of S1 does change, size of the bulk Eg grows, and the size of the data spike S3 grows. Parameters: NTK, // = 3/8, a = O, Sigmoid and MSE, and triple-index model targets.
8/9 (for original data). triple-index targets, ReLU
more isotropie, which can reduce the operator norm ||o' (X'WT) [|2 relative its Frobenius norm, potentially limiting the growth rate of | E2||,, || This predicted relative enhancement of S₁ and suppression of S2 is verified empirically. As discussed in Section (cf. Proposition B.Ip. ReLU can hinder residue spike Sj. However, Figure [ouure/]shore shows that with small amount of input noise 7 = 0.25, an initially suppressed Si re-emerges, while S is diminished relative to S/ and the bulk.
Jacobian penalization. Another form of regularization penalizes the sensitivity of the network output to changes in the inner weights. We consider the Jacobian penalty Lreg = X 22. Vi=:llöw f (x;))] To analyze this effect of Lreg on the gradient, we derive the gradient of Lreg with respect to W.
Proposition 4.3 (Gradient penalty). Let Diag( (n:{((=::1??) be the diagonal matrix, whose entries are | 2 If is twice differentiable, then
Lreg
Xmm (o' d"Wxx") Diag([I=:1{)xx:
n
The gradient of the regularizer factorizes into data-aligned rank-one spike S3 and error Eg:
S3
n -=
Eg
X V ( (o (o'(X"") o o (XW"))
Proposition 4.4. Given Assumptions I fraction of the entries of of (((WX) @
3
and 6 If |||l||2 = @(Vn). a < and constant are bounded away from 0, then
A n 20
1-3a
Z mm
||E2ll2 2 A (n:-033
30
S
IF> 1 + a then we have that asymptotically the residue spike does not escape the bulk for any 2 2 == If J. we see that increasing X suppresses the residue spike. For the data spike, we have that AS will grow grows. Hence this enhances the data spike. We empirically verify that increasing kills the residue spike while promoting the data spike (Figure 7
Real-Data validation. The identified low-rank spike-plus plus-bulk gradient structure and the discussed regularization effects are observable in two standard vision datasets MNIST and CIFAR 10. For MNIST, estimate v 0.784
1/2 and the data is highly ill-conditioned, suggesting a large effective a. Theorem 3.2 predicts a gradient dominated by data-aligned components (Panel (c) of Figure 8). Adding isotropic Gaussian noise with x² 100 (Panel (d)) suppresses the original data- aligned spike and enhances the residue-aligned spike Sj. consistent with the analysis in Section 4 For CIFAR-10 we use a pretrained ResNet-18 (on ImageNet) extract 512 12-dimmnnion embedding. We estimate v Re 0.3572 < 1/2 and a RJ 0.6. For these parameters Theorem 3.1 suggests SJ (residue-alige) alignedd) can be prominent. Panel (a) of Tigure&sshow dominant S₁. Applying Jacobian regularization with À = 10" (Panel (b)) suppresses S and promotes data-aligned spike (akin to S2), consistent with the behavior analyzed for Jacobian penalization in Section
5 Conclusion
This work shows that in two-layer neural networks, the hidden-layer gradient is approximately rank-two, driven by data-residual (S₁) and data-spike (S₂) components connected by an interpolant

---

CIFAR,
(b) CIFAR, A= 10
(c) MNIST, 0²
(d) MNIST, 0² 100
Figure 8: Gradient singular value spectra real datasets. Each panel displays the singular values of the gradient matrix G under the specified conditions.
(2) We show that activation function choice, scaling, and regularization can result qualitatively different gradients. In particular, we have the following rule of thumb for the number of spikes.
Gradient Sradient-spike rule-of thumb:
Which spike dominates initialization?
S (residue spike)
+
2μ m a, a} or Large isotropic input noise
{
(i) 2v > min 1, 33, or (ii) Strong Jacobian penalty, or (tii) ReLU and 2v > 1 1-a
S12 S2 S3 (data spike)
If none of the above holds, both spikes remain, and the gradient is typically rank-two.
The coexistence and interplay of the two spike components offer nuanced understanding of the gradient. We believe that the saiduee-aligned part propels the network towards fitting the current errors for the specific task, while the tata-aligned part reflects the network's 's adaptation to or influence by the inherent structure and biases present the input data distribution. This dual influence provides potential mechanism for reconciling how networks can be both task-specific and data-adaptive. This is an interesting avenue for future work.
Acknowledgments
GM has been supported in part by the NSF in NSF CCF-2212520 and NSF DMS RMS-2145630. GM has also been supported by DARPA in AIQ project HR001 12520014, DFG in SPP 2298 (FoDL) project 464109215, and BMFTR in DAAD project 57616814 (SECAI).
References
[1] Ben Adlam, Jake A. Levinson, and Jeffrey Pennington. A random matrix perspective on mixtures of nonlinearities in high dimensions. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 3434-3457. PMLR, 28-30 Mar 2022. URL ptttps://proceediii mlr
press/v151/as add1//dllmm2222 html
[2] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analy- sis of optimization and generalization for overparameterized two-layer neural networks. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Pro- ceedings of Machine Learning Research, pages 322-332. PMLR, 09-15 Jun 2019. URL
mlr www.s/pt//aaaeeeeeeeeeeep arroa/a19 html
[3] 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. JtRp:httpp://oopeeevviee het//forum2i-aakddueeoe
[4] Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in presence of low-dimensional structure: A spiked random matrix perspective. In seveth- Conference on Neural Information Processing Systems, 2023. URL https::tps.. h/ap::s://ppeeviiee pet//forum?idd-HIAAACO LACCHDW
10

---

[5] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(48):300033300777 2020.
[6] Lucas Benigni and Sandrine Péché. Eigenvalue distribution of some nonlinear models of random matrices. Electronic Journal of Probability, 26(none):1 37, 2021. URL https
7
[7] Chris Bishop. Training with noise is equivalent to Tikhonov regularization. Neural computa- tion, 7(1):108-116, 16,, 1995.
[8] Chen Cheng and Andrea Montanari. Dimension free ridge regression. The Annals of Statistics, 2266):287--2912, 2024.
[9] Lénaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neu- ral Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL http:::/proceddggg..oo neurip cc/ cc/paper_pp. fuaperffless ptt.le///ppprr/////1iiee pdf
alaf 41I13666645500077999316666----60
[10)
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in Neural Information Processing Systems, volume 32. Curran Associates. Inc., 2019. URL httpptt::/ppoceeddd
neurdd cc/ cccpaper_ra. f files/ pappr//011//t/ell hipppe__llle//pppp//////////11
pdf
[11] Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (ELUs). arXiv: 1.000729, 2015. Published as a conference paper ICLR 2016.
[12] 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 Proceedings of the 41st International Conference Machine Learning. volume 235 of Proceedings of Machine Learning Research, pages 9662-9699 PMLR, 21-27 Jul 2024. URL httpp:://rroceedingg. mlr press/ htsp//235////// cui24d. html
[13] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. Conference on Learning Theory, pages 54522. PMLR, 2022.
[14] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two- layer neural networks learn, one (giant) step at time. Journal of Machine Learning Research, 25(349):1 65, 2024. URL http:/// wttp:oogp/pp//////5p/5 223//3- 23- 1543. html
[15] Edgar Dobriban and Stefan Wager. High-dimensionall asymptotics of prediction: Ridge regres- sion and classification. The Annals of Statistics, 46(1):247-279, 2018.
[16] Zhou Fan and Zhichao Wang. Spectra of the conjugate kernel and neural tangent kernel for linear- width neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 7710-7721. Curran Associates, Inc., 2020. htt p/tpceedd/// neurips cct paper wtt.a///wwp.....ooooooooooooooo0
Paper pdff
[17] Spencer Frei, Gal Vardi, Peter Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky ReLU networks trained on high-dimensional data. In The Eleventh International Conference on Learn- ing Representations, 2023. URL
[18] Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. Annals of statistics, 50(2):949, 2022.
[19] Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions on Information Theory,
2022.
[20] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, page 8580-8589, Red Hook, NY, USA, 2018 Curran Associates Inc

---

[21] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019. URL aattps://openrevit
[22] Chinmaya Kausik, Kashvi Srivastava, and Rishi Sonthalia. Double descent and overfitting under noisy inputs and distribution shift for linear denoisers. Transactions on Machine Learning Re- search, 2024. ISSN 2835-8856. URL hattps://pperrviieee wttpiiewwweeeoooeeeeeeee wtt.://ooorii.icccc.a...........
[23] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https::/pproceednnggs http:///pppoocddinnnn.nee....pp. neurips. wwtpaepcccpppppppppeee.......
wwtpppppprtieeeeeeeee pdf
[24] Xinyue Li and Rishi Sonthalia. Least squares regression exhibit mnder-parameteriieed double descent. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL pattps://openrevieet wet/forummid=gghhnntte
[25] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. On the lincarity of large non-linear models: when and why the tangent kernel is constant. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829544
[26] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018. doi: 10.1073/pnas. pet.ppa.....691111111111 1806579115. URL hattps://wwe www.wwpww..wppp.wwwww......w.... pnas org/doi/
abs/10. 1073/ /tt////////.../////// pnas. 1806579115
[27] Hancheng Min, Enrique Mallada, and Rene Vidal. Early neuron alignment two-layer ReLU networks with small initialization. In The Twelfth International Conference Learning Representations, 2024. URL htttp:://pperrviieee
[28] 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, 2024. URL https 2ttg//bb//311..... 2tt.//2210..00911111
//arxiv.or/ wwwxvv.orr//b//231000
[29] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient-based feature learning under structured data. In Thirty Thirth-evennth Processing Systems, 2023. URL /ttp::ts:///peeeeieee nww/ ?id Fod
seventh Conference on Neural Information
[30] S. Péché. A note on the Dennington-Wordt distribution. Electronic Communications in Probability, 24(none):1-7, 2019. doi: 10.1214/19-ECP26.. URL https:://aoi.org/1 2214/19-ECP20
[31] Jeffrey Pennington and Yasaman Bahri. Geometry of neural network loss surfaces via random matrix theory. In Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 2798-2806. PMLR, 06-11 11 Aug 2017. URL htRL/tttt/////rrocddingg
mlr. http://ww........ee//ee///////// ptwppess//////eeemngnnooo......p
[32] Jeffrey Pennington and Pratik Worah. Nonlinear random matrix theory for deep learn- ing. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL WtR.uttpps//pprcceddini neurips.cnmm httpppppr_ffleee_ ppapprrffiees/ppppp///////111111 patpp/pp/eeeeereeeeee ofaa0fooealdd0000ccce20000022..3 Pappee pdf
[33] Mary Phuong and Christoph H Lampert. The inductive bias of ReLU networks on orthogonally separable data. In International Conference on Learning Representations, 2021. URL https
7 //openreview /
[34] Vanessa Piccolo and Dominik Schröder. Analysis one-hidden-layer neural networks via the resolvent method. In Advances in Neural Information Processing Systems, volume 34, pages 5225-5235. Curran Associates, Inc., 2021. URI /ttptttttt//////oeeeddeeo
neurips............e............ cc/
pdf
12

---

[35] Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv: 171000511 2017.
[36] Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.,
2018. URL httpss:/prroceedingss neurips. cc/ http:pper.ffllesspper///211/... (1955544aac63777777444f000088888 Paper. pdf
[37] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In Advances Neural Information Processing Sys- lems, volume 29, Curran Associates, Inc., 2016. URL URL /ttp//tttp:://prccedddie neurips. cc/ pppperr-df]
[38] Andrew M Saxe, James McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:/312.614 2013.
[39] Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: central limit theorem. Stochastic Processes and their Applications, (3203):1800-1822 2020. ISSN 0304-4149. doi: https://doi..rr///0110////////. 2019......0..000000 0000.0. 003. URL StLhttps: /ttpss:/wwwx //www
sciencedirect. cwt.coo//ccc///o httpa////c.aa/aaacat///////////.
[40] Rishi Sonthalia and Raj Rao Nadakuditi. Training data size induced double descent for de- noising feedforward neural networks and the role of training noise. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL http:://pppnreeee.eeeeeeee/eeeee
FdMWtpVT11
[41] Nadav Timor, Gal Vardi, and Ohad Shamir. Implicit regularization towards rank minimization in relu networks. In International Conference on Algorithmic Learning Theory, pages 1429-1459, PMLR, 2023.
[42] Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices, page 210-268. Cambridge University Press, 2012.
[43] Yutong Wang, Rishi Sonthalia, and Wei Hu. Near-interpolators: Rapid norm growth and the trade-off between interpolation and generalization. International Conference on Artificial Intelligence and Statistics, pages 4483 4491. PMLR, 2024.
[44] Zhichao Wang and Yizhe Zhu. Deformed semicircle law and concentration of nonlinear random matrices for ultra-wide neural networks. The Annals of Applied Probability, 1962) 2024.
[45] Zhichao Wang, Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propagation in deep neural networks. In The Thirty Seventh Annual Conference on Learning Theory, pages 4891-4957. PMLR, 2024.
[46] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In The Eleventh International Conference Learning Representations, 2023. URL
https:///opaarviie..ett///ee www.:sww.eeeefooeeeeeeee htt.://os.....eeeoooooooooossoos 5spDgkmpy6x
13

---

Contents
Introduction
Related work
2
Setup and Assumptions
3
Spiked Data Leads to a Low-Rank Gradient
3.1
Smalr-to-Moderate Spike E 00, 0.5)): C² Activations
3.2
Smade-to-Moderate Spike: ReLUU Activation
3.3
Large Spike (v 0.5): Non C² Activations and Dependence between W and X
Impact of the Scale Parameter: MF vs. NTK Scaling
Effect of Regularization
Conclusion
Proofs
16
Regularization Proots
16
A.2
Spikey Gradient Proof
20
A.2.1
Upper and Lower Bounds
21
A.2.2
Helper Results: Subgaussianity Concentration
27
A.3
ReLU Data Alignment
31
Assumption Discussion
32
B.
Activation Function Properties
32
B.I.I
Sigmoid Function
................................
32
B.1.2
Hyperbolic Tangent (Tanh) Function
33
B.1.3
Rectified Linear Unit (ReLU) Function
................................
33
B.1.4
Exponential Linear Unit (ELU) Function
................................
33
B. 1.5
Swish Function
34
B.1.6
Softplus Function
34
B.2
Loss Function Derivatives
35
B.
3 Residue Concentration
35
B.4 B Alignment
35
Empirical Details
35
Figure
36
Figure
36
Figure4
37
Later in Training Experiments
37
Real Data Experiments
37
14

---

D Later in Training
37
15

---

Table Notation
Symbol
Meaning
Where first defined / used
d
Input dimension
Assumption
Number of samples
Assumption
Hidden-layer width
Assumption
11 = n/d, 1/2 = m/d
Proportional-saaiin ratios
Assumption
Σ
Bulk covariance matrix
Assumption
D 299
Full covariance (bulk spike)
Assumption
A 0
Bulk eigen-spectrum
Assumption
a20
Spectral-decay exponent
Assumption
<= n", V20
Spike magnitude
Assumption
q NES CS--1
Spike direction
Assumption
2 2 RR
Latent coordinates of the spike
Equation
X = XB + Xs
Data split bulk + spike
Bquation
XB
Bulk part (0, Σ rows)
Equation
Xs (zq
Rank-1 spike part Inner-layer weight matrix
Equation
W E TRX
A Assumption
E 111)"
Outer weights (fixed)
Assumption
Ym
Width scale (NTK 1//m, MF 1/m) Activation and derivatives
Equation
0,0', o"
Assumption
= Eg (W)))
Mean derivative vector
Assumption
1
Contered derivative Residue vector
Assumption
Equation
8
Alignment exponent Irlla =
Assumption
S1
Residuc-aligned rank- term Data-spike-ali aliiged rank-1 term Interpolant rank-1 term Full gradient wrt W Eror term
Section
82
Section
S12
Section
VwC
Prop. 2
S12 S2
Thm.
Er
Error term S12 S2 (large-spike version)
Thm.
E2
Bulk error from Jacobian-per penalty gradient Data-aligned rank-1 term induced by Jacobian penalty Reg. strength and Jacobian penalty
Prop. 4.3
S1
Prop. 4.3
Legg
Section
Variance isotropic Gaussian noise
Section
@
Hadamard outer products
A
Proofs
Notation In the appendix, we shall use $ g to mean that f = O(g) with probability 1 - o(1).
A.1 Regularization Proofs
Proposition 4.1. Given Assumptions
and |||rl12 = 0 ((vn), then with probability 1 (1)
we have that || ll S O(ym),), |112 S
4), and ||2l|2 S Ol mm)mm)))
Proof. These bound immediately follow from Lemma Lemma and Lemma A.3
2 with independent X and W. Assume a satisfies Assumption ffoo the noisy data X'. Additionally, suppose the modified residues satisfy rf (11) with probability 1 o(1), and Assumption 6 holds for r' with scaling parameter B'. If 7" ni and (X'WT)
Proposition 4.2 (Isotropic Gaussian noise). Assume the setup of Assumptions
wT )|| o(ni), then with high probubility:
| ll w11;
|15 5(n"~5 8),
11
S s(n" 3 2-8).
E
|
E
Proof. We prove each bound in turn.
S Bound: Recall that S Imm (X')) T "r'". Since d > n, and X33 E RRaxd is full rank with probability 1, we have that with probability 1, for any vector v
(1(X)) ell2 2 Jman (X)]e
16

---

Since the smallest eigenvalue of E' is n², with probability 1 - o(1), we have that
o Miin (XB) 2 n
Applying to S1, we get
|| 2 7 1|a p² ||r/12
vn
Then using Assumption 4 the fact that the entries of are +1, and r = @(1) we get
|| Z Y'm n/2 V
Next, we have that E' = 2m (X ((r'a))r)))) o' (X'WT). Using the fact that with probability 1 - o(1), r = (11) and 0; = 1, we have that with probability 1 - o(1)
E' Bound:
II(r'a") of (X'W")l 2 1| (X'")ll))
Thus, we have that
||"
12
Tm
X X
n
/2n 11 (X'
n
Since d > n and X'3 if full rank with probability 1, we have that with probability 1,
[I(x))) (x ((''a") oo ((''"))'''("""))))))) ox )]1 2 mmin (X') o' (X'l
(XX
o min [x))ld(XX""
Hence we get
|E 12 2
Tm Vnno/21o ld
n
S Bound: Recall that
S2² = 2mm" n" qz g0ggT=(( ((''2))
Hence we get that
|1 = Ym
n [|llll|> ((r' o''' (X'WT)))
Ym
n
n"||=||2|(((' (x'W"))
n [
S/2 Bound: Recall that
S/2 =
7m nu
n qz r ('a
Thus, we have that
||1/2112
Ym nu "1 r 11)l22
n
Imm n n
x Im n nu- pu /2 |elllllllaa 22)l22
n
E Vmmmn"
17

---

Relative Bounds: Thus, we have that using || (X'WT)
1)2 o(n)
|
(X'WT)l2 1l =1(
For the upper bounds that
S l2 || EE 2
|
ll
1o (X'"))le n-- 90-)) -
n
11
|| F2
n
Proposition 4.3 (Gradient penalty). Let Ding ( be the diagonal matrix, whose entries are || If is twice differentiable, then
1 Amm (d(WX")@o" Diag((:::)) Vw Ireg n
Proof. Letting W and f. f(x;) then note
ft = a² (.,,,) = a" h4, and Ōn.Si a.
It follows that
0z.fi = 0r fo @ = a o =z)
Recall
OZrc
1 (c= Xrj:
[Wkj
then
0f:
0f: OZic
Of: Xis
oW kj
77ic OW kj
7ik
and therefore
Ow = @a d' d'W
Let ② (h:), then
1|ó fil? = |1 ? 11? =
Σ
L3jz = 11..11111.1133
j.k
Now
0 aw,
0
11
a
d' (wf
OW
Dw
(20 (wj 22)00"((0))))))) (wf x;)) Dic.
The term inside the brackets is independent of while the term outside the brackets is independent of r. As a result this is an outer product and
1l = 2 (a o o (Wxx) @ o" (Wv,)) T..
Note above 02 refers squaring operation being applied elementwise the vector a. Therefore
Ow R
Ow @ww fill
(9)
| 1ww0 119.11
(10)
= Σ
||2.112 (a 02 2 (W o..
(11)
= 1T o < o "oo Diag([|=:113
(12)
(o'(WX") d" o (WXT)) Diag([1x:11))x
(13)
18

---

Proposition 4,4. Given Assumptions /
3
and 6 |r = @V))), a < 1, and constant
fraction of the entries of o' (XW") 2
are bounded away from then
|| 2 A (n---9- + $4)
A nn 20 9- x
mm
||
Proof. We begin by noting that since o, o' are lipschitz, we have that a', o' are bounded. Hence
d' d'(X" o o (XW")
has an operator norm that is at most (n). Since a constant fraction p of the entries are at least some universal constant then in the proportional regime, we have that
||o( o" (XW") (Xl2)lla f (lo(X"")
2 Vmc (((Vn
Recall that
1 2 x Diag([< (o (d'(X")) o ox"()"))
E2
n
Then since d > X Diag( [|: (1²) is full rank with probability 1, we have that
1 2 m min (x2 Diag([la:1))) (l) @o" )ll $ F2ll2
n
and
| 2 -m maz (x Diag([1z;11 11o (ow('""))))) (xW"))""(XW"
Due to Assumption] 2 with high probability - o(1), we have that
maa (XB) S v and Jmin (XB) Z n 1-
Then since
|| concentrates to n 1 (for a < 1), we have that
E2ll2 S 2 Vn(n2" =0
and
[221 Z
n
+ (n"
Then using the Onn) upper bound on |XW( o" (x(X ll, in the proportional regime, with high probability 1 - o(1), we get that
||22|a S mn (n 20 n -a)
Using our ((Vn) lower bound on I|o'(XTT) 3 o"(X""
112, we get
| 2 m mm 2~=--3 3 - +
On the other hand, ||rl|/2 (vvn) we have that
Vmmmn" 2 || S Vmmm
For the NTK regime, we have that
n
+
n
1 3-30
2
mm
Eg l2 | 2|
n
2v
+ n
30
19

---

A.2 Spikey Gradient Proof
Proposition 2.1 (Gradient of the loss). If Assumpiioonnholds and is differentiable, then
G == Vwr ='mx" [(ra") o o'XW"") + WVTTR(W) gaxm
exists for almost every W in Rmxdd
Proof. The first thing we need to is to compute the gradient. To begin, we compute
f(x;)
Σ
ajo Σ
j=1
Thus, we see that
0 Owrs
I(((a))
1 n
Σ
(f(x:))
a Owrs
Σase ajo
m
Σ
d
Wyk( (xre))
71
(f(x)))
aj
wan (ai)a)
WyR
71
Owis
j=1
m
((f(z;)) >
ajod (wf 2z)
Σ wyylaz)a
Owrs
((((s:))a,o (w. (xi))e
n
( Xi.
We begin by decomposing the gradient
G Ym xT ((ra") o (XW"))
This algebraic decomposition holds for the current state (X, W,r, a). irrespective any statistical dependence between W and X. Recall the data decomposition
X XB Xs =B + Ç2q" E R=xd
where rows of Xn are from NO E), NN I). |læll = and the activation derivative de- composition o'(WT)) = Inj=" + + (XW*), where μ E, (o((wx)) depends on the current W.
Substituting these into the gradient expression yields:
Ym x ((ra") o [InjiT +o' +XW(X))))
n
im x" (r(a (1) "ra")oo((XW"))
n
Ymm
(xZ + x5) (r(a + Im (xT + xx) (raa))o' (XW(())))
Ym
x 11)T" 7m x1)
St
Si2
Im ( oooo (XWT))
Ym x xg((ra") T) oool
8/2
20

---

Using Xs = =2z9", we identify the components explicitly:
X
S 7m n
u))
TTT
S12 = Smm
a(a o u)²
So
Imm
q ((ra') d'L (XW"))
n
I' X o (XW"))
n
Note that S12 shares its right singular vector (a (s) with S (up to scaling) and left singular vector q with S2. Understanding the gradient structure requires bounding the norms of these terms, which depends on the properties of the current W, r, n4, and the data statistics.
A.2.1 Upper and Lower Bounds
Given our helper results, we provide bounds for the S12, S2, and E appearing in Section Lemma A.1 (S1 Bound). Let W be the weight matrix (e.g., step t) with unit norm rows, and let S1 = "m (a o p))². Suppose XB is from Assumption has fixed #1 entries (Assumption 3 r is the current residual, and μ = F2 [a' xW)) satisfies j4 = (11) (1) for all (Assumption 4). Assume d >> Then with high probability:
X
v mmmiammiall
S ||1:112 S mm//m/mmx/ (rr
where Umin
ming |pa (11 and Pmax max; |u81 O(1).
Proof. The operator norm is
|| = 7/ 112 Mll2-
First, consider μ4, where ag +1 and μ* Eg ((1) and /max max; |M = 01) (since o' is bounded). We have:
lo (w) x)]. By assumption, Umin = ming |ps|
|| 1ul22 =
1=1
Thus, we see that
Pmin m I lla nll2 S Mmax V
By Assumption if d > n we have that with high probability
- ||||| x |X S n² |r|l2.
Substituting the bounds for || and ||a æll2 into the expression for ||S ll2 Imm X 11rll2l1a
/112:
|| Z 7m (n [1ll)) MMMmmi)) = 7m v mwmmmmmmliiiii/mmmmmmmm.m ||
Lower:
Upper: || Ym (ná ( = 7m Vmnumax I |r|22n -1.
This completes the proof.
Lemma A.2 (S12 Bound), Let W be the weight matrix (e.g.. step t) with wnit norm rows. Let S12 = {( 7m
z ="T )q(a T.)T. Suppose 2, 3,9.C Ç = n' are from Assumption 2 a has fixed +1 entries (Assumption 3 a = [x(o)(Wxx)) satisfies μa = O(1) (Assumption 4). and the current residual 8/2-1/2)
satisfies @(1r1112 ([11:|1/n-8/2-112)) (Assumption 6 Assume d > n. Then whh.p.:
|| (1m 8
21

---

Proof. Since S12 is a rank- I matrix and l|q/12 = 1, its operator norm is:
r²"r)
12 /ll2 Tm R-
|S12l22 = 7mm
1
||a
Hll2.
By Assumption @applied to the current residual r, we have
θ (Irlllan 1laenm--55 8- zTr
7
Substituting this scaling, we get
||2ll = nm -
(1rl2n
9 ||a 1ll2 = -0('''
|Tllalla o n1ll)
As established the proof of Lemma A.1 using the assumptions on a and μ (specifically un = O()).
we have ||a ull2 @ (mm). Combining these gives the final result:
S12112 = O (mmm" - I1ll2((vm))
Lemma A.3 (S2 Bound). Let W be the weight matrix (e.g., at step 1) with unis norm rows. Let S2 I'm q²² [(ra") o o'1 (XWT)). Suppose 2, q, n²" are from Assumption 2 a has fixed 3) = E., (o' (d(Ww)) satisfies μix = (1) (Assumption 4 and the current 99/2-112)
n
#1 entries (Assumption
residual satisfies =I1
((11|12nn9//--1/) (Assumption 6 Then, w.h.p.:
Tmn"" 1rll200m (o (X'( 2 || 2 7m m |1] min(n", II 22
Where S hides universal constants C1, Cz.
Proof. The operator norm is
S Imm n² || ((ra(rr)) o (XWT)))l..
n
Upper Bound: Using Lemma A.8 and Assumption that 0; Unif(=1), we have the upper bound
||= ((ra" T oo'o (XWT) ))1l & |=lla////
||:||a]1//1<|||||XXXX (XW*))
Then with probability 1 o(1. since NO,, we have that || & Cyn.. Hence get that
||2 ((raT) o o' (XW7) 1) Cllrllscllo (XW l2vv.
Then since we have Assumption we can use Lemma A.7 to bound the norm ||o (PYT)) which gives that with probability o(1),
| (Ta((raa ooo (Xl(())))) X Cl Vi (1,, VallWE
C||rl Vn ||| 2n 21//2)
Thus, we get that
S22ll2 < Ym n" "'C'I Io min ||W +1/2)
CCmm 7m ||rl min(n", |[ 2| 2
where we used the proportional scaling of n and m, Assumption I the second line.
Lower Bound: For lower bound, start by writing
(ra") "a"")o'(XX'")) Diag(r) (XWT) Diag(a).
22

---

q=" ((ra*) o ooo (XWT)) y (= Diag(r)) of (XW")) Diag(a) Thus, we have that
= =((or)" dXWX"") Diag(a).
Taking the norm and recalling that Ç n", we get
Kkq= ((ra") o (XW")) ll2 = n" Ialll [| or) o' (XW*) Diag(a)) Since the entries of a #1 and has unit norm, have that this is the same as
n" |aaall|| r)" d' (XW*)) n" ((2 r)T o' (XW*)]
By Schwarz, we have using Assumption
//Tr/vnl |r122 = (/2)
that
| olll
Σ (2,,,)2
ii 2m
||=||||||2 2(n 1|1l112).
/
1==1"
Thus, we get that for some constant C
S 2 Cm
1 n
|||l|/00mmn (a (o1(X
Lemma A.4 (Upper Bound on E). Assuming Assumption Ty. Assumption 3 Assumption 4 we have that with probability at least o(1)
Assumption
and
|| S C/mimllllle min
(1,n2~$1ll11b
Proof. Recall
7m X T) o (XWT)). Using Lemma A.8 we have that
BBI2 z | (aal (XW"))a
Tnm
Then using Assumption 1 - -1) that
whereby the rows of XB are iid from (0, Σ), we have with probability
112 S CVa,
and using Assumption 3
we trivially have that
||allx = 1.
Thus, have that
E 2 Co (XW2))l2
Pim
Then using Lemma A.7
we have that with probability 1 - o(1
||o |Xo'( Ill I C' min (n, vall
Since || = n², we get the result in the proportional scaling of Assumption
Theorem 3.1 (Gradient approximation). Suppose Assumptions
are satisfied, X and Then, for all v, Rgo.
independent, and is C2 function. Define E G S1 S12
S2-
|| - S1 - S12l12 O (11Wllnn
G S - S12 S2ll2 o (IWllenn +)
(4)
Vm r//oo
Vmmmmll7 |/
with probability - o(1) as d, n, xc. Moreover, if < then with the same probability
|
3 -
| = n
Ω
(5)
E l
log W 1|
EE
logn
110 |X
|| |1a1l12
Ω
n - 4
Ω(n u 9 9)
||12212
< O(n" "
-
(6)
||
log n W|]2
|
23

---

Proof. We start with the gradient decomposition derived in Section
G = S1 + S12 + S2 + E
where
X (a)
S1 = 7m
S12 7mm
a(a o H²
S2 = Tm (z" ((ra*) ) oo (XW2))
E Im ( oo (XW")).
We assume the conditions of the theorem hold, including the scaling vm - O(1) and the residual concentration ||r||// ||rl// O(vn/ lo (Assumption 5).
Proof of Upper Bounds:
For the first upper bound, we have - S S12 - S2 = E. Using the upper bound on || from
Lemma And the assumption /m 0()):
|| S S12
S2
|
|
||
C Cymim min (1, (1,n2=1/lW11b
|
=0 [min(1,[W[laan"
For the second upper bound, we have G S S12 = S2 + E. Using the triangle inequality and the
upper bounds ||2||2 from Lemma A.3 and |Fll2 from Lemma A.4 along with /m O()):
|G
S1
S12 12
||21|2
1|122
||/
|
v mm |1/ls min(n", || 21) + /m 11 min (1, (1, (
-0 (min(n", Bomin(n",|[W||lll 20- + min(1, [W
Proof of Lower Bounds:
We establish lower bounds for the ratios ||5|1//||||| || and ||S2|1/||||| These rely on the lower bounds for [|S 1111 ||S and the upper bound for ||] We use the result |rrll//[/lll =
((vn/ log
Ratio //1ll/1 E||: Using Lemma (lower bound) and Lemma
(upper bound), we have that
S 2
Vmmim/mmn ||||2n "!
| 2
mm || min 11, 2 1|
||r| n (a+11))
|||| so min(1, W 1222-1122
vn n (a+11)/2
log min 1, - n a
log min(1, (
24

---

If y < 1/2 and we assume ||W 122" 1/2 (1) is the dominant term in the minimum, the ratio is
12122aa//2
log nul W
If 2 1/2 and assume ||W]|2n" (11), the minimum is O()). The ratio is
n a/2
log
Ratio ||S12|/||BB| Using Lemma
A.2 (lower bound) and Lemma A.4 (upper bound);
|
v mm | 9/2-1/2 1/2
| 2
mm || min (1 W
22 1/2)
||r|2
71"-0/2-1/2
22
|||| min(1, ||W 1222
vn n² 8/2-1/2112
log n min (1, W| 11/2)
1/2)
n2" n 8/2
log (1, W 122"-1/2)
< 1/2 and assume ||W||2n --12 O(1) dominates the minimum, the ratio
Ω n 12/2-1/2
log nl Wl22
v 2 1/2 and assume |W||2n"-1/2 2 (2) the minimum is 0(1). The ratio is
n²"3/2
n" 3/2
log n
Ratio ||21//| El: We have that
o ILcMM
S
Jmm |xllll| oz 10 (JMM
2mm n (=
[|]
n ||
1|r (0' (XW1))
n² -- ||r||2 x (0' (XW(X)
|rlbox
n"-
(x'(X)T))
logn
Relative Sizes Next, prove the relative bounds. First, we have that
|| |XTrll yall n² ||
|
|
ulll
rll
Then since
2 + ||r|2 < || S Vn|[rl1e.
we get that
n Pーー
S
< --
25

---

For the second relative bound, we have that
n= r) oo S
n² ||||2 1
nク |r|l [[ (XWT)II
=
|
For a lower bound, we get that
||11 S
|r|
n
1
2
C
| n
n²
For an upper bound, we have that
||
- |||||2 m (o' (XXT))
n
n T (a'' (XWT))
|| S2 n
Theorem 3.2 (Large data-spike gradient approximation). Suppose Assumptions 2 5 satisfied, and define EL G S12 S2. Then, with probability 1 o(1) have
and
|Eall2 /m/m |
|S12l2 F1,11
n2 nー
S
n" [1(=or) (XW")
(1),
Ω
logn
logn
[XWL(X
(8)
Proof.
This proof exactly the same as Theorem 3.1 In particular, we note that
E,. = E + S
Except we use the following upper bounds. We have already bounded Si, in the following we bound E.
Data Spike: The operator norm is
|S2 7mm n2"
|| o
Using Lemma A.8 and Assumption that a, ~ Unif(#1), we have the upper bound
I ((ra²) o o (XW*)) ))1 < ||2]||/>
$ [|alla////aald]11XX)))))) [l
Then with probability 1 - o(), since N(0, I), we have that ||=||2 $ CV... Hence we get that
(I=a( o (XWT)) S Cl (aliicllo((wwww) Then since we have Assumption we can bound the norm I (X1X Il by O(n)
||>" ((ra") oi (Xl"))l S Cl Vnn
Thus, we get that
||a||| S Cym Tm ||rl n",
where we used the proportional scaling of n and m, Assumption I in the second line.
Error Term: Recall E 7tm
XT (rra") o' (XWT)). Using Lemma A.8 we have that
n
F2 S | (X
Ymm
Then using Assumption 1 o[] that
whereby the rows of X are iid from from. (0, E), we have with probability
||x & CVn,
and using Assumption 3
trivially have that
||a loo =
Thus, we have that
2
2
Ym
Then
(XW S O(n).
Since || n²", we get the result in the proportional scaling of Assumption |
26

---

A.2.2 Helper Results: Subgaussianity and Concentration
Lemma A.S. Let Z E Raxd a matrix with standard normal IID entries. If < d, then as n/d + E (0, 1), we have that with probubility I, the eigenvalues of ZZ are ((1). Further,
min (Z) @(vd - vv), (Z) O(vd v
Proof. As 1Z is a Wishart matrix, the limiting empirical spectral distribution almost surely weakly converges to the Marchenko-Paun Pastur distribution supported on [[L vO)², (1+ vo)²].
Lemma A.6. Let XB Ruxd have IID rows from (0, Σ),, where A (E) k as per Assump-
tion 2
Then with probability 1 exp( -cn) for positive universal constants c, we have that
Ω n 1-a - X - 0 00000 (n2)
Proof. We can write XB 1 where Z Rnxd has IID standard normal entries. Using
Lemma A.5 we have that in the proportional regime (Assumption ) |z||2 ((vn). The result follows using the fact that
Mmin ($ |2 < | /|/21/2 S max (5112) |Z||2,
and noting that
min =
= 0 0/2) and max 1
Lemma A.7. Let W be a given fixed matrix indepedent of X. Assumption 4 is satisfied and a is C², then have with probability 1 Ceex(-ccn) for positive universal constants c, that
(oW E C² ('minn (n.villl
for some constant C' > 0. Here E & + C² the full data covariance from Assumption 2
Proof. Since is T.-Lipschitz (Assumption 4.. its derivative of is bounded by T.. As [[)))) the centered term a' (XWT) = o' j Injr² has entries bounded by some M (e.g., M 21.). Thus, using the relation between operator and Frobenius norms:
(XW II [XW II < Mnm.
Thus, we have that in the proportional regime
= |o (x)))
On the other hand, o' o' (XW7) represents mean-centered features and is Lipschitz, using Corol- lary A.2 with probability 1 - Ceepp-ccn) we have that
||o' (XW*) O - (V7
The overall bound follows by taking the minimum of the derived bounds.
Lemma A.8. For any vectors u, v and matrix A, we have that
min [us min boy ||A||| < [I(w=)) Alla |1]]=0 ||0| A||
Proof. This follows from the observation that (uoT) A = diaggu)Adiagli
where diag(u) is the diagonal matrix with u in the diagonal. Then using the fact that
miin (B) || < ||A [2 < 0max (B)]|A|]),
where min is allowed to be zero and noticing that
muxx (diag(1 uu) = ||u]|so and
Omin (diag( (diag(u)) = min |u:]-
The bounds follow from applying the matrix norm inequality twice.
27

---

Lemma A.9 (Sub-Gausianity), For x ~ -N(0,)) E), fixed vector w E R and an --Lipschitzz function f IR + R. the random variable f(w" x) is sub-gaussian with subgaussian norm at most Ilw /1/2 for some constant C. Furthermore,
E| =|] -
Cyllw /33 11122
Proof. Using Lipschitzness,
[f(x"w) w) - f(00(w))))))))) - Lyll²" w - 01 Lylæ² w/
The variable wf ((0,00) where oi 12/211 Thus, w x is -uub-aauaan For
t 1 0,
t2
Pr[l((="w) f(0) 2 Pr [C Tw 2 1 < 2exp
2c3 / E1/2
Thus, we see that
c)2
Pr[lf(2 w)l (x"ww) 2 t| S 2 exp
2C3 m 2
where c [f(0)). For the expectations, taking expectations, we get that
E [[f(x () - [(0)]] s EE [Syl=] w|] = -rv Iw 11/213
Using 1/( (f( [x)] < [f( x) - f(0)] + [(()] and the triangle inequality for expectations, [(0)
Eljf w =)]l S ((0
- f(0)) (f()
(( giving the result.
Lemma A.10 (Covariance Operator Norm Bound). Let W R mxd be a fixed matrix whose rows have unit norm and let x ~ (0,, Σ). Suppose that f R - R is C Lipschitz respectively. Define the population second moment matrix
[(W))
where J is applied element-wise to the vector Wæ E Rm. Then
|| | [f( + [WE
for some universal constants C1, Cz.
Proof. We note that & is the uncentered covariance matrix. However, bound the operator norm of & need to consider the centered covariance matrix &
= [[())) -
Then that
| 11
T Th
sup ||
TT tv (E 10 (B f(Wx))²
sup ||||==1
sup E (0 f(Wx)) (oT f(Wx))) (B [ ((Wx)))
||||==
sup Var (v" f(Wx)) ||||=1
We want to bound this using the Gaussian Poincare inequality. Which we recall here (Link). Let g R + R be a C2 function then
Var (an S E E [[
28

---

Since z ~ E), we can write it as x 21/22. Thus, define the function
o(z) = f(Wx) f /
Σ
W 1
kml
Let us then define
u= - (wf
vn [,12,:/)+)) 7
Then we see that
m
Vg(z)² Σ u [ (2/21//2)) (uf21//) =uTW/'/
Thus, we see that
3g
| (=)
Eg
/
II [1|u|3]]
Then using Lemma
A.9
and noting that f bounded by Ly, we get that
Ex
2
E
[('(w[=)))*
C3
km
Thus, we have that Thus, using the Gaussian Poincare inequality, we see that
E [ S
+ 2 - |
12c3
Thus, we see that
|| | I $ 2112 ||WE\/2114
Finally, we see that
| 2112 [[](Ww)]
[[(
Thus,
| < [[[ + 11WE>/1 23
We are going to instantiate a few corollaries for cases that we care about. Specifically, we shall f= o' as the non non-linearity. In this case we have that E [W[)] O.
Corollary A.1. Tf [[(xx)] = 0, we have that
|| | S |I
We shall also need to bound the norm of the expectation. In the case, when a is bounded, we get that the expectation
Lemma A.11 (Feature Norm Bound). Let I, (O, be IID i 1.. n. forming rows of X. Let W E Rexdd be fixed matrix whose rows wi have norm ||w/||2 = 1. Let f R TR. be -ipsschii Define the population second moment matrix
=x(((W))((W>)))
(as in Lemma A.10 Then with probability 1 - 2e -on
2en for some universal constant c > O,
f(XWT)
<1+C'
2
n
for some universal constant C².
29

---

Proof. Since F; are IID, we have the rows of f(XWT) E Raxm are IID. Additionally, by Lemma A.9 the entries are L 2 sub-gaussian entries. Thus, we have that
X Ej max
2 f(XW")
has IID rows whose sub-Gaussian norm is at most universal constant. Let
-
C max I /1/21 +
Then using Equation 5.26 from (42. there exists universal constant C, such that
1
X
x
max (1)
< 2e att
8
m
n
Thus, with probability
2e - we have that
x $ S maxx(6, (na)(6,83)1
Using the reverse triangle inequality, we have that
- x 1
x 1112 <
|/112
Thus, with probability at least
2e 20-ct et have that
11 S | + max(o, 1/
Thus, we get that
vn | 112 V |012 max(ō, /
Multiplying both sides by C, max u / we see that
//2 (WW,)))
$ L i=1..
max
|| x 12( + ('5)/166
43 < (1+
|| N1/21/2113 21
max mm
$1 C'd) 1
Using ( = Vm, we see that with probability 1 2cm
Vi ((w6T))
- <(1+0"
2
Hence, we can again instantiate some simple corollaries. Corollary A.2. [f(x)) = O, we have that
[(W6,>))| |1 S 12C
Another important case, if f is uniformly bounded. This the case, when we apply for o'.a". Here we either have the expectation is zero. In which Corollary A.2lapplies. If the mean in non-zero then we get the following.
Corollary A.3. If [[](()]]] M, have that
|(W6x*")||
1|, Sc I+ 1/IWE
30

---

A.3 ReLU Data Alignment
Lemma A.12. Let M wp" be non-zero rank I matrix, where R and E R Assume that all entries of u and are non-zero, i.e., u, + for all = 1...., m and v/ * for all = 1,..., n. Let M be the matrix with entries Mij = du Let M = M 0.5J, where J is the m matrix of all ones. Then, rank(M) =.
Proof. Let u', u² E {0, 13
and v', v" E {0, 13 indicator vectors defined as follows:
u 6u,00
u 6ms50
of 80,20
of 80
Since we assume u, + 0 and By # for all i. every entry in us either positive or negative, and similarly for v. This means 1m = u' + u" and 1, = v' + v", where 1 denotes vector of all ones of the appropriate dimension.
The entry Mi owny >0 is if and only if (u, 0 and oy 0) or (u, 0 and vy 0). This can be written as:
= + +uu u'"""
The all-ones matrix J can be written as J = 1m Using the property that 1 = u' + u" and 1 of v":
J J= (u' + +2" D')T
= =''(''' + (1" +1"
Now compute M M1 0.5J:
M = (u'(v')² + - +1'(1") + d'(o) + (o"()"
0.5u'(v')² + 0.55" 0.5'' - - (o)"
(1' - -0 "
0.5 (6'")( ( +
= -0.(u'- (')((( - ()")²)
-
= -0" - -
Let sig denote the vector with entries sign sign(u:), where sign(x) = 1 if > 0 and sign(r) - if O. Since no u; is zero, (w' v"); = 6u,20 du. (u) Similarly, (u' v²); sign(1:) Thus, we have shown:
M = 0.5 sign(u) (i""
Since M = uo- is non-zero, both u. and must be non-zero vectors. Because we assumed zero entries, the vectors sign(u) (containing only +1) and sign(v) (containing only +1) are non- zero vectors. The matrix M is expressed as the outer product of two non-zero vectors. Therefore. rank(M) 1.
Proposition 3.1 (ReLU gradient). 1f 1 Oz, and the row of W are i.i.d. from the unit sphere, then with probability 1 - o(1) we have that o', (XWT) sign (2,) sign FT.
Proof. Recall the data decomposition x; = Ç2ay + Tb,i, where the spike direction g E IRd unit-norm, (0, 1), the bulk component Tb,; has spectrum exponent a, and the spike magnitude scales
31

---

as Ç = nn" Since each row 3 of W is uniform on S [[al] 2 = /m/d m with high probability. Using standard concentration for random projections, with probability o(1),
(((d-a) <<
W v
C | = C
c j=1
j
(logd) 0(1) a>1
a= 1,
(14)
For the spike term [[ = |zil||/Woll> 2 n" V |zs] 2 n", since |ai| 2 c with
probability 1 o() for some universal Hence, whenever 2v 1 a, the spike contribution W({zgq) dominates the bulk, so that sign(Wx.)) sign(W((<zr)) Then Lemma A.12 then impliee for ReLU that
of (XW") = sig)(z;) sign(W ("q"".
B Assumption Discussion
B.1 Activation Function Properties
We verify the smoothness and lipschitzness conditions for several common activation functions.
B.1.1 Sigmoid Function
Let (u) = (1+
Smoothness:
The Sigmoid function is infinitely differentiable (C*) for all R.
d'(u)
d"(u) (w)(1-2o(u)) = - (1))) (1))((122(())
Both (') and ou)) exist for all u E R.
Lipschitzness: Since Sigmoid is bounded and all derivatives of the sigmoid be written polynomial of sigmoid, we see that the derivatives are bounded and hence lipschitz.
Non- Non-Vanishing Derivative Here we show that if the weight vector wy is drawn uniformly from the unit sphere gd-1 then the expected derivative μj = [o æ)] is 2(11 when < 1/2.
The derivative d'(u) (1) (u)) is bounded. We can see that the argument Uj w/ is Gaussian N(O,, ((0,)22)) with variance o² = uf Emy + n² (wf 2)² Then the behavior of μıj is such that if of (1), then μr = 21).. Specifically, if o2, 0, then μj - o' o'(0) 0.25. If o2 + oc. then μ4j O.
- Spike Contribution Vss - . (wo 2) For a fixed E gd-1 and random wj E 8d-1 the term (w 4)² concentrates around its mean [[( g) 1/d. With high probability for large d, (wf q)² O(1dd). Then in proportional regime,
have that, Vs n² O(1/n) O(n2v-1). Since 1/2. 2u - < 0, so VS = o(1) as m + ao.
Bulk Contribution VB wf Ewy; For random wy E gd-1, w Ewz concentrates around Blu Sw,, = Tr)).. The eigenvalues d (E) k
If Tr(E)) O(d). so =
If0 1: (r( =
so VB =(d-9)=
0(n-a) o(1).
=f Tx{) = (log so VB = O(log d)/d) ((logn)/n) = o(1.
Ifa 1: Tr(E) @(1). so VB ((1dd) 0(1/n) o(1).
Thus, VB is either (()) (for 0) or (11 (for > 0).
32

---

B.1.2 Hyperbolic Tangent (Tanh) Function
Let (u) = tanh(u).
Smoothness: The Tanh function is C for all a R.
d'(u) = tanh (u) sech?(u) d"(u) 2 tanh(i))scch (u)
Both d'(u) and o" exist for all E R.
Lipschitzness:
For (u): max [o'( d'()) Thus, (u)) is I-Lipschitz.
For of (u): max (u) occurs at arctanh(#1/v33) giving (o"(u) 3v3 ~ 0.7698.
Thus, o'(u) is Lipschitz with 0.77 (or L 1 as a looser bound).
serves common upper bound.
Non- Von-vanishing Derivative: Let (u) tanh(u). Its derivative is (') sech? This derivative always positive, d'(u) & 1, with a maximum o0() = 1, and o'(u) - as |u| + cc. The analysis of the expected derivative (4j Ex [o'f z)] parallels that of the Sigmoid function.
B.1.3 Rectified Linear Unit (ReLU) Function
Let (u) = max(0, u).
Smoothness: Here that the derivatives for + 0 are follows
d'(u) = {1 0 i < 0
o" (u) = O for + 0
1 if uf 0
Lipschitzness:
For (u): (u)) a.e. Thus, o(u) is 1-Lipschitz.
For o'(u): d'(u) step function. bounded, but not Lipschitz over due to the discontinuity at = 0. However, its values are 0 or 1.
Non-vanishing Derivative: Since W is symmetric, we get that the mean is 0.5.
B.1.4 Exponential Linear Unit (ELU) Function
uf 0 Let (u)) = Iu -c" if 5 0"
Smoothness:
The derivatives are as follows.
{ nd
if u >
10 if u > 0 e"
d'(u) d'u)=
d"(u)
if u < 0
if u < 0
Here we have that of is continuous, and a" is defined everywhere except for Lipschitzness:
For (w): For u > 0, o'()) = 1. For u s 0, 0')) = = E (0, 1]. Thus [o((u) s 1. So (u) is 1-Lipschitz.
For d'()): For u > 0, (u) = 0. For u < O. (w) = e" E (0, 1). On [-1, 1]. the function continuous. Hence lipschitz. Thus, we have global lipschitzness.
Non-vanishing Derivative: The derivative dominates the ReLU case. Hence μj is at least 0.5.
33

---

B.1.5 Swish Function
Let (u) a sigmoid(u) u(1 + ===-
Smoothness: This follows from smoothness of Sigmoid.
Lipschitzness: Let (w) sigmoid(u) = + Then o(u) = uS(u).
For (u): The first derivative is:
o'(u) ((u) - us (u) = S(u) + uS(u)(1 S(u))
This a continuous function that decays to zero. Hence is bounded.
For o' (u)): The second derivative of (u) is:
o" (u) = d- (S(uu) + us' (u)) = ('( + (S" + (u)(
du
= 25 + uus((u) (u)
This a continuous function that decays to zero. Hence is bounded.
For a" (u): The third derivative of (u)) is:
d" = (25'(u) uu = (u) + (5" - uS" us""(u))) du
= 35" (u) + (u)
This a continuous function that decays zero. Hence is bounded.
Therefore, au)), d'), and o"(u) are all Lipschitz for Swish with 3 = 1.
Non-vanishing Derivative: The expected derivative My is:
Mg = Elo'(uu = +
[E]5(,,))
We evaluate each term:
For (1)(u;)) ): The function g(u) Suu) 1/2 is an odd function. Since uy NO, oa,) has probability density function symmetric about 0, the expectation of any odd function of uj is 0. Thus, E[S(u,) 1/2] = O, which implies [1)(u,)] 1/2.
For (u, (u,) ): The derivative of sigmoid, S' S'(w) S(u)(1- is an even function: (' S( S(-u) S(-u)) = (1 (1-5((e (()))S(u) = S' S*(w). The product h(u) = u.)((u) is an odd function, being the product of an odd function (u) and symmetric PDF about Elu, 0.
function (S' (1)( Since uj N(0, o.,) has
Combining these results:
Pj 1/2+ - 1/2
The value 1/2 is positive constant, independent of other parameters such as d, m, v, or the specifics of Σ (provided it is positive definite) and wy (provided wy E 1d-1).
B.1.6 Softplus Function
Let (u) = log(1 +c").
Smoothness: The Softplus function is C for all e
e
o'(u) 1 1 e sigmoid(u)
o" (u) (1 enyz sigmoid(u) - siggmii))
Both α'()) and o"('' exist for all u E R. Lipschitzness: The lipschitzness follows from the boundedness and lipschitzness of sigmoid.
Non-vanishing Derívative: Following the argument presented for the Swish activation function, the mean is 0.5.
34

---

B.2 Loss Function Derivatives
Let us see what this is for some common loss functions.
For the Mean Squared Error (MSE) loss,
L(f(x)) = 1 (l((()) -
- = 1 2 Σ (f(x))-)) and L'(f(x)) = f(x) - y.
=
For the Binary Cross Entropy (BCE) loss, assume the network produces logits f(X) R with associated class-one probabilities p sigmoid(2) = Ite-x RR" computed component wise. Then, for given output data {0, 13""
n
L((())== - Σ ys In(p)) + - yi) In(1 pi) L'((())) - p-2 sigmoid - yy-
For the Hinge loss for binary classification with output data y f (-1, f(X) R and
L((()) = Σ max (0, - yi f(x)))
Then (')(()) is the vector whose ith entry is given by the subgradient
aT. 0, if f(x)) 2
11 -m if f(x
0f(z;)
B.3 Residue Concentration
1. Suppose the training labels satisfy y: (x,,) + E., where f. is Lipschitz and & are i.i.d. subgaussian random variables. Then, for independent W and X, lipschitz activation functions and for either the MSE Binary Cross Entropy (BCE) loss the residues are subgaussian variables and satisfy this assumption.
For binary classification with the hinge loss, then since a, Unif(+11 we have with probability 1 o(1) that at least constant fraction of the data points satisfy 1 - fxf( 2 0, and therefore T4 11. As a result the assumption holds at initialization.
B.4 3 Alignment
Here consider Sigmoid, ReLU, Tanh, ELU, Softplus, and Swish activation functions. For each activation function, we consider three different loss functions MSE, BCE, and Hinge. Then for for each activation and loss function combination, This gives us 96 scenarios. We do each each scenario for the Mean Field and NTK scalings. For each scenario let vh 0.75 and t/2 1.25. We consider n E {750, 1500, 2250, 3000, 3750). We triple index targets
consider (v. a) {1/8, 3/8, 3/8,5/8 5/8) x 1/2).
f(x) = sigmoid + tanh + relu
for three unit vectors 31, 82 83. For each value do 50 trials get the mean inner product |d|le| eTrl. Then we then estimate beta using linear regression.
Figure 9 presents the estimates fs. Here we see that 3 has mode around 1. Recall if #1, 22 are independent uniformly unit norm vectors. Then 2 bigger than 1. This suggest r, rapidly becoming orthogonal. Note that negative 3s are cases, where the alignment improves, so 2, r are becoming parallel. Eventually, the inner product will saturate 1 and 3 should be close zero. The reason get negative (s is due to the limited range of n used for the experiments.
22 ~ d Figure however, that many is are
C Empirical Details
All code for the experiments can be found at Link
The following details common for all experiments.
35

---

Distribution of Estimates
Distribution of Reta Fstimates
16
20.3
14
12
15.0
Frequency
10
12.5
8-
10.0
7.5
2.5
0.0
Estimates
Estimates
Figure 9: Estimated & values
Hardware: All experiments were Google Colab using an A100.
Data X: We sampled uniformly randomly from the unit sphere and we used a diagonal Σ.
Estimation: We estimate μ using 10000 samples.
Targets: The triple index model we used is as follows.
f(x) sigmod ( x:) + tanh(35 + relu (Bf
For three unit vectors 31, 12,32,1 93- When using MSE loss, we let
f(x) +
for standard gaussian noise e.
When using BCE loss,
y= f(x).
Note that these y are not necessarily in [0, 1]. However, the BCE loss is still well defined. When using Hinge loss,
sign((((z)) - 0.5).
Note this dataset can imbalanced.
Alignment determination: To plot the red and blue lines in Figures 1 1,2,3,7,8, we use the following procedure. We let 51 + S12 + S2 S3 for the gradient penalty). Then we compute its leading left singular vectors for B. We then check if with y and E Thus, how we get the associated singular value and we plot the corresponding lines.
C.1 Figure
For non-isotropic W, we generate Ws by sampling the rows i.i.d. from the unit sphere. We then introduce anisotropy, by adding n -1/4199 Ws and then renormalizing to unit norm. This results in the weight concentrating around
C.2 Figure
For Figure (b), we generate Ws by sampling the rows i.i.d. from the unit sphere. We then introduce anisotropy, by adding n 12g to Ws and then renormalizing unit norm. This results in the weight concentrating around q.
For Figure (c), we generate Ws by sampling the rows i.i.d. from the unit sphere. Then we project onto the ortho-complement of and renormalize the rows.
For Figure (d), we generate Ws by sampling the rows i.i.d. from the unit sphere. We then let W X X and renormalize the rows. This results in W that is highly dependent on X.
36

---

0.615
0.3
0.012
0.602
Am Neee
Pen Plar
0.002
0.3
0.947
0.933
0.0
86
80
150
20
83
100
20
40
63
83
300
Epoch
Epoch
Epoch
(a) Small v: 0.125
(b) Medium v v 0.4375
(c) Large v v 0.75
Figure 10: Evolution of Mmin and Pmar during training. Fixed parameters: MF scaling, Tanh activation, MSE loss, a = 0.
C.3 Figure 4
Here we use ( = 0, a 0. Hence applies for prior work from 3 28
We let n E (100, 200, 300, 400, 500, 600, 700, 8003 and use d = n/2 and m = n2/3.
C.4 Later in Training Experiments
Here both network are initialized with the same weight matrix for both the inner and outer layers. We step size of 17 Ym Additionally, after each iteration, re-normalize the rows of W to have unit norm.
For Figure [ic), the mean principal angle in the following quantity. Given orthonormal basis U1, ..., 0g and VJ,..., Up two subspaces, we form the matrix A via
Aij uf 0j
We the compute cos (A))). These are the principal angles between the subspaces. We then report the mean of angels.
C.5 Real Data Experiments
MNIST Dataset: We load the standard MNIST dataset, divide by 256 have all entries in
[0, 1]. We use 1000 centered and flattened MNIST images to form X E IR 180000881 781, We estimate
0.784 > 1/2. The data is highly ill-conditioned, suggesting large effective
CIFAR Dataset: We use n = 1000 CIFAR- 10 training images, processed through a pretrained ResNet-18 (on ImageNet) extract 512-dimensional penultimate-layer activations, forming X 11000x512 We estimate v =s 0.3572 < 1/2 and a Rs 0.6.
Spcifically, the code for the transformations are as follows.
resnet18 resnet18 Bettre11...iiittteeReeeee11..... DEFAULT)
transform = transforms. [[
transforms transformm.RResize2222 Resize
transforms. (oT.oroorm.TTTTTTT
transforms NNommalie bransforme.NNrrmaliee (mean= [0.485, 0.456, 0.406], ReaNet defaults
std= 229, 0. 224, 0. 0.2]5]
])
D Later in Training
For Theorem 3.1 and Theorem 3.2 to apply beyond initialization, during training require certain assumptions to hold. We begin by considering the common assumptions needed for both theorems.
Assumption concerns the proportional regime and hence holds during training.
2. Assumption concerns the data generation process and hence also holds during training.
37

---

0310
6:30
0.15200
02207
13201343
Reeidam
0.345
Sesid
615203
63145
03264
6:20
C15228
Spoch
Spoch
Epoch
(a) Small v: 0.125
(b) Medium v.v- 0.4375
(c) Large v
0.75
Figure 11: Evolution of [r|ls//||rl during training. Fixed parameters: MF scaling, Tanh activation,
MSE loss, ar = 0.
0.32
0.334
600820
0.34
Mgrene B0002s
606820
0.22
Alggmant
0335
0.34
0.33
Epach
Epoch
(a) Small Mi 0.125
(b) Medium MiV 0.4375
(c) Large v v 0.75
Figure 12: Evolution of 12fr/(Vnlln during training. Fixed parameters: MF scaling, Tanh activation, MSE loss, 0.
3. Assumption 3 concerns the network initialization, and scaling, hence the assumptions a and Y'm continue to hold during training. Moreover, through the use of weight normalization the assumption that the rows of W are the unit sphere also holds.
4. Assumption concerns the activation function, namely its smoothness and Lipschitzness which of course also hold during training. However, it is not clear that the assumption on the non-vanishing gradient is satisfied. Despite this, we empirically verify as per Figure 10 that it does hold during training at least for small v. For moderate v = 7/16 we observe that jmin appears to decrease, hence later in training this assumption may be violated. For large appears to hold for the first iteration. We remark that this results in the suppression of S1 and S12 but does not effect S2 E. As result, suspect that the data spike q remains dominant.
3/ 3/4,the assumption only
5. Assumption
This is the assumption that
|| = 0 ||r||2
(
(log n Vn
Figure 1I shows that while this ratio grows, the change is very small. Hence, we believe that this assumptions holds.
6. Assumption 6 This is about the alignment between and Figure 12 shows that while this ratio grows, the change very small. Hence, we believe that this assumptions holds.
For the additional assumptions required for Theorem 3.1 clearly if the activation is C2 initialization then it is also C throughout training. Finally, although clearly the independence of W, and X is violated, due to the near constant gradient direction, (at least for the MF scaling) the correlation between W and X remains small.
38

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction are about general conditions under which we see low rank gradients. Our theorems and experiments reflect these claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
The paper exactly defines the scope of the work and presents results within that scope. Additionally, we have experimental evidence for the phenomena holding beyond theoretical setting.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well sell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
39

---

Answer: [Yes]
All of the assumptions are very carefully listed and the detailed proofs are presented
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
All experimental details are presented. The code is also available
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
40

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The code and data are publicaly available
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
//nips //ipss.cc/pill wwtpcccpppppppppeeeeeeeeeeeeeeee htt.:///gg..eee/ooeSSuuiiiiisss Coooeooomisssssseeeeeeee wwwiiissssipicccyccc for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
All experimtnal details are presented
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: None of the experiments require errorbars. Most experiments single runs. For experiments that are averaged over many trials, we report the number trials.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
41

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All experiments were on Google Colab with an A100.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
The research conducts conforms with the NeurIPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This a theory paper exploring theoretical aspects of spectrum of the gradient. There are no societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
42

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such threats
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: No existing assets were used.
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
43

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
The code used to run the experiments provided at an anonymize github
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/modell as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects involved
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects involved
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
44

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: We only used an LLM to help write the paper
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy for what should or should not be described.
45


Reviews:
[{"review_id": "K1VqKZ6WKP", "reviewer": "Reviewer_Y78K", "summary": "This paper investigates the low-rank structure in gradients of the hidden-layer weights in two-layer neural networks. It considers a spiked covariance model with anisotropic and ill-conditioned bulk, extending the results of existing literature. The authors demonstrate that the gradient matrix is often approximately rank two. The paper characterizes how their relative influence is modulated by the data spectrum (via parameters ν and α), network scaling regimes, activation function, and regularization techniques (e.g., weight decay, input noise, Jacobian penalties). Theoretical predictions are validated through a comprehensive suite of experiments on both synthetic data and real-world embeddings (MNIST, CIFAR-10).", "questions": "See the Weakness part.", "limitations": "yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n- The paper consider the gradient structure of two-layer NN by relaxing the usual isotropy assumptions on\nthe training data and parameters. The identification and characterization of a dominant rank-two gradient structure under general anisotropic data conditions appears to be novel. It sheds light on why and how low-rank structures persist under realistic data conditions, which is critical for interpreting deep learning models.\n- The theoretical analysis is mathematically rigorous, particularly Theorems 3.1 and 3.2, which give precise scaling laws for the emergence of low-rank gradient structure, though I did not check all the proofs completely. The paper covers both mean-field and NTK scaling, expanding the applicability of results. Regularization effects are treated with nuance and technical depth (e.g., Propositions 4.1–4.4). Experiments are well-designed and carefully interpreted, providing convincing empirical support.\n\nWeaknesses\n\n- Limited by One-Step Gradient Analysis. Most of the theoretical results are framed around single-step or initialization-time analysis. While longer-term training dynamics are partially addressed empirically (e.g., Section 3.4), the analytical extension beyond early training is missing.\n\n- Real-World Implications. 1) Although experiments on CIFAR and MNIST are included, these use fixed embeddings and do not reflect end-to-end training. 2) The paper will be improved if additional discussion could be added regarding practical implications (e.g., how this informs data preprocessing, architecture design or optimizer design).", "was_revised": true, "final_justification": "I have checked the author's rebuttal and other reviewers' comments, and I still tend to weakly accept this paper.", "rebuttal": "Thank you for your thoughtful and constructive review.  We are particularly grateful for your recognition of the paper’s strengths:\n\n### **Positives**\n\n* **Novelty and generality:** you highlight the originality of our rank‑two gradient analysis under anisotropic, ill‑conditioned data, an essential generalisation of prior work.\n\n* **Theoretical depth:** you note the mathematical rigour of Theorems 3.1, 3.2 and the nuanced regularisation analysis.\n\n* **Experimental care:** you find the experiments well‑designed and consistent with theory on both synthetic data and real‑world embeddings.\n\nBelow we address your main concerns.\n\n### **Reviewer Concerns**\n\n> One step Gradient \n\nWe agree that extending theoretical insights beyond the first gradient step is a crucial direction. While our main theorems are stated at initialization, the following mechanisms allow our analysis to remain relevant later in training:\n\n**Large‑spike regime (Thm. 3.2).** Because Thm. 3.2 does *not* rely on independence between data and weights, its conclusions hold after several training steps.  Appendix D (p. 37) verifies that the remaining assumptions continue to hold during training, so the rank‑two structure persists.\n\n**Resampling argument for Theorem 3.1:** In the small-to-moderate spike setting, Theorem 3.1 does require $W$ and $X$ to be independent. However, this requirement can be satisfied at iteration $t$ by resampling a new batch $\\tilde{X}_t$, which is independent of the current weights $W_t$. In Appendix D (page 37), we empirically verify that the remaining assumptions (e.g., residue concentration and alignment) continue to hold later in training. This suggests that Theorem 3.1 still applies.\n\n> Practical Takeaway\n\nNote our current experiments on CIFAR-10 use fixed embeddings. However, we use the MNIST dataset with minimal whitening. \n\nWe appreciate the request to sharpen the practical relevance of our findings. Although a full end‑to‑end treatment is future work, the present analysis already yields concrete insights.\n\nA pragmatic rule‑of‑thumb emerging from our analysis is:\n\n* **Leading eigendirection helpful for the task** → use a non‑saturating activation such as ReLU and/or add a Jacobian penalty to keep the spike data‑aligned.  \n* **Leading eigendirection harmful / carries spurious signal** → use a smooth activation (sigmoid, GELU) and/or inject isotropic noise to push the spike into the residue.\n\nThese choices control whether training emphasizes or suppresses the dominant data mode, offering a knob to match the assumed structure of the task.\n\n**Extension of feature‑learning dynamics.** Building on Moniri et al. (2024), we show that the learned feature matrix satisfies\n\n$$  F\\_1 = F\\_0 + \\sum\\_{k=0}^L \\begin{bmatrix} c\\_k^{\\tau_1} \\\\\\\\ \\vdots \\\\\\\\ c\\_k^{\\tau_m} \\end{bmatrix} \\mathbf{1}\\_n^T \\circ \\left[ (1-\\lambda \\eta) \\tilde{X}\\_S W\\_0^T + \\eta \\tilde{X}(S\\_1 + S\\_{12} + S\\_2)\\right]^{\\circ k} + \\Delta $$\n\nwhere $\\Delta$ is small.  Whereas Moniri et al.’s Theorem 3.1 shows that features become **univariate polynomials of the spike, independent of weights**, our formula reveals **multivariate polynomials of both spikes and the evolving weight matrix** in the anisotropic setting.  This richer dependence explains why regularisers that favour a data‑aligned spike can improve generalisation after the outer layer is trained.\n\nWe also conducted a small experiment to test how different “spikes’’ affect performance.  With $\\nu=3/8$ we trained three networks that share the same architecture but differ in the first‑layer setup:\n\n* **ReLU** (data‑aligned spike)  \n* **Sigmoid** (residue‑aligned spike)  \n* **Sigmoid + Jacobian penalty** (data‑aligned spike)\n\nInner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer.  Results:\n\n| variant | train loss (inner) | test loss (inner) |\n|---------|-------------------|-------------------|\n| Sigmoid (no reg.) | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid + Jac. pen. | highest | highest |\n\nAfter subsequently fitting the outer layer (ridge regression), the ordering **reversed**:\n\n| variant | final train loss | final test loss |\n|---------|-----------------|-----------------|\n| Sigmoid + Jac. pen. | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid (no reg.) | highest | highest |\n\nThis illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts."}, {"review_id": "iANMuTwck3", "reviewer": "Reviewer_y7gX", "summary": "This paper studies the implicit gradient structure in two-layer feedforward neural networks under generalized data assumptions (anisotropy) and independence assumptions on weight initializations. When the input data distribution follows a rank-one Gaussian *spiked* model with a (possibly anisotropic) bulk covariance, then under some technical conditions, the loss gradient will exhibit a specific low-rank structure in weight space. In particular, it is shown that a rank-two spike emerges, composed of a \"data'' spike which is aligned to the spike in the input covariance, and a \"residue'' spike which is aligned to the loss derivative (e.g. $y - f(x)$ in the case of square loss). It is then shown the relative strength of the two spikes is determined essentially by three things: 1. the input spike strength, 2. the regularity properties of the input activation, and 3. the \"scale'' parameter, e.g. Neural Tangent Kernel or Mean-Field regimes. These predictions are corroborated in simulation. Lastly, the effect of various regularizers is studied.", "questions": "Beyond responding or clarifying the points in Weaknesses above, I have the following clarification questions:\n    \nCan the bounds be applied inductively for any-iteration guarantee? I do not see clearly stated whether the intention is for the theoretical bounds to hold beyond early-training (for example, Assumption 3 implies the output layer weights are random and not updated).\n\nWhat are predictions for how the observed structure behaves beyond a precisely rank-one spike, such as for low-rank or approximate spikes with rapidly decaying (but not bulk-separated) singular values? For example, for $r = O(1)$ spikes of sufficiently large and comparable strength, how many spikes do we expect to see in the gradient?\n\nWhy is the spike strength $\\zeta$ defined as a function of the data/batch size $n$? Maybe I'm missing something, but it would seem more natural to be a function of the data dimension $d$, as it is a population-level quantity.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Overall, I think the paper is clearly presented and concretely establishes some novel phenomena in the learning theory of neural networks. Therefore, I lean toward acceptance of the paper. My main concerns with the paper lie with the breadth and relevance of the results in the paper.\n\n## Strengths\n\nCompared to prior work studying feature learning on two-layer networks via gradient descent, this paper generalizes the setting in various ways. Notably, key prior work in feature learning considers isotropic gaussian inputs, and typically considers 1-step gradient updates from i.i.d. weight initializations. This work extends to a rank-one spiked model with anisotropic bulk under a general loss function, and computing the gradient at a weight matrix with possibly dependent entries. The key insight is that this added complexity leads to a quantifiable phenomenon in the residue and data spikes in the gradient, which otherwise does not manifest in the more ideal settings in prior work. I think documenting this phenomenon and providing some tools for analyzing feature evolution in broader settings are valuable.\n\nThe other technical assumptions are justified either from empirical observation or theoretical derivations in the appendix, which is appreciated.\n\n## Weaknesses\n\nIn general, my main concern with the paper, tying back to prior work on feature learning, is that it is not clear to me what demonstrating the rank-two structure (and how it interacts with choice of activation function and regularization) implies about *feature learning* itself, or similarly learning performance. For example, it is not clear to me after reading the paper whether both spikes are useful for learning, or if the data spike can be spurious, such as in the prototypical single-index regression setting.\n\nOn that note, I am not sure if the raw gradient (and ipso facto gradient descent) is necessarily a good object to study for feature learning in anisotropic settings. For example, [1] and [2] demonstrate that SGD can be severely suboptimal when the data distribution is anisotropic, proposing forms of preconditioning to alleviate the poor convergence of SGD. From my reading [2] considers a two-layer network setting. Now, given a generic anisotropic gaussian input and a single-index target, the proposed algorithm is to precondition the weight gradient by the *inverse (sample $\\approx$ population) covariance*. Therefore, it seems therein the authors recover an approximately *rank-one* gradient as in the seminal paper [3] (which considers isotropic inputs). If one places the spiked anisotropic covariance into the set-up of [2], it seems that a dichotomy occurs: the raw gradient should exhibit a rank-two structure as proven here but aligns very poorly with the correct direction, whereas a provably adjusted gradient recovers the *rank-one* structure from prior single-index learning papers.\n    \nTo re-iterate: it seems in the general anisotropic setting, the raw gradient can be very ill-aligned to a descent direction, and existing alternate algorithms that improve on SGD's performance need not respect the rank-two spike structure. This is even more so when considering popular adaptive descent methods that alter the gradient in a non-linear fashion such as Adam and Muon, hence my concerns about the ultimate relevance of these observations.\n\n[1] Amari et al, \"When does preconditioning help or hurt generalization?\"\n\n[2] Zhang et al. \"On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning\"\n\n[3] Ba et al. \"High-dimensional asymptotics of feature learning: How one gradient step improves the representation\"", "was_revised": true, "final_justification": "This paper presents some deep learning theory demonstrating the emergence of a low-rank structure in the gradient of two-layer networks in data-anisotropic settings. The paper presents some new results that may be interesting to the field, and therefore I lean toward acceptance (4). My main concern with the paper is that it does not concretely tie this new phenomena to generalization behavior, and thus does not yet present actionable takeaways.", "rebuttal": "Thank you very much for your thoughtful and constructive review.  We appreciate your positive assessment of the paper’s contributions and your insightful questions, particularly regarding the relevance of the observed rank-two gradient structure for feature learning. \n\n### Positives\n\n* **Generalization beyond prior work:** The paper extends existing feature learning analyses from isotropic Gaussian inputs and independent weights to a rank-one spiked covariance model with anisotropic bulk, under general loss functions and non-i.i.d. weight initializations.\n\n* **Novel gradient structure:** It identifies a new rank-two gradient phenomenon, with distinct residue- and data-aligned spikes, offering a richer understanding of feature evolution that does not arise in prior simplified settings.\n\n* **Sound theoretical and empirical grounding:** The technical assumptions are well-justified (empirically and analytically), and simulations corroborate the theoretical predictions.\n\n### Concerns of the reviewer.\n\n> On that note, I am not sure if the raw gradient (and ipso facto gradient descent) is necessarily a good object to study for feature learning in anisotropic settings.\n\nWe thank the reviewer for pointing us to papers [1] and [2], which indeed raise important considerations about optimization dynamics in anisotropic settings.\n\nHowever, we respectfully disagree with the assertion that studying the raw gradient is a “good object to study for feature learning in anisotropic settings”. On the contrary, we believe that both [1] and [2] reinforce the value of understanding the structure of unpreconditioned gradients:\n\n**Preconditioning need not kill rank 2.** \n\nIn [2] the feature‑learning update involves $P_g$ and $Q_G$, both positive‑semidefinite.  A preconditioner that acts similarly on the residue‑ and data‑directions will *retain* the rank‑two structure.  Hence knowing the raw gradient’s geometry is still informative when designing or analyzing preconditioned algorithms.\n\n**Need Dual Understanding to Isolate Important Algorithmic Aspects**\n\nSecond, while it is true that preconditioned SGD and adaptive methods like Adam may accelerate convergence or improve alignment in some settings, a theoretical understanding of why and when this occurs necessarily begins with understanding the raw gradient. Our work aims to provide this foundational understanding.\n\n**Beyond Single Index Model**\n\nThe alignment metric in [2] is tailored to single‑index models.  Our labels $y=f(x)$ are only assumed Lipschitz, so they may depend on *many* directions.  In this broader setting it is unclear whether collapsing to rank 1 is optimal; the second spike can carry additional tasks‑relevant signals that single‑index theory neglects.\n\n**Broader scope than [2].** \n\nThe core theorems of [2] (3.6 & 3.9) are themselves *gradient‑structure* results.  Theorem 3.9 is a **special case** of our Theorem 3.1: it is limited to mean‑field (MF) scaling, MSE loss, isotropic first‑layer weights, and a well‑conditioned covariance.  Our analysis covers both NTK *and* MF regimes, arbitrary losses, data‑dependent or anisotropic weights, and ill‑conditioned bulks via the spectral‑decay exponent $\\alpha$.  Consequently we capture phenomena, e.g. the emergence of a **data‑aligned spike**, that [2] cannot see.\n\n**Rank‑two can beat rank‑one even in single‑index cases.** \n\nA mis‑aligned rank‑one update is not necessarily better than a rank‑two update whose second direction compensates for the misalignment.  Our framework provides the tools to analyse this possibility, which [2] leaves unexplored.\n\n**Regularization**\n\nAdditionally, if the gradient is misaligned, other ways of fixing it besides natural gradients and Adam could be the use of regularization. Hence our study about the effect of regularization on the gradient is also important. \n\n**Comparison with [1]**\n\nFinally, we briefly comment on [1]. While that paper focuses on linear models and assumes well-conditioned data, its results are more nuanced than they may first appear. For example, Figure 3 of [1] shows that gradient descent (GD) initially reduces error faster than natural gradient descent (NGD), and that early-stopped GD can outperform early-stopped NGD. This suggests that the initial descent direction of GD, precisely the object we study, can be quite effective, and possibly superior in practice. Thus, understanding the structure of GD's updates remains valuable, even in light of competing optimization schemes.\n\n**Summary**\n\nIn summary, this paper aims to provide a systematic characterization of the raw gradient in anisotropic, spiked settings. We view this as a necessary step toward understanding and potentially improving feature learning in practical neural networks, whether by standard SGD or by more sophisticated algorithms.\n\n> In general, my main concern with the paper, tying back to prior work on feature learning, is that it is not clear to me what demonstrating the rank-two structure (and how it interacts with choice of activation function and regularization) implies about feature learning itself, or similarly learning performance. For example, it is not clear to me after reading the paper whether both spikes are useful for learning, or if the data spike can be spurious, such as in the prototypical single-index regression setting.\n\nThis is a great question that is currently being explored by the authors.\n\nBuilding on Moniri et al. (2024), we show that the learned feature matrix satisfies\n\n$$  F\\_1 = F\\_0 + \\sum\\_{k=0}^L \\begin{bmatrix} c\\_k^{\\tau\\_1} \\\\\\\\ \\vdots \\\\\\\\ c\\_k^{\\tau\\_m} \\end{bmatrix} \\mathbf{1}\\_n^T \\circ \\left[ (1-\\lambda \\eta) \\tilde{X}\\_S W\\_0^T + \\eta \\tilde{X}(S\\_1 + S\\_{12} + S\\_2)\\right]^{\\circ k} + \\Delta $$\n\nwhere $\\Delta$ is small.  Whereas Moniri et al.’s Theorem 3.1 shows that features become **univariate polynomials of the spike, independent of weights**, our formula reveals **multivariate polynomials of both spikes and the evolving weight matrix** in the anisotropic setting.  This richer dependence explains why regularisers that favour a data‑aligned spike can improve generalisation after the outer layer is trained.\n\nWe also conducted a small experiment to test how different “spikes’’ affect performance.  With $\\nu=3/8$ we trained three networks that share the same architecture but differ in the first‑layer setup:\n\n* **ReLU** (data‑aligned spike)  \n* **Sigmoid** (residue‑aligned spike)  \n* **Sigmoid + Jacobian penalty** (data‑aligned spike)\n\nInner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer.  Results:\n\n| variant | train loss (inner) | test loss (inner) |\n|---------|-------------------|-------------------|\n| Sigmoid (no reg.) | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid + Jac. pen. | highest | highest |\n\nAfter subsequently fitting the outer layer (ridge regression), the ordering **reversed**:\n\n| variant | final train loss | final test loss |\n|---------|-----------------|-----------------|\n| Sigmoid + Jac. pen. | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid (no reg.) | highest | highest |\n\nThis illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts.\n\n\n> Can the bounds be applied inductively for any-iteration guarantee? I do not see clearly stated whether the intention is for the theoretical bounds to hold beyond early-training (for example, Assumption 3 implies the output layer weights are random and not updated).\n\nYes in the following way.  \n\n*Large‑spike regime:* Theorem 3.2 does **not** depend on weight–data independence, so its conclusions hold for *all* iterations as long as the other assumptions continue to hold, which we verify empirically in App. D (p. 37).  \n\n*Small / moderate spikes:* At step $t$ we can draw a fresh mini‑batch $\\tilde X_t$ that is independent of the updated weights $W_t$.  Independence is then restored and the remaining assumptions still hold (again confirmed in App. D).  Hence the rank‑two guarantee can be applied inductively throughout training.\n\n\n> What are predictions for how the observed structure behaves beyond a precisely rank-one spike, such as for low-rank or approximate spikes with rapidly decaying (but not bulk-separated) singular values? \n\nIf the population covariance has $k$ well‑separated spikes, our theory predicts $k+1$ spikes in the gradient matrix.  Their relative magnitudes follow exactly the same formulas as in Theorems 3.1 (small/mid spikes) and 3.2 (large spike), evaluated at the corresponding eigenvalues.  \n\n\n> Why is the spike strength defined as a function of the data/batch size? Maybe I'm missing something, but it would seem more natural to be a function of the data dimension, as it is a population-level quantity.\n\nSince we consider the proportional scaling limit $n^\\nu = \\Theta(d^\\nu)$ they are equivalent."}, {"review_id": "PEBeQBrH2q", "reviewer": "Reviewer_4UNe", "summary": "This paper investigates the low-rank structure of training-loss gradients in two-layer neural networks. Under a spiked data model, the authors analyze both mean-field and Neural Tangent Kernel (NTK) scaling regimes and demonstrate that the input gradient can be closely approximated by a rank-two matrix.", "questions": "Given the independence assumption between $X$ and $W$ in Theorem 3.1, why focus exclusively on the first training update? Can the analysis be extended to subsequent training steps where this independence no longer holds?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The manuscript is well-organized and clearly written.  \n- The theoretical results are precise, rigorous, and elegantly presented.  \n- Comprehensive empirical evaluations substantiate the theoretical insights.\n\n## Weaknesses\n- Theorem 3.1 relies on the strong assumption that data $X$ and weights $W$ are independent, a condition that holds only at the very first training step. It is unclear how—or if—this analysis extends to later stages of training when $X$ and $W$ become correlated.", "was_revised": true, "final_justification": "They answered all my questions.", "rebuttal": "We thank the reviewer for the comments and feedback. In particular for highlighting:\n\n* The manuscript is well-organized and clearly written.\n* The theoretical results are precise, rigorous, and elegantly presented.\n* Comprehensive empirical evaluations substantiate the theoretical insights.\n\nWe now address the reviewers concerns \n\n>  Why Theorem 3.1 matters even at step 1  \n* It allows anisotropy in **both** the data \\(X\\) *and* the initial weights \\(W\\), and tolerates ill‑conditioned bulks—features absent from prior rank‑one analyses.  \n\n* It unifies **mean‑field and NTK** scalings, whereas earlier work treats only one regime.\n\nUnderstanding this richer first‑step geometry is essential: it sets the initial condition for subsequent feature evolution. \n\n>  Extending beyond the first update  \n\n1. **Large‑spike regime.** Theorem 3.2, which governs large spikes, assumes *no* independence between $X$ and $W$.  Its rank‑two conclusion therefore holds throughout training.  In particular, in Appendix D (p. 37), we empirically verify that the rest of the assumptions still hold later in training. \n\n2. **Small / moderate spikes.** At iteration $t$ we can draw a fresh mini‑batch $\\tilde X_t$ that is independent of the current weights $W_t$.  This restores the required hypothesis  of Theorem 3.1.  Appendix D (p. 37) shows that the remaining conditions (residue concentration, alignment, etc.) still hold empirically, so the theorem applies to $(\\tilde X_t,W_t)$ at every step.  \nResampling is a standard device in stochastic gradient methods and matches practical minibatch training.\n\n**Summary**\n  \nIndependence is needed only for Theorem 3.1 and only to control *small* spikes.  Large spikes are already covered by Theorem 3.2 for the full trajectory, and small spikes can be re‑analysed at any step via minibatch resampling.  We will add a remark making this explicit in the final version."}, {"review_id": "V1Z1woYOPj", "reviewer": "Reviewer_k44k", "summary": "This paper presents a theoretical and empirical investigation into the low-rank structure of gradients in two-layer neural networks without assuming data isotropy. Under the spike Gaussian ensemble with small, moderate, and large spike regimes, the high-dimensional asymptotic behaviors with proportional diverging n, m, d (as described in Assumption 1) are derived. It shows that gradients are dominated by two rank-one components related to the bulk data–residue and the input data structure. The influence of activation functions, scaling regimes (both mean field (MF) and neural tangent kernel (NTK) regimes), and regularization methods is also characterized.", "questions": "- How does the gradient influence the training process? Empirically, is it possible to simultaneously show the change in loss and the change in gradient rank during training? For example, Figure 6(b) shows a phase transition in gradient rank—does the loss exhibit any corresponding behavior?\n- Could the authors provide a guideline or criterion for selecting different training strategies under varying conditions?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n- The paper provides a rigorous theoretical framework for understanding low-rank gradients under anisotropic and ill-conditioned data. \n- The paper analyzes how common regularizers (e.g. weight decay, noise, Jacobian penalty) affect gradient structure. Meanwhile, behaviors across both MF and NTK scalings are taken into account. \n\n**Weaknesses:**\n- The paper primarily focuses on gradients at or near initialization, with less on full training trajectories.\n- Assumption 2 requires that the input data $x_i$ follow a (spike) multivariate Gaussian distribution. \n- Only standard two-layer architectures are considered; extensions on the effects in CNNs, transformers, etc., are not explored.", "was_revised": true, "final_justification": "After reviewing the initial assessments, the other reviewers’ comments, and the authors’ rebuttal, I incline to recommend acceptance. The work extends prior analyses beyond isotropic settings, derives asymptotic results under a spiked Gaussian model, and characterizes the roles of activation, scaling regime (MF/NTK), and regularization. These extensions are novel relative to prior literature, and the theoretical findings are corroborated with simulations. In addition, the rebuttal was clear and addressed my concerns. Overall, the paper’s clarity, theoretical rigor, and extension of existing frameworks make it a valuable contribution to the theory of learning under non-isotropic data. I lean toward acceptance.", "rebuttal": "We thank the reviewer for their thoughtful and positive assessment. \n\n**Rebuttal TLDR:** While our theory is developed at (or very near) initialization for a 2 layer network, we explain below how the rank‑2 structure it reveals continues to shape learning dynamics and how the analysis can be broadened beyond the current Gaussian‑spike setting and to CNN setting.\n\n### **Positives from the Review**\n\nWe would also like to highlight the positives from the review \n\n* Develops a **rigorous asymptotic theory** that explains why gradients collapse to rank $\\le 2$ under anisotropic (spiked) data.\n* **Unifies MF and NTK regimes** within the same analytical framework, showing how each scaling modulates the low‑rank structure.\n* **Dissects the impact of common regularizers** weight decay, isotropic noise, and Jacobian penalty, and links each to specific changes in gradient rank.\n* **Combines theory with targeted experiments**, validating the predicted phase transition in gradient rank across spike strengths.\n\n### **Responses to the reviewers concerns**\n\n> The paper primarily focuses on gradients at or near initialization, with less on full training trajectories.\n\nWe acknowledge this limitation. Establishing the rank‑2 structure at initialization is only a first step, but it is the prerequisite for any later‑time analysis; extending the theory along the entire trajectory is left for future work.\n\n> Assumption 2 requires that the input data $x_i$  follow a (spike) multivariate Gaussian distribution.\n\nThe Gaussian assumption is mainly technical.\n* **Large‑spike regime:** our arguments rely only on sub‑Gaussian concentration and therefore carry over unchanged to generic sub‑Gaussian inputs, with a potential path toward sub‑exponential or even sub‑Weibull tails.\n* **Small‑spike regime:** the obstacle is the use of the Gaussian Poincaré inequality; since analogous inequalities hold for many log‑concave distributions, the same proof strategy should apply there as well.\n\n> Only standard two-layer architectures are considered; extensions on the effects in CNNs, transformers, etc., are not explored.\n\nWe note that a convolutional network can be represented as a feed‑forward two‑layer network with a sparse, weight‑tied matrix $W$.  Assuming the convolutional filter $w$ is $\\ell_2$‑normalized, our theory applies.\n\nLet $x\\in\\mathbb{R}^d$ be the input, $w=(w_1,\\dots,w_k)\\in\\mathbb{S}^{k-1}$ the filter, and $a\\in\\mathbb{R}^m$ the second‑layer weights, where $m=d-k+1$.  Define\n\n$$\nW_{i,j}=\n\\begin{cases}\nw_{\\,j-i+1}, & 1\\le j-i+1\\le k,\\\\\\\\\n0, & \\text{otherwise},\n\\end{cases}\n\\qquad W\\in\\mathbb{R}^{m\\times d},\n$$\n\nso the network output is\n\n$$\ny = a^{T}\\sigma(Wx).\n$$\n\nTreating the entries of $W$ as independent parameters yields the unfolded gradient\n\n$$\nG = \\frac{\\partial L}{\\partial W}\\in\\mathbb{R}^{m\\times d},\n$$\n\nto which Theorems 3.1–3.2 apply, implying $\\operatorname{rank}(G)=2$.\n\nBecause $W$ is tied through $w\\mapsto W$, the true filter gradient is\n\n$$\n\\frac{\\partial L}{\\partial w_\\ell}\n= \\sum_{i=1}^m G_{i,i+\\ell-1},\n\\quad\n\\ell=1,\\dots,k.\n$$\n\nWriting $G = u^{(1)}(v^{(1)})^{T} + u^{(2)}(v^{(2)})^{T}$ with $u^{(r)}\\in\\mathbb{R}^m$ and $v^{(r)}\\in\\mathbb{R}^d$, and letting $\\tilde{v}^{(r)}\\_i = v^{(r)}\\_{d-i+1}$, we have\n\n$$\n\\nabla_w L\n= u^{(1)} * \\tilde v^{(1)} + u^{(2)} * \\tilde v^{(2)},\n$$\n\nso $\\nabla_w L$ lies in a subspace of dimension at most $2$.  Preliminary CNN experiments confirm this prediction; we will include the plots in the appendix.\n\n\n> How does the gradient influence the training process?\n\nUsing the learning‑rate prescription $\\eta=\\gamma_m^{-1}$, we consistently observe lower loss in the MF regime than in the NTK regime.  \n\nFor the run shown in Fig. 6 (b), we logged the loss alongside the gradient alignment. However, the loss curve remains smooth, so no dramatic visual “kink’’ appears.  We will add the combined plot to the supplement.\n\nWe also conducted a small experiment to test how different “spikes’’ affect performance.  With $\\nu=3/8$ we trained three networks that share the same architecture but differ in the first‑layer setup:\n\n* **ReLU** (data‑aligned spike)  \n* **Sigmoid** (residue‑aligned spike)  \n* **Sigmoid + Jacobian penalty** (data‑aligned spike)\n\nInner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer.  Results:\n\n| variant | train loss (inner) | test loss (inner) |\n|---------|-------------------|-------------------|\n| Sigmoid (no reg.) | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid + Jac. pen. | highest | highest |\n\nAfter subsequently fitting the outer layer (ridge regression), the ordering **reversed**:\n\n| variant | final train loss | final test loss |\n|---------|-----------------|-----------------|\n| Sigmoid + Jac. pen. | **lowest** | **lowest** |\n| ReLU | medium | medium |\n| Sigmoid (no reg.) | highest | highest |\n\nThis illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts.\n\n\n> Could the authors provide a guideline or criterion for selecting different training strategies under varying conditions?\n\nA pragmatic rule‑of‑thumb emerging from our analysis is:\n\n* **Leading eigendirection helpful for the task** → use a non‑saturating activation such as ReLU and/or add a Jacobian penalty to keep the spike data‑aligned.  \n* **Leading eigendirection harmful / carries spurious signal** → use a smooth activation (sigmoid, GELU) and/or inject isotropic noise to push the spike into the residue.\n\nThese choices control whether training emphasizes or suppresses the dominant data mode, offering a knob to match the assumed structure of the task."}]

Author rebuttal:

Meta-review: The paper considers two-layer neural networks under both mean-field and NTK scaling regimes and shows that the gradient of the training loss exhibits an approximately low-rank structure. This analysis relaxes the usual assumptions of isotropic data and parameters, as well as their independence. In particular, when the data follow a multivariate normal distribution with a rank-one spiked anisotropic covariance and an ill-conditioned bulk, the gradient is theoretically shown to be approximately rank-two, with its two dominant components aligned with the bulk residue and the input data spike. The relative influence of these components is modulated by the data spectrum, network scaling regime, activation function, and regularization techniques. The theoretical claims are supported by experiments.

The reviewers agree that the paper makes a meaningful theoretical contribution, noting that: 
1) the insights into early training dynamics are interesting, while generalizing the commonly considered conditions for the data, weights, and their dependence, and considering both MF and NTK regimes, as well as the implications of regularization techniques
2) the theoretical results are rigorous, sound, well-presented, and well-supported by the experiments
3) the paper is well-written with potential for further impact. 

A few concerns were raised, particularly regarding the implications of the analysis on feature learning and generalization, the fact that the analysis focuses on early training, and some technical clarifications. Most of these points were addressed in the rebuttal. I therefore recommend acceptance and encourage the authors to consider the feedback and update the final version of the paper accordingly.
Decision: Accept (poster)