Source-linked AI summary
Data Schedules and Learning Rate Schedules Should Be Designed Together
Haodong Wen, Kairong Luo, Tingkai Yan, Wenguang Chen, Kaifeng Lyu
TL;DR
The paper asks whether the usual all-late transition from PT to MT minimizes final MT loss under fixed compute and MT-token budgets. It compares data schedules under different LR schedules and studies decay timing, finding that schedule rankings can reverse and that delayed LR decay often performs best. A quadratic model explains how the preferred data schedule can depend on the LR schedule.
Problem
With compute and MT-token budgets fixed, it is unclear whether reserving all MT data for the end minimizes final MT loss.
Source: §1 · p. 1Method
The paper compares data schedules under different LR schedules, varies LR-decay timing, and analyzes schedule choice with a quadratic SGD model.
Source: §1 · p. 3Results
The ranking of data schedules can shift or reverse across LR schedules; across 19 sweeps, the best tested data–LR pairs beat constant LR and decay beginning at MT entry.
Source: §1 · p. 3Takeaways & Limitations
Data and LR schedules should be chosen together rather than assessing a data schedule independently of the LR schedule.
Source: §1 · p. 3The theoretical analysis focuses on a spectral regime with b0 ≥ b1 and MT data capped at an α fraction of total training data.
Source: §3.1 · p. 7, §3.3 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper asks whether the usual complete transition from PT to MT is optimal when compute and MT-token budgets are fixed. It shows that data-schedule rankings depend on the LR schedule, then studies joint scheduling experimentally and theoretically.
1 INTRODUCTION
- The paper treats the usual PT-to-MT transition as one data schedule and asks whether it minimizes final MT loss under fixed training compute and MT-token budgets.The broader schedules specify the fraction of MT tokens in each batch over training.Source: §1 · p. 1
- Under fixed compute and MT-token budgets, replaying PT while introducing MT earlier lowers final MT loss with full-range LR decay, but increases it with constant LR.With constant LR, the all-MT final stage gives the lowest tested loss.Source: §1 · p. 2
- The ranking of data schedules by final MT loss can shift or reverse when only the LR schedule changes.The reversal persists across the paper’s experimental setups.Source: §1 · p. 3
- The authors vary when LR decay begins relative to MT entry to explore joint data and LR schedules that minimize final MT loss.They call the time MT data first enters training the MT entry time.Source: §1 · p. 3
- A quadratic SGD model finds that all-late MT is optimal under constant LR but can lose optimality under sufficiently strong linear decay.Under additional spectral assumptions and sufficiently long training, the unique optimum starts with PT, passes through a mixture, and ends with MT.Source: §1 · p. 3
2 DATA SCHEDULES AND LEARNING RATE SCHEDULES SHOULD BE DESIGNED TOGETHER
Experiments across dense and MoE models show that data-schedule rankings depend on the LR regime, and that LR decay should generally begin after MT enters training. The section compares schedule families and regimes, then sweeps decay timing for fixed data schedules.
2 DATA SCHEDULES AND LEARNING RATE SCHEDULES SHOULD BE DESIGNED TOGETHER
- Across dense and MoE experiments, earlier MT mixing lowers final MT loss under LR decay, whereas constant LR favors placing all MT data in a final pure-MT stage.The section reports over 250 runs; across LR-decay sweeps, the lowest final MT loss comes from keeping LR high after MT entry and decaying near training’s end.Source: §2 · p. 3, §2.2 · pp. 4–5
2.1 EXPERIMENTAL SETUP
- The experiments use dense and MoE models, three dense MT budgets, and Const, FDR, and MDR LR regimes to compare data schedules under controlled MT-token budgets.TS switches from PT-only data to a fixed MT ratio; ZL increases the MT ratio linearly after an initial PT-only period.Source: §2.1 · pp. 3–4
2.2 DIFFERENT LR SCHEDULES PREFER DIFFERENT DATA SCHEDULES
- Under FDR, replaying PT improves MT loss, while under Const and MDR the lowest tested loss occurs with pure MT in the second TS stage.The figure reports this pattern across all three MT budgets; MDR’s comparison with Const differs between dense and MoE settings.Source: Figure 2 · p. 4, §2.2 · p. 5
2.3 LEARNING RATE DECAY SHOULD BEGIN WELL AFTER MT DATA ENTERS TRAINING
- The LR-duration sweep holds each data schedule fixed and varies when LR begins its linear decrease to the same final value.Relative duration is measured against the interval from MT entry to training completion; Const is 0% and MDR is 100%.Source: §2.3 · pp. 5–6
- Across 19 decay-timing scans, the best tested relative LR decay durations fall between 10% and 40% and beat both Const and MDR for dense and MoE models.For dense models, the best tested durations are 40% for TS and 20% for ZL; ZL’s best duration is about half TS’s across the reported experiments.Source: Figure 3 · p. 5, §2.3 · p. 6
- The preferred decay timing is stable across reported MoE training lengths and earlier PT checkpoints, and persists when FineMath is replaced by StarCoderData.The reported best tested durations remain 10% for ZL and 20% for TS in the MoE sweeps.Source: §2.3 · p. 6
3 THEORETICAL ANALYSIS
The theoretical analysis shows that the optimal allocation of a fixed MT data budget depends on the LR schedule. It derives when an abrupt switch is optimal, when a PT–MT mixture is preferred, and why optimal LR decay begins after MT entry under stated conditions.
3.1 SETUP
- The analysis models SGD on a block-diagonal quadratic loss with a continuous-time SDE approximation, then minimizes expected final MT loss under a fixed MT budget.The model represents stochastic gradients as the mean gradient plus independent zero-mean noise and constrains MT data to at most a fixed fraction of training.Source: §3 · p. 6, §3.1 · p. 7
3.2 OPTIMAL DATA SCHEDULES DEPEND ON THE LR SCHEDULE
- With constant LR, the optimal schedule uses PT first and MT only at the end; under sufficiently strong linear decay, the optimum instead has PT-only, mixed, and MT-only stages.The three-stage result is derived under an additional power-law spectral assumption and sufficiently long training.Source: §3.2 · p. 8
- With constant LR, a final pure-MT stage minimizes the variance term, whereas under decay, moving some MT data earlier can lower final loss by reducing bias and noise contributions.The earlier allocation also leaves less subsequent MT training to suppress noise introduced between the shifted update times.Source: §3.2 · pp. 8–9
3.3 LR SHOULD DECAY AFTER MT DATA ENTERS TRAINING
- Under suitable power-law spectral conditions, for sufficiently small fixed MT budgets and sufficiently long training, optimal LR decay begins strictly after MT entry and ends before training ends.The analysis attributes this timing to keeping LR high for initial MT updates, then decaying it to reduce final variance.Source: §3.3 · pp. 9–10
4 ADDITIONAL EXPERIMENTS
The additional experiments show that joint schedule choices shape the PT–MT loss trade-off, while gradual MT entry can reduce MoE routing imbalance during distribution shifts. Under abrupt entry, expert-load rankings also change more substantially.
4.1 JOINT SCHEDULING ALSO SHAPES THE PT–MT LOSS FRONTIER
- ZL usually yields a narrower range of PT and MT losses across ending MT ratios, while TS varies more, especially in PT loss, making TS require more careful tuning.For dense models, MDR mostly dominates Const and FDR; under MDR, TS dominates ZL on both losses, while ZL has a MoE advantage over TS under constant LR.Source: §4.1 · p. 10
4.2 CONTINUOUS DATA SCHEDULES REDUCE LOAD IMBALANCE CAUSED BY DISTRIBUTION SHIFTS
- Under a 10B-token MoE run with a fixed 25% MT budget, TS produces sharp load-balancing-loss spikes at MT entry, while gradual ZL largely avoids them.Under Const, both schedules share the same LR trajectory, isolating the difference to the data schedule; under MDR, a similar pattern appears despite different decay timing.Source: §4.2 · p. 10, Figure 6 · p. 11
- Under ZL, experts with higher loads before MT entry tend to remain more heavily used afterward; TS instead reshuffles expert-load rankings across the shift.The authors note that this substantial reshuffling may make subsequent training less stable.Source: §4.2 · p. 11
5 CONCLUSION
The paper concludes that data and LR schedules should be designed together, drawing on experiments in dense and MoE models and a quadratic theoretical model.
5 CONCLUSION
- The authors combine dense and MoE experiments with a quadratic theoretical model to show that data and LR schedules should be designed together.They identify determining jointly optimal schedules under fixed compute and MT-token budgets, and evaluating PT and MT losses and downstream metrics, as future work.Source: §5 · p. 11
AI USE STATEMENT
The paper reports using generative AI tools for analysis and manuscript-related tasks, with the authors reviewing the assisted work. Its related-work discussion distinguishes prior mixture-ratio and scheduling research from its focus on temporal allocation of a fixed MT budget across LR schedules.
AI USE STATEMENT
- The authors report using generative AI for checking mathematical claims, analyzing results, drafting, and other manuscript tasks, and say they reviewed the AI-assisted work.They state that AI tools were not used to design experiments or implement methods.Source: AI USE STATEMENT · p. 11
A RELATED WORKS
- Unlike prior work that optimizes mixture amounts or adapts proportions, this study fixes compute and cumulative MT tokens and tests whether LR schedules change the preferred temporal allocation.The related-work discussion also contrasts its within-training mixture decision with analyses that prescribe task sequences or use fixed mixing ratios within stages.Source: A RELATED WORKS · pp. 20–21
B TRAINING SETUP AND SCHEDULE DEFINITIONS
The study compares data and learning-rate schedules during continued training from shared PT checkpoints, using dense and MoE models, fixed training setups, and defined validation procedures. It tests abrupt and gradual MT introduction alongside constant and decaying learning rates.
Data and models
- The dense model uses DCLM for PT and FineMath for MT, and both models are pretrained on DCLM before continued-training comparisons.The final PT checkpoint’s dense DCLM and FineMath validation losses are 2.8147 and 2.1472, respectively.Source: B TRAINING SETUP AND SCHEDULE… · pp. 21–22
Models and optimization
- Both model families start continued-training runs from the same final PT weights, with a freshly initialized optimizer and no LR warmup.The dense model has 630M parameters; the MoE model has 2.99B.Source: B TRAINING SETUP AND SCHEDULE… · pp. 21–22
Data schedules
- TS jumps to a chosen MT ratio and holds it, whereas ZL ramps MT upward after an initial PT-only interval.ZL is feasible when 2α ≤ γend ≤ 1; at a shared budget and final ratio, its MT data window is twice TS’s.Source: B TRAINING SETUP AND SCHEDULE… · p. 22
Implementation and evaluation
- Batch composition is reconstructed deterministically from the prescribed schedule, while all final PT/MT comparisons use validation cross-entropy at the final checkpoint.Dense and MoE absolute losses are not directly comparable because the families use different tokenizers and training implementations.Source: B TRAINING SETUP AND SCHEDULE… · pp. 21, 24
LR schedules
- Const retains the peak LR, FDR decays throughout continued training, and MDR begins decay when MT first enters and reaches the final LR at training end.The relative decay duration is measured against the MT data window; 100% corresponds to MDR.Source: B TRAINING SETUP AND SCHEDULE… · p. 24
C COMPLETE DATA SCHEDULE RESULTS FOR THE DENSE MODEL
Dense-model results show that the preferred final MT ratio depends on the LR rule and MT budget. Comparing TS and ZL also depends on whether schedules are evaluated at a shared endpoint or at each schedule’s best tested endpoint.
Endpoint sensitivity
- For the 47.75% MT budget, FDR’s minimum at a 70% endpoint is only 0.0016 below its loss at 60%, making the broad flattening more informative than a precise endpoint preference.ZL’s tested final ratios span only 95.5% to 100%, and its FDR losses vary by only 0.001.Source: C COMPLETE DATA SCHEDULE RESULTS FOR… · p. 25, Figure 10 · p. 29
Data schedule and LR comparisons
- Under Const and MDR, TS’s best tested final MT ratio is 100% at every budget, while under FDR it is 20%, 40%, and 70% as the MT budget increases.The FDR loss range contracts from approximately 0.068 to 0.031 to 0.007; lowering the endpoint does not always improve FDR loss.Source: Table 4 · p. 25, C COMPLETE DATA SCHEDULE RESULTS FOR… · p. 25, Figure 8 · p. 27
PT/MT trade-off
- At the 25% budget under Const, TS has lower MT loss than ZL but substantially higher PT loss.The result illustrates a trade-off between MT and PT validation losses.Source: C COMPLETE DATA SCHEDULE RESULTS FOR… · p. 25
LR-weighted MT exposure
- The paper’s exposure measure weights each MT token by the LR at which it is used, but it omits later updates and optimization noise and does not itself predict final loss.The authors report that ZL can have higher FDR loss despite greater exposure.Source: C COMPLETE DATA SCHEDULE RESULTS FOR… · pp. 25–26
TS versus ZL
- At a shared 100% endpoint, ZL has lower observed MT loss under FDR, while TS has lower loss under Const and MDR.Selecting each schedule’s best endpoint instead favors TS under FDR at all three dense-model budgets; the table reports differences, not significance estimates.Source: C COMPLETE DATA SCHEDULE RESULTS FOR… · p. 25, Table 5 · p. 26
D LEARNING RATE DECAY DURATION SCANS
Across the tested LR-decay scans, delayed decay gives the lowest observed MT loss. The preferred tested duration differs by model and data schedule, and some nearby or incompletely sampled minima remain uncertain.
Additional MoE scans
- In further MoE scans, the best tested durations remain 20% for TS and 10% for ZL, but two earlier-checkpoint scans test only two additional durations, leaving the nearby loss-curve shape uncertain.For 4B-token runs, the minima are 1.6080 at 47 steps for TS and 1.6000 at 48 steps for ZL; for 10B-token runs, both best durations last 119 steps.Source: D LEARNING RATE DECAY DURATION SCANS · pp. 30, 32, Figure 15 · p. 37, Figure 16 · p. 38
Overall scan results
- A delayed LR decay gives the lowest tested loss in all 19 scans, improving on Const and MDR by 0.0094 to 0.0660 relative to MDR.The best tested dense-model durations are 40% for TS and 20% for ZL; for MoE they are 20% and 10%, respectively.Source: D LEARNING RATE DECAY DURATION SCANS · p. 30, Table 6 · p. 30, Table 7 · p. 31
Dense-model scans
- At a 25% dense-model MT budget, TS reaches 1.4275 at 40% and ZL reaches 1.4392 at 20%, with nearby short durations giving similar losses.These configurations have no repeated runs.Source: Figure 12 · p. 34
- At a 47.75% dense-model MT budget, TS reaches its lowest tested loss of 1.3690 at 40%, while ZL reaches 1.3857 at 20%.The figures plot loss against decay duration and show the corresponding LR trajectories.Source: Figure 14 · p. 36
E DATA STREAM REPLICAS AND NOISE ESTIMATES
The replication study repeats six selected TS configurations using altered training examples or example order. Its estimates have limited scope because runs share PT checkpoints and, for dense runs, many examples.
Replication setup
- Six selected configurations are repeated twice, changing training examples or their order while keeping the PT checkpoint and training setup fixed.Dense repeats share about 36% of DCLM examples with the original run; MoE repeats use non-overlapping PT and MT subsets.Source: E DATA STREAM REPLICAS AND NOISE… · p. 39
Scope of noise estimates
- The repeat-based estimates do not capture variation across independently pretrained models, and dense repeats also share many training examples.At the 47.75% MT budget, dense repeats mainly differ in the order of MT examples.Source: E DATA STREAM REPLICAS AND NOISE… · p. 39, Table 9 · p. 39
Interpretation
- The repeats do not establish that small differences between nearby dense ZL decay durations or the best 10B-token MoE TS and ZL settings reliably favor one configuration.Repeating MDR alone also does not estimate variation at other decay durations or data schedules.Source: E DATA STREAM REPLICAS AND NOISE… · p. 39
F MOE RESULTS AND ADDITIONAL ABLATIONS
Across the MoE ablations, preferred data schedules depend on the LR rule, while short final LR decay windows improve tested losses across multiple settings. The findings persist across larger MT budgets, earlier PT checkpoints, and a change from mathematics to code.
LR-dependent preferences persist, but MDR is not uniformly better than Const
- 4B: MDR’s best TS loss is worse than Const’s and its minimum is at 50% MT; at 10B, the MDR–Const gap at 100% shrinks to 0.0115 from 0.0310.At 10B and 50% MT, MDR is 0.0014 below Const; these comparisons do not isolate a single mechanism.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 39, Figure 17 · p. 41, Figure 18 · p. 42
TS versus ZL
- ZL’s best scanned loss is lower than TS’s at 4B, but TS’s is lower at 10B, so neither shape is generally preferable for MoE training.The respective values are 1.6000 versus 1.6080 at 4B, and 1.4853 versus 1.4872 at 10B.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 40, Table 12 · p. 40
A short final decay improves both data shapes
- 1.4007 (TS, 20% window) and 1.4110 (ZL, 10%) are the lowest tested losses; both decay windows span 228 steps and beat their Const baselines.No data-stream replicas were run for this budget.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 44, Figure 21 · p. 46
A larger MT budget at the 10B horizon
- 1.4292 (Const, 100% MT) is lower than 1.4528 (FDR, 47.75% MT), while FDR favors the 47.75% endpoint and Const favors 100%.At 100% MT, MDR reaches 1.4390, 0.0098 above Const.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 44, Figure 20 · p. 45
Earlier checkpoints from the same PT trajectory
- At all three PT checkpoints, Const favors 100% MT while FDR favors 25%; the reversal therefore appears before the full 40B PT run.The checkpoints represent approximately 20B, 30B, and 40B PT tokens from the same PT trajectory.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 47, Figure 22 · p. 47
The preference for a short final decay also persists
- 20% decay gives the lowest tested loss at all three checkpoints, improving on Const, while MDR increases loss at each checkpoint.The earlier-checkpoint scans test only 20% and 40% windows beyond Const, MDR, and FDR, so they do not locate a continuous optimum.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 47, Figure 23 · p. 48, Figure 24 · p. 49
Changing the MT dataset
- 20% for TS and 10% for ZL are the lowest tested relative decay durations with both FineMath and StarCoderData Python; absolute losses are dataset-specific.The two durations correspond to 119 LR-decay steps in the code-data scans.Source: F MOE RESULTS AND ADDITIONAL ABLATIONS · p. 49, Figure 25 · p. 50
G ROUTER DYNAMICS AND INTERPRETATION
The router analysis finds transient imbalance around abrupt TS data entry, examines how expert loads adapt, and compares redistribution patterns across schedules. The additional theoretical setup models SGD with data and LR schedules and uses an SDE approximation to study expected validation loss.
G ROUTER DYNAMICS AND INTERPRETATION
- At both 4B and 10B horizons, abrupt TS entry causes a load-balancing spike that subsides over about 80 steps; larger MT-ratio jumps produce larger spikes.ZL avoids a comparable peak as MT enters gradually, and the effect is concentrated around the distribution change.Source: G ROUTER DYNAMICS AND INTERPRETATION · p. 51, Figure 27 · p. 53, Figure 28 · p. 54
- With the same LR at every step, TS still produces a sharp spike while ZL avoids a comparable one across both training lengths and all three tested final MT ratios.This comparison under Const separates the data-transition contrast from LR decay.Source: G ROUTER DYNAMICS AND INTERPRETATION · p. 51, Figure 26 · p. 52
- Before the switch, routing is nearly balanced on DCLM but uneven on FineMath; during MT training, FineMath loads become more balanced and DCLM loads somewhat less balanced.The authors suggest that PT-trained routing initially assigns MT tokens unevenly across experts.Source: G ROUTER DYNAMICS AND INTERPRETATION · p. 51, Table 13 · p. 52, Figure 29 · p. 55
- TS+Const and ZL+MDR show comparable net redistribution, but ZL+MDR has less scatter and preserves more of the initial expert ordering.Because the comparison uses different LR schedules and similar but nonidentical starting routers, it describes redistribution patterns rather than isolating data shape.Source: G ROUTER DYNAMICS AND INTERPRETATION · p. 56, Table 14 · p. 56, Figure 30 · p. 57
H ADDITIONAL THEORETICAL SETUP
- The theoretical setup samples MT data with probability γ(k), uses an LR schedule in online SGD, and approximates its expected validation-loss objective through a continuous-time SDE.The approximation replaces randomness from gradient noise and data sampling with Brownian motion and data-sampling randomness.Source: H ADDITIONAL THEORETICAL SETUP · p. 58
I FORMAL STATEMENTS OF THEORETICAL RESULTS
The theoretical results show that the optimal data schedule depends on the learning-rate schedule: the constant-rate optimum need not remain optimal under linear decay. A power-law spectral model gives a closed-form characterization under linear decay, with proofs based on convexity and KKT conditions.
I FORMAL STATEMENTS OF THEORETICAL RESULTS
- Under a constant learning rate, the optimal schedule is the 0–1 data schedule.The theorem identifies this result under Condition 3.1.Source: I FORMAL STATEMENTS OF THEORETICAL… · p. 58
- Changing only the learning-rate schedule to linear decay can make the constant-rate 0–1 schedule non-optimal.Theorem I.2 gives sufficient conditions for this change but does not characterize the exact linear-decay optimum.Source: I FORMAL STATEMENTS OF THEORETICAL… · pp. 58–59
- A power-law spectral model enables a closed-form optimal-schedule result for linearly decaying learning rates.The theorem assumes 0 < ρ1 < 1, a fixed data fraction below x⋆, and sufficiently large T.Source: I FORMAL STATEMENTS OF THEORETICAL… · p. 59
- The power-law spectral assumption gives eigenvalues density proportional to λαl exp(−Klλ), with −1 < αl < 0 and Kl > 0.The noise variance conditioned on eigenvalue λ is Flλρl exp(−Glλ), with Fl and Gl positive and independent of the learning-rate schedule.Source: I FORMAL STATEMENTS OF THEORETICAL… · p. 59
- The spectral density places more mass near zero, representing many slowly learned directions, while its exponential tail makes very large eigenvalues increasingly rare.The noise model also makes average noise variance decay for very large eigenvalues.Source: I FORMAL STATEMENTS OF THEORETICAL… · p. 59
- Under linear decay, the characterized schedule uses two transition times, each scaling as Θ(T).The theorem minimizes expected validation loss under Assumption I.1 and Condition 3.1.Source: I FORMAL STATEMENTS OF THEORETICAL… · p. 59
J PROOF OF THEOREM I.1
- The proof derives the validation loss from the SDE and establishes convexity of the objective and feasible set, making KKT conditions necessary and sufficient.The proof also establishes uniqueness of an optimal schedule when one exists.Source: J PROOF OF THEOREM I.1 · p. 60, §J.1 · p. 60
J.2 KKT CONDITION FOR EQUATION (2)
- The KKT conditions compare the marginal value of MT data with a threshold to determine whether the schedule is zero, interior, or one.The marginal value of MT data is defined before the pointwise conditions are stated.Source: §J.2 · p. 62
J.3 PROOF OF THEOREM I.1
The proof establishes the constant-learning-rate optimum by comparing schedules under a fixed MT budget, then shows why linear decay can invalidate that optimum. It derives the objective and uses KKT conditions to characterize feasible optimal schedules.
J.3 PROOF OF THEOREM I.1
- For the constant-rate proof, the candidate schedule attains the upper bound on remaining MT mass at every time, yielding the global comparison with feasible schedules.Equality forces the candidate schedule to match almost everywhere.Source: §J.3 · p. 62
- With a constant learning rate, the optimal schedule is the 0–1 schedule, and adding MT mass where the schedule is below one strictly lowers the objective.Except at the boundary budgets α = 0 and α = 1, an optimum exhausts the MT budget; the 0–1 schedule is unique up to measure zero.Source: §J.3 · p. 62
K PROOF OF THEOREM I.2
- Changing the learning rate from constant to linear decay can make that 0–1 schedule non-optimal, even with other settings fixed.The contradiction proof uses the changing marginal value around the transition time.Source: K PROOF OF THEOREM I.2 · pp. 63–64
L.1 DERIVATIONS OF THE MINIMIZATION OBJECTIVE
- The minimization objective retains only the shared and MT-specific loss contributions, and both schedule-dependent terms are nonincreasing as MT data is added.For 0 < α < 1, an optimum therefore exhausts the MT budget.Source: §L.1 · pp. 64–65
L.2 KKT CONDITION
- The pointwise KKT conditions set the schedule to zero below the marginal-value threshold, to one above it, and allow an interior value at equality.The threshold is the multiplier for the MT-budget constraint.Source: §L.2 · p. 66
- The derived objective is convex in the data schedule, so its KKT conditions are necessary and sufficient for optimality.The optimal schedule is unique if it exists.Source: §L.1 · p. 65, §L.2 · p. 66
L.4 PROOF OF THEOREM I.3
The proof constructs and verifies a two-transition optimal schedule under the stated power-law model. Numerical checks support both the SDE approximation and the closed-form schedule, while schedule shape changes with learning-rate decay strength.
L.4 PROOF OF THEOREM I.3
- For sufficiently large T, the proof verifies a valid schedule with zero MT data before t1, an interior segment, and full MT data after t2.The transition times satisfy 0 < t1 < t2 < T, and the MT budget is active.Source: §L.4 · pp. 67–68
M NUMERICAL VALIDATION OF THE THEORETICAL MODEL
- The numerical validation uses the quadratic model, power-law spectrum, linearly decaying learning rate, and constrained MT budget specified by the theorem.Expected validation loss is calculated using the closed form of the SDE approximation.Source: M NUMERICAL VALIDATION OF THE… · p. 70, §M.1 · p. 71
M.1 ADDITIONAL DETAILS ON THE SIMULATION EXPERIMENTS
- The closed-form schedule is computed by solving a KKT sufficient condition and the active MT-budget equation for its two transition times.The interior schedule is then evaluated between those times, with values zero before t1 and one after t2.Source: §M.1 · p. 71
M.2 CHECKS OF THE SDE APPROXIMATION
- The SDE approximation, exact discrete SGD second-moment recursion, and Monte Carlo estimates are almost identical across four representative schedules.Figure 31 reports constant learning rate and linear decay with c = 0.8; its error bars are two standard errors of the Monte Carlo mean.Source: Figure 31 · p. 72, §M.2 · p. 72
M.3 NUMERICAL VALIDATION OF THEORETICAL PREDICTIONS
- The closed-form schedule and directly optimized solutions agree across both MT budgets, including their three-stage structure and transition locations.The stages are PT only, a continuously increasing MT ratio, and MT only.Source: §M.3 · p. 72, Figure 32 · p. 73
M.4 HOW THE OPTIMAL DATA SCHEDULE CHANGES WITH THE LEARNING RATE SCHEDULE
- Under constant or weakly decaying learning rates, the numerical optimum is PT-only followed by MT-only; sufficiently strong decay yields a three-stage schedule with increasing MT share.The reported examples use c = 0.2 for weak decay and c = 0.5 or 0.8 for strong decay.Source: §M.4 · p. 73
M.5 VERIFICATION OF THE KKT CONDITION OF THE OPTIMAL DATA SCHEDULE
This section verifies optimal data-schedule conditions under constant and linearly decaying learning rates, and proves that optimal learning-rate decay begins after MT data enters training under the stated assumptions. It also shows that the objective cannot attain its minimum at W = 0 for sufficiently large T.
M.5 VERIFICATION OF THE KKT CONDITION OF THE OPTIMAL DATA SCHEDULE
- For the linear-decay optimum, Mγ(t) increases during PT-only and MT-only stages and is constant during the mixed-data stage, satisfying the KKT conditions.The multiplier κ is set to the rescaled Mγ(t) value in the mixing stage.Source: §M.5 · p. 73, Figure 33 · p. 74
- The two-stage schedule satisfies the KKT condition under constant LR but violates it under linear decay, so its optimality depends on the LR schedule.Figure 33 illustrates the switching condition around (1−α)T; the prose states the schedule is optimal for constant LR and not for linear decay.Source: §M.5 · p. 73, Figure 33 · p. 74
N.1 SETUP AND FORMAL STATEMENT
- Under the stated assumptions and for fixed α∈(0,αc), Theorem N.1 guarantees a global minimizer for sufficiently large T, with LR decay beginning strictly after MT data enters.The theorem assumes 1<b1<2 and b0≥b1.Source: §N.1 · p. 75
N.2 DERIVATION OF THE EXPECTED VALIDATION LOSS
- The expected final MT loss is rewritten as an objective combining shared- and MT-specific bias terms, variance integrals, and an additional schedule-dependent term.The derivation retains the shared-block contribution because it depends on the LR schedule.Source: §N.2 · pp. 75–76
N.3 AUXILIARY LEMMAS
- The auxiliary lemmas establish that αc is positive, A_W(u) and Iβ,g(H,W) are nonincreasing in W, and the bias contribution is nondecreasing in W.A further bound shows P(W) is at least the smaller of its endpoint values P(L) and P(T).Source: §N.3 · pp. 77–79
N.4 PROOF OF THEOREM N.1
- The proof establishes continuity of the objective on [0,T] and compares schedules with W≥L against the reference W°=κL to rule out a minimum in that range.The supplied proof passage introduces this comparison as the second step of Theorem N.1’s proof.Source: §N.4 · p. 81
- For sufficiently large T, J_T(W◦) < J_T(0), ruling out W = 0 as a minimizer.The proof establishes a threshold T0 such that the strict inequality holds for all T > T0.Source: §N.4 · pp. 85–86
- The comparison uses W◦ = καT, with H = L or T satisfying H ≥ αT > W◦.Source: §N.4 · p. 85
- For H ≥ αT, the terms B_l[ K_l/2 + A_W(H)]^a_l vanish as T → ∞ for W = 0 and W = W◦.The proof bounds A_W(H) below by mαT for both choices of W.Source: §N.4 · p. 85