Learn AI Create the Future — with Pranay
Main site

The BERT Family

Fine-tuning BERT for classification

Fine-tuning bolts a small, freshly-initialised classification head onto a pretrained BERT and retrains the whole thing briefly on your own labelled examples, reusing everything BERT already knows about language.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. Where you have already seen it
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Fine-tuning takes a BERT model that already understands language broadly. It briefly retrains the model on your labelled examples, so it learns your specific task without starting from zero.

Think about hiring an experienced cook and asking them to learn one restaurant's menu. You do not teach them how to hold a knife, how heat works, or what salt does. They already know all of that from years of general cooking. You only need to show them this restaurant's specific recipes. They pick them up fast, because the underlying skill was already there.

Fine-tuning BERT works the same way. The model already learned grammar, word meaning, and how sentences fit together, during the masked language modelling pretraining covered earlier in this section. Fine-tuning only has to teach it your specific task — is this review positive or negative. It needs far less data and far less training time than teaching it language from scratch would need.

Why it exists

Training a language model from nothing needs enormous amounts of text and compute. That is what BERT's original pretraining did. Most businesses do not have that much unlabelled text. Even if they did, retraining a full model for every new task would be wasteful. So much of what a model needs to know — grammar, word meaning, sentence structure — is identical across many different tasks.

Fine-tuning reuses the expensive part: general language understanding. It only trains the cheap part fresh. That means a small classification layer bolted onto the model's output, plus a light adjustment of the network's existing weights. This is why fine-tuning a sentiment classifier can work well with a few hundred labelled examples. Training an equally capable model from scratch would need millions.

How it works

  BEFORE fine-tuning:
    BERT (pretrained)  ->  [CLS] token's output vector
                            (a general-purpose sentence summary,
                             not yet connected to your task)

  ADD:
    A small, freshly-initialised classification layer on top,
    with one output per label -- e.g. 2 outputs for "positive" / "negative"

  DURING fine-tuning:
    Show labelled examples: "Great service!" -> positive
                              "Terrible, avoid this." -> negative
    Both the new classification layer AND BERT's existing weights
    adjust slightly, specialising toward this specific task.

  AFTER fine-tuning:
    New, unseen text  ->  BERT  ->  [CLS] vector  ->  classification layer
                                                        ->  "positive" or "negative"

Most of the useful knowledge was already there before fine-tuning started. Fine-tuning is a comparatively small, fast nudge, not a rebuild.

Where you have already seen it

  • App store review classifiers sort reviews into categories like bug reports, feature requests and praise. Almost always a fine-tuned BERT-family model sits under the hood.
  • Support ticket routing, automatically tagging incoming tickets by urgency or department, is a standard fine-tuning use case.
  • Spam and content moderation filters fine-tune a pretrained encoder on examples of the exact content a specific platform wants to catch.

Remember this

  • Fine-tuning reuses a pretrained model's general language understanding and only teaches it a specific new task.
  • It needs far less labelled data and training time than training an equally capable model from scratch.
  • Both a small new classification layer and the pretrained model's existing weights adjust during fine-tuning, not only the new layer alone.

What to learn next

Developer — Code and libraries.

Hugging Face's Trainer API handles the training loop. This example fine-tunes on ten tiny, made-up product reviews — enough to see real fine-tuning happen, not enough to build a production classifier from.

Setup

bash
pip install transformers torch

Minimal runnable code

finetune_classifier.py
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
from torch.utils.data import Dataset

texts = [
    "The delivery was fast and the packaging was excellent.",
    "Terrible service, my order never arrived.",
    "I love how easy this app is to use.",
    "Completely broken, crashes every time I open it.",
    "Great value for the price, highly recommend.",
    "Waste of money, stopped working after a day.",
    "Support team resolved my issue in minutes.",
    "Rude staff and a dirty store, never going back.",
    "Fast shipping and great customer support.",
    "The product broke on the first use, very disappointed.",
]
labels = [1, 0, 1, 0, 1, 0, 1, 0, 1, 0]  # 1 = positive, 0 = negative

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
enc = tok(texts, truncation=True, padding=True, return_tensors="pt")

class ReviewDataset(Dataset):
    def __len__(self): return len(labels)
    def __getitem__(self, i):
        item = {k: v[i] for k, v in enc.items()}
        item["labels"] = torch.tensor(labels[i])
        return item

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

args = TrainingArguments(
    output_dir="./tmp_out", num_train_epochs=8, per_device_train_batch_size=2,
    logging_steps=100, report_to=[], disable_tqdm=True, use_cpu=True,
)
trainer = Trainer(model=model, args=args, train_dataset=ReviewDataset())
train_out = trainer.train()
print("final training loss:", train_out.training_loss)

test_texts = ["This is amazing, works perfectly.", "Awful experience, avoid this."]
test_enc = tok(test_texts, truncation=True, padding=True, return_tensors="pt")
model.eval()
with torch.no_grad():
    logits = model(**test_enc).logits
preds = torch.argmax(logits, dim=-1)
for t, p in zip(test_texts, preds):
    print(f"{t!r} -> {'positive' if p.item()==1 else 'negative'}")
Output
[transformers] DistilBertForSequenceClassification LOAD REPORT from: distilbert-base-uncased
Key                     | Status     |
------------------------+------------+
vocab_transform.bias    | UNEXPECTED |
vocab_transform.weight  | UNEXPECTED |
vocab_projector.bias    | UNEXPECTED |
vocab_layer_norm.weight | UNEXPECTED |
vocab_layer_norm.bias   | UNEXPECTED |
classifier.bias         | MISSING    |
classifier.weight       | MISSING    |
pre_classifier.bias     | MISSING    |
pre_classifier.weight   | MISSING    |

{'train_runtime': '11.92', 'train_samples_per_second': '6.713', 'train_steps_per_second': '3.356', 'train_loss': '0.3679', 'epoch': '8'}
final training loss: 0.3679457902908325
'This is amazing, works perfectly.' -> positive
'Awful experience, avoid this.' -> negative

The exact loss number, the timing figures in the summary dict, and the row order in the LOAD REPORT table will all differ on your machine — the classification head starts from random weights, training runtime depends on your hardware, and the summary dict's exact fields depend on your installed transformers version. What should stay consistent is both test sentences ending up correctly classified.

Line by line

The "LOAD REPORT" at the top is expected, not an error. distilbert-base-uncased was pretrained for masked language modelling, so its checkpoint has weights for predicting masked words (vocab_transform, vocab_projector) — UNEXPECTED here because a classification model does not use them. The MISSING rows (classifier, pre_classifier) are the new, randomly-initialised classification head, since no pretrained checkpoint could possibly already have weights for a task-specific head it has never seen.

use_cpu=True forces training onto CPU explicitly. Without it, Trainer will use a GPU automatically if one is available, and a model trained partly on GPU tensors will error if you later run inference with CPU tensors, unless you are careful to keep everything on the same device — being explicit here avoids that entirely.

ReviewDataset wraps pre-tokenized tensors, not raw text. Tokenizing once, upfront (enc = tok(texts, ...)), rather than inside __getitem__, is faster for a dataset this small — for a much larger dataset, tokenizing on the fly or caching to disk becomes worth the added complexity.

The classification head sits on top of the [CLS] token's final representation (DistilBERT calls this the "pooled" or pre-classifier output) — the same summary-of-the-whole-sentence vector position BERT itself was originally trained to specialise for this exact purpose.

Common mistakes

Forgetting use_cpu=True (or an equivalent device setting) and mixing devices between training and inference. This produces a RuntimeError about tensors being on different devices — Trainer silently moves the model to GPU if one is present, and code written assuming CPU-only execution can fail confusingly on a different machine.

Fine-tuning on ten examples and expecting production-grade accuracy. This example works because the test sentences closely resemble the training pattern. A real sentiment classifier needs hundreds to thousands of varied, representative labelled examples per class.

Ignoring the "LOAD REPORT" warning entirely, or panicking over it. Both reactions are wrong. Reading it tells you exactly which weights are new (and therefore need training) versus reused from pretraining — genuinely useful diagnostic information, not noise to suppress.

Training for too many epochs on too little data. Eight epochs over ten examples risks memorising them outright rather than learning a generalisable pattern — acceptable for this demonstration, a real warning sign on an actual project. Watch validation loss on held-out data, not only training loss, to catch this.

Try it yourself

Add a genuinely ambiguous review — "It's okay, does the job but nothing special." — to test_texts without a label, and check which class the model assigns it. Ambiguous cases like this are where a small, ten-example fine-tuned model's limits are most visible.

What to learn next

Researcher — Mathematics and papers.

The fine-tuning objective

Given a pretrained encoder with parameters theta_enc and a freshly-initialised classification head with parameters theta_cls, fine-tuning minimises standard cross-entropy loss over labelled examples (x_i, y_i):

text
L = -(1/N) * sum over i=1..N of log P(y_i | x_i; theta_enc, theta_cls)
  • P(y_i | x_i; theta_enc, theta_cls) = softmax(W · h_CLS + b), where h_CLS is the encoder's final-layer representation at the [CLS] position, and W, b are the classification head's weights.
  • Both theta_enc and theta_cls receive gradient updates — this is full fine-tuning, distinct from feature extraction (freezing theta_enc, training only theta_cls) or parameter-efficient fine-tuning methods that update a small subset of theta_enc.

Why fine-tuning needs far less data than pretraining

Pretraining learns a general-purpose language representation from an unsupervised objective (masked language modelling) over a very large corpus, entirely unconstrained by any specific downstream task. Fine-tuning starts from that already-structured representation space and only needs to learn a comparatively simple mapping — a hyperplane, or a small number of them, separating classes within an already well-organised representation space — which is a much lower-capacity, lower-data-requirement problem than learning both the representation and the classification boundary simultaneously from random initialisation. Howard & Ruder (2018), ULMFiT, provide an early, systematic empirical demonstration of this effect for text classification specifically, predating BERT.

Catastrophic forgetting and learning rate choice

Full fine-tuning risks catastrophic forgetting: aggressive weight updates on a small, narrow fine-tuning dataset can overwrite general language understanding the model gained during pretraining, degrading performance on anything outside the fine-tuning distribution. The standard mitigation is a small learning rate for fine-tuning (typically 1e-5 to 5e-5, one to two orders of magnitude smaller than typical pretraining learning rates) and few epochs — Devlin et al. (2019) report 2 to 4 epochs as sufficient for most GLUE tasks, well below the number needed to fit a randomly-initialised model to the same data. Sun et al. (2019) provide a systematic study of fine-tuning hyperparameters for BERT specifically, finding layer-wise learning-rate decay (lower rates for earlier layers, higher for later layers closer to the task-specific head) further reduces forgetting.

Parameter-efficient alternatives

Full fine-tuning updates every parameter in theta_enc, which becomes costly to store per-task at scale — a separate full copy of the encoder for every fine-tuned task. Adapters (Houlsby et al., 2019) insert small trainable bottleneck layers between frozen pretrained layers, updating only the adapters. LoRA (Hu et al., 2021) freezes the original weight matrices and learns a low-rank update Delta W = B A, where A and B are much smaller than the original matrix, requiring far fewer trainable parameters and far less storage per fine-tuned task — see QLoRA: fine-tuning a large model on one GPU for the practical mechanics.

Key references

  • Howard, J. & Ruder, S. (2018). Universal Language Model Fine-tuning for Text Classification. arXiv:1801.06146 — ULMFiT.
  • Devlin, J. et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers. arXiv:1810.04805
  • Sun, C., Qiu, X., Xu, Y. & Huang, X. (2019). How to Fine-Tune BERT for Text Classification? arXiv:1905.05583
  • Houlsby, N. et al. (2019). Parameter-Efficient Transfer Learning for NLP. arXiv:1902.00751
  • Hu, E. et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685

Current state and open problems

Full fine-tuning of an encoder-sized model (under roughly 500M parameters) remains cheap and standard practice, and is the correct default for most classification tasks with a dedicated encoder — the compute cost is small enough that parameter-efficient methods offer little practical benefit at this scale, unlike for much larger decoder LLMs where full fine-tuning becomes prohibitively expensive. How much labelled data is genuinely needed for a target accuracy remains task-dependent and not well predicted by any general theory, though empirically, a few hundred well-distributed examples per class is a reasonable starting estimate for straightforward classification tasks, with harder or more ambiguous tasks needing considerably more. The open, actively researched question is less about the fine-tuning mechanism itself, which is well understood, and more about data efficiency and reliably estimating in advance how much labelled data a given task will require.

What to learn next