Low Rank Gradients and Where to Find Them

We thank the reviewer for their thoughtful and positive assessment.

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.

Positives from the Review

We would also like to highlight the positives from the review

• Develops a rigorous asymptotic theory that explains why gradients collapse to rank $\le 2$ under anisotropic (spiked) data.
• Unifies MF and NTK regimes within the same analytical framework, showing how each scaling modulates the low‑rank structure.
• Dissects the impact of common regularizers weight decay, isotropic noise, and Jacobian penalty, and links each to specific changes in gradient rank.
• Combines theory with targeted experiments, validating the predicted phase transition in gradient rank across spike strengths.

Responses to the reviewers concerns

The paper primarily focuses on gradients at or near initialization, with less on full training trajectories.

We 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.

Assumption 2 requires that the input data $x_i$ follow a (spike) multivariate Gaussian distribution.

The Gaussian assumption is mainly technical.

• 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.
• 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.

Only standard two-layer architectures are considered; extensions on the effects in CNNs, transformers, etc., are not explored.

We 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.

Let $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

so the network output is

Treating the entries of $W$ as independent parameters yields the unfolded gradient

to which Theorems 3.1–3.2 apply, implying $\operatorname{rank}(G)=2$.

Because $W$ is tied through $w\mapsto W$, the true filter gradient is

Writing $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

so $\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.

How does the gradient influence the training process?

Using the learning‑rate prescription $\eta=\gamma_m^{-1}$, we consistently observe lower loss in the MF regime than in the NTK regime.

For 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.

We 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:

• ReLU (data‑aligned spike)
• Sigmoid (residue‑aligned spike)
• Sigmoid + Jacobian penalty (data‑aligned spike)

Inner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer. Results:

After subsequently fitting the outer layer (ridge regression), the ordering reversed:

This illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts.

Could the authors provide a guideline or criterion for selecting different training strategies under varying conditions?

A pragmatic rule‑of‑thumb emerging from our analysis is:

• 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.
• Leading eigendirection harmful / carries spurious signal → use a smooth activation (sigmoid, GELU) and/or inject isotropic noise to push the spike into the residue.

These choices control whether training emphasizes or suppresses the dominant data mode, offering a knob to match the assumed structure of the task.

We thank the reviewer for the comments and feedback. In particular for highlighting:

• The manuscript is well-organized and clearly written.
• The theoretical results are precise, rigorous, and elegantly presented.
• Comprehensive empirical evaluations substantiate the theoretical insights.

We now address the reviewers concerns

Why Theorem 3.1 matters even at step 1

• 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.

• It unifies mean‑field and NTK scalings, whereas earlier work treats only one regime.

Understanding this richer first‑step geometry is essential: it sets the initial condition for subsequent feature evolution.

Extending beyond the first update

1. 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.

2. 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.
   Resampling is a standard device in stochastic gradient methods and matches practical minibatch training.

Summary

Independence 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.

We thank the reviewer for their comment. We answer the reviewers questions below.

However, the proof only provides an upper bound for $E$ instead of $E_L$, and there should be an extra part $S_1$ inside $E_L$. Can you elaborate on this?

In the proof for Theorem 3.2 we focused on the terms whose scaling is different between Theorem 3.1 and Theorem 3.2

In particular, Lemma A.1 still tells us that $\|\|S_1\|\|_2 \lesssim \sqrt{m} \gamma_m \|\|r\|\|_2 n^{-1/2}$. The proof for Theorem 3.2 gives us that $\|\|E\|\|_2 \lesssim \sqrt{m} \gamma_m \|\|r\|\|_2$. Hence here the $S_1$ is much smaller.

This then gives us the bound that $\|\|E_L\|\|_2 = \|\|E + S_1\|\|_2 = O(\sqrt{m} \gamma_m \|\|r\|\|_2)$

As discussed in 191-195, if certain conditions hold, then we are guarateed to have the $S_2$ spike escape the bulk and hence the gradient will be rank 1 not rank 2. This is empirically verified in Figure 3.

Besides, I noticed that the key point of the proof is that you require $W$‘s columns to be on the unit sphere. On the other hand, the paper considers the GD updates. How can you ensure that both conditions are satisfied?

The current Theorem 3.2 is presented as such for simplicity. However, this is not needed for Theorem 3.2 ($\nu > 1/2$).

For the Theorem 3.1 ($\nu < 1/2$), we need the assumption to get explicit tighter bounds on $E$ and $S_2$, as these terms depend on $\sigma’_{\perp}(XW^T)$. In particular, we need the entries of $XW^T$ to be gaussian with equal variance when conditioned on $W$. Hence we need all rows of $W$ to be unit norm.

However, for Theorem 3.2 ($\nu > 1/2$), we do not use these concentration inequalities. Instead, we note that since $\sigma$ is Lipschitz, $\sigma’$ is bounded, hence $\sigma’_{\perp}(XW^T)$ has uniformly bounded entries. Thus, the operator norm cannot be larger than $O(n)$. We use this in 652 and 660 in the proof of Theorem 3.2.

Furthermore, when we run experiments, this assumption is still satisfied for later steps. As discussed in line 199, we employ the common weight normalization technique from [27]. This projects the weights back onto the sphere.

[27] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In Advances in Neural Information Processing Systems

Furthermore, in your experiments, you consider the vision models. Especially, you extract the hidden states before the last layer and use those as the input $X$. But in the theory part, you consider a two-layer NNs with fixed weights $\gamma_m$ and learnable parameters $X$. Can you give more explanation on this? Like, why are they related in some sense?

These experiments were to highlight that our theory provides intuition even in cases where the theory does not strictly apply. In particular, let $F$ be a pretrained network and let $F_\ell$ be its output at some layer $\ell$. Then given image data $Z$, we consider $X = F_\ell(Z)$.

We then use this $X$ as the input to train another network $f(X) = a^T \sigma(WX^T)$ and show that the gradient of $W$ is as predicted and that regularization has the predicted qualitative effect as well.

A small typo is that you didn’t add $\gamma_m$ in your gradient formula in Proposition 2.1. Besides, you didn’t consider the effect of the regularization in your theory. It’s better to get rid of the term in your gradient formula.

We thank the reviewer for pointing out this typo. We shall fix this.

We have theory results on regularization. While Theorem 3.1 and Theorem 3.2 do not talk about the regularization, Propositions 4.1, 4.2, 4.3 analyze the effect of regularization. In particular, Propositions 4.1 and 4.3 consider regularizers of the form in Proposition 2.1. Hence we presented the gradient jointly.

We hope that our answers have cleared any doubts the reviewer had. If the reviewers has further questions. We would be happy to answer them.

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.

Positives

• 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.

• 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.

• Sound theoretical and empirical grounding: The technical assumptions are well-justified (empirically and analytically), and simulations corroborate the theoretical predictions.

Concerns of the reviewer.

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.

We thank the reviewer for pointing us to papers [1] and [2], which indeed raise important considerations about optimization dynamics in anisotropic settings.

However, 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:

Preconditioning need not kill rank 2.

In [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.

Need Dual Understanding to Isolate Important Algorithmic Aspects

Second, 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.

Beyond Single Index Model

The 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.

Broader scope than [2].

The 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.

Rank‑two can beat rank‑one even in single‑index cases.

A 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.

Regularization

Additionally, 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.

Comparison with [1]

Finally, 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.

Summary

In 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.

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.

This is a great question that is currently being explored by the authors.

Building on Moniri et al. (2024), we show that the learned feature matrix satisfies

$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$

where $\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.

We 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:

• ReLU (data‑aligned spike)
• Sigmoid (residue‑aligned spike)
• Sigmoid + Jacobian penalty (data‑aligned spike)

Inner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer. Results:

After subsequently fitting the outer layer (ridge regression), the ordering reversed:

This illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts.

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).

Yes in the following way.

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).

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.

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?

If 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.

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.

Since we consider the proportional scaling limit $n^\nu = \Theta(d^\nu)$ they are equivalent.

The rebuttal has solved most of the questions. I still keep my score to weakly accept this paper.

Thank you for your thoughtful and constructive review. We are particularly grateful for your recognition of the paper’s strengths:

Positives

• Novelty and generality: you highlight the originality of our rank‑two gradient analysis under anisotropic, ill‑conditioned data, an essential generalisation of prior work.

• Theoretical depth: you note the mathematical rigour of Theorems 3.1, 3.2 and the nuanced regularisation analysis.

• Experimental care: you find the experiments well‑designed and consistent with theory on both synthetic data and real‑world embeddings.

Below we address your main concerns.

Reviewer Concerns

One step Gradient

We 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:

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.

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.

Practical Takeaway

Note our current experiments on CIFAR-10 use fixed embeddings. However, we use the MNIST dataset with minimal whitening.

We 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.

A pragmatic rule‑of‑thumb emerging from our analysis is:

• 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.
• Leading eigendirection harmful / carries spurious signal → use a smooth activation (sigmoid, GELU) and/or inject isotropic noise to push the spike into the residue.

These choices control whether training emphasizes or suppresses the dominant data mode, offering a knob to match the assumed structure of the task.

Extension of feature‑learning dynamics. Building on Moniri et al. (2024), we show that the learned feature matrix satisfies

$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$

where $\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.

We 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:

• ReLU (data‑aligned spike)
• Sigmoid (residue‑aligned spike)
• Sigmoid + Jacobian penalty (data‑aligned spike)

Inner‑layer weights were trained in the NTK regime for 100 epochs while freezing the outer layer. Results:

After subsequently fitting the outer layer (ridge regression), the ordering reversed:

This illustrates that a residue‑aligned spike can speed early optimization but degrade generalization once the readout layer adapts.