Research map / Architectures and efficiency / RNNs and normalization
The Optimization Landscape of SGD Across the Feature Learning Strength
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.
Empirically and theoretically maps how the feature-learning strength gamma (final-layer down-scaling in muP) interacts with learning rate in online SGD training, finding optimal learning rates scale as gamma^2 (lazy) and gamma^{2/L} (ultra-rich), with a linear-network toy model explaining the phases.
Problem
How the feature-learning strength gamma, which down-scales the output layer in muP networks, interacts with the learning rate and affects optimization and online generalization across the lazy, rich and ultra-rich regimes.
Approach
Sweeps gamma and learning rate jointly across MLPs, CNNs, ResNets and ViTs in the online (single-pass) setting, identifies scaling regimes eta* ~ gamma^2 (small gamma) and eta* ~ gamma^{2/L} (large gamma), and explains them with a one-parameter deep linear network toy model; also studies silent alignment, staircase loss drops, and progressive sharpening at large gamma.
Key results
- Optimal learning rate scales as eta* ~ gamma^2 for gamma << 1 and eta* ~ gamma^{2/L} for gamma >> 1 for depth-L feed-forward networks.
- Large-gamma networks match or exceed gamma=1 test accuracy when trained long enough in the online setting, with marginal returns past some point.
- Two regimes of sharpness at end of training: lambda_max scales as gamma^{-2} for small gamma and gamma^{-2/L} for large gamma in MLPs and CNNs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a larger γ with the learning rate scaled appropriately | Online-trained neural networks · MNIST-1M, CIFAR-5M, TinyImageNet, and other studied tasks | test accuracy | Helped |
| Train sufficiently long at large γ rather than stopping early | Online-trained neural networks · Studied tasks | test accuracy | Helped |
| Scale the learning rate with γ in the ultra-rich regime | Feed-forward networks trained with SGD · Large γ (γ ≫ 1) | maximum convergent learning rate: \eta_* \propto \gamma^{2/L} | Helped |
| Scale the learning rate with γ in the lazy regime | Neural networks trained with SGD · Small γ (γ ≪ 1) | maximum convergent learning rate: \eta_* \propto \gamma^2 | Helped |
| Increase γ while keeping training time insufficient | Online-trained neural networks · Studied tasks | test accuracy | Mixed |
| Use smaller γ in the lazy regime | Neural networks trained with SGD · Small γ (γ ≪ 1) | catapult loss increase | Hurt |
| Add learning-rate warmup and cosine decay | SGD networks · CIFAR-5M, MSE and cross-entropy | γ–η scaling relationships | No clear effect |
| Replace SGD with Adam | Neural networks and the toy model · Lazy and ultra-rich γ regimes | learning-rate scaling: \eta_\text{min}, \eta_\text{max} \sim \gamma when \gamma \ll 1; \eta_\text{min}, \eta_\text{max} \sim \gamma^{1/L} when | Mixed |
| Replace SGD with SignSGD | Toy deep linear network · Ultra-rich γ (γ ≫ 1) | minimum and maximum viable learning-rate scaling: \eta_\text{min}, \eta_\text{max} \sim \gamma^{1/L} | Helped |
| Replace SGD with SignSGD or Adam | Neural networks · Large γ | silent alignment | Hurt |
Surprising findings
- Large-gamma networks show no catapult effects and instead exhibit silent alignment and staircase loss drops even in realistic networks.
- Offline-setting findings that strong feature learning degrades performance do not hold in the online setting; large gamma can be best.
Limitations
- Theory relies on a width-one deep linear network trained on a single example, with a heuristic one-parameter model.
- Authors observe that the lazy regime needs double precision for gamma below 1e-3, a numerical stability issue.
- Analysis restricted to SGD; Adam results only from SignSGD toy model.
Experimental setup
Datasets: MNIST-1M CIFAR-5M CIFAR-10 TinyImageNet
Benchmarks: MNIST CIFAR-10 CIFAR-5M TinyImageNet
Metrics: final test accuracy online training loss hessian top eigenvalues kernel-target alignment
Baselines: lazy limit (small gamma) gamma=1 networks
Builds on: muP neural tangent kernel catapult effect silent alignment edge of stability
Closest papers by meaning
- A Dynamical Model of Neural Scaling Laws (ICML 2024 · similarity 0.90)
- Emergence and scaling laws in SGD learning of shallow neural networks (NeurIPS 2025 · similarity 0.90)
- Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks (NeurIPS 2026 · similarity 0.89)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.89)
- Feature Learning Dynamics in Infinite-Depth Neural Networks (NeurIPS 2026 · similarity 0.89)
- How Feature Learning Can Improve Neural Scaling Laws (ICLR 2025 · similarity 0.89)
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.