---
title: Learning Rate Decay in LLM Pre-training
url: https://www.emergentmind.com/topics/learning-rate-decay-in-llm-pre-training
type: topic
---

# Learning Rate Decay in LLM Pre-training

Learning rate decay is a central technique for managing optimization dynamics during large language model (LLM) pre-training. It aims to balance the need for rapid initial progress and robust convergence by manipulating the step size of the optimizer over the course of training. Multiple theoretical frameworks, empirical studies, and recent algorithmic advances have refined the understanding and practice of learning rate decay in LLM pre-training. This article synthesizes the main paradigms, theoretical justifications, modern alternatives—including adaptive and decay-free methods—their impact on training and downstream tasks, and best practices grounded in recent arXiv literature.

## 1. Paradigms and Schedule Taxonomy

The dominant learning rate decay paradigms in LLM pre-training reflect attempts to approximate optimal bias-variance tradeoffs for stochastic optimization and to manage the nonconvex structure of neural loss landscapes:

- **Cosine and Step Decay**: The learning rate is annealed according to a fixed curve (e.g., cosine, linear, multi-step), typically following a brief linear warmup phase. For cosine:
  $$
  \eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max}-\eta_{\min})\left[1 + \cos\left(\pi \frac{t-T_w}{T_{\text{tot}}-T_w}\right)\right],~ T_w < t \leq T_{\text{tot}}
  $$
- **Linear Decay-to-Zero (D2Z)**: The learning rate is decreased linearly from $\eta_{\max}$ to zero post-warmup. D2Z consistently outperforms 10$\times$ cosine or step decays in large-scale empirical studies, yielding lower pre-training loss and substantial compute savings when the pre-training budget is known [2502.15938].
- **Warmup-Stable-Decay (WSD)**: Training proceeds by warmup $\to$ long constant plateau (“stable”) at $\eta_0$ $\to$ fast decay near the end. WSD and its variants match or outperform budget-tuned cosine schedules and allow checkpointing at flexible positions [2410.05192].
- **Power-Law and Intrinsic Optimal Schedules**: Theoretically optimal decay is either a sharp power-law tail or a WSD profile depending on task and model exponents, as formalized via the Functional Scaling Law (FSL) framework [2509.19189, 2602.06797].
- **Constant LR with Model Averaging**: Instead of decaying the learning rate, checkpoint merging (uniform or weighted) simulates the effect of decay in a “decay-free” regime, providing flexibility and often improving downstream/fine-tuned task performance [2507.17634, 2511.18903].
- **Infinite/Indefinite Schedules**: “Infinite” cosine or inverse-sqrt schedules eschew a fixed training budget, enabling seamless continuation or continual pre-training without losses due to LR re-warming [2503.02844, 2403.08763].

A summary table of canonical schedule formulas:

| Schedule   | Formula (post-warmup)                  | Key Properties                   |
|------------|----------------------------------------|----------------------------------|
| Cosine     | $\eta(t)=\eta_{min}+\frac{\eta_{max}-\eta_{min}}{2}(1+\cos(\pi (t-T_{wu})/(T_{tot}-T_{wu})))$ | Smooth decay; fixed budget |
| Linear D2Z | $\eta(t) = \eta_{max}(1 - (t-T_{up})/(T-T_{up}))$ | Outperforms cosine when $T$ known |
| WSD        | Stable $\eta_0$ then decay near end    | Budget-agnostic branching        |
| Constant+Merge | $\eta(t)=\eta_{0}$; merge last $k$ checkpoints | Decay-free; maximizes adaptability |
| Infinite   | Warmup, cooldown, constant plateau     | Flexible continual training      |

## 2. Theoretical Foundations

The mathematical rationale for decaying learning rates is articulated through continuous-time and SDE treatments of stochastic gradient descent:

- **Functional Scaling Laws (FSL)**: Under general kernel-regression and SGD models, FSL expresses the evolution of risk/loss as a sum of approximation, full-batch, and stochastic-noise terms, where each term’s exponent is controlled by model capacity $\beta$ and task hardness $s$:
  $$
  \mathbb{E}[R(v_t)] \sim M^{-s\beta} + t^{-s} + \int_0^t \mathcal{K}(t-r)\gamma(r)dr
  $$
  Here $\gamma(t)$ encodes the effective gradient noise, coupling learning rate and batch size [2509.19189, 2602.06797].
- **Phase Transition: WSD vs. Power Decay**: In the easy-task regime ($s \ge 1-1/\beta$), the optimal LR schedule is a sharp power-law decay:
  $$
  \eta^*(z) = \eta_{\mathrm{peak}} (1 - z/N)^{2\beta - 1}
  $$
  For harder tasks ($s < 1-1/\beta$), the minimax-optimal solution is a WSD profile: rapid warmup to the largest stable LR, hold constant almost all training, then decay over a small final fraction of steps [2602.06797].

- **Loss-Landscape Perspective**: WSD’s behavior is explained by the river-valley analogy: in a highly anisotropic (valley-river) loss landscape, a high stable LR enables rapid “down the river” progress (low-curvature), while final decay eliminates loss due to oscillations in “hill” directions (high-curvature). The Mpemba effect justifies using a high plateau, ensuring that convergence in flat directions accelerates during the decay phase [2507.04206, 2410.05192].

## 3. Adaptive and Decay-Free Methods

The rigidity of traditional decay schedules can be mitigated or replaced by dynamic adaptation or checkpoint merging:

- **Adaptive Learning Rate Search (AdaLRS)**: AdaLRS tunes the LR on-line by optimizing loss descent velocity, using a sliding window least-squares slope estimate. Scaling factors $\alpha$ (up) and $\beta$ (down) adaptively multiply/divide the current LR based on observed slope improvements. This approach provably converges to the optimal LR and corrects mis-specified LRs within $10$–$40$\% of training, outperforming standard schedules in wall-clock step savings [2506.13274].
- **Checkpoint Merging (WSM)**: Any monotonically decreasing decay profile can be emulated by forming a weighted average of recent model checkpoints. Merge duration is the dominant hyperparameter. Mean or “1-sqrt” weighting outperforms exponential moving average, and WSM matches or exceeds WSD and cosine decay on multiple LLM benchmarks [2507.17634].
- **Model Averaging in Curriculum Schedules**: In curriculum-based LLM pre-training, aggressive LR decay nullifies the value of placing high-quality data late. Averaging the last few checkpoints under a constant LR recovers significant downstream benchmark gains (+0.74%–1.64%) [2511.18903].

## 4. Empirical Findings and Schedule Comparisons

Recent large-scale studies and benchmarks delineate the regimes where each schedule excels:

- **Linear D2Z**: Outperforms 10$\times$ cosine decay at scale—yielding lower perplexities and up to $60$\% compute savings at compute-optimal token-per-parameter (TPP) budgets [2502.15938].
- **WSD and WSD-S**: WSD consistently matches “oracle” cosine decay tuned to different pre-training budgets, while WSD-S, reusing branches, further improves loss by $10^{-3}$ over WSD in multi-budget runs [2410.05192].
- **WSM**: On math, code, and professional benchmarks, constant LR with merge-based annealing (WSM) achieves +2–5% gains over WSD, and post-fine-tuning improvements transfer downstream [2507.17634].
- **No-Decay Schedules for Fine-Tuning**: Recent work shows that omitting decay (WSO: Warmup-Stable-Only schedule) during pre-training yields flatter minima in the loss landscape, leading to superior supervised fine-tuning (SFT) downstream—even when pre-training validation loss is slightly worse [2603.16127].
- **Continual and Curriculum Regimes**: For continual pre-training, infinite schedules or quick re-warm/re-decay procedures (followed by replay) enable Chinchilla-scale LLMs to match full re-training using a fraction of the compute [2403.08763, 2503.02844]. For data curricula, moderate decay or model averaging unlocks the benefit of ascending-quality ordering [2511.18903].

Performance comparison (selected results):

| Study                  | Schedule(s) Compared      | Metric / Result                         |
|------------------------|--------------------------|-----------------------------------------|
| [2502.15938]           | Linear D2Z, 10$\times$Cosine    | D2Z achieves $>$1% lower loss, 60% less compute at scale       |
| [2506.13274]           | AdaLRS vs. tuned Cosine  | AdaLRS matches/betters PPL, 30–50% fewer steps if LR mis-set   |
| [2603.16127]           | WSO (no decay) vs. others| WSO yields highest SFT scores despite higher pre-training loss |
| [2507.17634]           | WSM (merge) vs. WSD      | +3.5% MATH, +5.5% MMLU-Pro eval gains                   |
| [2511.18903]           | Const+merge vs. decay    | +0.74–1.64% benchmark gain with SMA/EMA merging             |

## 5. Special Topics: Batch Size, Curriculum, and Continual Training

- **Batch Size Scheduling (Seesaw)**: In variance-dominated regimes (Adam, large-scale LLMs), halving LR is equivalent to multiplying batch size by $\sqrt{2}$ (for SGD, equivalently by $2$). Seesaw schedules combine batch-ramps with LR decay, achieving the same loss curve with $~36$\% reduced wall-clock [2510.14717].
- **Curriculum & Data Quality**: Aggressively decaying LR interferes with quality-driven curricula by shrinking update magnitude precisely when the best data arrives. Solutions include moderating LR decay, model averaging, or constant LR [2511.18903].
- **Continual Pre-Training & Infinite Schedules**: Repeated LR re-warming upon data shifts leads to forgetting “spikes.” Infinite schedules—warmup → cooldown → plateau—avoid sharp re-adaptation and support seamless extension [2503.02844, 2403.08763, 2410.04103].

## 6. Practical Guidelines and Future Directions

- Estimate task hardness ($s$) and model capacity ($\beta$) via power-law fitting on small runs. Use power decay if $s \ge 1-1/\beta$, WSD if $s < 1-1/\beta$ [2602.06797].
- Prefer D2Z where the training budget is known and over-training is planned [2502.15938].
- For maximum adaptability (fine-tuning/SFT), consider Warmup-Stable-Only (constant LR after warmup) [2603.16127].
- For curriculum or high-quality data late in training, use moderate decay or constant LR plus checkpoint/model averaging [2511.18903, 2507.17634].
- For continual or versioned LLM training, apply learning rate path switching: high constant LR for checkpointing, decayed LR for new data “branches” [2410.04103].
- Tune decay phase to $5$–$15$\% of training for WSD, exploit flexible branching (WSD-S), or abandon decay in favor of model merging/averaging where architecture and application permit [2410.05192, 2507.17634].

Emerging trends suggest further theoretical refinement of optimal schedule shapes under nonconvex and adversarial pre-training, as well as the adoption of plug-in adaptive algorithms (e.g., AdaLRS) robust to mis-specified initial hyperparameters. The interplay between learning rate decay, batch size schedules, data quality curriculum, and model averaging remains an active area of optimization theory and LLM engineering.

Source: https://www.emergentmind.com/topics/learning-rate-decay-in-llm-pre-training