Train Step with Warmup Schedule

Take one SGD step using a learning rate chosen by linear warmup. The rate starts at zero, rises to peak_lr at warmup_steps, then stays at the peak. This exercise has no decay phase.

Choose the rate, then update

For a nonnegative step count:

current_lr = peak_lr * min(step / warmup_steps, 1)
new_params = params - current_lr * grads

With peak_lr = 0.1 and warmup_steps = 100, the rates at steps 0, 50, 100 and 150 are 0.0, 0.05, 0.1 and 0.1 respectively. For a parameter of 1.0 and gradient of 0.2, step 50 gives 1.0 - 0.05 * 0.2 = 0.99.

Each call uses the supplied step to choose a rate and updates the supplied parameters once. Calling at step 50 does not perform fifty updates first. The gradients are already computed, and do not affect the schedule's value.

Using Optax

Create optax.linear_schedule(0.0, peak_lr, int(warmup_steps)), evaluate it at int(step), and use that value as the learning rate for optax.sgd. Initialize the optimizer, calculate its update and apply it once.

This exercise uses ordinary scalar step arguments outside jax.jit. Converting a concrete schedule result to Python float is allowed here, but is not required to prevent a tracer leak. Under JIT, converting a traced value to float or int can instead cause a concretization error.

For a training loop, Optax can accept a schedule directly as an optimizer's learning rate. Its state then tracks the update count. That differs from this function, where the caller explicitly supplies the step to evaluate.