---
title: Warmup–Stable–Merge Learning Rate Schedule
url: https://www.emergentmind.com/topics/warmup-stable-merge-learning-rate-schedule
type: topic
---

# Warmup–Stable–Merge Learning Rate Schedule

A Warmup–Stable–Merge (WSM) learning rate schedule, introduced as a decay-free alternative for large-scale language model 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:
\[
\hat\theta = \sum_{j=0}^{K} c_j \, \theta_{t_0 + j \Delta t}
\]
where $\{\theta_{t_0},\ldots,\theta_{t_0+K\Delta t}\}$ are consecutive saved models, $c_j$ are merge weights with $\sum_j c_j=1$, $\Delta t$ is the checkpoint interval, and $T_m = K \Delta t$ is the total merge window duration [2507.17634].

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 $K+1$ checkpoints generated under SGD, a merge defined by weights $\{c_j\}$,
\[
\hat\theta_{n+K} = \sum_{j=0}^K c_j \theta_{n+j}
\]
can be re-expressed as a monotone reweighting of the trailing $K$ stochastic gradients:
\[
\hat\theta_{n+K} = \theta_n - \sum_{i=1}^K w_i\, g_{n+i-1}, \quad w_i = \sum_{j=i}^K c_j
\]
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 [2507.17634]. Explicit matching between desired decay curves $\{w_i\}$ and merge weights $\{c_j\}$ is given by
\[
c_K = w_K,\quad c_j = w_j - w_{j+1} \ (1 \le j < K),\quad c_0 = 1 - w_1
\]
Examples:
- **Linear decay:** $w_i = 1 - \frac{i-1}{K}$, $c_j = 1/(K+1)$ (uniform average).
- **Cosine decay:** $w_i = \frac{1}{2}(1+\cos\frac{\pi(i-1)}{K})$, $c_j = w_j - w_{j+1}$, $c_0 = 1 - w_1$.
- **Inverse-sqrt decay:** $w_i = 1-\sqrt{\frac{i-1}{K}}$, $c_j = w_j - w_{j+1}$, $c_0 = 1 - w_1$.

## 3. Key Hyperparameters and Practical Implementation

Empirical results establish that *merge duration* $T_m = K \Delta t$ is the single most sensitive and impactful hyperparameter. A typical recipe for large-scale LLM pre-training is:
- **Merge duration ($T_m$):** $T_m \approx 100$ billion tokens. Performance saturates beyond $T_m \approx 200$B tokens.
- **Checkpoint interval ($\Delta t$):** $[5\text{B}, 25\text{B}]$ tokens, with $\Delta t = 25$B as a robust default.
- **Number of checkpoints ($K$):** $4 \leq K \leq 16$; $K$ too small (<4) undersamples, too large (>16) offers limited gain.

Granularity ($\Delta t$) and merge shape (weight profile) have modest effects compared to $T_m$. The merge window should cover late stable-phase training, maximizing effective annealing without redundant snapshots [2507.17634].

| Hyperparameter        | Typical Range  | Primary Importance             |
|----------------------|---------------|-------------------------------|
| Merge duration $T_m$ | 100–200B toks | Dominant; controls annealing  |
| Checkpoint interval $\Delta t$ | 5–25B toks | Moderate; tradeoff vs. storage |
| # of checkpoints $K$ | 4–16          | Minor beyond $K\ge8$          |

In practice, the simplest *linear-average* merge ($c_j=1/(K+1)$) 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 [2507.17634]:
- **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_m$ provides more annealing and thus better performance, with clear saturation beyond a threshold.
- **Merge algorithm:** 1–√ $>$ uniform mean $>$ EMA.
- **Granularity:** Finer checkpointing marginally outperforms coarser, but is distinctly less important than $T_m$.
- **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 [2507.17634].

Further, WSM maintains its advantage after supervised fine-tuning, with $\sim$+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$ at the stability limit for nearly all training, and apply annealing only at the very end [2602.04774; 2602.06797].
- In the WSM framework, the “annealing” is entirely via model averaging (gradient reweighting), and no explicit late-phase decay is required. This eliminates the $\log$-penalty in generalization curves associated with exponential or power-law decays, achieving power-law optimality in both data- and compute-limited settings [2509.19189; 2602.04774].

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 $T_1$ (warmup), $T_2$ (stable), merge window $T_m$, and snapshotting frequency $\Delta t$.
- No decay hyperparameters, curve fitting, or power exponents are needed.
- Once $T_m$ and $K$ 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 [2507.17634].

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 [2503.12811; 2602.04774; 2602.06797]. In the prevalent “hard phase” for LLMs (small signal exponent $s$, large $\beta$), 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 [2507.17634], with supporting insights from [2602.04774], [2509.19189], [2602.06797], and [2503.12811].

Source: https://www.emergentmind.com/topics/warmup-stable-merge-learning-rate-schedule