LM6-L2 — Fine-Tune a Language Model for Classification

Objective

Implement a classification head with last-real-token pooling, a function that freezes layers by mode, the confusion matrix with precision, recall and F1, and an error finder, so that all five self-checks pass, and explain why the head-only model fails on the imbalanced task although its accuracy looks high.

Why a notebook? This exercise trains or runs a real PyTorch model. PyTorch does not run in the browser Python used by the site’s other labs, so you run this one on your own computer or in a hosted notebook service. The interactive browser labs of this module cover the surrounding ideas without PyTorch.

Setup and requirements

Hardware: CPU only, about 2 GB of RAM, a few minutes (tiny model); the book path with GPT-2 small is a GPU or patient-CPU option. CPU only; about 2 GB of RAM; the whole notebook takes about one minute. The backbone is a 107,000-parameter GPT pretrained for a few seconds inside the notebook on synthetic text, so nothing is downloaded. The link below gives the book’s chapter that fine-tunes the pretrained GPT-2 small, which needs more time and, ideally, a GPU.

Download the notebook (.ipynb) Reference solution (try the lab first)

On your own machine (Python 3.10 or newer):

python3 -m venv llm-env
source llm-env/bin/activate        # on Windows: llm-env\Scripts\activate
pip install torch numpy jupyterlab
jupyter lab fine-tune-a-classifier.ipynb

In a hosted notebook service such as Google Colab or Kaggle, upload the file and run it there; PyTorch is usually preinstalled. Matplotlib is used for optional plots when it is available; the notebook runs without it.

What you will do

1. Data. Read the synthetic imbalanced task (12% positives) and the short pretraining of the backbone (given).
2. Classifier. Add a linear head and pool the hidden state of the last real token of each sentence; check that padding after it has no effect.
3. Freezing. Write set_trainable for head-only, last-block and all-layer training and count the trainable parameters.
4. Metrics. Write the confusion matrix, precision, recall and F1; see why always answering 0 gets 88% accuracy and F1 of 0.
5. Compare. Train five variants and read the table, including training from scratch and pooling the wrong token.
6. Errors. List the sentences the partly fine-tuned model gets wrong and look for a pattern.

Cells that you must complete are marked TODO in the notebook and have a self-check right after them (an assert with a message). Run the cells from the top; a failing check tells you what is wrong.

Expected behaviour

• Self-check 1: the classifier returns logits of shape (batch, 2), uses the hidden state at position length − 1 and ignores the padding that follows.
• Self-check 2: head only trains 130 parameters; last block + head trains 50,050; all layers trains 105,410 (everything except the language-model head).
• Self-check 3: the confusion matrix and scores match a hand-computed example, and an always-0 predictor gets precision, recall and F1 of 0.
• Self-check 4: fine-tuning all layers reaches F1 of at least 0.9 on the held-out examples; pooling the first token gives F1 below 0.2. In the run used to prepare this lab, head-only training predicted 0 for everything (accuracy 0.892, F1 0), the last block reached F1 0.91, all layers F1 0.96 and training from scratch F1 0.91; with a single seed and 500 test sentences such differences are small and should not be over-read.
• Self-check 5: the error finder returns exactly the misclassified held-out sentences. In the prepared run, the nine errors of the last-block variant were all missed positives, and all involved the object “bin”: a pattern worth investigating before collecting more data.

How to know you are done

• All five self-checks print “passed” when the notebook is run from the top.
• You can explain why accuracy misleads on imbalanced data and what the F1 score adds.
• You can say why the last real token, and not the first, represents the sentence in a causal model.

The site cannot run or grade the notebook for you. The self-checks inside the notebook are your verification. The reference solution was executed with PyTorch on a CPU when this module was prepared. The task and the pretraining text are synthetic and the pretraining is a few seconds long, so the comparison between pretrained and from-scratch models says little about real tasks with large pretrained models.

If something goes wrong

• Index error when pooling: positions are lengths - 1 (zero-based). Use h[torch.arange(B), positions] to select one position per sequence.
• The output depends on the padding: you are using the last padded position instead of the last real token.
• Head-only training gives F1 = 0: that is expected on this task with so few trainable parameters; it is the point of the comparison.
• Results differ a little from the numbers above: random seeds and PyTorch versions change the last decimals; the assertions have margins.

Theory for this lab

• LM6-T1: Pretraining Versus Supervised Fine-Tuning
• LM6-T2: Classification Heads and Sequence Representations
• LM6-T3: Frozen and Trainable Parameters

Reference material

• Build a Large Language Model (From Scratch), Chapter 6: Finetuning for Classification

Fine-Tuning for Classification labs

1. LM6-L1: Classification Metrics and Failure Cases
2. LM6-L2: Fine-Tune a Language Model for Classification: this lab.

Next lab: LM7-L1: Format Instruction Data with Templates.