KAHIBARO
6.4. Learning Rate Schedules
Table of Contents
What a Learning Rate Schedule Does
A learning rate schedule changes the optimizer learning rate during training, usually as a function of epoch or step. The goal is to get fast progress early and more careful updates later. In practice, a schedule is often the difference between a model that plateaus early and one that keeps improving, even when everything else in the training loop is correct.
A schedule does not replace choosing a reasonable base learning rate. You still start with a learning rate that makes training stable, then you decide how it should evolve over time.
A learning rate schedule only changes the learning rate. It cannot fix a fundamentally wrong model, broken data pipeline, or incorrect loss setup.
Where Schedulers Fit in the Training Loop
In PyTorch, schedules are typically handled by torch.optim.lr_scheduler. You create a scheduler that wraps an optimizer, then you call scheduler.step() at the correct time.
The key detail is that different schedulers expect to be stepped at different times. Many schedules are stepped once per epoch. Some are stepped every optimizer update. Some are stepped based on a validation metric.
Calling scheduler.step() at the wrong frequency, or at the wrong point in the loop, is a common silent bug that can ruin training.
The Most Common Scheduler Types
Step based decay
Step decay reduces the learning rate by a constant factor every fixed number of epochs. Conceptually, you start with $\eta_0$ and periodically multiply it by $\gamma$.
If the learning rate starts at $\eta_0$, then after $k$ drops the learning rate is $\eta_k = \eta_0 \gamma^k$.
In PyTorch this is commonly done with StepLR or MultiStepLR. StepLR uses a fixed period, MultiStepLR uses a list of milestone epochs.
This is a strong default for beginners because it is easy to reason about and often works well.
Exponential decay
Exponential decay multiplies the learning rate by a constant factor each epoch. It is like a step schedule with a step every epoch. It can be useful when you want a smooth decay that is easy to tune.
In PyTorch this is ExponentialLR.
Cosine annealing
Cosine schedules smoothly decrease the learning rate following a cosine curve from an initial value to a minimum value. They are popular because they often give strong results without much tuning, especially for vision models and larger trainings.
In PyTorch, look at CosineAnnealingLR and CosineAnnealingWarmRestarts. Warm restarts periodically raise the learning rate again, which can help the optimizer explore new regions.
Reduce on plateau
Sometimes you do not want to reduce the learning rate on a fixed timetable, you want to reduce it when validation performance stops improving. This is what ReduceLROnPlateau does. You call scheduler.step(metric) after validation, where metric is often validation loss.
This is a practical choice when you have no strong idea how long training will take, or when datasets vary in difficulty.
ReduceLROnPlateau is stepped with a validation metric after validation, not after every training batch. Most other schedulers are stepped without a metric.
Warmup schedules
Warmup gradually increases the learning rate from a small value to the target value over the first few hundred or few thousand steps. Warmup is common with Transformers and larger batch sizes, where jumping to the full learning rate immediately can cause instability.
PyTorch has some warmup like behavior via LinearLR, ConstantLR, or by composing schedulers with SequentialLR. Many projects also implement warmup manually by updating optimizer.param_groups[i]["lr"].
How to Implement a Scheduler in Code
Epoch stepped schedulers
A typical pattern is to step once at the end of each epoch. This aligns with schedulers designed around epochs.
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
for epoch in range(num_epochs):
model.train()
for xb, yb in train_loader:
optimizer.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()
scheduler.step() # step per epochStep stepped schedulers
Some schedulers, including popular one cycle policies, are meant to be stepped every optimizer update. A typical pattern is to call scheduler.step() after optimizer.step().
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader)*num_epochs)
for epoch in range(num_epochs):
model.train()
for xb, yb in train_loader:
optimizer.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()
scheduler.step() # step per batch
The important part is that the scheduler you choose must match how you step it. If T_max is in steps, you should step in steps.
For schedulers stepped per batch, compute schedule lengths using number of optimizer updates, not number of examples. If you use gradient accumulation, the number of optimizer updates changes.
ReduceLROnPlateau pattern
ReduceLROnPlateau is usually called after validation:
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode="min", factor=0.5, patience=2
)
for epoch in range(num_epochs):
train_one_epoch(...)
val_loss = evaluate_on_validation(...)
scheduler.step(val_loss)Here, the schedule reacts to validation behavior, not training loss.
Inspecting and Logging the Current Learning Rate
When you use a scheduler, you should log the learning rate alongside loss. The most direct way is to read it from the optimizer:
current_lr = optimizer.param_groups[0]["lr"]If you have multiple parameter groups, each group can have its own learning rate, so you may want to log all of them.
If you use multiple parameter groups, do not assume param_groups[0]["lr"] represents the whole model. Different groups can follow different effective learning rates.
Practical Guidance for Choosing a Scheduler
If you want a simple starting point, use a fixed learning rate first to confirm training works. Then add one of these depending on your situation. For general small to medium problems, StepLR or ReduceLROnPlateau are reliable. For longer runs where you expect steady improvement, cosine annealing is a strong option. If your training is unstable at the start, add warmup.
A schedule that decays too early can freeze learning, and a schedule that stays high for too long can keep the model bouncing around without settling. When you look at training curves, a helpful mental model is that lowering the learning rate usually reduces noise in updates and can help validation loss improve after a plateau.
If training loss decreases but validation loss stagnates or worsens, reducing learning rate might help, but it can also simply make overfitting happen more slowly. Always check validation metrics, not just training loss.
Common Mistakes to Avoid
One frequent mistake is stepping the scheduler before the optimizer update for schedulers that assume the opposite. Another is mixing epoch based intuition with step based parameters, for example setting T_max as number of epochs but stepping every batch. Another common issue is forgetting that ReduceLROnPlateau needs the metric argument, so it never updates if you call step() without passing the value.
Finally, if you resume training from a checkpoint, you need to restore the scheduler state as well as the optimizer state, otherwise the learning rate may jump to an incorrect value relative to the training progress.
When resuming training, save and load both optimizer.state_dict() and scheduler.state_dict(). Otherwise your learning rate schedule will not continue correctly.
Views: 82