Warmup–Stable–Decay (WSD) in Model Training

• Warmup–Stable–Decay (WSD) is a three-phase learning rate schedule featuring warmup, stable, and decay phases to optimize training dynamics.
• It decouples training from a fixed endpoint, allowing flexible checkpointing and adaptability in large-scale language model training.
• Empirical studies indicate that tuning the cooldown phase directly influences validation perplexity and the bias–variance trade-off.

Searching arXiv for papers on Warmup–Stable–Decay and closely related analyses. Warmup–Stable–Decay (WSD) is a three-phase learning-rate schedule in which the learning rate is first increased during warmup, then held constant during a stable phase, and finally reduced during a decay or cooldown phase. In contemporary language-model training, WSD is used as an alternative to schedules such as warmup-plus-cosine because it decouples most of training from a precommitted endpoint while preserving a late annealing stage that can strongly improve the final model. Recent work treats WSD not merely as an engineering convenience but as a schedule with distinctive optimization dynamics, geometric structure, and asymptotic implications for scaling, noise control, and checkpoint reuse (Hu et al., 2024, Dremov et al., 2 Aug 2025, Li et al., 6 Feb 2026).

1. Definition and canonical forms

WSD is conventionally described by three phases.

Phase	Role
Warmup	Increase the learning rate from $0$ or a small value to a peak value
Stable	Keep the learning rate constant at the peak
Decay / cooldown	Reduce the learning rate from the stable value to zero or a lower endpoint

A widely used explicit formulation sets, for total steps $T_{end}$, warmup steps $T_w$, and decay start $T_c$,

$$\eta(t) = \begin{cases} \frac{t}{T_w} & \text{if } t \leq T_w, \ 1 & \text{if } T_w < t \le T_c, \ 1-\frac{t-T_c}{T_{end}-T_c} & \text{if } t > T_c, \end{cases}$$

assuming peak learning rate $=1$. The corresponding warmup-cosine baseline instead begins decaying immediately after warmup (Belloni et al., 13 Jan 2026). In the more general formulation introduced in “MiniCPM” (Hu et al., 2024), WSD is written as

$$WSD(T; s) = \begin{cases} \frac{s}{W}\eta, & s < W \ \eta, & W < s < T \ f(s-T)\eta, & T < s , \end{cases}$$

where $f(s-T)$ is any decreasing function satisfying $0 < f(s-T) \le 1$. This definition makes the stable/decay split primary and leaves the decay shape open.

That openness has produced several concrete WSD variants. Practical studies analyze linear, cosine-like, exponential, inverse-square-root-like, and power-like terminal decays, while some papers note the older term “trapezoid scheduler” for a closely related shape (Belloni et al., 13 Jan 2026, Hu et al., 2024). A useful distinction is therefore between WSD as a phase structure and any single decay curve used inside that structure.

2. Emergence and practical role in large-scale training

WSD became prominent because standard cosine schedules require the final training horizon to be fixed in advance, whereas large-scale pretraining is often extended, resumed, or repurposed as more data or compute becomes available. “MiniCPM” introduced WSD as a schedule “tailored for continuous training with an un-predefined pre-training token number” and argued that it “makes the reusing of model intermediate checkpoints highly feasible” (Hu et al., 2024). The schedule’s long constant-learning-rate branch permits continued training without redesigning the full trajectory, and a later decay branch can be launched once a stopping point or domain transition is chosen.

This operational flexibility was one of the main reasons WSD spread in language-model practice. A later cross-architecture study states that WSD “has become popular especially in LLM training” because it can perform “as well as or better than warmup + cosine decay,” is convenient for training continuation, and avoids the need to commit early to a final training horizon (Belloni et al., 13 Jan 2026). The same work emphasizes that WSD also avoids issues associated with restarting cosine schedules, such as “loss spikes” or “unlearning” caused by rewarming.

The warmup and decay components also have partially separable historical motivations. Earlier empirical work on warmup and cosine restarts found that warmup’s effect is to prevent “the deeper layers from creating training instability,” whereas the benefits often attributed to cosine annealing appear to come more from smooth decay itself than from restart events (Gotmare et al., 2018). This supports the later WSD decomposition into an early stabilization phase, a high-learning-rate working phase, and a terminal annealing phase.

3. Cooldown as the decisive stage

The most detailed empirical account of WSD’s terminal stage appears in “Training Dynamics of the Cooldown Stage in Warmup-Stable-Decay Learning Rate Scheduler” (Dremov et al., 2 Aug 2025). That paper isolates the final cooldown stage in transformer language-model training and argues that it is not a minor finishing touch: although cooldown is “typically around 20% of the total training duration,” it is responsible for a characteristic sharp drop in validation perplexity. In the main setup, training proceeds for 26,400 iterations before cooldown, the cooldown fraction is 20% of total training, and the learning rate is always reduced to zero during cooldown.

The paper parameterizes cooldown by normalized progress $x\in[0,1]$, with $T_{end}$0 at cooldown start and $T_{end}$1 at cooldown end. The tested shape functions are

$T_{end}$2

$T_{end}$3

The same study also defines “lowered linear” cooldowns whose starting learning rate is a fraction of the stable-phase learning rate; “lowered linear 0.7” starts at 70% of the stable learning rate and then decays linearly to zero (Dremov et al., 2 Aug 2025).

Its central empirical conclusion is that cooldown shape induces a genuine exploration–exploitation, or bias–variance, trade-off. Using $T_{end}$4 cooldown runs from the same pre-cooldown checkpoint and different data permutations, the paper decomposes final validation loss into bias and variance terms and finds that the strongest single-run shapes are not the most exploratory and not the most conservative. Rather, the best trade-off shapes are “sqrt” and “lowered linear 0.7”; “mirror cosine” and “square” are more exploratory and higher variance, while aggressively lowered linear schedules are more exploitative and more biased (Dremov et al., 2 Aug 2025).

The same work shows that optimizer tuning during cooldown matters nearly as much as schedule-shape choice. With AdamW baseline hyperparameters

$T_{end}$5

and gradient clipping at $T_{end}$6, changing cooldown-stage optimizer settings can alter performance by about as much as changing cooldown shape. Its strongest optimizer recommendation is stage-specific: higher values of $T_{end}$7 during cooldown consistently improve results, and poor beta choices can erase the gains from a favorable cooldown shape (Dremov et al., 2 Aug 2025). This makes the decay stage a separate hyperparameter-tuning problem rather than a passive tail of the main run.

An important exception concerns model averaging. When four independent cooldown runs from the same pre-cooldown model are trained on non-overlapping data portions and their weights are averaged, higher-variance shapes such as mirror cosine become more attractive; however, the paper still reports that a single longer run beats these averages when matched for total compute and data (Dremov et al., 2 Aug 2025).

4. Geometric interpretations and cross-architecture evidence

A major interpretive line treats WSD through loss-landscape geometry. “Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective” argues that WSD produces an elevated loss during the stable phase and a sharp decline during decay because pretraining loss behaves like a river valley landscape: the stable high learning rate drives rapid progress along the river direction while sustaining oscillations across steeper directions, and the fast decay suppresses those oscillations, revealing optimization progress in the mountain directions (Wen et al., 2024).

The cooldown-focused transformer study supports a closely related picture using explicit landscape visualization (Dremov et al., 2 Aug 2025). It constructs a two-dimensional plane in weight space from $T_{end}$8, the vector from the pre-cooldown checkpoint to the final post-cooldown model, and $T_{end}$9, the vector corresponding to ten optimizer steps from a given point. Early in cooldown, the observed validation-perplexity landscape shows a clear valley along the global direction; midway through cooldown the valley remains but is less pronounced; by the end, the final basin dominates. The same paper reports that the cosine similarity between gradients and AdamW’s momentum state increases over cooldown, moving from negative alignment to slightly positive, while gradient norms also evolve toward more directionally stable behavior (Dremov et al., 2 Aug 2025).

“Universal Dynamics of Warmup Stable Decay” extends this geometric line beyond transformers (Belloni et al., 13 Jan 2026). Comparing a roughly 160M-parameter Pythia-like decoder-only LLM trained on SlimPajama with a 334K-parameter CNN trained on CIFAR-10, it finds qualitatively similar WSD signatures in both settings: modest decreases during the stable phase, sudden larger gains during cooldown, valley-shaped interpolation near the end of the stable phase, monotone interpolation across cooldown, and increasing sharpness during cooldown. In both architectures, the first principal component explains at least approximately 40% of the variance in sampled stable-phase and decay-phase iterates, and the early-decay direction is more aligned with high-curvature subspaces than the stable-phase direction (Belloni et al., 13 Jan 2026). The main architectural difference reported there is quantitative rather than qualitative: the LLM used a 20% cooldown, whereas the CNN needed a 35% decay phase to better match cosine performance.

A more speculative geometric account comes from a minimal “valley-river” model motivated by the Mpemba effect (Liu et al., 6 Jul 2025). There, a sufficiently high plateau learning rate can prepare a state whose overlap with the slowest post-decay relaxation mode is smaller than that of a lower plateau; at a “strong Mpemba point,” that slowest mode vanishes entirely. This does not establish real-LLM WSD optimality, but it provides a principled rationale for why a high plateau can accelerate later decay rather than merely delay convergence.

5. Theory, scaling laws, and optimality claims

Recent theory places WSD inside functional scaling-law analyses rather than treating it as a purely empirical heuristic. “Optimal Learning-Rate Schedules under Functional Scaling Laws” derives a sharp phase transition at

$T_w$0

where $T_w$1 is the source exponent and $T_w$2 the capacity exponent (Li et al., 6 Feb 2026). In the easy-task regime $T_w$3, the optimal schedule is a power decay from the start: $T_w$4 In the hard-task regime $T_w$5, the optimal schedule is WSD-like: $T_w$6 with decay fraction

$T_w$7

The notable caveat is terminological: in this theory, “warmup” is not derived as an explicit increasing phase. The mathematically optimal asymptotic form begins immediately at the maximal admissible stable rate, so the result is an idealized stable-decay analogue of practical WSD rather than a literal three-piece schedule (Li et al., 6 Feb 2026).

A related FSL analysis makes the same intuition more explicit in risk decomposition form (Li et al., 23 Sep 2025). There the schedule affects training through a convolutional noise term, and for a WSD-like schedule the final excess risk scales as

$T_w$8

where $T_w$9 and $T_c$0 are the intrinsic-time contributions of the stable and decay phases. The interpretation is that the long stable phase enlarges intrinsic optimization time, while the decay phase controls the final noise term; because earlier noise is discounted by the forgetting kernel, WSD can outperform direct exponential decay by logarithmic factors and, in hard regimes, remove the logarithmic inefficiency entirely (Li et al., 23 Sep 2025).

Empirical schedule-law work also points toward WSD-like optima. “A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules” models loss as a dominant power law in cumulative learning rate plus an additional loss-reduction term induced by decay events, then numerically optimizes schedules and repeatedly finds WSD-like optima with a long stable phase and a terminal power-like decay (Luo et al., 17 Mar 2025). In that work, the discovered decay is approximately

$T_c$1

which outperforms cosine and tuned linear/exponential WSD variants on the reported pretraining-loss comparisons (Luo et al., 17 Mar 2025).

WSD also changes batch-size theory. “How to Set the Batch Size for Large-Scale Pre-training?” argues that under WSD’s long stable phase the classical single critical batch-size picture fails; instead, one obtains a minimum feasible batch size $T_c$2 and a token-optimal batch size $T_c$3, both of which typically increase as loss decreases. This motivates dynamic batch growth during the stable phase rather than a single fixed batch choice (Zhou et al., 8 Jan 2026).

6. Variants, alternatives, and active debates

Although WSD is often treated as a strong default, several papers challenge specific parts of the recipe. One line questions whether explicit decay is necessary at all. “Warmup-Stable and Merge (WSM)” interprets WSD’s decay phase as something that can be emulated by checkpoint merging: a weighted average over stable-phase checkpoints induces an effective decay over update contributions, and in the reported MoE pretraining experiments WSM outperforms a WSD baseline, with merge duration emerging as the most critical hyperparameter (Tian et al., 23 Jul 2025). “Through the River” makes a related argument for Schedule-Free AdamW: when tuned well, the schedule-free method appears to follow the river structure without explicit decay or explicit weight averaging, though the paper also reports strong sensitivity to momentum and weaker large-batch behavior for vanilla schedule-free methods (Song et al., 14 Jul 2025).

Another line questions WSD when the downstream goal is supervised fine-tuning rather than best pretraining loss. “Pre-training LLM without Learning Rate Decay Enhances Supervised Fine-Tuning” compares WSD-style decay-based schedules to Warmup-Stable-Only (WSO), which removes the terminal decay while keeping warmup and the stable phase. Across 1B and 8B models, including mid-training and over-training regimes, WSO gives worse pretraining validation loss but better post-SFT performance. The same paper reports that decay-based schedulers lead to regions with about 2–3× higher sharpness than WSO and finds a negative correlation between pretraining sharpness and SFT performance with Pearson $T_c$4 (Yano et al., 17 Mar 2026). This does not refute WSD as a pretraining schedule, but it narrows its claim: optimizing the cooldown may improve the raw base model while reducing downstream adaptability.

A third line argues that WSD is structurally mismatched when high-quality data arrives late. “How Should LLMs Consume High-Quality Data?” studies midtraining with a discrete quality transition and reports that a batch-scheduling alternative, Drop–Stable–Rampup, outperforms WSD by $T_c$5 average accuracy on a 15B MoE model midtrained on 108B tokens, with particularly large gains on GSM8K $T_c$6 and MATH $T_c$7 (Zhu et al., 25 May 2026). The critique is not that late decay is always bad, but that decaying update intensity precisely when cleaner data becomes available misses a signal-amplification opportunity.

Warmup itself is also being re-theorized independently of full WSD. Recent work explains warmup through curvature that decreases with the loss gap or suboptimality gap and derives warmup-like schedules from generalized smoothness assumptions (Alimisis et al., 3 Oct 2025, Riabinin et al., 5 Feb 2026). Related empirical work argues that the main benefit of warmup is to enable larger target learning rates by forcing the model into better-conditioned regions, and proposes loss-catapult-based initialization of the warmup learning rate as well as modified Adam initialization that reproduces some warmup benefits (Kalra et al., 2024). These results strengthen the “W” in WSD, but they do not yet provide a unified theorem for the full practical three-phase schedule.

Taken together, the recent literature places WSD in a more precise position than the original heuristic framing. It is a schedule family with a clear practical advantage for open-ended training, a well-supported empirical signature in which the final decay does disproportionate work, and several increasingly sharp theoretical explanations for why long plateaus plus late annealing can be effective. At the same time, the literature no longer treats any single WSD instantiation as universally optimal. Cooldown shape, optimizer state, batch-size policy, downstream objective, and data-quality transitions all materially change the answer.