Skip to content

← Flax NNX step 58 of 100

Medium Primitives

NNX LR Schedule

Jump to the task

Learning-rate schedules are functions from a step number to a learning rate. The same Optax schedule can be used with NNX, Linen or a model written directly in JAX. This exercise evaluates a schedule; it does not create an NNX model or update any parameters.

Warmup followed by cosine decay

Write W = warmup_steps, T = total_steps and P = peak_lr. The learning rate follows three rules:

  • Before W, it rises linearly: P * step / W.
  • From W to T, it follows P * (1 + cos(pi * u)) / 2, where u = (step - W) / (T - W).
  • At T and beyond, it stays at zero.

Warmup reaches the peak at step W, and cosine decay lasts T - W steps. For P=0.01, W=100 and T=1000, steps 0, 50, 100, 550 and 1000 give 0, 0.005, 0.01, 0.005 and 0 respectively. These are zero-based schedule coordinates. No step counter changes when you call the schedule.

Warmup limits the learning rate early in training, while cosine decay reduces it later. This is one schedule choice, not a guarantee of stable training or a rule followed by every language model.

The Optax argument includes warmup

For this constructor, decay_steps means the whole schedule length, including warmup. Pass total_steps, not total_steps - warmup_steps:

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 calculates the cosine duration internally. Subtracting warmup before passing the argument makes the schedule end early and can produce an invalid cosine duration when warmup is more than half the total.

To use a schedule in training, pass it as the learning rate to an Optax optimizer, such as optax.sgd(schedule). The optimizer manages the update count used to query the function. Here you supply the step explicitly.

Problem

Implement lr_at_step(step, peak_lr, warmup_steps, total_steps) using the constructor 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 the requested step. Return a floating-point JAX array of shape (1,) containing that value. This is a JAX array, not a NumPy host array. Do not return the callable or perform an optimizer update.

Loading visualization…