Research map / Optimization / Nonconvex and smooth optimization
The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training
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 LLM and vision learning-rate schedules (cosine, WSD, linear cooldown) match the last-iterate suboptimality bound of non-smooth stochastic convex SGD, derives a log-free bound for WSD, and uses the bound to transfer schedules and learning rates for continued training.
Problem
Learning-rate schedules for large model training are chosen empirically with trial-and-error; there is little theory explaining why schedules like cosine and WSD work, why cooldown helps, and how optimal base learning rates transfer across schedules and horizons.
Approach
Use a last-iterate convergence bound for SGD with general step-size schedules from non-smooth convex optimization (Orabona-style Thm 10) as a proxy for training loss curves; the bound reproduces the shape of cosine and WSD loss curves, the WSD cooldown removes logarithmic terms, and the bound is used to adapt schedules for continued training and to transfer base learning rates across schedules.
Key results
- Theory-guided schedule adaptation for continued training gives a slight but noticeable validation-loss improvement for 124M and 210M Llama models at 100k and 200k steps, with a sudden drop after decreasing the schedule at 40k steps.
- WSD bound removes the logarithmic H_{T-1} term present in the constant-schedule bound, and the optimal base learning rate for cosine is roughly twice that for WSD (0.92/0.47), matching empirical observations.
- Learning-rate transfer from 20% cooldown to linear decay via the bound's predicted factor exp(0.7) gives final validation loss 2.9535 for linear decay vs 2.9660 for the best 20% cooldown run (124M, 50k steps).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Decrease the learning rate after 40k steps when extending WSD training | WSD continued training · SlimPajama, 124M and 210M Llama-style models, 100k and 200k steps | validation loss: around 0.01 | Helped |
| Transfer the tuned learning rate from 20% cooldown to linear decay | 124M Llama-style training on SlimPajama · SlimPajama, 50k steps | final validation loss: 2.9660 -> 2.9535 | Helped |
| Transfer the learning rate to linear decay | ResNet50 training · ImageNet | validation accuracy | Helped |
| Use WSD with 20% cooldown instead of cosine | Language-model training · OpenWebText2, three model sizes | loss | No clear effect |
| Use a last-iterate bound instead of a minimum-suboptimality bound (ablation) | Theoretical schedule comparison · Convex SGD bound | bound shape match to empirical loss curves | Helped |
| Use WSD cooldown instead of a constant schedule | Convex SGD bound · Non-smooth stochastic convex optimization | suboptimality bound | Helped |
| Fully tune the base learning rate and set cooldown fraction to 1 (ablation) | WSD schedule · Theoretical bound simulation | suboptimality bound | Helped |
| Reduce the cooldown fraction below 1 while holding the base learning rate fixed (ablation) | WSD schedule · Theoretical bound simulation | suboptimality bound | Mixed |
| Use non-vanishing rather than decaying gradient-norm bounds (ablation) | WSD schedule · Theoretical bound simulation | cooldown loss-drop visibility | Helped |
| Use a cosine cycle length of one instead of shorter cycles (ablation) | Cosine schedule · Theoretical schedule comparison | suboptimality bound | Helped |
Surprising findings
- The convex SGD bound, derived for non-smooth stochastic convex problems, reproduces the shape of non-convex LLM loss curves, including the sudden WSD drop during cooldown.
- Linear-decay schedule with transferred learning rate outperforms the best tuned 20% cooldown WSD run without any additional sweep.
Limitations
- Theory is for SGD and convexity, but experiments mainly use AdamW on non-convex transformers; the bound cannot capture optimizer differences.
- The bound is an upper bound and matching is qualitative; test loss is used as proxy, which assumes negligible generalization gap.
- Gradient norm bounds G_t and D are unknown in practice, and the bound-minimizing schedule assumes gradient norms independent of the schedule.
Experimental setup
Datasets: SlimPajama SlimPajama-6B ImageNet OpenWebText2
Benchmarks: validation loss on SlimPajama ImageNet top-1 validation accuracy OpenWebText2 validation loss
Metrics: validation loss suboptimality bound top-1 accuracy
Baselines: cosine schedule constant schedule linear-decay schedule 1/sqrt schedule 1-sqrt schedule
Builds on: Orabona last-iterate bound for SGD WSD schedule cosine learning-rate schedule AdamW Chinchilla scaling laws
Closest papers by meaning
- The Road Less Scheduled (NeurIPS 2024 · similarity 0.88)
- Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers (NeurIPS 2026 · similarity 0.87)
- Capacity-Constrained Online Convex Optimization with Delayed Feedback (NeurIPS 2026 · similarity 0.86)
- Information Complexity of Stochastic Convex Optimization: Applications to Generalization, Memorization, and Tracing (ICML 2024 · similarity 0.86)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.86)
- A Reduction from Delayed to Immediate Feedback for Online Convex Optimization with Improved Guarantees (NeurIPS 2026 · 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.