Warmup-Stable-Only (WSO): Stability in Neural Training

Updated 3 July 2026

WSO is a training paradigm that replaces conventional warmup and decay with strategies that ensure spectral, curvature, and representational stability in large-scale models.
It employs both schedule-centric and stability-centric methods to maintain a constant learning rate after a brief warmup, preventing early-stage instabilities like gradient explosions and entropy collapse.
Empirical evidence shows that WSO enhances downstream transfer performance and training efficiency across Transformers, vision models, and large language models.

Warmup-Stable-Only (WSO) refers to a family of principles, learning-rate schedules, and optimization strategies that replace or reinterpret the conventional learning-rate warmup and decay phases in large-scale neural network training. WSO has become central in contemporary training of Transformers and LLMs, providing both empirical and theoretical improvements in stability, efficiency, and downstream performance. The unifying theme is the emphasis on achieving stability—whether spectral, curvature, or representational—without prolonged manual warmup or late-phase decay.

1. Definitions and Core Schedules

WSO denotes two related but distinct ideas:

Schedule-centric WSO: A learning-rate schedule in which the learning rate is linearly ramped from zero to a set maximum value over an initial “warmup” period, then held constant through the remainder of training, with no subsequent decay phase. Formally, for learning rate $\eta_{\max}$ , warmup steps $T_{\mathrm{warmup}}$ , and total steps $T_{\mathrm{pre}}$ :

$\text{WSO}(t) = \begin{cases} \eta_{\max} \,\frac{t}{T_{\mathrm{warmup}}}, & t \le T_{\mathrm{warmup}} \[6pt] \eta_{\max}, & T_{\mathrm{warmup}} < t \le T_{\mathrm{pre}} \end{cases}$

This definition appears, for example, in the context of LLM pre-training (Yano et al., 17 Mar 2026).

Stability-centric WSO: An optimization principle where, instead of using a preset warmup curve, updates are dynamically regulated so that weight changes remain spectrally or curvature-wise “small”, preventing early-stage instabilities (e.g., entropy collapse in attention, gradient explosions). This can be realized by bounding the largest singular value of weight updates relative to weights, or via curvature-aware depth-unfreezing (Qi et al., 28 May 2025, Ramasinghe et al., 15 Jun 2026).

The table below summarizes schedule-centric and stability-centric variants:

Variant	Mechanism	Decay Phase	Underlying Rationale
Linear WSO	Linear ramp, constant LR	Absent	Simplicity, stability
Spectral WSO	Adaptive LR bounds	Absent	Spectral smoothness
Curvature WSO	Depth/curvature control	Absent	Hessian-driven stability

2. Theoretical Foundations

WSO arises as an optimal or nearly-optimal policy in several theoretical frameworks:

Functional Scaling Laws (FSL): Under FSL, optimal learning rate schedules for ill-conditioned (“hard task”) regimes asymptotically approach a “warmup–stable–only” two-phase structure: a power-law warmup to the maximum stable learning rate, then hold that value constant, with decay occupying a vanishing fraction of training (Li et al., 6 Feb 2026). Formally, for capacity exponent $\beta$ and source exponent $s$ with $s < 1 - 1/\beta$ ,

$\eta_{\mathrm{WSO}}(z) = \begin{cases} \bigl(\frac{z}{T_{\mathrm{warmup}}}\bigr)^{2\beta-1}\,\eta_{\mathrm{peak}}, & 0 \le z \le T_{\mathrm{warmup}} \ \eta_{\mathrm{peak}}, & T_{\mathrm{warmup}} < z \le N \end{cases}$

Spectral Stability: In Transformer training, instability such as entropy collapse is provably linked to spectral energy concentration of the attention matrices (notably $W_q^{\top} W_k$ ). If weight updates remain small in spectral norm compared to the weights, one prevents excessive growth in top singular values, thus averting pathologies like near-zero attention entropy and rank collapse (Qi et al., 28 May 2025).
Curvature Control: The Edge-of-Stability phenomenon motivates bounding the product $\eta_t \lambda_{\max}(G_t)$ , where $T_{\mathrm{warmup}}$ 0 is the preconditioned Hessian, to remain under a critical threshold. WSO-inspired strategies progressively unfreeze network depth or adapt learning rates to keep curvature within a stable zone, widening the effective learning-rate window (Ramasinghe et al., 15 Jun 2026).

3. Empirical Evidence and Schedule Comparisons

WSO alternatives have been compared to classical warmup+decay schedules, cosine/linear decay, and variance-rectified optimizers across vision and language tasks:

Adaptive Optimization Baselines: Linear WSO–for example, Adam with untuned warmup length $T_{\mathrm{warmup}}$ 1–matches or outperforms RAdam and other heuristic alternatives in ImageNet, Wikitext-103, and WMT benchmarks for both stability and final accuracy. Empirical loss and perplexity curves for WSO and RAdam are indistinguishable after warmup; vanilla Adam without warmup exhibits catastrophic instability (Ma et al., 2019).
Vision Transformer Experiments: Vision Transformer (ViT), Swin Transformer, and GPT-Small trained with AdamW² (spectral WSO) and no explicit warmup achieve equivalent or superior performance compared to baseline recipes with canonical warmup, and with consistently smooth training curves (Qi et al., 28 May 2025).
LLM Pre-training and Transfer: Constant-LR WSO pre-training of LLMs yields worse pre-training perplexity but superior downstream supervised fine-tuning (SFT) performance on Tulu-3, AlpacaEval, TruthfulQA, and MMLU benchmarks, across models up to 8B parameters. Pre-trained WSO models achieve lower Hessian trace (“sharpness”), which is strongly negatively correlated to SFT task average (Yano et al., 17 Mar 2026).

Empirical outcomes are robust across data modalities, batch sizes, architectures, and scaling regimes.

4. Mechanistic Interpretations and Analysis

WSO achieves stability and transfer via distinct mechanisms depending on the operationalization:

Representational Stability: For large-batch settings, WSO ensures the deepest layers (top fully-connected stack) are not destabilized by abrupt parameter shifts; layerwise similarity measured by SVCCA stays $T_{\mathrm{warmup}}$ 2 with warmup or explicit freezing, but collapses to $T_{\mathrm{warmup}}$ 3 for large-batch no-warmup (Gotmare et al., 2018).
Spectral Smoothing: By enforcing

$T_{\mathrm{warmup}}$ 4

where $T_{\mathrm{warmup}}$ 5 is the top singular value and $T_{\mathrm{warmup}}$ 6 is a small threshold, one bounds the growth of spectral energy and avoids collapse into low-rank/low-entropy attention, maintaining stable optimization trajectory (Qi et al., 28 May 2025).

Curvature Tracking and Controlled Depth: Real-time estimation of the largest preconditioned Hessian eigenvalue allows progressive depth unfreezing only when the stability margin is sufficient ( $T_{\mathrm{warmup}}$ 7), eliminating characteristic loss spikes at billion-parameter scale (Ramasinghe et al., 15 Jun 2026).
Loss Landscape Geometry: WSO (constant LR after warmup) yields flatter minima by pre-training end (lower Hessian trace), in contrast to decay-based schedules, which lead to sharper (higher curvature) minima, resulting in lower adaptability to supervised gradients during SFT (Yano et al., 17 Mar 2026).

5. Practical Implementations and Prescriptions

WSO-inspired schedules or stabilization mechanisms are straightforward to deploy:

Schedule-centric:

Select AdamW or similar optimizer.
Set warmup length: $T_{\mathrm{warmup}}$ 8 (e.g., 2000 steps for $T_{\mathrm{warmup}}$ 9).
Linear ramp up to target LR over $T_{\mathrm{pre}}$ 0 steps, then hold constant; omit late-phase decay.
Optionally, retain whatever LR decay is customary for final convergence, but for transferability, pure WSO is preferred (Ma et al., 2019, Yano et al., 17 Mar 2026).

Spectral/curvature-centric:

For each weight matrix or parameter group, estimate the current largest singular value ( $T_{\mathrm{pre}}$ 1) and update norm ( $T_{\mathrm{pre}}$ 2) via a few (e.g., 3) power iterations.
Restrict the effective learning rate $T_{\mathrm{pre}}$ 3 to $T_{\mathrm{pre}}$ 4.
For model depth, activate layers progressively after the initial warmup, conditioned on curvature criteria (Qi et al., 28 May 2025, Ramasinghe et al., 15 Jun 2026).

Empirical Tuning: For most settings, the WSO regimen requires minimal tuning; stability is robust w.r.t. warmup length and schedule details. Piecewise WSO is the canonical regime for ill-conditioned (hard) tasks, where decay can often be omitted entirely (Li et al., 6 Feb 2026).

6. Implications for Transfer, Adaptability, and Model Release

Downstream Task Robustness: WSO pre-training produces models that fine-tune more effectively and robustly on downstream supervised objectives, with less need for hyperparameter retuning or scheduler redesign. This effect persists even for multi-stage (pre + mid + over) training pipelines (Yano et al., 17 Mar 2026).
Model Distribution and Open-Source: WSO is recommended as a model release default; practitioners adopting such checkpoints for transfer or SFT can expect more consistent adaptability and reduced overfitting to pre-training loss minimization.
Efficiency and Scaling: WSO removes the need for manual warmup-length tuning and avoids the training slowdowns introduced by both long initial ramps and aggressive late-phase decay. This yields more efficient large-scale training at both hardware and wall-clock levels (Qi et al., 28 May 2025, Ramasinghe et al., 15 Jun 2026).

7. Limitations and Comparisons to Alternative Stabilization

Comparison to RAdam and QK-Norm/Spectral Clipping: Untuned WSO matches or outperforms RAdam and attention-specific heuristics (QK-Norm, QK-Clip, Softcap) in both vision and LLMs, especially for large peak learning rates. Alternative curvature-controlling methods (adaptive $T_{\mathrm{pre}}$ 5 via Hessian estimation) impose higher computational costs (>20 HVPs/step), whereas WSO-inspired strategies require only a handful (Qi et al., 28 May 2025, Ramasinghe et al., 15 Jun 2026, Ma et al., 2019).
Open Problems: While WSO effectively resolves instability in most large-batch or high-LR regimes, marginal benefits of late-stage decay for final convergence, and its effect under certain pathologically hard optimization landscapes, remain active research topics. If downstream adaptability is not required, decay-based or multi-phase schedules may still yield lower pre-training loss (Yano et al., 17 Mar 2026).

WSO recasts the stabilizing role of traditional warmup and decay as enforcing small, stable, and smooth steps in parameter space, achievable through either simple schedules or explicit spectral/curvature control. It is now an empirically validated, theoretically grounded strategy for achieving both robust large-scale training and superior downstream transfer in contemporary neural network practice (Gotmare et al., 2018, Ma et al., 2019, Qi et al., 28 May 2025, Li et al., 6 Feb 2026, Yano et al., 17 Mar 2026, Ramasinghe et al., 15 Jun 2026).