Research map / Optimization / Optimizers
Surge Phenomenon in Optimal Learning Rate and Batch Size 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.
Derives a scaling law for optimal learning rate vs batch size under Adam-style sign-of-gradient updates, showing the optimum rises then falls (peak at the noise scale), and validates it with grid searches on CNN, ResNet-18, DistilGPT2 and MoE workloads.
Problem
Existing batch-size/learning-rate scaling rules (square-root, linear, critical batch size models) were derived for SGD and do not describe the optimal learning rate for Adam-style sign-based optimizers, which are the dominant choice in practice.
Approach
Model Adam-style updates as sign(mini-batch gradient) with Gaussian gradient noise, derive a closed-form optimal learning rate as a function of batch size that first increases then decreases ("surge"), with the peak at the gradient noise scale B_noise, and link the peak to the training-speed/data-efficiency trade-off.
Key results
- Theory predicts the optimal Adam learning rate rises with batch size to a peak at B_noise and then declines, with epsilon_opt approx 2*epsilon_max/sqrt(B_noise)*sqrt(B) for small B.
- Estimated B_noise increases as training loss decreases, so the peak optimal learning rate shifts to larger batch sizes during training (observed on ResNet-18/TinyImageNet).
- Experiments on CNN-FashionMNIST, ResNet-18-TinyImageNet, DistilGPT2-ELI5 and MoE-RedPajama-v2 show the observed optimum follows the predicted surge curve better than SGD-style alpha=0.5 or alpha=1 fits.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the proposed Adam optimal-learning-rate scaling law instead of the fitted SGD scaling curves | Adam · CNN-FashionMNIST | optimal learning rate vs. batch size | Helped |
| Use the proposed Adam optimal-learning-rate scaling law instead of the fitted SGD scaling curves | Adam · ResNet18-TinyImageNet | optimal learning rate vs. batch size | Helped |
| Use the proposed Adam optimal-learning-rate scaling law | Adam · DistilGPT2-ELI5-Category | optimal learning rate vs. batch size | Helped |
| Use the proposed Adam optimal-learning-rate scaling law | Adam · MoE-RedPajama-v2 | optimal learning rate vs. batch size | Helped |
| Use the fitted SGD curve with α=0.5 to model Adam optimal learning rates | Adam · ResNet18-TinyImageNet | optimal learning rate vs. batch size | Hurt |
| Use the fitted SGD curve with α=0.5 to model Adam optimal learning rates | Adam · DistilGPT2-ELI5-Category | optimal learning rate vs. batch size | No clear effect |
| Vary Adam configuration from β₁=0.0, β₂=0.0 to β₁=0.9, β₂=0.999 (ablation) | Adam · DistilGPT2-ELI5-Category | agreement with the predicted optimal-learning-rate curve | No clear effect |
| Increase batch size beyond the peak-learning-rate batch size | Adam · Adam-style optimization | optimal learning rate | Hurt |
| Continue training as training loss decreases | Adam · ResNet18-TinyImageNet | batch size at peak optimal learning rate | Helped |
Surprising findings
- The optimal learning rate for Adam-style optimizers is not monotone in batch size: it declines after peaking, contradicting prior square-root scaling claims.
- The stable large-batch learning rate may exceed the local peak in later training stages, since gradient variance is small there.
Limitations
- Theory relies on a quadratic loss approximation and sign-of-gradient approximation of Adam, which may not hold for large learning rates.
- Assumes Gaussian gradient distributions per parameter via CLT; does not account for weight decay, gradient clipping or other configurations.
- Experiments are limited to four workloads with a single A100 per run and grid-search-based fitting.
Experimental setup
Datasets: Fashion-MNIST Tiny-ImageNet ELI5-Category RedPajama-v2
Benchmarks: Fashion-MNIST Tiny-ImageNet ELI5-Category RedPajama-v2
Metrics: training loss reduction optimal learning rate vs batch size steps S and examples E to reach target loss
Baselines: SGD-style scaling with alpha=1 (linear scaling) Adam scaling with alpha=0.5 (square-root scaling)
Builds on: empirical model of large-batch training (McCandlish et al.) Adam critical batch size scaling laws
Closest papers by meaning
- Understanding the Generalization of Stochastic Gradient Adam in Learning Neural Networks (NeurIPS 2025 · similarity 0.89)
- Scaling Exponents Across Parameterizations and Optimizers (ICML 2024 · similarity 0.88)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.88)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.87)
- Adam-mini: Use Fewer Learning Rates To Gain More (ICLR 2025 · similarity 0.87)
- A Comprehensive Framework for Analyzing the Convergence of Adam: Bridging the Gap with SGD (ICML 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.