llm_from_scratch / lessons/12 · classificationlesson 13 / 16

Part V — Fine-tuning

Classification fine-tuning — a spam detector

Lesson 11 handed us a real base model: OpenAI's GPT-2 weights loaded into our own modules, sampling text fluently with temperature and top-k. But a base model only continues text — ask it "spam or not?" and it will happily write a paragraph about spam rather than answer. This lesson repurposes that pretrained network into a decision-maker by swapping out one small piece: the head. We build a spam detector that reuses almost all of GPT-2's ~124M learned parameters and adds only a few thousand new ones.

What the last step left broken
The model from lesson 11 is a fluent completer, but completion is not a decision. Its output layer is the LM head: a matrix mapping each position's d-dimensional vector to V = 50,257 logits over the whole vocabulary. That machinery answers "what token comes next?", not "which of two categories is this message?" We could prompt it and parse the free text, but that is brittle and wasteful — we would be asking a text generator to role-play a classifier. We want a clean, calibrated p(spam) out of the same understanding the model already has.
Linear position
Forced by: a pretrained base model (11) completes text but will not make a decision — its d→V LM head only ranks next tokens, it has no notion of a class label.
New idea: repurpose the network by surgery on the head alone — freeze most of the pretrained body, delete the d→V LM head, and bolt on a tiny d→#classes classification head; read the last token's final vector (the only one that has attended to the whole message) and train it with the same cross-entropy loss from lesson 09.
Forces next: a spam detector solves exactly one fixed task with a fixed label set. We do not want to bolt a new head on for every job — we want one model that follows arbitrary instructions in plain language. Lesson 13 must teach the base completer to answer instructions directly.
The plan
Four moves. (1) Name the two families of fine-tuning — classification and instruction — and place this lesson in the first. (2) Make the surgical move: freeze the body, replace the LM head with a d→2 classification head. (3) Decide which token's vector feeds that head, and prove that for a causal model it must be the last one. (4) Train it — cross-entropy over classes, accuracy as the metric, a balanced labeled split, and why a handful of epochs is enough. Then run the head-swap yourself and watch the read-token choice decide whether the classifier is right or muddled.

1 · Two families of fine-tuning

Pretraining gives you a foundation model — a network that has read a great deal of text and, in the course of learning to predict the next token, has absorbed grammar, facts, and a surprising amount of world knowledge. Fine-tuning is the much cheaper second act: you take that foundation and nudge it toward a specific purpose. There are two broad ways to do it, and they answer different questions.

Family 1 · Classification
Map a text to one of a fixed set of labels
Spam vs. ham, positive vs. negative sentiment, which of five topics. The output is a category, not free text. The model's job collapses from "produce the next of 50,257 tokens" to "pick one of a few classes." This lesson.
Family 2 · Instruction
Follow an instruction phrased in natural language
"Summarize this," "translate that," "answer this question." The output is open-ended text, and one model handles arbitrarily many tasks by reading the instruction. This is how a base completer becomes an assistant — lesson 13.

Classification is the narrower, more contained problem, which makes it the right place to introduce fine-tuning: the change to the model is small and the success metric is unambiguous (you either labeled the message correctly or you did not). We will use the classic example — a spam detector trained on short text messages, each labeled spam (unwanted) or ham (legitimate). The mechanics you learn here — freeze the body, swap the head, read the right token, minimize cross-entropy — transfer directly to any classification task, and the reason it works so cheaply is the whole point: the base model already understands the language; we only teach it to route that understanding into a decision.

2 · The move: freeze the body, swap the head

Recall the shape of the assembled GPT from lesson 08. Text becomes token + positional embeddings, flows through L stacked transformer blocks and a final LayerNorm, and arrives as a stream of vectors: one d-dimensional vector per input position. The very last operation is the LM head, a linear map Whead of shape (d, V) that turns each position's vector into V next-token logits.

That head is the only part of the network that is intrinsically tied to the vocabulary. Everything below it — embeddings, attention, feed-forward, norms — produces general-purpose representations of the text. So the surgery is precise: keep the body, discard the LM head, and attach a new classification head in its place.

LM head:   Whead ∈ ℝ(d × V)   ⟶   class head:   Wcls ∈ ℝ(d × C),   C = 2

For a two-class spam detector, C = 2, so the new head is a (d, 2) matrix. Plug in GPT-2's d = 768 and it holds 768 × 2 = 1{,}536 weights (plus 2 biases). Compare that to the ~124 million parameters already in the body, or the 768 × 50{,}257 ≈ 38.6 million we just threw away with the old head. The trainable surface of the task can be almost nothing — a rounding error against the pretrained model.

Freeze most, unfreeze a little
To exploit the pretraining, we freeze the bulk of the body: its gradients are switched off, its weights never move, so the general language understanding is preserved intact. In practice you also unfreeze a little — commonly the final transformer block and the final LayerNorm, in addition to the new head — because the top of the network benefits from adjusting to the new objective. But the embeddings and the lower blocks stay put. The spectrum runs from "train only the head" (cheapest, sometimes enough) to "unfreeze everything" (full fine-tuning, most expressive, most expensive). The lesson's move is deliberately near the cheap end.

Why not just retrain the whole thing on spam labels? Two reasons. First, cost: updating 124M parameters needs far more compute and far more labeled data than updating a few thousand. Second, and more subtly, catastrophic forgetting — if you let every weight move freely on a tiny spam dataset, the model can overwrite the very language understanding that made it useful, overfitting the handful of examples and losing its grip on English. Freezing the body is not only cheaper; it is a regularizer that protects what pretraining paid for.

3 · Which token do we read?

Here is the question that trips people up. The body emits one vector per input position — for a message of T tokens, a (T, d) block of vectors. But the classification head wants a single d-vector to turn into two logits. Which of the T vectors do we hand it? First token? An average of all of them? The last token? The choice is not cosmetic — it decides whether the classifier can work at all.

The answer follows directly from the causal mask we built in lesson 06. In a causal (decoder-only) transformer, position t may attend only to positions ≤ t — never to the future. That constraint propagates through every layer. So consider what each position's final vector has actually "seen":

read strategywhat that vector has attended toverdict for a causal model
first tokenonly itself — position 0 can attend to nothing before itblind to the rest of the message; useless as a summary
mean of all tokensa blur: early positions saw little, late positions saw a lot, averaged togetherdilutes the one informed vector with many uninformed ones
last tokenevery position ≤ T-1 — i.e. the entire messagethe only vector that has integrated the whole sequence ✓

Only the final token's vector has, by the causal wiring, attended over all the words that came before it. It is the sequence's built-in summary slot — the one place where information from the whole message is allowed to pool. So for a causal LLM classifier, we read the last token's final-layer vector and feed it to the head:

logits = Wcls · hT-1   ∈ ℝ2   →   p = softmax(logits)   →   p(spam), p(ham)

where hT-1 is the last position's output vector. Two contrasts sharpen the intuition. In an encoder model like BERT — bidirectional, no causal mask, every token sees every other — people prepend a special [CLS] token at the front and read that, because in a bidirectional model even the first position sees the whole input. Our GPT is not bidirectional, so the front position is blind and the informed slot is at the back. And note we could pad each message to a fixed length; then "the last token" means the last real token before padding, since padding carries no message content.

The trap
Reading the first token (or, worse, a fixed padding token at the end) in a causal model feeds the head a vector that never saw the message. The classifier will train to near-chance and you will blame the data or the learning rate. The read-token choice is a correctness decision forced by the mask, not a hyperparameter to sweep. The widget below lets you flip between first / mean / last and watch the prediction go from muddled to confident.

4 · Training the classifier

With the head swapped and the read-token fixed, training is almost entirely machinery you already built. The forward pass runs a message through the frozen body, reads the last vector, and produces two logits. The loss is the same one from lesson 09 — cross-entropy — only now the target distribution is one-hot over two classes instead of over the vocabulary:

loss = − log p(true class)   =   −log softmax(logits)y

Everything else is lesson 10's loop: batches of messages, forward, cross-entropy, backprop (gradients flow only into the unfrozen parts), an AdamW step, repeat. Because the body is frozen and already fluent, the model does not have to learn language from scratch — it only has to learn a linear boundary in a representation it already possesses — so a handful of epochs is usually enough, where pretraining took thousands of GPU-hours. That asymmetry is the entire economic argument for fine-tuning.

Data, metric, and the splits

A few practical points that decide whether the classifier is trustworthy:

Labeled data
Fine-tuning is supervised: each message needs a human (or curated) label, spam or ham. Far less data than pretraining — a few thousand messages can suffice — but it cannot be free next-token labels; someone had to say "this one is spam."
Balanced classes
If 90% of messages are ham, a lazy model scores 90% accuracy by always guessing "ham." Balance the training set (downsample the majority or upweight the minority) so accuracy actually measures skill, not the base rate.
Train / val / test
Split three ways: train to fit, validation to choose when to stop and tune, test (untouched until the end) for the honest final number. Reusing the test set to tune quietly inflates your reported accuracy.

The natural metric here is accuracy — the fraction of messages classified correctly — because the task has a single right answer per example. (For skewed real-world spam, precision and recall matter too: a false positive silences a real message, a false negative lets junk through. But start with accuracy on a balanced set.) Cross-entropy remains the training objective; accuracy is the human-readable report. They agree in direction — driving loss down drives accuracy up — but loss is what the gradient can act on, so loss is what we minimize.

5 · Swap the head yourself

The widget stages the whole move on one short, spammy message. The frozen body has already turned each word into a small final-layer vector (fixed, deterministic — no randomness here). The old LM head is drawn greyed-out and disconnected; the new 2-unit classification head is live. The one control that matters is read which token: switch between first, mean, and last and watch which vector gets routed into the head, the two class logits it produces, the softmax, and the verdict. Only last — the vector that attended over the whole message — yields a confident, correct "spam." First is blind; mean is muddy. The KPIs report the tiny size of the new head against the frozen body, and the resulting p(spam).

Head-swap classifier
Proves two things at once: fine-tuning can be a surgical change (a d×2 head, ~1.5k weights, against ~124M frozen), and in a causal model the last token is the sequence summary — read the first or the mean and the classifier is blind or blurred.
new head params (d×2)
1,536
frozen body params
≈124M
p(spam)
prediction

Failure modes & checklist

Failure modes

  • Reading the wrong token. The first token (or a trailing pad token) has not seen the message in a causal model. Signal: accuracy stuck near chance even as loss for other setups falls.
  • Unfreezing everything on tiny data. Full fine-tuning on a few thousand messages overfits and can cause catastrophic forgetting of the pretrained language ability.
  • Imbalanced classes. A 90/10 split lets "always predict the majority" score 90%; balance the set or the accuracy number is a lie.
  • Tuning on the test set. Peeking at test accuracy to pick epochs or thresholds silently inflates the final report; keep test untouched.
  • Keeping the LM head. Trying to parse a class out of 50,257 next-token logits is brittle; replace the head with a clean d→C map.

Checklist

  • Load pretrained weights (lesson 11) into the body; freeze most of it.
  • Replace the (d, V) LM head with a (d, C) classification head.
  • Read the last real token's final-layer vector → head → logits.
  • Cross-entropy over classes (lesson 09); AdamW loop (lesson 10); a few epochs.
  • Balanced classes; train / val / test split; report accuracy on held-out test.

Checkpoint

Try it
Start with the read-token on last and note the confident spam verdict and the high p(spam). Now switch to first: the head is now fed the opening token's vector, which — under the causal mask — never saw the words "win," "prize," or "click," and the prediction slides toward a coin-flip. Switch to mean: better than first, because a few informed positions are in the average, but still blurred by all the uninformed early ones. You have just demonstrated, by hand, why every causal-LLM classifier reads the last token. Ask yourself: which of lesson 00's three properties — accurate, trainable, steerable — is this head-swap exploiting?

Where this points next

The spam detector is a genuine win: near-free reuse of a pretrained model, turned into a calibrated decision by a head with barely a thousand weights. But look at its ceiling. It does exactly one thing, with a hard-wired label set of size two. Want sentiment instead? Bolt on a different head and retrain. Want it to also summarize, translate, and answer questions? You cannot — a fixed d→C head has no way to express open-ended output, and you would need a new head per task. What we actually want is a single model that reads a task described in plain language and does it. That means going back to text generation, but teaching the completer to treat "instruction → response" as the pattern to continue. That is lesson 13, Instruction fine-tuning — completer to assistant.

Takeaway
Classification fine-tuning repurposes a pretrained GPT with minimal surgery: freeze the body to preserve its language understanding, replace the (d, V) LM head with a small (d, C) classification head (just d×C new weights — ~1,536 for a 2-class GPT-2), and read the last token's final vector, which is the only one the causal mask lets attend over the entire message. Train it with the same cross-entropy loss and AdamW loop as pretraining, on a balanced, labeled, three-way-split dataset; because the base already "understands" language, a few epochs suffice and accuracy is the metric. The whole point: fine-tuning a huge model into a task can cost almost nothing — but only if you read the right token and don't clobber the pretraining.

Interview prompts

Companion reads: gpt_mini · 03 SFT and cs336 · 16 SFT cover supervised fine-tuning of an assumed architecture; this lesson is the from-scratch version of the same head-swap-and-train idea.