Home › Neural network training & inference › Optimizers & schedules
algorithm low

Learning-rate warmup and cosine decay

Problem Training large models, especially transformers with Adam-style optimizers, can diverge at the start or end with a poor final loss if the learning rate is wrong. Engineers need a schedule that is stable and converges well.

Ramp the learning rate up linearly for the first ~1% of steps, then decay it along a half cosine to ~10% of peak.

O(1) per step; one scalar per parameter group
learning-rateschedulewarmupcosinewsd
Overview

Early in training Adam's second-moment estimate has seen only a few gradients and the weights sit in a high-curvature region, so a large lr diverges. Warmup ramps lr linearly from near 0 to the peak. Late in training, gradient noise sets a loss floor proportional to lr; decaying lr along a half cosine removes that noise and buys a large final loss drop.

Key facts

Use when

  • Any AdamW transformer run, and large-batch SGD.
  • Language models (LLaMA), image classifiers (Swin Transformer: 20 warmup epochs then cosine; ResNet with 5 warmup epochs), speech recognition (Whisper: 2048 warmup updates, then linear decay).

Avoid when

  • Training continues past the schedule: a run decayed to about 0 is hard to resume; use WSD or re-warm.

Deep dive

How it works

In: optimizer step number, peak, floor, warmup W, total T · Out: the learning rate for this step · State: none (derive from the global step) · Cost: O(1) per step.

  1. Before W, ramp linearly from near 0 to the peak.
  2. After W, follow the half cosine from peak to floor over T − W steps (or hold the peak and cool down for WSD).

With peak η_max, floor η_min, warmup W and total T optimizer steps:

step < W:   η = η_max · (step + 1) / W
step ≥ W:   p = (step − W) / (T − W)                  progress in [0, 1]
            η = η_min + ½ (η_max − η_min) (1 + cos(π p))

 η_max ┤   ╭──╮
       │  ╱    ╲___                cosine
       │ ╱         ╲___
       │╱              ╲____
 η_min ┤                    ‾‾‾‾
       └──W──────────────────── T

WSD instead holds η_max until step T − D, then decays to η_min over D steps (linear, or 1 − √p, which does best in Hägele et al.).

Why warmup works, mechanically:

  • Adam's v after k steps averages only k squared gradients. The variance of 1/√v is huge for small k, and RAdam is a closed-form fix for exactly this.
  • Post-LN transformers have large gradients near the output at initialization and diverge without warmup. Pre-LN needs much less (Xiong et al.).
  • For large-batch SGD, the linear-scaling rule assumes weights change slowly, which does not hold at step 0. Goyal et al. warm up for 5 epochs.

Code

use std::f64::consts::PI;

pub enum Decay { Cosine, Wsd { decay_steps: u64 }, InvSqrt }

pub struct Schedule { pub peak: f64, pub min: f64, pub warmup: u64, pub total: u64, pub decay: Decay }

impl Schedule {
    /// `step` counts optimizer steps (not micro-batches), starting at 0.
    pub fn lr(&self, step: u64) -> f64 {
        if step < self.warmup {
            return self.peak * (step + 1) as f64 / self.warmup as f64; // never exactly 0
        }
        let s = step.min(self.total);
        match self.decay {
            Decay::Cosine => {
                let p = (s - self.warmup) as f64 / (self.total - self.warmup).max(1) as f64;
                self.min + 0.5 * (self.peak - self.min) * (1.0 + (PI * p).cos())
            }
            Decay::Wsd { decay_steps } => {
                let start = self.total.saturating_sub(decay_steps).max(self.warmup);
                if s < start { return self.peak; }
                let p = (s - start) as f64 / decay_steps.max(1) as f64;
                self.min + (self.peak - self.min) * (1.0 - p.sqrt()) // "1 − sqrt" cooldown
            }
            // Transformer (Vaswani et al.): peak at step = warmup, then ∝ 1/√step. Ignores total.
            Decay::InvSqrt => self.peak * (self.warmup as f64 / (step + 1) as f64).sqrt(),
        }
    }
}

fn main() {
    let cos = Schedule { peak: 3e-4, min: 3e-5, warmup: 2000, total: 100_000, decay: Decay::Cosine };
    let wsd = Schedule { decay: Decay::Wsd { decay_steps: 20_000 }, ..cos };
    assert!((cos.lr(1999) - 3e-4).abs() < 1e-12);
    assert!((cos.lr(51_000) - 1.65e-4).abs() < 1e-9);      // halfway: midpoint of peak and min
    assert!((cos.lr(100_000) - 3e-5).abs() < 1e-12);
    assert_eq!(wsd.lr(79_999), 3e-4);
    for s in [0, 1000, 2000, 25_000, 50_000, 80_000, 90_000, 99_999] {
        println!("{s:>6}  cosine {:.2e}  wsd {:.2e}", cos.lr(s), wsd.lr(s));
    }
}

In PyTorch, wrap it as LambdaLR(opt, lambda s: sched.lr(s) / peak) and call sched.step() after opt.step().

Complexity and constants

  • O(1) per step. The cost is all in choosing the numbers.
  • Warmup length in practice: BERT 10k of 1M steps (1%); GPT-3 375M tokens (about 0.1% of 300B); LLaMA 2000 steps; ImageNet SGD 5 of 90 epochs.
  • Final lr: 10% of peak (LLaMA, Chinchilla) or linear to 0 (BERT). In WSD runs, most of the schedule's loss gain shows up during the short final cooldown.

Efficiency

  • Cost: one scalar per optimizer step; the schedule is free at run time. What it spends is training compute: a wrong length or floor wastes steps. Bound: none; it is a scalar function.
  • Levers:
    • cosine length equal to the real run length — overshooting by more than about 25% hurts (Chinchilla) — fixed-budget runs
    • warmup of 0.1-1% of steps — avoids early divergence, so no restarted runs (LLaMA 2000 steps, BERT 1%) — Adam-style optimizers
    • WSD: constant lr, short final cooldown — one run supports many lengths, scaling sweeps need fewer, reusable runs (Hägele et al.) — open-ended training
    • save a checkpoint before each cooldown — branches cost only the cooldown — WSD runs
    • floor at 10% of peak instead of 0 — simpler resume (LLaMA, Chinchilla) — continued training
  • Naive → standard → state of the art:
    • constant lr · diverges early with Adam, loss floor from gradient noise
    • warmup + cosine to 10% · standard for LLMs; the length is fixed in advance
    • WSD with reusable cooldowns, or schedule-free iterate averaging (no T needed) · extendable without a restart

Implementation notes

  • Count optimizer steps. With gradient accumulation k, total = micro_batches / k.
  • On resume, derive lr from the global step stored in the checkpoint, not from a fresh scheduler.
  • Per-group lr (e.g. a lower lr for embeddings) composes as a multiplier on the global schedule.
  • For WSD runs, save the checkpoint just before each cooldown. Any later cooldown can branch from it, which gives cheap scaling-law sweeps over training length.
  • Log lr every step. Schedule bugs (wrong T, off-by-k accumulation) are the most common silent training bug.

Pitfalls

  • total set too large (a multi-epoch estimate that never runs): lr never decays, and the final loss is worse than a correctly sized schedule.
  • Warmup to step W with lr = peak · step / W gives lr = 0 at step 0. Harmless for SGD; with Adam it wastes the step and still updates m and v.
  • Changing batch size mid-run without rescaling lr: linear for SGD, roughly √ for Adam.
  • Continued pretraining from a cosine-to-zero checkpoint: re-warm, and expect some initial forgetting.
  • Too-short warmup with Post-LN, or a high peak lr without QK-norm/logit soft-capping: divergence after 1-10k steps, not immediately.

Variants and alternatives

Step decaydivide by 10 at fixed epochs (ResNet 30/60/80). Simple, worse than cosine.
Linear decay to 0BERT and many fine-tuning recipes.
Inverse square rootoriginal Transformer, d^−0.5 · min(t^−0.5, t·W^−1.5) with W = 4000. Never anneals hard.
WSD / trapezoidalMiniCPM, DeepSeek LLM-style multi-stage. Supports open-ended training.
Cosine with warm restarts (SGDR)periodic restarts. Useful for snapshot ensembles.
One-cyclewarmup to a high lr, then anneal below the start (fast convnet training).
Schedule-free (Defazio et al. 2024)iterate averaging replaces decay, so no T is needed.

In production

LLaMA 1-3 (2000-step warmup, cosine to 10%), Chinchilla, GPT-3 (cosine to 10% over 260B tokens), BERT (10k warmup + linear), MiniCPM (WSD), Hugging Face get_cosine_schedule_with_warmup, Optax warmup_cosine_decay_schedule; Swin Transformer (ImageNet, AdamW with cosine decay), Whisper (speech).

Implementations

namelangrepoinstallverdictlicencehealthas_ofnotes
burnRusttracel-ai/burncargo add burn————cosine, linear and Noam LR schedulers
PyTorchPythonpytorch/pytorchpip install torch————torch.optim.lr_scheduler (LambdaLR, CosineAnnealingLR, OneCycleLR)
TransformersPythonhuggingface/transformerspip install transformers————get_cosine_schedule_with_warmup
OptaxPythongoogle-deepmind/optaxpip install optax————warmup_cosine_decay_schedule
schedule-freePythonfacebookresearch/schedule_freepip install schedulefree————optimizers without a decay schedule

Sources

  • Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR 2017. paper
  • Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017. paper
  • Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla), 2022. paper
  • Liu et al., On the Variance of the Adaptive Learning Rate and Beyond (RAdam), ICLR 2020. paper · code
  • Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020. paper
  • Hu et al., MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies, 2024. paper
  • Hägele et al., Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations, NeurIPS 2024. paper · code
  • Liu et al., Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, ICCV 2021 (AdamW, 20 warmup epochs, cosine decay). paper
  • Radford et al., Robust Speech Recognition via Large-Scale Weak Supervision (Whisper), 2022 (2048 warmup updates, linear decay). paper

Used by

Image model trainingHow do you set up the training of an image classifier or a diffusion model so it works well? This guide helps pick optimizer, augmentation, resolution schedules and other settings that are known to work.LLM fine-tuningHow do you adapt an existing pretrained language model to a specific task, style or preference, cheaply and without damaging what it already knows? This guide covers when to fine-tune and how to check it worked.LLM pretrainingHow is a large language model trained from scratch in practice? This guide covers deciding data mix, scale, schedules and parallel layout so a months-long run on thousands of GPUs stays stable and efficient, and whether to pretrain at all.
neural network training · plannedTraining a neural network at scale: optimizer and LR schedule, mixed precision, sharded optimizer state across GPUs, gradient all-reduce and a fast input pipeline

Systems (5):

llm pretrainingnnue trainingsurrogate modelingml weather modelrobot policy

Alternatives and related

AdamWSomeone training a transformer or similar large model needs an optimizer that adapts step sizes per parameter and still applies regularization correctly, so training is stable and generalises well.SGD with momentumWhen training a model with noisy minibatch gradients, plain gradient descent is slow and zig-zags. Engineers training vision models want faster, steadier progress without paying for heavy optimizer memory.Mixed-precision trainingNeural network training is slow and memory-hungry in full 32-bit arithmetic. Engineers want to use faster, smaller number formats on modern GPUs without losing accuracy or hitting numerical failures.