WSD Schedules for Efficient Deep Learning

• WSD schedules are a learning rate strategy that divides training into warmup, stable, and decay phases to optimize convergence.
• They mitigate gradient instabilities in the warmup phase, maintain exploration during a constant plateau, and fine-tune parameters through decay.
• These schedules enhance convergence, generalization, and compute efficiency across diverse deep learning domains including vision, language, and speech.

A Warmup-Stable-Decay (WSD) schedule is a learning rate scheduling strategy commonly used in the training of deep neural networks, particularly in large-scale pretraining and large-batch regimes. WSD schedules divide the evolution of the learning rate (LR) into three explicit phases: a warmup phase with increasing LR, a stable phase with constant LR, and a decay phase with decreasing LR. This structure offers both empirical robustness and theoretical advantages for convergence and generalization in substantial deep learning workloads across vision, language, and speech domains.

1. Formal Definition and Schedule Structure

A WSD schedule is mathematically characterized by sequential phases:

• Warmup: The LR increases monotonically from a low initial value (often zero) to a prescribed maximum over $W$ steps. This is commonly implemented as linear or sub-exponential increase:

$$\eta(s) = \begin{cases} (s/W)\cdot\eta_{\text{max}}, & s < W \ \eta_{\text{max}}, & W \leq s < T \ f(s-T)\cdot\eta_{\text{max}}, & s \geq T \end{cases}$$

where $s$ is the step, $W$ is the number of warmup steps, $T$ the transition to decay, and $f(\cdot)$ is a monotonic decreasing function (e.g., linear, exponential, or cosine decay) (Hu et al., 9 Apr 2024).

• Stable (Plateau): For a prolonged interval, the LR is held constant at its peak $\eta_{\text{max}}$.
• Decay: LR decreases according to the chosen rule (linear, cosine, polynomial, etc.) until it reaches a low terminal value or zero.

Schedules fitting this generic form appear across diverse work on LLM pretraining (Hu et al., 9 Apr 2024, Wen et al., 7 Oct 2024, Schaipp et al., 31 Jan 2025, Bergsma et al., 21 Feb 2025), speech-to-text (Gaido et al., 29 May 2025), and robust training (Shi et al., 2021).

2. Motivations and Theoretical Explanations

The combination of these three phases is supported by several empirical and theoretical insights:

• Warmup: Prevents instability—especially sharp gradient updates in deeper network layers—when starting training with a high LR. Gradually increasing the LR allows adaptive optimizers (like Adam) or SGD to achieve stable early updates before entering high-LR regimes (Gotmare et al., 2018, Ma et al., 2019, Kalra et al., 13 Jun 2024).
• Stable phase: Enables efficient “exploration” of the loss landscape at maximal LR. In LLMs and other high-dimensional systems, progress along the flattest (“river”) directions is governed by a sufficiently large LR; too small a LR would stall optimization along those directions, while too large can cause divergence (Wen et al., 7 Oct 2024, Liu et al., 6 Jul 2025). Thermodynamic analogies suggest that this phase “preheats” the system, setting the stage for rapid convergence during decay via the Mpemba effect (Liu et al., 6 Jul 2025).
• Decay: As optimization transitions from exploring to “fine-tuning,” reducing the LR allows smaller, more precise parameter updates, mitigates the effect of gradient noise, and helps the system settle in flatter minima, improving generalization (Hu et al., 9 Apr 2024, Schaipp et al., 31 Jan 2025, Bergsma et al., 21 Feb 2025).

Theoretical convergence analyses support these empirical phenomena. For instance, a WSD schedule (constant plateau with linear cooldown) yields performance bounds that omit the usual logarithmic suboptimality penalty seen in pure constant schedules, thus optimizing both practical and theoretical convergence (Schaipp et al., 31 Jan 2025).

3. Empirical Properties and Training Dynamics

Empirical investigation reveals a set of characteristic behaviors:

Phase	Loss/Weight Dynamics	LR Behavior
Warmup	Layerwise stabilization, especially in deep layers; mitigates large update magnitudes and sharpness spikes	LR ramp-up (linear/sub-exp)
Stable	Loss plateaus but parameters traverse the “river valley” of the loss surface; rapid progress along slowest modes; high “temperature”	LR constant at $\eta_{\max}$
Decay	Sudden, sharp decrease in validation/training loss; oscillations (“bouncing”) in “valley” directions dampen; parameter updates become more refined	LR decreases quickly

Notably, the greatest loss reduction often occurs during the decay phase following a prolonged stable plateau. High LR during stability phase is crucial for accelerating long-term convergence during decay (the “Mpemba effect”) (Liu et al., 6 Jul 2025). Deviating from this by decaying too early or keeping the plateau too low can degrade performance (Wen et al., 7 Oct 2024). Conversely, excessive or poorly controlled warmup risks divergence—or, if too conservative, can simply waste compute by prolonging ineffective training (Gaido et al., 29 May 2025, Kalra et al., 13 Jun 2024).

4. Practical Implementations and Variations

WSD schedules are instantiated with various detailed hyperparameters and functional forms depending on modality and scale:

• Warmup: Typically a linear increase, but sub-exponential (Gaido et al., 29 May 2025), exponential, or piecewise-linear (“double linear”) warmups are also used. In adaptive optimizers, linear warmup over $2/(1-\beta_2)$ steps is recommended for Adam (Ma et al., 2019).
• Stable phase: Length can be flexibly chosen. The plateau’s height should not be too small; empirical and theoretical analyses suggest that a high plateau is generally beneficial (Liu et al., 6 Jul 2025).
• Decay: Can follow linear decay-to-zero (D2Z) (Bergsma et al., 21 Feb 2025), cosine, or more tailored time-dependent schedules. D2Z, in particular, has shown systematic improvements over cosine decay to a fixed fraction of the maximum (e.g., 10%) in LLMs, yielding lower final losses and greater compute efficiency, especially at high tokens-per-parameter (TPP) (Bergsma et al., 21 Feb 2025).

Several works recommend checkpointing at the end of the stable phase for easy resumption or domain adaptation (“continual” or “continued” pretraining) (Hu et al., 9 Apr 2024, Gupta et al., 2023). The WSD-S (“simplified”) variant reuses the decayed checkpoint as a new base for continued high-LR training, reducing unnecessary forking of training branches (Wen et al., 7 Oct 2024).

5. Landscape and Loss Curve Interpretation

Recent theoretical interpretations posit that the loss surface during large-scale pretraining has a “river valley” geometry—broad, flat directions (rivers) coupled with steep orthogonal valleys (Wen et al., 7 Oct 2024, Liu et al., 6 Jul 2025). The WSD schedule exploits this geometry:

• The stable (high LR) phase propels the optimizer rapidly along the river, despite producing large oscillations in “hill” (steep) directions. This manifests as elevated loss curves due to off-river deviations.
• The decay phase quells oscillations, causing the optimizer’s trajectory to collapse toward the true valley bottom, as revealed by a sharp drop in loss once the LR starts to decrease.

This viewpoint accounts for the empirically observed nonmonotonic loss dynamics and sharp “cooldown” effect and is buttressed by minimal analytical models inspired by thermodynamics (e.g., Mpemba point selection) (Liu et al., 6 Jul 2025).

6. Extensions, Automatic Tuning, and Integration

Several extensions and methodological improvements to WSD-type schedules have been proposed:

• Adaptive decay and warmup triggers: Strategies such as ABEL use signals like weight norm “bounces” to automatically transition between phases instead of fixed epochs (Lewkowycz, 2021).
• Quality-driven adaptation: Scheduling learning rate (and possibly weight decay) in synchrony with intrinsic network metrics such as knowledge gain or mapping condition improves data and layer utilization (Hosseini et al., 2020).
• Scheduled weight decay: Coordinating the adaptation of weight decay (e.g., SWD) with the learning rate phase to control gradient norm growth, especially in Adam-type optimizers (Xie et al., 2020), naturally fits within the WSD scheduling logic.
• Gradient transformations: Gradient preconditioning (e.g., GradPower) further enhances the efficacy of WSD schedules by amplifying slow directions and mitigating oscillations in the stable phase; optimal exponent settings depend on signal-to-noise ratios (Wang et al., 30 May 2025).

The empirical multi-power law (MPL) provides a quantitative tool for predicting loss curves under various WSD schedules and for optimizing schedule parameters with minimal trial runs (Luo et al., 17 Mar 2025).

7. Advantages, Limitations, and Recommendations

WSD schedules offer several critical advantages:

• Robustness and flexibility: Effective under wide-ranging compute budgets since the stable phase can, in principle, be extended indefinitely and the transition to decay can be checkpointed or resumed as needed (Wen et al., 7 Oct 2024, Hu et al., 9 Apr 2024).
• Theoretical and empirical convergence: Convergence bounds are improved compared to pure constant LR schedules, matching those of convex optimization with linear cooldown and often yielding superior empirical results over cosine or fixed-decay schedules (Schaipp et al., 31 Jan 2025).
• Improved compute efficiency: Full decay-to-zero (e.g., D2Z) leads to lower final loss at substantially reduced TPP, offering up to 60% compute savings over cosine “10x decay” in LLM pretraining (Bergsma et al., 21 Feb 2025).
• Clarity for scaling law studies: The schedule’s explicit separation of exploration and fine-tuning phases allows efficient experimentation along the data axis and direct derivation of optimal data-to-model ratio (Hu et al., 9 Apr 2024).

However, certain practical challenges remain:

• Warmup duration and schedule shape must be tuned to avoid instability or excessive convergence delay (Gaido et al., 29 May 2025, Kalra et al., 13 Jun 2024).
• The theoretically “optimal” plateau height (the “strong Mpemba point”) is problem-dependent and may not be easy to estimate in general (Liu et al., 6 Jul 2025).
• In low-noise or lazy training regimes, warmup and plateau phases may yield diminishing returns, with much of the gain arising from the decay phase alone (Lewkowycz, 2021).

References to Key Papers

Below is a selection of key references that document the development, analysis, and application of WSD schedules and their variants:

Topic/Concept	Reference
Basic WSD definition, training dynamics	(Hu et al., 9 Apr 2024, Wen et al., 7 Oct 2024)
Theoretical convergence, suboptimality bounds	(Schaipp et al., 31 Jan 2025)
D2Z (linear decay-to-zero) efficiency in LLMs	(Bergsma et al., 21 Feb 2025)
Gradient warmup necessity and regimes	(Kalra et al., 13 Jun 2024)
Knowledge gain/mapping condition adaptation	(Hosseini et al., 2020)
Scheduled weight decay synergy	(Xie et al., 2020)
River valley/Mpemba thermodynamic analogy	(Liu et al., 6 Jul 2025, Wen et al., 7 Oct 2024)
Power-law prediction of losses across LRS	(Luo et al., 17 Mar 2025)
Robust phase and schedule optimization	(Gupta et al., 2023, Kim et al., 2021)
Speech-to-text warmup schedule comparison	(Gaido et al., 29 May 2025)
Gradient transformation during stable phases	(Wang et al., 30 May 2025)

An understanding and appropriate application of WSD schedules is now considered essential in designing robust, efficient pretraining and fine-tuning pipelines for LLMs and deep networks more broadly.