Title: Low Rank Gradients and Where to Find Them
Venue: NeurIPS 2025
Accepted: yes
Keywords: One step gradient

low-rank gradients

two-layer neural networks

spiked covariance

regularization

high dimensional probability
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-21"}

comments: 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.

STRENGTHS AND WEAKNESSES:
## Strengths
- The manuscript is well-organized and clearly written.  
- The theoretical results are precise, rigorous, and elegantly presented.  
- Comprehensive empirical evaluations substantiate the theoretical insights.

## Weaknesses
- 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.

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
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, 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\"\n\nQUESTIONS:\nBeyond 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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).\n\nQUESTIONS:\nSee the Weakness part.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-04"}