# Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler

arXiv: https://arxiv.org/abs/2508.01483 · HTML: https://mlsanity.com/paper/2508.01483/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 49 references that are in the map.
- First posted on arXiv 2 Aug 2025.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [DeepSeek LLM: Scaling Open-Source Language Models with Longtermism](https://mlsanity.com/paper/2401.02954/) (2024)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Continual Pre-Training of Large Language Models: How to (re)warm your model?](https://mlsanity.com/paper/2308.04014/) (2023)
- [Simple and Scalable Strategies to Continually Pre-train Large Language Models](https://mlsanity.com/paper/2403.08763/) (2024)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [On the SDEs and Scaling Rules for Adaptive Gradient Algorithms](https://mlsanity.com/paper/2205.10287/) (2022)
- [Understanding Emergent Abilities of Language Models from the Loss Perspective](https://mlsanity.com/paper/2403.15796/) (2024)
- [Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective](https://mlsanity.com/paper/2410.05192/) (2024)
- [The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares](https://mlsanity.com/paper/1904.12838/) (2019)
- [JetMoE: Reaching Llama2 Performance with 0.1M Dollars](https://mlsanity.com/paper/2404.07413/) (2024)
- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)
- [Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks](https://mlsanity.com/paper/2305.17212/) (2023)
- [Cautious Optimizers: Improving Training with One Line of Code](https://mlsanity.com/paper/2411.16085/) (2024)
- [Why Warmup the Learning Rate? Underlying Mechanisms and Improvements](https://mlsanity.com/paper/2406.09405/) (2024)

## Built on by (most cited first, 12 of 12)

- [How Learning Rate Decay Wastes Your Best Data in Curriculum-Based LLM Pretraining](https://mlsanity.com/paper/2511.18903/) (2025)
- [Sparse Layers are Critical to Scaling Looped Language Models](https://mlsanity.com/paper/2605.09165/) (2026)
- [Compute-Optimal Quantization-Aware Training](https://mlsanity.com/paper/2509.22935/) (2025)
- [MapPFN: Learning Causal Perturbation Maps in Context](https://mlsanity.com/paper/2601.21092/) (2026)
- [Anytime Training with Schedule-Free Spectral Optimization](https://mlsanity.com/paper/2605.23061/) (2026)
- [Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors](https://mlsanity.com/paper/2606.25971/) (2026)
- [K-EXAONE Technical Report](https://mlsanity.com/paper/2601.01739/) (2026)
- [Progressive Residual Warmup for Language Model Pretraining](https://mlsanity.com/paper/2603.05369/) (2026)
- [Synthetic Persona Pretraining: Alignment from Token Zero](https://mlsanity.com/paper/2608.13482/) (2026)
- [Towards Understanding Momentum Acceleration in River-Valley Loss Landscape](https://mlsanity.com/paper/2609.30957/) (2026)
- [Raising Bars, Not Parameters: LilMoo Compact Language Model for Hindi](https://mlsanity.com/paper/2603.03508/) (2026)
- [Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models](https://mlsanity.com/paper/2610.13524/) (2026)

## Abstract (arXiv)

Learning rate scheduling is essential in transformer training, where the final annealing plays a crucial role in getting the best performance. However, the mechanisms behind this cooldown phase, with its characteristic drop in loss, remain poorly understood. To address this, we provide a comprehensive analysis focusing solely on the cooldown phase in the Warmup-Stable-Decay (WSD) learning rate scheduler. Our analysis reveals that different cooldown shapes reveal a fundamental bias-variance trade-off in the resulting models, with shapes that balance exploration and exploitation consistently outperforming alternatives. Similarly, we find substantial performance variations unicodex2013 comparable to those from cooldown shape selection unicodex2013 when tuning AdamW hyperparameters. Notably, we observe consistent improvements with higher values of β2 during cooldown. From a loss landscape perspective, we provide visualizations of the landscape during cooldown, supporting the river valley loss perspective empirically. These findings offer practical recommendations for configuring the WSD scheduler in transformer training, emphasizing the importance of optimizing the cooldown phase alongside traditional hyperparameter tuning.