Chapter 11 · Fine-Tuning BERT

Chapter goals: Using representation models (BERT) as the core, systematically practice five types of fine-tuning methods: supervised classification, layer freezing strategies, few-shot classification (SetFit), masked language modeling (MLM) continued pre-training, and named entity recognition (NER).

11.1 Environment Preparation (Optional)

If you are viewing this notebook on Google Colab (or any other cloud platform), you need to uncomment and run the code block below to install the dependencies for this chapter:

💡 NOTE: Running the examples in this chapter requires a GPU. In Google Colab, go to Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4.

python

# %%capture
# !pip install "datasets>=2.18.0,<3" transformers>=4.38.2 sentence-transformers>=2.5.1 setfit>=1.0.3 accelerate>=0.27.2 seqeval>=1.2.2

11.2 Data Preparation

We use the classic Rotten Tomatoes movie review dataset, which is built into the Hugging Face datasets library:

python

from datasets import load_dataset

# Prepare data and splits
tomatoes = load_dataset("rotten_tomatoes")
train_data, test_data = tomatoes["train"], tomatoes["test"]

11.3 Supervised Classification

11.3.1 HuggingFace Trainer

First load the bert-base-cased model and its tokenizer; num_labels=2 indicates this is a binary classification task:

python

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load Model and Tokenizer
model_id = "bert-base-cased"
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_id)

Tokenize the data. To keep sequence lengths consistent within each batch, we use DataCollatorWithPadding to dynamically pad to the longest sequence in the batch:

python

from transformers import DataCollatorWithPadding

# Pad to the longest sequence in the batch
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

def preprocess_function(examples):
   """Tokenize input data"""
   return tokenizer(examples["text"], truncation=True)

# Tokenize train/test data
tokenized_train = train_data.map(preprocess_function, batched=True)
tokenized_test = test_data.map(preprocess_function, batched=True)

Define the evaluation metric — here we use the F1 score:

python

import numpy as np
import evaluate


def compute_metrics(eval_pred):
    """Calculate F1 score"""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    load_f1 = evaluate.load("f1")
    f1 = load_f1.compute(predictions=predictions, references=labels)["f1"]
    return {"f1": f1}

Configure training arguments and build the Trainer to execute the training process:

python

from transformers import TrainingArguments, Trainer

# Training arguments for parameter tuning
training_args = TrainingArguments(
   "model",
   learning_rate=2e-5,
   per_device_train_batch_size=16,
   per_device_eval_batch_size=16,
   num_train_epochs=1,
   weight_decay=0.01,
   save_strategy="epoch",
   report_to="none"
)

# Trainer which executes the training process
trainer = Trainer(
   model=model,
   args=training_args,
   train_dataset=tokenized_train,
   eval_dataset=tokenized_test,
   tokenizer=tokenizer,
   data_collator=data_collator,
   compute_metrics=compute_metrics,
)

Start training:

python

trainer.train()

Evaluate results:

python

trainer.evaluate()

11.3.2 Freezing All Layers Except the Classification Head

Full fine-tuning is expensive; a more economical approach is to train only the classification head while freezing the encoder backbone. First reload the model, then print all parameter names to inspect the structure:

python

# Load Model and Tokenizer
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_id)

python

# Print layer names
for name, param in model.named_parameters():
    print(name)

Parameters starting with classifier belong to the classification head — keep them trainable; freeze everything else:

python

for name, param in model.named_parameters():

     # Trainable classification head
     if name.startswith("classifier"):
        param.requires_grad = True

      # Freeze everything else
     else:
        param.requires_grad = False

Check whether the model was updated correctly:

python

# We can check whether the model was correctly updated
for name, param in model.named_parameters():
     print(f"Parameter: {name} ----- {param.requires_grad}")

Rerun training and evaluation with the same configuration as before:

python

from transformers import TrainingArguments, Trainer

# Trainer which executes the training process
trainer = Trainer(
   model=model,
   args=training_args,
   train_dataset=tokenized_train,
   eval_dataset=tokenized_test,
   tokenizer=tokenizer,
   data_collator=data_collator,
   compute_metrics=compute_metrics,
)
trainer.train()

python

trainer.evaluate()

11.3.3 Freezing Encoder Blocks 1–5

A step beyond "train only the classification head": unfreeze selected encoder blocks. First print the parameter list with indices to determine the index range of the encoder blocks:

python

# We can check whether the model was correctly updated
for index, (name, param) in enumerate(model.named_parameters()):
     print(f"Parameter: {index}{name} ----- {param.requires_grad}")

Experiments show that index 165 is the starting position of encoder block 10; freeze everything before it and train only the last few blocks:

python

# Load model
model_id = "bert-base-cased"
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Encoder block 10 starts at index 165 and
# we freeze everything before that block
for index, (name, param) in enumerate(model.named_parameters()):
    if index < 165:
        param.requires_grad = False

# Trainer which executes the training process
trainer = Trainer(
   model=model,
   args=training_args,
   train_dataset=tokenized_train,
   eval_dataset=tokenized_test,
   tokenizer=tokenizer,
   data_collator=data_collator,
   compute_metrics=compute_metrics,
)
trainer.train()
trainer.evaluate()

11.3.4 [BONUS] Systematic Experiment: Progressive Block Freezing

The commented-out code below is a complete ablation experiment: progressively unfreeze encoder blocks from "fully frozen" to "fully unfrozen," record the F1 score for each configuration, and finally plot the "number of trainable encoder blocks vs. F1" curve:

python

# scores = []
# for index in range(12):
#     # Re-load model
#     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
#     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

#     # Freeze encoder blocks 0-index
#     for name, param in model.named_parameters():
#         if "layer" in name:
#             layer_nr = int(name.split("layer")[1].split(".")[1])
#             if layer_nr <= index:
#                 param.requires_grad = False
#         else:
#             param.requires_grad = True

#     # Train
#     trainer = Trainer(
#       model=model,
#       args=training_args,
#       train_dataset=tokenized_train,
#       eval_dataset=tokenized_test,
#       tokenizer=tokenizer,
#       data_collator=data_collator,
#       compute_metrics=compute_metrics,
#     )
#     trainer.train()

#     # Evaluate
#     score = trainer.evaluate()["eval_f1"]
#     scores.append(score)

python

# scores

The plotting code is also preserved (commented), recording the measured F1 values under 12 configurations — you can see that performance stabilizes after unfreezing about 4 blocks:

python

# import matplotlib.pyplot as plt
# import numpy as np

# # Create Figure
# plt.figure(figsize=(8,4))

# # Prepare Data
# x = [f"0-{index}" for index in range(12)]
# x[0] = "None"
# x[-1] = "All"
# y = [
#     0.8541862652869239,
#     0.8525519848771267,
#     0.8514664143803217,
#     0.8506616257088847,
#     0.8398104265402844,
#     0.8391345249294448,
#     0.8377358490566037,
#     0.8433962264150944,
#     0.8258801141769743,
#     0.816247582205029,
#     0.7917485265225934,
#     0.7019400352733686
# ][::-1]

# # Stylize Figure
# plt.grid(color='#ECEFF1')
# plt.axvline(x=4, color="#EC407A", linestyle="--")
# plt.title("Effect of Frozen Encoder Blocks on Training Performance")
# plt.ylabel("F1-score")
# plt.xlabel("Trainable encoder blocks")

# # Plot Data
# plt.plot(x, y, color="black")

# # Additional Annotation
# plt.annotate(
#     'Performance stabilizing',
#     xy=(4, y[4]),
#     xytext=(4.5, y[4]-.05),
#     arrowprops=dict(
#         arrowstyle="-|>",
#         connectionstyle="arc3",
#         color="#00ACC1")
# )
# plt.savefig("multiple_frozen_blocks.png", dpi=300, bbox_inches='tight')

11.4 Few-Shot Classification

When labeled data is scarce (only a dozen samples per class), traditional fine-tuning struggles to converge. SetFit solves this problem through contrastive learning. First simulate a few-shot setting with sample_dataset — sample 16 examples per class:

python

from setfit import sample_dataset

# We simulate a few-shot setting by sampling 16 examples per class
sampled_train_data = sample_dataset(tomatoes["train"], num_samples=16)

Load a pre-trained SentenceTransformer model as the backbone:

python

from setfit import SetFitModel

# Load a pre-trained SentenceTransformer model
model = SetFitModel.from_pretrained("sentence-transformers/all-mpnet-base-v2")

SetFit training has two stages: first contrastive learning on text pairs, then training the classification head. When defining training arguments, control how many text pairs are generated via num_iterations:

python

from setfit import TrainingArguments as SetFitTrainingArguments
from setfit import Trainer as SetFitTrainer

# Define training arguments
args = SetFitTrainingArguments(
    num_epochs=3, # The number of epochs to use for contrastive learning
    num_iterations=20  # The number of text pairs to generate
)
args.eval_strategy = args.evaluation_strategy

# Create trainer
trainer = SetFitTrainer(
    model=model,
    args=args,
    train_dataset=sampled_train_data,
    eval_dataset=test_data,
    metric="f1"
)

The older SetFit syntax does not require explicitly passing args; it is kept in comments for reference:

python

# from setfit import SetFitTrainer

# # Create trainer
# trainer = SetFitTrainer(
#     model=model,
#     train_dataset=sampled_train_data,
#     eval_dataset=test_data,
#     metric="f1",
#     num_epochs=3, # The number of epochs to use for contrastive learning
# )

Execute the training loop:

python

# Training loop
trainer.train()

Evaluate on the test set:

python

# Evaluate the model on our test data
trainer.evaluate()

Inspect the trained model head structure:

python

model.model_head

11.5 MLM Continued Pre-training

If your domain data differs significantly from BERT's pre-training corpus, you can first continue pre-training using Masked Language Modeling (MLM) for domain adaptation. Load AutoModelForMaskedLM instead of a sequence classification model:

python

from transformers import AutoTokenizer, AutoModelForMaskedLM

# Load model for Masked Language Modeling (MLM)
model = AutoModelForMaskedLM.from_pretrained("bert-base-cased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

MLM requires no labels; simply remove the label column:

python

def preprocess_function(examples):
   return tokenizer(examples["text"], truncation=True)

# Tokenize data
tokenized_train = train_data.map(preprocess_function, batched=True)
tokenized_train = tokenized_train.remove_columns("label")
tokenized_test = test_data.map(preprocess_function, batched=True)
tokenized_test = tokenized_test.remove_columns("label")

Use DataCollatorForLanguageModeling to automatically mask tokens at a 15% probability:

python

from transformers import DataCollatorForLanguageModeling

# Masking Tokens
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15
)

In addition to token-level masking, whole-word masking is also available (code kept in comments):

python

# from transformers import DataCollatorForWholeWordMask

# # Masking Whole Words
# data_collator = DataCollatorForWholeWordMask(
#     tokenizer=tokenizer,
#     mlm=True,
#     mlm_probability=0.15
# )

MLM training typically runs for more epochs than classification tasks (10 epochs here):

python

# Training arguments for parameter tuning
training_args = TrainingArguments(
   "model",
   learning_rate=2e-5,
   per_device_train_batch_size=16,
   per_device_eval_batch_size=16,
   num_train_epochs=10,
   weight_decay=0.01,
   save_strategy="epoch",
   report_to="none"
)


# Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    tokenizer=tokenizer,
    data_collator=data_collator
)

Save the tokenizer → train → save the updated model:

python

# Save pre-trained tokenizer
tokenizer.save_pretrained("mlm")

# Train model
trainer.train()

# Save updated model
model.save_pretrained("mlm")

Predictions on the fill-mask task using the original BERT before training:

python

from transformers import pipeline

# Load and create predictions
mask_filler = pipeline("fill-mask", model="bert-base-cased")
preds = mask_filler("What a horrible [MASK]!")

# Print results
for pred in preds:
    print(f">>> {pred['sequence']}")

After domain adaptation, run the same prediction with the mlm model and compare the output differences:

python

# Load and create predictions
mask_filler = pipeline("fill-mask", model="mlm")
preds = mask_filler("What a horrible [MASK]!")

# Print results
for pred in preds:
    print(f">>> {pred['sequence']}")

11.6 Named Entity Recognition

Here are some NER datasets worth exploring:

tner/mit_movie_trivia
tner/mit_restaurant
wnut_17
conll2003

Import NER-related modules:

python

from transformers import AutoModelForTokenClassification, AutoTokenizer
from transformers import DataCollatorWithPadding
from transformers import TrainingArguments, Trainer
import numpy as np

Load the classic CoNLL-2003 named entity recognition dataset:

python

# The CoNLL-2003 dataset for NER
dataset = load_dataset("conll2003", trust_remote_code=True)

Inspect one sample:

python

example = dataset["train"][848]
example

Define the label mapping (O/B-PER/I-PER/B-ORG, etc., 9 classes total):

python

label2id = {
    'O': 0, 'B-PER': 1, 'I-PER': 2, 'B-ORG': 3, 'I-ORG': 4,
    'B-LOC': 5, 'I-LOC': 6, 'B-MISC': 7, 'I-MISC': 8
}
id2label = {index: label for label, index in label2id.items()}
label2id

Load the token classification model and pass in id2label/label2id so that inference results have readable labels:

python

from transformers import AutoModelForTokenClassification

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

# Load model
model = AutoModelForTokenClassification.from_pretrained(
    "bert-base-cased",
    num_labels=len(id2label),
    id2label=id2label,
    label2id=label2id
)

The core challenge in NER: words are split into sub-tokens, but labels are word-level. First split the words to inspect the sub-token situation:

python

# Split individual tokens into sub-tokens
token_ids = tokenizer(example["tokens"], is_split_into_words=True)["input_ids"]
sub_tokens = tokenizer.convert_ids_to_tokens(token_ids)
sub_tokens

The align_labels function aligns labels: special tokens are marked -100 (not participating in loss computation); the first sub-token of a word inherits the original label; subsequent sub-tokens that are B-XXX are changed to I-XXX:

python

def align_labels(examples):
    token_ids = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = examples["ner_tags"]

    updated_labels = []
    for index, label in enumerate(labels):

        # Map tokens to their respective word
        word_ids = token_ids.word_ids(batch_index=index)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:

            # The start of a new word
            if word_idx != previous_word_idx:

                previous_word_idx = word_idx
                updated_label = -100 if word_idx is None else label[word_idx]
                label_ids.append(updated_label)

            # Special token is -100
            elif word_idx is None:
                label_ids.append(-100)

            # If the label is B-XXX we change it to I-XXX
            else:
                updated_label = label[word_idx]
                if updated_label % 2 == 1:
                    updated_label += 1
                label_ids.append(updated_label)

        updated_labels.append(label_ids)

    token_ids["labels"] = updated_labels
    return token_ids

tokenized = dataset.map(align_labels, batched=True)

Compare original and aligned labels:

python

# Difference between original and updated labels
print(f"Original: {example['ner_tags']}")
print(f"Updated: {tokenized['train'][848]['labels']}")

NER uses the seqeval library for sequence-level evaluation, collecting non -100 predictions and true labels per document per token:

python

import evaluate

# Load sequential evaluation
seqeval = evaluate.load("seqeval")

def compute_metrics(eval_pred):
    # Create predictions
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=2)

    true_predictions = []
    true_labels = []

    # Document-level iteration
    for prediction, label in zip(predictions, labels):

      # token-level iteration
      for token_prediction, token_label in zip(prediction, label):

        # We ignore special tokens
        if token_label != -100:
          true_predictions.append([id2label[token_prediction]])
          true_labels.append([id2label[token_label]])

    results = seqeval.compute(predictions=true_predictions, references=true_labels)
    return {"f1": results["overall_f1"]}

Token classification tasks require a dedicated DataCollatorForTokenClassification padding collator:

python

from transformers import DataCollatorForTokenClassification

# Token-classification Data Collator
data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

Training configuration is the same as before; initialize the Trainer and train:

python

# Training arguments for parameter tuning
training_args = TrainingArguments(
   "model",
   learning_rate=2e-5,
   per_device_train_batch_size=16,
   per_device_eval_batch_size=16,
   num_train_epochs=1,
   weight_decay=0.01,
   save_strategy="epoch",
   report_to="none"
)

# Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)
trainer.train()

Evaluate the test set:

python

# Evaluate the model on our test data
trainer.evaluate()

Save the fine-tuned model and run inference on a new sentence using the token-classification pipeline:

python

from transformers import pipeline

# Save our fine-tuned model
trainer.save_model("ner_model")

# Run inference on the fine-tuned model
token_classifier = pipeline(
    "token-classification",
    model="ner_model",
)
token_classifier("My name is Maarten.")

Chapter Summary

The standard supervised fine-tuning pipeline: load model → tokenize (dynamic padding) → define metrics → Trainer train and evaluate;
Freezing strategies can significantly reduce training cost: training only the classification head is cheapest; unfreezing selected encoder blocks trades off performance — ablation experiments show F1 stabilizes after unfreezing about 4 blocks;
For few-shot scenarios, SetFit is the preferred choice: even with only 16 samples per class, contrastive learning achieves notable results;
When domain data differs significantly, perform MLM continued pre-training first, then downstream fine-tuning; DataCollatorForLanguageModeling auto-masks at 15% probability;
The key to NER is sub-word to word label alignment: special tokens marked -100; first sub-token inherits the label; subsequent sub-tokens B→I.

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. In the HuggingFace Trainer fine-tuning flow, what is the role of DataCollatorWithPadding?

2. Regarding layer freezing strategies, which statement is correct?

3. What are the two training stages of SetFit few-shot classification?

4. During NER label alignment, why are special tokens assigned the label -100?

🛠️ Hands-on Practice

Reverse the freezing strategy in Section 11.3.2: freeze only the first half of parameters outside the classifier head (e.g. parameters starting with bert.encoder.layer.0. through layer.4.); after training, compare the F1 score with "train only the classification head."
Change SetFit's num_samples from 16 to 8 and 32 respectively; keep everything else unchanged; plot the F1 score curves across the three settings to验证 the impact of few-shot sample count on performance.
Following the flow in Section 11.6, swap the CoNLL-2003 dataset for wnut_17; inspect the difference in label sets versus label2id; adjust the mapping and complete training and evaluation.

Chapter 11 · Fine-Tuning BERT ​

11.1 Environment Preparation (Optional) ​

11.2 Data Preparation ​

11.3 Supervised Classification ​

11.3.1 HuggingFace Trainer ​

11.3.2 Freezing All Layers Except the Classification Head ​

11.3.3 Freezing Encoder Blocks 1–5 ​

11.3.4 [BONUS] Systematic Experiment: Progressive Block Freezing ​

11.4 Few-Shot Classification ​

11.5 MLM Continued Pre-training ​

11.6 Named Entity Recognition ​

Chapter Summary ​

🧪 随堂测验

🛠️ Hands-on Practice ​

Chapter 11 · Fine-Tuning BERT

11.1 Environment Preparation (Optional)

11.2 Data Preparation

11.3 Supervised Classification

11.3.1 HuggingFace Trainer

11.3.2 Freezing All Layers Except the Classification Head

11.3.3 Freezing Encoder Blocks 1–5

11.3.4 [BONUS] Systematic Experiment: Progressive Block Freezing

11.4 Few-Shot Classification

11.5 MLM Continued Pre-training

11.6 Named Entity Recognition

Chapter Summary

🛠️ Hands-on Practice