Learning-rate warmup is a standard technique in language-model training, yet its duration remains largely heuristic. Common approaches use either a fixed number of updates or a fixed fraction of the training horizon, two choices that imply very different scaling as training gets longer. When should warmup stay fixed, and when should it grow with the horizon? We address this question with a quadratic model whose modes respond differently to the peak learning rate. Warmup slows progress in directions that already contract well at the peak rate, but can remove persistent error in directions near the stability edge, with higher peak rates shifting the balance toward longer warmup durations. This yields a compact horizon scaling law that captures regimes ranging from essentially no warmup, through fixed-duration warmup, to durations that grow with the training horizon, and explains how the preferred regime changes with peak learning rate. Because the law captures the tradeoff between giving up early progress and improving the trajectory that follows, it can be fit using shorter runs and used to predict warmup at substantially longer horizons. Together, our results explain several familiar properties of warmup through a single tradeoff and suggest treating warmup duration as a horizon-dependent hyperparameter rather than a fixed training heuristic.
Figures & tables
Figure 1: Overview. (a) Longer warmup can trade slower early progress for better performance at aggressive peak learning rates. (b) Schematic of the loss-optimal warmup regimes, from little or bounded warmup to durations that grow with the training horizon, depending on peak learning rate. The star marks the approximate location of the best measured runs
Figure 2: Peak learning rate changes how optimal warmup evolves with horizon. (a,b) Lower rates favor short warmup, while higher rates increasingly favor longer durations as training continues. (c) At high peak rate, a longer warmup can start worse but become better later, so its benefit persists beyond the warmup duration.
Figure 3: Fixed and horizon-dependent warmup regimes coexist. Lines summarize how the measured loss-optimal warmup changes with training horizon within each family.
Figure 4: Shortest successful warmup durations. Shortest tested successful warmup versus (a) peak learning rate, (b) batch size, and (c) model size.
Model
Peak LR
Warmup (k)
Fraction of horizon (T=128k)
Loss
60M
0.004
30
23.4%
3.404
100M
0.005
50
39.1%
3.246
350M
0.004
40
31.3%
2.983
Table 1: The best tested runs use long warmups. Learning rate and warmup are jointly tuned.
Figure 5: Long warmups improve the loss attainable at high peak rates. Validation loss at 128k with warmup tuned freely or capped at 1k. Loss is relative to the best joint configuration at each scale; capped curves end when no successful short-warmup run remains. Additional learning rates used for the stability comparison are omitted because they used a sparser warmup grid.
Figure 6: Warmup trades early progress for lower persistent error. (a,b) A longer duration slows progress in one direction but reduces error in another that improves only slowly at the peak learning rate. (c) As a result, the longer warmup is worse early in training but better later, matching the ranking reversal observed empirically.
Evaluation
Median R2
RMSE
Mean regret
All horizons (descriptive)
0.996
4.92
2.00±1.51
Alternating holdout
0.996
4.98
2.05±1.61
Table 2: Absolute-fit accuracy and selection quality. R2 and RMSE are medians across families; regret is the mean loss above the best tested warmup. RMSE and regret are in 10−3 loss units.
Figure 7: The loss law captures qualitatively different warmup regimes. Selected illustrative full-data fits span low and high peak learning rates across model scales.
Figure 8: Growth and scale summarize the fitted warmup regimes. (a) Warmup growth exponent β=(p+1−q)/(s+1) from the fitted loss law. (b) Corresponding warmup scale M32=(2sK/(Ap))1/(s+1)32β for the interior approximation W∗≈M32(T/32)β . (c) Representative 100M families comparing the continuous, model-selected and measured optimum.
Selector
Parameters
Alternating
Four horizons
128k
Absolute fit
6
2.05±1.61
2.26±0.66
3.05
Difference fit
5
1.45±2.22
1.58±0.77
2.61
1k target
0
15.47±10.80
13.79±0.22
13.56
10% target
0
13.39±4.10
12.58±0.90
11.72
Table 3: Fits through 32k select low-regret warmups at later horizons. Mean regret in 10−3 loss units. Both fits use checkpoints through 32k updates for the four-horizon and 128k evaluations; the alternating column uses interleaved held-out checkpoints.
Figure 9: Warmup selection extrapolates beyond the fit window. Bars average regret over 50k, 75k, 100k, and 128k after fitting through 32k, then equally over all fitted peak learning rates within each scale and indicated range.
Fit horizon
Absolute
Difference
Fixed dur.
Fixed frac.
1k
10%
20k (29 families)
2.83±0.16
4.22±0.29
9.31±1.21
11.72±0.82
14.28±0.30
13.42±0.96
32k (33 families)
2.51±0.27
2.62±0.36
7.61±1.12
8.47±0.62
13.79±0.22
12.58±0.90
Table 4: Three short runs support longer-horizon warmup prediction. Mean regret over 50k, 75k, 100k, and 128k, in 10−3 loss units. Fixed dur. and fixed frac. keep the best fitting-run duration or its fraction of the fit horizon.
Figure 10: Three-run selection by scale and peak learning rate. Absolute-fit selection and the two baselines after fitting through 32k; regret averaged over 50k, 75k, 100k, and 128k.
LR =0.5×10−3
LR =2×10−3
LR =5×10−3
T (k)
250
500
1000
250
500
1000
250
500
1000
1k-warmup loss
3.370
3.337
3.311
3.433
3.421
3.402
3.371
3.347
3.322
10% loss
3.375
3.347
3.339
3.418
3.367
3.340
3.363
3.335
3.314
W50
0.6
0.7
0.8
47.7
77.5
125.9
129.5
241.7
451.1
Loss
3.372
3.336
3.311
3.412
3.362
3.338
3.359
3.325
3.303
W75
0.8
1.1
1.2
56.6
95.2
160.0
116.0
213.5
392.9
Table 5: Long-horizon validation-loss comparison. Three-run absolute-fit forecasts from 50k or 75k prefixes, compared with 1k-warmup and 10% baselines. Horizons and predicted warmups are in thousands of updates.
Appendix figures & tables41 assets
Supplementary material from the paper’s appendix.
Appendix
Setting
Value
Training corpus
C4 English; 1,024 shards; 184,565,529,484 tokens
Tokenizer / storage
t5-base SentencePiece / uint16
Sequence length
256
Batch size
512, 131,072 tokens per update
Training budget
128,000 updates; 16,777,216,000 tokens
AdamW (β1,β2) / ϵ
(0.9,0.99) / 10−6
Appendix
Table 6: Base training configuration.
Setting
60M
100M
350M
Layers
8
12
24
Hidden dimension
512
640
1,024
MLP intermediate dimension
1,376
1,708
2,736
Attention heads / head dimension
8 / 64
10 / 64
16 / 64
Vocabulary size
32,000
32,100
32,000
Total parameters
58,073,600
100,117,120
367,969,280
Appendix
Table 7: Model and execution settings. All models use Llama decoder blocks with SwiGLU, rotary position embeddings, RMSNorm with ϵ=10−6 , initialization standard deviation 0.02, and untied input embeddings and output heads. The configured maximum sequence length is 1,024; training uses 256 tokens.
Scale
Peak LRs ( 10−3 )
Families
Warmup range (updates)
60M
0.1, 0.3, 0.5, 0.7, 1, 2, 3, 4, 5, 6, 7, 8
12
0–50,000
100M
0.1, 0.3, 0.5, 0.7, 1, 2, 3, 4, 5, 6, 7, 8
12
0–50,000
350M
0.5, 0.7, 1, 3, 4, 5, 6, 7, 8
9
0–64,000
Appendix
Table 8: Experimental grid. Peak learning rates and warmup ranges for the 33 model-size and peak-rate families; warmup grids vary with model size and peak learning rate.
Setting
Choice
Fitting optimizer
Trust-region reflective (SciPy least_squares )
Huber transition
0.02 loss units
Warmup offset
w0=0.032
Exponent bounds
p,q,s∈[0.05,1]
Absolute amplitudes
A,K/w0s∈[10−9,105]
Difference amplitude
C∈[10−9,104]
Appendix
Table 9: Fitting settings. Time is measured in thousands of updates. Widening the exponent bounds to [0.05,2] , which allows β≤0 , gave slightly higher selection regret.
Figure 11: What warmup does to one mode. ρ(α) : log-contraction of one step at the peak rate. g(α) : its average over a linear warmup from 0 to α ( Equation 13 ). Dots: αcross , where g=ρ , and αmax , where g=0 .
Figure 12: The two-mode optimum ( Theorem 1 ) with the curvatures of Figure 6 , he=1 and hs/he=0.04 , so ηc≃1.923 . (a) At η=1.995 and W=xT , the loss decays at rate min(ρs−dsx,ρe+aex) (solid), which is largest where the two modal rates (dashed) are equal, at x=xη . (b) Wquad∗ against T ( Ae/As=100 ); dotted: W∗=T . (c) xη against the peak rate. Ticks: ηc . Circles: 2hs/he .
Figure 13: Warmup above the stability limit ( Theorem 2 ). (a) C/T=ρ+aW/T against the warmup fraction; dots mark κ(α) . (b) The threshold κ(α) .
Figure 14: Sampling flips the warmup ranking. The two-mode example of Equation 35 with ϵ=0.02 , ηh=0.4 , B=2 , and W=20 ; losses are relative to the initial loss. (a) Warmup empties the sharp mode more slowly. (b) Warmup leaks less error into the flat mode. (c) Warmup is worse until step 25, then settles at about a quarter of the no-warmup loss. The exact recursion in Equation 34 gives the same picture, with the crossing at step 24.
Peak LR ( 10−3 )
Source warmups (updates)
0.1
0, 256, 1,000, 4,000, 16,000
0.3
0, 128, 512, 2,000, 8,000
1
0, 512, 2,000, 8,000, 30,000
3
512, 1,500, 8,000, 16,000, 30,000
5
768, 2,000, 6,000, 16,000, 30,000
7
1,000, 3,000, 8,000, 20,000, 50,000
Appendix
Table 10: Decay-branching design. Model: 100M; AdamW; seed 0. Fork horizons C∈{8,16,32,64,128} k updates; linear decay to zero lasts 5k updates, with endpoint C+5 k. Only W≤C−2 k is eligible. The comparison uses the 128 branches that completed.
Figure 15: Warmup ordering before and after decay. Left: Spearman correlation between warmup rankings at the fork and after 5k decay updates, within each peak-rate family ( 1 : identical ordering; −1 : reversed). Right: post-decay regret of the warmup selected before decay, in 10−3 loss units. Each cell compares the same warmups before and after decay.
Fork (k)
Branches
Preserved pairs
Same best
Mean regret
Max regret
8
18
19/20
6/6
0.00
0.00
16
22
27/31
3/6
1.38
4.65
32
29
50/56
4/6
1.61
8.78
64
29
49/56
4/6
0.53
2.81
128
30
57/60
5/6
0.02
0.10
Appendix
Table 11: Ordering agreement and selection after decay. Preserved pairs have the same strict loss ordering before and after decay. “Same best” counts peak-rate families whose best warmup is unchanged by decay. Mean regret weights families equally; maximum regret is the largest family regret. Both are in 10−3 loss units.
Setting
Value
Muon parameters
Hidden matrices: 59,013,120
Auxiliary AdamW parameters
Embeddings, output head, norms: 41,104,000
Muon momentum / Nesterov
0.95 / enabled
Newton–Schulz iterations
5
Muon update scaling
0.2max(dout,din) after approximate orthogonalization
Muon peak LR ( 10−3 )
0.1, 0.2, 0.5, 1, 2
Appendix
Table 12: Muon protocol and controls. Model: 100M; seed 0; 128k updates. All remaining data, batch, precision, clipping, and evaluation settings match Tables 6 and 7 .
Muon peak LR ( 10−3 )
8k
16k
32k
64k
128k
0.1
0
512
512
512
2,000
0.2
512
2,000
2,000
8,000
8,000
0.5
512
2,000
8,000
8,000
8,000
1
512
2,000
8,000
16,000
16,000
2
2,000
2,000
8,000
16,000
50,000
Appendix
Table 13: Best measured Muon warmup (updates). Columns give the training horizon; only completed runs with W<T are eligible. Ties favor shorter warmup. The loss curves in Figure 16 show how sharply each minimum is determined.
Figure 16: Warmup dependence with Muon. Loss gaps are relative to the best tested warmup at each horizon and Muon peak rate, using only W<T from the 30 completed runs.
Shape
m(x)
∫01m(x)dx
τ=∫0Tηt/ηdt
Linear
x
1/2
T−W/2
Half-cosine
(1−cos(πx))/2
1/2
T−W/2
Concave quadratic
2x−x2
2/3
T−W/3
Appendix
Table 14: Warmup shapes and design. Here x=t/W during warmup, ηt=ηm(x) , and m=1 after warmup. The last columns use continuous integrals, with T≥W . Model: 60M; AdamW; peak LRs {0.5,1,5}×10−3 ; nonlinear warmups {500,1000,5000,15000,50000} ; horizon 128k.
Figure 17: Warmup learning-rate profiles. Relative learning rate ηt/η against relative time t/W . All shapes reach the same peak at the end of warmup and remain constant afterward.
Figure 18: Horizon-dependent warmup under different shapes. Each panel compares durations within a fixed shape and peak rate. Rows show linear, half-cosine, and concave-quadratic warmup; columns show peak learning rates. Loss gaps are relative to the best tested duration within each shape, peak rate, and horizon, restricted to W<T . Each shape uses its own tested duration grid; linear curves follow the base-grid eligibility rule in Section A.1 . Lines connect measured candidates.
Shape
Peak LR ( 10−3 )
8k
16k
32k
64k
128k
Linear
0.5
256
256
256
384
768
Linear
1
128
128
768
768
1,500
Linear
5
768
2,000
4,000
20,000
30,000
Half-cosine
0.5
500
500
500
500
500
Half-cosine
1
500
500
500
500
500
Half-cosine
5
1,000
1,000
5,000
15,000
15,000
Appendix
Table 15: Best measured warmup across horizons within each shape (updates). Each row tracks the loss-minimizing duration for a fixed shape and peak rate. Columns give the training horizon; only W<T is eligible, and exact ties favor shorter warmup.
Selector
Parameters
Alternating
Four horizons
128k
Absolute fit
6
2.05±1.61
2.26±0.66
3.05
Difference fit
5
1.45±2.22
1.58±0.77
2.61
Direct warmup scaling
2
1.34±4.94
1.99±0.69
2.86
Fixed best duration
0
–
3.00±1.75
3.87
Fixed best fraction
0
–
1.73±0.28
2.06
1k target
0
15.47±10.80
13.79±0.22
13.56
Appendix
Table 16: Selection results including the direct warmup-scaling baseline. All methods are evaluated on the same retained candidate sets. For the tuned rules, the alternating column would measure one checkpoint ahead rather than extrapolation, so it is omitted.
Tolerance ( 10−3 )
Multiple
Median count
Median fraction of grid
1
62 / 132
1
16.7%
2
92 / 132
2
20.0%
5
124 / 132
4
40.8%
Appendix
Table 17: Near-optimal measured warmups. The table summarizes all retained family–horizon comparisons used by the later-horizon benchmark. “Multiple” counts comparisons with more than one tested duration within the stated loss tolerance.
Item
Choice
Target fitting warmups
2k, 8k, 16k updates
Fit horizons
20k and 32k updates
Target horizons
50k, 75k, 100k, 128k updates
Loss-law fits
absolute and difference forms
Primary practical predictor
absolute fit
Target candidates
same retained warmup grid as the full-grid benchmark
Appendix
Table 18: Three-run extrapolation protocol.
Selector
Through 20k (29 families)
Through 32k (33 families)
Absolute fit
2.83±0.16
2.51±0.27
Difference fit
4.22±0.29
2.62±0.36
Direct warmup scaling
9.09±1.25
7.08±0.82
Fixed best duration
9.31±1.21
7.61±1.12
Fixed best fraction
11.72±0.82
8.47±0.62
1k target
14.28±0.30
13.79±0.22
Appendix
Table 19: Complete three-run comparison. Mean regret on the same later-horizon candidate grids as the main-text experiment, in 10−3 loss units.
Figure 19: Where three-run warmup recommendations fall relative to the observed warmups. Left: number of family–horizon selections that coincide with an observed fitting warmup, fall between observed warmups, or lie below or above their range. Right: mean measured regret within each group. All selections are frozen before target losses are scored.
Figure 20: Sensitivity to the warmup coverage of three fitting runs. Left: mean four-horizon regret across families. Right: the same comparison split by model size. Random-triplet results are first averaged within family, so families with more admissible triplets receive no additional weight.
Fitting-run design
Log-range positions / rule
Mean regret
Median family
Worst family
Default triplet
Main-experiment runs
2.51
1.67
8.88
Wide coverage
0,0.5,1
2.15
1.47
6.80
Short-end cluster
0,0.1,0.2
2.32
1.47
12.34
Middle cluster
0.4,0.5,0.6
3.55
2.13
13.44
Long-end cluster
0.8,0.9,1
5.19
1.67
49.95
Random triplets
Up to 8 admissible triplets
2.63
2.21
6.72
Appendix
Table 20: Choice of three fitting runs and selection results. All fits use shared post-warmup checkpoints through 32k, with at least four shared checkpoints. Target positions are normalized coordinates in log(W+w0) over each family’s available warmup range. The default row uses the exact warmup durations from the main experiment. Family values first average regret over the four target horizons; random triplets are averaged within family before aggregation. Regret is in 10−3 loss units.
Figure 21: Selection as the number of fitting runs increases. At a cap n , each family uses up to n available trajectories. Fitting uses checkpoints through 32k. Purple shows the nested coverage order; cyan shows the mean of three fixed random orders, with shading spanning their aggregate curves. Dotted lines show the full-grid reference. The same 33 families are used throughout; their available counts range from 3 to 19. Regret averages the four later target horizons.
Reference
Fitting runs
Checkpoints used
Mean regret
Main three-run fit
3
Shared post-warmup
2.51
Count-sweep fit
3
All selected post-warmup
3.32
Full-grid fit
All available
All post-warmup
2.26
Appendix
Table 21: Reference fits for varying the number of runs. The main three-run protocol uses shared post-warmup checkpoints, while the count sweep uses all post-warmup observations from each selected trajectory. The full-grid row reproduces the corresponding full-prefix benchmark. Regret is in 10−3 loss units.
Figure 22: Sensitivity to fit horizon and short-run budget. The same three fitting trajectories are used at every fit horizon, from 20k to 64k. All comparisons use the same 29 families and the 128k target horizon. Total fitting-run updates equal three times the fit horizon, excluding the target run. Left: mean final-horizon regret against the total updates observed across the fitting runs. Right: the same comparison by model size as the fit horizon increases.
Item
Setting
Model / batch / sequence length
100M / 128 sequences / 256 tokens
Fit horizons
50k and 75k updates
Target horizons
250k, 500k, and 1M updates
Fitting warmups, LR 0.5×10−3
0, 512, 8k warmup updates
Fitting warmups, LR 2×10−3
8k, 16k, 32k warmup updates
Fitting warmups, LR 5×10−3
12k, 16k, 32k warmup updates
Appendix
Table 22: Long-horizon protocol.
Figure 23: Warmup behavior across model scales. The larger model uses peak LR 0.003 ; the smaller model uses the same LR (gray) or a fitted higher LR (orange). We average log(W+32) , with W in steps, over seed-0 warmup sweeps with weights proportional to exp(−L/0.002) , then fit monotonically in log LR. The multiplier minimizes weighted squared log-warmup differences across horizons. The 100M-to-350M fit reaches the LR limit, 0.007 .
Metric
Free s
Fixed s=1/2
Full-data absolute-fit median RMSE
4.92
5.94
Held-out absolute-fit median RMSE
4.98
5.96
Four-horizon absolute-fit regret
2.26±0.66
2.00±0.35
Three-run absolute-fit regret, 20k (29 families)
2.83±0.16
2.87±0.32
Three-run absolute-fit regret, 32k
2.51±0.27
2.64±0.23
Four-horizon difference-fit regret
1.58±0.77
1.58±0.67
Appendix
Table 23: Sensitivity to the warmup exponent. RMSE and regret are in 10−3 loss units. Lower values are better.
Difference fit
Absolute fit
w0 (updates)
Later-horizon regret
Final-horizon regret
Later-horizon regret
Final-horizon regret
8
1.50
2.58
2.32
3.54
16
1.84
2.54
2.09
2.83
32
1.85
2.61
2.49
3.05
64
1.61
2.67
2.41
3.16
128
1.68
2.92
2.28
3.06
Appendix
Table 24: Sensitivity to the fixed warmup offset. Each row refits all families with the indicated global offset. Here, later-horizon regret averages the dense 33k–128k grid, rather than the four target horizons used in the main comparison. Regret is in 10−3 loss units.
Figure 24: Freely fitted and shared persistent exponent sums. Points show qη+sη from independent full-data fits at each peak learning rate. Dashed lines show the single sum fitted jointly across all peak-rate families of the same model size. Points at 2.0 lie on the parameter bound ( q=s=1 ).
Figure 25: Effect of sharing q+s across peak learning rates. Left: median descriptive RMSE for free and shared-sum fits at each model size. Right: later-horizon selection regret after fitting only through the 32k fit window. The shared sum is learned from the same fitting data and is not transferred from the full-data fit.
Model
Fit
Full RMSE
Full R2
Later regret
60M
Free
6.59
0.993
2.83
60M
Shared
6.92
0.993
2.24
100M
Free
4.71
0.997
2.73
100M
Shared
4.89
0.997
3.02
350M
Free
4.81
0.997
0.88
350M
Shared
4.88
0.997
1.77
Appendix
Table 25: Shared-exponent-sum ablation. RMSE and regret are in 10−3 loss units. RMSE and R2 are medians over peak-rate families. Full-data fit quality is descriptive; later regret averages families and the four later target horizons after fitting through 32k.
Learning rate warm-up -- increasing the learning rate at the beginning of training -- has become a ubiquitous heuristic in modern deep learning, yet its theoretical foundations remain poorly understood. In this work, we provide a principled explanation for why warm-up improves training. We rely on a generalization of the (L0,L1)-smoothness condition, which bounds local curvature as a linear function of the loss suboptimality and exhibits desirable closure properties. We show -- both theoretically and empirically -- that this condition is satisfied by common neural architectures and accurately captures the curvature of the optimization landscape early in training. Adapting the learning rate in response to this curvature condition naturally induces a warm-up-like schedule, and we show that this choice yields provably faster convergence guarantees than using a fixed learning rate. Experiments on language and vision models show that the resulting one-parameter warm-up schedule can match tuned linear warm-up and improve over no warm-up.
Standard learning rate schedules such as cosine annealing are tied to a fixed training horizon, limiting their ability to accommodate post hoc horizon extension. Warmup-stable-decay (WSD) partially addresses this issue by maintaining a long constant-rate phase before a short linear cooldown, allowing training to resume from a pre-decay checkpoint. However, its peak learning rate is still tuned based on the original training horizon and can become suboptimal when training is extended. Motivated by stochastic convex optimization, we propose WSqD (Warmup with Square-root base and linear Decay), a learning rate schedule that replaces WSD's constant stable phase with a shifted inverse-square-root base while retaining the final linear cooldown. In the stochastic convex setting, WSqD provably attains the minimax-optimal O(1/T) last-iterate convergence rate. Importantly, its base learning rate schedule is horizon-independent, and the training horizon is needed only to determine when to begin the final cooldown. Empirically, on language-model pretraining using the SlimPajama corpus, WSqD matches or outperforms carefully tuned WSD and other baselines across multiple training horizons while reusing a single peak learning rate.
Jianhao Ma, Yuxin Chen
Department of Statistics and Data Science, University of Pennsylvania
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.