Skip to content

← Optax step 8 of 25

Medium Primitives

Warmup + Cosine Decay

Jump to the task

This schedule starts at zero, rises linearly to a peak learning rate, then follows a cosine curve back to zero. Warmup makes the initial learning rates smaller, but does not by itself guarantee stable training.

Where each phase ends

In optax.warmup_cosine_decay_schedule, decay_steps includes warmup. It is the step where the entire schedule reaches its final value. The cosine portion lasts decay_steps - warmup_steps steps.

For peak_lr = 0.01, warmup_steps = 100 and decay_steps = 1000:

Step Position in the schedule Learning rate
0 Start 0.0
50 Halfway through warmup 0.005
100 Peak; cosine decay begins 0.01
550 Halfway through the 900-step cosine portion 0.005
1000 End of decay 0.0
1100 After the schedule ends 0.0

Do not add the warmup length to decay_steps. In this example the rate reaches zero at step 1000, not step 1100.

The values between milestones

During warmup, the rate is peak_lr * step / warmup_steps. During decay, let progress = (step - warmup_steps) / (decay_steps - warmup_steps). The rate is peak_lr * 0.5 * (1 + cos(pi * progress)) until progress reaches one. After that, the rate stays zero; the cosine does not start rising again.

Inputs

Implement warmup_cosine_lr_at(step, peak_lr, warmup_steps, decay_steps) using Optax. Explicitly set init_value=0.0 and end_value=0.0, pass peak_lr as peak_value, and leave the cosine exponent at its default of 1.

  • step: a nonnegative whole-number step count, supplied as a scalar.
  • peak_lr: a positive peak learning rate.
  • warmup_steps: a positive whole-number warmup duration.
  • decay_steps: the whole-number end step, greater than warmup_steps.

Convert the three step-count arguments to Python integers before using them.

Output

Return the learning rate at step as a float32 JAX scalar.

The chart uses the same 100-step warmup and 1000-step endpoint as the table, but a smaller peak of 0.001. Move the warmup slider to see how its length changes the time left for cosine decay.

Loading visualization…