Research map / Architectures and efficiency / RNNs and normalization
On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling
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 under cross-entropy loss, standard parameterization with large learning rates (n^{-1/2}) sits in a "controlled divergence" regime where logits diverge but hidden features still learn, giving the first practical infinite-width limit with feature learning for SP.
Problem
Infinite-width theory predicts SP becomes unstable at learning rates above O(1/n) and loses feature learning at O(1/n), yet optimal learning rates in practice decay like ~1/sqrt(n) and networks train stably with feature learning at large width.
Approach
Splits the previously 'unstable' regime under CE loss into a catastrophic regime and a controlled divergence regime where logits diverge but gradients and activations stay stable; at the edge (eta ~ n^{-1/2} for hidden layers) hidden-layer features evolve in the infinite-width limit. Uses refined coordinate checks to disentangle effective and propagating updates.
Key results
- Under CE loss, SP hidden layers learn features width-independently at learning rate eta ~ n^{-1/2}, while logits diverge as Theta(n^{1-alpha}); MSE loss instead stays stable only with n^{-1} rates.
- Empirical maximal stable and optimal learning rate exponents in MLPs and GPT under SGD closely follow alpha=1/2 under CE loss.
- Under MSE loss in SP, feature learning is lost and training is suboptimal, whereas under muP the MSE loss becomes competitive with CE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CE instead of MSE loss | SP training of deep nonlinear networks · Deep MLPs, width scaling | maximal stable learning-rate scaling: n^{-1} -> n^{-1/2} | Helped |
| Use CE instead of MSE loss | SP training · Deep MLPs, controlled divergence regime | hidden-layer feature learning: O(n^{-1/2}) -> Θ(n^{1/2-α}) | Helped |
| Set the learning-rate exponent to α = 1/2 | SP with CE · Deep MLPs, infinite-width limit | hidden-layer activation-update scaling: Θ(1) | Helped |
| Use MSE instead of CE loss | SP training · Deep MLPs, α < 1 | training stability | Hurt |
| Remove trainable Layernorm parameters (ablation) | GPT trained with Adam in SP · Language data | maximal stable learning-rate scaling: η_n≈Θ(n^{-1/2}) -> η_n=O(n^{-1}) | Hurt |
| Use μP learning rates rather than SP global learning rates | Image MLPs trained with MSE · MNIST, CIFAR-10, Fashion-MNIST | training accuracy | Helped |
| Use MSE instead of CE loss | SP · Image datasets | training accuracy | Hurt |
| Use MSE instead of CE loss | μP · Image datasets | training accuracy | Helped |
Surprising findings
- Infinite-width alignment predictions between weights and incoming activations hold at moderate width under sufficiently refined coordinate checks, contrary to prior hypotheses.
- Finite-width effects, depth, and catapult/edge-of-stability dynamics alone do not explain the stability of large learning rates in SP.
Limitations
- Results focus on single-epoch training; multi-epoch setting left to future work.
- Layerwise SP-full-align learning-rate transfer breaks on image datasets and the mechanism is not fully understood.
- Numerical precision is assumed sufficient; floating point range may be exceeded at moderate width under fast divergence.
Experimental setup
Datasets: CIFAR-10 MNIST FashionMNIST DCLM-Baseline
Benchmarks: CIFAR-10 MNIST FashionMNIST DCLM-Baseline
Metrics: optimal learning rate exponent maximal stable learning rate exponent training loss training accuracy update scaling (RMS norm)
Baselines: maximal update parameterization SP-full-align standard parameterization with MSE loss neural tangent parameterization
Builds on: Tensor Programs Maximal Update Parameterization Edge of stability Catapult mechanism
Closest papers by meaning
- Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks (NeurIPS 2026 · similarity 0.91)
- A Dynamical Model of Neural Scaling Laws (ICML 2024 · similarity 0.90)
- Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs (NeurIPS 2025 · similarity 0.90)
- The Optimization Landscape of SGD Across the Feature Learning Strength (ICLR 2025 · similarity 0.89)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.89)
- Feature Learning Dynamics in Infinite-Depth Neural Networks (NeurIPS 2026 · 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.