Learning Rate Drives Everything
You can get the architecture right, the data pipeline right, the initialization right, and still fail if the learning rate schedule is wrong. I've watched identical models with identical data train to 5% accuracy with one schedule and 90% with another. The learning rate isn't just a hyperparameter — it's the single biggest lever you have over training dynamics.
Warmup: The Non-Negotiable Phase
For transformer training, warmup isn't optional. Starting with a large learning rate when the model is randomly initialized produces gradient spikes that can permanently corrupt early layer representations. The optimizer's adaptive statistics (Adam's second moment) haven't accumulated enough history to normalize these spikes.
Warmup lets the optimizer build up accurate gradient statistics while the learning rate is low enough that individual noisy updates don't cause damage. In practice, 1-5% of total training steps as warmup is sufficient. For a 100K-step training run, I'd use 1000-5000 warmup steps.
def get_warmup_scheduler(optimizer, warmup_steps, total_steps):
def lr_lambda(step):
if step < warmup_steps:
return float(step) / float(max(1, warmup_steps))
return 1.0
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
Cosine Annealing
After warmup, cosine annealing decays the learning rate following a half-cosine curve from the peak to some minimum (usually 0 or 1/10th of the peak). The decay is slow at the start, fast in the middle, and slow again near the end.
The slow initial decay matters. Right after warmup, the model is making large representational changes. A high learning rate during this phase helps escape bad local minima. The rapid middle-phase decay then refines the solution, and the slow final decay allows fine-tuning without overshooting.
def cosine_with_warmup(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
import math
def lr_lambda(step):
if step < warmup_steps:
return float(step) / float(max(1, warmup_steps))
progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
return min_lr_ratio + (1.0 - min_lr_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
I've found min_lr_ratio=0.1 (minimum LR is 10% of peak) works better than decaying to zero. Models that train for many epochs benefit from maintaining some learning capacity in the late phase, especially when there's label noise or distribution shift in the data.
We covered a related topic in Batch Normalization vs Layer Normalization in Production Tra.
Linear Decay
Linear decay is what it sounds like: a straight line from peak LR to minimum after warmup. It's the default in HuggingFace's Trainer and works surprisingly well across a wide range of tasks. Compared to cosine, linear decay is more aggressive in the early phase (drops faster) and less aggressive in the late phase.
For fine-tuning pre-trained models on small datasets (under 100K examples), linear decay consistently outperforms cosine in my experiments. The logic: fine-tuning doesn't need a long high-LR exploration phase because the pre-trained weights are already near a good solution. Quick decay to a low LR preserves the pre-trained features while adapting the top layers.
Cyclic Learning Rates
Cyclic schedules oscillate the learning rate between a minimum and maximum over repeating cycles. The idea is that periodic high-LR phases help the model escape sharp local minima, leading to solutions that generalize better.
Smith's 1cycle policy is the most practical cyclic variant. It does a single cycle: ramp up from low LR to peak over the first 30% of training, then cosine decay for the remaining 70%. Combined with momentum cycling (high momentum when LR is low, low momentum when LR is high), it often finds good solutions in fewer epochs than static schedules.
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=3e-4,
total_steps=total_steps,
pct_start=0.3, # 30% warmup
anneal_strategy='cos',
div_factor=25.0, # initial_lr = max_lr / 25
final_div_factor=1e4, # final_lr = initial_lr / 10000
)
I use 1cycle for training from scratch when I'm not sure about the optimal peak learning rate. It's robust to peak LR choices — the ramp-up phase acts like an implicit LR range test, and the schedule self-corrects somewhat if the peak is slightly too high.
For a related perspective, see Text Classification Pipeline Architecture for Multi-Label Pr.
Warmup-Stable-Decay (WSD)
This schedule, popularized by some recent LLM training papers, has three explicit phases: warmup to peak, hold at peak for a long stable phase, then rapid decay. The stable phase can be 80%+ of training.
The advantage for LLM training is predictability. During the stable phase, loss decreases approximately linearly on a log scale. This makes it easy to estimate when the model will reach a target loss and budget GPU time accordingly. The decay phase typically lasts 10-20% of training and squeezes out the last few points of perplexity.
Our team used WSD for a 3B parameter language model, and the training loss curves were remarkably smooth compared to cosine annealing. The total training time to reach the same final loss was about 5% longer, but the predictability saved us from two checkpoint-and-restart cycles we would've needed with cosine (where you can't easily extend training without restarting the schedule).
Per-Parameter Group Scheduling
Not all parameters should follow the same LR schedule. When fine-tuning, lower layers need lower learning rates than upper layers to preserve pre-trained features. The "discriminative fine-tuning" approach assigns different LR multipliers to different layer groups.
def get_layer_lrs(model, base_lr, decay_factor=0.9):
param_groups = []
layers = list(model.named_parameters())
num_layers = len(set(n.split('.')[1] for n, _ in layers if 'layer' in n))
for name, param in model.named_parameters():
if 'embedding' in name:
lr = base_lr * (decay_factor ** num_layers)
elif 'layer' in name:
layer_idx = int(name.split('.')[1])
lr = base_lr * (decay_factor ** (num_layers - layer_idx))
else:
lr = base_lr
param_groups.append({'params': [param], 'lr': lr})
return param_groups
This technique consistently improves fine-tuning results. I've measured 1-3% accuracy improvements on text classification tasks and 0.5-1.0 BLEU improvement on translation, simply by giving lower layers a 10× smaller learning rate.
See also: CI/CD Pipelines for Machine Learning: Testing Models Before .
Practical Recommendations
For pre-training large models: cosine with warmup (1-5% warmup, min_lr_ratio 0.1). It's what works reliably at scale, and everyone's hyperparameter grids are calibrated around it.
For fine-tuning: linear decay with very short warmup (0.5-1% of steps). The model doesn't need extensive exploration — just careful descent from the pre-trained initialization.
For small-dataset training from scratch: 1cycle. Its implicit LR range test and aggressive schedule make it robust to hyperparameter misspecification.
And regardless of schedule, always log the learning rate alongside your loss curves. When training goes wrong, the LR schedule is often the first thing to check.