Module 5 · Sequences and Transformers

Fine-Tuning a Pretrained Transformer

Lesson 20 of 32 11 min read Updated

Fine-tuning a pretrained transformer means taking a model that already learned useful language representations and continuing training it on a narrower task. In this lesson, the outcome is practical: given a pretrained encoder such as BERT, RoBERTa, DistilBERT, or a similar PyTorch transformer, you should be able to attach an appropriate task head, choose which weights to update, build batches with attention masks, train without destroying the pretrained signal, and diagnose common failures.

This chapter sits in the sequences and transformers section because fine-tuning is where attention-based representation learning becomes useful for ordinary supervised problems. Instead of training word embeddings and sequence layers from scratch, you reuse parameters learned from large corpora and spend your labeled examples teaching the model a task-specific boundary, score, or generation style.

What Changes During Fine-Tuning

A transformer encoder maps token ids to contextual vectors. The same word receives different vectors depending on surrounding tokens because self-attention mixes information across sequence positions. Pretraining sets the initial values of embeddings, attention projections, feed-forward layers, layer normalization parameters, and sometimes a pooling layer. Fine-tuning starts from those values rather than random initialization.

The usual supervised fine-tuning path has five moving parts. The tokenizer converts text to token ids using the vocabulary expected by the checkpoint. The collator pads examples in a batch and builds an attention mask so padding tokens do not participate as real content. The transformer body produces hidden states with shape like batch, sequence length, hidden size. A task head converts those hidden states to logits, spans, token labels, or another task output. The loss compares those outputs with labels and backpropagates through the head and some or all of the transformer.

The important internal detail is that gradients flow through the whole computation graph unless you deliberately freeze parameters. The classifier head usually begins randomly initialized, so its gradients are large and task-specific. The lower transformer layers often encode broad lexical and syntactic patterns, so updating them too aggressively can cause catastrophic forgetting. Optimizers such as AdamW are common because decoupled weight decay handles large parameter sets better than plain L2 regularization in many transformer training recipes.

API Anatomy

In PyTorch, fine-tuning is still ordinary module training. A pretrained model is an nn.Module with named parameters. A head is another module. A batch is a dictionary or tuple containing input_ids, attention_mask, optional token_type_ids, and labels. The training loop switches the model to train(), computes logits, computes loss, calls backward(), clips gradients if needed, steps the optimizer, steps the learning-rate scheduler, and clears gradients.

Libraries such as Hugging Face Transformers wrap those pieces in model classes like AutoModelForSequenceClassification, but the mechanics are the same. The configuration records label count, hidden size, dropout probabilities, maximum positions, vocabulary size, and task mappings. Tokenizer and checkpoint must match; using a tokenizer from one family with weights from another usually produces valid tensors with poor semantics.

Example 1: Tokenization, Truncation, and Masks

The first example shows why the tokenizer is part of the model, not a preprocessing afterthought. Transformer checkpoints expect special tokens and integer ids from the vocabulary used during pretraining. The attention mask distinguishes real tokens from padding so batched computation can ignore padded positions.

class TinyTokenizer:
    def __init__(self, vocab, unk="[UNK]", pad="[PAD]", cls="[CLS]", sep="[SEP]"):
        self.vocab = vocab
        self.unk = unk
        self.pad = pad
        self.cls = cls
        self.sep = sep

    def encode(self, text, max_length):
        tokens = [self.cls] + text.lower().split() + [self.sep]
        ids = [self.vocab.get(tok, self.vocab[self.unk]) for tok in tokens]
        ids = ids[:max_length]
        mask = [1] * len(ids)
        while len(ids) < max_length:
            ids.append(self.vocab[self.pad])
            mask.append(0)
        return ids, mask

vocab = {"[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "pytorch": 4, "tuning": 5, "works": 6}
tokenizer = TinyTokenizer(vocab)
print(tokenizer.encode("PyTorch tuning works", max_length=6))

The deterministic output is ([2, 4, 5, 6, 3, 0], [1, 1, 1, 1, 1, 0]). The leading 2 is the class token, 3 is the separator token, and the final 0 is padding. The matching mask has a final 0, which later becomes a padding mask for attention. In a real pretrained checkpoint, subword tokenization is more complex, but the same contract holds: ids, masks, and special tokens must be consistent with the checkpoint.

Example 2: Freezing Most of the Encoder

When the labeled dataset is small, a conservative design is to train the head and perhaps the last transformer block while freezing earlier layers. This reduces memory use, speeds training, and lowers the chance of overwriting general language features. It may underfit if the target domain differs strongly from pretraining text.

import torch
from torch import nn

class ToyEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.embeddings = nn.Embedding(100, 8)
        self.layers = nn.ModuleList([nn.Linear(8, 8) for _ in range(2)])

    def forward(self, input_ids):
        x = self.embeddings(input_ids).mean(dim=1)
        for layer in self.layers:
            x = torch.tanh(layer(x))
        return x

class Classifier(nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder
        self.head = nn.Linear(8, 2)

model = Classifier(ToyEncoder())
for name, param in model.encoder.named_parameters():
    param.requires_grad = name.startswith("layers.1")

trainable = [name for name, param in model.named_parameters() if param.requires_grad]
print(trainable)

The expected output is ['encoder.layers.1.weight', 'encoder.layers.1.bias', 'head.weight', 'head.bias']. The embeddings and first layer are frozen because their parameters have requires_grad=False. The last encoder layer and classifier head still receive gradients. In a real transformer, you would freeze by named layer groups such as embeddings, encoder block numbers, or adapter modules.

Example 3: One Fine-Tuning Step

This example uses PyTorch’s transformer encoder to demonstrate the training step that a pretrained model class hides. The input is already embedded to keep the code small. The key padding mask is the inverse of the attention mask: PyTorch expects True where positions should be ignored.

import torch
from torch import nn

torch.manual_seed(7)
encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=8, nhead=2, batch_first=True, dropout=0.0),
    num_layers=1,
)
head = nn.Linear(8, 2)
optimizer = torch.optim.AdamW(list(encoder.parameters()) + list(head.parameters()), lr=1e-3)
input_embeddings = torch.randn(3, 5, 8)
labels = torch.tensor([0, 1, 0])
attention_mask = torch.tensor([[1, 1, 1, 1, 0], [1, 1, 1, 0, 0], [1, 1, 1, 1, 1]], dtype=torch.bool)
key_padding_mask = ~attention_mask

pooled = encoder(input_embeddings, src_key_padding_mask=key_padding_mask)[:, 0]
logits = head(pooled)
loss = nn.CrossEntropyLoss()(logits, labels)
loss.backward()
optimizer.step()
print(tuple(logits.shape))
print(loss.ndim == 0)
print(head.weight.grad is not None)

The deterministic output is (3, 2), then True, then True. That means the model produced three examples with two class logits each, the loss is a scalar, and the classification head received gradients. What matters is the shape and gradient path: the loss updates both the encoder and the classification head because both parameter groups are passed to AdamW.

Design Choices and Trade-Offs

The first choice is the base checkpoint. A larger model often improves accuracy but costs more memory, has higher latency, and needs more careful regularization. A smaller distilled model can be preferable for batch inference or online prediction. A domain-specific checkpoint can help when vocabulary and phrasing are specialized, but it can also inherit narrower biases.

The second choice is how much to update. Full fine-tuning gives maximum flexibility but consumes more memory because gradients and optimizer states are stored for every parameter. Feature extraction freezes the body and trains only the head, which is cheap but less adaptable. Intermediate options include unfreezing only later layers, gradual unfreezing, layer-wise learning-rate decay, adapters, and low-rank updates. These methods trade implementation complexity for lower memory or lower overfitting risk.

The third choice is sequence length. Longer sequences preserve context but self-attention cost grows roughly with the square of sequence length. Truncating text can silently remove decisive evidence, especially for legal, medical, or support-ticket tasks. Sliding windows, hierarchical pooling, or summarization can help, but each adds assumptions that must be evaluated.

The fourth choice is evaluation. Accuracy alone is weak for imbalanced classification. Use precision, recall, F1, calibration, confusion matrices, and slice metrics that match the cost of mistakes. Always compare against a simple baseline such as TF-IDF plus logistic regression or a frozen embedding model; fine-tuning is justified by measured improvement, not by popularity.

Failure Modes and Troubleshooting

Symptom: training loss falls but validation performance is unstable. A common cause is overfitting from updating too many parameters on a small dataset. Diagnose by plotting train and validation loss, checking label distribution, and repeating runs with different seeds. Correct by freezing more layers, lowering the learning rate, increasing dropout, adding weight decay, using early stopping, or collecting more representative examples.

Symptom: predictions look random even after training. The cause may be a tokenizer-checkpoint mismatch, swapped label ids, or labels not aligned with the intended text. Diagnose by decoding a few tokenized examples, printing the id2label mapping, and overfitting ten examples until near-zero loss. Correct the tokenizer, label mapping, or dataset transform before scaling up training.

Symptom: GPU memory is exhausted. Transformers store activations for every layer and token during backpropagation. Diagnose by reducing batch size to one, logging sequence lengths, and checking whether gradients are enabled for frozen layers. Correct with shorter max length, dynamic padding, gradient accumulation, mixed precision, gradient checkpointing, or parameter-efficient tuning.

Symptom: masked tokens affect outputs. The cause is often an inverted or missing mask. Diagnose with a batch where one example has extra padding and compare logits with and without that padding. Correct the mask convention for the API you call; some functions use 1 for keep, while PyTorch padding masks use True for ignore.

Security, Performance, and Reliability

Fine-tuned transformers can memorize rare training strings, so avoid training on secrets, credentials, or private records without a retention and redaction policy. For user-facing systems, record model version, tokenizer version, label mapping, max length, and threshold choices with every evaluation artifact. This is reliability work specific to models: a checkpoint without its preprocessing and label metadata is not reproducible.

Performance bottlenecks usually come from tokenization, sequence length, and batch shape. Dynamic padding reduces wasted attention on padding tokens. Batching improves throughput but increases latency for individual requests. Quantization, distillation, and smaller checkpoints can help at inference time, but they should be evaluated against the task metric because compression can disproportionately hurt minority classes or long examples.

Hands-On Lab: Fine-Tune a Sentiment Classifier

Prerequisites: Python with PyTorch, a transformers-style pretrained encoder package if you use one, a small labeled text dataset, and access to a CPU or GPU. Keep the lab small: a few hundred examples is enough to verify mechanics.

  1. Create train, validation, and test splits before tokenization. Stratify if labels are imbalanced.
  2. Load the tokenizer that belongs to the checkpoint. Tokenize text with a fixed max_length, truncation, and padding performed by a collator.
  3. Load a sequence-classification model with the exact number of labels. Set id2label and label2id explicitly.
  4. Run a smoke test on two examples and assert that logits have shape (2, num_labels).
  5. Train for a small number of epochs with AdamW, a low learning rate, gradient clipping, and validation after each epoch.
  6. Save the model, tokenizer, label mappings, metric report, and the command or notebook parameters used for the run.

Verification: confirm that training loss decreases on a tiny subset, validation metrics beat the baseline, and a hand-written positive and negative example receive sensible labels. Also verify that changing only padding length does not change the predicted class.

Cleanup: delete temporary checkpoints that are not selected, clear cached datasets if disk space matters, and roll back to the baseline model if validation performance or slice metrics regress. Keep the final checkpoint only with its tokenizer and metadata.

Assessment Exercises

  1. You have 800 labeled examples and a large pretrained encoder. Which parameters would you train first, and what evidence would make you unfreeze more layers?
  2. A model achieves 94 percent accuracy but misses most rare positive cases. Which metrics and data slices would you inspect before changing architecture?
  3. Explain why padding masks must be tested with examples of different lengths. What bug would this reveal?
  4. Design a small experiment comparing full fine-tuning, frozen-encoder training, and last-layer-only fine-tuning. What would you hold constant?
  5. A teammate wants to raise max_length from 128 to 512. Name two benefits, two costs, and a validation check.

Summary

Fine-tuning a pretrained transformer reuses a language model’s learned representations while adapting its outputs to a supervised task. The core mechanics are tokenizer consistency, attention masks, a task head, gradient flow, optimizer configuration, and evaluation against a baseline. Good practice is conservative at first: verify shapes, overfit a tiny batch, watch validation behavior, and increase trainable capacity only when the evidence supports it.

Progress is saved in this browser.