Fine-Tuning BERT for Downstream Tasks

Jr Codex NLP & LLM Notes

Level: Advanced Prerequisites: Chapter 2: BERT's Pretraining — MLM + NSP Time to complete: ~25 minutes

The Fine-Tuning Pattern
Task 1: Text Classification
Task 2: Named Entity Recognition (Token Classification)
Task 3: Extractive Question Answering
Task 4: Sentence-Pair Tasks
Fine-Tuning Hyperparameters — BERT-Specific Guidance
Summary & Next Steps

1. The Fine-Tuning Pattern

This chapter directly applies DL Notes, Module 7's transfer learning framework to BERT specifically: take the pretrained encoder (Chapters 1-2), add a small, task-specific head, and fine-tune the whole thing (or just the head) on labeled data for a SPECIFIC downstream task.

The General Pattern, Across Every Task in This Chapter
─────────────────────────────────────────
  1. Load PRETRAINED BERT (Chapter 1, Section 5)
  2. ADD a small task-specific layer on top of BERT's output
       (which specific output — [CLS] vs per-token — depends
       on the task)
  3. FINE-TUNE (DL Notes, Module 7, Ch.2) — typically the
       ENTIRE model, at a low learning rate, since downstream
       datasets are usually small relative to pretraining
─────────────────────────────────────────

2. Task 1: Text Classification

Directly reusing Chapter 1, Section 3's [CLS] token — its final representation summarizes the entire input, making it a natural input to a simple classification head.

from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
 
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)      # e.g. binary sentiment
 
dataset = load_dataset("imdb")
 
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=256)
 
tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
training_args = TrainingArguments(
    output_dir="./bert-sentiment",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,      # DL Notes, Module 7, Ch.2, Sec.3's low fine-tuning rate
    weight_decay=0.01,          # DL Notes, Module 3, Ch.4
)
 
trainer = Trainer(model=model, args=training_args,
                   train_dataset=tokenized_dataset["train"], eval_dataset=tokenized_dataset["test"])
trainer.train()

What AutoModelForSequenceClassification Adds
─────────────────────────────────────────
  A single nn.Linear layer (DL Notes, Module 2, Ch.1) on TOP
  of BERT's POOLED [CLS] output (Chapter 1, Section 3) —
  mapping BERT's 768-dimensional representation down to
  num_labels output scores, exactly the classification head
  pattern from DL Notes, Module 7, Chapter 2, Section 2.
─────────────────────────────────────────

3. Task 2: Named Entity Recognition (Token Classification)

Unlike classification (one label for the WHOLE input), NER requires a label for every individual token — directly the "many-to-many, aligned" pattern from DL Notes, Module 5, Chapter 2, Section 6.

from transformers import AutoModelForTokenClassification
 
# Labels like: O (outside any entity), B-PER (beginning of a person name),
# I-PER (inside a person name), B-ORG, I-ORG, B-LOC, I-LOC, etc.
label_list = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
 
model = AutoModelForTokenClassification.from_pretrained(model_name, num_labels=len(label_list))
 
text = "Marie Curie worked at the University of Paris"
inputs = tokenizer(text, return_tensors="pt")
 
import torch
with torch.no_grad():
    outputs = model(**inputs)
 
predictions = torch.argmax(outputs.logits, dim=-1)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
for token, pred in zip(tokens, predictions[0]):
    print(f"{token}: {label_list[pred]}")

What Changed From Classification (Section 2)
─────────────────────────────────────────
  Instead of using ONLY [CLS]'s output, a classification head
  is applied to EVERY token position's final hidden state
  INDEPENDENTLY — each token gets its OWN prediction, using
  the CONTEXTUAL representation self-attention already computed
  for it (directly benefiting from Chapter 1, Section 6's
  bidirectional context).
─────────────────────────────────────────

4. Task 3: Extractive Question Answering

Given a question and a passage, extractive QA predicts the exact span (start and end position) within the passage that answers the question.

from transformers import AutoModelForQuestionAnswering
 
qa_model = AutoModelForQuestionAnswering.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
qa_tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased-whole-word-masking-finetuned-squad")
 
question = "Who developed the theory of relativity?"
passage = "Albert Einstein developed the theory of relativity, revolutionizing physics."
 
inputs = qa_tokenizer(question, passage, return_tensors="pt")      # Chapter 1, Sec.4's TWO-segment input
 
with torch.no_grad():
    outputs = qa_model(**inputs)
 
start_idx = torch.argmax(outputs.start_logits)      # the MOST likely START of the answer span
end_idx = torch.argmax(outputs.end_logits)              # the MOST likely END of the answer span
 
answer_tokens = inputs["input_ids"][0][start_idx:end_idx+1]
answer = qa_tokenizer.decode(answer_tokens)
print(f"Answer: {answer}")      # "Albert Einstein"

Two Independent Classification Heads
─────────────────────────────────────────
  A "start" classifier and an "end" classifier, EACH applied
  to every token position (similar to Section 3's per-token
  approach) — predicting, for EVERY position, how likely it is
  to be the START (or END) of the answer span.
─────────────────────────────────────────

5. Task 4: Sentence-Pair Tasks

Tasks like Natural Language Inference (NLI) — given two sentences, predict whether the first entails, contradicts, or is neutral toward the second — directly reuse Chapter 1, Section 4's segment embeddings for encoding two sentences at once.

model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)      # entailment/neutral/contradiction
 
premise = "A man is playing guitar on stage."
hypothesis = "A person is performing music."
 
inputs = tokenizer(premise, hypothesis, return_tensors="pt")      # TWO sentences, ONE input, via segment embeddings
outputs = model(**inputs)
prediction = torch.argmax(outputs.logits, dim=-1)
labels = ["entailment", "neutral", "contradiction"]
print(f"Prediction: {labels[prediction.item()]}")

This is structurally identical to Section 2's classification setup — the only difference is the input contains two sentences (via BERT's native segment-embedding support) instead of one.

6. Fine-Tuning Hyperparameters — BERT-Specific Guidance

Practical Defaults That Work Well for BERT (Directly Extending
DL Notes, Module 7, Chapter 4's Recipe)
─────────────────────────────────────────
  Learning rate:      2e-5 to 5e-5 — NOTABLY smaller than
                         typical CNN fine-tuning rates (DL Notes,
                         Module 7), since BERT's pretrained
                         representations are especially easy to
                         disrupt with large updates
  Epochs:                 2-4 — BERT fine-tuning typically needs
                              FAR fewer epochs than training from
                              scratch would, since most of the
                              "learning" already happened during
                              pretraining (Chapter 2)
  Batch size:                  16-32, often the LARGEST that fits
                                   in available GPU memory (DL
                                   Notes, Module 9, Ch.2's
                                   gradient accumulation is a
                                   common workaround for memory
                                   constraints)
  Warmup:                          a SHORT warmup period (DL
                                       Notes, Module 3, Ch.5) is
                                       standard practice, typically
                                       10% of total training steps
─────────────────────────────────────────

7. Summary & Next Steps

Key Takeaways

Fine-tuning BERT follows a consistent pattern: load pretrained weights, add a small task-specific head, and fine-tune at a low learning rate — directly applying the DL Notes' transfer learning framework.
Text classification uses the [CLS] token's pooled representation; NER (token classification) applies a classifier to every token position independently.
Extractive QA uses two independent classification heads (start and end position) applied per-token; sentence-pair tasks like NLI reuse BERT's native segment embeddings for two-sentence input.
BERT fine-tuning typically uses notably lower learning rates (2e-5 to 5e-5) and fewer epochs (2-4) than training from scratch, since most learning already occurred during pretraining.

Concept Check

Why does text classification use the [CLS] token's representation, while NER applies a classifier to every token position?
How does extractive question answering use two separate classification heads to identify an answer span?
Why do BERT fine-tuning recipes typically use much lower learning rates than training a CNN from scratch would?

Next Chapter

→ Chapter 4: The BERT Family — RoBERTa, ALBERT, DistilBERT, ELECTRA

Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to NLP & LLM Index