---
title: Warmup-Stable-Decay (WSD) Paradigm
url: https://www.emergentmind.com/topics/warmup-stable-decay-wsd
type: topic
---

# Warmup-Stable-Decay (WSD) Paradigm

Warmup-Stable-Decay (WSD) is a learning rate scheduling paradigm for optimizing deep networks—including large language models (LLMs), vision transformers, and kernel regression systems—structured into three sequential phases: a warmup interval during which the learning rate increases from a low (often zero) value to a target peak; a stable phase maintaining the maximum learning rate for most steps; and a final decay or cooldown phase where the learning rate is annealed, often sharply, to facilitate convergence. WSD has emerged as a robust, theoretically justified alternative to schedules such as cosine decay, and is empirically shown to improve convergence dynamics, stability, loss minimization, and scaling law efficiency in both data- and compute-limited regimes.

## 1. Mechanics and Theoretical Rationale

The essential mechanics of WSD rely on the partitioning of training into warmup, stable, and decay periods:

- **Warmup phase**: The learning rate $\eta_t$ is ramped from zero or a small initial value to its target value $\eta_{\max}$ over a predetermined interval, typically using a linear schedule:
  $$
  \eta(t) =
    \begin{cases}
      (\eta_{\max}/t_\mathrm{warm}) \cdot t, & t \leq t_\mathrm{warm} \\
      \eta_{\max}, & t > t_\mathrm{warm}
    \end{cases}
  $$
  This modulates the effective step size and prevents large, unstable parameter updates—particularly in sensitive deeper layers [1810.13243], [2406.09405], [2510.03164].

- **Stable phase**: Maintains the peak learning rate for the majority of training to maximize “intrinsic time” (cumulative effective learning progress), enhancing fast progress along low-curvature (“river”) directions of the loss landscape [2410.05192], [2509.19189]. The choice of plateau height, initial warmup duration, and stable phase length are critical; analytical conditions for the optimal plateau (strong Mpemba point) exist for the valley–river model [2507.04206].

- **Decay (cooldown) phase**: Initiates a rapid reduction in learning rate (“annealing”)—often exponential, linear, or inverse-time shaped—at a scheduled step to eliminate oscillations along “sharp” or “mountain” directions, suppress gradient noise, and descend into a well-conditioned basin of the loss surface [2410.05192], [2508.01483].

The theoretical support for WSD leverages several lines of analysis:

- **Smoothness and Convergence Theory**: The $(H_0, H_1)$-smoothness condition—$\|\nabla^2 f(w)\|_2 \le H_0 + H_1(f(w)-f^*)$—shows that curvature and convergence rate are determined by sub-optimality, motivating an adaptive schedule that increases effective step-size in early (high-loss) regime [2510.03164].

- **Loss Landscape Conditioning**: Warmup triggers “catapult” events that forcibly reduce sharpness (top Hessian eigenvalue) via transient loss spikes, guiding optimization into flatter regions and enabling safe, high learning rates subsequently [2406.09405].

- **Functional Scaling Laws (FSL)**: The FSL framework models risk as a sum of approximation error, full-batch progress (scaling as $1/t^s$ in “intrinsic time”), and a convolutional noise term; a prolonged stable high-LR phase extends $t$, while the decay phase cleans up accumulated noise [2509.19189].

- **Mpemba Effect in Valley–River Model**: Starting with a “hot” (high-LR) plateau then quenching (decaying) the LR yields faster convergence (“strong Mpemba point”), explained by analytical conditions on the amplitude of slowest relaxation modes in the Fokker–Planck operator [2507.04206].

## 2. Empirical Findings and Training Dynamics

Extensive empirical studies confirm several characteristic features and advantages of WSD schedules:

- **Loss Evolution**: During the stable phase, validation loss remains elevated due to oscillatory behavior; the decay phase induces a sharp drop, revealing underlying progress made along flat directions [2410.05192], [2508.01483].

- **Layerwise Stability**: CCA (Canonical Correlation Analysis) shows warmup especially stabilizes deeper, discriminative layers. Freezing final layers mimics the effect of warmup, confirming its protective role against instability in high-LR, large-batch regimes [1810.13243].

- **Scaling Law Efficiency**: WSD schedules facilitate efficient empirical measurement of data-model scaling laws, with MiniCPM experiments identifying a much higher compute-optimal data-to-model ratio (192×) than previous “Chinchilla Optimal” estimates (20×), enabling more data-efficient training of small models [2404.06395].

- **Bias–Variance Trade-off**: The cooldown phase shape mediates the bias–variance balance—aggressive cooldowns provide more exploration (higher variance, lower bias, better model souping), whereas faster decay suppresses variance but risks suboptimal minima [2508.01483].

- **Optimizer Interactions**: Untuned linear or exponential warmup for Adam-based optimizers is empirically validated to match or outperform more algorithmically complex approaches (e.g., RAdam), and GI-Adam initialization can further automate warmup benefits [1910.04209], [2406.09405].

## 3. Practical Implementation and Schedule Selection

Implementation of WSD schedules typically follows these phases:

| Phase   | Learning Rate Rule       | Objective                |
|---------|-------------------------|--------------------------|
| Warmup  | Linear or exponential   | Stability, safe ramp-up  |
| Stable  | Constant $\eta_{\max}$  | Fast progress, "explore" |
| Decay   | Linear, exponential, or sqrt/inv-time | Fine-tuning, bias-variance balance |

- Warmup duration is frequently set by “rule-of-thumb” formulas (e.g., linear warmup over $2/(1-\beta_2)$ steps in Adam; exponential warmup $\omega_t=1-\exp[-(1-\beta_2)t]$) [1910.04209].

- The stable phase dominates runtime (often 80–90% of steps) and should not be truncated prematurely [2404.06395], [2509.19189].

- Decay phase length and shape are critical; “sqrt” or “lowered linear 0.7” shapes minimize bias–variance, and higher AdamW $\beta_2$ values during cooldown increase stability [2508.01483].

- Practical recommendations include checkpointing at stable-LR, then branching into decay (WSD-S scheme), allowing flexible adaptation to different compute budgets and facilitating continual or domain-adaptive training [2410.05192].

## 4. Advanced Variants, Alternatives, and Model Averaging

Recent developments question the necessity of an explicit online decay phase:

- **WSM: Warmup-Stable-Merge** ([2507.17634]): Here, model merging of periodically checkpointed states post-stable training can emulate decay strategies as weighted averages of model weights: $\hat{\theta}_{n+k} = \sum_{j=0}^k c_j \theta_{n+j}$, with the $c_j$ designed to match gradient decay sequences. Extensive benchmarking shows WSM delivers higher performance than online decay schedules (WSD), particularly when merge duration is optimized.

- **Linear Decay to Zero (D2Z)** ([2502.15938]): Empirical and theoretical analysis (via AdamW’s EMA interpretation) reveals that linear decay to zero outperforms partial decay (e.g., cosine to 10%), offering improved bias reduction and variance suppression; D2Z provides compute savings over standard schedules, especially at high tokens-per-parameter.

## 5. Interactions with Scaling, Initialization, and Feature Alignment

- **Weight Decay vs. Maximal Update Parameterization (μP)** ([2510.19093]): Independent weight decay, rather than μP learning rate rules, governs stable feature update dynamics across widths after early training; μP effectively serves as an implicit warmup mechanism, with the main stabilizing factor being the product $\eta \lambda$.

- **Metric-Based Update Control** ([2410.23922]): Monitoring update $\ell_2$-norm, angular change, and Relative Representation Change (RRC) enables adaptive warmup and reduces fixed warmup requirements. Optimizer modifications (e.g., LionA/LionAR) can normalize updates such that angular and RRC metrics are controlled, further reducing early training instabilities.

## 6. Loss Landscape and Scaling Law Perspectives

- **River Valley Landscape** ([2410.05192], [2507.04206], [2508.01483]): The “river valley” metaphor models loss surfaces as valleys with steep sides (fast modes) and a flat river (slow global modes). WSD’s stable phase maximizes progression along the river, while decay phase reduces oscillations and ensures effective descent into the basin. Mathematically, SGD dynamics can be decomposed into “river” (primary progress) and “hill” (oscillation due to large LR and noise) components.

- **Functional Scaling Law Analysis** ([2509.19189]): The optimal schedule balances intrinsic time (by maximizing stable LR period) versus final risk minimization (by aggressive noise suppression during decay). Quantitative risk decomposition via FSL provides explicit guidelines for hyperparameter settings, with rapid decay phases dominating final error terms.

## 7. Practical Recommendations and Outlook

- Optimal training with WSD requires careful but not fragile hyperparameter settings due to its robustness; plateau height can be chosen via analytical conditions such as the strong Mpemba point.
- Most training should reside in the stable phase for efficiency; the decay phase should be tuned (typically 10–20% of total steps, shape depending on bias–variance trade-off); model merging and interpolative checkpointing offer new avenues for decay-free finishing [2507.17634].
- Theoretical and empirical analyses affirm WSD’s superiority over constant learning rate, pure decay, or cosine schedules for language model pretraining, scaling law estimation, and transfer across model sizes.

WSD thus represents both a practical and theoretically grounded framework for large-scale model optimization, incorporating insights from modern loss landscape analysis, SGD stochasticity, scaling law efficiency, and advanced techniques in model averaging and optimizer design.

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