7 — Teaching it one job

Chapter 6 ended with a model that writes fluent English and is useful for nothing in particular. It continues text the way web pages continue. Ask it a question and the reply is whatever tends to follow that question online.

The simplest way to make it useful is to give it one narrow job with a clear right answer. Here the job is deciding whether a text message is spam. The obvious first attempt is to ask. Each message goes into a prompt like this:

Text message: "Yeah whatever lol"
Is this message spam? Answer yes or no.
Answer:

Then compare the probability the model gives to " yes" against " no" as the next token. On 262 held-out messages, half spam and half not, this scores 48.9%, worse than a coin flip. GPT-2 says “yes” to 259 of the 262. With the prompt worded “Question: Is this text message spam?” it says “yes” to all 262. Left to generate, it doesn’t answer at all; it echoes the message back in quotation marks. Nothing in its training taught it that a question followed by “Answer:” should get an answer.

So this chapter changes the model instead of the prompt. It removes the output head that scores 50,257 possible next tokens, puts in one that scores two classes, spam and not spam, and trains on labeled examples. This is classification fine-tuning: start from weights that already know a lot about language, and adjust them for one decision.

Everything ran on an Apple M3 Pro with 18 GB of memory, on its GPU through PyTorch’s MPS backend, with Python 3.12, PyTorch 2.15.0 and tiktoken 0.14.0. The model is chapter 6’s: the chapter 4 GPT filled by load_gpt2 with OpenAI’s GPT-2 124M weights. Seeds are 123 throughout.

The data

The labeled examples come from the SMS Spam Collection, assembled by Tiago Almeida and José María Gómez Hidalgo in 2011 and published in the UCI Machine Learning Repository under a Creative Commons Attribution 4.0 license. It is one 477,907-byte text file, one message per line, a label, a tab, and the message:

ham	Ok lar... Joking wif u oni...
spam	Free entry in 2 a wkly comp to win FA Cup final tkts 21st May 2005. Text FA to 87121 to receive entry question(std txt rate)T&C's apply 08452810075over18's

“Ham” is the dataset’s word for a legitimate message. The 5,574 lines hold 4,827 ham and 747 spam. Most of the ham comes from a corpus of messages donated by students in Singapore and most of the spam from a UK forum where people reported it, so the two classes differ in more than intent: the spam is mostly British, from the 2000s, full of premium-rate numbers and prize claims, and the ham is full of Singaporean English and texting shorthand. Some messages contain <#> and <DECIMAL>, HTML-escaped placeholders that appear to stand in for numbers removed from the original messages. These are left as they are; the model sees whatever tokens the raw text produces.

Two problems have to be fixed before training.

Duplicates. Some messages appear more than once, word for word: only 5,171 of the 5,574 are distinct. If a message is in both the training set and the test set, getting it right on the test is memory, not skill. With the same seed and proportions but no deduplication, 27 of the 300 test messages also appear, word for word, in the training set. Dropping repeats first fixes that. No message appears with both labels, so nothing is lost. That leaves 4,518 distinct ham and 653 distinct spam.

Imbalance. 87% of the messages are ham. A model that says “ham” to everything would score 87% accuracy while catching no spam at all, and training on such data pushes a model toward that answer. The fix used here is the bluntest one: keep all 653 spam and a random 653 of the ham. That throws away most of the ham, which is wasteful, but it makes 50% the score of a model that knows nothing, and makes accuracy easy to read.

The 1,306 balanced messages are shuffled and split 70/10/20: 914 for training, 130 for validation, used to watch training, and 262 for the test, used once at the end.

def load_splits(seed=123):
    rows, seen = [], set()
    for line in open(PATH, encoding="utf-8"):
        label, text = line.rstrip("\n").split("\t", 1)
        if text in seen:                      # the file repeats some messages verbatim
            continue
        seen.add(text)
        rows.append((text, 1 if label == "spam" else 0))
    spam = [r for r in rows if r[1] == 1]
    ham = [r for r in rows if r[1] == 0]
    rng = random.Random(seed)
    ham = rng.sample(ham, len(spam))          # undersample ham to match spam
    data = spam + ham
    rng.shuffle(data)
    n = len(data)
    n_train, n_val = int(0.7 * n), int(0.1 * n)
    return data[:n_train], data[n_train:n_train + n_val], data[n_train + n_val:]

Label 1 is spam and 0 is ham. random.Random(seed) is a random number generator of its own, so the split is the same on every run regardless of anything else that uses randomness.

The bar to beat

Before training anything, it helps to know how far simple rules get, because a fine-tuned 124-million-parameter model that only matches a one-line rule hasn’t earned its cost.

Spam is long: the median training spam is 42 tokens and the median ham is 15. The rule “spam if longer than 28 tokens”, with 28 picked as the best threshold on the training set, gets 84.0% on the test set. Spam also asks to be called or texted at a number. The rule “spam if it contains five digits in a row” gets 89.3%. Those are the numbers to beat, not 50%.

A new head

The chapter 4 model ends in an output head: a linear layer from the 768 numbers of the final vector at each position to 50,257 scores, one per token. GPT-2 ties this head to the token embedding, so it is the same 50,257 × 768 matrix.

For classification the head becomes a linear layer from 768 numbers to 2. On a tiny scale, with a 4-number final vector h = [0.5, -1.0, 2.0, 0.0], a two-class head is a 2 × 4 weight matrix W and a bias b:

$\text{logits} = W h + b, \qquad W = \begin{bmatrix} 0.1 & 0.2 & -0.3 & 0.0 \\ -0.2 & 0.1 & 0.4 & 0.5 \end{bmatrix}, \quad b = \begin{bmatrix} 0 \\ 0 \end{bmatrix}$

Row one gives 0.05 − 0.2 − 0.6 + 0 = −0.75, the score for ham. Row two gives −0.1 − 0.1 + 0.8 + 0 = 0.6, the score for spam. Softmax turns [−0.75, 0.6] into probabilities [0.21, 0.79], so this head says spam with 79% confidence. Training adjusts W and b, and possibly the layers underneath, until the scores land on the right side for the training messages.

At the real size, W is 2 × 768 and b has 2 entries, 1,538 numbers in all, replacing a head that had 38.6 million.

model = load_gpt2("data/gpt2/model.safetensors", dropout=0.0)
torch.manual_seed(123)
model.head = nn.Linear(768, 2)          # new output head: 2 classes instead of 50,257 tokens
model.to("mps")

Assigning a new module to model.head replaces the old one; since the old head was the token embedding under another name, the embedding itself stays, and only the link to it is gone. nn.Linear(768, 2) creates a fresh layer with small random weights and biases. The model now returns, at every position, two scores instead of 50,257.

Which position to read

A message of 20 tokens gives 20 output vectors, so 20 pairs of scores. The classification should come from one of them, and the right one is the last.

The reason is the causal mask from chapter 3. Take a three-token message, “Call now free”. The attention weights form a lower triangle:

           Call   now   free
Call        ✓
now         ✓      ✓
free        ✓      ✓     ✓

The vector at “Call” has only seen “Call”. The vector at “now” has seen two tokens. Only the vector at “free”, the last, has attended to the whole message, so only it can summarize all of it. During pretraining that last vector was shaped to predict the next token, which means it already had to carry whatever about the message matters for what comes next. The new head reads that vector.

Messages in a batch have different lengths, so shorter ones are padded on the right with the end-of-text token, 50,256, to the length of the longest. Two things keep the padding harmless. Because it comes after the real tokens and attention is causal, no real token ever attends to it. And the scores are read at each message’s own last real token, not at the end of the padded row:

def batches(rows, bs, shuffle, gen=None):
    order = torch.randperm(len(rows), generator=gen).tolist() if shuffle else list(range(len(rows)))
    for i in range(0, len(rows), bs):
        chunk = [rows[j] for j in order[i:i + bs]]
        ids = [enc.encode(t)[:128] for t, _ in chunk]
        T = max(len(x) for x in ids)
        x = torch.full((len(ids), T), PAD)
        for r, s in enumerate(ids):
            x[r, :len(s)] = torch.tensor(s)
        last = torch.tensor([len(s) - 1 for s in ids])
        y = torch.tensor([lab for _, lab in chunk])
        yield x.to(device), last.to(device), y.to(device)

def class_logits(x, last):
    logits = model(x)                                            # (B, T, 2)
    return logits[torch.arange(len(x), device=device), last]     # (B, 2) at the last real token

torch.randperm(n, generator=gen) is a shuffled list of 0 to n − 1 drawn from a seeded generator, so the order of training examples is reproducible. Messages are cut at 128 tokens to bound memory; only one message in the balanced set is longer, at 132. torch.full(shape, value) makes a tensor filled with one value. The last line uses two index tensors together: logits[[0, 1, 2], [14, 31, 7]] picks row 0 at position 14, row 1 at position 31, row 2 at position 7, one pair of scores per message.

Before any training, with the new head’s random weights, the model scores 50.0% on the test set. It knows a great deal about language and nothing about which of its two outputs means spam.

Which layers to train

Every parameter in PyTorch has a flag, requires_grad, that says whether backpropagation computes a gradient for it. A parameter with the flag off is frozen: the optimizer can’t change it, and PyTorch doesn’t store the intermediate values needed to compute its gradient, which saves memory and time. Freezing is how a fine-tuning run chooses what to change.

for p in model.parameters():
    p.requires_grad = False
trainable = [model.head]
if mode in ("last", "all"):
    trainable += [model.final_norm, model.blocks[-1]]
if mode == "all":
    trainable = [model]
for m in trainable:
    for p in m.parameters():
        p.requires_grad = True

Three settings are worth comparing:

• Head only. Everything from chapter 6 stays exactly as it was, and only the 1,538 numbers of the new head learn. This tests whether the final vectors GPT-2 already produces separate spam from ham.
• Last block and head. The final transformer block and final layer norm train as well, 7,090,946 numbers. The lower blocks, which handle more general features of text, stay fixed; the top block can reshape its output for the new job.
• All layers. All 124,441,346 numbers train: GPT-2’s 124,439,808 plus the new head’s 1,538. The old head added nothing to that count, because it was the embedding matrix under a second name.

The training loop

Training is the loop from chapter 5, with two classes instead of 50,257. Cross-entropy is still the loss: the negative log of the probability the model gives to the correct class. In the tiny example above, if the message really is spam, the loss is −ln 0.794 = 0.23; if it’s ham, −ln 0.206 = 1.58. A model that always gives 50/50 has loss ln 2 = 0.693.

opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.1)
gen = torch.Generator().manual_seed(123)
for ep in range(3):
    model.train()
    for x, last, y in batches(train, 8, True, gen):
        loss = nn.functional.cross_entropy(class_logits(x, last), y)
        opt.zero_grad(); loss.backward(); opt.step()

Only parameters with requires_grad set are handed to the optimizer. Batches are 8 messages, so one pass over the 914 training messages, an epoch, is 115 steps, and each run is 3 epochs, 345 steps.

The learning rate used for the first comparison is 0.00005, a common choice for fine-tuning: small, because the weights are already good and large steps would wreck what pretraining built.

GPU memory is the total the MPS driver had allocated at the end of the run. Times are wall-clock training time, excluding evaluation.

The three runs at 0.00005 tell a clean story: the more layers train, the better. The last row breaks it. Training only the head reaches 97.7% once its learning rate is raised to 0.001. At 0.00005 it simply hadn’t finished learning in 345 steps: its validation accuracy was still rising, from 44.6% to 63.1% to 78.5% across the three epochs. A freshly initialized layer of 1,538 weights, with no pretrained weights in training to disturb, can take steps 20 times larger.

So GPT-2’s final vectors already separate spam from ham well enough that a straight line through 768-dimensional space, which is all a linear head is, gets 97.7% of test messages right. The model learned that during pretraining, from text that wasn’t labeled spam or anything else. Fine-tuning the head only finds the line.

Training more layers still helps a little, and the reason it’s only a little is in the size of the test set. With 262 test messages, each one is 0.38 percentage points. 98.9% is 3 errors, 97.7% is 6, 97.3% is 7. Differences of three or four messages between single runs with one seed are not strong evidence of anything. The difference between any of them and 89.3% for the digit rule, or 50% for prompting, is.

The training loss shows the same thing from another side. With all layers training, loss falls almost immediately and ends near 0.02. With only the last block training, it starts far higher, above 4, because the new head’s random weights produce large, confident wrong scores, and it takes most of the first epoch to come down, to about 0.4 by step 115 and 0.14 by the end. The head-only run at the higher learning rate is down to 0.4 within 40 steps and then levels off near 0.15: a single linear layer runs out of things it can change.

What it learned

Detailed results from here use the last-block model, the middle setting: 7 million trainable parameters, 38 seconds of training, 97.3%. A confusion matrix splits its 262 test answers by the true label and the predicted one:

It caught every spam message and wrongly flagged 7 real ones. For a spam filter that’s the costly direction, since a flagged real message may never be read. The balance between the two errors can be moved after training: the model outputs a probability, and requiring, say, 90% instead of 50% before calling a message spam trades missed spam for fewer false alarms. Which balance is right depends on what a mistake costs, not on the model.

The 7 false alarms:

You please give us connection today itself before  <DECIMAL>  or refund the bill
WOT STUDENT DISCOUNT CAN U GET ON BOOKS?
If you want to mapquest it or something look up "usf dogwood drive", that's the tiny street where the parking lot is
Doing my masters. When will you buy a bb cos i have for sale and how's bf
Check wid corect speling i.e. Sarcasm
You have registered Sinco as Payee. Log in at icicibank.com and enter URN  <#>  to confirm. Beware of frauds. Do NOT share or disclose URN to anyone.
Your account has been refilled successfully by INR  <DECIMAL> . Your KeralaCircle prepaid account balance is Rs  <DECIMAL> . Your Transaction ID is KR <#> .

Most are understandable. Two are automated messages from a bank and a phone company, which share spam’s vocabulary of accounts, payments and instructions to log in, and differ from it mainly in being genuine. Several others are about buying, selling, discounts or bills. All-caps shouting is common in the spam. These are the messages a person might pause over too.

The errors also barely depend on how much of the model trains. The all-layers model’s 3 errors are “WOT STUDENT DISCOUNT”, the ICICI bank message, and one spam it let through:

Do you ever notice that when you're driving, anyone going slower than you is an idiot and everyone driving faster than you is a maniac?

That message is labeled spam in the dataset. It reads as an ordinary forwarded joke. Perhaps it was part of a chain-message campaign; the label can’t be checked. Every labeled dataset has some of these, and once a model gets 97% right, a meaningful share of the remaining “errors” are disagreements with labels that are themselves arguable.

New messages, written for this test and run through the last-block model, give spam probabilities:

0.995  Congratulations! You have been selected to receive a free cruise. Call 09061701461 now to claim.
0.014  Are we still meeting for lunch at 1?
0.994  Your parcel could not be delivered. Reply YES to reschedule.
0.370  Free tonight? Call me when you get this
0.018  I won! The race was close but I won

The parcel message is a kind of scam that became common years after this dataset was collected, and the model flags it anyway; it shares the shape of the old spam, an unexpected notice demanding a reply. “Free tonight? Call me” has the words “free” and “call” and gets 0.37, uncertain but on the right side. “I won!” has “won” and is confidently ham.

The model computes two scores at every position, not just the last, so it is possible to watch the estimate change as the cruise message is read token by token. These intermediate readings come with a caveat: training only ever looked at the last position, so the head was never taught what to say mid-message, and it shows. The probability swings from 0.04 at “Congratulations” to 0.985 at " been", back to 0.05 at " selected", and stays above 0.9 only from " a" onward, and above 0.97 from " Call" to the end. The final answer is reliable; the path to it isn’t meant to be read.

Where this leaves things

GPT-2 can now do one useful thing. In under a minute of training on 914 labeled messages it tells spam from real messages with about 97% accuracy on messages it never saw, where asking it the same question in words scored worse than chance and the best one-line rule scored 89%. Most of the ability was already inside the pretrained model. A head with two outputs was enough to read it out.

The price is that it can do nothing else. Its output is two numbers now. It can’t explain a decision, can’t handle a third category without a new head and new training, and can’t write a sentence: the head that turned vectors into words is gone. Every new job means a new labeled dataset, a new head and a new run, and the model has to be told in advance which job it’s doing.

What would be more useful is a model that keeps producing text and does whatever a request in plain words asks: classify this, summarize that, answer this question. That needs a different kind of training example: not a message and a label, but an instruction and the response a person would want. Chapter 8 keeps GPT-2’s original head and trains it on a thousand of those.