Capstone Lesson 38: Classifier Fine-Tuning by Head Swap

A head swap takes a pretrained model, discards its output layer, and attaches a new one for a different task while keeping the body's weights. Here is the case for it. You pre-trained a small transformer on a generic corpus, and its output head projects the last hidden state to the pretraining vocabulary (say 1000 tokens, for a hypothetical pretrained model; the demo's own body uses a 260-id byte vocabulary). Now you have 800 SMS messages labelled spam or ham and you want a binary classifier. Three options exist.

The wrong option is to train a fresh classifier from scratch on 800 examples. The body of the pretrained model already encodes useful structure: word identity, position, simple co-occurrence. Throwing it away wastes the compute that built it.

The two right options are head swap with the body frozen, and head swap with the body trainable. Head-only training is fast, almost free in memory, and rarely overfits with this little data. Full fine-tuning is slower and can overfit on small data, but it reaches higher accuracy when the downstream domain drifts from the pretraining corpus.

We build both and compare them on the same fixture.

Fine-tuning by head swap means keeping the pretrained body and replacing the function that sits on top of it. The body is a function f_theta(tokens) -> hidden_states. The head is a function g_phi(hidden) -> logits. Swapping heads keeps theta and replaces g_phi. The body's parameters are the expensive part, while the head is a single linear layer.

Two trainable parameter sets matter:

• theta (the body): tens of thousands of weights per attention block.
• phi (the head): hidden_dim * num_classes weights plus a bias.

Head-only training computes gradients against phi and none against theta. In PyTorch you get this by setting requires_grad=False on the body parameters, so the optimiser sees only the head and the body stays frozen.

Full fine-tuning lets gradients flow back through the whole stack, and the body's weights drift to fit the classification objective. The risk on small data is catastrophic forgetting: the body's pretraining gets washed out as it overfits noise.

A classifier needs one vector per sequence, not one vector per token, and pooling is the step that collapses the per-token hidden states into that single vector. Three common choices:

• Mean pool: average the hidden states across the sequence, weighted by the attention mask.
• CLS pool: prepend a special token and use only its output. This is what BERT does.
• Last-token pool: use the last non-padding token. This is what GPT-class classifiers do.

We use mean pooling with explicit attention-mask weighting. It is the simplest option, gives a stable signal across sequence lengths, and does not require pretraining a CLS token.

The data is a synthetic SMS fixture of eight hundred messages, balanced 400 spam and 400 ham, generated deterministically in code/main.py. The generator uses a fixed seed, picks templates, substitutes slot fillers, and emits short messages. The byte-level tokenizer encodes one id per byte and truncates or pads each message to a max length of 32. Real datasets have noise this fixture does not, but reproducibility is the point here.

The split is 80/20: 640 train, 160 test. It is stratified, so the test set keeps the 50/50 balance. Because the balance of the held-out set is known, precision and recall can be read as honest numbers.

These are binary classification metrics with class 1 (spam) as the positive class. Counts are:

• TP: predicted spam, was spam.
• FP: predicted spam, was ham.
• FN: predicted ham, was spam.
• TN: predicted ham, was ham.

The three headline metrics:

• precision = TP / (TP + FP). Of the messages flagged spam, what fraction actually are?
• recall = TP / (TP + FN). Of the actual spam, what fraction did the model flag?
• F1 = 2 * P * R / (P + R). The harmonic mean of the two.

A confusion matrix prints the four counts as a 2x2 grid. The demo writes one to stdout for each training regime.

The body is a deliberately tiny transformer: vocab 260, hidden 64, 4 heads, 2 blocks, max sequence 32. It is small enough to train both regimes on a CPU. It is not pretrained in the lesson. Instead, the pretrain_quick helper runs five epochs of LM training on the train split's text (never the test split) to give the body a non-trivial starting point, which keeps the lesson self-contained.

The implementation is one main.py plus one test module (code/tests/test_main.py).

1. ByteTokenizer: maps bytes to ids, reserves a pad id.
2. Block: a transformer block with multi-head attention and a feed-forward layer. Pre-norm.
3. LMBody: token + position embeddings plus a stack of blocks. Returns hidden states.
4. mean_pool: a function (not a class) that takes the mask-weighted average over the sequence axis.
5. Classifier: body, pool, linear head. Both regimes start from the same pretrained weights: _clone_body deep-copies the body so each regime trains its own copy.
6. freeze_body and unfreeze_body: toggle requires_grad on body parameters.
7. train_classifier(model, train_loader, epochs, lr, seed): one shared loop. It builds an Adam optimiser over the parameters with requires_grad set, and raises ValueError if none are trainable.
8. evaluate: runs the test set and returns Metrics(precision, recall, f1, tp, fp, fn, tn), whose confusion() method formats the 2x2 grid.
9. run_demo: pretrains the body briefly on the train split texts, then trains and evaluates head-only, then full, prints both reports, and returns 0. If either regime's F1 is not above 0.5 it prints an error to stderr and returns 1.

Head-only training usually trains faster and underfits more gracefully. The demo's only hard check is that both regimes beat random, meaning F1 above 0.5. Beyond that, typical runs land at strong precision and recall for both regimes, with exact numbers depending on the seed. Full fine-tuning costs more per epoch because gradients flow through the whole body, and it typically lands close to head-only on this fixture.

I don't declare a winner here. What matters is reading the numbers next to the cost. On 800 examples and a tiny body, head-only is the right call. On 80,000 examples and a bigger body, full fine-tuning starts to pay off. The contract to take away is the API: the same train_classifier function handles both regimes, and the toggle is one call.

• Add a third regime that unfreezes only the last block. This is sometimes called partial fine-tuning. It costs less than full FT and learns more than head-only.
• Add a learning-rate scheduler. A cosine schedule on the head plus a smaller constant rate on the body is a common production setup.
• Replace mean pooling with a learned attention pool: a small attention layer with one learned query. This often beats mean pool on longer sequences.

The implementation gives you the hooks. The tests pin the contract. The numbers are yours to push.