Learning Rate Decay › Implementation
Language AI HandbookAll chapters Learning Rate Decay › ImplementationOne section · Sections in this chapter9 / 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
Let us implement all four schedules using PyTorch's scheduler API and verify they behave as expected.
import torch.nn as nn
import torch.optim as optim
class TinyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
model = TinyModel()
optimizer_step = optim.SGD(model.parameters(), lr=0.1)
scheduler_step = optim.lr_scheduler.StepLR(
optimizer_step, step_size=30, gamma=0.1
)
optimizer_exp = optim.SGD(model.parameters(), lr=0.1)
scheduler_exp = optim.lr_scheduler.ExponentialLR(optimizer_exp, gamma=0.955)
optimizer_cos = optim.SGD(model.parameters(), lr=0.1)
scheduler_cos = optim.lr_scheduler.CosineAnnealingLR(
optimizer_cos, T_max=100, eta_min=1e-4
)
optimizer_cosr = optim.SGD(model.parameters(), lr=0.1)
scheduler_cosr = optim.lr_scheduler.CosineAnnealingWarmRestarts(
optimizer_cosr, T_0=50, T_mult=1
)def record_lr_schedule(optimizer, scheduler, n_steps=100):
lrs = []
for _ in range(n_steps):
lrs.append(optimizer.param_groups[0]["lr"])
scheduler.step()
return lrs
torch_step_lrs = record_lr_schedule(optimizer_step, scheduler_step)
torch_exp_lrs = record_lr_schedule(optimizer_exp, scheduler_exp)
torch_cos_lrs = record_lr_schedule(optimizer_cos, scheduler_cos)
torch_cosr_lrs = record_lr_schedule(optimizer_cosr, scheduler_cosr)Step decay LRs at epochs [0, 30, 60, 90]: Epoch 0: 0.100000 Epoch 30: 0.010000 Epoch 60: 0.001000 Epoch 90: 0.000100 Cosine LRs at epochs [0, 25, 50, 75, 99]: Epoch 0: 0.100000 Epoch 25: 0.085370 Epoch 50: 0.050050 Epoch 75: 0.014730 Epoch 99: 0.000125
The step decay drops exactly at epoch 30 and 60 as configured. The cosine decay smoothly interpolates from 0.1 to near 0.0001 over 100 epochs.
For the inverse square root schedule with warmup, PyTorch does not provide a built-in, but it is straightforward to implement using LambdaLR:
def make_inv_sqrt_lambda(warmup_steps):
def lr_lambda(t):
if t < warmup_steps:
return t / warmup_steps
return (warmup_steps**0.5) / (t**0.5)
return lr_lambda
model_inv = TinyModel()
optimizer_inv = optim.AdamW(model_inv.parameters(), lr=0.001)
scheduler_inv = optim.lr_scheduler.LambdaLR(
optimizer_inv, lr_lambda=make_inv_sqrt_lambda(warmup_steps=100)
)
torch_inv_lrs = record_lr_schedule(optimizer_inv, scheduler_inv, n_steps=500)Inverse sqrt LRs at steps [0, 50, 100, 200, 499]: Step 0: 0.000000 Step 50: 0.000500 Step 100: 0.001000 Step 200: 0.000707 Step 499: 0.000448
The learning rate climbs linearly during warmup (steps 0 to 100), peaks at step 100, and then decays as the inverse square root of the step count.
Key Parameters
The key parameters to configure for each schedule are:
- Step decay:
step_size(epochs between drops),gamma(multiplicative factor, typically 0.1 or 0.5) - Exponential decay:
gamma(per-step multiplier, typically 0.99 to 0.999 for step-wise or 0.1 to 0.5 per epoch) - Inverse square root:
warmup_steps(typically 1 to 10 percent of total steps),eta_peak(the peak learning rate) - Cosine annealing:
T_max(period length),eta_min(minimum value, often ) - Warm restarts:
T_0(initial period),T_mult(cycle growth factor, 1 for fixed cycles, 2 for doubling)
Using Hugging Face Schedulers
The Hugging Face transformers library provides convenient schedule constructors that are particularly useful for fine-tuning language models:
from transformers import (
get_cosine_schedule_with_warmup,
get_linear_schedule_with_warmup,
)
model_hf = TinyModel()
optimizer_hf = optim.AdamW(model_hf.parameters(), lr=2e-5)
hf_warmup_steps = 100
hf_total_steps = 1000
scheduler_hf_linear = get_linear_schedule_with_warmup(
optimizer_hf,
num_warmup_steps=hf_warmup_steps,
num_training_steps=hf_total_steps,
)
model_hf2 = TinyModel()
optimizer_hf2 = optim.AdamW(model_hf2.parameters(), lr=2e-5)
scheduler_hf_cosine = get_cosine_schedule_with_warmup(
optimizer_hf2,
num_warmup_steps=hf_warmup_steps,
num_training_steps=hf_total_steps,
)
hf_linear_lrs = record_lr_schedule(
optimizer_hf, scheduler_hf_linear, n_steps=hf_total_steps
)
hf_cosine_lrs = record_lr_schedule(
optimizer_hf2, scheduler_hf_cosine, n_steps=hf_total_steps
)HuggingFace linear schedule at key steps: Step 0: 0.00000000 Step 50: 0.00001000 Step 100: 0.00002000 Step 500: 0.00001111 Step 999: 0.00000002
These functions handle the warmup transition automatically. The num_warmup_steps parameter sets the warmup duration, and num_training_steps sets the total. The schedule is then applied via scheduler.step() after each optimizer step, exactly like any PyTorch scheduler.
