OneCycleLR schedule, contrasted against warmup + cosine decay

The problem, from first principles

In [08-lr-warmup-cosine-decay] you built the schedule almost every large language model uses: ramp up, then decay down to (near) zero, spending most of training near the peak learning rate. Leslie Smith's "super-convergence" research asked a different question, aimed at a different regime: for models small enough to train in a handful of epochs, could you get away with a much more aggressive schedule, spending only a small fraction of training near the peak and instead using a MUCH larger peak learning rate than you'd normally dare use, tolerating some instability in the middle of the run in exchange for a shorter total training time? The answer was yes, and the resulting policy, "one-cycle," became the default in libraries like fastai for exactly this reason: it often reaches a target accuracy in fewer total epochs than a conventional schedule.

The mechanism is a single, symmetric-in-shape cycle across the whole run rather than a long flat plateau at the peak. Learning rate starts low, rises to a much higher-than-usual peak partway through training, then falls all the way down to a value far below where it started. The rising phase acts like an aggressive form of warmup (letting the model tolerate large steps once gradients have stabilized), and the falling phase acts like an aggressive annealing that settles the model into a sharp minimum by the very end.

From theory to code

You'll implement annealing_cos(start, end, pct), a helper that smoothly interpolates from start to end as pct sweeps from 0 to 1 along a cosine curve (this is the same shape as the second half of [08-lr-warmup-cosine-decay]'s cosine_decay_lr, but generalized to an arbitrary start and end value instead of assuming the start is always base_lr). Then onecycle_lr(step, total_steps, max_lr, pct_start, div_factor, final_div_factor) computes two derived boundary values, initial_lr = max_lr / div_factor and min_lr = initial_lr / final_div_factor, and a phase boundary step_up = pct_start * total_steps. For step <= step_up, use annealing_cos to interpolate from initial_lr up to max_lr. For step > step_up, use annealing_cos again to interpolate from max_lr down to min_lr, over the remaining total_steps - step_up steps.

Constraints

• annealing_cos(start, end, 0) must return exactly start, and annealing_cos(start, end, 1) must return exactly end.
• onecycle_lr at step == 0 must return initial_lr (max_lr / div_factor), at step == step_up must return exactly max_lr, and at step == total_steps must return min_lr (initial_lr / final_div_factor).
• min_lr ends up FAR below initial_lr in practice: with the PyTorch defaults div_factor=25 and final_div_factor=1e4, min_lr is max_lr / 250000, four orders of magnitude smaller than where the schedule started.
• Assume 0 < pct_start < 1 and total_steps > 0.

Hints