Skip to content

← Flax step 67 of 100

Medium Primitives

Warmup-Cosine LR at Step

Jump to the task

A learning-rate schedule maps a step number to a learning rate. This exercise uses linear warmup followed by cosine decay: start at zero, reach the peak at the end of warmup, then decrease to zero at the total schedule length. Different training tasks may benefit from different schedules; warmup-cosine is one choice, not a requirement for training.

Read the step axis

Write W = warmup_steps, T = total_steps and P = peak_lr.

  • For 0 <= step < W, the learning rate is P * step / W.
  • For W <= step < T, set u = (step - W) / (T - W) and use P * (1 + cos(pi * u)) / 2.
  • For step >= T, the learning rate is zero.

These are zero-based schedule coordinates. Warmup reaches the peak at W; the cosine portion has length T - W. Repeated calls at the same step return the same value. Calling the schedule does not advance an internal counter.

Optax includes warmup in decay_steps

For optax.warmup_cosine_decay_schedule, the argument decay_steps is the length of the whole schedule, including warmup. Pass total_steps directly:

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=peak_lr,
    warmup_steps=warmup_steps,
    decay_steps=total_steps,
    end_value=0.0,
)

Optax subtracts the warmup length internally to obtain the cosine duration. Passing total_steps - warmup_steps would subtract warmup twice and finish early. It can also make the cosine duration invalid when warmup occupies more than half of the requested total.

For P=0.001, W=10 and T=100:

Step Learning rate
0 0
5 0.0005
10 0.001
55 0.0005
100 and later 0

You can pass a schedule to an Optax optimizer as its learning rate. The optimizer tracks the update count used to call it; that count is separate from the schedule function itself. Here you only evaluate the schedule at a supplied step, without creating a model or optimizer.

Problem

Implement lr_at_step(step, peak_lr, warmup_steps, total_steps) using the Optax schedule above with its default cosine exponent.

step, warmup_steps and total_steps are integer-valued scalars that may arrive as floats. Cast them to int. Inputs satisfy step >= 0 and 0 < warmup_steps < total_steps. peak_lr is finite and positive.

Evaluate the schedule at step and return a floating-point JAX array of shape (1,) containing that learning rate. Do not return the schedule function or run an optimizer update. Steps beyond the total remain at zero.

Loading visualization…