Skip to content

← Optax step 10 of 25

Easy Primitives

Exponential Decay Schedule

Jump to the task

A learning-rate schedule maps a step count to a learning rate. Exponential decay reduces the rate by the same proportion over each interval, rather than subtracting a fixed amount.

The rule

With the settings used here:

learning_rate = init * decay_rate ** (step / transition_steps)

Use ordinary division in the exponent. For example, halfway through the first interval, the exponent is 0.5, not zero. Python uses ** for exponentiation; ^ is a different operator.

For init = 0.1, decay_rate = 0.9 and transition_steps = 100:

Step Calculation Learning rate
0 0.1 * 0.9**0 0.1
50 0.1 * 0.9**0.5 about 0.094868
100 0.1 * 0.9**1 0.09
200 0.1 * 0.9**2 0.081

Every 100 steps multiplies the rate by another 0.9. The rate also changes between those milestones. It does not wait until step 100 to start falling. A factor of 1 would hold the rate constant; a factor greater than 1 would increase it for a positive starting rate.

Inputs

Implement expdecay_at(step, init, decay_rate, transition_steps) with optax.exponential_decay. Use its defaults: start at step zero, keep staircase=False, and do not set a final floor value.

  • step: a nonnegative whole-number step count, supplied as a scalar.
  • init: a positive initial learning rate.
  • decay_rate: a factor strictly between zero and one.
  • transition_steps: a positive whole-number interval, supplied as a scalar.

Convert step and transition_steps to Python integers before calling the schedule, as required by this exercise's scalar interface.

Output

Return the learning rate at the requested step as a float32 JAX scalar. You are evaluating a schedule, not updating model parameters.

The chart below compares schedule shapes using separate demonstration settings. Its exponential curve starts at 0.001 and reaches one percent of that value at step 1000; it does not plot the table above.

Loading visualization…