PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Law with Learning Rate Annealing

Howe Tissue, Venus Wang, Lu Wang

arXiv:2408.11029v2cs.CLcs.AIcs.LG

TL;DR

Typical scaling laws focus on final training loss and do not model loss changes during LR annealing. This paper formulates a full-curve scaling law using LR area and an annealing term; it reports accurate prediction across unseen schedules and fitting compute below 1% of a typical scaling-law comparison. The authors place its scope primarily in pretraining and identify post-training extensions as future work.

  • Problem

    Typical scaling laws fit only final training loss, leaving intermediate losses and local drops from LR annealing unmodeled.

    Source: §1 · p. 1

  • Method

    The paper models loss at each step with a power-law term based on forward LR area S1 and an additional annealing term based on S2.

    Source: §1 · p. 2

  • Results

    The formulation accurately predicts full loss curves across unseen learning-rate schedulers and requires less than 1% of the compute to fit compared with typical scaling laws.

    Source: §1 · p. 3, §8 · p. 20

  • Takeaways & Limitations

    The formulation offers a way to study LLM training dynamics and select training recipes, with demonstrated application primarily to pretraining and LR re-warmup in continual pretraining.

    Source: §1 · p. 3, §7.3 · p. 19

    The work focuses primarily on pretraining; extending the equation to post-training, including distribution shift, data mixtures, alignment, and downstream evaluations, remains future work.

    Source: §7.3 · p. 19

Abstract

from arXiv · show

We find that the cross-entropy loss curves of neural language models empirically adhere to a scaling law with learning rate (LR) annealing over training steps: $$L(s) = L_0 + A\cdot S_1^{-\alpha} - C\cdot S_2,$$ where $L(s)$ is the validation loss at step $s$, $S_1$ is the area under the LR curve, $S_2$ is the LR annealing area, and $L_0$, $A$, $C$, $\alpha$ are constant parameters. This formulation takes into account two factors: (1) power-law scaling over data size, and (2) the additional loss reduction during LR annealing. Therefore, this formulation can describe the full loss curve at each step, rather than the single loss point at the end of training. Applying the scaling law with LR annealing and fitting only one or two training curves, we can accurately predict the loss at any given step across any learning rate scheduler (LRS). This approach significantly reduces computational cost in formulating scaling laws while providing more accuracy and expressiveness for training dynamics. Extensive experiments demonstrate that our findings hold across a range of hyper-parameters and model architectures, and our equation can extend to scaling effect of model sizes. Moreover, our formulation provides accurate theoretical verification and explanation for empirical results observed in numerous previous studies, particularly those focusing on LR schedule and annealing. We believe that this work is promising to enhance the understanding of LLM training dynamics while greatly democratizing scaling laws, and it can guide researchers in refining training strategies (e.g. critical LRS) for further LLMs.

1 INTRODUCTION

Typical scaling laws fit only end-of-training loss, missing intermediate loss changes during learning-rate annealing. The paper models the full loss curve using forward LR area and an annealing-area term, and reports accurate prediction across unseen schedules with substantially lower fitting compute than typical scaling laws.

1 INTRODUCTION

  • Earlier scaling laws fit only final loss, leaving intermediate losses and local reductions from LR annealing unmodeled and underusing training-run compute.A full-curve formulation can predict future training dynamics and offer insight into the learning process.Source: §1 · p. 1
  • The formulation models validation loss at each training step as a power-law term in S1 plus an annealing-related reduction term, −C·S2.S1 is the forward area under the LR curve; S2 represents the LR annealing area.Source: §1 · p. 2
  • The proposed equation fits full loss curves across varied hyper-parameters and architectures, extends to model-size scaling, and explains findings on LR schedules and annealing.The authors also report accurate fits to open-sourced model loss curves.Source: §1 · p. 3
  • The approach is more general than typical scaling laws such as Chinchilla and requires less than 1% of their compute to fit.The authors present this as lowering the compute burden of developing scaling laws.Source: §1 · p. 3

2 PRELIMINARY

The preliminary section reviews scaling laws that generally fit final training loss, then introduces learning-rate annealing as a training dynamic associated with additional loss changes. It motivates modeling loss across the full training curve rather than only at its endpoint.

2.1 SCALING LAWS

  • Traditional scaling laws model validation loss with power-law terms for dataset and model size, but fit only the final training loss.Obtaining a loss value for a new data size requires another training run with the same LR schedule.Source: §2.1 · p. 3

2.2 LEARNING RATE ANNEALING

  • Learning-rate annealing progressively lowers LR; cosine schedules decay across training, while WSD schedules hold LR constant for most steps and anneal near the end.WSD annealing is described as occurring in the final 10%–20% of steps.Source: §2.2 · pp. 3, 5
  • Prior work qualitatively associated stronger LR decreases with sharper validation-loss drops, while this paper aims to formulate those changes quantitatively.Figure 4 reports similar shapes for LR, gradient-norm, and validation-loss curves across schedules.Source: §2.2 · p. 5, Figure 4 · p. 5

3 THEORY

The section develops a full-curve scaling law by separating power-law progress from loss changes associated with LR annealing. Experiments test the formulation across schedulers and settings and extend it to model-size scaling.

3.1 SIMILARITY BETWEEN LEARNING RATE, GRADIENT NORM, AND LOSS

  • Across learning-rate schedules, LR, gradient-norm, and validation-loss curves have similar shapes; under constant LR, loss follows a power law over training steps.The authors use these observations to motivate an additional LR-related term for local loss changes during annealing.Source: §3.1 · p. 5

3.2 LR ANNEALING MOMENTUM

  • Validation-loss turning points lag behind LR turning points; the delay grows with steeper LR changes and is longer for decreasing-then-increasing schedules than the reverse.The authors model this historical effect as annealing momentum, defining S2 from LR changes with a decay factor λ that retains past information.Source: §3.2 · pp. 6–7

3.3 FINAL FORMULATION

  • The final formulation models validation loss at each training step using power-law progress in S1 and an additional loss-reduction term based on S2.It uses one extra parameter, C, relative to the Chinchilla scaling law and supports fitting simpler, shorter schedules to predict more complex, longer ones.Source: §3.3 · p. 7

3.4 EXPERIMENTS

  • Using 20K-step constant and cosine schedules for fitting, the equation achieved R² > 0.999; it predicted four unseen 60K-step schedules with mean error as low as 0.2%.It also accurately predicted schedules with multiple re-warmup stages, which were absent from the fitting curves.Source: §3.4 · p. 8

3.5 EXTENSION TO MODEL SIZE SCALING

  • The loss drop from LR annealing increases with both annealing steps and model size, motivating an S2 ∝ N^γ term in the model-size extension.The extension incorporates model size alongside the original annealing term.Source: §3.5 · p. 9
  • The model-size-extended law achieved R² > 0.998 when fitting and predicting full loss curves for an unseen larger model.Additional experiments across different setups also supported the robustness of the N-extended formulation.Source: §3.5 · p. 9

4 TAKEAWAYS: EXPERIMENTAL FINDINGS VERIFICATION AND EXPLANATION

The section uses the scaling law to interpret how forward area S1 and annealing area S2 balance across learning-rate schedules. It applies that account to explain observed loss behavior and identify scheduler settings associated with lower final loss.

4.1 IT VERIFIES AND EXPLAINS WHY LOSS DROPS MORE SHARPLY WHEN LR ANNEALS.

  • During WSD annealing, the negative S2 term affects loss more than the S1 term, explaining the sharper validation-loss drop than under constant LR.Figure 8 distinguishes the loss drops attributed to S1 and S2 during annealing.Source: §4.1 · pp. 9–10, Figure 8 · p. 10

4.2 DETERMINING COSINE CYCLE LENGTH AND MINIMUM LR IN COSINE LRS.

  • For cosine LR, setting cycle length T equal to total steps and minimum LR to 0 achieves the predicted optimal final loss.T greater than total steps leaves annealing incomplete; T smaller leads to early annealing and smaller S1.Source: §4.2 · p. 10, Figure 9 · p. 10

4.3 DETERMINING COSINE CYCLE LENGTH AND MINIMUM LR IN COSINE LRS.

  • With fewer training steps, constant LR can yield lower final loss than cosine, but cosine performs better with more steps because its scaling slope is larger.The explanation shifts from S1 dominating at small steps to S2 dominating later, making annealing more valuable.Source: §4.3 · pp. 10–11, Figure 10 · p. 11

4.4 IT VERIFIES AND EXPLAINS WSD AND MULTI-STEP COSINE LRS HAVE MATCHED OR EVEN LOWER LOSS THAN COSINE LRS.

  • WSD and multi-step cosine achieve lower validation loss overall because their slightly smaller S2 is outweighed by a substantially larger S1 than cosine LR.The comparison analyzes the predicted loss-drop contributions of the two terms.Source: §4.4 · pp. 11–12, Figure 11 · p. 11

4.5 DETERMINING OPTIMAL ANNEALING RATIO OF WSD LRS.

  • A moderate WSD annealing ratio, typically 10%–20%, balances S1 and S2 and is predicted to minimize validation loss across total training steps.The equation can predict the optimal ratio without large-scale experiments.Source: §4.5 · p. 12, Figure 12 · p. 12

4.6 DETERMINING THE OPTIMAL ANNEALING FUNCTION IN WSD LRS.

  • At 10% annealing, 1-sqrt yields lower final loss than cosine; at 50%, cosine is lower, so the preferred WSD annealing function depends on the ratio.The experiments align with these predictions; the account attributes the reversal to the changing balance between S1 and S2.Source: §4.6 · pp. 12–13, Figure 13 · p. 13

4.7 IT VERIFIES AND EXPLAINS THAT IN CONTINUAL PRE-TRAINING, A HIGHER MAX LEARNING RATE DURING RE-WARMUP LEADS TO A HIGHER INITIAL PEAK IN THE LOSS, FOLLOWED BY A MORE RAPID DECREASE.

  • In continual pre-training, a higher re-warmup maximum LR raises the initial loss peak and accelerates subsequent loss decline, as the predicted curves reproduce prior observations.The analysis assumes no data-distribution shift between training stages.Source: §4.7 · p. 13, Figure 14 · p. 14

4.8 IT VERIFIES AND EXPLAINS THAT IN CONTINUAL PRE-TRAINING, THE STEPS OF RE-WARMUP HAVE LITTLE IMPACT ON THE FINAL LOSS.

  • Longer re-warmup smooths the loss-curve transition but has little effect on final validation loss because final S1 and S2 remain similar across re-warmup lengths.The re-warmup period is a small part of the full training run.Source: §4.8 · p. 14

5 COMPARISON WITH CHINCHILLA SCALING LAW

The section shows that the proposed full-curve scaling law is compatible with Chinchilla-style endpoint scaling, while requiring fewer training runs to fit and predict losses.

5.1 REDUCTION TO CHINCHILLA SCALING LAW

  • Chinchilla scaling law fits the final losses predicted by the proposed equation well, supporting its interpretation as a generalization reducible to Chinchilla scaling.The comparison fits Chinchilla scaling to synthetic endpoint losses generated for cosine and WSD schedules.Source: §5.1 · pp. 14–15, Table 1 · p. 15, Figure 15 · p. 15
  • The synthetic parameter sampling includes tuples that the authors acknowledge may be unrealistic or unlikely in real experiments.This caveat applies to the generated fitting setups used in the endpoint comparison.Source: §5.2 · p. 15, §5.1 · p. 15

5.2 SCALING LAW FITTING AND PREDICTION DEMOCRATIZATION

  • The proposed equation uses less than 1% of Chinchilla scaling law’s computational cost and predicts full loss curves for unseen schedules.It can be fit using one or two moderate-length training curves rather than many separate endpoint runs.Source: §5.2 · pp. 15–16, Table 2 · p. 16

6 DISCUSSION

The discussion examines how the decay factor affects annealing predictions, considers explanations for delayed loss responses, and tests alternative scaling-law forms.

6.1 THE IMPACT OF DECAY FACTOR λ

  • A larger decay factor, such as λ = 0.999 rather than 0.99, tends to lengthen delay steps and favor higher annealing ratios; the authors recommend testing λ against empirical data.They interpret λ as setting the balance between S1 and S2, with λ = 0.999 retaining information from approximately 1,000 previous steps.Source: §6.1 · pp. 16–17, Figure 16 · p. 16, Figure 17 · p. 17

6.2 POSSIBLE ROOT REASONS OF DELAY PHENOMENON IN LEARNING RATE ANNEALING

  • The authors find that AdamW momentum has little to do with the delay phenomenon, while S1 affects delay length but is not its root cause; they leave the underlying reason unresolved.Their analysis describes possible cases but says it has not yet identified the root reasons.Source: §6.2 · pp. 17–18, Figure 18 · p. 18

6.3 OTHER POSSIBLE SCALING LAW FORMATS WITH LR ANNEALING

  • Adding an LR-weighted S2 term produced similar fits with ε close to zero, while adding a power ζ yielded only marginal improvement and creates complications when S2 is negative.The authors therefore retain the original form in their experiments.Source: §6.3 · pp. 18–19, Figure 19 · p. 19

7 FUTURE WORKS

The authors consider the equation’s format practical across scenarios and describe applications to predict training dynamics and guide recipe selection. Their work focuses primarily on LLM pretraining, with post-training extensions left for future research.

7.1 TUNING SCALING LAW FORMAT

  • The final equation adds one parameter and is reported to fit and predict well across scenarios, including future training steps and model sizes.The authors describe the format as the best tested so far and invite further research on scaling laws with LR annealing.Source: §7.1 · p. 19

7.2 MORE APPLICATIONS VIA OUR SCALING LAW

  • The authors propose using the equation to predict future training dynamics without experiments and to analyze and select recipes for specific scenarios.They describe these applications as presented in Section 4.Source: §7.2 · p. 19

7.3 EXTENSION TO POST-TRAINING

  • The study focuses primarily on LLM pretraining and applies the equation to guide LR re-warmup in continual pretraining.Extensions to post-training topics such as data shifts, mixtures, alignment, and downstream evaluations are left for future work.Source: §7.3 · p. 19

8 CONCLUSION

The paper concludes that neural language model loss curves empirically follow a scaling law with LR annealing over training steps. It reports prediction across unseen schedulers, extension to model-size scaling, and accuracy across varied setups.

8 CONCLUSION

  • The proposed equation is reported to predict full loss curves across unseen learning rate schedulers.The paper says the formulation captures scaling with learning rate annealing over training steps.Source: §8 · pp. 19–20
  • The formulation can be extended to capture the scaling effect of model sizes.The conclusion includes this among the equation’s demonstrated capabilities.Source: §8 · p. 20
  • Experiments show that the scaling law holds across various experimental setups and offers insights into LLM training dynamics.The authors also report that it explains phenomena observed in previous studies.Source: §8 · p. 20

A IMPACT OF WARMUP STEPS

In experiments using cosine scheduling and 20K total steps, 500 warmup steps accelerated convergence and achieved lower validation loss than 100 or no warmup. The authors use 500 warmup steps for scratch training and distinguish it from continual-training re-warmup.

A IMPACT OF WARMUP STEPS

  • 500 warmup steps accelerated convergence and achieved the lowest validation loss compared with 100 or no warmup.Figure 20 shows validation loss curves for different warmup steps using cosine scheduling with 20K total steps.Source: A IMPACT OF WARMUP STEPS · p. 23, Figure 20 · p. 23
  • The authors use 500 warmup steps in their main experiments when training models from scratch.They distinguish this warmup from LR re-warmup in continual training.Source: A IMPACT OF WARMUP STEPS · p. 23
  • Continual-training re-warmup is treated as a special learning rate schedule with S2 < 0 in the equation.The passage contrasts this with warmup for randomly initialized models.Source: A IMPACT OF WARMUP STEPS · p. 23

B EXPERIMENTAL SETUPS

The experiments use multiple setups to test the equation across different conditions, with most experiments using Setting A. The authors also fit loss curves from BLOOM and OLMo despite their substantially different training settings.

B EXPERIMENTAL SETUPS

  • Most experiments use Setting A, while the paper summarizes all experimental setups in Table 3.The passage does not specify the individual settings in the supplied text.Source: B EXPERIMENTAL SETUPS · p. 23
  • The equation was successfully fit to loss curves from BLOOM and OLMo.Their experimental settings differed significantly from the authors’ settings.Source: B EXPERIMENTAL SETUPS · p. 23
  • The paper directs readers to the respective model papers for details of the BLOOM and OLMo experimental setups.Source: B EXPERIMENTAL SETUPS · p. 23

C.1 ANOTHER SET OF TRAINING HYPER-PARAMETERS

The scaling law predicts full loss curves across varied experimental settings, including MoE models, much longer training runs, open-source models, and an extension to model size.

C.1 ANOTHER SET OF TRAINING HYPER-PARAMETERS

  • Across a different hyper-parameter setting, the scaling law with LR annealing continued to fit and predict loss curves well.The prediction errors across unseen learning-rate schedulers were as low as ∼0.2%.Source: §C.1 · p. 23, Table 3 · p. 24, Figure 22 · p. 25

C.2 EXPERIMENTS ON THE MIXTURE OF EXPERTS ARCHITECTURE

  • The equation still works well on the Mixture of Experts architecture when the learning-rate schedule and total steps are changed to 60K WSD with 10K annealing steps during fitting, testing effectiveness under varying circumstances.Source: §C.2 · pp. 25–26

C.3 SCALING UP FOR MUCH LONGER STEPS

  • The fitted equation accurately predicted the WSD annealing-stage loss curve after 10x longer steps, using fits from 30K-step constant and cosine runs.The prediction targeted a 350K-step run with a 1.7-billion-parameter model trained on 1,400 billion tokens.Source: §C.3 · p. 26

C.4 OPEN-SOURCED MODELS

  • The equation fit loss curves from open-source BLOOM-176B and OLMo-1B models, including OLMo training over 740K steps and 2000B tokens.Both models use cosine learning-rate schedules.Source: §C.4 · p. 26

C.5 OUR N-EXTENDED SCALING LAW ON ANOTHER EXPERIMENTS SETUPS

  • The N-extended scaling law with LR annealing also worked across different experimental settings when modeling the effect of model size.The reported fitting used cosine learning rates over 143K steps.Source: §C.5 · p. 26

D WSD SCHEDULER AND ANNEALING FUNCTIONS

The section describes WSD learning-rate schedules and annealing functions, then examines loss-curve fitting and prediction across schedulers, model types, and sizes. The reported results include accurate predictions for unseen WSD schedules and longer runs.

WSD scheduler and annealing functions

  • WSD uses warmup, stable, and decay stages; prior work reports that 1-sqrt annealing with a moderate ratio such as 20% can further decrease final loss.The section also identifies 1-square annealing as a baseline.Source: D WSD SCHEDULER AND ANNEALING FUNCTIONS · p. 26
  • For MoE models, the fitted equation predicted full loss curves for learning-rate schedules (LRS) unseen during fitting, with mean prediction errors across different LRS as low as ∼0.2%.Source: Figure 24 · p. 27
  • For unseen WSD schedules, the fitted equation accurately predicted the annealing-stage loss curve after 10x longer steps.Figure 25 fits constant and cosine schedules over 30K steps; Figure 26 predicts the WSD schedule with a 10% cosine annealing ratio.Source: Figure 26 · p. 28, Figure 25 · p. 28
  • The equation fits open-source BLOOM and OLMo loss curves, while its N-extended form fits curves across different model sizes.Both open-source models use cosine schedules; the size-extension experiment uses cosine over 143K steps.Source: Figure 27 · p. 28, Figure 28 · p. 29

Related papers

Loading 2408.11029v2…