Learning rate scheduling: warmup and cosine decay

The problem, from first principles

Picture training a large transformer from a freshly initialized set of weights. At step 0, every weight is essentially random noise, and the gradients you compute from the first few batches are enormous and unreliable: the model hasn't seen enough data yet to know which direction is actually "downhill." If you slam in the full learning rate immediately, Adam's second-moment estimate v is still close to zero (it hasn't accumulated enough history), so the bias-corrected update m_hat / (sqrt(v_hat) + eps) can spike to a huge value on the very first few steps. That spike can permanently damage the weights: a large early update can push a layer into a bad region it never recovers from, or blow up activations to inf/nan before training has barely started. This is a well documented, reproducible failure mode, not a theoretical concern.

The fix that essentially every modern large-model training run uses is to not start at the target learning rate at all. Instead, ramp the learning rate up from (near) zero over the first few hundred or thousand steps ("warmup"), giving Adam's moment estimates time to stabilize before the model takes full-sized steps. Then, once training is underway, the learning rate can't just stay at its peak forever either: late in training, you want small, careful updates that fine-tune the model into a good minimum rather than continuing to bounce around with large steps. So after warmup, the learning rate decays, typically following a smooth cosine curve down toward a small final value (often zero) by the end of training.

From theory to code

You'll implement the three pieces separately, then compose them. First, linear_warmup_lr(step, warmup_steps, base_lr): a straight linear ramp from 0 up to base_lr, reaching base_lr exactly when step == warmup_steps and holding there if called for a larger step. Second, cosine_decay_lr(step, total_steps, base_lr, min_lr): takes a fraction ("progress") through [0, total_steps] and maps it through one half-period of a cosine curve from base_lr down to min_lr, landing exactly on min_lr at step == total_steps and holding there afterward. Third, warmup_cosine_lr(step, warmup_steps, total_steps, base_lr, min_lr): the schedule an actual training loop calls once per step. It should reuse the two functions above rather than reimplementing their formulas: for step < warmup_steps, delegate to linear_warmup_lr; for step >= warmup_steps, delegate to cosine_decay_lr, but re-based so the decay phase's own internal step counter starts at 0 the moment warmup ends (i.e. pass step - warmup_steps as the decay step, and total_steps - warmup_steps as the decay phase's own length).

Constraints

• linear_warmup_lr must return exactly 0.0 at step == 0 and exactly base_lr at step == warmup_steps (and stay at base_lr for any step > warmup_steps, since the function may be called past its own valid warmup range by a caller that forgot to branch, guard against that).
• cosine_decay_lr must return exactly base_lr at step == 0 and exactly min_lr at step == total_steps, holding at min_lr afterward.
• warmup_cosine_lr must be continuous at the step == warmup_steps boundary: the value returned for step == warmup_steps should equal base_lr, matching both what warmup ends at and what decay starts at.
• Assume warmup_steps < total_steps and all inputs are non-negative.

Hints