---
title: Warmup-Stable-Decay LRS
url: https://www.emergentmind.com/topics/warmup-stable-decay-learning-rate-scheduler-wsd-lrs
type: topic
---

# Warmup-Stable-Decay LRS

The Warmup-Stable-Decay (WSD) Learning Rate Scheduler is a three-phase adaptive learning rate strategy that has become foundational in the training of large-scale neural language models, particularly transformers. WSD divides training into a brief warmup, a long stable-phase with constant learning rate, and a cooldown or decay phase, during which the learning rate is annealed to a low value. The structure, mathematical formulation, theoretical justification, practical implications, and empirical tradeoffs of WSD have been extensively studied across recent literature, providing a unified perspective on scheduler design and optimization in deep learning.

## 1. Mathematical Structure and Phase Definition

WSD is defined via three contiguous training phases:

- **Warmup**: Learning rate increases linearly from zero (or a small initial value) to a peak value $\eta_{\max}$ over $T_{\rm warm}$ steps,
  $$
  \eta(t) = \eta_{\max} \frac{t}{T_{\rm warm}} \quad\quad 0 \leq t < T_{\rm warm}
  $$
- **Stable**: Learning rate is held constant,
  $$
  \eta(t) = \eta_{\max} \quad\quad T_{\rm warm} \leq t < T_{\rm stab}
  $$
- **Decay/Cooldown**: Learning rate is annealed smoothly to zero or a minimum value,
  $$
  \eta(t) = \eta_{\max}\,f\left(\frac{t-T_{\rm stab}}{T_{\rm tot}-T_{\rm stab}}\right), \quad T_{\rm stab} \leq t \leq T_{\rm tot}
  $$
where $f(\cdot)$ is a monotonically decreasing decay function, such as linear, cosine, exponential, square-root, or custom power-law variants. This template supports both single-phase runs and more advanced multi-budget (“branching”) strategies [2508.01483][2410.05192][2507.17634][2404.06395][2509.19189].

### Phase Timing and Hyperparameters

Typical settings emphasize a minimal warmup fraction (0.5–2%), a long plateau (80–90%), and a final decay of 10–20% of total steps. The decay function can be selected from several families, with square-root and 1-sqrt laws often empirically dominating [2508.01483][2507.17634][2404.06395].

| Phase      | Interval                                              | Typical Fraction | Main Role          |
|------------|------------------------------------------------------|------------------|--------------------|
| Warmup     | $[0, T_{\rm warm})$                                  | 0.5–2%           | Stabilization      |
| Stable     | $[T_{\rm warm}, T_{\rm stab})$                       | 80–90%           | Exploration        |
| Decay      | $[T_{\rm stab}, T_{\rm tot}]$                        | 10–20%           | Exploitation/final convergence |

## 2. Theoretical Justification and Loss Landscape

The rationale for WSD arises from both empirical phenomena and formal analysis of optimization in ill-conditioned, anisotropic loss landscapes characteristic of modern LLM pretraining. The “river valley” theory posits that the loss surface comprises sharp “valley” directions (with large Hessian eigenvalues) and flat “river” directions (with small eigenvalues), yielding disparate timescales for optimization.

- **Stable phase (high constant LR)**: Large steps cause rapid diffusion and exploration along the wide river manifold, while valley directions equilibrate quickly. This phase accumulates most progress along the loss landscape’s principal directions.
- **Cooldown/decay phase**: By reducing the learning rate, the stochastic variance in sharp (“hill”) directions contracts, focusing optimization tightly into the loss basin and enabling sharp drops in validation loss [2410.05192][2508.01483][2507.04206].

This theoretical grounding is supported by SDE modeling and principal component visualizations that highlight how performance gains and loss drops are realized almost exclusively during the final decay, as the optimizer transitions from broad exploration to fine convergence [2410.05192][2508.01483][2509.19189].

## 3. Decay Functional Forms and Bias-Variance Tradeoff

The cooldown phase's effect is highly sensitive to the specific decay function. Common choices include:

- **Linear**: $f_{\rm linear}(\tau) = 1-\tau$
- **Cosine**: $f_{\rm cosine}(\tau) = \frac{1}{2}(1+\cos(\pi\tau))$
- **Quadratic**: $f_{\rm square}(\tau) = 1-\tau^2$
- **Square-root**: $f_{\rm sqrt}(\tau) = 1-\sqrt{\tau}$
- **Exponential**: $f_{\rm exp}(\tau) = \exp(-\alpha\tau)$

The choice of $f(\cdot)$ reflects a fundamental bias–variance tradeoff in the resulting model ensemble [2508.01483]:

- **Aggressive (high-variance, low-bias) shapes** (e.g., mirror-cosine) yield a wide spread in final model quality, occasionally reaching better minima but with less reproducibility.
- **Gentle (high-bias, low-variance) shapes** (e.g., linear) yield highly repeatable but potentially suboptimal solutions.
- **Intermediate shapes** (notably square-root decay or lowered linear) balance these tendencies and deliver the lowest validation perplexity, achieving 1–2 point gains (2–3% loss improvement) relative to naive linear decay on large transformer training runs [2508.01483][2507.17634]. 

Empirically, 1-sqrt or square-root decay schedules are preferred in modern fine-tuning and LLM pretraining setups [2507.17634].

## 4. Optimizer Interactions and Training Dynamics

The interplay between the learning rate schedule and optimizer hyperparameters becomes critical during cooldown. Specifically:

- **AdamW’s second moment parameter $\beta_2$** is identified as a key driver of cooldown performance. Increasing $\beta_2$ (lengthening the EMA half-life) smooths gradient noise and can improve validation perplexity by up to 2 points. Retuning $\beta_1$ and $\beta_2$ together can move runs from worst to best within the same decay shape [2508.01483].
- **Batch size** can be increased during cooldown, with mild gains, though matching the optimizer’s token half-life to batch dynamics may degrade final loss.
- **Weight decay** has nuanced effects: resetting or removing it may harm best-performing shapes but occasionally helps especially high-variance decays.

Consecutive step cosine similarities, gradient norm shrinkage, and aligned directional derivatives are observed during decay, consistent with rapid “funneling” into locally convex minima and confirming the theoretical “river valley” model [2404.06395][2508.01483].

## 5. Empirical Benchmarks and Scaling Law Implications

Comparative studies demonstrate WSD’s strong practical performance across data/model scales and learning regimes:

- **Benchmarks**: On transformer models (e.g., MiniCPM) and standard tasks, WSD with optimal decay shapes matches or outperforms cosine and linear schedulers, requiring only 10% of total steps for decay to see maximal final loss drop [2404.06395][2507.17634].
- **Scaling Laws**: WSD’s structure is theoretically justified and empirically verified to facilitate compute/data/model scaling explorations. In particular, WSD enables $O(m)$ studies of the data–model law, where data-to-model ratios an order of magnitude higher than the Chinchilla prescription are found optimal under WSD [2404.06395][2509.19189].
- **Continual/Domain Adaptation**: WSD and its variants (e.g., WSD-S) support continual training and modular checkpoint usage across multiple budgets, outperforming cyclic or cosine-rewarm methods and simplifying code [2410.05192].

## 6. Model Merging and Recent Alternatives

Recent work establishes a formal equivalence between WSD’s decay phase and model averaging (merging) methods. Instead of decaying the learning rate online, one may maintain a constant learning rate and retrospectively form convex combinations of model checkpoints, assigning weights that emulate any target decay law [2507.17634]:

$$
\hat\theta_{n+k} = \sum_{j=0}^k c_j\,\theta_{n+j}
$$
where the sequence $w_i = \sum_{j=i}^k c_j$ corresponds to the effective learning rate profile. The WSM (Warmup-Stable and Merge) framework outperforms WSD baselines by focusing optimization on the merge duration, and evidence suggests that the main critical hyperparameter is the window for merging, not the decay law per se [2507.17634].

## 7. Implementation, Best Practices, and Practical Guidelines

Pseudocode for integrating WSD into PyTorch- or similar training frameworks requires, at each step, recomputation of $\eta(t)$ prior to optimizer step execution. Key best practices include [2507.17634][2508.01483]:

- Warm up for a small fraction of total steps (0.5–2%), sufficient for gradient stabilization.
- Hold $\eta_{\max}$ until late training, devoting 10–20% of tokens to cooldown.
- Deploy square-root or 1-sqrt decay as the default; these yield robust bias–variance balance.
- Tune AdamW’s $\beta_2$ higher (e.g., $0.95^{0.5} \approx 0.974$) in cooldown.
- For fine-tuning, use a short decay over 5–10% of new steps with aggressive annealing.
- If training must be extended post-decay, revert to the last stable checkpoint or switch to decay-free/merging schedulers for continuous operation.

| Hyperparameter        | Typical Value / Range         | Role                      |
|----------------------|------------------------------|---------------------------|
| Warmup length        | 0.5–2% of steps              | Early stabilization       |
| Stable phase         | 80–90% of steps              | Exploration, progress     |
| Decay/cooldown       | 10–20% of steps              | Convergence, exploitation |
| AdamW $\beta_2$      | $0.95^{0.5} \approx 0.974$   | Stabilize variance        |
| Decay shape          | 1-sqrt or square-root        | Bias-variance tradeoff    |

The warmup–stable–decay paradigm, with its explicit separation of optimization roles by subphase, remains foundational. Recent advances in model merging (WSM) and river valley landscape theory inform nuanced adaptations and underline WSD’s ongoing central role in the calibration of high-performance language models [2508.01483][2410.05192][2509.19189][2507.04206][2404.06395][2507.17634].

Source: https://www.emergentmind.com/topics/warmup-stable-decay-learning-rate-scheduler-wsd-lrs