Warmup–Stable–Merge Learning Rate Schedule

• The paper demonstrates that WSM replaces explicit decay with a principled checkpoint merging operation that emulates late-stage annealing.
• WSM divides training into three phases—warmup, stable, and merge—reducing hyperparameter complexity and streamlining schedule design.
• Empirical and theoretical results show that WSM enhances both convergence and downstream performance in large-scale language model pre-training.

A Warmup–Stable–Merge (WSM) learning rate schedule, introduced as a decay-free alternative for large-scale LLM pre-training, is a framework that eschews the traditional explicit decay phase present in schedules such as Warmup–Stable–Decay (WSD). Instead, WSM achieves effective late-stage “annealing” via a principled model-merging operation over a window of checkpoints at the end of training. Theoretical and empirical research demonstrates clear connections between gradient-based learning rate decay and post-hoc model averaging, and highlights practical advantages for both convergence and downstream performance in LLM pre-training and supervised fine-tuning.

1. Structure and Formulation of the WSM Schedule

WSM divides training into three distinct phases, each characterized by a specific learning rate or post-processing operation:

a. Warmup ($t < T_1$):

A strictly linear ramp from $\eta(t) = 0$ at step $t = 0$ up to the peak value $\eta_{\rm peak}$ at $t = T_1$: $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$

b. Stable ($T_1 \le t < T_1+T_2$):

A prolonged plateau with fixed learning rate: $\eta(t) = \eta_{\rm peak}$

c. Merge ($t \ge T_1+T_2$):

No further training or explicit decay. Instead, the final model is produced by merging the last $K+1$ checkpoints from the stable phase using a specific convex combination determined by the desired “annealing” curve: $\eta(t) = 0$0 where $\eta(t) = 0$1 are consecutive saved models, $\eta(t) = 0$2 are merge weights with $\eta(t) = 0$3, $\eta(t) = 0$4 is the checkpoint interval, and $\eta(t) = 0$5 is the total merge window duration (Tian et al., 23 Jul 2025).

The resulting live learning rate is strictly non-decaying in the final phase; all annealing is achieved by the merge operator.

2. Formal Connection to Classical Learning Rate Decay

WSM provides a theoretical bridge between learning rate schedules and checkpoint merging. Any classical decaying learning rate schedule can be emulated by a properly chosen checkpoint-average. Given a sequence of $\eta(t) = 0$6 checkpoints generated under SGD, a merge defined by weights $\eta(t) = 0$7,

$\eta(t) = 0$8

can be re-expressed as a monotone reweighting of the trailing $\eta(t) = 0$9 stochastic gradients: $t = 0$0 This equivalence (Theorem 3.1) enables WSM to exactly reproduce the effect of any target decay profile (e.g., linear, cosine, inverse square root) via discrete model averaging (Tian et al., 23 Jul 2025). Explicit matching between desired decay curves $t = 0$1 and merge weights $t = 0$2 is given by

$t = 0$3

Examples:

• Linear decay: $t = 0$4, $t = 0$5 (uniform average).
• Cosine decay: $t = 0$6, $t = 0$7, $t = 0$8.
• Inverse-sqrt decay: $t = 0$9, $\eta_{\rm peak}$0, $\eta_{\rm peak}$1.

3. Key Hyperparameters and Practical Implementation

Empirical results establish that merge duration $\eta_{\rm peak}$2 is the single most sensitive and impactful hyperparameter. A typical recipe for large-scale LLM pre-training is:

• Merge duration ($\eta_{\rm peak}$3): $\eta_{\rm peak}$4 billion tokens. Performance saturates beyond $\eta_{\rm peak}$5B tokens.
• Checkpoint interval ($\eta_{\rm peak}$6): $\eta_{\rm peak}$7 tokens, with $\eta_{\rm peak}$8B as a robust default.
• Number of checkpoints ($\eta_{\rm peak}$9): $t = T_1$0; $t = T_1$1 too small (<4) undersamples, too large (>16) offers limited gain.

Granularity ($t = T_1$2) and merge shape (weight profile) have modest effects compared to $t = T_1$3. The merge window should cover late stable-phase training, maximizing effective annealing without redundant snapshots (Tian et al., 23 Jul 2025).

Hyperparameter	Typical Range	Primary Importance
Merge duration $t = T_1$4	100–200B toks	Dominant; controls annealing
Checkpoint interval $t = T_1$5	5–25B toks	Moderate; tradeoff vs. storage
# of checkpoints $t = T_1$6	4–16	Minor beyond $t = T_1$7

In practice, the simplest linear-average merge ($t = T_1$8) is nearly optimal for linear decay, while the 1–√ weight schedule empirically outperforms EMA and uniform for LLMs.

4. Comparative Analysis: WSM vs. Warmup–Stable–Decay (WSD)

WSM consistently surpasses WSD (explicit decay) schedules across diverse tasks. In extensive MoE transformer pre-training experiments (Tian et al., 23 Jul 2025):

• Model: 16.3B-parameter MoE transformer (1.43B active), trained on 10.2T tokens mixed English/Chinese.
• Benchmark results (WSM vs. WSD):
  • MATH: 58.81% vs. 57.49% (+3.5% relative)
  • HumanEval: 65.58% vs. 64.88% (+2.9%)
  • MMLU-Pro: 45.2% vs. 43.3% (+5.5%)
  • Overall: +2% relative

Ablation studies confirm:

• Merge duration dominates—longer $t = T_1$9 provides more annealing and thus better performance, with clear saturation beyond a threshold.
• Merge algorithm: 1–√ $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$0 uniform mean $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$1 EMA.
• Granularity: Finer checkpointing marginally outperforms coarser, but is distinctly less important than $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$2.
• Hybrid or stacked schedules (e.g., live decay plus post-hoc merge) confer no advantage, validating the formal equivalence between WSM-type merging and decay-based anneal (Tian et al., 23 Jul 2025).

Further, WSM maintains its advantage after supervised fine-tuning, with $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$3+1.8–1.9% improvements on base and instruction-following tasks.

5. Theoretical Rationale and Scaling Behavior

WSM is grounded in loss-dynamics theory and directly connects with functional scaling law (FSL) and optimal control perspectives:

• Classical FSL and optimal control analyses establish that, in the “hard” regime (slow signal learning, high variance), the optimal policy is to maximize $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$4 at the stability limit for nearly all training, and apply annealing only at the very end (Bordelon et al., 4 Feb 2026, Li et al., 6 Feb 2026).
• In the WSM framework, the “annealing” is entirely via model averaging (gradient reweighting), and no explicit late-phase decay is required. This eliminates the $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$5-penalty in generalization curves associated with exponential or power-law decays, achieving power-law optimality in both data- and compute-limited settings (Li et al., 23 Sep 2025, Bordelon et al., 4 Feb 2026).

This theoretical unity explains the empirical parity between decay-based and merge-based approaches: both achieve the optimal minimax rates, and merging recovers all the variance reduction and bias control of explicit decay.

6. Practical Implications and Guidelines

WSM is optimizer-agnostic and requires no complicated schedule design:

• The user must specify only $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$6 (warmup), $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$7 (stable), merge window $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$8, and snapshotting frequency $\eta(t) = \eta_{\rm peak} \frac{t}{T_1}$9.
• No decay hyperparameters, curve fitting, or power exponents are needed.
• Once $T_1 \le t < T_1+T_2$0 and $T_1 \le t < T_1+T_2$1 are fixed, merges can be updated online or retrospectively, with minimal compute and storage overhead.
• Model merging as annealing is robust to model scale, batch size, and optimizer variant, and consistently outperforms or matches all tuned explicit-decay alternatives across tested LLM benchmarks (Tian et al., 23 Jul 2025).

Repeated findings indicate that, when model merging is permitted, the entire purpose of late-stage decay is effectively realized via a tailored ensemble of recent solutions, obviating the need for decaying the step size in the live run.

7. Broader Context and Extensions

Recent literature on loss prediction and optimal schedules via multi-power laws and SDE-based FSL frameworks corroborate that WSM and WSD are functionally equivalent solutions to the late-stage annealing problem in large-scale training (Luo et al., 17 Mar 2025, Bordelon et al., 4 Feb 2026, Li et al., 6 Feb 2026). In the prevalent “hard phase” for LLMs (small signal exponent $T_1 \le t < T_1+T_2$2, large $T_1 \le t < T_1+T_2$3), optimal schedules universally take a warmup-stable-decay/merge form, with decay/merge restricted to a vanishing late-stage tail. This vindicates the practical dominance of WSM- and WSD-type schedules in large-scale pre-training, and foregrounds model averaging as a simple, resource-efficient, and robust protocol for annealing.

All claims and results are derived from (Tian et al., 23 Jul 2025), with supporting insights from (Bordelon et al., 4 Feb 2026, Li et al., 23 Sep 2025, Li et al., 6 Feb 2026), and (Luo et al., 17 Mar 2025).