Transfer learning

A new task with ten labels

Now a new question about the same kind of images: does the digit have a closed loop? (0, 6, 8 and 9 do.) Pretend we've labeled only ten of the 797 kept-aside images, five with loops and five without, and test on the rest. Three ideas compete:

• From scratch: a fresh network, trained on the ten examples.
• Frozen backbone, new head: load the saved backbone, freeze it so its weights can't change (requires_grad_(False)), put a new two-class head on top, and train only the head.
• Next section: fine-tuning, which lets the backbone move too, gently.

Ten examples is a small, lucky-or-unlucky sample, so the cell repeats the experiment with six different sets of ten and averages.

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import load_digits

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
order = np.random.default_rng(0).permutation(len(X))
old, new = order[:1000], order[1000:]
loops = torch.tensor(np.isin(digits.target, [0, 6, 8, 9]).astype(np.int64))   # 1 if the digit has a loop

def make_backbone():
    return nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU())

def train(model, X, y, epochs=100, lr=0.01):
    optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=lr)
    for _ in range(epochs):
        optimizer.zero_grad()
        nn.functional.cross_entropy(model(X), y).backward()
        optimizer.step()

def accuracy(model, X, y):
    model.eval()
    with torch.no_grad():
        return (model(X).argmax(dim=1) == y).float().mean().item()

torch.manual_seed(0)                                         # the first cell's backbone, trained again
source = make_backbone()
train(nn.Sequential(source, nn.Linear(32, 10)), X[old], torch.tensor(digits.target[old]), epochs=150)
saved = {name: t.clone() for name, t in source.state_dict().items()}

scores = {"from scratch": [], "frozen backbone": []}
for seed in range(6):
    rng = np.random.default_rng(seed)
    with_loop, without = new[loops[new].numpy() == 1], new[loops[new].numpy() == 0]
    few = np.concatenate([rng.permutation(with_loop)[:5], rng.permutation(without)[:5]])
    test = np.setdiff1d(new, few)
    torch.manual_seed(seed)

    backbone = make_backbone()
    backbone.load_state_dict(saved)
    backbone.requires_grad_(False)                           # freeze: the optimizer won't touch it
    frozen = nn.Sequential(backbone, nn.Linear(32, 2))       # a new head for the new task
    train(frozen, X[few], loops[few])

    scratch = nn.Sequential(make_backbone(), nn.Linear(32, 2))
    train(scratch, X[few], loops[few])

    scores["frozen backbone"].append(accuracy(frozen, X[test], loops[test]))
    scores["from scratch"].append(accuracy(scratch, X[test], loops[test]))

for name, accs in scores.items():
    print(f"{name:<16} mean {np.mean(accs):.3f}  runs {[round(a, 2) for a in accs]}")

With ten labels the frozen backbone gets about 78% right and a network trained from scratch about 68%: transfer is worth ten points here. Look at the runs too. From scratch swings by 16 points (0.59 to 0.75) depending on which ten images were labeled, while the frozen backbone stays between 0.74 and 0.80: steadier results are another benefit of transfer. Either way, one lucky run proves nothing. Notice what the head learned from: not pixels, but the backbone's features, which already separate a 0 from a 7.

import torch.nn as nn


def freeze_except(model, keep):
    return sum(p.numel() for p in model.parameters())


backbone = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU())
model = nn.Sequential(backbone, nn.Linear(32, 2))
print(freeze_except(model, ("1.",)))

Fine-tuning: let the backbone move, slowly

Freezing is safe but rigid: the features stay exactly as they were for digits. Fine-tuning trains the whole network on the new task, so the backbone can adjust its features too. The risk is that big updates wreck what it learned, especially with only a few examples to learn from. The usual fix is two learning rates: small for the backbone (a gentle nudge), larger for the new head (which starts from random weights and has everything to learn).

PyTorch optimizers take parameter groups for this: a list of dictionaries, each with its own params and lr.

optimizer = torch.optim.Adam([
    {"params": backbone.parameters(), "lr": 0.001},   # pretrained: small steps
    {"params": head.parameters(), "lr": 0.01},        # new: bigger steps
])

Run on twenty labels with the same method as the cell above (six draws each, averaged), the three approaches come out at about:

Fine-tuning edges ahead because the features can bend toward loops. With only ten labels the order flips: freezing scored 0.78 and fine-tuning 0.76, because there's too little data to move the backbone's weights sensibly. So freeze when labels are very few, and fine-tune when you have more. Two learning rates are a sensible default, not a law: at 150 labels in the same experiment, a single learning rate did slightly better. Measure on your own data.

One more detail when reusing a real network: call model.eval() before computing features you'll keep, and think about it during fine-tuning too. Layers like dropout and batch normalization behave differently in training mode, and batch normalization's running averages would otherwise drift toward your few new examples. Freezing doesn't prevent that: requires_grad_(False) stops gradient updates, but in training mode batch norm still updates its running averages, so put frozen layers in evaluation mode.

import torch
import torch.nn as nn


def fine_tune_optimizer(backbone, head, lr):
    return torch.optim.Adam(list(backbone.parameters()) + list(head.parameters()), lr=lr)


backbone = nn.Sequential(nn.Linear(64, 32), nn.ReLU())
head = nn.Linear(32, 2)
opt = fine_tune_optimizer(backbone, head, 0.01)
print([g["lr"] for g in opt.param_groups])

The cheapest transfer: features once, a simple classifier

If the backbone stays frozen, its features for each image never change, so there's no need to run it inside a training loop at all. Compute the features once, then train any simple classifier on them: logistic regression, from the classification lessons, fits in milliseconds. This is exactly how teams use an LLM's embeddings: send each text through the embedding model once, store the vectors, and fit a small classifier on top.

The cell compares logistic regression on raw pixels with logistic regression on the backbone's 32 features as the number of labels grows: a learning curve.

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
order = np.random.default_rng(0).permutation(len(X))
old, new = order[:1000], order[1000:]
loops = np.isin(digits.target, [0, 6, 8, 9]).astype(np.int64)

torch.manual_seed(0)
backbone = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU())
model = nn.Sequential(backbone, nn.Linear(32, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for _ in range(150):
    optimizer.zero_grad()
    nn.functional.cross_entropy(model(X[old]), torch.tensor(digits.target[old])).backward()
    optimizer.step()

backbone.eval()
with torch.no_grad():
    features = backbone(X).numpy()            # every image's 32 features, computed once
pixels = X.numpy()

print(" labels   pixels   features")
for n in [10, 20, 50, 150]:
    row = []
    for seed in range(6):
        rng = np.random.default_rng(seed)
        few = np.concatenate([rng.permutation(new[loops[new] == 1])[:n // 2], rng.permutation(new[loops[new] == 0])[:n // 2]])
        test = np.setdiff1d(new, few)
        row.append([LogisticRegression(max_iter=3000).fit(data[few], loops[few]).score(data[test], loops[test])
                    for data in (pixels, features)])
    print(f"{n:>7}   {np.mean(row, axis=0)[0]:.3f}    {np.mean(row, axis=0)[1]:.3f}")

Features beat pixels at every size, by the most when labels are scarcest: with ten labels, about 0.80 against 0.72. With more labels the gap shrinks to a few points, because with enough examples a model can learn useful features from the new data itself. That's the shape to expect: transfer matters most when labels are few. Be honest about this example, too: "has a loop" is just a grouping of the digits the backbone already knows, which makes it about the friendliest transfer there is. Real projects are rarely that kind.

When transfer hurts

Transfer works when the old task taught features the new task needs. When it didn't, reused features can be worse than none. Train the backbone only on digits 0 to 4 until it's perfect at them, then try to transfer it to telling 5, 6, 7, 8 and 9 apart with five labels each:

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import load_digits

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target)
low = np.flatnonzero(digits.target < 5)
high = np.random.default_rng(0).permutation(np.flatnonzero(digits.target >= 5))
few = np.concatenate([high[digits.target[high] == d][:5] for d in range(5, 10)])
test = np.setdiff1d(high, few)

def make_backbone():
    return nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU())

def train(model, X, y, epochs):
    optimizer = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=0.01)
    for _ in range(epochs):
        optimizer.zero_grad()
        nn.functional.cross_entropy(model(X), y).backward()
        optimizer.step()

def accuracy(model, X, y):
    model.eval()
    with torch.no_grad():
        return round((model(X).argmax(dim=1) == y).float().mean().item(), 3)

torch.manual_seed(0)
backbone = make_backbone()
train(nn.Sequential(backbone, nn.Linear(32, 5)), X[low], y[low], epochs=150)     # digits 0-4 only
backbone.requires_grad_(False)
frozen = nn.Sequential(backbone, nn.Linear(32, 5))
train(frozen, X[few], y[few] - 5, epochs=150)
scratch = nn.Sequential(make_backbone(), nn.Linear(32, 5))
train(scratch, X[few], y[few] - 5, epochs=150)
print("frozen backbone from digits 0-4:", accuracy(frozen, X[test], y[test] - 5))
print("from scratch:                   ", accuracy(scratch, X[test], y[test] - 5))

The transferred backbone loses badly, 0.57 against 0.90. Trained only to tell 0 to 4 apart, and trained until its loss was nearly zero, it kept exactly the features that separate those five digits and threw away the rest. Two different 7s and a 9 can look the same to it. This is negative transfer, and two habits guard against it: always compare with a from-scratch baseline (as the experiments here do), and prefer backbones trained on broad, varied data, which is why big pretrained models are trained on so much of it.