Skip to content

← Optax step 7 of 25

Easy Primitives

Linear Schedule

Jump to the task

A linear schedule changes the learning rate by a fixed amount per step until it reaches a chosen endpoint. After that, it holds the final value. It can rise for warmup or fall for decay.

Interpolate, then hold

For a nonnegative step t and positive duration T = transition_steps:

fraction = min(t / T, 1)
learning_rate = init_value + (end_value - init_value) * fraction

For init_value = 1.0, end_value = 0.0 and T = 100:

Step Learning rate
0 1.0
50 0.5
100 0.0
150 0.0

The schedule reaches its endpoint at step 100. It does not continue below zero or restart. With init_value = 0.0 and end_value = 0.01 instead, it rises to 0.005 at step 50 and reaches 0.01 at step 100.

Inputs

Implement linear_decay_at(step, init_value, end_value, transition_steps) with optax.linear_schedule. Keep the default transition start at step zero. Despite the function's name, it must also support increasing schedules.

  • step: a nonnegative whole-number step count, supplied as a scalar.
  • init_value: the starting learning rate.
  • end_value: the final learning rate.
  • transition_steps: a positive whole-number duration, supplied as a scalar.

Convert step and transition_steps to Python integers for this exercise's scalar interface. Use ordinary division when thinking through the formula.

Output

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

The chart below starts with a separate example: a linear decrease from 0.001 to zero over 1000 steps. It illustrates the shape rather than the specific values in the table.

Loading visualization…