Fine-tune GPT-2 (125M) for Text Classification | PyTorch & Python

This chapter provides a comprehensive, step-by-step guide to fine-tuning a pre-trained GPT-2 small (124M/125M parameter) model for binary text classification (spam vs. not spam). Using PyTorch, we will modify the base causal language model architecture, freeze the majority of its weights to perform parameter-efficient fine-tuning on a local CPU, and evaluate its performance.

1. The Core Problem & Fine-Tuning Paradigm

A base large language model (LLM) is trained on a massive corpus of text to predict the next token. Because it is optimized for next-token prediction (autocomplete), it cannot natively classify text out of the box.

To transform this autocomplete engine into a high-performance classifier, we must perform supervised fine-tuning (SFT). We do this by:

1. Appending a custom classification head (a linear layer) to the end of the model.
2. Feeding labeled training data (SMS messages labeled as spam or not spam) into the model.
3. Updating a small, targeted subset of the model's weights using backpropagation.

2. Dataset Analysis & Train-Validation Splitting

The dataset used for this task is the open-source SMS Spam Collection dataset from Hugging Face. It is a lightweight dataset (~359 KB) containing 5,500 SMS messages.

Dataset Schema

• SMS: The raw text message.
• label: A binary integer where 0 represents Not Spam (Ham) and 1 represents Spam.

Splitting and Shuffling

Because the raw dataset only contains a training split, we must manually partition it into a 90% training set and a 10% validation set. Shuffling the dataset prior to splitting is critical to ensure that spam and ham labels are evenly distributed across both sets.

# Conceptual split distribution calculation
total_rows = 5574
train_size = int(0.9 * total_rows)  # ~5,016 rows
val_size = total_rows - train_size   # ~558 rows

Label Distribution Verification

To prevent class imbalance issues during validation, we verify the percentage of spam in both splits:

This distribution is highly consistent, ensuring that our validation accuracy will be a reliable metric for generalization.

3. Tokenization & Data Preparation

Because neural networks process numerical tensors rather than raw text, we must convert our text into token IDs. We use the pre-trained GPT-2 Tokenizer to ensure that the tokenization matches the vocabulary on which the model was originally trained.

PyTorch Tokenization Pipeline

The following Python function iterates over our dataset, encodes the text sequences, and pairs them with their corresponding integer labels:

def tokenize_dataset(dataset, tokenizer):
    tokenized_data = []
    for row in dataset:
        # Encode the SMS text into token IDs
        encoded_sms = tokenizer.encode(row['SMS'], truncation=True, max_length=1024)
        label = row['label']
        tokenized_data.append((encoded_sms, label))
    return tokenized_data

# Example of a tokenized pair:
# ( [1024, 342, 532], 0 ) -> (Token IDs, Label)

To verify the integrity of our tokenization pipeline, we can decode the token IDs back into text:

decoded_text = tokenizer.decode(tokenized_training_data[0][0])
print(f"Decoded: {decoded_text} | Label: {tokenized_training_data[0][1]}")

4. GPT-2 Base Architecture & The Need for Modification

The standard GPT-2 small model contains approximately 124 million parameters. Its architecture consists of:

The Dimensionality Problem

The final layer of the base GPT-2 model outputs a tensor of shape:

Where:

• $B$ = Batch Size
• $T$ = Sequence Length (Time steps)
• $C$ = Hidden Dimension ($768$ for GPT-2 small)

Because this output is designed to map back to the vocabulary size ($50,257$) for next-token prediction, it cannot directly output a binary classification probability ($[B, 2]$). We must modify this architecture.

5. Designing the Custom Classification Head (PyTorch)

To adapt GPT-2 for classification, we wrap the pre-trained transformer and append a custom linear layer (nn.Linear(768, 2)) to map the hidden dimension to our two target classes (Spam vs. Ham).

The Causal Model Trick

In a causal (autoregressive) model like GPT-2, information flows from left to right. The representation of the entire sentence accumulates at the very last token of the sequence. Therefore, instead of pooling or averaging all token representations, we extract only the hidden state of the final token ($T-1$) and pass it to our classification head.

PyTorch Implementation

import torch
import torch.nn as nn

class GPT2ForSentimentAnalysis(nn.Module):
    def __init__(self, base_gpt2_model):
        super().__init__()
        self.transformer = base_gpt2_model
        # Custom linear layer: maps 768 hidden dimensions to 2 output classes
        self.output_layer = nn.Linear(768, 2)
        
    def forward(self, input_ids):
        # Pass inputs through the base GPT-2 model
        transformer_outputs = self.transformer(input_ids)
        
        # Extract the hidden states from the last layer
        # Shape: [Batch Size, Sequence Length, 768]
        hidden_states = transformer_outputs.last_hidden_state
        
        # Extract the hidden state of the last token in the sequence
        # Shape: [Batch Size, 768]
        last_token_state = hidden_states[:, -1, :]
        
        # Pass through our custom classification head to get logits
        # Shape: [Batch Size, 2]
        logits = self.output_layer(last_token_state)
        return logits

6. Parameter-Efficient Fine-Tuning: Freezing Weights

Training all 124 million parameters of GPT-2 on a local CPU is computationally expensive and unnecessary. Because the base model already has a deep, pre-trained understanding of English syntax and semantics, we can freeze the majority of the network and only train:

1. The very last attention block (Block 11, index -1).
2. Our newly added custom output layer.

This strategy dramatically reduces memory usage and prevents catastrophic forgetting.

# 1. Freeze all parameters in the model
for param in model.parameters():
    param.requires_grad = False

# 2. Unfreeze the final transformer block (h[-1])
for param in model.transformer.h[-1].parameters():
    param.requires_grad = True

# 3. Unfreeze the custom linear classification head
for param in model.output_layer.parameters():
    param.requires_grad = True

Trainable Parameter Analysis

By freezing the early layers, we reduce the computational footprint significantly:

• Total Model Parameters: $124,439,808$ (~124M)
• Trainable Parameters: $7,087,874$ (~7M)
• Percentage of Trained Weights:

We are only updating 5.7% of the model's total parameters, making CPU-based training highly feasible.

7. The PyTorch Training Loop from Scratch

We implement the training loop from scratch using PyTorch, bypassing high-level wrappers to maintain granular control over batching and backpropagation.

Training Hyperparameters

• Epochs: 2 (The model sees the entire training set twice)
• Batch Size: 8
• Optimizer: Adam
• Loss Function: Cross-Entropy Loss (calculated via logits)
• Evaluation Interval: Every 50 batches (400 samples)

PyTorch Training Loop Implementation

import torch.optim as optim

# Initialize optimizer only for parameters that require gradients
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)
loss_fn = nn.CrossEntropyLoss()

epochs = 2
batch_size = 8
model.train()

for epoch in range(epochs):
    running_loss = 0.0
    batch_inputs = []
    batch_labels = []
    
    for i, (tokens, label) in enumerate(tokenized_training_data):
        batch_inputs.append(torch.tensor(tokens))
        batch_labels.append(torch.tensor(label))
        
        # Process when batch is full
        if len(batch_inputs) == batch_size:
            # Pad sequences to the same length within this batch
            inputs_tensor = nn.utils.rnn.pad_sequence(batch_inputs, batch_first=True, padding_value=tokenizer.eos_token_id)
            labels_tensor = torch.stack(batch_labels)
            
            # Reset gradients
            optimizer.zero_grad()
            
            # Forward pass
            logits = model(inputs_tensor)
            
            # Compute loss
            loss = loss_fn(logits, labels_tensor)
            
            # Backward pass (Backpropagation)
            loss.backward()
            
            # Update weights
            optimizer.step()
            
            running_loss += loss.item()
            
            # Clear batch lists
            batch_inputs, batch_labels = [], []
            
            # Print loss every 50 batches (400 samples)
            if (i + 1) % 400 == 0:
                print(f"Epoch {epoch+1} | Step {i+1} | Average Loss: {running_loss / 50:.4f}")
                running_loss = 0.0

8. Performance Evaluation: Before vs. After

Loss Convergence

During training, the loss value drops rapidly, demonstrating that the model is successfully adapting to the classification task:

Accuracy Evaluation

We evaluate the model on the validation set (which it has never seen during training) to measure generalization.

def evaluate_accuracy(model, validation_data):
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for tokens, label in validation_data:
            input_tensor = torch.tensor(tokens).unsqueeze(0) # Add batch dimension
            logits = model(input_tensor)
            prediction = torch.argmax(logits, dim=-1).item()
            
            if prediction == label:
                correct += 1
            total += 1
            
    return (correct / total) * 100

Performance Comparison

The massive jump from 16.12% to 99.28% validation accuracy proves that fine-tuning just 5.7% of the model's weights is highly effective for specialized downstream tasks.

9. Real-World Inference & Edge Case Validation

To deploy the model for real-time inference, we write a helper function that tokenizes raw text inputs, runs a forward pass, and applies an argmax over the output logits to return the predicted class.

def predict_sms(text, model, tokenizer):
    model.eval()
    with torch.no_grad():
        # Tokenize input text
        inputs = torch.tensor(tokenizer.encode(text)).unsqueeze(0)
        # Forward pass
        logits = model(inputs)
        # Get class with highest probability
        prediction = torch.argmax(logits, dim=-1).item()
        
        return "Spam" if prediction == 1 else "Not Spam"

Live Test Examples

# Test 1: Spam message containing typical urgency triggers
print(predict_sms("Last chance to win $100,000! Call 5676 now!", model, tokenizer))
# Output: Spam

# Test 2: Standard conversational message
print(predict_sms("Hey, how are you doing? Are we still meeting today?", model, tokenizer))
# Output: Not Spam

# Test 3: Financial spam message
print(predict_sms("Get home loans at 2% interest rates. Click this link to apply.", model, tokenizer))
# Output: Spam

# Test 4: Informational update
print(predict_sms("I will be reaching your location in 5 minutes.", model, tokenizer))
# Output: Not Spam

10. Summary and Key Takeaways

• Base Models are Autocomplete Engines: Pre-trained base models (like GPT-2 small) are optimized for next-token prediction and require structural modification and fine-tuning to perform classification tasks.
• Custom Classification Head: By appending a simple linear layer (nn.Linear(768, 2)) to the output of the transformer, we can map high-dimensional hidden states to discrete class probabilities.
• The Causal Last-Token Trick: Because GPT-2 is a causal language model, the semantic context of an entire sequence accumulates in the final token's hidden state. We extract only this last token's state for classification.
• Parameter-Efficient Fine-Tuning (PEFT): Freezing the majority of the model and training only the final attention block and the custom classification head allows us to train on local CPUs while updating only 5.7% of the total parameters.
• High Performance on Small Data: Fine-tuning on a small dataset (~5,500 rows) for just 2 epochs yielded a validation accuracy jump from 16.12% to 99.28%.
• Local CPU Feasibility: Specialized fine-tuning does not always require massive GPU clusters; small-scale models can be successfully adapted on standard consumer hardware.