---
title: Warmup-Stable-Decay Scheduling
url: https://www.emergentmind.com/topics/warmup-stable-decay-wsd-scheduling
type: topic
---

# Warmup-Stable-Decay Scheduling

The Warmup-Stable-Decay (WSD) learning rate scheduling paradigm is designed to optimize the training dynamics of large-scale neural architectures, notably transformers, by splitting training into three distinct continuous phases: a linear warmup ramp, a long plateau (“stable” phase) at peak rate, and a short, controlled annealing or decay (“cooldown”) phase. WSD originated as a compute-agnostic alternative to classic cosine annealing, offering flexibility in compute budget allocation, checkpoint reuse, and principled integration of domain adaptation and model scaling studies [2410.05192][2404.06395][2508.01483]. Its widespread adoption in LLM pretraining pipelines reflects both empirical and theoretical advances elucidating its optimization structure, generalization behavior, and practical trade-offs.

## 1. Formal Definition and Parameterization

The WSD schedule is characterized by the following phases and parameterization:

- **Warmup:** Linear ramp from zero to peak rate $\eta_{\max}$ over $T_w$ steps,
  $$
  \eta(t) = \eta_{\max} \frac{t}{T_w} \qquad 0 \leq t \leq T_w
  $$
- **Stable Plateau:** Constant learning rate for $T_s$ steps,
  $$
  \eta(t) = \eta_{\max} \qquad T_w < t \leq T - T_d
  $$
- **Cooldown (Decay):** Monotonically decreasing schedule shape $S(x)$ over $T_d$ steps, where
  $$
  x = \frac{t - (T - T_d)}{T_d}, \quad x \in [0, 1]
  $$
  $$
  \eta(t) = \eta_{\max} S(x)
  $$
Canonical decay functions include:
- **Linear:** $S_{\rm lin}(x) = 1 - x$
- **Square:** $S_{\rm sq}(x) = 1 - x^2$
- **Sqrt:** $S_{\rm \sqrt{}}(x) = 1 - \sqrt{x}$
- **Cosine:** $S_{\rm cos}(x) = \frac{1 + \cos(\pi x)}{2}$
- **Lowered linear:** $S_\alpha(x) = (1-\alpha) + \alpha(1-x)\,,\ \alpha \in [0,1]$  
[2508.01483][2404.06395][2410.05192][2507.17634][2509.19189]

The decay phase typically occupies $8$–$20\%$ of total steps depending on task hardness and performance sensitivity [2508.01483][2410.05192][2404.06395][2509.19189]. The schedule is continuous at phase boundaries.

## 2. Training Dynamics and Loss Landscape Perspective

WSD’s distinct loss curve arises from the interplay between high learning rate exploration and final-phase annealing. During the stable phase, large step sizes induce oscillations in the “hill” directions of the loss landscape while enabling rapid traversal downstream along the “river valley”—a one-dimensional manifold characterized by low curvature [2410.05192][2508.01483]. The rapid cooldown suppresses oscillations, projecting the iterate closer to the true minimum and thereby revealing accumulated optimization progress. This mechanism is formally modeled:
- **Gradient flow in the valley**: $\frac{d}{dt}\theta = -\nabla L(\theta)$  
- **SGD stationary variance (hill-loss):** Proportional to $\frac{1}{2}(d-1)\eta\sigma^2$
- **Cooldown phase:**  
  Suppresses $\eta$, reducing variance and sharp-drop in observed loss [2410.05192][2508.01483].

Loss landscape coordinates $(e_1, e_2)$—global downstream and local-gradient directions—exhibit empirical “river valley” geometry, with optimal decay shapes striking the best balance between exploration and exploitation [2508.01483].

## 3. Cooldown Shape, Bias-Variance Trade-offs, and Model Selection

Cooldown shape $S(x)$ governs a critical bias-variance trade-off in final model quality:
- **Aggressive decay (linear, square):** Promotes exploration, higher variance, but can yield inconsistent or suboptimal fits.
- **Conservative decay (lowered linear, small $\alpha$):** Yields low variance, high bias toward the pre-cooldown state.
- **Intermediate shapes ($S_{\rm \sqrt{}}$, $S_{0.7}$):** Empirically minimize combined bias and variance, producing the best perplexity and generalization performance [2508.01483].

Table: Cooldown shape vs. bias-variance regime

| Shape           | Exploration    | Bias        | Variance    |
|-----------------|---------------|-------------|-------------|
| Linear/Square   | High          | Low         | High        |
| Lowered Linear  | Low           | High        | Low         |
| Sqrt/$\alpha\sim0.7$ | Balanced      | Minimal      | Minimal      |

Empirical recommendations favor the $\sqrt{}$-decay or lowered-linear ($\alpha\approx0.7$), robust across model and dataset scales with $\approx$1.5 PPL improvement over pure linear or square decays [2508.01483].

## 4. Optimizer Hyperparameter Effects During Cooldown

AdamW’s exponential moving average coefficients ($\beta_1, \beta_2$) interact substantially with the cooldown regime:
- Raising $\beta_2$ to $0.99$ (with $p\approx0.3$–$0.5$) improves final perplexity by $\sim0.2$; full tuning of both betas can induce swings comparable to scheduler shape effects.
- The ordering of cooldown shapes by bias-variance remains stable under reasonable $\beta$ hyperparameters.
- Weight decay ($0.1$ is standard) and batch size have secondary effect; aggressive upsampling or disabling weight decay produces marginal gains or setbacks depending on decay shape [2508.01483].

## 5. Theoretical Foundation and Functional Scaling Laws

The WSD schedule has rigorous SDE and kernel regression analyses, culminating in the Functional Scaling Law (FSL) that predicts the evolution of population risk for general LRSs [2509.19189]:
$$
\mathbb{E}[R_K] - \frac{\sigma^2}{2} \approx M^{-s\beta} + T^{-s} + \frac{\sigma^2}{B} \left[b + (a-b) \frac{\min\{M,T_2^{1/\beta}\}}{T_2}\right]
$$
Here, $M$ is model size, $B$ batch size, $a$ peak rate, $b$ final rate, $T_2$ cooldown intrinsic time. WSD achieves optimal scaling exponents and removes logarithmic factors found in direct decay schedules (e.g. exponential), especially in compute- or data-limited regimes [2509.19189]. FSL supports zero-shot prediction and optimization of loss trajectories for unseen schedules and is consistently validated across model sizes and architectures.

## 6. Applications in LLM Training and Domain Adaptation

WSD is widely employed for training transformers and small language models, including MiniCPM, where it facilitates efficient scaling law studies (`data–model scaling`) and continuous/incremental training. It enables checkpoint reuse: a stable-phase checkpoint plus a fixed-length decay matches full-length cosine baselines. In MiniCPM, 10% decay completes convergence and enables mixing new data (domain adaptation or SFT) strictly during the decay phase [2404.06395].

Empirical advantages:
- No precommitment to total compute.
- Checkpoint agnostic: any stable checkpoint can be finalized via short decay.
- Domain adaptation and data-mix-in handled precisely during cooldown.

## 7. Extensions, Model Merging, and Practical Guidelines

Research on WSM (Warmup-Stable and Merge) demonstrates that WSD’s decay phase can be replaced by merging a tail of recent constant-LR checkpoints, yielding superior performance [2507.17634]. Theoretical correspondences map monotone decay schedules to principled merging weights, supporting linear, cosine, and inverse-sqrt schemes.
Key guidelines:
- Merge duration (tail window) is the dominant factor: longer tail merging yields steadier improvement.
- Decay phase or merging window should occupy $8$–$20\%$ of the total budget for optimal annealing benefits.

Practical recommendations:
- **Cooldown length:** $8$–$20\%$.
- **Shape:** Sqrt ($1-\sqrt{x}$) or lowered-linear ($\alpha\sim0.7$).
- **AdamW $\beta_2$:** Increase to $0.99$ during cooldown.
- **Decay window/merge duration:** Extend as far as resource constraints allow.

WSD, underpinned by landscape modeling, scaling law theory, and robust empirical validation, remains a foundational tool for LLM, SLM, and transformer optimization, checkpoint management, and data-efficient scaling research [2508.01483][2410.05192][2404.06395][2507.17634][2509.19189].

Source: https://www.emergentmind.com/topics/warmup-stable-decay-wsd-scheduling