Learning Rate Decay › Inverse Square Root Decay

Inverse square root decay, popularized by the original Transformer paper "Attention Is All You Need," has a distinctive profile: it decreases quickly at first and then levels off gradually. This behavior makes it well-matched to how optimization often proceeds: rapid early progress followed by slower refinement.

The schedule is named for its decay rate: the learning rate is proportional to the inverse of the square root of the step count. This means the rate falls rapidly at first (the difference between step 1 and step 4 is large) and then changes slowly at large step counts (the difference between step 10,000 and step 10,001 is tiny). On a log-log scale, the decay appears as a straight line with slope $-1/2$.

The schedule is:

$$\eta_t = \eta_{\text{peak}} \cdot \frac{\sqrt{t_{\text{warmup}}}}{\sqrt{\max(t,\, t_{\text{warmup}})}}$$

where:

• $\eta_{\text{peak}}$: the peak learning rate reached after warmup
• $t_{\text{warmup}}$: the number of warmup steps
• $t$: the current step

This form describes the decay phase while capping the rate at $\eta_{\text{peak}}$ for $t \leq t_{\text{warmup}}$. In practice, it is usually paired with a separate linear warmup that increases the learning rate from zero to the peak. After warmup, the $\max(t, t_{\text{warmup}})$ term equals $t$, and the schedule follows a true inverse square root curve.

A common combined form that includes the linear warmup is the original Transformer formulation:

$$\eta_t = d_{\text{model}}^{-0.5} \cdot \min\!\left(t^{-0.5},\; t \cdot t_{\text{warmup}}^{-1.5}\right)$$

where:

• $d_{\text{model}}$: the model hidden dimension (used as a scaling factor in the original Transformer)
• The first argument $t^{-0.5}$ governs the decay phase
• The second argument $t \cdot t_{\text{warmup}}^{-1.5}$ governs the linear warmup phase

The minimum of the two quantities ensures a smooth transition: during warmup the linear term is smaller (it is small because $t$ is small), and after warmup the inverse square root term takes over (it is smaller because $t^{-0.5}$ shrinks faster than the linear term grows). The crossover happens exactly at $t = t_{\text{warmup}}$, where both terms equal $t_{\text{warmup}}^{-0.5}$.

To understand why inverse square root is preferred over exponential for long training runs, it helps to compare how aggressively each schedule reduces the learning rate. With exponential decay at $\gamma = 0.999$ per step, after 100,000 steps the learning rate has fallen to $0.999^{100000} \approx 4.5 \times 10^{-44}$, which is effectively zero. With inverse square root, after 100,000 steps from a warmup of 4,000 steps, the learning rate is $\sqrt{4000} / \sqrt{100000} \approx 0.2$ times the peak, a reduction to 20% of peak. For pre-training runs that span hundreds of thousands of steps, this gentler decay means the model is still making meaningful learning rate-scaled gradient steps even late in training.

This is particularly important for large language model pre-training, where the model is learning from a dataset so large that many distinct learning signals appear throughout training. An exponential schedule would reduce the learning rate to near zero long before the dataset is exhausted, effectively stopping learning. The inverse square root schedule allows the model to continue learning from later portions of the data with a non-negligible step size.

The warmup period addresses a specific training instability. At the very start of training, the model parameters are random, the gradient estimates are highly noisy, and the loss surface is poorly characterized. A large learning rate in this regime can cause the loss to explode or send the model into a region of parameter space that is very difficult to escape.

Linear warmup starts the learning rate near zero and increases it steadily over the first several thousand steps. By the time the learning rate reaches its peak value, the model has already made some initial progress, gradient estimates have stabilized, and the optimizer's momentum and variance estimates (in Adam) have had time to accumulate. The peak learning rate can then safely be much larger than if training had started at full speed.

The warmup is especially important for Adam. Adam estimates the first and second moments of the gradient using exponential moving averages. At step one, these estimates are initialized to zero and have not yet had time to reflect the true gradient distribution. Early updates with a large learning rate are based on these poorly calibrated estimates, which can cause wild oscillations. Warmup gives the moment estimates time to converge before the learning rate reaches its full value.

The warmup length is typically set to 4,000 to 10,000 steps for transformer training, though the optimal value depends on batch size and model size. Larger models and larger batches generally benefit from longer warmup. With very large batch sizes (as in distributed training), the signal-to-noise ratio in gradient estimates is higher because more samples are averaged, but the learning rate is also larger due to linear scaling, so the need for warmup does not disappear.

An underappreciated aspect of warmup is how it interacts with batch size. When you double the batch size, the linear scaling rule (introduced by Goyal et al. in 2017 for ImageNet training with SGD) recommends doubling the learning rate to maintain the same relative update magnitude. But doubling the learning rate without extending warmup often causes instability. The convention in large-batch training is to also scale the warmup duration: if you multiply batch size by $k$, multiply warmup steps by $k$ as well. This gives the optimizer time to accumulate reliable statistics before taking the larger steps that the scaled learning rate demands.