Data Schedules and Learning Rate Schedules Should Be Designed Together

Haodong Wen; Kairong Luo; Tingkai Yan; Wenguang Chen; Kaifeng Lyu

cs.LGOct 21, 2026

Data Schedules and Learning Rate Schedules Should Be Designed Together

Authors: Haodong Wen, Kairong Luo, Tingkai Yan, Wenguang Chen, Kaifeng Lyu

Organizations: Tsinghua University · Peng Cheng Laboratory

Abstract

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

Panel of Figure 1 — 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

\gamma_{\mathrm{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.

Panel of Figure 8 — 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.

Panel of Figure 9 — 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.

Panel of Figure 10 — 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.

MT budget	Schedule	$\gamma_{\mathrm{end}}$	MT window (steps)	Additional relative LR decay durations (%)
Dense: 200B-token PT; 20.97B-token continuation; FineMath
10%	TS	100%	1,000	20, 40, 50, 80, 120
	TS	80%	1,250	20, 40, 80, 90, 110, 120
	ZL	100%	2,000	5, 10, 20, 40, 80, 90, 110, 120
25%	TS	100%	2,500	20, 40, 80, 120
	ZL	100%	5,000	5, 10, 20, 40, 80

Appendix

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.

Dense: 200B-token PT; 20.97B-token continuation; FineMath
MT budget	Schedule, $\gamma_{\mathrm{end}}$	Best relative LR decay duration	Best loss	Const	MDR	Gain over MDR
10%	TS, 100%	40%	1.5156	1.5420	1.5394	0.0238
	TS, 80%	40%	1.5215	1.5529	1.5454	0.0239
	ZL, 100%	20%	1.5203	1.5485	1.5682	0.0479
25%	TS, 100%	40%	1.4275	1.4717	1.4475	0.0200
	ZL, 100%	20%	1.4392	1.4810	1.4802	0.0410

Appendix

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.

Dense: 200B-token PT; 20.97B-token continuation
MT budget	$\gamma_{\mathrm{end}}$	Relative LR decay duration or rule	Primary	Repeat 1	Repeat 2	$s_{3}$
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

s_{3}=\sqrt{\frac{1}{3}\sum_{j=1}^{3}(\ell_{j}-\bar{\ell})^{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

\gamma_{\mathrm{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

\gamma_{\mathrm{end}}

in percent. Both MDR minima after 4B continued training tokens occur at

\gamma_{\mathrm{end}}=50\%

, so MDR does not select

\gamma_{\mathrm{end}}=100\%

at both continued training token budgets.

Continuation tokens and LR rule	$\gamma_{\mathrm{end}}=100\%$	$\gamma_{\mathrm{end}}=75\%$	$\gamma_{\mathrm{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.

Panel of Figure 17 — 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%.

Panel of Figure 18 — 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

\pm s_{3}

, computed from unrounded losses relative to the primary Const loss; they do not estimate uncertainty at other durations or in the Const baseline.

Panel of Figure 20 — 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_{\ell}

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.

Panel of Figure 29 — 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.

Panel of Figure 30 — 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

\alpha=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

\alpha=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-\alpha)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.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 344 papers

May 18, 2026

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

Schedule-Free Learning has shown promise as a practical anytime training method for machine learning, showing success across dozens of standard benchmark problems. However, strong performance for LLM training has only been demonstrated at small scales. We identify a number of fixes necessary to scale up Schedule-Free Learning to larger batch sizes and model sizes, and present a learning-rate-free and schedule-free method (ScheduleFree+) for training large language models which greatly outperforms Warmup-Stable-Decay (WSD) schedules. We also demonstrate that Schedule-Free Learning is most effective for long duration training, and at 1000 tokens per parameter, it outperforms SOTA schedules by 31%. Schedule-Free Learning provides a theoretical foundation for the use of model averaging and checkpoint merging during pretraining.

Aaron Defazio

FAIR at Meta Super-Intelligence Labs

Language Model Pretraining Learning Rate Scheduling LLM Training

Ask PDF

May 25, 2026

How Should LLMs Consume High-Quality Data? Optimal Data Scheduling via Quality-Aware Functional Scaling Laws

High-quality data is scarce in large language model (LLM) training, yet how to schedule its use jointly with training dynamics lacks theoretical guidance. We extend functional scaling laws by incorporating a data-quality dimension, and solve the joint data-quality and batch-size scheduling problem in asymptotic closed form. The solution reveals two regimes and a dual role of high-quality data. In the noise-limited regime, high-quality data should be used as a signal amplifier: lowering the batch size converts cleaner data into more signal without amplifying noise. In the signal-limited regime, it should be used as a noise suppressor: late placement reduces terminal noise without sacrificing signal accumulation. Existing curriculum-style pipelines primarily exploit the second role by placing cleaner data late, but miss the first role because conventional decay schedules reduce update intensity exactly when high-quality data becomes available. Guided by this, we propose Drop-Stable-Rampup for LLM midtraining: upon the quality transition, drop the batch size, hold it stable to accumulate signal, then ramp up to suppress terminal noise. On a 15B Mixture-of-Experts model midtrained on 108B tokens, Drop-Stable-Rampup improves average accuracy over Warmup-Stable-Decay (WSD) by +1.70 and over Cosine-decay by +2.98, with particularly large gains on mathematical reasoning benchmarks such as GSM8K (+4.23) and MATH (+2.80).

Zhitao Zhu, Xili Wang, Shizhe Wu +2

Peking University · Meituan

Language Model Scaling Laws LLM Training

Ask PDF

Sep 30, 2026

From Spectra to Joint Schedules in LLM Pre-training: 3+3(+2) Scaling-Law Regimes

Power-law learning curves are often treated as fixed properties of a model and its data, although learning-rate and batch-size schedules can change the observed loss. We study this dependence in noisy online SGD with linear random features. Conditional on the representation, an exact Volterra equation separates two response components: a forcing term that propagates unresolved target error and a memory kernel that propagates stochastic-error injections. We prove that either component follows a power law if and only if its cumulative weighted spectral mass has the corresponding low-spectrum scaling; individual eigenvalues and target coefficients need not obey coordinatewise power laws. Under a joint schedule, intrinsic time

T_t=\sum_{s<t}η_s

controls optimization progress, while

r_t=B_t/η_t

controls noise injection. Their interaction yields sharp conditions under which a schedule preserves, changes, or destroys the clean power law, together with a memory ceiling on noise reduction. The power-law random-feature model realizes this mechanism in

3+3(+2)

propagation regimes with phase-dependent compute rates. Controlled nanoGPT experiments show that (1) learning-rate and batch-size schedules with matched

B/η

paths are nearly equivalent in intrinsic time, (2) a forcing-memory surrogate accurately predicts loss across schedules, and (3) its fitted exponents across real-world datasets identify the regime of LLMs in

3+3(+2)

map. The source code for reproducing our experiments is available on https://github.com/yichenblue/spectra-to-schedules-in-pretraining.

Yichen Wang, Fanghui Liu, Yudong Chen

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.

Language Model Pretraining Deep Learning Optimization Learning Rate Scheduling Language Model Scaling Laws Stochastic Gradient Descent Learning Curves

Ask PDF

Abstract

LLM Training

Explore similar work

ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models

How Should LLMs Consume High-Quality Data? Optimal Data Scheduling via Quality-Aware Functional Scaling Laws

From Spectra to Joint Schedules in LLM Pre-training: 3+3(+2) Scaling-Law Regimes