PaperlayerView exact revision on arXiv

Source-linked AI summary

Balancing Early Performance Sacrifices with Long-Term Gains: Scaling Learning-Rate Warmup Duration Across Training Horizons

Kristi Topollai, Anna Choromanska

arXiv:2609.33041v1cs.LG

TL;DR

Warmup is widely used, but its optimal duration across training horizons remains unclear because fixed-step and fixed-fraction choices imply different scaling. This paper uses a quadratic model and a compact loss law to explain how the preferred duration shifts with peak learning rate, and reports that the law can guide choices at longer horizons from short runs.

  • Problem

    Although warmup is standard, there is little agreement on its duration, and prior explanations for why it helps do not determine the optimal balance between persistent benefits and lost early progress.

    Source: §1 · p. 1

  • Method

    The paper analyzes a quadratic model whose curvature directions respond differently to warmup and peak rates, then fits a compact loss law to short trajectories for later-horizon predictions.

    Source: §1 · p. 2, §3 · p. 4, §5 · p. 7

  • Results

    Lower peak rates often favor short or bounded warmup, whereas higher rates favor durations that grow with the training horizon.

    Source: §2.2 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    Warmup duration is a horizon-dependent hyperparameter that should be tuned jointly with peak learning rate and training budget.

    Source: §6 · p. 9

    The fixed-quadratic analysis isolates warmup duration from schedule-ordering effects, although ordering can matter in neural-network training.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Warmup duration presents a tradeoff between stabilizing training at larger peak learning rates and preserving early optimization progress. The paper studies this tradeoff across horizons and proposes a model that explains why optimal duration depends on peak rate and training budget.

1 INTRODUCTION

  • Too little warmup can destabilize training at larger peak rates, while too much sacrifices early progress, and the balance depends on the peak rate itself.A rate that appears suboptimal with short warmup can become competitive when longer warmup is allowed.Source: §1 · p. 1
  • Prior explanations of why warmup helps do not determine its optimal duration, which requires balancing persistent benefits against progress lost to smaller initial steps.The cited accounts include optimizer statistics, loss curvature, parameter and representation changes, and suboptimality-dependent smoothness.Source: §1 · p. 1
  • The empirical study finds that optimal warmup can stay nearly fixed at lower peak rates, grow with the horizon at larger rates, and change which learning rates are competitive.The study varies peak learning rate and warmup across training horizons.Source: §1 · pp. 1–2
  • A quadratic model motivates a compact loss law that can be fit to short trajectories and used to predict warmup choices at longer horizons.The paper reports predictions using as few as three short runs.Source: §1 · p. 2
  • Warmup duration has its own horizon dependence, shaped by peak learning rate, placing it among training hyperparameters whose preferred values change with budget.The paper frames this dependence as a scaling law rather than a fixed startup heuristic.Source: §1 · p. 2

2 WARMUP REGIMES ACROSS LEARNING RATES AND TRAINING HORIZONS

Across training horizons, loss-optimal warmup can remain nearly fixed or grow with the horizon, with longer durations favored more often at higher peak learning rates. Longer warmups can improve attainable loss and change which learning rate appears best, so stability and loss optimization impose distinct requirements.

2.2 FROM SHORT WARMUP TO HORIZON-DEPENDENT WARMUP

  • At conservative peak rates, little or no warmup can suffice; at larger rates, the loss-optimal duration shifts longer as training continues, with growth ranging from weak to nearly proportional.The study evaluates validation loss across horizons under linear warmup followed by a constant rate; one warmup–stable trajectory can be assessed at multiple horizons.Source: §2.2 · p. 2, §2.1 · p. 2, Figure 2 · p. 3
  • Even among successful runs, a longer warmup can start worse and later become better, so horizon dependence is not only about avoiding instability.Across families, longer optimal warmups become more common at larger peak learning rates.Source: §2.2 · p. 3, Figure 3 · p. 3
  • The shortest tested warmup that permits convergence increases with higher peak rates, smaller batches, and larger models, but is often far shorter than the loss-optimal warmup.Figure 4 compares this threshold against peak learning rate, batch size, and model size.Source: §2.2 · p. 3, Figure 4 · p. 3

2.3 LONG WARMUPS UNLOCK COMPETITIVE HIGH LEARNING RATES

  • At larger peak rates, long warmups often occupy the best-performing region, and restricting warmup can make otherwise competitive rates appear suboptimal.The best tested runs jointly tune learning rate and warmup (Table 1).Source: §2.3 · p. 3, Table 1 · p. 4
  • Allowing longer warmups improves attainable validation loss after tuning warmup at each peak rate and can shift the preferred learning rate.Figure 5 compares unrestricted tuning with warmup capped at 1k; capped curves end when no successful short-warmup run remains.Source: §2.3 · p. 4, Figure 5 · p. 4
  • The warmup rule can change which learning rate looks optimal at a given training horizon.The paper relates this to prior work identifying warmup duration as a source of confounding in scaling results.Source: §2.3 · p. 4

3 A MECHANISM FOR HORIZON-DEPENDENT WARMUP

The quadratic model explains horizon-dependent warmup as a tradeoff between slowing progress in well-contracting directions and reducing error that persists near the stability edge. It predicts that higher peak rates can shift the optimal duration from zero or bounded warmup to a growing fraction of the horizon.

3 A MECHANISM FOR HORIZON-DEPENDENT WARMUP

  • Near the stability edge, peak-rate contraction approaches zero while average warmup contraction approaches 2, so longer warmup can remove otherwise persistent error.For small α, the relationship reverses: warmup slows progress because g(α) ≈ α while ρ(α) ≈ 2α.Source: §3 · p. 5
  • For 0 < hs < he/2 and ηhe < 2, the long-horizon optimum is zero below ηc, bounded at ηc, and a growing fraction of T above it.The transition occurs at ηc = 2/(hs + he).Source: §3 · p. 5
  • The model also predicts that higher peak rates favor longer loss-optimal warmup and, past the stability edge, longer warmup to avoid net growth by a given horizon.Above the stability limit, the constant-rate phase amplifies error in the near-stability mode.Source: §3 · p. 5
  • The fixed-quadratic analysis isolates curvature-dependent responses to warmup and peak rates, but does not capture schedule-ordering effects that can matter in neural-network training.Finite-batch sampling and evolving curvature are identified as additional mechanisms for early–late tradeoffs.Source: §3 · p. 5

4 A LOSS LAW FOR WARMUP DURATION

The loss law combines the cost of delayed progress with the benefit of reducing persistent error to describe loss over warmup duration and training horizon. Its fitted form covers bounded, sublinear-growing, and horizon-proportional warmup regimes, and captures measured loss curves across families.

4 A LOSS LAW FOR WARMUP DURATION

  • The loss law combines a delayed-progress cost with a persistent-error term whose dependence on warmup increases near the stability edge.The delayed-progress term follows from peak-rate-equivalent updates and a spectral power-law assumption; the persistent-error term uses a power-law density of effective decay rates near zero.Source: §4 · p. 6

4.1 WHY OPTIMAL WARMUP DEPENDS ON THE HORIZON

  • For 0 < β < 1, optimal warmup grows as a sublinear power of T; the law also covers bounded or zero warmup when β ≤ 0 and warmup proportional to T when β ≥ 1.The balance between the cost and benefit of another warmup step determines the horizon scaling.Source: §4.1 · p. 6

4.2 FITTING THE MEASURED LOSS CURVES

  • The fitted functional form captures short- and long-warmup families, explains most loss variation, and selects low-regret warmups; near-fixed low-rate warmup appears through a near-zero scale.The reported fit uses a bounded multistart Huber procedure; Table 2 summarizes absolute-fit accuracy and selection quality using median R2 and RMSE and mean regret.Source: §4.2 · p. 6, Table 2 · p. 7, Figure 7 · p. 7

5 PREDICTING WARMUP AT LATER HORIZONS

The loss law is fit on shorter trajectories and used to rank candidate warmups at later horizons. The reported fits select low-regret durations, including from three short runs, and the three-run forecasts remain competitive at horizons up to one million updates.

5 PREDICTING WARMUP AT LATER HORIZONS

  • The selection procedure fits the loss law on shorter trajectories, freezes its coefficients, and chooses the available warmup with the lowest predicted loss at later horizons.With a full warmup grid this tests the law; with three short runs it forms a practical protocol.Source: §5 · p. 7, §5.1 · p. 8

5.1 FROM THE LOSS LAW TO A SELECTION RULE

  • Subtracting a reference warmup preserves the minimizing duration, allowing the difference form to fit loss differences without the asymptotic loss term.The shortest successful warmup serves as the reference.Source: §5.1 · pp. 7–8

5.2 HORIZON EXTRAPOLATION

  • Both loss-law fits select low-regret warmups within the fit window and at later horizons, with the difference fit best in every column of Table 3.Fits use checkpoints through 32k updates, and regret is evaluated at later horizons as well as within the fit window.Source: §5.2 · p. 8, Table 3 · p. 8

5.3 PRACTICAL EXTRAPOLATION FROM THREE SHORT RUNS

  • With only three separated stable warmup runs, both loss-law fits outperform practical baselines and naive extrapolation, and the absolute fit is the primary practical predictor.The difference fit has only two loss-difference curves to constrain five parameters, while the absolute fit retains all three loss curves.Source: §5.3 · pp. 8–9, Table 4 · p. 9

5.4 LONG-HORIZON EXTRAPOLATION

  • Three-run absolute-fit forecasts have the lowest loss at every tested horizon for the two higher peak rates and match the 1k baseline at the lowest rate, through one million updates.The comparison uses 250k, 500k, and one-million-update horizons; the fixed 1k warmup trains without divergence at every tested peak rate.Source: §5.4 · p. 9, Table 5 · p. 9

6 CONCLUSION

The conclusion frames warmup duration as a horizon-dependent hyperparameter whose preferred value depends on peak learning rate and training budget. A quadratic tradeoff motivates a compact loss law that can guide longer-horizon choices from three short trajectories.

6 CONCLUSION

  • In the experiments, lower peak rates often favor short or bounded warmup, while higher rates favor durations that grow with the training horizon.Warmup is neither fixed across budgets nor tied to a fixed fraction of the run.Source: §6 · p. 9
  • The quadratic model explains this shift as a balance between sacrificing early progress and reducing error that would otherwise persist later.Source: §6 · p. 9
  • The compact loss law can be fit using three short trajectories to guide warmup choices at longer horizons.The paper recommends tuning warmup jointly with peak learning rate and training budget.Source: §6 · p. 9

A.1 TRAINING SETUP, OBSERVATIONS, AND UNITS

The appendix documents the language-model training data and experimental grid, fitting and baseline procedures, evaluation metrics, and supporting theory.

A.1 TRAINING SETUP, OBSERVATIONS, AND UNITS

  • Experiments use Llama-style models pretrained on C4 English, with validation cross-entropy, 33 model-size and peak-rate families, and runs retained only when successful and W < T.The configuration tables specify model and execution settings; the batch-size ablation varies batches of 256, 512, and 1,024 sequences without rescaling peak learning rate.Source: §A.1 · p. 16, Table 7 · p. 16, Table 6 · p. 16

A.2 FITTING AND BASELINES

  • Absolute and difference loss laws are fit independently, while baselines include directly scaling observed best warmup, fixing its duration, or fixing its fraction of the horizon.The fits use bounded Huber optimization with deterministic multistart initialization; Table 9 reports the fitting settings.Source: §A.2 · p. 17, Table 9 · p. 17

A.3 CANDIDATE SELECTION AND METRICS

  • Methods select among the same successful warmup candidates, and regret is measured against the best observed candidate, averaged across horizons and families.RMSE and R2 use unmodified validation loss; reported ± values are sample standard deviations across horizon-wise family means.Source: §A.3 · p. 17

B THEORY SUPPORTING THE MAIN TEXT

  • The supporting theory develops the quadratic warmup tradeoff, extends it to warmup above the stability limit, and connects many modes to a fitted loss law and its optimal duration.A complementary discrete result states that sampling can break order invariance and change horizon-optimal warmup even with fixed population curvature.Source: B THEORY SUPPORTING THE MAIN TEXT · p. 17

B.1 QUADRATIC WARMUP TRADEOFF

The quadratic analysis models warmup as a tradeoff between modal contraction rates: it may slow modes that already contract well while improving modes near the stability edge. In the two-mode model, the optimal duration changes regime at a peak-rate threshold.

Approximation and contraction rates

  • The model uses a continuous-warmup approximation, whereas exact integer schedules can show isolated cancellations when a step lands at ηth = 1.Figure 11 defines ρ(α) as one-step log-contraction at the peak rate and g(α) as its average over linear warmup.Source: §B.1 · p. 18, Figure 11 · p. 18

When warmup helps a mode

  • Warmup changes a mode’s log-loss by ρ−g; its effect changes sign at αcross ≃ 1.27846, where g and ρ are equal.For α < 1, the average warmup contraction g is below the peak-rate contraction ρ; above αcross, g exceeds ρ.Source: §B.1 · p. 19

Two-mode warmup tradeoff

  • The two-mode loss combines a low-curvature mode whose warmup adds cost with a near-stability-edge mode whose warmup removes persistent error.The model expresses these contributions as As e^-ρsT+dsW and Ae e^-ρeT-aeW, with ds = ρs−gs > 0 and ae = ge−ρe > 0.Source: §B.1 · p. 19

Two-mode optimum and peak-rate transition

  • The two-mode optimum changes at ηc = 2/(hs+he): below it the asymptotic warmup is zero, at it the optimum is bounded, and above it W* grows proportionally to T.For η > ηc, the asymptotic fraction xη increases with η and approaches 2hs/he as η approaches 2/he.Source: §B.1 · pp. 19–20

B.2 WARMUP ABOVE THE STABILITY LIMIT

The appendix derives when warmup can prevent net growth above the stability limit, then builds a fitted loss law whose optimal duration spans no, bounded, sublinear, and proportional warmup regimes. It also shows how minibatch sampling can make warmup worse early but better at later horizons.

B.2 WARMUP ABOVE THE STABILITY LIMIT

  • For 2 < α < αmax, warmup prevents net amplification exactly when T ≥ κ(α); for α ≥ αmax, no W < T prevents it.κ increases from 0 to 1 as α goes from 2 to αmax; at α = αmax, W = T gives C = 0.Source: §B.2 · p. 20

B.3 FROM MANY MODES TO THE FITTED LOSS LAW

  • The persistent-error model assumes effective decay-rate density near zero proportional to z^(γ−1); its asymptotic power law requires mass at arbitrarily small rates.Finitely many modes can approximate the power law over an intermediate range of T.Source: §B.3 · p. 22
  • The local product approximation uniquely matches the loss and its log-τ and log-W derivatives at a reference point, with error vanishing at a fixed warmup fraction.Its exponents divide γ according to the two terms’ shares of accumulated contraction at that point.Source: §B.3 · pp. 22–23
  • As peak rate rises toward the stability edge, the local exponent shifts from progress toward warmup: sη increases toward γ while qη decreases toward zero.Source: §B.3 · p. 23

B.4 OPTIMAL DURATION UNDER THE FITTED LAW

  • For L_T, strict convexity on [0, T] gives a unique finite-horizon minimizer; when q ≤ p+1, the optimal warmup strictly increases with T wherever it is interior.Source: §B.4 · pp. 23–24
  • The exponent β maps to no warmup, bounded warmup, sublinear growth, or proportional growth as it moves across the stated regimes.For β > 1, the proportional limit is min{1, 2s/(q+s)}.Source: §B.4 · p. 24

B.5 A STOCHASTIC MECHANISM FOR HORIZON DEPENDENCE

  • In the two-mode sampling example, the exact recursion’s warmup ranking crosses at step 24, after which warmup settles at about one-quarter of the no-warmup loss.The figure reports losses relative to initial loss; its plotted example has ϵ = 0.02, ηh = 0.4, B = 2, and W = 20.Source: Figure 14 · p. 25
  • Sampling makes large steps leak more error into the slowly shrinking flat mode, so warmup pays an early sharp-mode cost but can yield a later benefit.Without sampling, no error leaks between modes and warmup is worse at every horizon in this example.Source: §B.5 · pp. 25–26

C ROBUSTNESS OF WARMUP BEHAVIOR

Across decay, optimizer, and schedule-shape tests, warmup’s horizon and peak-rate dependence largely persists. The results also show that decay usually preserves warmup rankings and that schedule shape can shift the duration tradeoff.

C.1 WARMUP ORDERING AFTER DECAY

  • Decay leaves the best warmup unchanged in 22 of 30 fork–family cases and in 5 of 6 families at the latest fork.When the best duration changes, decay favors longer warmup at intermediate forks and higher peak rates.Source: Figure 15 · p. 26, §C.1 · p. 27, Table 11 · p. 27

C.2 HORIZON AND PEAK-RATE DEPENDENCE WITH MUON

  • Muon shows the same pattern as AdamW: lower peak rates favor short warmup, while higher rates favor durations that grow with the horizon.The preferred durations and learning-rate scales differ between the optimizers.Source: §C.2 · p. 27, Figure 16 · p. 28

C.3 HORIZON DEPENDENCE UNDER DIFFERENT WARMUP SHAPES

  • Concave-quadratic warmup has a smaller progress penalty per warmup step, shifting the balance toward longer durations; linear and half-cosine have the same accumulated learning rate.Persistent-error coefficients may also depend on warmup shape.Source: §C.3 · p. 28
  • Across warmup shapes, the shift toward longer durations is clearest at the highest peak rate, while lower rates favor shorter durations.Linear and half-cosine schedules have similar qualitative trends; the concave-quadratic optimum continues moving longer at the latest horizon under the highest rate.Source: §C.3 · p. 28, Figure 18 · p. 29, Table 15 · p. 30

D ADDITIONAL SELECTION RESULTS

This section compares direct warmup extrapolation with loss-law selection and examines how fitting data and run choices affect later-horizon recommendations. It also describes the long-horizon evaluation protocol.

D.1 DIRECT WARMUP-SCALING BASELINE

  • With only three fitting runs, direct scaling rules incur several times the regret of the absolute loss-law fit, which can select durations between or beyond the fitted runs.With the full warmup grid, direct extrapolation and tuned rules are competitive because the grid supplies accurate horizon–optimum pairs.Source: §D.1 · p. 30, Table 16 · p. 30

D.2 NEAR-OPTIMAL MEASURED WARMUPS

  • Several measured warmups can lie within a loss tolerance of the best, so selectors may choose visibly different durations yet incur similar regret.The near-optimal set contains only measured warmups; gaps are not filled by interpolation, and the set expands with the tolerance.Source: §D.2 · p. 30, Table 17 · p. 31

E.1 THREE-RUN EXTRAPOLATION

  • The three-run extrapolation protocol fits the loss law to three separated warmup trajectories and freezes coefficients before evaluating later horizons.The runs share checkpoints after the longest warmup has ended.Source: §E.1 · p. 31, Table 18 · p. 31, Table 19 · p. 31

E.2 RECOMMENDATIONS BEYOND THE OBSERVED WARMUPS

  • Many frozen recommendations use durations not included in fitting, and recommendations above the observed range remain low-regret.Selections are scored against measured target-horizon trajectories at the selected durations.Source: Figure 19 · p. 31, §E.2 · p. 32

E.3 SENSITIVITY TO THE CHOICE OF FITTING RUNS

  • Broad log-space coverage performs similarly or better than the default fitting-run design, while clustering all three warmups near the long end is substantially less reliable.Random triplets remain competitive on average; the comparison also varies the shared checkpoint window when the longest warmup changes.Source: §E.3 · p. 32, Figure 20 · p. 32

E.4 SENSITIVITY TO THE NUMBER OF FITTING RUNS

  • Moving beyond very sparse trajectory coverage helps, followed by smaller, nonmonotonic changes as more fitting runs are added.The effect depends on which trajectories are added, and the same pattern appears across model sizes.Source: §E.4 · p. 33, Figure 21 · p. 33, Table 21 · p. 33

E.5 FIT HORIZON AND SHORT-RUN BUDGET

  • Longer fitting horizons do not improve final selection monotonically; error changes most when the observed prefix is short, then remains in a similar range.The comparison holds families and the final target horizon fixed while changing the amount of short-run data.Source: Figure 22 · p. 33, §E.5 · p. 34

E.6 LONG-HORIZON EXTRAPOLATION PROTOCOL

  • The long-horizon protocol fits the loss law independently for each peak rate using three runs, then minimizes the frozen absolute fit over warmup duration at each target horizon.The predicted schedule is trained from initialization, and validation loss is reported directly rather than regret to an oracle warmup.Source: §E.6 · p. 34, Table 22 · p. 34

E.7 MODEL-SCALE EXTRAPOLATION

The paper reports similar horizon dependence across some model scales when their peak learning rates differ. It also finds that exponent flexibility improves fit quality without consistently changing selection, while the fixed zero-warmup offset has little effect.

E.7 MODEL-SCALE EXTRAPOLATION

  • A larger model at a lower peak learning rate can show a preferred-warmup horizon dependence similar to a smaller model at a higher rate.A shared learning-rate multiplier improves overall agreement between model pairs, though differences remain at individual horizons.Source: §E.7 · pp. 34–35, Figure 23 · p. 34

F.1 WARMUP EXPONENT

  • The free warmup exponent improves fit quality, while selection performance remains similar across the reported comparisons.The comparison is against a fixed-s ablation at s = 1/2.Source: §F.1 · p. 35, Table 23 · p. 35

F.2 ZERO-WARMUP OFFSET

  • Across offsets from 8 to 128 updates, later-horizon regret changes by at most 0.4×10^-3, so the fits are insensitive to the fixed offset.The paper uses w0 = 32 updates to keep the law finite at zero warmup.Source: §F.2 · p. 35, Table 24 · p. 35

F.3 SHARING THE PERSISTENT EXPONENT SUM ACROSS PEAK RATES

Sharing the persistent exponent sum across peak-rate families retains most descriptive fit quality with fewer exponent parameters, but selection effects vary by model size. The accompanying gallery and related-work discussion situate the fitted loss law within broader optimization and scaling research.

F.3 SHARING THE PERSISTENT EXPONENT SUM ACROSS PEAK RATES

  • Sharing the exponent sum preserves most descriptive fit quality while reducing exponent parameters, but later warmup selection improves in some model-size groups and degrades in others.The largest departures of free sums from the shared value occur at the smallest peak rates, where some fits hit the upper bound and are not data-determined.Source: §F.3 · p. 36, Figure 25 · p. 36, Figure 24 · p. 36

G COMPLETE FIT GALLERY

  • The plots show full-data absolute fits for every analyzed family, with measured validation-loss points and model-prediction curves at displayed horizons across the cited model-size and learning-rate panels.Panels cover 60M, 100M, and 350M models at lower peak learning rates, and 60M and 100M models at higher peak learning rates.Source: G COMPLETE FIT GALLERY · p. 37, Figure 26 · p. 37, Figure 27 · p. 37

H RELATED WORK

  • Related work spans warmup mechanisms, hyperparameter scaling, schedule-aware loss models, and quadratic optimization analysis; these lines motivate the paper’s horizon-dependent loss law.The paper frames its focus as complementary to prior studies of warmup mechanisms and hyperparameter scaling.Source: H RELATED WORK · pp. 39–40

Related papers

Loading 2609.33041v1…