Learning Rate Schedules

Optax includes a rich set of learning rate schedules that compose naturally with optimizers. A schedule is simply a function that maps step count to a learning rate value.

import optax
import jax.numpy as jnp

# Cosine decay: starts at init_value, decays to alpha over decay_steps
schedule = optax.cosine_decay_schedule(
    init_value=1e-3,
    decay_steps=10000,
    alpha=0.0,        # minimum learning rate
)

# Check values at different steps
print(f"Step 0: {schedule(0):.6f}")      # 0.001000
print(f"Step 5000: {schedule(5000):.6f}") # 0.000500
print(f"Step 10000: {schedule(10000):.6f}")# 0.000000

# Warmup + cosine decay (very common in practice)
schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,        # start from 0
    peak_value=1e-3,       # warm up to this
    warmup_steps=1000,     # linear warmup for 1000 steps
    decay_steps=50000,     # total steps including warmup
    end_value=1e-5,        # minimum LR at end
)

# Use schedule with an optimizer
optimizer = optax.adamw(learning_rate=schedule, weight_decay=0.01)

# Or compose with chain
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule),
)

More Schedule Options

# Exponential decay
schedule = optax.exponential_decay(
    init_value=1e-3,
    transition_steps=1000,
    decay_rate=0.96,
)

# Piecewise constant (manual step schedule)
schedule = optax.piecewise_constant_schedule(
    init_value=1e-3,
    boundaries_and_scales={
        5000: 0.1,   # multiply LR by 0.1 at step 5000
        8000: 0.1,   # multiply again at step 8000
    }
)

# Warm restarts (SGDR)
schedule = optax.sgdr_schedule([
    dict(init_value=1e-3, peak_value=1e-3,
         decay_steps=5000, warmup_steps=500),
    dict(init_value=1e-3, peak_value=5e-4,
         decay_steps=5000, warmup_steps=500),
])