Capstone Lesson 38: Classifier Fine-Tuning by Head Swap
Track B's first capstone. A pretrained language model is a stack of self-attention blocks ending in a token-prediction head. When you want spam vs ham, the head is wrong but the body is mostly right. This lesson rips the head off, glues a two-class linear layer onto the pooled representation, and trains the classifier two different ways: final-layer only, and full fine-tuning. The eval is precision, recall, and F1 on a held-out split. You learn what each strategy buys you and what it costs. Replace a language-model head with a classification head without re-initialising the body. Implement two training regimes: frozen body (head-only) and full fine-tuning, sharing one training loop. Build a tokeniser-aware data pipeline that pads, masks padding, and pools attention output. Compute precision, recall, F1, and a confusion matrix from raw logits. Reason about the trade-off between parameter count, training time, and head-room. You pre-trained a small transformer on a generic corpus. The output head projects the last hidden state to a 1000-token vocabulary. You now 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…
Capstone Lesson 38: Classifier Fine-Tuning by Head Swap: Track B's first capstone. A pretrained language model is a stack of self-attention blocks ending in a…
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.