PaperlayerView exact revision on arXiv

Source-linked AI summary

A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules

Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Zhiyuan Liu, Maosong Sun, Kaifeng Lyu, Wenguang Chen

arXiv:2503.12811v1cs.LGcs.AIcs.CLstat.ML

TL;DR

Large-scale pretraining makes it costly to understand how learning-rate schedules shape loss, while existing scaling laws do not provide schedule-specific guidance. The paper proposes MPL, which combines cumulative learning-rate scaling with a decay-related loss correction to predict loss curves. Those predictions support schedules that outperform cosine and tuned WSD variants in the reported comparisons.

  • Problem

    Existing scaling laws do not address how to set learning rates over time, while testing schedules at large training scales is costly.

    Source: §1 · p. 1

  • Method

    MPL predicts the loss curve from cumulative learning rates and adds a correction for loss reduction associated with learning-rate decay.

    Source: §1 · p. 3

  • Results

    Optimized schedules beat cosine schedules by a margin exceeding 0.02 across model sizes and training steps, and outperform optimally tuned WSD and WSDLD variants.

    Source: §6.2 · p. 14

  • Takeaways & Limitations

    Minimizing MPL-predicted final loss can identify a schedule that outperforms the standard cosine schedule.

    Source: §1 · p. 3

    The theoretical analysis is limited to quadratic loss cases; extending it to landscapes that better capture practical deep learning remains future work.

    Source: §4.2 · p. 11

Abstract

from arXiv · show

Training large models is both resource-intensive and time-consuming, making it crucial to understand the quantitative relationship between model performance and hyperparameters. In this paper, we present an empirical law that describes how the pretraining loss of large language models evolves under different learning rate schedules, such as constant, cosine, and step decay schedules. Our proposed law takes a multi-power form, combining a power law based on the sum of learning rates and additional power laws to account for a loss reduction effect induced by learning rate decay. We extensively validate this law on various model sizes and architectures, and demonstrate that after fitting on a few learning rate schedules, the law accurately predicts the loss curves for unseen schedules of different shapes and horizons. Moreover, by minimizing the predicted final pretraining loss across learning rate schedules, we are able to find a schedule that outperforms the widely used cosine learning rate schedule. Interestingly, this automatically discovered schedule bears some resemblance to the recently proposed Warmup-Stable-Decay (WSD) schedule (Hu et al, 2024) but achieves a slightly lower final loss. We believe these results could offer valuable insights for understanding the dynamics of pretraining and designing learning rate schedules to improve efficiency.

1 INTRODUCTION

The paper addresses the lack of scaling-law guidance for how learning-rate schedules shape training loss, a costly question to study through full pretraining runs. It proposes the Multi-Power Law (MPL) to predict loss curves across schedules and use those predictions to search for improved schedules.

1 INTRODUCTION

  • Existing scaling laws do not address optimal learning-rate schedules and may generalize poorly beyond the schedule family used for fitting.This is especially costly to investigate at trillion-token training scales, where trying multiple schedules is impractical.Source: §1 · pp. 1–2
  • The schedule-aware prediction problem asks whether a formula can predict the full loss curve for a given schedule, assuming a fixed warmup and a nonincreasing post-warmup learning rate.The warmup shape and peak rate are treated as already selected; the focus is on decay schedules after warmup.Source: §1 · p. 2
  • MPL models loss using a power law of cumulative learning rates plus a correction term for additional loss reduction induced by learning-rate decay.The cumulative-rate terms alone give only a crude approximation; the decay correction accounts for schedule shape.Source: §1 · pp. 2–3
  • Fitting MPL on at most three pretraining runs predicts loss curves for unseen schedules, using entire training curves rather than only final losses.The evaluations include unseen WSD and two-stage schedules, including schedules with a different second-stage learning rate.Source: §1 · p. 3, Figure 2 · p. 3
  • Optimized schedules consistently outperform cosine schedules across model sizes and training steps, with a margin exceeding 0.02, and also beat optimally tuned WSD and WSDLD variants.The optimized schedule’s shape resembles WSD, but its decay is optimized rather than selected by grid search.Source: Figure 1 · p. 2, §1 · p. 3
  • The paper also reports that MPL emerges theoretically for quadratic losses when the Hessian and noise-covariance matrices have certain power-law structures.The theoretical analysis is specifically for quadratic loss functions.Source: §1 · p. 3

2 PRELIMINARY

This section defines the learning-rate schedules and warmup setup used in the analysis, then describes a data-scaling law and why it can mispredict intermediate losses.

2 PRELIMINARY

  • A learning-rate schedule specifies the learning rate at each training step; cosine and WSD are examples with distinct decay patterns.WSD holds the peak rate stable before a final linear or exponential decay.Source: §2 · p. 4
  • The analysis starts after a linear warmup of W steps, whose cumulative learning-rate sum is SW = 1/2ηmaxW.Post-warmup step counting begins at t = 1, and the last warmup rate η0 equals the peak rate ηmax.Source: §2 · p. 4
  • The Chinchilla Data Scaling Law is fitted to final losses for a schedule family, but applying it at intermediate steps introduces bias because the schedule has undergone less decay.The resulting intermediate loss trajectories differ from those of the full schedule.Source: §2 · p. 4

3 EMPIRICAL DERIVATION OF THE MULTI-POWER LAW

The empirical derivation treats loss as a constant-learning-rate baseline matched by cumulative learning-rate sum, plus a residual for decay; it builds that residual from intermediate loss reductions and combines the terms into the Multi-Power Law.

3.1 OUR APPROACH: LEARNING RATE SUM MATCHING

  • Matching cumulative learning-rate sums gives a constant-rate reference curve, with the residual representing loss reduction due to learning-rate decay.The reference step is the equivalent step in a constant-rate process with the same cumulative learning-rate sum.Source: §3.1 · p. 5

3.2.1 CONSTANT PROCESS LOSS APPROXIMATION

  • The constant-process loss is modeled as a power law in cumulative learning-rate sum, including the warmup contribution.The authors report validation across model sizes, training horizons, and learning rates.Source: §3.2.1 · p. 6

3.2.2 LOSS REDUCTION APPROXIMATION

  • A linear loss-reduction approximation tracks learning-rate reduction but predicts an abrupt loss drop at step changes, unlike the smoother observed decline.The finer-grained approach instead decomposes total loss reduction into a telescoping sum between adjacent auxiliary processes.Source: §3.2.2 · pp. 6–7

3.3.1 CASE 1: TWO-STAGE LEARNING RATE SCHEDULE

  • The two-stage loss-reduction curve follows a power form that fits better than an exponential form.The fitted amplitude is tied to the learning-rate drop, while the other coefficient follows a power-law pattern in the later-stage learning rate; parameter sensitivity to prefix length is weak.Source: Figure 4 · p. 7, §3.3.1 · pp. 7–8

3.3 BOTTOM-UP DERIVATION: TWO-STAGE, MULTI-STAGE, AND GENERAL SCHEDULES

  • The bottom-up derivation addresses weakly measurable intermediate reductions and dependence on the preceding learning-rate prefix by moving from two-stage to multi-stage and general schedules.The section first uses larger learning-rate reductions in simple schedules, then checks prefix effects in multi-stage schedules.Source: §3.3 · p. 7

3.3.2 CASE 2: MULTI-STAGE LEARNING RATE SCHEDULE

  • In multi-stage schedules, stage-wise loss reductions retain the power form, with parameters approximately tied to the adjacent learning-rate drop and the stage learning rate.The observed parameter patterns are similar to those in the two-stage case, and intermediate reductions are relatively insensitive to the learning-rate prefix.Source: §3.3.2 · pp. 8–9, Figure 6 · p. 9

3.3.3 CASE 3: GENERAL LEARNING RATE SCHEDULE

  • The Multi-Power Law combines a power-law approximation for the constant process with a loss-reduction term built from intermediate schedule changes.For general schedules, the paper extends the intermediate loss-reduction form inferred from two-stage and multi-stage cases.Source: §3.3.3 · p. 9

4 HOW MIGHT THE MULTI-POWER LAW ARISE?

The paper examines whether the Multi-Power Law can arise from SGD on a quadratic loss with noisy gradients. Under specified spectral assumptions and theorem conditions, the expected loss admits a related estimate, but the analysis does not yet extend beyond quadratic cases.

4.1 SETUP

  • The analysis uses SGD on a quadratic loss with Gaussian gradient noise and assumes power-law structure in the Hessian spectrum and related quantities.The paper presents this as a simplified setting that does not fully capture deep-learning complexity.Source: §4 · p. 9, §4.1 · pp. 9–10

4.2 MAIN THEOREM

  • Under Assumption 1, sufficiently small maximum learning rate, and sufficiently large cumulative learning-rate sum, the theorem estimates the expected SGD loss curve.The theoretical estimate has an exponent linked to Hessian eigenvalue decay and initial-parameter variance decay.Source: §4.2 · p. 10
  • The theoretical loss-reduction factor differs from the empirical form, and extending the analysis beyond quadratic cases remains future work.The authors suggest the quadratic loss may not capture local-loss-landscape changes associated with learning-rate changes.Source: §4.2 · p. 11

5 EMPIRICAL VALIDATION OF THE MULTI-POWER LAW

MPL predicts full loss curves for unseen learning-rate schedules using only a few training curves, and generalizes across schedule shapes and longer horizons. It also compares favorably with Chinchilla-style and Momentum Law baselines, while remaining robust across several model settings.

5 EMPIRICAL VALIDATION OF THE MULTI-POWER LAW

  • With training curves from at most three schedules, MPL predicts unseen schedules with different shapes and extended horizons.The validation uses 25M, 100M, and 400M models, training on a 24,000-step constant and cosine pair and a 16,000-step two-stage schedule.Source: §5 · p. 11, §5.1 · p. 11, §5.2 · p. 12

5.2 RESULTS

  • MPL extrapolates loss curves beyond three times the training horizon, from about 22,000 post-warmup training steps to test curves of up to 70,000.The experiments also report accurate predictions for long-horizon cyclic and random-polyline schedules despite deriving MPL for monotonic decay schedules.Source: §5.2 · p. 12

5.3 COMPARISON WITH BASELINES

  • MPL reduces final loss prediction error to less than 1/3 that of CDSLs while requiring about 1/5 the compute budget.It also fits the open-source 7B OLMo linear-schedule curve accurately.Source: §5.3 · p. 12, Figure 8 · p. 13
  • MPL consistently outperforms the Momentum Law in average and worst-case prediction accuracy across unseen schedules, including discontinuous schedules.For linear-decay WSD schedules, MPL better captures loss reduction during the decay stage.Source: §5.3 · pp. 12–13

5.4 ABLATION EXPERIMENTS

  • Ablations report high prediction accuracy across GPT-2 and OLMo, model sizes, batch sizes, and random seeds, while accuracy degrades as peak learning rate increases.For peak learning rates of 4×10^-4 and 6×10^-4, average R2 values are 0.9965 and 0.9940, respectively; across tested batch sizes R2 exceeds 0.9970.Source: §5.4 · pp. 13–14

6 THE MULTI-POWER LAW INDUCES BETTER LR SCHEDULES

The paper uses MPL’s predicted final loss as a surrogate for optimizing learning-rate schedules. The resulting schedules outperform cosine and tuned WSD variants, with a stable-then-decay pattern and a power-shaped decay.

6.1 METHOD

  • MPL’s predicted final loss is optimized as a surrogate over the schedule vector, subject to monotonicity constraints.The schedule is treated as a T-dimensional vector and optimized using MPL-estimated parameters.Source: §6.1 · p. 14

6.2 RESULTS

  • The optimized schedule resembles WSD’s stable-then-decay pattern despite no WSD-like schedule being included in the training set.Its post-warmup phases are a constant peak learning rate followed by decay.Source: §6.2 · p. 14
  • Optimized schedules consistently outperform cosine schedules by a margin exceeding 0.02 across tested model sizes and training steps.The authors also report overall downstream-task improvements over cosine schedules.Source: §6.2 · p. 14, Figure 9 · p. 15
  • The optimized schedules achieve smaller final loss than optimally tuned WSD and WSDLD schedules, even when their decay ratios are tuned.For a 400M model, the optimized 24,000-step schedule’s decay step is close to the grid-searched optimum of about 6,000 steps for those WSD schedules.Source: §6.2 · p. 14
  • The optimized schedule’s decay roughly follows ηt ≈ ηmax ·(1−τ)^1.5, motivating WSDSC, which outperforms WSD but remains behind the optimized schedule.The optimized schedule decays below 1/20 of peak LR, whereas WSD schedules decay to 1/10.Source: §6.2 · p. 15, Figure 9 · p. 15

7 RELATED WORK

Prior scaling-law work models loss across factors such as model size and data but often neglects learning-rate schedules. This paper situates MPL as a schedule-aware alternative that predicts full curves and supports schedule optimization.

Scaling Laws and Loss Curve Prediction

  • Earlier empirical scaling-law research established power-law trends and refined their fitting, evaluation, and applicability across training settings.The related work also describes theoretical accounts based on data manifolds, memorization distributions, and linear or simple neural-network models.Source: §7 · pp. 15–16
  • Most existing scaling-law work neglects learning-rate effects, while schedule-specific final-loss laws can require more than 10 full training runs.MPL instead predicts full loss curves across schedules using fewer than three training curves.Source: §7 · p. 15
  • MPL differs from the Momentum Law by modeling power-law rather than exponential loss reduction, and from Xie et al. (2024) by predicting full curves across varied schedules.The cited account says the Momentum Law yields collapsed schedules, while Xie et al.’s method is limited to final-loss prediction and manually defined phases.Source: §7 · pp. 15–16

Optimal Learning Rate Schedule

  • Prior learning-rate work includes warmup, cyclical and adaptive schedules, batch-size scaling, and WSD; schedule-optimality studies often impose specific constraints.Unlike those constrained approaches, this work integrates schedules into scaling laws to enable gradient-based optimization over all possible schedules.Source: §7 · pp. 16–17

8 CONCLUSIONS

The conclusion presents MPL as a schedule-aware law for predicting full loss curves and optimizing schedules, while noting prediction deviations under some settings.

8 CONCLUSIONS

  • MPL predicts entire loss curves with fewer training runs than existing scaling laws and supports optimization of schedules that outperform commonly used ones.The conclusion notes slight deviations from actual curves, especially for long horizons and high peak learning rates.Source: §8 · p. 17
  • The derivation assumes constant β across learning-rate scales, prefix-independent intermediate loss reduction, and no warmup-rate variation.These simplifications are identified as likely sources of prediction deviations.Source: §8 · p. 17
  • Future work aims to strengthen MPL’s theoretical foundation, study varying peak learning rates and other hyperparameters, and improve prediction accuracy and generalizability..Source: §8 · p. 17

A FORMULA COMPONENT ABLATION

The ablations test which parts of the Multi-Power Law are needed for accurate fitting and schedule optimization. They show that the loss-reduction component and power-based decay formulation matter, while simpler alternatives lose accuracy or produce unsuitable schedules.

A FORMULA COMPONENT ABLATION

  • Removing the loss-reduction term causes significant fitting degradation, showing that matching cumulative learning-rate sums alone is insufficient.The One-Power Law sets LD(t)=0, omitting the correction term.Source: A FORMULA COMPONENT ABLATION · p. 24
  • LLDL fits better than OPL but worse than full MPL, and its schedule optimization collapses to a final-step learning-rate drop.Figure 10 reports a mean R2 of 0.9980 for the near-linear relationship between loss reduction and learning-rate reduction across schedules.Source: A FORMULA COMPONENT ABLATION · p. 24, Figure 10 · p. 25
  • Setting γ=0 causes a slight fitting drop, so the final MPL retains γ for improved accuracy at minimal added computational cost.Source: A FORMULA COMPONENT ABLATION · p. 24
  • The step-based SPL reduces prediction accuracy and implies continued loss reduction even after the learning rate reaches zero.Source: A FORMULA COMPONENT ABLATION · p. 24
  • Exponential loss-reduction approximations underperform the power-based MPL; MTL is a step-based exponential variant and also lags behind MEL.The exponential form is inconsistent with the observed power form, while step-based approximations have limitations.Source: A FORMULA COMPONENT ABLATION · pp. 24–25

C.1 CONSTANT PROCESS LOSS APPROXIMATION (SECTION 3.1)

  • Constant-schedule fits reached MSE 1.55 × 10−5 and R2 0.9976 in peak-LR ablations, and average MSE 8.04 × 10−5 and R2 0.9947 on 72,000-step curves.The appendix also documents model series, parameter-count conventions, and training hyperparameters in Tables 4–5.Source: B EXPERIMENT SETTING · p. 25, §C.1 · p. 26, Table 4 · p. 26

C.2 TWO-STAGE EXPERIMENTS (SECTION 3.3.1)

  • Fixing β=0.4 retained feasible fitting margins across ηA and ηB ablations, supporting derivation of the two-stage coefficient patterns.Figure 12 presents the refitted curves and ground truths.Source: §C.2 · p. 27, Figure 12 · p. 27

C.3 MULTI-STAGE EXPERIMENTS (SECTION 3.3.2)

The multi-stage experiments test how loss reduction depends on the preceding learning-rate sequence and examine MPL predictions across schedules and horizons. The reported setup includes multi-stage runs and held-out schedules, while the accompanying figure shows GPT-2 schedule comparisons and a longer-horizon prediction.

C.3 MULTI-STAGE EXPERIMENTS (SECTION 3.3.2)

  • Figure 13 compares 24,000-step GPT-2 cosine, WSD, and WSDSC loss curves and shows a prediction for a 72,000-step cosine curve.Source: Figure 13 · p. 28
  • Across a nine-stage schedule, fitted loss-reduction parameters vary with adjacent-stage learning-rate changes: B̃(i) with η(i−1)−η(i), and C̃(i) with η(i).The schedule begins with 8,000 steps at 3 × 10−4, followed by eight 90-step stages.Source: §C.3 · pp. 27–28

D.1 TRAINING SET AND TEST SET

  • The training set uses 24,000-step constant and cosine schedules plus a 16,000-step two-stage schedule; tests include longer horizons and unseen schedules.The test set includes 72,000-step constant and cosine curves, 24,000-step WSD and WSDLD curves, and two-stage schedules with different ηB/ηA ratios.Source: §D.1 · p. 28, Table 6 · p. 28

D.2 FITTING THE LAW

The section fits the loss law by minimizing Huber loss over validation steps, using the area under an interpolated learning-rate curve as a proxy for cumulative learning rate. It also compares the fit with alternative laws and examines fitting choices and schedule configurations.

D.2 FITTING THE LAW

  • The fitting objective sums Huber loss across validation steps, using interpolated learning-rate area as a surrogate for the learning-rate sum.This reduces computation to validation-step numbers, learning rates, and losses.Source: §D.2 · pp. 28–29
  • MPL is fit with Adam in two optimization stages, with lower learning rates in the second stage and the lowest training-loss result retained.Each optimization runs for 5×10^4 steps.Source: §D.2 · p. 29, Table 7 · p. 29
  • During decay on WSDLD, MPL fits more accurately and matches loss-curve curvature, whereas MTL predicts a counterfactual concave curve.MTL fits the stable stage well.Source: §D.2 · p. 29, Figure 14 · p. 29
  • MPL’s prediction error is 0.007 versus 0.024 for CDSL, while requiring 4.8 times less compute.CDSL uses six cosine curves; MPL uses two 24000-step curves.Source: §D.2 · pp. 29–30
  • The optimization-method discussion reports that L-BFGS is sensitive to initialization, while Adam lacks theoretical convergence guarantees.Parameter constraints and grid-search initialization can mitigate some L-BFGS issues; greater robustness and stability remain future work.Source: §D.2 · p. 30

E DETAILS OF LR SCHEDULE OPTIMIZATION (SECTION 6)

The section optimizes schedules against the fitted law using learning-rate reductions as optimization variables and validates the resulting schedules across horizons, model sizes, and schedule variants. The optimized schedules outperform cosine and tested WSD variants in the reported comparisons.

E DETAILS OF LR SCHEDULE OPTIMIZATION (SECTION 6)

  • The learning-rate schedule is reparameterized by its stepwise reductions, with clipping constraints used to stabilize optimization.The reduction variables are restricted to [0, η0], and set to zero when the learning rate is at or below 10^-10.Source: E DETAILS OF LR SCHEDULE OPTIMIZATION… · p. 34
  • The optimized schedules outperform cosine schedules for models from 25M to 400M at 72000 steps and for a 1B model.For 1B models, the MPL-induced schedule also improves the average downstream score by 1.03 over cosine across the evaluated tasks.Source: E DETAILS OF LR SCHEDULE OPTIMIZATION… · pp. 34–35, Figure 19 · p. 35
  • For models from 25M to 400M, schedules optimized for 72,000-step training using the training-set MPL fit exhibited a WSD-like stable phase followed by decay and outperformed cosine schedules.Source: E DETAILS OF LR SCHEDULE OPTIMIZATION… · p. 34
  • Optimized schedules outperform WSD variants with near-zero ending learning rates, and lower ending rates do not consistently improve final loss.The passage notes that zero-ending WSD exceeds baseline loss in one example.Source: E DETAILS OF LR SCHEDULE OPTIMIZATION… · p. 35, Figure 20 · p. 35
  • The decay phase of optimized schedules is analyzed across Llama2 models from 25M to 400M, and symbolic regression is applied to averaged normalized learning rates.The resulting WSDSC schedule beats cosine on 1B Llama2 and beats standard WSD and cosine on GPT, but falls short of the MPL-optimized schedule on 1B Llama2.Source: E DETAILS OF LR SCHEDULE OPTIMIZATION… · pp. 35–36

F OPTIMAL LEARNING RATE SCHEDULE FOR MOMENTUM LAW

The section contrasts the Momentum Law’s mathematically characterized optimal schedules with the empirically effective WSD-like schedules induced by MPL. It also outlines the proof strategy for the paper’s expected-loss results on quadratic losses.

F OPTIMAL LEARNING RATE SCHEDULE FOR MOMENTUM LAW

  • The Momentum Law’s optimal schedules remain at peak learning rate and then drop directly to zero in no more than two steps, unlike the empirically effective WSD-like schedules induced by MPL.The Momentum Law’s optimizer is characterized by a convex optimization problem.Source: F OPTIMAL LEARNING RATE SCHEDULE FOR… · p. 37

G PROOF OF THEOREM 1

  • The proof of Theorem 1 first derives expected loss after t steps for quadratic losses, then takes expectations over noise covariance, curvature, and initialization.The derivation’s first part establishes a theorem characterizing expected loss for SGD on quadratic losses.Source: G PROOF OF THEOREM 1 · p. 40, §G.1 · p. 40

G.1 GENERAL THEOREM FOR ALL QUADRATIC LOSS FUNCTIONS

  • The general quadratic-loss proof decomposes expected terminal loss into a telescoping sum of auxiliary expectations.The proof bounds the terms under the condition ηmax ≤ 1/λmax.Source: §G.1 · pp. 40, 41, 43

G.2 PROOF FOR THEOREM 1

The proof of Theorem 1 takes expectations over the curvature, noise covariance, and initialization under the stated distributional assumptions. It derives a bound using incomplete-gamma expressions and relates the exact function G(x) to an approximate power form.

G.2 PROOF FOR THEOREM 1

  • The proof assumes specified distributions for curvature and noise, with exponents defined from their parameters.It sets α := 2 − ν − κ and β := 1 − ν − ρ.Source: §G.2 · p. 45
  • The proof derives the result first for the terminal step, then applies it to schedules truncated at each earlier step.It takes expectations over curvature, noise covariance, and initialization to obtain a bound from the estimate in Theorem 3.Source: §G.2 · p. 45
  • The proof uses incomplete-gamma expressions to bound terms, and reports that the gap between G(x) and its approximate power form converges to zero as x grows.The figure comparison uses the approximate form with parameter C = 2.Source: §G.2 · p. 46, Figure 21 · p. 47

Related papers

Loading 2503.12811v1…