Learning Rate Schedulers¶
A learning rate scheduler is a function that maps the current epoch $t$ to a learning rate $\eta_t$, modifying how quickly the optimizer traverses the loss landscape over time. Properly annealing the learning rate is one of the most impactful techniques for improving training convergence.
Why Schedule the Learning Rate?¶
The learning rate $\eta$ controls the step size in the parameter update rule. For gradient descent:
$$ \boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta_t \nabla_{\boldsymbol{\theta}} \mathcal{L}(\boldsymbol{\theta}_t) $$
A constant $\eta$ presents a trade-off:
- Too large: the optimizer oscillates around the optimum and may diverge.
- Too small: convergence is slow, especially early in training.
Scheduling addresses this by using a large $\eta$ early (coarse exploration) and a small $\eta$ late (fine convergence).
StepLR¶
StepLR multiplies the learning rate by a fixed factor $\gamma \in (0, 1)$ every $s$ epochs:
$$ \eta_t = \eta_0 \cdot \gamma^{\lfloor t / s \rfloor} $$
where $\eta_0$ is the initial learning rate, $s$ is the step_size, and $\gamma$ is the gamma parameter. The resulting schedule is a staircase function — the learning rate is constant within each interval of $s$ epochs and drops abruptly at multiples of $s$.
ExponentialLR¶
ExponentialLR applies exponential decay at every epoch:
$$ \eta_t = \eta_0 \cdot \gamma^t $$
This is the continuous analogue of StepLR with $s = 1$. The schedule is strictly monotonically decreasing and smooth.
CosineAnnealingLR¶
CosineAnnealingLR anneals the learning rate following a half-cosine curve over $T_{\max}$ epochs, cycling between $\eta_0$ (initial rate) and $\eta_{\min}$ (minimum rate):
$$ \eta_t = \eta_{\min} + \frac{1}{2}(\eta_0 - \eta_{\min})\left(1 + \cos\left(\frac{\pi\, t}{T_{\max}}\right)\right) $$
At $t = 0$ this gives $\eta_0$; at $t = T_{\max}$ it gives $\eta_{\min}$.
!!! warning "The cosine is periodic" The formula has period $2T_{\max}$. If you keep stepping past $T_{\max}$ the learning rate rises back towards $\eta_0$ (a warm restart). Choose $T_{\max}$ equal to your number of epochs, or stop stepping the scheduler once it is reached.
The cosine schedule has become a standard baseline in deep learning because it:
- Starts with a relatively large learning rate for exploration.
- Decays smoothly (no abrupt drops).
- Naturally reaches a very small value near convergence.
ReduceLROnPlateau¶
Unlike the previous schedulers (which are epoch-driven), ReduceLROnPlateau is metric-driven: it monitors a scalar quantity (e.g., validation loss) and reduces the learning rate by a factor $\rho \in (0,1)$ when no improvement has been observed for patience consecutive epochs:
$$ \eta_{t+1} = \begin{cases} \max(\rho\, \eta_t,\ \eta_{\min}) & \text{if no improvement for `patience' epochs} \\ \eta_t & \text{otherwise} \end{cases} $$
Improvement is defined as: in mode='min', the metric must decrease by at least threshold to count as an improvement; in mode='max' it must increase.
Where the Scheduler Fits in the Training Loop¶
A scheduler never touches gradients or parameters — it only rewrites optimizer.param_groups[*]['lr'], which optimizer.step() reads on its next call. So the forward and backward passes are completely unaffected; only the step size applied to the already-computed gradients changes.
The call order matters. scheduler.step() goes after optimizer.step(), and it is called once per epoch, not once per mini-batch:
for epoch in range(n_epochs):
for X_batch, y_batch in loader:
optimizer.zero_grad() # clear old gradients
loss = criterion(model(X_batch), y_batch) # forward pass
loss.backward() # backward pass → p.grad for every parameter
optimizer.step() # p ← p − η_t · update(p.grad), with the current η_t
scheduler.step() # choose η_{t+1} for the next epoch
Stepping the scheduler inside the batch loop instead would advance the schedule len(loader) times per epoch, decaying the learning rate far faster than intended.
Interface¶
All schedulers (except ReduceLROnPlateau) expose:
.step()— advance by one epoch, updatingoptimizer.param_groups[*]['lr']..get_last_lr()— return current learning rates (one per param group)..state_dict()/.load_state_dict()— serialisation for checkpointing.
ReduceLROnPlateau.step(metric) takes the monitored metric value as argument instead.
Every epoch-driven scheduler computes $\eta_t$ in closed form from the initial learning rates $\eta_0$ (stored as base_lrs) and the epoch counter last_epoch, rather than by repeatedly multiplying the optimizer's current value. That is what makes state_dict() / load_state_dict() restore the learning rate exactly.
# Uncomment the next line and run this cell to install sorix
#!pip install 'sorix @ git+https://github.com/Mitchell-Mirano/sorix.git@develop'
import math
import numpy as np
import sorix
from sorix.optim.lr_scheduler import StepLR, ExponentialLR, CosineAnnealingLR, ReduceLROnPlateau
def make_optimizer(lr=1e-3):
"""Helper: a small model + Adam optimizer."""
layer = sorix.nn.Linear(4, 4)
return sorix.optim.Adam(layer.parameters(), lr=lr)
StepLR — Staircase Decay¶
The learning rate drops by a factor of gamma every step_size epochs.
opt = make_optimizer(lr=1.0)
sched = StepLR(opt, step_size=3, gamma=0.1)
for epoch in range(7):
lr = opt.param_groups[0]['lr']
print(f"Epoch {epoch:2d} | lr = {lr:.6f}")
sched.step()
Epoch 0 | lr = 1.000000 Epoch 1 | lr = 1.000000 Epoch 2 | lr = 1.000000 Epoch 3 | lr = 0.100000 Epoch 4 | lr = 0.100000 Epoch 5 | lr = 0.100000 Epoch 6 | lr = 0.010000
ExponentialLR — Smooth Exponential Decay¶
The learning rate is multiplied by gamma at every epoch, producing smooth continuous decay.
opt = make_optimizer(lr=0.1)
sched = ExponentialLR(opt, gamma=0.95)
for epoch in range(5):
lr = opt.param_groups[0]['lr']
print(f"Epoch {epoch:2d} | lr = {lr:.6f}")
sched.step()
Epoch 0 | lr = 0.100000 Epoch 1 | lr = 0.095000 Epoch 2 | lr = 0.090250 Epoch 3 | lr = 0.085737 Epoch 4 | lr = 0.081451
CosineAnnealingLR — Cosine Curve¶
The learning rate follows a half-cosine from base_lr down to eta_min over T_max epochs.
opt = make_optimizer(lr=0.1)
T_max = 50
eta_min = 1e-6
sched = CosineAnnealingLR(opt, T_max=T_max, eta_min=eta_min)
lrs = []
for epoch in range(T_max + 1):
lrs.append(opt.param_groups[0]['lr'])
sched.step()
print(f"Epoch 0 | lr = {lrs[0]:.6f} (initial)")
print(f"Epoch 25 | lr = {lrs[25]:.6f} (mid-point ≈ η_min + (η0−η_min)/2)")
print(f"Epoch 50 | lr = {lrs[50]:.6f} (at T_max → η_min)")
Epoch 0 | lr = 0.100000 (initial) Epoch 25 | lr = 0.050001 (mid-point ≈ η_min + (η0−η_min)/2) Epoch 50 | lr = 0.000001 (at T_max → η_min)
ReduceLROnPlateau — Adaptive Reduction¶
The scheduler monitors a validation metric. Once no improvement is seen for patience epochs, the learning rate is multiplied by factor.
opt = make_optimizer(lr=0.1)
sched = ReduceLROnPlateau(opt, mode='min', patience=3, factor=0.5)
val_losses = [0.5, 0.5, 0.5, 0.5, 0.3, 0.3]
prev_lr = opt.param_groups[0]['lr']
for epoch, val_loss in enumerate(val_losses):
sched.step(val_loss)
lr = opt.param_groups[0]['lr']
reduced = " ← LR reduced!" if lr < prev_lr else ""
prev_lr = lr
print(f"Epoch {epoch:2d} | val_loss={val_loss:.4f} | lr={lr:.6f}{reduced}")
Epoch 0 | val_loss=0.5000 | lr=0.100000 Epoch 1 | val_loss=0.5000 | lr=0.100000 Epoch 2 | val_loss=0.5000 | lr=0.100000 Epoch 3 | val_loss=0.5000 | lr=0.050000 ← LR reduced! Epoch 4 | val_loss=0.3000 | lr=0.050000 Epoch 5 | val_loss=0.3000 | lr=0.050000
Checkpointing¶
state_dict() captures last_epoch and base_lrs; load_state_dict() restores them
and re-applies the corresponding learning rate to the optimizer, so a resumed run
continues exactly on the original schedule.
opt = make_optimizer(lr=1.0)
sched = StepLR(opt, step_size=2, gamma=0.5)
# Train for 5 epochs
for _ in range(5):
sched.step()
# Save
state = sched.state_dict()
print(f"Saved scheduler state at epoch 5: last_epoch={state['last_epoch']}")
# Restore into a new scheduler instance
opt2 = make_optimizer(lr=1.0)
sched2 = StepLR(opt2, step_size=2, gamma=0.5)
sched2.load_state_dict(state)
print(f"Restored scheduler epoch: {sched2.last_epoch}")
sched2.step()
print(f"Next lr after step: {opt2.param_groups[0]['lr']:.6f}")
Saved scheduler state at epoch 5: last_epoch=5 Restored scheduler epoch: 5 Next lr after step: 0.125000
End-to-End: Scheduler Inside a Real Training Loop¶
A concrete run on a linear regression problem. Note that the scheduler only changes the
step size used by optimizer.step() — the forward pass, the loss and loss.backward()
are identical to a run without a scheduler.
sorix.manual_seed(0)
# Synthetic regression data
X = sorix.tensor(np.random.randn(256, 8).astype(np.float32))
w_true = np.random.randn(8, 1).astype(np.float32)
y = sorix.tensor(X.data @ w_true + 0.1 * np.random.randn(256, 1).astype(np.float32))
model = sorix.nn.Linear(8, 1)
opt = sorix.optim.Adam(model.parameters(), lr=0.1)
sched = CosineAnnealingLR(opt, T_max=40, eta_min=1e-4)
for epoch in range(40):
lr = opt.param_groups[0]['lr'] # the lr this epoch's updates will use
opt.zero_grad() # clear gradients from the previous epoch
loss = ((model(X) - y) ** 2).mean() # forward pass
loss.backward() # backward pass: fills W.grad and b.grad
opt.step() # W -= lr * Adam_update(W.grad)
sched.step() # sets the lr for the next epoch
if epoch % 8 == 0 or epoch == 39:
print(f"Epoch {epoch:2d} | lr = {lr:.5f} | loss = {float(loss.data):.6f}")
Epoch 0 | lr = 0.10000 | loss = 8.417645 Epoch 8 | lr = 0.09046 | loss = 2.107695 Epoch 16 | lr = 0.06549 | loss = 0.601116 Epoch 24 | lr = 0.03461 | loss = 0.203377 Epoch 32 | lr = 0.00964 | loss = 0.114271 Epoch 39 | lr = 0.00025 | loss = 0.103842