Learning Rate Schedules: Warmup, Decay & Cycling

Indexing. Here a schedule is a list of exactly $T$ rates, one for each optimizer update $t=0,\ldots,T-1$. Assume integers $T\ge 2$ and $0\le T_w<T$, with $0\le \eta_{\min}\le \eta_{\max}$. This counting is the convention of this page, not a universal scheduler API convention: libraries differ over whether a boundary value is used at an update or is the scheduler's state before or after one. The warmup skips the rate zero, so its first rate is positive when $\eta_{\max}>0$; $\eta_{\max}=0$ gives the all-zero schedule. Only updates that are applied count: an attempted update that is skipped, for example because its gradient overflowed, does not advance $t$ here. Training libraries differ on this as well.

For a decay coordinate $u\in[0,1]$, define

Choose $D=D_{\mathrm{linear}}$ or $D=D_{\mathrm{cosine}}$. When $T_w=0$, there is no warmup, and $D$ spans all $T$ updates:

When $T_w>0$, the first $T_w$ updates form a linear warmup that skips zero, followed by decay:

Update $T_w-1$ reaches $\eta_{\max}$; it is both the last warmup update and the $u=0$ start of decay. When $\eta_{\min}<\eta_{\max}$, update $T_w$ uses $u=1/(T-T_w)$ and is already below the peak, so the peak is not repeated. When $\eta_{\min}=\eta_{\max}$, the post-warmup segment instead remains constant at $\eta_{\max}$. The final update, $t=T-1$, uses $u=1$ and reaches $\eta_{\min}$. So $\eta_{\min}$ is where decay ends, not necessarily the smallest rate in the run: an early warmup rate can be smaller.

Inverse-square-root decay is the decay half of the original Transformer schedule (Vaswani et al., 2017, §5.3). With step number $s=t+1$, model width $d_{\text{model}}$ and $w=4000$ warmup steps, that schedule was

which rises linearly for $w$ steps, peaks at $s=w$, and then falls as $1/\sqrt{s}$. The demo's inverse-square-root option drops the warmup and starts the $1/\sqrt{\cdot}$ decay at its peak $\eta_0$ at update $t=0$:

For every schedule, the demo reports the mean over all $T$ updates,

not the mean of the fewer points used to draw the chart. The cosine decay above is one half-cycle of the SGDR annealing curve, without restarts.

Same budget, different timing

The mean, or equivalently the sum $\sum_t\eta_t$, is a schedule's total budget. Two schedules can spend the same budget at different times. Take $T=6$, $T_w=2$, $\eta_{\max}=1$ and $\eta_{\min}=0$, so the decay coordinate is $u=(t-1)/4$ for $t=2,\dots,5$:

They differ only at $t=2$, where the cosine is higher by $\delta=(\sqrt2-1)/4\approx0.104$, and at $t=4$, where it is lower by the same $\delta$, so both sums are $3$ and both means $1/2$. Now move a parameter from $\theta_0=0$ with $\theta_{t+1}=\theta_t-\eta_t q_t$, where the direction $q_t$ is fixed in advance. After six updates $\theta_6=-\sum_t\eta_t q_t$. With $q_t=1$ throughout, both schedules end at $\theta_6=-3$, by different paths: after three updates the cosine run is at $-(8+\sqrt2)/4\approx-2.354$ and the linear one at $-9/4$. With $q=(0,0,1,0,-1,0)$ the endpoints separate, $-1/2$ for linear and $-\sqrt2/2\approx-0.707$ for cosine, because the endpoint weights each rate by the direction it meets. In general the two endpoints differ by $-\delta\,(q_2-q_4)$, so they agree whenever $q_2=q_4$, whatever the other directions are.

In training the directions are not fixed in advance: they come from gradients that depend on where earlier steps went, so changing the schedule also changes them. The calculation holds them fixed to isolate timing; it does not say which schedule trains better.