Research map / Optimization / Implicit bias and generalization
A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that one gradient step with a learning rate growing as n^alpha on a two-layer network's first layer adds one spike per polynomial degree to the feature matrix spectrum, and that these spikes precisely characterize the limiting train/test error improvements, including learning quadratic components.
Problem
Single-step gradient feature learning in two-layer networks with a constant step size only learns the linear component of the teacher; it is unknown which non-linear features are learned when the step size grows with sample size, and how they affect training and test error.
Approach
Use step size eta ~ n^alpha; the updated feature matrix is approximated by the untrained features plus rank-one spikes aligned with the monomials (X beta)^{∘k}, one per k up to a number that depends on alpha. Equivalence theorems replace the trained features with this approximation, and Gaussian equivalence is then applied to the untrained part to derive exact high-dimensional training and test er
Key results
- For eta ~ n^alpha with alpha in ((l-1)/(2l), l/(2l+2)), the feature matrix has l added spikes aligned with (X beta)^{∘k}, k=1..l, up to o(sqrt n) error in operator norm.
- The training-loss improvement for l=1 is Delta_1 = psi*lambda*c_{*,1}^4*m_2 / (phi[c_{*,1}^2+phi(c_*^2+sigma_eps^2)]) and l=2 adds a term proportional to c_{*,2}^2, so quadratic components are learned only when l>=2.
- Simulations with eta=n^0.29 (l=2) show two isolated spikes aligned with X beta and (X beta)^{∘2}; the train/test error curves diverge between teacher settings with and without a quadratic component.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase step size from constant (eta=O(1)) to eta ~ n^alpha with alpha in ((ell-1)/(2ell), ell/(2ell+2)) | one-step gradient update of first layer in two-layer network · high-dimensional proportional asymptotics | number of separated spikes in feature matrix spectrum: ell spikes added for alpha in that range | Helped |
| Step size eta ~ n^alpha with ell=1 regime (linear feature only) | training loss of ridge regression on learned features · theory, single-index teacher | training loss reduction Delta_1: Delta_1 >= 0, increasing in c_{*,1} | Helped |
| Step size moved into ell=2 regime (quadratic features learned) | training loss of ridge regression on learned features · theory, single-index teacher | additional training loss reduction Delta_2 vs Delta_1: additional positive term depending on c_{*,2} | Helped |
| Step size moved into ell=2 regime (quadratic features learned) | test error of ridge regression on learned features · theory, single-index teacher | test error improvement Lambda_2 vs Lambda_1: Lambda_2 = Lambda_1 + positive term in c_{*,2}^2 | Helped |
| Teacher with information exponent >1 (c_{*,1}=0) with gradient step | limiting training loss · theory | training loss change | No clear effect |
| Constant step size (eta=O(1)) one-step update | non-linear teacher components · theory | learning of non-linear components | No clear effect |
| Shifted-ReLU experiments: Setting 2 (H1 + H2/sqrt2 target) vs Setting 1 (linear only) across log(eta)/log(n) (ablation) | train/test error of one-step updated network · synthetic, n=1000, N=500, d=300 | train/test error vs log(eta)/log(n) | Mixed |
Surprising findings
- With a constant step size (l=1 regime) the one-step update does not improve learning of non-linear teacher components, only the linear part.
Limitations
- Analysis restricted to alpha < 1/2 and to step sizes not at the boundaries alpha = (l-1)/(2l), leaving boundary case open.
- Uses sample splitting between the gradient step and ridge regression, not covering the case where both use the same data.
- Requires smooth activations with c_0 = 0 (zero-mean) and Hermite coefficient decay conditions; teacher is restricted to single-index models.
Experimental setup
Datasets: synthetic Gaussian data with single-index teacher
Metrics: training error test error singular value spectrum
Baselines: constant step size one-step GD (Ba et al. 2022) random features model (eta = 0) linear and kernel predictors
Builds on: Gaussian equivalence property random features models one-step gradient feature learning (Ba et al.) Wedin's theorem
Closest papers by meaning
- How does Gradient Descent Learn Features --- A Local Analysis for Regularized Two-Layer Neural Networks (NeurIPS 2024 · similarity 0.90)
- Asymptotics of feature learning in two-layer networks after one gradient-step (ICML 2024 · similarity 0.89)
- Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks (NeurIPS 2025 · similarity 0.88)
- Asymptotic Analysis of Two-Layer Neural Networks after One Gradient Step under Gaussian Mixtures Data with Structure (ICLR 2025 · similarity 0.88)
- Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.87)
- Low Rank Gradients and Where to Find Them (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.