Table of Contents
What Fine-Tuning Means for Text Classification
Fine tuning for text classification means taking a transformer that has already learned general language patterns from large scale pretraining, then continuing training on your labeled dataset so the model predicts a discrete label such as sentiment, topic, intent, or toxicity. In practice, you keep most of the pretrained transformer intact and attach a small classification head that maps the final hidden representation to class logits. During fine tuning, you update the transformer weights, the classification head weights, or both, using your task loss, usually cross entropy.
This chapter focuses on the workflow and the PyTorch specific training details that are unique to fine tuning transformers for classification. Tokenization, attention masks, and batching are handled in the tokenizer and batching chapter, and attention internals are handled in earlier transformer building blocks chapters.
Model Setup: Base Transformer plus Classification Head
A common pattern is to use an encoder style model such as BERT, RoBERTa, or DistilBERT, then classify using the pooled representation. Many encoder models expose either a pooled output or you can use the hidden state at the first token position, often called the CLS token representation. If your model outputs last_hidden_state with shape (batch, seq_len, hidden), then last_hidden_state[:, 0, :] yields a (batch, hidden) tensor suitable for a linear layer to get logits of shape (batch, num_classes).
In pure PyTorch, the classification head is typically a dropout layer followed by a linear layer. Some libraries provide ready made classes that include this head, but it is still useful to understand the tensor shapes because most fine tuning errors come from mismatched shapes or incorrect label formats.
For single label classification with $C$ classes, your logits must have shape $(N, C)$ and your labels must have shape $(N,)$ with integer class indices in $[0, C-1]$. Use nn.CrossEntropyLoss, which expects raw logits, not probabilities.
Choosing Which Layers to Train: Full Fine-Tune vs Frozen Encoder
You have two main strategies. Full fine tuning updates all transformer parameters plus the classification head. This usually gives the best accuracy when you have enough labeled data and reasonable regularization. Feature extraction freezes the transformer and trains only the classification head. This is faster and can work well with small datasets, but it may underfit if the pretrained features are not aligned with your task.
A practical compromise is partial unfreezing. You freeze the embedding layers and early encoder blocks, then unfreeze the last few blocks and the head. This reduces compute and can improve stability on small datasets. In PyTorch, freezing is done by setting param.requires_grad = False for parameters you do not want to update, and then constructing the optimizer over only parameters with requires_grad=True.
If you freeze parameters after creating the optimizer, the optimizer may still hold references to those parameters. Freeze first, then build the optimizer.
Loss Function and Label Handling
For most text classification tasks, use cross entropy. In PyTorch, nn.CrossEntropyLoss combines log_softmax and nll_loss, so you should pass logits directly. Labels must be integer class indices. If your dataset stores labels as strings, map them to integers once and keep a consistent mapping for training and inference.
If you are doing multi label classification where each example can have multiple active classes, you usually switch to nn.BCEWithLogitsLoss and represent targets as float vectors of shape (N, C) with values 0 or 1. That is a different target format and different metrics, so do not mix it with the single label setup.
Learning Rate Choices and Parameter Groups
Transformers are sensitive to learning rate. For fine tuning, learning rates are typically much smaller than what you would use when training a model from scratch. A common pattern is to use a smaller learning rate for the pretrained encoder and a slightly larger one for the new classification head, because the head starts from random initialization.
In PyTorch, you can use optimizer parameter groups to set different learning rates. Conceptually, you create two groups, one for encoder parameters and one for head parameters, each with its own lr. Weight decay is often applied to most weights but not to bias terms and normalization parameters.
Applying weight decay to LayerNorm or bias parameters can hurt fine tuning. Use parameter grouping to exclude them from weight decay when possible.
Schedulers and Warmup for Stable Fine-Tuning
Fine tuning often benefits from a learning rate schedule that includes warmup. Warmup means starting with a very small learning rate and increasing it linearly for a short fraction of training steps, then decaying. This helps avoid early training instability when gradients are large because the classification head is random and the model is adapting quickly.
The key implementation detail is that scheduler steps should usually be based on optimizer steps, not epochs. If you use gradient accumulation, the number of optimizer steps is smaller than the number of batches, and your warmup and total step counts should reflect that.
If you call scheduler.step() every batch but you only call optimizer.step() every $k$ batches due to gradient accumulation, your schedule will run $k$ times faster than intended. Tie scheduler stepping to optimizer stepping.
Training Loop Details Specific to Transformers
A transformer classification forward pass typically takes input_ids and attention_mask, and sometimes token_type_ids, then returns logits. Your training loop must ensure that all tensors are on the same device, and that labels are of the correct dtype, usually torch.long for cross entropy.
Mixed precision is commonly used to speed up fine tuning on GPUs. The important fine tuning detail is that loss scaling and gradient clipping interact with stability. If you clip gradients, clip after unscaling the gradients when using mixed precision.
Another transformer specific detail is that sequence lengths affect memory heavily. If you use dynamic padding per batch, you can reduce wasted computation. This is mainly a data pipeline choice, but it directly impacts whether fine tuning fits in memory.
Evaluation: Metrics Beyond Loss
Loss is useful for optimization, but classification quality is often judged with accuracy, F1, precision, recall, or AUC depending on the task. During evaluation you should disable dropout by switching the model to evaluation mode and avoid gradient tracking.
For single label classification, predictions come from argmax over logits. For probabilistic outputs, convert logits to probabilities with softmax only for reporting or thresholding, not for the loss.
Do not apply softmax before nn.CrossEntropyLoss. It will worsen numerical stability and can slow learning.
Handling Class Imbalance in Fine-Tuning
Text datasets often have imbalance, such as many neutral examples and few rare classes. If you ignore imbalance, a model can achieve high accuracy by predicting the majority class. In PyTorch, cross entropy supports class weights, where you provide a weight vector of length C that upweights rare classes. Another approach is to adjust sampling in the DataLoader, which is covered elsewhere, but the fine tuning specific point is that imbalance can appear as apparently good training loss but poor minority class metrics.
Checkpointing and Best Model Selection for Fine-Tuning
Fine tuning runs are often short, and overfitting can happen quickly. It is common to save the best checkpoint based on validation metric rather than training loss. Save both the model weights and the label mapping used by your dataset. If you use a tokenizer with special settings, save those settings too, so inference uses the exact same preprocessing.
A practical habit is to store a checkpoint containing model state_dict, optimizer state_dict, scheduler state, current epoch or step, and the random seed state when you need exact resumption.
Common Failure Modes
If the loss does not decrease, the most common causes are incorrect label format, a learning rate that is too high, forgetting to call model.train(), or accidentally freezing all parameters. If training accuracy is high but validation is poor, you are likely overfitting, so consider smaller learning rate, more dropout, weight decay, earlier stopping, or freezing more layers.
If you see NaNs, reduce learning rate, enable gradient clipping, verify that your inputs and labels are valid, and check mixed precision settings. Transformers can also become unstable if the batch size is too small, especially without warmup, because gradient estimates are noisy.
Always verify a single forward pass: logits shape, label dtype, and loss value should all look sensible before running a full fine tuning job.
A Minimal PyTorch Skeleton for Fine-Tuning
A typical fine tuning loop has these moving parts: a model that returns logits, an optimizer with parameter groups, an optional scheduler with warmup, and an evaluation function that computes task metrics. The concrete implementation details vary with the transformer source, but the structure stays the same: forward pass with tokenized inputs, compute cross entropy loss, backpropagate, optionally clip gradients, step optimizer, step scheduler, and periodically evaluate on the validation set to select the best checkpoint.
The main goal is consistency: consistent label mapping, consistent tokenization settings, consistent device placement, and consistent stepping logic for optimizer and scheduler. When those are correct, transformer fine tuning for text classification becomes a reliable, repeatable recipe.