Warmup-Stable-Decay Learning Rates and the River Valley Loss Landscape
Nov 14, 2026
This episode examines the Warmup-Stable-Decay learning rate schedule and the "river valley" loss landscape conjecture meant to explain it. The conjecture pictures a canyon with steep walls and a gently sloping floor. A large learning rate makes the optimizer bounce between the walls, which keeps training loss looking high. The same large rate still makes fast progress along the floor, so the sharp drop during the final decay phase reveals progress already made. The hosts contrast WSD with cosine decay, which fixes the endpoint in advance and forces a fresh run for every compute budget. WSD instead keeps a constant-rate main branch and forks short decay runs to get checkpoints for many budgets, which helps continual training and scaling-law studies. The hosts also stress that the canyon is a conjecture and not a measured property of full-scale LLMs, and they set up two open questions: whether the theory holds for Adam and whether anyone has seen this structure in a real model.
Interactive Visualization: Warmup-Stable-Decay Learning Rates and the River Valley Loss Landscape
Sources:
1. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective — Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma, 2024
http://arxiv.org/abs/2410.05192v1
2. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective — Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma, 2024
https://scholar.google.com/scholar?q=Understanding+Warmup-Stable-Decay+Learning+Rates%3A+A+River+Valley+Loss+Landscape+Perspective
3. Adaptive Gradient Methods at the Edge of Stability — Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E. Dahl, Justin Gilmer, 2022
https://scholar.google.com/scholar?q=Adaptive+Gradient+Methods+at+the+Edge+of+Stability
4. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability — Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, Ameet Talwalkar, 2021
https://scholar.google.com/scholar?q=Gradient+Descent+on+Neural+Networks+Typically+Occurs+at+the+Edge+of+Stability
5. An Investigation into Neural Net Optimization via Hessian Eigenvalue Density — Behrooz Ghorbani, Shankar Krishnan, Ying Xiao, 2019
https://scholar.google.com/scholar?q=An+Investigation+into+Neural+Net+Optimization+via+Hessian+Eigenvalue+Density
6. MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies — Shengding Hu, Yuge Tu, Xu Han, et al., 2024
https://scholar.google.com/scholar?q=MiniCPM%3A+Unveiling+the+Potential+of+Small+Language+Models+with+Scalable+Training+Strategies
7. SGDR: Stochastic Gradient Descent with Warm Restarts — Ilya Loshchilov, Frank Hutter, 2017
https://scholar.google.com/scholar?q=SGDR%3A+Stochastic+Gradient+Descent+with+Warm+Restarts
8. Training Compute-Optimal Large Language Models — Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, et al., 2022
https://scholar.google.com/scholar?q=Training+Compute-Optimal+Large+Language+Models
9. The Road Less Scheduled — Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, 2024
https://scholar.google.com/scholar?q=The+Road+Less+Scheduled
10. The Large Learning Rate Phase of Deep Learning: the Catapult Mechanism — Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, Guy Gur-Ari, 2020
https://scholar.google.com/scholar?q=The+Large+Learning+Rate+Phase+of+Deep+Learning%3A+the+Catapult+Mechanism
11. Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability — Alex Damian, Eshaan Nichani, Jason D. Lee, 2022
https://scholar.google.com/scholar?q=Self-Stabilization%3A+The+Implicit+Bias+of+Gradient+Descent+at+the+Edge+of+Stability
12. Three Factors Influencing Minima in SGD — Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Amos Storkey, Yoshua Bengio, Yoshua Bengio (with colleagues), 2017
https://scholar.google.com/scholar?q=Three+Factors+Influencing+Minima+in+SGD
13. Don't Decay the Learning Rate, Increase the Batch Size — Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, Quoc V. Le, 2017
https://scholar.google.com/scholar?q=Don%27t+Decay+the+Learning+Rate%2C+Increase+the+Batch+Size
14. Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations — Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro von Werra, Martin Jaggi, 2024
https://scholar.google.com/scholar?q=Scaling+Laws+and+Compute-Optimal+Training+Beyond+Fixed+Training+Durations
15. Simple and Scalable Strategies to Continually Pre-train Large Language Models — Adam Ibrahim et al., 2024
https://scholar.google.com/scholar?q=Simple+and+Scalable+Strategies+to+Continually+Pre-train+Large+Language+Models
16. Scaling Laws for Learning Rate Schedules / Loss curve scaling law under LR annealing — Tissue et al., 2024
https://scholar.google.com/scholar?q=Scaling+Laws+for+Learning+Rate+Schedules+%2F+Loss+curve+scaling+law+under+LR+annealing
17. The Road Less Scheduled (Schedule-Free Learning) — Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, 2024
https://scholar.google.com/scholar?q=The+Road+Less+Scheduled+%28Schedule-Free+Learning%29
18. Why Warmup the Learning Rate? / On the Training Instability of Shuffling SGD, and Why Adam Needs Warmup — Kalra & Barkeshli; Gilmer et al., 2023-2024
https://scholar.google.com/scholar?q=Why+Warmup+the+Learning+Rate%3F+%2F+On+the+Training+Instability+of+Shuffling+SGD%2C+and+Why+Adam+Needs+Warmup
19. Does your LLM truly unlearn? / Why Transformers Need Adam: A Hessian Perspective — Yushun Zhang et al., 2024
https://scholar.google.com/scholar?q=Does+your+LLM+truly+unlearn%3F+%2F+Why+Transformers+Need+Adam%3A+A+Hessian+Perspective
20. Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws — Zeyuan Allen-Zhu, Yuanzhi Li, 2024
https://scholar.google.com/scholar?q=Physics+of+Language+Models%3A+Part+3.3%2C+Knowledge+Capacity+Scaling+Laws
21. Beyond neural scaling laws / DeepSeek LLM: Scaling Open-Source Language Models with Longtermism — DeepSeek-AI, 2024
https://scholar.google.com/scholar?q=Beyond+neural+scaling+laws+%2F+DeepSeek+LLM%3A+Scaling+Open-Source+Language+Models+with+Longtermism
22. Language Models are Few-Shot Learners / Chinchilla (Training Compute-Optimal LLMs) — Hoffmann et al., 2022
https://scholar.google.com/scholar?q=Language+Models+are+Few-Shot+Learners+%2F+Chinchilla+%28Training+Compute-Optimal+LLMs%29
Interactive Visualization: Warmup-Stable-Decay Learning Rates and the River Valley Loss Landscape
1. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective — Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma, 2024
http://arxiv.org/abs/2410.05192v1
2. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective — Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma, 2024
https://scholar.google.com/scholar?q=Understanding+Warmup-Stable-Decay+Learning+Rates%3A+A+River+Valley+Loss+Landscape+Perspective
3. Adaptive Gradient Methods at the Edge of Stability — Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E. Dahl, Justin Gilmer, 2022
https://scholar.google.com/scholar?q=Adaptive+Gradient+Methods+at+the+Edge+of+Stability
4. Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability — Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, Ameet Talwalkar, 2021
https://scholar.google.com/scholar?q=Gradient+Descent+on+Neural+Networks+Typically+Occurs+at+the+Edge+of+Stability
5. An Investigation into Neural Net Optimization via Hessian Eigenvalue Density — Behrooz Ghorbani, Shankar Krishnan, Ying Xiao, 2019
https://scholar.google.com/scholar?q=An+Investigation+into+Neural+Net+Optimization+via+Hessian+Eigenvalue+Density
6. MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies — Shengding Hu, Yuge Tu, Xu Han, et al., 2024
https://scholar.google.com/scholar?q=MiniCPM%3A+Unveiling+the+Potential+of+Small+Language+Models+with+Scalable+Training+Strategies
7. SGDR: Stochastic Gradient Descent with Warm Restarts — Ilya Loshchilov, Frank Hutter, 2017
https://scholar.google.com/scholar?q=SGDR%3A+Stochastic+Gradient+Descent+with+Warm+Restarts
8. Training Compute-Optimal Large Language Models — Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, et al., 2022
https://scholar.google.com/scholar?q=Training+Compute-Optimal+Large+Language+Models
9. The Road Less Scheduled — Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, 2024
https://scholar.google.com/scholar?q=The+Road+Less+Scheduled
10. The Large Learning Rate Phase of Deep Learning: the Catapult Mechanism — Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, Guy Gur-Ari, 2020
https://scholar.google.com/scholar?q=The+Large+Learning+Rate+Phase+of+Deep+Learning%3A+the+Catapult+Mechanism
11. Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability — Alex Damian, Eshaan Nichani, Jason D. Lee, 2022
https://scholar.google.com/scholar?q=Self-Stabilization%3A+The+Implicit+Bias+of+Gradient+Descent+at+the+Edge+of+Stability
12. Three Factors Influencing Minima in SGD — Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Amos Storkey, Yoshua Bengio, Yoshua Bengio (with colleagues), 2017
https://scholar.google.com/scholar?q=Three+Factors+Influencing+Minima+in+SGD
13. Don't Decay the Learning Rate, Increase the Batch Size — Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, Quoc V. Le, 2017
https://scholar.google.com/scholar?q=Don%27t+Decay+the+Learning+Rate%2C+Increase+the+Batch+Size
14. Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations — Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro von Werra, Martin Jaggi, 2024
https://scholar.google.com/scholar?q=Scaling+Laws+and+Compute-Optimal+Training+Beyond+Fixed+Training+Durations
15. Simple and Scalable Strategies to Continually Pre-train Large Language Models — Adam Ibrahim et al., 2024
https://scholar.google.com/scholar?q=Simple+and+Scalable+Strategies+to+Continually+Pre-train+Large+Language+Models
16. Scaling Laws for Learning Rate Schedules / Loss curve scaling law under LR annealing — Tissue et al., 2024
https://scholar.google.com/scholar?q=Scaling+Laws+for+Learning+Rate+Schedules+%2F+Loss+curve+scaling+law+under+LR+annealing
17. The Road Less Scheduled (Schedule-Free Learning) — Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, 2024
https://scholar.google.com/scholar?q=The+Road+Less+Scheduled+%28Schedule-Free+Learning%29
18. Why Warmup the Learning Rate? / On the Training Instability of Shuffling SGD, and Why Adam Needs Warmup — Kalra & Barkeshli; Gilmer et al., 2023-2024
https://scholar.google.com/scholar?q=Why+Warmup+the+Learning+Rate%3F+%2F+On+the+Training+Instability+of+Shuffling+SGD%2C+and+Why+Adam+Needs+Warmup
19. Does your LLM truly unlearn? / Why Transformers Need Adam: A Hessian Perspective — Yushun Zhang et al., 2024
https://scholar.google.com/scholar?q=Does+your+LLM+truly+unlearn%3F+%2F+Why+Transformers+Need+Adam%3A+A+Hessian+Perspective
20. Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws — Zeyuan Allen-Zhu, Yuanzhi Li, 2024
https://scholar.google.com/scholar?q=Physics+of+Language+Models%3A+Part+3.3%2C+Knowledge+Capacity+Scaling+Laws
21. Beyond neural scaling laws / DeepSeek LLM: Scaling Open-Source Language Models with Longtermism — DeepSeek-AI, 2024
https://scholar.google.com/scholar?q=Beyond+neural+scaling+laws+%2F+DeepSeek+LLM%3A+Scaling+Open-Source+Language+Models+with+Longtermism
22. Language Models are Few-Shot Learners / Chinchilla (Training Compute-Optimal LLMs) — Hoffmann et al., 2022
https://scholar.google.com/scholar?q=Language+Models+are+Few-Shot+Learners+%2F+Chinchilla+%28Training+Compute-Optimal+LLMs%29
Interactive Visualization: Warmup-Stable-Decay Learning Rates and the River Valley Loss Landscape