Summary (Overview)

• This paper investigates the joint design of data schedules (how PT and MT token proportions change over training) and learning rate (LR) schedules in LLM training pipelines, showing they must be designed together rather than independently.
• The authors demonstrate a ranking reversal phenomenon: under LR decay throughout training (FDR), mixing PT data with MT data lowers final MT loss, but under constant LR, placing all MT data in a final pure MT stage is optimal.
• Through ablation studies, they find that LR decay should begin well after MT data enters training — the best tested relative LR decay durations lie between 10% and 40% of the MT data window, improving on both constant LR and decay matched to MT entry.
• The paper develops a theoretical model based on SGD on noisy quadratic objectives that explains why optimal data schedules depend on LR schedules, recovering the empirical ranking reversal and deriving a three-stage optimal schedule under LR decay.
• Additional findings show that gradual data schedules (ZL) reduce load imbalance spikes in MoE models compared to abrupt transitions (TS), and that joint scheduling shapes the PT-MT loss frontier.

---

Introduction and Theoretical Foundation

Background

LLM training pipelines typically consist of two phases:

• Pretraining (PT): uses broad web corpora
• Mid-training (MT): uses smaller, curated mixtures enriched with mathematics, code, and other targeted data

The standard approach transitions completely from PT data to MT data at some point. The paper formalizes this as a data schedule function $\gamma(t)$, which denotes the fraction of MT tokens in the batch at step $t$. The complete transition corresponds to $\gamma(t) = \mathbf{1}\{t > \tau_0\}$.

Key Question

With fixed total training compute and MT token budgets, is this complete transition to MT data optimal for minimizing final MT loss?

Schedule Families

The paper compares two schedule families:

Two-Stage (TS): jumps directly to a chosen MT ratio and holds it fixed:

$$\gamma_{\mathrm{TS}}(t) = \begin{cases} 0, & 0 \leq t < \tau_{\mathrm{TS}}, \\ \gamma_{\mathrm{end}}, & \tau_{\mathrm{TS}} \leq t \leq T, \end{cases} \quad \tau_{\mathrm{TS}} = T\left(1 - \frac{\alpha}{\gamma_{\mathrm{end}}}\right)$$

Zero-then-Linear (ZL): increases the MT ratio gradually:

$$\gamma_{\mathrm{ZL}}(t) = \begin{cases} 0, & 0 \leq t < \tau_{\mathrm{ZL}}, \\ \gamma_{\mathrm{end}} \frac{t - \tau_{\mathrm{ZL}}}{T - \tau_{\mathrm{ZL}}}, & \tau_{\mathrm{ZL}} \leq t \leq T, \end{cases} \quad \tau_{\mathrm{ZL}} = T\left(1 - \frac{2\alpha}{\gamma_{\mathrm{end}}}\right)$$

The MT token budget is fixed via the constraint $\int_0^T \gamma(t) \, dt = \alpha T$.

LR Schedule Regimes

Three LR regimes are compared:

• Const: LR remains constant throughout MT
• FDR (Full Decay Range): LR decays throughout the entire continued training period
• MDR (Matching Decay Range): LR decay starts when MT data is introduced

---

Methodology

Experimental Setup

Dense Model: 630M parameters, follows Qwen2.5-0.5B architecture with untied embeddings. Pretrained on 200B DCLM tokens, then trained on 20.97B tokens during MT phase with MT token budgets of 10%, 25%, or 47.75%.

MoE Model: OLMoE design with 2.99B total parameters, 64 experts per layer, top-8 routing. PT phase: 40B tokens; MT phase: 4B or 10B tokens.

Data: DCLM-Baseline as PT data, FineMath-4plus subset as MT data.

Theoretical Framework

The paper models training as SGD on a block-diagonal quadratic loss. The covariance matrices for PT and MT data are:

$$\mathcal{H} = \mathrm{diag}(H_0, H_1, \mathbf{0}_{d_2 \times d_2}), \quad \mathcal{G} = \mathrm{diag}(H_0, \mathbf{0}_{d_1 \times d_1}, H_2)$$

where $H_0$, $H_1$, $H_2$ represent shared, MT-specific, and PT-specific feature blocks respectively.

The training dynamics follow an SDE approximation:

$$\mathrm{d}\boldsymbol{\theta}_t = -\eta(t)\nabla\mathcal{L}_t(\boldsymbol{\theta}_t)\mathrm{d}t + \eta(t)\sqrt{h\boldsymbol{\Sigma}}\mathrm{d}\boldsymbol{B}_t \tag{1}$$

The optimization problem is:

$$\min_{\gamma: [0,T] \to [0,1]} \mathbb{E}[\mathcal{L}(\boldsymbol{\theta}_T)] \quad \text{s.t.} \quad \frac{1}{T}\int_0^T \gamma(t) \, \mathrm{d}t \leq \alpha \tag{2}$$

---

Empirical Validation / Results

Ranking Reversal Phenomenon

Under FDR (LR decay): Replaying PT data improves MT loss. As $\gamma_{\mathrm{end}}$ increases (more MT data placed later), both PT and MT loss increase. Mixing PT and MT data gives lower final MT loss than pure MT.

Under Const LR: The advantage reverses. As $\gamma_{\mathrm{end}}$ increases, MT loss decreases while PT loss increases. Placing all MT data in a final pure MT stage gives the lowest MT loss.

Key finding: The same data schedule can be optimal or suboptimal depending solely on the LR schedule.

Optimal LR Decay Timing

Across all 19 sweeps, the best tested relative LR decay durations (percentage of the MT data window during which LR decays) are:

The ZL schedule requires approximately half the relative decay duration of TS because its MT data window is twice as long (to reach the same final MT ratio with the same budget).

Theoretical Results

Theorem 3.1 (Informal):

1. Const LR: If $\eta(t) \equiv \eta_{\mathrm{max}}$, the optimal data schedule is $\gamma^*(t) = \mathbb{1}_{\{t \geq (1-\alpha)T\}}$ (TS with $\gamma_{\mathrm{end}} = 1$).

2. Linear decay (FDR): Under certain conditions on decay strength $c$ and budget $\alpha$, the constant-LR optimum is no longer optimal.

3. Three-stage optimum under FDR: With power-law spectral assumptions, the optimal schedule is:

$$\gamma^*(t) = \begin{cases} 0, & 0 \leq t < t_1, \\ C\eta(t)^{-\frac{2(b-1)}{b+1}}, & t_1 \leq t < t_2, \\ 1, & t_2 \leq t \leq T, \end{cases}$$

where $b > 1$ is determined by the power-law spectral assumption.

Theorem 3.2 (Informal): For TS with $\gamma(t) = \mathbb{1}_{\{t \geq (1-\alpha)T\}}$ and stable-decay LR schedules, any optimal decay duration $W^*$ satisfies $0 < W^* < \alpha T$, meaning optimal LR decay begins strictly after MT entry.

Theoretical Interpretation

The expected final MT loss decomposes into:

$$\mathbb{E}[\mathcal{L}(\boldsymbol{\theta}_T)] = A + \underbrace{\frac{1}{2}\sum_{i=1}^{d_1}\lambda_i^{(1)}\theta_i^{(1)}(0)^2 e^{-2\lambda_i^{(1)}R(0)}}_{\text{bias term}} + \underbrace{\frac{h}{2}\sum_{i=1}^{d_1}\lambda_i^{(1)}\nu_i^{(1)}\int_0^T \eta(s)^2 e^{-2\lambda_i^{(1)}R(s)} \mathrm{d}s}_{\text{variance term}}$$

where $R(s) = \int_s^T \eta(t)\gamma(t) \, dt$ is the cumulative LR of MT updates after time $s$.

• Under constant LR: All schedules have the same bias term (fixed $R(0)$); placing MT data last maximizes $R(s)$ at every time, minimizing variance.
• Under LR decay: Moving MT data earlier increases $R(0)$ (reducing bias) but leaves less MT training to suppress later noise (increasing variance). The optimal schedule balances these effects.

MoE-Specific Findings

• Load balancing: TS schedules produce sharp spikes in load balancing loss at the data shift; ZL largely avoids this (Figure 6).
• Expert load rankings: ZL better preserves expert load rankings across the data shift; TS causes substantial reshuffling.

---

Theoretical and Practical Implications

Practical Implications

1. Data and LR schedules must be designed jointly — evaluating a data schedule in isolation is insufficient; the LR schedule determines which data schedule is optimal.

2. Delayed LR decay is beneficial: Keeping LR high after MT data enters, followed by a short final cooldown (10-40% of the MT window), outperforms both constant LR and decay matched to MT entry.

3. Gradual data mixing reduces MoE instability: ZL schedules avoid load-balancing spikes and preserve expert load rankings, making training more stable.

4. TS vs. ZL tradeoffs: TS exhibits larger PT-MT loss variation and requires more careful tuning; ZL offers a narrower range but may be preferred under constant LR for MoE models.

Theoretical Implications

1. The quadratic model provides a tractable framework for understanding data-LR coupling in optimization.

2. The bias-variance decomposition explains why the optimal data schedule is LR-dependent: constant LR favors late MT placement (variance minimization), while LR decay favors earlier placement (bias reduction).

3. The three-stage optimal schedule (PT → PT+MT mixture → MT) under FDR provides a theoretical foundation for practical data mixing strategies.

4. The power-law spectral assumption connects scaling law analyses to data scheduling theory.

---

Conclusion

This work demonstrates that data schedules and LR schedules should be designed together in LLM training pipelines. The key findings are:

1. Ranking reversal: The optimal data schedule depends critically on the LR schedule — PT replay is beneficial under LR decay but harmful under constant LR.

2. Joint optimization: LR decay should begin well after MT data enters training (10-40% of the MT window), balancing the benefits of high LR for MT updates against final noise reduction.

3. Theoretical grounding: A quadratic SDE model recovers both the ranking reversal and the optimal timing of LR decay, with a three-stage optimal data schedule under FDR.

Future Directions

1. Determining the form of jointly optimal schedules under fixed compute and MT token budgets remains theoretically challenging.

2. Future work should consider PT and MT losses jointly and incorporate downstream task metrics when evaluating and choosing schedules.

3. The interaction between continuous data schedules and MoE load balancing warrants further investigation.