Research map / Optimization / Implicit bias and generalization
Asymptotics of feature learning in two-layer networks after one gradient-step
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.
Models a two-layer network after one large gradient step as a spiked random features (sRF) model and derives exact replica-based asymptotics of its test error, with a conditional Gaussian equivalence for the features.
Problem
Understanding how two-layer networks learn features in one large-learning-rate gradient step and beat kernel/random-features (lazy regime) generalization in the high-dimensional proportional regime with single-index targets.
Approach
Map the first-layer weights after one large gradient step to a spiked random features model (bulk plus rank-one spike), derive scalar parameters (c, r, gamma) for it, and give a conditional Gaussian equivalent feature map conditioned on the projection along the spike, yielding replica-method asymptotic test error formulas.
Key results
- Asymptotic sRF test-error characterization closely matches learning curves of two-layer networks trained with one gradient step (d=2000 experiments).
- Even small spike strengths (r=0.2) give test errors sizably lower than vanilla random features (r=0).
- Test error decreases as alpha_0 grows, approaching the lowest error at gamma=1 as alpha_0 goes to infinity.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train the first layer with one large gradient step instead of keeping it at initialization | two-layer network · single-index targets with isotropic Gaussian inputs | test error | Helped |
| increase the spike strength | spiked Random Features · single-index targets | test error | Helped |
| use a nonzero spike strength rather than an unspiked Random Features model | spiked Random Features · single-index targets | test error: r=0.2 | Helped |
| increase the spike-target alignment | spiked Random Features · single-index sign target | test error | Helped |
| increase the sample complexity of the first-step representation-learning batch | two-layer network trained with one gradient step · single-index targets | test error | Helped |
| use random features at initialization (zero first-layer training) instead of training the first layer | two-layer network · sine target with isotropic Gaussian inputs | test error | Hurt |
| replace the spiked Random Features model with the conditional Gaussian feature map | spiked Random Features · high-dimensional proportional regime | test error | No clear effect |
| use non-uniform readout initialization with a finite vocabulary instead of uniform readout initialization | two-layer network after one gradient step · single-index targets | expressivity: 2|V| | Helped |
Surprising findings
- Tuning the learning rate can make the upper bound on test error zero for erf targets (perfect learning), per the discussion of approximation subspaces.
Limitations
- Replica method is non-rigorous; rigorous proof left open.
- Bulk approximated as random rows on the sphere, ignoring possible row covariance structure.
- Main result covers uniform readout initialization, one gradient step, and single-index targets.
Experimental setup
Datasets: synthetic single-index Gaussian data
Metrics: test error (generalization error) mean squared error
Baselines: random features (lazy regime, eta=0) kernel/linear methods lower bound h2*-(h1*)^2 NTK
Builds on: Ba et al. 2022 one-step gradient feature learning Gaussian equivalence principle for random features replica method
Closest papers by meaning
- Asymptotic Analysis of Two-Layer Neural Networks after One Gradient Step under Gaussian Mixtures Data with Structure (ICLR 2025 · similarity 0.92)
- A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks (ICML 2024 · similarity 0.89)
- Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks (NeurIPS 2025 · similarity 0.89)
- Asymptotics of Learning with Deep Structured (Random) Features (ICML 2024 · similarity 0.88)
- Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization (NeurIPS 2024 · similarity 0.88)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.87)
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.