Unit 07.04: Fine-tuning schedule and learning rates

The backbone knows something already. The head knows nothing. Training both at the same rate is wrong for one of them.

Different rates for different parts

A pretrained backbone needs small, careful updates — its weights are already close to useful. A randomly initialised head needs large ones. PyTorch supports this directly through parameter groups, each with its own learning rate.

A scheduler then decays all of them together, preserving the ratio:

import torch
from torch import nn

torch.manual_seed(0)
backbone = nn.Sequential(nn.Linear(10, 32), nn.ReLU())
head = nn.Linear(32, 1)

# Discriminative learning rates: the pretrained body moves slowly, the fresh
# head moves quickly. One global rate either wrecks the body or starves the head.
opt = torch.optim.Adam([
    {"params": backbone.parameters(), "lr": 1e-4},
    {"params": head.parameters(), "lr": 1e-2},
])
for group in opt.param_groups:
    print("group lr:", group["lr"], " tensors:", len(group["params"]))

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
x, y = torch.randn(64, 10), torch.randn(64, 1)

print("\nepoch  backbone_lr   head_lr")
for epoch in range(10):
    opt.zero_grad()
    nn.MSELoss()(head(backbone(x)), y).backward()
    opt.step()
    scheduler.step()
    if epoch % 3 == 0:
        print(f"{epoch:>5}  {opt.param_groups[0]['lr']:.2e}   {opt.param_groups[1]['lr']:.2e}")

# Both rates decay together, keeping their ratio. Call scheduler.step() once
# per epoch, after the optimiser -- calling it per batch decays ten times too fast.

The printed rates show both groups decaying along a cosine curve while the head stays roughly two orders of magnitude above the backbone throughout.

Note where scheduler.step() is called: once per epoch, after the optimiser. Calling it per batch decays the rate as many times as you have batches, which collapses it to nearly zero within the first epoch.

The mistake this prevents

Calling scheduler.step() inside the batch loop. The learning rate decays far too quickly, training stalls early, and the cause is invisible unless you print the rate.

Takeaway

Small rate for pretrained weights, large for new ones, and step the scheduler once per epoch.

Mark complete