---
title: 'FG-WSD: Fine-Grained Warmup–Stable–Decay'
url: https://www.emergentmind.com/topics/fine-grained-warmup-stable-decay-fg-wsd-aaab6e14-922f-4b0b-97d0-972f6834f5f4
type: topic
---

# FG-WSD: Fine-Grained Warmup–Stable–Decay

Fine-Grained Warmup–Stable–Decay (FG-WSD) is a pretraining scheduler for large language models that jointly modulates learning rate and training data mixture. FG-WSD introduces a fine-grained design in the stable (plateau) regime, progressively increasing the proportion of high-quality (HQ) data over multiple sub-phases, and maintains a clear decoupling between data-mixture transitions and learning-rate changes. This approach was introduced in the pretraining procedure of the Nanbeige4-3B family of language models, where it demonstrated significant gains over conventional schedulers, particularly on reasoning-heavy tasks [2512.06266].

## 1. Motivation and Theoretical Rationale

FG-WSD arises from two key limitations in existing learning-rate and data-mixing schedules. First, while a standard Warmup–Stable–Decay (WSD) learning-rate profile empirically outperforms Warmup–Cosine–Decay when using very high-quality datasets, the uniform data mixture during the stable plateau is suboptimal. The model is effectively exposed to the same data-quality distribution (often 1:1 HQ vs. medium-quality) throughout this period. FG-WSD targets these limitations by introducing multiple sub-phases during the stable plateau, with each sub-phase using a higher proportion of HQ data and by decoupling learning-rate shifts from abrupt changes in data mixture.

This design enables an early-stage "exploration" period—when the model benefits from broader diversity—followed by a later "exploitation" phase prioritizing only the best data. This prevents undesirable optimization instabilities that can be caused by simultaneous drastic changes in learning rate and data distribution.

## 2. Formal Description and Mathematical Framework

Let $t$ denote the cumulative number of training tokens. FG-WSD divides training into four principal stages, parameterized by token horizons:
- $T_{\mathrm{w}}$: warm-up horizon;
- $T_{\mathrm{s}1}$: diversity-enriched stable phase;
- $T_{\mathrm{s}2}$: HQ-only stable phase;
- $T_{\mathrm{d}}$: decay horizon.

The maximum plateau learning rate is $\eta_{\max}$ and the minimum end-of-decay learning rate is $\eta_{\min}$. Schedules are as follows:

**Learning-rate schedule:**
- $0 \le t \le T_{\mathrm{w}}$: $\eta(t) = \eta_{\max} \cdot (t/T_{\mathrm{w}})$ (linear warmup).
- $T_{\mathrm{w}} < t \le T_{\mathrm{w}} + T_{\mathrm{s}1} + T_{\mathrm{s}2}$: $\eta(t) = \eta_{\max}$ (constant).
- $T_{\mathrm{w}} + T_{\mathrm{s}1} + T_{\mathrm{s}2} < t \le T_{\mathrm{w}} + T_{\mathrm{s}1} + T_{\mathrm{s}2} + T_{\mathrm{d}}$:
  $$
  \eta(t) = \eta_{\min} + \frac{1}{2} (\eta_{\max} - \eta_{\min}) \big[1 + \cos\big(\pi \frac{t - (T_{\mathrm{w}} + T_{\mathrm{s}1} + T_{\mathrm{s}2})}{T_{\mathrm{d}}}\big)\big]
  $$
- Training may terminate beyond this or continue at $\eta_{\min}$.

**Data-mixture schedule:** Let $w_\mathrm{HQ}(t)$ denote the probability of drawing a sample from HQ data at time $t$:
- $0 \le t \le T_{\mathrm{w}}$: $w_\mathrm{HQ}(t) = \alpha_\mathrm{warmup}$ (often close to the first stable phase's proportion).
- $T_{\mathrm{w}} < t \le T_{\mathrm{w}} + T_{\mathrm{s}1}$: $w_\mathrm{HQ}(t) = \alpha_\mathrm{S1}$ (e.g., 50% HQ, 50% medium-quality).
- $T_{\mathrm{w}} + T_{\mathrm{s}1} < t \le T_{\mathrm{w}} + T_{\mathrm{s}1} + T_{\mathrm{s}2}$: $w_\mathrm{HQ}(t) = 1.0$ (HQ only).
- During decay, maintain $w_\mathrm{HQ} = 1$.

This staged, piecewise-constant schedule ensures smooth transitions for both learning rate and data mixture.

## 3. Hyperparameter Settings and Practical Instantiation

For Nanbeige4-3B-Base, FG-WSD was instantiated with the following parameters (token units: trillions, $T=10^{12}$):

| Stage                | Tokens | Learning Rate Behavior         |
|----------------------|--------|-------------------------------|
| Warmup               | 0.1 T  | linear $0 \rightarrow 4.5 \times 10^{-4}$ |
| Diversity-Enriched   | 12.4 T | constant $4.5 \times 10^{-4}$ |
| HQ Stable            | 6.5 T  | constant $4.5 \times 10^{-4}$ |
| Decay                | 4.0 T  | $4.5 \times 10^{-4} \rightarrow 1.5 \times 10^{-6}$ |

- $T_\mathrm{w} = 0.1$ T (warm-up)
- $T_{\mathrm{s}1} = 12.4$ T (diversity-enriched)
- $T_{\mathrm{s}2} = 6.5$ T (HQ stable)
- $T_\mathrm{d} = 4.0$ T (decay)
- $\eta_{\max} = 4.5 \times 10^{-4}, \eta_{\min} = 1.5 \times 10^{-6}$
- Data mixture: diversity-enriched phase uses roughly $50\%$ HQ, $50\%$ other filtered data; HQ-stable and decay phases use only HQ.

Hyperparameters were derived from ablations on a 1 B-parameter model with a toy 1 T-token corpus: splitting the stable phase ($0.75$ T mixed medium/HQ, $0.25$ T HQ only) yielded noticeable gains, motivating proportional scaling to larger runs.

## 4. Data-Mixture Schedule and Evolution

FG-WSD specifies the evolution of training data mixture alongside learning rate. The four main phases are:
- **Warmup ($0.1$ T):** Data mixture approximates the ensuing diversity-enriched phase.
- **Diversity-Enriched Stable ($12.4$ T):** Mini-batches sampled from the full corpus with approximately $50\%$ HQ and $50\%$ medium/filtered data.
- **HQ Stable ($6.5$ T):** Sampling exclusively from the HQ subset.
- **Decay ($4.0$ T):** Continued exclusive HQ sampling, with a decaying learning rate.

Earlier ablations compared vanilla WSD (uniform 1:1 shuffle of medium and high-quality data) to FG-WSD, confirming benefits from prolonging the mixed phase and carefully transitioning to HQ only.

## 5. Pseudocode and Training Workflow

The FG-WSD scheduler is implemented as follows:

```python
initialize model parameters θ
set ηmax = 4.5e-4, ηmin = 1.5e-6
set Tw = 0.1*T_total,  Ts1 = 12.4*T_total,  Ts2 = 6.5*T_total,  Td = 4.0*T_total
define data pools: Pool_HQ, Pool_MQ (and lower)
for token_count t from 1 to T_total:
  # 1. Determine LR
  if t <= Tw:
    η = ηmax * (t / Tw)
  else if t <= Tw + Ts1 + Ts2:
    η = ηmax
  else if t <= Tw + Ts1 + Ts2 + Td:
    τ = (t - (Tw + Ts1 + Ts2)) / Td
    η = ηmin + 0.5*(ηmax-ηmin)*(1 + cos(pi * τ))
  else:
    η = ηmin
  # 2. Determine data-mixture weight wHQ
  if t <= Tw:
    wHQ = Ts1 > 0 ? Ts1/(Ts1+Ts2) : 0.0   # or preset αwarmup
  else if t <= Tw + Ts1:
    wHQ = Pool_HQ_size / (Pool_HQ_size + Pool_MQ_size)   # ~0.5
  else:
    wHQ = 1.0
  # 3. Sample a mini-batch
  for each sample in batch:
    if random() < wHQ:
      x ← sample_from(Pool_HQ)
    else:
      x ← sample_from(Pool_MQ ∪ others)
  # 4. Take optimizer step at learning rate η
  loss = compute_loss(θ, x)
  θ ← θ - η * ∇_θ loss
```

## 6. Experimental Results and Empirical Impact

On a 1 B-parameter model with 1 T tokens of pretraining, FG-WSD achieved substantial improvements over vanilla WSD, most notably on mathematical reasoning tasks:

| Metric    | Vanilla WSD | FG-WSD | Absolute Gain |
|-----------|-------------|--------|--------------|
| GSM8K     | 27.1        | 34.3   | +7.2         |
| CMath     | 34.5        | 39.5   | +5.0         |
| BBH       | 29.3        | 31.6   | +2.3         |
| MMLU      | 49.2        | 50.6   | +1.4         |
| CMMLU     | 50.3        | 51.9   | +1.6         |
| MMLU-Pro  | 16.87       | 18.64  | +1.77        |

Gains were especially pronounced on reasoning benchmarks, reflecting the positive effect of concentrating on HQ, reasoning-dense data in the latter stages. These improvements motivated the scaling and full deployment of FG-WSD in the 23 T-token pretraining of Nanbeige4-3B, contributing to state-of-the-art few-/zero-shot performance among $3$ B-parameter models [2512.06266].

## 7. Context and Implications

FG-WSD represents a shift in scheduler design for large-scale LM pretraining, emphasizing decoupled and staged optimization of both learning-rate and data-mixture profiles. The results in Nanbeige4-3B suggest that the joint fine-graining of plateau phase data curation and smooth learning-rate transitions is particularly effective for maximizing generalization and reasoning ability at sub-10B parameter scales. A plausible implication is that future schedulers may increasingly incorporate granular, data-aware control at the epoch or token level to further leverage high-quality data advantages. The use of such schedulers is predicated on rigorous data curation and quality filtering capabilities.

Source: https://www.emergentmind.com/topics/fine-grained-warmup-stable-decay-fg-wsd-aaab6e14-922f-4b0b-97d0-972f6834f5f4