--- title: Warmup–Stable–Decay (WSD) LR Schedule url: https://www.emergentmind.com/topics/warmup-stable-decay-wsd-pattern type: topic --- # Warmup–Stable–Decay (WSD) LR Schedule The Warmup–Stable–Decay (WSD) Pattern is a three-phase learning rate (LR) scheduling paradigm that has become fundamental in large-scale training of deep models, especially large language models (LLMs) and transformers. It decomposes the learning rate trajectory into (1) a warmup phase to suppress divergence, (2) a sustained plateau (“stable”) at the peak learning rate, and (3) a cooldown or decay phase that anneals the LR toward zero or a small floor. WSD outperforms traditional cosine or direct-decay schedules in compute/data efficiency, loss convergence, and practical flexibility. Its mechanistic foundations are now understood: WSD exploits the geometric structure of high-dimensional loss surfaces (“valley–river” landscapes), formal scaling-law optimality, and even statistical physics analogies such as the Mpemba effect. ## 1. Mathematical Formulation and Variants The canonical WSD learning rate schedule with $T$ total steps is defined by three parameters: number of warmup steps $N_w$, plateau (“stable”) steps $N_s$, and decay steps $N_d$, with $T = N_w + N_s + N_d$. Let $\eta_{\max}$ denote the plateau learning rate, and $\eta_{\min}$ the terminal value. \[ \eta_t = \begin{cases} \eta_{\max} \frac{t}{N_w}, & 0 \leq t \leq N_w \text{ (warmup)} \\ \eta_{\max}, & N_w < t \leq N_w+N_s \text{ (stable)} \\ \eta_{\min} + (\eta_{\max} - \eta_{\min})\, f\left(\frac{t-(N_w+N_s)}{N_d}\right), & N_w+N_s < t \leq T \text{ (decay)} \end{cases} \] with $f(u)$ a monotonically decreasing function, common choices being $f(u) = 2^{-u}$ (exponential), $f(u) = 0.5[1+\cos(\pi u)]$ (cosine), or $f(u) = (1-u)^p$ (power/polynomial) ([2404.06395], [2507.09846], [2601.09000]). Tabulated examples of decay functions: | Decay type | $f(u)$ form | Typical behavior | |-----------------|-----------------------|---------------------------------------| | Linear decay | $1-u$ | Uniform decrease to zero | | Cosine decay | $0.5[1+\cos(\pi u)]$ | Slow – fast – slow; smooth tail-off | | Exponential | $2^{-u}$ | Fastest at start of decay | | Power | $(1-u)^p$ | Tunable shape (e.g., $p=1.5$ optimal) | ## 2. Theoretical Foundations and Scaling Laws WSD’s three-phase structure is motivated by functional scaling law (FSL) theory and convex/nonconvex optimization analysis. Under the FSL framework, for teacher–student kernel regression and SGD with data/model power-law spectra, the optimal LR schedule for hard tasks (capacity exponent $\beta > 1$, source exponent $s < 1 - 1/\beta$) is WSD: a long plateau at $\eta_{\max}$ (subject to stability), followed by a short vanishing polynomial decay of exponent $2\beta - 1$ ([2602.06797], [2509.19189], [2503.12811]). Key results: - The “stable” phase builds up intrinsic time (aggregate step magnitude) required for learning the signal. - By delaying decay, one concentrates all loss-reduction into the final phase, which the SGD forgetting kernel $K(t)$ renders effective for noise suppression. - The decay window may shrink as $r_N = N^{-\delta}$, so for large $N$, almost all steps are spent at peak LR ([2602.06797]). - Empirically, losses predicted by FSL for WSD match actual LLM pretraining to $<0.01$ loss units ([2509.19189]). For norm-constrained optimizers and optimization under suboptimality-dependent smoothness, WSD arises as the step-size $\eta^t = \frac{\Delta^t}{D(K_0 + K_1\Delta^t + K_\rho (\Delta^t)^\rho)}$ naturally transitions from a ramp (“warmup”), through a plateau (peak), and finally decays as $1/t$ once the gap $\Delta^t$ is small ([2602.05813]). ## 3. Geometric and Dynamical Mechanisms: The Valley–River Model Modern analyses interpret WSD through the “valley–river” or “river valley” landscape, where: - Optimization trajectories decompose into stiff (valley, $x$) and flat (river, $y$) directions, i.e., $L(x, y) = c(y) + \frac{1}{2}a(y)x^2$ with $a(y)\gg 1$. - The SGD/Langevin dynamics see fast equilibration in $x$ directions, while slow progress and heavy stochastic oscillations persist in “river” directions at high LR ([2507.04206], [2410.05192]). - During the stable phase, the model oscillates above the true valley floor but moves rapidly downstream in $y$. When LR decays, these oscillations collapse and the parameter falls onto the true minimum, causing a pronounced loss drop (“loss cliff” phenomenon). The “Mpemba effect” analogy further clarifies the role of the strong plateau: like a hotter system cooling faster when quenched, choosing the “strong Mpemba point” $\eta^* = 2/(\lambda_{\max}+\lambda_{\min})$ maximally accelerates convergence during decay by canceling the slowest mode ([2507.04206]). ## 4. Practical Structure, Tuning, and Empirical Observations WSD schedules in practice use: - Warmup: $1-5\%$ of steps, linear or polynomial ramp to $\eta_{\max}$ to avoid divergence and permit optimizer statistics (e.g., AdamW’s moments) to stabilize. - Stable phase: $80-95\%$ of steps at constant $\eta_{\max}$. The learning rate is chosen by estimating Hessian extreme eigenvalues (Lanczos/PCA) and using $\eta_{\rm plat} \approx 2/(\lambda_{\max}+\lambda_{\min})$ ([2507.04206]). - Decay: Final $5-15\%$ of steps; schedule may be cosine, polynomial, or exponential, decaying to a small value or zero ([2404.06395], [2507.09846]). Characteristics: - The loss remains nearly flat during the plateau, with most actual progress in the river (low-curvature) direction being hidden by hill-direction variance ([2508.01483], [2410.05192]). - Upon entering decay, loss drops sharply. Decay shape and duration significantly influence final perplexity—sqrt or “lowered linear” are optimal for bias-variance trade-off ([2508.01483]). - The phase fractions are robust across scales and models ([2601.09000]). Optimizing batch size for WSD: The E(S) theory reveals the existence of $B_{\min}$ (minimum batch for progress) and $B_{\opt}$ (batch minimizing total token consumption), both increasing with training depth ([2601.05034]). ## 5. Applications and Extensions: Scaling, Domain Adaptation, and Checkpoint Management WSD’s “forkable” architecture enables practical flexibility: - Continuous training: Any plateau checkpoint can be decayed later, obviating the need for schedules tied to total compute a priori ([2410.05192], [2404.06395]). - Domain adaptation: Mixing domain-relevant data in decay improves downstream generalization relative to post-hoc fine-tuning ([2404.06395]). - Scaling laws: Training to a long plateau and decaying from multiple token levels covers a grid of model/data pairs efficiently, permitting data-model scaling-law measurement with $O(m)$ instead of $O(m^2)$ runs ([2404.06395]). - WSD-S: Single-branch variants that sequence successive decays for multiple budgets, matching per-budget tuned cosine baselines ([2410.05192]). Recent developments in checkpoint-averaging (WSM) reinterpret the decay phase as model merging, demonstrating that any desired decay can be emulated by merging a window of stable-phase checkpoints, improving both performance and generalization ([2507.17634]). ## 6. Comparative Analysis and Limitations Empirical studies show WSD outperforms cosine decay for the same compute, especially for large LLMs and in hard-signal regimes ([2602.06797], [2509.19189], [2503.12811], [2601.09000]). However, WSD schedules require a manual or heuristic trigger to start decay, and stable-phase checkpoints do not reveal final performance until decay is applied, complicating convergence diagnostics ([2507.09846]). Some alternatives such as Schedule-Free AdamW ([2507.09846]) and WSM checkpoint merging ([2507.17634]) offer similar or improved efficiency and flexibility by avoiding explicit decay phases. ## 7. Generalization Beyond Transformers and Universal Dynamics WSD’s effectiveness and underlying river–valley geometry are not specific to transformer LMs. Experiments on standard CNNs (CIFAR-10) reveal qualitatively identical optimizer path features (two dominant directions, sharpness increase during decay, similar quasi-convexity indices) ([2601.09000]). This suggests that the WSD pattern reflects universal geometric properties of high-dimensional loss landscapes in deep learning. --- In summary, the Warmup–Stable–Decay pattern is a theoretically justified, empirically validated, and practically flexible learning rate strategy that exploits timescale separation in optimization dynamics, optimally balances signal extraction and noise forgetting, and adapts naturally to scaling and checkpointing workflows ([2507.04206], [2602.06797], [2410.05192], [2601.09000], [2404.06395], [2509.19189], [2507.09846], [2508.01483], [2503.12811], [2601.05034], [2507.17634], [2602.05813]). Source: https://www.emergentmind.com/topics/warmup-stable-decay-wsd-pattern