Ramp Up, Ease Down

Modern training runs rarely keep one learning rate from start to finish. They warm up, starting below the target rate and climbing to it over the first few updates so that the wild gradients of freshly initialised weights cannot wreck the run. Then they decay along a smooth cosine curve down to a small floor, so the end of training takes the fine, settling steps it needs.

Number the updates $s = 1, 2, \dots, T$, where $T$ is total_steps, and let $W$ be warmup_steps.

Warmup, for $s \le W$. The rate climbs in equal increments and reaches the peak exactly on update $W$:

$\text{lr}_s = \text{peak} \times \frac{s}{W}$

Cosine decay, for $s > W$. First measure how far through the decay phase this update is:

$\text{progress} = \frac{s - W}{T - W}$

then ease from the peak down to the floor:

$\text{lr}_s = \text{min} + (\text{peak} - \text{min}) \times \frac{1 + \cos(\pi \cdot \text{progress})}{2}$

The final update, $s = T$, has progress $1$ and lands exactly on the floor. The cosine takes its argument in radians.

Task: write warmup_cosine(peak_lr, min_lr, warmup_steps, total_steps), returning the learning rate of every update in order: a list of total_steps floats, each rounded to 4 decimal places.

You can rely on total_steps $\ge 1$, $0 \le$ warmup_steps $\le$ total_steps and $0 \le$ min_lr $\le$ peak_lr. warmup_steps can be $0$, meaning no warmup at all, and it can equal total_steps, meaning the run ends before any decay begins.

Worked through, warmup_cosine(0.4, 0.0, 2, 6):

The answer is [0.2, 0.4, 0.3414, 0.2, 0.0586, 0.0].

A schedule that is slightly wrong still trains, which is exactly why schedule bugs survive for so long. Get the conventions right: where the ramp starts, what the decay's progress is measured against, and where the floor enters the formula.