Here is a frustrating fact about the cosine schedule. To use it you must decide, before step one, exactly how long the run will be, because the curve's shape depends on its end point. If you later find you have budget for 20 percent more tokens, you cannot just keep going. The learning rate is already near its floor, and stretching the curve after the fact means restarting. Chinchilla's authors found the length matters, too. Setting the cosine cycle much longer than the actual run made the final loss noticeably worse (Hoffmann et al. 2022), which means every point on a scaling-law plot needed its own run from scratch.

This section is about the three numbers that set the pace of training, the learning rate, how it changes over time, and the batch size, and about the theory that tells us how they should relate. We will derive warmup length from Adam's second-moment estimate, see why the loss falls sharply when the learning rate decays, meet the warmup-stable-decay schedule that removes the fixed-length problem, and derive the critical batch size, beyond which more GPUs stop buying speed.

Throughout, $\eta_t$ is the learning rate at step $t$, $\eta_{\max}$ its peak, $T$ the total number of steps, $B$ the batch size, $\beta_2$ Adam's second-moment decay, and $\mathbf{g}_t$ the minibatch gradient. Warmup and cosine decay themselves were introduced in §0.6 and used for our GPT in §6.3.

Warmup, derived from Adam's second moment

Why does Adam need warmup at all? §5.7 gave the qualitative answer. In the first steps, the second-moment estimate $\hat{\mathbf{v}}_t$ rests on very few gradients, so it is noisy, and dividing by its square root makes the steps large and erratic. Let's make that quantitative.

Adam's bias-corrected second moment is a weighted average of past squared gradients,

How many gradients is that average really based on? For a weighted average of independent samples with equal variance, the variance is $\sum_i w_i^2$ times the variance of one sample, so the effective sample size is $n_{\text{eff}} = 1/\sum_i w_i^2$ (with weights summing to one). Summing the geometric series,

In plain words, early on $n_{\text{eff}} \approx t$, every gradient so far counts fully, and as $t \to \infty$ it saturates at

which is 1,999 for $\beta_2 = 0.999$ and 39 for $\beta_2 = 0.95$. (This is exactly the constant $\rho_\infty$ in RAdam, Liu et al. 2020, who reached it by approximating the exponential average with a simple moving average.)

Now the variance. If the gradient coordinate has mean near zero and variance $\sigma^2$, as it nearly does early in training, then $\hat v_t \approx \sigma^2\chi^2_{n}/n$ with $n = n_{\text{eff}}$, and the step size is proportional to $1/\sqrt{\hat v_t}$. A chi-square with $n$ degrees of freedom has $\mathbb{E}[1/\chi^2_n] = 1/(n - 2)$, so

which is infinite for $n \le 2$ and close to 1 only when $n$ is in the tens. In plain words, until Adam has seen a few dozen gradients, the variance of its step size is huge, because the denominator can by chance be tiny. RAdam turns this into an algorithm, using plain momentum SGD until its approximation of $n_{\text{eff}}$ exceeds 4, and afterwards multiplying Adam's step by a rectification factor chosen so that the step's variance is the same as it would be with an infinite average. Ma and Yarats (2021) showed that ordinary linear warmup does the same job if it lasts about $2/(1 - \beta_2)$ steps, 2,000 steps for $\beta_2 = 0.999$ and 40 for $\beta_2 = 0.95$, the time it takes $n_{\text{eff}}$ to approach its limit.

With $\beta_2 = 0.999$, after 2 steps one coordinate in a hundred takes a step almost 10 times the learning rate, eight times the median, which is the "large and erratic" early behavior in numbers. After 200 steps the tail is tame. The first line also shows what the warmup rule means. At $t = 2/(1-\beta_2)$, $\beta_2^t \approx e^{-2}$, so $n_{\text{eff}}$ has reached $(1 - e^{-2})/(1 + e^{-2}) \approx 76$ percent of its limit.

Why, then, do large runs warm up for thousands of steps even with $\beta_2 = 0.95$? LLaMA 3 405B warmed up for 8,000 steps, DeepSeek-V3 and OLMo 2 for 2,000. Adam's variance is one reason, and the other reasons are about the network rather than the optimizer. At initialization, gradients are large and change direction quickly, the loss surface is sharp in some directions, and post-LN or deep models have the gradient profiles of §5.7. A recent account (Kalra & Barkeshli 2024) argues that warmup's main effect is to let the network move into flatter regions that tolerate the peak learning rate. In practice, warmup length is set empirically, somewhere around 0.1 to 1 percent of the run, and longer warmup also reduces the final loss's sensitivity to the peak learning rate (Wortsman et al. 2023).

import numpy as np

def effective_samples(t, beta2):
    return np.asarray(t, dtype=float)          # placeholder: every gradient counts fully, forever

def warmup_steps_for(beta2, frac=0.95):
    raise NotImplementedError

Why the loss falls when the learning rate decays

Anyone who has watched a training run has seen it. The loss creeps down slowly at a constant learning rate, then plunges as the learning rate decays. The simplest model that explains it is one coordinate of a quadratic with a noisy gradient. Let the loss be $\mathcal{L}(\theta) = \tfrac{h}{2}\theta^2$ with curvature $h$, and let the minibatch gradient be $h\theta + \xi$, with noise $\xi$ of mean zero and variance $\sigma^2/B$. One SGD step is

Take the variance of both sides. The noise is independent of $\theta_t$, so in the steady state, with $V = \operatorname{Var}(\theta)$,

and the expected excess loss is

In plain words, at a fixed learning rate the parameters never settle. They bounce around the minimum with a spread that grows with $\eta$ and with the gradient noise, and the loss carries a floor proportional to $\eta/B$. Lowering $\eta$ lowers the floor within about $1/(\eta h)$ steps, which is the sudden drop. A large learning rate makes fast progress along directions where the loss is still falling, and pays for it with noise in directions where it has already bottomed out. A good schedule spends most of the run at a high rate, for progress, and the end at a low one, to remove the noise. (Real loss surfaces have both kinds of directions at once, which some authors picture as a river flowing down a narrow valley. The high rate moves you along the river, and the decay drops you to its bed.)

The constant-rate losses match the formula, and decaying the rate at the end removes most of the floor in just 100 steps. That is the whole case for decay schedules, and the reason the end of training matters so much for data, too (annealing in §8.4).

Cosine, and its fixed-length problem

Recall the cosine schedule from §0.6 (Loshchilov & Hutter 2017, adopted by GPT-3, Chinchilla, and LLaMA). It warms up and then follows half a cosine from $\eta_{\max}$ to a floor $\eta_{\min}$, usually 10 percent of the peak,

It works well, and it has the problem from the start of this section, that $T$ is baked in. Halfway through, the rate is already down to 55 percent of the peak, so the run has started paying the "decay" price long before the end. You cannot extend it without a second warmup or a strange-shaped tail, and every run length on a scaling-law plot needs its own run.

The floor itself is a choice worth questioning. Bergsma et al. (2025) compared decaying to 10 percent of the peak with decaying linearly all the way to zero, at the best peak rate for each, across model sizes, batch sizes, and datasets. Decay to zero won consistently, and the gap grew with the number of tokens per parameter. In one comparison a 610M model trained for 80 tokens per parameter with linear decay to zero beat the same model trained for 200 tokens per parameter with a 10 percent floor. The noisy quadratic of the last section suggests why. Whatever learning rate is left at the end sets the noise floor the model finishes on.

Warmup-stable-decay

The fix is to separate the two jobs. Hold the learning rate constant at its peak for most of the run, and do all the decaying in a short cooldown at the end. This warmup-stable-decay schedule (WSD) was named in MiniCPM (Hu et al. 2024), which decayed over the last 10 percent of tokens and found that 2.5 percent was too short. DeepSeek LLM used a step version, dropping to 31.6 percent of the peak ($1/\sqrt{10}$) after 80 percent of the tokens and to 10 percent after 90 percent (DeepSeek-AI 2024). Hägele et al. (2024) studied the idea systematically and found the following.

• A constant rate followed by a cooldown matches cosine's final loss, and it scales the same way with model size.
• Gains from a longer cooldown level off at about 20 percent of the run. On long runs, 5 to 10 percent is nearly as good.
• The shape of the cooldown matters. A $(1 - \sqrt{x})$ shape, which drops the rate quickly at first and then slowly, beat a linear one, and the gap grew with run length.
• Because the stable phase is shared, one long run with cooldowns branched off at several points gives models of several lengths, which cut the cost of their scaling-law experiments by about half.

With a cooldown that starts at step $t_d$ and lasts until $T$, writing $x = (t - t_d)/(T - t_d) \in [0, 1]$, the three common shapes are

In plain words, climb, cruise, then land, and the landing can be any monotone curve from the peak to the floor. DeepSeek-V3 combined the ideas. Its learning rate warmed up to $2.2\times10^{-4}$ over 2,000 steps, stayed constant for the first 10 trillion tokens, followed a cosine down to $2.2\times10^{-5}$ over 4.3 trillion tokens, and then ran two short constant stages at $2.2\times10^{-5}$ and $7.3\times10^{-6}$ for the final 500 billion (DeepSeek-AI 2024).

The toy loss in the designer is the noisy quadratic above with a slowly improving "progress" term added, so it reproduces the drop at cooldown and the advantage of a long stable phase, but it is too simple to reproduce the fine differences between cooldown shapes that Hägele et al. measured on real models. Use it to build intuition about when loss falls, not to pick a shape.

import math

def wsd_lr(step, total, peak, warmup, decay_frac, shape="linear", min_lr=0.0):
    raise NotImplementedError

The gradient noise scale and the critical batch size

Now batch size. With more GPUs you can process a bigger batch in the same time. Does that make training proportionally faster? Up to a point. McCandlish et al. (2018) gave the cleanest account of where that point is, and the derivation is short enough to do in full.

Let $\mathbf{G}$ be the true gradient, $\mathbf{H}$ the Hessian, and $\boldsymbol\Sigma$ the covariance of the per-example gradient, so a batch of $B$ examples estimates $\mathbf{G}$ with covariance $\boldsymbol\Sigma/B$. Take one step of size $\varepsilon$ along the minibatch gradient $\mathbf{G}_{\text{est}}$, and expand the loss to second order,

using $\mathbb{E}[\mathbf{G}_{\text{est}}^\top\mathbf{H}\mathbf{G}_{\text{est}}] = \mathbf{G}^\top\mathbf{H}\mathbf{G} + \operatorname{tr}(\mathbf{H}\boldsymbol\Sigma)/B$. This is a quadratic in $\varepsilon$, minimized at

with $\varepsilon_{\max} = \lVert\mathbf{G}\rVert^2/\mathbf{G}^\top\mathbf{H}\mathbf{G}$, the best step with an exact gradient, $\Delta\mathcal{L}_{\max}$ the corresponding improvement, and the gradient noise scale

In plain words, a batch much smaller than $B_{\text{noise}}$ is dominated by noise, so the best step and the progress per step both grow linearly with $B$, and a batch much larger than $B_{\text{noise}}$ is almost an exact gradient, so more examples barely help. If the Hessian is a multiple of the identity, the noise scale simplifies to $B_{\text{simple}} = \operatorname{tr}(\boldsymbol\Sigma)/\lVert\mathbf{G}\rVert^2$, the total gradient variance over the squared gradient norm, which can be measured cheaply.

Two famous rules fall out of $\varepsilon_{\text{opt}}$ at once. For $B \ll B_{\text{noise}}$, $\varepsilon_{\text{opt}} \approx \varepsilon_{\max}B/B_{\text{noise}}$, so the best learning rate grows linearly with the batch size, the linear scaling rule that Goyal et al. (2017) used to train ResNet-50 with batches of 8,192 images. For $B \gg B_{\text{noise}}$ it saturates, which is why the linear rule eventually fails.

Steps against examples

Now ask how many steps $S$ and how many examples $E = BS$ a run needs to reach a fixed loss. Suppose each step makes the optimal progress $\Delta\mathcal{L}_{\text{opt}}(B)$, and let $S_{\min}$ be the number of steps with an infinite batch. Since each step makes $1/(1 + B_{\text{noise}}/B)$ of the infinite-batch progress,

with $E_{\min} = S_{\min}B_{\text{noise}}$, the examples needed with a tiny batch. Then $S/S_{\min} - 1 = B_{\text{noise}}/B$ and $E/E_{\min} - 1 = B/B_{\text{noise}}$, so

In plain words, there is a hyperbolic trade-off between time (steps) and compute (examples). The critical batch size $B_{\text{crit}} = E_{\min}/S_{\min} = B_{\text{noise}}$ is the balanced point, where a run takes twice the minimum steps and twice the minimum examples. Below it, doubling the batch nearly halves the steps for almost no extra compute. Above it, doubling the batch nearly doubles the compute for almost no saving in time. McCandlish et al. measured critical batch sizes from about 20 for a small autoencoder to more than 10 million for Dota 5v5, and found that the noise scale grows as training progresses, because $\lVert\mathbf{G}\rVert$ shrinks as the loss falls while the noise does not. That is the reason for batch-size ramps.

How do you measure $B_{\text{simple}}$ during a run, without per-example gradients? With two batch sizes. A minibatch gradient's squared norm satisfies $\mathbb{E}\lVert\mathbf{G}_B\rVert^2 = \lVert\mathbf{G}\rVert^2 + \operatorname{tr}(\boldsymbol\Sigma)/B$, a line in $1/B$. Measure it at a small batch $B_s$ (for example one GPU's gradient before the all-reduce) and a large one $B_b$ (the all-reduced gradient), and solve the two equations,

each unbiased. Their ratio is a biased estimate of $B_{\text{simple}}$, so in practice both are averaged over many steps before dividing (McCandlish et al. 2018, Appendix A).

import numpy as np

def noise_scale(G_small, G_big, B_small, B_big):
    raise NotImplementedError

import numpy as np

def steps_and_examples(B, B_noise, S_min):
    raise NotImplementedError

def batch_for_slowdown(B_noise, max_extra_compute):
    raise NotImplementedError

Batch-size ramps, and how the learning rate should move with the batch

Because the noise scale is small early in training and grows as the loss falls, a fixed large batch wastes data at the start. So large runs ramp the batch size. GPT-3 increased its batch linearly from 32,000 tokens to its full 3.2 million over the first 4 to 12 billion tokens, depending on model size (Brown et al. 2020). PaLM went from 512 to 1,024 to 2,048 sequences (1 to 4 million tokens) at steps 50,000 and 115,000. LLaMA 3 405B started at 4 million tokens with sequences of 4,096, moved to 8 million tokens and sequences of 8,192 after 252 million tokens, and to 16 million after 2.87 trillion. DeepSeek-V3 ramped from 3,072 to 15,360 sequences over its first 469 billion tokens.

How should the learning rate change when the batch does? For SGD in the small-batch regime, the derivation above says linearly. For Adam the answer is different. Malladi et al. (2022) derived from a stochastic differential equation approximation of Adam that when the batch is multiplied by $\kappa$, the dynamics are preserved by the square-root rule,

In plain words, with a batch $\kappa$ times larger the run takes $\kappa$ times fewer steps, so each step must make $\kappa$ times the progress while adding no more noise per example. When noise dominates Adam's denominator, $\sqrt{\hat v}$ shrinks by $\sqrt\kappa$ along with the noise, so the same $\eta$ already makes each step's drift $\sqrt\kappa$ times larger, and a learning rate $\sqrt\kappa$ times larger supplies the rest of the factor $\kappa$. The same $\sqrt\kappa$ keeps the noise added per example unchanged, which is why the two requirements agree. The betas must shorten their memory in proportion, because each step now covers $\kappa$ times as many examples, so the averages should span the same number of examples rather than the same number of steps. The rule holds only below the critical batch size, like the linear rule for SGD.

At the frontier, teams fit learning rate and batch size directly as power laws in compute. DeepSeek LLM found the near-optimal values for its models followed $\eta_{\text{opt}} = 0.3118\,C^{-0.1250}$ and $B_{\text{opt}} = 0.2920\,C^{0.3271}$, with $C$ the training compute and $B$ in tokens (DeepSeek-AI 2024). For $C = 10^{24}$ FLOPs, that is a batch of about 21 million tokens and a learning rate of about $3.1\times10^{-4}$, close to what they used for their 67B model. The optimal batch grows with compute, because bigger runs reach lower losses where the noise scale is larger, and the optimal learning rate shrinks slowly. How such hyperparameters transfer across model sizes is the subject of §9.5.

Two recent results refine this picture. Zhang et al. (2024) separated the effects of model size and data size on the critical batch size and found that it grows mainly with the amount of data trained on, not with the number of parameters, so a longer run can use a larger batch even at a fixed model size. And Marek et al. (2025) looked at the other end, very small batches, and found that they train language models well, down to a batch of one sequence, provided Adam's $\beta_2$ is adjusted so that the second moment's half-life measured in tokens stays fixed. With $\beta_2^{\,k}$ the decay after $k$ steps of $B$ tokens each, holding the half-life fixed when the batch goes from $B$ to $B'$ gives $\beta_2' = \beta_2^{\,B'/B}$, which to first order is the $1 - \beta_2' = \kappa(1 - \beta_2)$ of the square-root rule above. Both results say the same thing from different sides. Batch size is a statement about how many tokens each optimizer step averages, and the other hyperparameters should be set in tokens too.

In PyTorch

In PyTorch, any schedule is a LambdaLR that multiplies the base learning rate by a function of the step. Here is WSD with a $(1 - \sqrt{x})$ cooldown, checked against the formula.

Read more

• McCandlish et al. (2018). An Empirical Model of Large-Batch Training. The gradient noise scale, the critical batch size, and how to measure them.
• Hägele et al. (2024). Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations. Constant learning rate with cooldown, cooldown shapes, and cheaper scaling laws.
• Hu et al. (2024). MiniCPM. The warmup-stable-decay schedule, with its decay-length experiments.
• Liu et al. (2020). On the Variance of the Adaptive Learning Rate and Beyond, with Ma and Yarats (2021), On the adequacy of untuned warmup for adaptive optimization. Why Adam needs warmup, and how long.
• Malladi et al. (2022). On the SDEs and Scaling Rules for Adaptive Gradient Algorithms. The square-root scaling rule for Adam.
• DeepSeek-AI (2024). DeepSeek LLM, Scaling Open-Source Language Models with Longtermism. The multi-step schedule and fitted scaling laws for batch size and learning rate.
• Bergsma et al. (2025). Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs. The case against a 10 percent floor.
• Zhang et al. (2024). How Does Critical Batch Size Scale in Pre-training?, with Marek et al. (2025), Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful. The critical batch grows with data, and small batches work when $\beta_2$ is set in tokens.

Watch next

▶

▶

▶