Skip to main content

Fine-Tuning a Text Classifier

Fine-tuning updates some or all parameters of a pretrained model for a downstream objective. The learning rate, pooling rule, frozen layers, class weights, validation policy, and tokenizer are part of the experiment. This lesson demonstrates the mechanics locally; it does not pretend that a randomly initialized toy encoder is pretrained.

Prerequisites​

Complete Pretraining and Transfer Learning and understand binary logits and gradient updates.

Learning goals​

  • separate encoder parameters from a task-specific classification head.
  • mask padding during pooling and train a head before optionally unfreezing an encoder.
  • keep pretrained Hugging Face loading outside deterministic CI.

Worked example​

The fixed embedding table acts as a known local encoder. We freeze it and train only the linear head on six examples. This is a mechanical analogue of head-first fine-tuning, not a quality benchmark.

import torch
from torch import nn

torch.manual_seed(45)
input_ids = torch.tensor([[2, 4, 0], [2, 5, 0], [4, 2, 0], [3, 5, 0], [3, 4, 0], [5, 3, 0]])
labels = torch.tensor([1, 1, 1, 0, 0, 0], dtype=torch.float32)
encoder = nn.Embedding(6, 2, padding_idx=0)
with torch.no_grad():
    encoder.weight.copy_(torch.tensor([
        [0.0, 0.0], [0.0, 0.0], [1.0, 0.2], [-1.0, 0.2], [0.0, 0.1], [0.0, -0.1]
    ]))
encoder.weight.requires_grad_(False)
head = nn.Linear(2, 1)
optimizer = torch.optim.SGD(head.parameters(), lr=0.5)

mask = input_ids.ne(0)
for _ in range(60):
    pooled = (encoder(input_ids) * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True)
    logits = head(pooled).squeeze(1)
    loss = nn.functional.binary_cross_entropy_with_logits(logits, labels)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

with torch.inference_mode():
    probabilities = torch.sigmoid(head(pooled).squeeze(1))
    predictions = probabilities.ge(0.5).to(torch.int64)

print(f"ids={tuple(input_ids.shape)} pooled={tuple(pooled.shape)} logits={tuple(logits.shape)}")
print(f"encoder_trainable={encoder.weight.requires_grad} head_trainable={head.weight.requires_grad}")
print(f"predictions={predictions.tolist()} accuracy={predictions.eq(labels.to(torch.int64)).float().mean():.1f}")

Output:

ids=(6, 3) pooled=(6, 2) logits=(6,)
encoder_trainable=False head_trainable=True
predictions=[1, 1, 1, 0, 0, 0] accuracy=1.0

Real fine-tuning must select hyperparameters and checkpoints on training/validation data. Unfreezing can improve adaptation but increases memory, compute, and overfitting risk. Class imbalance may require weighting, sampling, threshold selection, or a different objective; each changes what a probability or default threshold means.

This optional Hugging Face sketch is not executed. It may download artifacts and requires a separately pinned transformers==5.20.0 environment plus an explicit dataset implementation.

from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments

model_id = "distilbert/distilbert-base-uncased"
revision = "PIN_A_REVIEWED_COMMIT_SHA"
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
model = AutoModelForSequenceClassification.from_pretrained(model_id, revision=revision, num_labels=2)
arguments = TrainingArguments(output_dir="artifacts/text-classifier", eval_strategy="epoch")
trainer = Trainer(model=model, args=arguments, train_dataset=train_dataset, eval_dataset=validation_dataset)

That block is a manual/download boundary and intentionally has no fabricated output.

Common mistakes​

  • Tokenizing before splitting or fitting label/preprocessing state on all data.
  • Selecting epochs, learning rate, threshold, or checkpoint using test results.
  • Pooling padding positions or assuming every architecture uses the same pooled output.
  • Using accuracy alone with class imbalance or asymmetric harms.
  • Fine-tuning every layer by default without a resource and overfitting comparison.

Practice​

Unfreeze the local encoder only after the head has trained. Use a smaller encoder learning rate than head learning rate, and explain which validation evidence would justify the added complexity.

Next steps​

Generation, Decoding, and Randomness moves from classification logits to repeated token selection.

Primary sources​

Checked 2026-10-16: Hugging Face text-classification guide, Hugging Face Trainer, Hugging Face tokenizer API, PyTorch Embedding, and PyTorch binary cross-entropy with logits.