Research map / Architectures and efficiency / RNNs and normalization
Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning
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.
Derives exact gradient-flow solutions for a single-neuron two-layer linear network with layer-specific learning rates and conserved quantity delta, extends to wide/deep linear and piecewise-linear nets, and shows upstream (earlier-layer-faster) initializations accelerate rich feature learning.
Problem
Understanding how layer-specific initialization variances and learning rates determine the transition between lazy and rich feature learning in finite-width networks.
Approach
A conserved quantity delta = eta_w a^2 - eta_a ||w||^2 governs the geometry of gradient-flow trajectories; exact solutions via hyperbolic/spherical coordinates show upstream initializations (delta>0) give rapid rich learning in nonlinear nets, while in linear nets only balanced initializations do.
Key results
- Upstream initialization (delta>0) with small overall scale tau gives the smallest test loss and rapid rich learning in the teacher-student ReLU setup with d=100.
- Decreasing the variance of the embedding in a single-layer transformer (<6% of parameters) significantly reduces time to grokking in modular arithmetic.
- In linear networks rapid feature learning occurs only from balanced initializations; in nonlinear networks upstream initializations accelerate rich learning.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a small-overall-scale upstream initialization, promoting faster learning in earlier layers | Two-layer ReLU network · XOR-like task | activation-pattern change | Helped |
| Use an upstream initialization instead of a downstream initialization | Two-layer ReLU network · XOR-like task | parameter movement required to fit the data | Helped |
| Adjust the first layer's learning speed to promote faster early-layer learning | CNNs · Image classification tasks | early-layer filter interpretability | Helped |
| Modify the relative initialization scale toward an upstream initialization | Deep learning models · Random Hierarchy Model | sample complexity | Helped |
| Modify the relative initialization scale | Deep learning models · Random Hierarchy Model | sample complexity | Mixed |
| Decrease the embedding initialization variance | Single-layer transformer · Modular arithmetic | time to grokking | Helped |
| Use an upstream initialization | Wide two-layer student ReLU network · Teacher-generated data | test loss | Helped |
| Use a downstream rather than upstream initialization | Wide two-layer student ReLU network · Teacher-generated data | alignment to the teacher | Hurt |
Surprising findings
- Upstream initializations can drive rich learning via changing activation patterns while parameters move little, challenging the view that rich learning requires large parameter movement.
- Downstream initializations at small scale yield a delayed rich regime, with the delay set by the magnitude of delta.
Limitations
- Extending the theory to deeper nonlinear networks is difficult.
- Analysis omits discretization and stochastic effects of SGD, which break the conservation laws.
- Exact solutions are for a single-neuron linear model and two-layer piecewise-linear networks.
Experimental setup
Datasets: synthetic teacher-student ReLU data XOR-like task Random Hierarchy Model modular arithmetic
Benchmarks: Random Hierarchy Model modular arithmetic grokking
Metrics: test loss kernel distance Hamming distance in activation patterns parameter distance time to grokking sample complexity
Baselines: standard parameterization NTK parameterization mean-field parameterization
Builds on: Saxe et al. exact linear dynamics Chizat and Bach lazy training Kunin et al. balanced/unbalanced initialization Woodworth et al. kernel vs rich regimes Atanasov et al. infinite-width phase diagram
Closest papers by meaning
- From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks (ICLR 2025 · similarity 0.92)
- Three Mechanisms of Feature Learning in a Linear Network (ICLR 2025 · similarity 0.88)
- Feature Learning beyond the Lazy-Rich Dichotomy: Insights from Representational Geometry (ICML 2025 · similarity 0.87)
- Feature Learning Dynamics in Infinite-Depth Neural Networks (NeurIPS 2026 · similarity 0.87)
- From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning (ICML 2025 · similarity 0.87)
- The Optimization Landscape of SGD Across the Feature Learning Strength (ICLR 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.