Scaling and Transferability of Annealing Strategies in Large Language Model Training

AI summary

Overview

Research area: Large language model pre-training — learning rate scheduling, training-dynamics scaling laws, and hyperparameter transfer.

Technical level: Advanced. The paper assumes familiarity with AdamW optimization, learning-rate schedulers (cosine and WSD), Mixture-of-Experts architectures, and scaling-law parameter fitting.

Scope in one sentence: The paper proposes an improved "forward-momentum" scaling law for predicting LLM training loss curves and shows that the optimal annealing ratio in a Warmup-Steady-Decay (WSD) schedule can be transferred across model sizes, maximum learning rates, training steps, and datasets.

What This Paper Is About

Standard scaling laws predict final LLM loss from model size and token count, but they cannot explain why the same model trained on the same number of tokens produces different loss trajectories under different batch sizes or learning-rate schedules. The authors study what governs the shape of the training loss curve beyond tokens and model size, and whether the best length for the learning-rate decay (annealing) phase can be predicted rather than found by exhaustive hyperparameter search. Their goal is a practical, model-agnostic recipe for choosing annealing strategies that carries over from small models to large ones.

Key Contributions

1. A model-agnostic predictive model. The authors develop a scaling law combining a cumulative "forward" term (the integral of the learning rate over steps), an annealing momentum term (an Adam-style accumulated momentum), and a model-size term. It is validated on both Dense and MoE models and shown to predict the effect of learning-rate schedules.

2. Analysis of factors influencing annealing dynamics. They examine how batch size, maximum learning rate, model architecture, and model size affect annealing, and argue that training steps are a more reliable tracker of the loss curve than token count once batch size exceeds the optimal batch size.

3. Transferability of annealing strategy. They report that the optimal annealing ratio follows power-law relationships with maximum learning rate and with total training steps, and that it transfers across model sizes and across training/validation datasets.

4. A refinement of prior fitting practice. They replace multiplicative accumulation of momentum (used by Tissue et al. 2024) with integrals and an Adam-style momentum, which they argue is more robust to batch size variation, irregular step counts, and unstable training dynamics.

Main Findings

• Steps beat tokens as a progress metric (within a batch-size range). When batch sizes are 256, 512, and 1024 (sequence length 8192), loss curves look different when plotted against tokens but converge when plotted against training steps; the standard deviation of loss values for these three batch sizes is approximately 0.00847. The authors note that batch sizes should not be excessively large, citing the batch size 2048 case.

• Optimal batch size scales as a power law with loss. Using B_opt = λ_B / L^{α_B}, they estimate the optimal batch size for the 500M Dense model at about 293 (sequence length 8192) and for the 700M MoE model at about 366 (sequence length 8192). Fitted coefficients: MoE α_B ≈ 3.430 × 10⁰ and λ_B ≈ 4.390 × 10⁷; Dense α_B ≈ 2.710 × 10⁰ and λ_B ≈ 2.711 × 10⁷.

• Loss curves fit across model sizes. With batch size 256 and a cosine scheduler, Dense models (50M, 500M, 1B) at maximum learning rate 5e-5 and MoE models (100M, 700M, 1.5B) at maximum learning rate 2e-4 both fit with average MAPE below 2%. Loss shows power-law dependence on model size.

• Fits transfer across schedulers. A model fitted on the cosine scheduler's loss curve can predict the WSD curve and vice versa. MAPE (± std over 3 seeds): 50M — 0.456 ± 0.061 (cosine) and 0.720 ± 0.132 (WSD); 100M — 0.232 ± 0.020 and 0.410 ± 0.019; 500M — 0.453 ± 0.038 and 0.798 ± 0.014.

• Optimal annealing ratio scales with maximum learning rate. R_opt = λ_η · η_max^{α_η}. For Dense models λ_η ≈ 5.996 × 10³, α_η ≈ 7.09 × 10⁻¹; for MoE models λ_η ≈ 4.675 × 10⁴, α_η ≈ 1.056. This pattern is consistent across model sizes and architectures.

• Optimal annealing ratio transfers across model sizes. In Figure 6 experiments (Dense: batch 512, sequence length 8192, 24k steps, max LR 4e-5; MoE: same batch and steps, max LR 2e-4), the optimal annealing ratio converges to the same value across model sizes. The performance gap ΔLoss between different annealing ratios grows with model size, even though the optimum does not move.

• Optimal annealing ratio transfers across datasets. For a 1B MoE model (WSD, max LR 2e-4, batch 640), R_opt is consistent between training and validation sets, with absolute final loss error within 0.003. Loss gaps are larger on validation sets, and gaps shrink as training steps increase. The validation sets used are MMLU, OpenWebtext2, CN Books, and EN Books.

• Optimal annealing ratio shrinks as training gets longer. R_opt = λ_T · T^{α_T}; for the 1B MoE model, λ_T ≈ 5.987 × 10⁵ and α_T ≈ -9.46 × 10⁻¹. Figure 8 reports this for total steps of 30k, 100k, and 300k at fixed batch size 640 and maximum learning rate 2e-4. The authors note that annealing too early risks overfitting.

Methodology in Plain English

The authors start from the observation that loss-versus-steps curves are more comparable across batch sizes than loss-versus-tokens curves, so they use training steps as the unit of progress. They then write the training loss as the sum of three additive pieces: a term that decays as a power law in the accumulated learning rate over steps (the "forward effect" of training), a term that decays as a power law in model size, and a linear term in accumulated momentum (the "annealing effect"). The learning-rate integral S = ∫η(t)dt is straightforward. The momentum integral M is computed by treating the sequence of learning-rate changes Δη_t as data fed into an Adam-style update: first and second moments m_t and v_t with bias correction, accumulated as M_t = M_{t-1} + m̂_t / sqrt(v̂_t + ε). This replaces the multiplicative accumulation used in prior work and, the authors argue, makes the term stable and less sensitive to batch size.

To fit the coefficients, they use robust nonlinear regression with the L-BFGS-B optimizer from scipy.optimize.minimize, minimizing a Huber loss with threshold δ = 10⁻³ and a maximum of 1000 iterations. Loss curves are smoothed with a 50-step moving average before fitting, results are based on single-run trajectories unless otherwise stated, and accuracy is reported as Mean Absolute Percentage Error (MAPE). For the transferability studies, they sweep annealing ratios, compute the final loss at each via the fitted model, define ΔLoss as the gap to the best observed final loss at R_opt, and check whether the minimizing ratio lands in the same place across configurations.

Models span Dense (47.7M, 113.3M, 487.6M, 936.3M, and 7.03B parameters) and MoE (183M, 485M, 692M, 987M, 1.45B, and 6.8B active parameters, all top-3 routing). Training data is The Pile. AdamW uses β1 = 0.9 and β2 = 0.95; the cosine scheduler uses a 10x learning rate decay, with maximum learning rates of 1.5 × 10⁻³ for smaller models and 2 × 10⁻⁴ for larger ones, following Hoffmann et al. (2022). The WSD experiments on the 1B MoE use maximum learning rate 2 × 10⁻⁴ and 0.1 annealing ratios, with total steps ranging from 29k to 227k. Sequence length is 8192 throughout.

Why This Matters

Impact on research. The paper extends scaling-law work from predicting a single final loss number to predicting the whole loss trajectory under a scheduler, and it argues that the annealing ratio is a transferable rather than task-specific quantity. The theory appendix connects this to last-iterate convergence bounds for linear decay schedules and to bounds comparing constant and WSD schedules, where the constant-schedule bound carries a harmonic-number term H_{T-1} ≈ ln T that the WSD bound does not.

Real-world applications:

• Pre-training budget planning: picking an annealing ratio from a power law instead of running a full sweep at scale.
• Hyperparameter transfer: tuning annealing on a small Dense or MoE proxy model and applying the result to a larger model.
• Continued training and re-warming: choosing where to enter the decay phase without degrading the model.
• MoE training pipelines: the paper shows the same patterns hold for Mixture-of-Experts architectures, which have different compute/parameter tradeoffs than Dense models.

Industry relevance. Annealing schedules are one of the more expensive hyperparameters to tune because each trial requires a full training run. A predictive rule that transfers from small to large models, and across maximum learning rates and step counts, directly cuts the compute spent on schedule search.

Future Directions

• Whether the fitted coefficients (λ_η, α_η, λ_T, α_T, λ_B, α_B) hold outside the model sizes tested, since the largest configurations reported here are 1B–1.5B in the main transfer experiments and 7B architectures are listed in the model tables.
• Whether the transferability claims extend to non-linear decay shapes; the WSD formulation analyzed uses β = 1 (linear decay).
• Whether the momentum formulation stays accurate for batch sizes below the optimal batch size or for irregular scheduling, which the paper explicitly frames as regimes the integral formulation is meant to handle but does not fully sweep.
• Whether the shrinking optimal annealing ratio at longer training runs continues to follow the fitted power law well beyond the 30k–300k step range examined, or whether it saturates.

Target Audience

Researchers and engineers working on large-scale LLM pre-training, particularly those responsible for learning-rate schedule design, hyperparameter transfer, or scaling-law-based training planning. It will also interest optimization researchers studying last-iterate convergence and decay schedules, and practitioners at organizations that train Dense or MoE models where a single schedule sweep is expensive. Readers without background in Adam-family optimizers and scaling-law fitting will find the technical sections difficult, though the transferability results and the recommended power laws are usable on their own.

Authors’ abstract

Learning rate scheduling is crucial for training large language models, yet understanding the optimal annealing strategies across different model configurations remains challenging. In this work, we investigate the transferability of annealing dynamics in large language model training and refine a generalized predictive framework for optimizing annealing strategies under the Warmup-Steady-Decay (WSD) scheduler. Our improved framework incorporates training steps, maximum learning rate, and annealing behavior, enabling more efficient optimization of learning rate schedules. Our work provides a practical guidance for selecting optimal annealing strategies without exhaustive hyperparameter searches, demonstrating that smaller models can serve as reliable proxies for optimizing the training dynamics of larger models. We validate our findings on extensive experiments using both Dense and Mixture-of-Experts (MoE) models, demonstrating that optimal annealing ratios follow consistent patterns and can be transferred across different training configurations.