Unit 09.03: Fine-tuning boundaries and evaluation

Fine-tuning a pretrained encoder works — at a learning rate far below what feels natural.

Two orders of magnitude apart

The head is randomly initialised and needs large updates. The encoder holds pretrained structure and needs small ones. A single learning rate serves one of them badly.

Here the head runs at 1e-2 and the encoder at 1e-5 — a thousandfold difference, which is typical:

import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(240, 6, 16)
y = (X[:, 0, 0] + X[:, 1, 1] > 0).float().unsqueeze(1)
Xtr, ytr, Xva, yva = X[:180], y[:180], X[180:], y[180:]


def evaluate(unfreeze, steps=120):
    torch.manual_seed(0)
    encoder = nn.TransformerEncoderLayer(16, 4, 32, batch_first=True)
    head = nn.Linear(16, 1)
    for p in encoder.parameters():
        p.requires_grad = unfreeze
    groups = [{"params": head.parameters(), "lr": 1e-2}]
    if unfreeze:
        groups.append({"params": encoder.parameters(), "lr": 1e-5})   # small!
    opt = torch.optim.Adam(groups)
    for _ in range(steps):
        opt.zero_grad()
        nn.BCEWithLogitsLoss()(head(encoder(Xtr).mean(1)), ytr).backward()
        opt.step()
    with torch.no_grad():
        pred = torch.sigmoid(head(encoder(Xva).mean(1))) > 0.5
        return round((pred.float() == yva).float().mean().item(), 3)


print("frozen encoder     :", evaluate(unfreeze=False))
print("fine-tuned encoder :", evaluate(unfreeze=True))
print("majority baseline  :", round(max(yva.mean().item(), 1 - yva.mean().item()), 3))

# Fine-tune at a learning rate one or two orders of magnitude below the head's,
# or the pretrained weights are destroyed in the first few steps. And report
# the baseline: without it neither number means anything.

Compare the frozen and fine-tuned scores against the majority baseline. All three numbers are printed together, because either model in isolation is uninterpretable.

On a small dataset the frozen version often wins, and that is a legitimate result rather than a failed experiment.

The mistake this prevents

Fine-tuning at the head's learning rate. The pretrained weights are overwritten within a few dozen steps, and you have spent the compute to arrive at a randomly initialised model.

Takeaway

Fine-tune the encoder one to two orders of magnitude below the head, and always report the baseline.

Mark complete