Mid-training (MT) is a training phase in LLM training pipelines that continues pretraining (PT) on smaller datasets of higher value. Taking the PT and MT phases as a whole, training usually starts with only PT data and then transitions completely to MT data. In this work, we view this complete transition from PT to MT data as a special case of a broader class of data mixing schedules that specify how the proportions of PT and MT tokens in each batch change over training. Under fixed total training compute and MT token budgets, we ask: Is this complete transition to MT data optimal for minimizing final MT loss? We show that the answer depends on the learning rate (LR) schedule: one data schedule may outperform another under one LR schedule, but changing only the LR schedule can reverse the outcome. Specifically, when the LR decays throughout continued training, introducing MT data earlier and mixing it with PT data lowers the final MT loss compared with placing all MT data at the end of training, but this result is reversed if we use constant LR. We therefore study how to jointly choose when to introduce MT data and when to decay the LR through ablation studies, and find that it is better to decay the LR only after MT data enters training. Finally, we develop a theoretical model that explains why the optimal data schedule for a fixed MT token budget depends on the LR schedule.
Figures & tables
Figure 1: Replaying PT data can be both beneficial and harmful under different LR schedules. The compute and MT data budgets are fixed (the shaded areas in the left panels show the MT budgets). With F ull LR D ecay R ange throughout continued training (FDR, top), mixing PT and MT data in the second TS stage gives lower final MT loss than using MT data alone. With constant LR (Constant LR, bottom), the pure MT stage instead gives the lowest tested MT loss. Results use the dense model at a 10% MT budget. Training-step coordinates in the schedule panels are schematic: the PT checkpoint is drawn at 100k steps, rather than its actual step 95,368.
Figure 2: Across all three MT budgets, replaying PT data improves MT loss under FDR, but not under Const or MDR. The three panels show final MT loss as the final MT ratio in the second TS stage varies. Under the FDR regime, the lowest tested loss occurs when this stage mixes PT and MT data. Under the Const and MDR regimes, it occurs with pure MT data (100%). Large circles mark these minima, while other hollow markers show repeated runs. See Section 2.2 and Appendix C for the setup and full results.
Figure 3: LR should decay well after MT enters training for TS and ZL data schedules. The bottom row shows how final MT loss changes with relative LR decay duration for TS and ZL across three MT budgets. ZL reaches its lowest MT loss at 20%, and TS at 40%, improving on both Const (0%) and MDR (100%). The top row illustrates TS (left) and ZL (right) at a 10% MT budget, paired with LR schedules of different relative LR decay durations , measured as percentages of the time from MT entry to training completion. See Section 2.3 and Appendix D for the setup and further discussion.
Figure 4: The optimal data schedule can have three stages under LR decay. The first two columns show an abrupt switch from pure PT to pure MT. With stronger LR decay in the last two columns, the teal curves instead include a mixing stage in which the MT ratio increases continuously. This change occurs at both MT budgets shown. See Theorem 3.1 for the theoretical result and Appendix M for the numerical setup.
Figure 5: LR-dependent data-schedule preferences also occur in MoE. Left to right: 4B continuation tokens with a 25% MT budget; 10B tokens with 25% MT; and 10B tokens with 47.75% MT. All three settings use the same 2.99B-parameter model. At a final MT ratio of 100%, Const can outperform MDR. See Section 2.2 and Appendix F for the setup and full results.
Figure 6: Gradual MT entry reduces router imbalance under data distribution shift. TS produces sharp spikes when the MT ratio jumps, while ZL avoids comparable spikes. Under Const (left), both schedules use the same LR trajectory, so this difference comes from the data schedule. MDR (right) shows a similar pattern, although LR decay starts at different times for TS and ZL. These MoE runs use 10B continuation tokens and a fixed 25% MT budget. Dotted lines mark MT entry. See Appendix G for setup and diagnostic details.
Appendix figures & tables41 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Architecture
Reference architecture
Qwen2.5-0.5B
Total parameters
630M
Transformer layers
24
Hidden size
896
Attention heads (query / KV)
14 / 2
Appendix
Table 1: Architecture and training settings for the 630M dense model. The final LR applies to schedules with decay; Const retains the peak LR.
Hyperparameter
Value
Architecture
Reference architecture
OLMoE
Total parameters
2.99B
Active non-embedding parameters
454M
Input/output embedding parameters
155M combined
Transformer layers
12
Appendix
Table 2: Architecture and training settings for the 2.99B-parameter MoE model. Unlike the original OLMoE recipe, PT retains a constant LR after warmup, and normalization weights have no weight decay. The final LR applies only to schedules with decay.
Figure 7: Illustration of different LR schedules and data schedules. Rows show TS and ZL, and columns show Const, FDR, and MDR. Solid lines and shading show the MT ratio on the left axes. Dash dotted lines show LR relative to its peak on the right axes. Both data schedules use a 25% MT budget and end at 100% MT, so ZL introduces MT earlier. MDR therefore begins LR decay earlier for ZL than for TS. Time starts at the shared PT checkpoint.
MT budget
TS final MT ratios (%)
ZL final MT ratios (%)
10%
10, 20, 40, 80, 100
20, 40, 60, 80, 100
25%
25, 40, 50, 70, 75, 85, 100
50, 60, 70, 80, 90, 100
47.75%
48, 50, 60, 70, 80, 90, 95, 100
95.5, 97, 98, 99, 100
Appendix
Table 3: Complete data schedule runs for the dense model. Total and MT token budgets are fixed within each group.
MT budget
Data schedule
Const
FDR
MDR
10%
TS
1.5420 (100)
1.6108 (20)
1.5394 (100)
ZL
1.5485 (100)
1.6144 (20)
1.5682 (100)
25%
TS
1.4717 (100)
1.5024 (40)
1.4475 (100)
ZL
1.4810 (100)
1.5064 (50)
1.4802 (100)
47.75%
TS
1.4217 (100)
1.4194 (70)
1.3826 (100)
ZL
1.4335 (100)
1.4246 (97)
1.4221 (100)
Appendix
Table 4: Lowest final MT loss among the tested final MT ratios for each MT budget, data schedule, and LR schedule. Parentheses give the final MT ratio γend in percent.
MT budget
Const
FDR
MDR
10%
+0.0065
−0.0185
+0.0288
25%
+0.0093
−0.0138
+0.0327
47.75%
+0.0118
−0.0006
+0.0395
Appendix
Table 5: MT loss difference in the dense model, ZL minus TS, with both schedules ending at 100% MT. Negative values favor ZL. Differences use the reported final losses and are not significance estimates.
Figure 8: FDR favors earlier MT mixing, while Const and MDR favor ending with MT data alone. For both TS and ZL, the lowest tested MT loss occurs at a final MT ratio of 20% under FDR and 100% under Const and MDR. These dense model runs use a 10% MT budget over 20.97B continued training tokens, including 2.10B FineMath tokens.
Figure 9: Changing the LR schedule changes the preferred final MT ratio. FDR gives the lowest tested MT loss at 40% MT for TS and 50% for ZL; Const and MDR favor 100% for both. These dense model runs use a 25% MT budget over 20.97B continued training tokens, including 5.24B FineMath tokens.
Figure 10: ZL losses under FDR change little across its tested final MT ratios. They vary by only 0.001, while Const and MDR favor a final MT ratio of 100%. TS covers eight final MT ratios; ZL covers five from 95.5% to 100%. These dense model runs use a 47.75% MT budget over 20.97B continued training tokens, including 10.01B FineMath tokens. Hollow overlays show selected data sampling repeats.
Table 6: LR decay scan settings. Relative LR decay durations are percentages of each schedule’s MT data window. Const (0%) and MDR (100%) are tested in every scan; FDR is also tested except for StarCoderData ZL.
Table 7: Final MT losses in all 19 scans. Best loss and its relative LR decay duration refer to the lowest-loss primary run among the tested schedules. Gain over MDR is MDR minus best loss, calculated from the displayed values. Group headings specify the model, PT budget, continuation budget, and MT dataset.
MoE: 10B continuation tokens 47.75% MT
MoE: 4B continuation tokens 25% MT, TS
Relative LR decay duration or rule
TS
ZL
PT 20B
PT 30B
Const (0%)
1.4292
1.4424
1.6666
1.6409
5%
1.4091
1.4131
—
—
10%
1.4046
1.4110
—
—
20%
1.4007
1.4137
1.6562
1.6260
40%
1.4040
1.4268
1.6578
1.6303
Appendix
Table 8: Final MT losses at every tested relative LR decay duration for the four additional scans, including the reference LR schedules. All data schedules end at 100% MT. The 10B-token continuations start from the 40B-token PT checkpoint, while the last two columns use earlier checkpoints from the same PT trajectory. Bold entries mark the lowest tested loss in each column. Dashes indicate durations that were not tested.
Figure 11: A short final LR decay improves on Const and MDR. The lowest tested MT losses are 1.5215 at a relative LR decay duration of 40% for TS (top) and 1.5203 at 20% for ZL (bottom). These dense model runs use a 10% MT budget. TS ends at 80% MT with a 1,250-step MT window; ZL ends at 100% MT with a 2,000-step window. Hollow TS markers show repeats at 20%; the 40% minimum was not repeated.
Figure 12: TS and ZL favor different relative LR decay durations. TS (top) reaches its lowest tested loss of 1.4275 at 40%, and ZL (bottom) reaches 1.4392 at 20%. Nearby short decay durations give similar losses. Both dense model runs use a 25% MT budget and end at 100% MT, with MT windows of 2,500 steps for TS and 5,000 for ZL. These configurations have no repeated runs.
Figure 13: The best tested relative LR decay duration remains 40% when TS ends at lower MT ratios. At a final MT ratio of 80% (top), this gives a loss of 1.3831, with repeats at 1.3834 and 1.3846. At 60% MT (bottom), it gives 1.4039, with similar losses at relative LR decay durations between 20% and 80%. These dense model runs use a 47.75% MT budget and MT windows of 5,969 and 7,958 steps, respectively. The scan ending at 100% MT appears in Figure 3 .
Figure 14: TS favors a relative LR decay duration of 40%, while ZL favors 20%. Their lowest tested losses are 1.3690 for TS (top) and 1.3857 for ZL (bottom). For ZL, MDR and FDR differ by only 450 decay steps and yield similar losses of 1.4221 and 1.4249, respectively. Both dense model runs use a 47.75% MT budget and end at 100% MT, with MT windows of 4,775 steps for TS and 9,550 for ZL. This ZL configuration has no repeated runs.
Figure 15: The best tested TS and ZL decays span nearly the same number of training steps. TS (top) reaches 1.6080 at a relative LR decay duration of 20% (47 steps), while ZL (bottom) reaches 1.6000 at 10% (48 steps). These MoE runs use 4B continued training tokens, a 25% MT budget, and a final MT ratio of 100%. Their MT windows last 238 and 477 steps, respectively. Hollow TS markers repeat MDR; neither minimum was repeated.
Figure 16: The best tested TS and ZL decays both last 119 steps. TS (top) reaches 1.4853 at a relative LR decay duration of 20%, while ZL (bottom) reaches 1.4872 at 10%. These MoE runs use 10B continued training tokens, a 25% MT budget, and a final MT ratio of 100%. Their MT windows last 596 and 1,192 steps, respectively. Hollow TS markers repeat MDR; the best TS and ZL configurations have no matched data repetitions.
MT budget
γend
Relative LR decay duration or rule
Primary
Repeat 1
Repeat 2
s3
Dense: 200B-token PT; 20.97B-token continuation
10%
80%
20%
1.5227
1.5213
1.5212
0.000706
47.75%
100%
MDR
1.3826
1.3828
1.3830
0.000149
47.75%
70%
FDR
1.4194
1.4217
1.4198
0.001008
47.75%
80%
40%
1.3831
1.3834
1.3846
0.000620
MoE: 40B-token PT; 4B-token continuation
Appendix
Table 9: Final MT losses for the six repeated TS configurations, all using FineMath. The descriptive dispersion s3=31∑j=13(ℓj−ℓˉ)2 is computed from the original losses before rounding; its denominator is three, rather than the denominator-two sample standard deviation.
Data schedule
Final MT ratio
First MT step 4B-token continuation
First MT step 10B-token continuation
TS
100%
716
1,788
75%
636
1,589
50%
477
1,192
40%
358
894
25%
0
0
ZL
100%
477
1,192
Appendix
Table 10: Tested data schedules for the MoE model at a 25% MT budget. The first MT step denotes the TS switch or ZL ramp start. The TS γend=100% data windows are 238 and 596 steps.
Continuation tokens
Data schedule
Const
FDR
MDR
4B
TS
1.6176 (100)
1.6349 (25)
1.6343 (50)
ZL
1.6131 (100)
1.6548 (50)
1.6548 (50)
10B
TS
1.5110 (100)
1.5343 (25)
1.5225 (100)
ZL
1.5101 (100)
1.5548 (50)
1.5508 (100)
Appendix
Table 11: Lowest final MT loss among the tested final MT ratios for each training duration, data schedule, and LR schedule in the MoE model at a 25% MT budget. Parentheses give the final MT ratio γend in percent. Both MDR minima after 4B continued training tokens occur at γend=50% , so MDR does not select γend=100% at both continued training token budgets.
Continuation tokens and LR rule
γend=100%
γend=75%
γend=50%
4B Const
−0.0045
−0.0029
−0.0028
4B FDR
−0.0298
−0.0282
−0.0233
4B MDR
+0.0101
+0.0155
+0.0206
10B Const
−0.0010
+0.0010
+0.0015
10B FDR
−0.0264
−0.0251
−0.0183
10B MDR
+0.0283
+0.0282
+0.0303
Appendix
Table 12: MoE MT loss difference, ZL minus TS, at matched final MT ratios. Differences are computed from unrounded losses and reported to four decimal places. Negative values favor ZL. Each comparison uses one primary run per schedule, without significance testing.
Figure 17: Const gives the lowest tested MT loss among the three baseline LR schedules. MDR favors a final MT ratio of 50% for both TS and ZL. These MoE runs use 4B continued training tokens with a 25% MT budget. Hollow overlays repeat TS+MDR at a final MT ratio of 100%.
Figure 18: Longer MoE continued training narrows the TS loss gap between MDR and Const. At a final MT ratio of 100%, TS+MDR is closer to TS+Const than after 4B tokens, while ZL+MDR remains worse than ZL+Const. These runs use 10B continued training tokens with a 25% MT budget. Hollow overlays repeat TS+MDR at 100% MT.
Figure 19: Longer MoE continued training narrows the MDR–Const loss gap. For TS, extending training from 4B to 10B tokens reduces the primary-run gap from 0.0310 to 0.0115 at a final MT ratio of 100%, while the best tested relative LR decay duration remains 20%. The upper panels show final MT loss and its change from each horizon’s Const baseline; the lower panels vary the final MT ratio. All runs use a 25% MT budget. In the difference panel, circles and local error bars at MDR show the three-run mean and ±s3 , computed from unrounded losses relative to the primary Const loss; they do not estimate uncertainty at other durations or in the Const baseline.
Figure 20: Const and FDR favor opposite TS endpoints at a larger MoE MT budget. Const favors a final MT ratio of 100%, while FDR favors 47.75%. At 100% MT, ZL improves on TS under FDR but has higher MT loss under Const and MDR. These runs use 10B continued training tokens with a 47.75% MT budget, starting from the 40B-token PT checkpoint. ZL is tested only at 100% MT. Each point is one run.
Figure 21: Increasing the MoE MT budget preserves the best tested relative LR decay durations. TS (top) favors 20%, and ZL (bottom) favors 10%; both correspond to 228 decay steps. These runs use 10B continued training tokens with a 47.75% MT budget and end at 100% MT. Their MT windows last 1,138 and 2,277 steps, respectively. These configurations have no data stream repeats.
Figure 22: Const and FDR favor opposite TS endpoints even from earlier PT checkpoints. After either 20B PT tokens (left) or 30B (right), Const gives its lowest tested MT loss when TS ends at 100% MT, while FDR favors 25%. Both use 4B continued training tokens and a 25% MT budget. The checkpoints come from the same PT run.
Figure 23: Delaying LR decay improves on Const and MDR from both earlier PT checkpoints. The lowest tested MT losses are 1.6562 after 20B PT tokens (top) and 1.6260 after 30B (bottom), both at a relative LR decay duration of 20% (47 steps). Each MoE TS run uses 4B continued training tokens, a 25% MT budget, and a final MT ratio of 100%, giving a 238-step MT window. Only 20% and 40% are tested beyond Const, MDR, and FDR.
Figure 24: Changing the PT checkpoint preserves the preferred data and LR schedules among the shared configurations. Const and FDR favor opposite TS endpoints (left). A relative LR decay duration of 20% gives the lowest tested loss, while MDR remains worse than Const (right). The right panel shows loss changes relative to each checkpoint’s own Const baseline at a final MT ratio of 100%. These MoE runs start from 20B, 30B, and 40B checkpoints of the same PT trajectory and use 4B continued training tokens with a 25% MT budget. The 40B checkpoint has additional sampled durations.
Figure 25: Changing MT data from mathematics to code preserves the best tested relative LR decay duration for both TS (left) and ZL (right). Both datasets use the same MoE model and 10B-token continuation with a 25% MT budget and final MT ratio 100%. The horizontal axis gives decay duration as a percentage of each schedule’s MT data window, with zero denoting Const and 100% denoting MDR. The vertical axis gives final MT loss minus the corresponding Const loss for that dataset and data schedule. Points show measured values, and rings mark the lowest tested losses at 20% for TS and 10% for ZL. Lines connect tested values.
Figure 26: Gradual MT entry reduces load balancing spikes. Under Const (left), TS and ZL use the same LR at every step, allowing us to compare the data schedules without also changing the LR. TS produces a sharp spike when the MT ratio suddenly increases, while ZL avoids a comparable spike by introducing MT data gradually. MDR (right) shows a similar pattern, although LR decay begins at different times for TS and ZL. Rows show MoE experiments with 4B (top) and 10B (bottom) tokens of continued training. Curves compare final MT ratios of 100%, 75%, and 50%. Dotted lines mark when MT data first enters. The 10B results also appear in Figure 6 .
Figure 27: TS produces a transient load balancing spike, while gradual MT entry avoids a comparable peak. The upper panels compare TS at several final MT ratios, including Const and MDR at 100% MT; the lower panels compare TS and ZL under MDR at final MT ratios of 100%, 75%, and 50%. These MoE runs use 4B continued training tokens. Training loss (left) and the logged load balancing statistic (right) are shown every 20 training steps.
Figure 28: Smaller TS data shifts produce smaller load balancing spikes, while ZL avoids a comparable peak. The upper panels compare TS switches at steps 1,788, 1,589, and 1,192 for final MT ratios of 100%, 75%, and 50%; the lower panels compare TS and ZL under MDR at the same endpoints. These MoE runs use 10B continued training tokens. Training step is horizontal, with training loss on the left and logged load balancing on the right. Final validation losses are reported separately in Tables 11 and 12 .
Checkpoint (step)
DCLM imbalance
FineMath imbalance
1,750 (before switch)
1.005
1.216
2,250 (462 steps after switch)
1.029
1.006
2,384 (final)
1.043
1.015
Appendix
Table 13: Expert load imbalance on DCLM and FineMath for the MoE TS+Const run with 10B continuation tokens and a final MT ratio of 100%. Each value is the average of Iℓ across the 12 layers, as defined in Equation 3 . The first checkpoint is 38 steps before the data switch. These measurements evaluate each domain separately on validation data, whereas the training logs summarize routing on the changing mixture of PT and MT data.
Figure 29: The router initially assigns MT tokens unevenly across experts and becomes more balanced during MT training. Checkpoint routing diagnostics for the MoE TS+Const run with 10B continuation tokens and a final MT ratio of 100%. Top: imbalance by layer on DCLM and FineMath at steps 1,750, 2,250, and 2,384. Bottom: FineMath expert shares before the switch, at training end, and their difference, with expert and layer indices as coordinates. Shares are relative to uniform routing. These diagnostics describe one run, without replica averaging.
Schedule
Final imbalance FineMath / DCLM
Net routing change FineMath TV
Regression residual std.
Rank retention Spearman
TS+Const
1.015 / 1.043
0.156
0.119
0.30
ZL+MDR
1.008 / 1.039
0.140
0.042
0.74
Appendix
Table 14: MoE routing comparison over 10B continuation tokens with a 25% MT budget and final MT ratio 100%. TS+Const compares steps 1,750 and 2,384, and ZL+MDR compares steps 1,150 and 2,384, from before each schedule’s MT entry time to training completion. TV reports the net change in FineMath assignment shares. Regression residuals use shares relative to uniform routing, and Spearman correlations are averaged across layers. Figure 6 provides the separate control with identical LR trajectories.
Figure 30: ZL+MDR preserves expert load rankings better than TS+Const, although both reach nearly balanced MT loads. (a) Both schedules bring FineMath (MT) loads close to uniform in most layers; one denotes equal assignment shares across experts. (b) Initially more heavily used experts lose share, while less used experts gain it. Each point represents one expert in one layer. The horizontal axis shows its share before MT entry time minus the uniform share, and the vertical axis shows its final minus initial share, both in units of 1/64 . Points on the dotted line end at the uniform share. (c) Despite similar final balance, the ordering of experts changes more under TS+Const. The higher rank correlations under ZL+MDR mean that experts receiving relatively more FineMath assignments initially tend to remain more heavily used at the end. All checkpoint measurements use held-out data. These 10B-token MoE continuations use different LR schedules, so the comparison does not isolate the effect of the data schedule. See Table 14 for the setup.
Figure 31: Consistency of the SDE calculation, the exact discrete SGD second moment recursion, and Monte Carlo simulation for α=0.25 . Results are shown for a constant learning rate (left) and linear decay with c=0.8 (right), across four representative data schedules. Error bars give two standard errors of the Monte Carlo mean.
Figure 32: Numerical validation of Theorem I.3 . Solid curves show the schedules predicted by the closed form solution, while markers show the optima obtained by directly optimizing the expected validation loss using different numerical methods. The two solutions agree for both MT data budgets considered.
Figure 33: Numerical illustration of the KKT mechanism. Left: At the two stage schedule with α=0.25 , the marginal value of MT data is increasing under a constant learning rate, whereas linear decay with c=0.8 violates the switching condition around (1−α)T . Right: At the numerical optimum under linear decay, the normalized marginal value equals the KKT multiplier throughout the interior mixing phase and satisfies the appropriate inequalities on the two boundary phases.
Department of Computer Sciences, University of Wisconsin–Madison, USA. · School of Mathematical Sciences, Institute of Natural Sciences and MOE-LSC, Shanghai Jiao Tong University, China.