Learning rate schedules: Warmup & Cosine decay

This toy problem just wants to minimize $f(x)=x^2$ starting from $x=10$. The base rate here is deliberately too large to use constantly. Watch $\log_{10}(\text{loss})$ over 12 steps — it doesn't wobble, it climbs, forever.

Two schedules address this from different angles. Step decay reduces the rate at fixed checkpoints:

Warmup + cosine instead starts the rate at (near) zero and ramps it up before easing it back down:

• $\eta_0$ — the base (peak) learning rate.
• $\gamma$ — the step-decay shrink factor applied every $k$ steps.
• $T_w$ — the number of warmup steps, ramping linearly from $0$ up to $\eta_0$.
• $T$ — the total number of training steps, by which the cosine term reaches its floor.

1. Step decay still starts at the full (too-large) rate

   Until the first checkpoint at $t=k$, $\eta_{\text{step}}(t)=\eta_0$ — exactly the same too-large rate as never decaying at all. Whatever damage a too-large rate does in those first $k$ steps still happens.

2. Warmup never applies the full rate until x has already moved toward the optimum

   $\eta_{\text{warmup}}(0)=0$, and it only reaches $\eta_0$ at $t=T_w$ — by which point gradient descent has already taken several small, safe steps. The large rate arrives once it's less likely to cause an overshoot.

Switch between all three. Step decay eventually gets the loss down too — but only after a huge overshoot in the middle, while it's still using the full base rate. Warmup + cosine reaches a comparably tiny final loss without ever letting the loss climb above where it started.

1. Constant: diverges immediately and never recovers

   At the base rate, $x_{t+1}=x_t(1-2\eta_0)=x_t\cdot(-1.4)$ — each step grows in magnitude by 40%. Since every step multiplies by the same $-1.4$, this is a plain geometric sequence: $x_t = 10\cdot(-1.4)^t$. At $t=12$: $x_{12}=10\cdot(-1.4)^{12}\approx566.9$, and the loss $x_{12}^2\approx321{,}000$ — up from $100$ at the start.

2. Step decay: the same explosion, temporarily

   For the first 4 steps (before the first decay), step decay is identical to constant, so the same closed form applies: $x_4 = 10\cdot(-1.4)^4 \approx 38.4$, and loss $x_4^2\approx1476$ — nearly $15\times$ its starting value, before the rate ever drops. Once it does, the trajectory recovers quickly and the final loss ends up tiny — but only after that mid-run spike.

3. Warmup + cosine: no spike at all

   During warmup, $\eta(t)=1.2\cdot t/4$ is a different rate every step, so each step needs its own multiplication:

   • $t=0$: $\eta=0$, so $x_1=10\cdot(1-0)=10$ — completely unchanged
   • $t=1$: $\eta=0.3$, so $x_2=10\cdot(1-0.6)=4$
   • $t=2$: $\eta=0.6$, so $x_3=4\cdot(1-1.2)=-0.8$
   • $t=3$: $\eta=0.9$, so $x_4=-0.8\cdot(1-1.8)=0.64$

   By the time the rate ramps up to its peak at $t=4$, $x$ has already shrunk to $0.64$ on its own from those small early steps. The loss never exceeds its starting value of $100$ anywhere along the run, and still ends near $6\times10^{-4}$.