6 Finetuning for Classification

This chapter shows how to adapt a pretrained language model to a focused prediction task via classification finetuning, using SMS spam detection as a running example. It contrasts two common approaches—instruction finetuning and classification finetuning—and explains why a specialized classifier that predicts among fixed labels is often simpler, needs less data and compute, and can achieve strong accuracy on a narrow task, while being less flexible than instruction-following models. The workflow is framed as a practical, end-to-end pipeline: prepare data, adapt the model, finetune and evaluate, then deploy for inference.

The data preparation centers on the SMS Spam Collection. After loading the messages and observing a class imbalance (many more “ham” than “spam”), the dataset is balanced by undersampling, labels are mapped to integers, and the data is split into training, validation, and test subsets (70/10/20). Because messages vary in length, texts are tokenized with a GPT-2 tokenizer and padded (or truncated) to a uniform length using a special token ID, enabling efficient batching. A custom Dataset encapsulates encoding and padding logic, and PyTorch DataLoaders serve fixed-size batches of token ID sequences with corresponding labels, establishing a clean input pipeline for training and evaluation.

Model adaptation starts by loading pretrained GPT-2 weights and then replacing the large vocabulary-sized output layer with a compact two-class head. Most parameters are frozen to retain general language knowledge, while the new head, the final transformer block, and the last normalization layer are made trainable. Using causal masking as motivation, only the logits of the last token in each sequence are fed into a cross-entropy loss and used for accuracy computation. After initial, near-chance performance, finetuning with AdamW over a few epochs drives losses down and yields high accuracy (around the mid-to-high 90s) on validation and test sets with little overfitting. The chapter concludes with a small inference utility to classify new messages and guidance on saving/loading the finetuned model.

A mental model of the three main stages of coding an LLM, pretraining the LLM on a general text dataset and finetuning it. This chapter focuses on finetuning a pretrained LLM as a classifier.

Illustration of two different instruction-finetuning scenarios. At the top, the model is tasked with determining whether a given text is spam. At the bottom, the model is given an instruction on how to translate an English sentence into German.

Illustration of a text classification scenario using an LLM. A model finetuned for spam classification does not require further instruction alongside the input. In contrast to an instruction-finetuned model, it can only respond with "spam" and "not spam."

Illustration of the three-stage process for classification-finetuning the LLM in this chapter. Stage 1 involves dataset preparation. Stage 2 focuses on model setup. Stage 3 covers the finetuning and evaluation of the model.

Preview of the SMSSpamCollection dataset in a pandas DataFrame, showing class labels ("ham" or "spam") and corresponding text messages. The dataset consists of 5,572 rows (text messages and labels).

An illustration of the input text preparation process. First, each input text message is converted into a sequence of token IDs. Then, to ensure uniform sequence lengths, shorter sequences are padded with a padding token (in this case, token ID 50256) to match the length of the longest sequence.

An illustration of a single training batch consisting of 8 text messages represented as token IDs. Each text message consists of 120 token IDs. In addition, a class label array stores the 8 class labels corresponding to the text messages, which can be either 0 (not spam) or 1 (spam).

Illustration of the three-stage process for classification-finetuning the LLM in this chapter. After completing stage 1, preparing the dataset, this section focuses on initializing the LLM we will finetune to classify spam messages.

This figure illustrates adapting a GPT model for spam classification by altering its architecture. Initially, the model's linear output layer mapped 768 hidden units to a vocabulary of 50,257 tokens. For spam detection, this layer is replaced with a new output layer that maps the same 768 hidden units to just two classes, representing "spam" and "not spam."

The GPT model we developed in earlier chapters, which we loaded previously, includes 12 repeated transformer blocks. Alongside the output layer, we set the final LayerNorm and the last transformer block as trainable, while the remaining 11 transformer blocks and the embedding layers are kept non-trainable.

An illustration of the GPT model with a 4-token example input and output. The output tensor consists of 2 columns due to the modified output layer. We are only focusing on the last row corresponding to the last token when finetuning the model for spam classification.

Illustration of the causal attention mechanism as discussed in chapter 3, where the attention scores between input tokens are displayed in a matrix format. The empty cells indicate masked positions due to the causal attention mask, preventing tokens from attending to future tokens. The values in the cells represent attention scores, with the last token, "time," being the only one that computes attention scores for all preceding tokens.

Illustration of the three-stage process for classification-finetuning the LLM in this chapter. This section implements the last step of stage 2, implementing the functions to evaluate the model's performance to classify spam messages before, during, and after the finetuning.

The model outputs corresponding to the last token are converted into probability scores for each input text. Then, the class labels are obtained by looking up the index position of the highest probability score. Note that the model predicts the spam labels incorrectly because it has not yet been trained.

A typical training loop for training deep neural networks in PyTorch consists of several steps, iterating over the batches in the training set for several epochs. In each loop, we calculate the loss for each training set batch to determine loss gradients, which we use to update the model weights to minimize the training set loss.

This graph shows the model's training and validation loss over the five training epochs. The training loss, represented by the solid line, and the validation loss, represented by the dashed line, both sharply decline in the first epoch and gradually stabilize towards the fifth epoch. This pattern indicates good learning progress and suggests that the model learned from the training data while generalizing well to the unseen validation data.

Both the training accuracy (solid line) and the validation accuracy (dashed line) increase substantially in the early epochs and then plateau, achieving almost perfect accuracy scores of 1.0. The close proximity of the two lines throughout the epochs suggests that the model does not overfit the training data much.

Illustration of the three-stage process for classification-finetuning the LLM in this chapter. This section implements the final step of stage 3, using the finetuned model to classify new spam messages.

• There are different strategies for finetuning LLMs, including classification-finetuning (this chapter) and instruction-finetuning (next chapter)
• Classification-finetuning involves replacing the output layer of an LLM via a small classification layer.
• In the case of classifying text messages as "spam" or "not spam," the new classification layer consists of only 2 output nodes; in previous chapters, the number of output nodes was equal to the number of unique tokens in the vocabulary, namely, 50,256
• Instead of predicting the next token in the text as in pretraining, classification-finetuning trains the model to output a correct class label, for example, "spam" or "not spam."
• The model input for finetuning is text converted into token IDs, similar to pretraining.
• Before finetuning an LLM, we load the pretrained model as a base model.
• Evaluating a classification model involves calculating the classification accuracy (the fraction or percentage of correct predictions).
• Finetuning a classification model uses the same cross entropy loss function that is used for pretraining the LLM.

FAQ