Adding a Classification Head · Build an LLM From Scratch

With the data ready, it’s time for model surgery: start from the pretrained GPT, freeze its hard-won knowledge, and graft on a tiny new output layer that speaks “spam / not spam” instead of “next token”.

6.4 Initializing a model with pretrained weights

The starting point is chapter 5’s ending: a GPTModel filled with OpenAI’s GPT-2 weights.

</> Initialize and load — the ch5 machinery, reused

model = GPTModel(BASE_CONFIG)  #A
load_weights_into_gpt(model, params)  #B
model.eval();

#A BASE_CONFIG mirrors GPT_CONFIG_124M with two tweaks visible in the architecture print below: qkv_bias=True (the GPT-2 checkpoint HAS query/key/value biases — the flag threaded through since ch3 exists for this moment) and dropout 0.0 (regularization off for fine-tuning).
#B The §5.5 custom mapping pours OpenAI's tensors into our module tree.

print(model) shows the whole machine we built across chapters 2–4 — worth a look now that surgery is about to start:

</> The printed architecture (excerpt)

GPTModel(
(tok_emb): Embedding(50257, 768)
(pos_emb): Embedding(1024, 768)
(drop_emb): Dropout(p=0.0, inplace=False)
(trf_blocks): Sequential(
  (0): TransformerBlock(
    (att): MultiHeadAttention(
      (W_query): Linear(in_features=768, out_features=768, bias=True)
      (W_key): Linear(in_features=768, out_features=768, bias=True)
      (W_value): Linear(in_features=768, out_features=768, bias=True)
      (out_proj): Linear(in_features=768, out_features=768, bias=True)
      (dropout): Dropout(p=0.0, inplace=False)
    )
    (ff): FeedForward(
      (layers): Sequential(
        (0): Linear(in_features=768, out_features=3072, bias=True)
        (1): GELU()
        (2): Linear(in_features=3072, out_features=768, bias=True)
      )
    )
    (norm1): LayerNorm()
    (norm2): LayerNorm()
    (drop_resid): Dropout(p=0.0, inplace=False)
  )
  ... # 11 more identical blocks
)
(final_norm): LayerNorm()
(out_head): Linear(in_features=768, out_features=50257, bias=False)  #A
)

#A The 50,257-way next-token head — the one module about to be replaced.

6.5 The surgery: freeze, swap, unfreeze

Three moves turn the text generator into a classifier. Step through them on the model diagram:

Freeze the whole model

for param in model.parameters(): param.requires_grad = False — every pretrained weight becomes read-only. The language knowledge is locked in place; gradient descent can no longer touch (or damage) it.

1 / 3

The book's p. 99 figure as an operation: freeze everything, replace the head, then make the top of the network trainable again.

</> The three moves in code

# 1) freeze everything
for param in model.parameters():
  param.requires_grad = False

# 2) replace the output layer
torch.manual_seed(123)
num_classes = 2
model.out_head = torch.nn.Linear(
  in_features=BASE_CONFIG["emb_dim"],  # 768
  out_features=num_classes)

# 3) unfreeze the last block + final norm
for param in model.trf_blocks[-1].parameters():
  param.requires_grad = True
for param in model.final_norm.parameters():
  param.requires_grad = True

Why the LAST token carries the verdict

The model still outputs one vector per input token — now 2-dimensional. For the 4-token input “Do you have time”, the output is a 4×2 tensor. Which row should the classifier read? Chapter 3’s causal mask answers:

Left: the causal mask (ch3) — each row shows what that token may attend to. Right: the book’s p. 100 output tensor for “Do you have time”. Row 4 is the only representation informed by the entire input, so logits[:, -1, :] is the classification read-out.

Key idea — a thin trainable cap on a frozen foundation

The classification head (768 → 2) plus the unfrozen last block and final norm are the only parameters that learn; layer freezing protects everything else. And because causal attention makes the last token the sequence’s summary, classification is just: run the frozen tower, read the last token’s 2 logits, argmax. The next unit trains exactly that.

📝 Check yourself: the classification head

0 / 5

1.Why start fine-tuning from PRETRAINED weights instead of a random GPTModel?
2.What does the freeze-then-unfreeze strategy train, exactly?
3.The replaced head is Linear(768 → 2) instead of Linear(768 → 50257). What do the two outputs mean?
4.Why is the LAST token's output the one fed to the classifier?
5.The printed architecture shows the attention Linears with bias=True, unlike ch4's qkv_bias=False. Why?