Learning Rate Decay › Worked Example: Effect on Convergence
Language AI HandbookAll chapters Learning Rate Decay › Worked Example: Effect on ConvergenceOne section · Sections in this chapter10 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
To see concretely how schedule choice affects training, let us train a small neural network on a synthetic regression task using four different schedules and compare convergence.
torch.manual_seed(42)
np.random.seed(42)
n_samples = 400
X_data = torch.randn(n_samples, 10)
true_weights = torch.randn(10)
y_data = X_data @ true_weights + 0.1 * torch.randn(n_samples)
dataset = torch.utils.data.TensorDataset(X_data, y_data)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)def train_with_schedule(scheduler_type, n_epochs=60, lr=0.05):
net = nn.Linear(10, 1)
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9)
criterion = nn.MSELoss()
if scheduler_type == "step":
scheduler = optim.lr_scheduler.StepLR(
optimizer, step_size=20, gamma=0.1
)
elif scheduler_type == "exponential":
gamma_e = (1e-4 / lr) ** (1 / n_epochs)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma_e)
elif scheduler_type == "cosine":
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=n_epochs, eta_min=1e-5
)
else:
scheduler = None
losses = []
for epoch in range(n_epochs):
epoch_loss = 0.0
for X_batch, y_batch in dataloader:
optimizer.zero_grad()
pred = net(X_batch).squeeze()
loss = criterion(pred, y_batch)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss / len(dataloader))
if scheduler is not None:
scheduler.step()
return losses
n_epochs = 60
losses_step = train_with_schedule("step", n_epochs=n_epochs)
losses_exp = train_with_schedule("exponential", n_epochs=n_epochs)
losses_cosine = train_with_schedule("cosine", n_epochs=n_epochs)
losses_constant = train_with_schedule("constant", n_epochs=n_epochs)Final training loss (epoch 60): Step decay: 0.00985 Exponential: 0.01002 Cosine annealing: 0.00997 Constant LR: 0.01206
The constant learning rate model shows the highest final loss because its updates remain large enough to keep it moving around the minimum. All three decay schedules finish lower in this run. The logarithmic axis makes those late-stage differences visible without hiding the much larger loss reduction during the first few epochs.
Interpreting the Convergence Curves
The convergence plot reveals more than just final loss values. Notice how each schedule's characteristic shape shows up in the training loss.
For step decay, scheduled cuts at epochs 20 and 40 reduce the size of subsequent updates. In this noisy mini-batch run, those transitions are subtler than the learning-rate staircase itself; the clearest effect is the tighter band of losses after the first cut. A large loss drop at a schedule boundary can indicate that the previous rate was causing oscillation, but it is not guaranteed in every run.
For exponential decay, the loss falls smoothly throughout. There are no phase transitions visible in the loss curve, which is a double-edged sword: smooth training is easy to monitor, but you lose the visual feedback that phase boundaries provide. If exponential decay is decaying too fast, you will see the loss plateau early, while the learning rate is still non-negligible, with no obvious cause.
For cosine annealing, the loss curves similarly smoothly. The cosine schedule's gradual start means the learning rate at the very beginning is nearly identical to the initial value, so there is no initial benefit from decay. The benefit accumulates through the middle and end of training.