217 Fine-Tuning BERT for Downstream Tasks
Transfer learning transformed natural language processing in 2018 when Devlin et al. demonstrated that a single pretrained bidirectional encoder could be adapted to over a dozen NLP benchmarks by appending a lightweight task head and training end-to-end on modest labeled datasets. The resulting model, BERT (Bidirectional Encoder Representations from Transformers), established state-of-the-art results on the GLUE benchmark, SQuAD question answering, and named-entity recognition simultaneously, making fine-tuning the dominant paradigm for applied NLP. This chapter develops the theory and practice of BERT fine-tuning in full: the pretraining objectives that create transferable representations, the tokenization pipeline that maps raw text to model inputs, the classification head architecture, and the complete training workflow including hyperparameter selection justified by empirical evidence.
217.1 1. BERT Architecture and Pretraining
217.1.1 1.1 Bidirectional Encoding
Standard language models factorize the joint probability of a token sequence left-to-right: \(p(x_1, \ldots, x_T) = \prod_{t=1}^T p(x_t \mid x_1, \ldots, x_{t-1})\). This autoregressive constraint means each token representation attends only to its left context, losing information from future tokens. BERT abandons this constraint entirely. Instead of modeling token sequences as a directed process, BERT trains an encoder that conditions on the entire context simultaneously. Every layer of the transformer stack applies multi-head self-attention across all token positions bidirectionally, yielding contextualized representations that incorporate both left and right context from the first layer onward.
BERT-base consists of 12 transformer encoder layers, each with hidden dimension \(d = 768\), \(h = 12\) attention heads (head dimension \(d_k = 64\)), and a feed-forward sublayer with inner dimension \(d_{ff} = 3072\). The feed-forward sublayer applies two linear transformations with a GELU nonlinearity:
\[\text{FFN}(\mathbf{x}) = W_2 \, \text{GELU}(W_1 \mathbf{x} + b_1) + b_2\]
where \(W_1 \in \mathbb{R}^{3072 \times 768}\) and \(W_2 \in \mathbb{R}^{768 \times 3072}\). The total parameter count for BERT-base is approximately 110 million. BERT-large scales to 24 layers, \(d = 1024\), 16 heads, and 340 million parameters.
217.1.2 1.2 Pretraining Objectives
BERT is pretrained on BookCorpus (800M words) and English Wikipedia (2.5B words) using two self-supervised objectives that together teach the model syntactic structure, semantic similarity, and cross-sentence coherence.
Masked Language Modeling (MLM). Fifteen percent of WordPiece tokens are selected at random. Of these, 80% are replaced with [MASK], 10% with a random vocabulary token, and 10% left unchanged. The model must predict the original token at every masked position. Let \(\mathcal{M}\) denote the set of masked positions. The MLM loss is:
\[\mathcal{L}_{\text{MLM}} = -\frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \log p_\theta(x_i \mid \tilde{\mathbf{x}})\]
where \(\tilde{\mathbf{x}}\) is the corrupted sequence. Because the corruption is stochastic and the model must also explain unmasked tokens, it cannot memorize the input and must build genuine contextual representations.
Next Sentence Prediction (NSP). Given two text segments \(A\) and \(B\), the model predicts whether \(B\) is the actual continuation of \(A\) in the corpus (label IsNext) or a random sentence from a different document (label NotNext). NSP is supervised via the [CLS] token, a special token prepended to every input whose final-layer representation aggregates information about the full input pair and is used as the sentence-level embedding. The NSP loss is binary cross-entropy on this aggregate representation. While subsequent work (Liu et al. 2019 on RoBERTa) found NSP contributes less than MLM to downstream performance, it serves a pedagogical function: it teaches the model to represent inter-sentence relationships, which is essential for tasks like natural language inference.
The combined pretraining loss is \(\mathcal{L} = \mathcal{L}_{\text{MLM}} + \mathcal{L}_{\text{NSP}}\).
217.1.3 1.3 The [CLS] Token as Aggregate Representation
The input sequence to BERT always begins with the special token [CLS] and separates segment pairs with [SEP]. The final-layer hidden state \(\mathbf{h}_{[\text{CLS}]} \in \mathbb{R}^{768}\) accumulates global information about the entire input through bidirectional attention across all 12 layers. For classification tasks, this vector is the natural choice for attaching a prediction head, because it has no positional bias toward any particular span of text — every position can write information into it via attention.
217.2 2. The Fine-Tuning Paradigm
217.2.1 2.1 Why Pretrained Representations Transfer
The pretraining corpus is orders of magnitude larger than any task-specific labeled dataset. During pretraining, BERT’s weights learn distributional properties of language: morphological inflection patterns, syntactic constituency, coreference, and semantic role labeling all emerge as implicit solutions to the MLM objective (Tenney et al. 2019). When fine-tuning on a downstream task, these representations provide a powerful initialization. The labeled data need only teach the model to map existing linguistic knowledge onto task-specific categories rather than learning language from scratch.
This is formalized as follows. Let \(\theta_0\) denote pretrained BERT weights and \(\phi\) the randomly initialized task head parameters. The fine-tuning objective on a labeled dataset \(\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N\) is:
\[\min_{\theta, \phi} \frac{1}{N} \sum_{i=1}^N \mathcal{L}_{\text{task}}(f_\phi(\text{BERT}_\theta(x_i)), y_i)\]
All parameters \(\theta\) and \(\phi\) are updated jointly via gradient descent, but because \(\theta\) starts from a strong initialization, only small updates (learning rate \(\approx 2 \times 10^{-5}\)) are needed to specialize the representations.
217.2.2 2.2 Classification Head Architecture
For sequence classification, the task head is minimal. The [CLS] embedding passes through dropout (rate 0.1 by default), then a single linear layer projects to the number of class logits:
\[\hat{\mathbf{y}} = W \, \text{dropout}(\mathbf{h}_{[\text{CLS}]}) + b, \quad W \in \mathbb{R}^{K \times 768}\]
where \(K\) is the number of classes. The cross-entropy loss over a batch of \(B\) examples is:
\[\mathcal{L}_{\text{CE}} = -\frac{1}{B} \sum_{i=1}^B \log \frac{\exp(\hat{y}_{i, y_i})}{\sum_{k=1}^K \exp(\hat{y}_{i,k})}\]
The entire stack — pretrained encoder plus linear head — is differentiable end-to-end, so backpropagation updates both the encoder and the head simultaneously. AutoModelForSequenceClassification from HuggingFace Transformers instantiates exactly this architecture: it loads pretrained BERT weights and attaches the classification head with randomly initialized \(W\) and \(b\).
217.3 3. Tokenization in Depth
217.3.1 3.1 WordPiece Subword Tokenization
BERT uses WordPiece tokenization (Schuster and Nakamura 2012), which constructs a vocabulary of subword units by iteratively merging character pairs that maximize the likelihood of the training corpus under a unigram language model. The vocabulary size is 30,522 for BERT-base-uncased. Words that appear frequently in the corpus remain as single tokens; rare or compound words are split into subwords, with continuation tokens marked by a ## prefix.
For example, “tokenization” may be split into ['token', '##ization'] and “BERT” remains a single token ['bert'] in the uncased model (which lowercases input). Unknown characters fall back to [UNK], though this is rare because the vocabulary covers essentially all Unicode code points via individual characters.
217.3.2 3.2 Encoding Pipeline
AutoTokenizer converts raw text to the three integer arrays BERT expects:
input_ids: integer indices of each token in the vocabulary, including[CLS]at position 0 and[SEP]at the end of each segment.attention_mask: binary array, 1 for real tokens and 0 for padding positions. The self-attention mechanism multiplies attention logits by this mask, preventing the model from attending to padding.token_type_ids: 0 for tokens in segment \(A\) (including[CLS]and first[SEP]), 1 for tokens in segment \(B\). This signals to BERT which sentence each token belongs to.
BERT accepts a maximum sequence length of 512 tokens. Text longer than this limit must be truncated. The tokenizer handles this with truncation=True, max_length=512. Shorter sequences are padded to a common length within a batch using padding='max_length' or padding=True (pad to longest in batch). Padding to the longest sequence in the batch is preferred for efficiency, since it reduces the average sequence length.
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
examples = [
"The film is a masterpiece of narrative economy.",
"I fell asleep halfway through the second act.",
]
encoding = tokenizer(
examples,
padding=True,
truncation=True,
max_length=64,
return_tensors="pt",
)
for key, val in encoding.items():
print(f"{key}: shape={tuple(val.shape)}")
print(f" first example: {val[0].tolist()}")The output reveals how padding works: both sequences are padded to the length of the longer one, with attention mask zeros covering the padding positions.
217.4 4. Custom Dataset Class
Before constructing the full training loop, it is useful to see the low-level dataset abstraction. Subclassing torch.utils.data.Dataset requires implementing three methods: __init__ to store tokenized data and labels, __len__ to report dataset size, and __getitem__ to return a dictionary of tensors for a single example. The dictionary keys must match the argument names expected by the model’s forward method.
import torch
from torch.utils.data import Dataset
class TextClassificationDataset(Dataset):
def __init__(self, encodings, labels):
self.input_ids = encodings["input_ids"]
self.attention_mask = encodings["attention_mask"]
self.labels = torch.tensor(labels, dtype=torch.long)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return {
"input_ids": self.input_ids[idx],
"attention_mask": self.attention_mask[idx],
"labels": self.labels[idx],
}This class is agnostic to tokenizer choice and works with any HuggingFace model that accepts input_ids, attention_mask, and labels.
217.5 5. CPU-Feasible Demo: DistilBERT on a Tiny Subset
DistilBERT (Sanh et al. 2019) is a knowledge-distilled version of BERT with 6 layers instead of 12, 40% fewer parameters (66M), and 60% faster inference while retaining 97% of BERT’s GLUE score. Because it fits comfortably in CPU memory and runs in minutes on small datasets, it is the ideal vehicle for demonstrating the full fine-tuning pipeline end-to-end.
The following demo uses 100 training examples and 20 validation examples from SST-2 (binary sentiment: positive/negative). It trains for 2 epochs using plain PyTorch, printing loss and accuracy at each epoch.
import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset
import evaluate
torch.manual_seed(42)
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased", num_labels=2
)
raw = load_dataset("sst2", split="train[:120]")
train_raw = raw.select(range(100))
val_raw = raw.select(range(100, 120))
def tokenize(batch):
return tokenizer(
batch["sentence"],
padding="max_length",
truncation=True,
max_length=64,
)
train_enc = train_raw.map(tokenize, batched=True)
val_enc = val_raw.map(tokenize, batched=True)
train_enc.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
val_enc.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
def collate_fn(batch):
return {
"input_ids": torch.stack([b["input_ids"] for b in batch]),
"attention_mask": torch.stack([b["attention_mask"] for b in batch]),
"labels": torch.tensor([b["label"] for b in batch]),
}
train_loader = DataLoader(train_enc, batch_size=16, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_enc, batch_size=16, collate_fn=collate_fn)
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
accuracy_metric = evaluate.load("accuracy")
device = torch.device("cpu")
model.to(device)
for epoch in range(2):
model.train()
total_loss = 0.0
for batch in train_loader:
batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()
total_loss += loss.item()
model.eval()
all_preds, all_labels = [], []
with torch.no_grad():
for batch in val_loader:
batch = {k: v.to(device) for k, v in batch.items()}
logits = model(**batch).logits
preds = logits.argmax(dim=-1)
all_preds.extend(preds.tolist())
all_labels.extend(batch["labels"].tolist())
acc = accuracy_metric.compute(predictions=all_preds, references=all_labels)
print(f"Epoch {epoch+1} | loss={total_loss/len(train_loader):.4f} | val_acc={acc['accuracy']:.3f}")Running this on CPU completes in roughly two to four minutes and yields validation accuracy above 0.80 on the 20-example split, demonstrating that even minimal fine-tuning of distilled BERT representations achieves reasonable sentiment discrimination.
217.6 6. Full Fine-Tuning Workflow with HuggingFace Trainer
For production training on a full dataset, HuggingFace Trainer automates learning rate scheduling, gradient accumulation, distributed training, and checkpointing. The code below shows the complete workflow for SST-2 using BERT-base-uncased on GPU.
# GPU required — shown for reference
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
)
from datasets import load_dataset
import evaluate
import numpy as np
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased", num_labels=2
)
dataset = load_dataset("glue", "sst2")
def tokenize(batch):
return tokenizer(
batch["sentence"],
padding="max_length",
truncation=True,
max_length=128,
)
tokenized = dataset.map(tokenize, batched=True)
tokenized = tokenized.rename_column("label", "labels")
tokenized.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=-1)
return {
"accuracy": accuracy.compute(predictions=preds, references=labels)["accuracy"],
"f1": f1.compute(predictions=preds, references=labels, average="binary")["f1"],
}
args = TrainingArguments(
output_dir="./bert-sst2",
learning_rate=2e-5,
num_train_epochs=3,
per_device_train_batch_size=32,
per_device_eval_batch_size=64,
weight_decay=0.01,
warmup_steps=500,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
metric_for_best_model="accuracy",
max_grad_norm=1.0,
)
trainer = Trainer(
model=model,
args=args,
train_dataset=tokenized["train"],
eval_dataset=tokenized["validation"],
compute_metrics=compute_metrics,
)
trainer.train()This configuration achieves approximately 92-93% accuracy on the SST-2 validation set (67,349 training examples) within three epochs, consistent with the original BERT paper results. The load_best_model_at_end=True flag combined with metric_for_best_model="accuracy" implements early stopping in practice: the checkpoint with the highest validation accuracy is restored at the end of training.
217.7 7. Hyperparameter Selection
217.7.1 7.1 Learning Rate
The learning rate is the most sensitive hyperparameter in BERT fine-tuning. Sun et al. (2019) systematically investigated learning rates from \(10^{-5}\) to \(10^{-4}\) and found that rates above \(5 \times 10^{-5}\) consistently degrade performance, often causing catastrophic forgetting where the fine-tuned model loses pretrained linguistic knowledge within the first few steps. Rates below \(10^{-5}\) are too conservative: the model changes too slowly and may not converge within a few epochs.
The theoretical explanation is that BERT’s pretrained loss landscape is flat near the initialization point \(\theta_0\): gradient magnitudes are small because the model already fits general language patterns well. A large learning rate creates oscillations around \(\theta_0\) without finding task-specific minima; a small rate navigates the flat landscape but converges slowly. The empirically recommended range is \(\{1 \times 10^{-5}, 2 \times 10^{-5}, 3 \times 10^{-5}, 5 \times 10^{-5}\}\).
217.7.2 7.2 Learning Rate Scheduling with Linear Warmup
A linear warmup schedule increases the learning rate from 0 to the peak value over a specified number of steps, then linearly decays it back toward 0 over the remaining training steps. Let \(t\) denote the current step, \(T_{\text{warm}}\) the warmup steps, and \(T_{\text{total}}\) the total training steps:
\[\eta(t) = \eta_{\max} \cdot \begin{cases} t / T_{\text{warm}} & t \leq T_{\text{warm}} \\ (T_{\text{total}} - t) / (T_{\text{total}} - T_{\text{warm}}) & t > T_{\text{warm}} \end{cases}\]
Warmup prevents the large gradient magnitudes that occur at initialization from corrupting pretrained weights before the optimizer has seen enough data to compute reliable gradient statistics. Devlin et al. used 10% of total steps as warmup; 500 fixed warmup steps is also a common choice for datasets of moderate size.
217.7.3 7.3 Batch Size and Epochs
BERT fine-tuning is effective with batch sizes of 16 or 32. Larger batches reduce variance in gradient estimates but increase memory requirements; with gradient accumulation, effective batch sizes of 256 or 512 can be simulated on single-GPU hardware. The number of fine-tuning epochs is typically 2-4: more epochs risk overfitting because task-specific labeled datasets are small relative to the pretraining corpus.
217.7.4 7.4 Gradient Clipping
Gradient clipping caps the \(\ell_2\) norm of the gradient vector at a threshold \(\tau\):
\[\tilde{g} = g \cdot \min\left(1, \frac{\tau}{\|g\|_2}\right)\]
With \(\tau = 1.0\), clipping prevents individual large-gradient steps from overwriting pretrained representations. This is especially important in early fine-tuning steps when the random task head produces large loss values and consequently large gradients that propagate deep into the pretrained layers.
217.7.5 7.5 Weight Decay
L2 regularization via weight decay (\(\lambda = 0.01\)) applies to all parameters except bias terms and layer normalization parameters, which are left unregularized. HuggingFace AdamW handles this correctly by default.
217.8 8. Practical Considerations
Sequence length. Many classification tasks require sequences far shorter than 512 tokens. SST-2 sentences average 20 tokens. Using max_length=128 rather than 512 reduces memory by \(4\times\) and computation by \(16\times\) (attention complexity is \(O(L^2)\)), enabling larger batch sizes and faster iteration.
Layer-wise learning rate decay. A refinement of fixed learning rate scheduling applies smaller learning rates to lower layers and larger rates to upper layers, reflecting the fact that lower layers encode more general features that should change less. Layer \(\ell\) receives learning rate \(\eta_\ell = \eta_{\text{max}} \cdot \alpha^{L - \ell}\) for decay factor \(\alpha \in [0.9, 0.95]\). This technique is particularly valuable when fine-tuning data is scarce.
Multiple fine-tuning runs. Because BERT fine-tuning from the same pretrained checkpoint but different random seeds can yield validation accuracy varying by 1-2% (Dodge et al. 2020), best practice is to report the median or mean over three to five random seeds rather than a single run. The HuggingFace Trainer accepts a seed argument to control reproducibility.
Task-specific architectures beyond classification. While this chapter focused on the [CLS]-based classification head, other tasks use different output layers. Token classification (NER, POS tagging) attaches a linear layer to every token’s final hidden state. Extractive question answering (SQuAD) uses two linear layers over all token positions to predict start and end span logits. Natural language inference uses the same [CLS] head with three classes (entailment, neutral, contradiction). The pretrained encoder is reused identically across all these variants; only the head architecture changes.
217.9 References
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2019). BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT 2019, pp. 4171-4186. arXiv:1810.04805. https://arxiv.org/abs/1810.04805
Sun, C., Qiu, X., Xu, Y., and Huang, X. (2019). How to fine-tune BERT for text classification? In China National Conference on Chinese Computational Linguistics, pp. 194-206. arXiv:1905.05583. https://arxiv.org/abs/1905.05583
Sanh, V., Debut, L., Chaumond, J., and Wolf, T. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv:1910.01108. https://arxiv.org/abs/1910.01108
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. (2019). RoBERTa: A robustly optimized BERT pretraining approach. arXiv:1907.11692. https://arxiv.org/abs/1907.11692
Tenney, I., Das, D., and Pavlick, E. (2019). BERT rediscovers the classical NLP pipeline. In Proceedings of ACL 2019, pp. 4593-4601. arXiv:1905.05950. https://arxiv.org/abs/1905.05950
Dodge, J., Ilharco, G., Schwartz, R., Farhadi, A., Hajishirzi, H., and Smith, N. (2020). Fine-tuning pretrained language models: Weight initializations, data orders, and early stopping. arXiv:2002.06305. https://arxiv.org/abs/2002.06305