BERT Fine-tuning › Classification Fine-tuning
Language AI HandbookAll chapters BERT Fine-tuning › Classification Fine-tuningOne section · Sections in this chapter3 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Sentiment analysis, spam detection, topic classification: these tasks require mapping an entire text to a single label. BERT handles classification through its [CLS] token.
The Classification Architecture
To classify a sentence, we need to reduce BERT's variable-length sequence of token representations into a single fixed-size vector. But which tokens should we use? We could average all token representations, but that treats every word equally, giving "the" as much weight as "terrible" in a movie review. We could use the final token, but that's arbitrary.
BERT's solution is elegant: it reserves a special [CLS] token at the beginning of every input specifically for sequence-level tasks. During pre-training, this token learned to aggregate information from the entire sequence for the Next Sentence Prediction objective. By the time pre-training finishes, the [CLS] representation has become a 768-dimensional summary of the sequence's meaning.
For classification, we add a single linear layer on top of this representation. The layer maps the 768-dimensional [CLS] vector into a vector of class scores:
where:
- : the output logits vector with dimension (number of classes)
- : the weight matrix of shape that projects the hidden state to class space
- : the 768-dimensional hidden state of the
[CLS]token from BERT's final layer - : the bias vector of dimension
Each element of represents the model's confidence that the input belongs to that class, called a logit or unnormalized score. Higher values mean higher confidence. To convert these logits into proper probabilities that sum to 1, we apply the softmax function. During training, we minimize the cross-entropy loss between these predicted probabilities and the true labels, which pushes the model to assign high probability to the correct class and low probability to incorrect ones.
from transformers import BertModel
class BertForClassification(nn.Module):
"""BERT with a classification head for sequence classification."""
def __init__(
self,
num_classes: int,
model_name: str = "bert-base-uncased",
dropout: float = 0.1,
):
super().__init__()
self.bert = BertModel.from_pretrained(model_name)
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)
def forward(self, input_ids, attention_mask=None, token_type_ids=None):
# Get BERT outputs
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
)
# Use [CLS] token representation (first token)
cls_output = outputs.last_hidden_state[:, 0, :]
# Apply dropout and classify
cls_output = self.dropout(cls_output)
logits = self.classifier(cls_output)
return logitsfrom transformers import BertTokenizer
# Create a sentiment classifier
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForClassification(num_classes=2) # Binary: positive/negative
# Example inference
text = "This movie was absolutely fantastic!"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)model.eval()
with torch.no_grad():
logits = model(**inputs)
probs = F.softmax(logits, dim=-1)Input: 'This movie was absolutely fantastic!' Logits shape: torch.Size([1, 2]) Probabilities: Negative=0.5580, Positive=0.4420
The probabilities are nearly equal, which is expected: the classifier produces random predictions because its weights are randomly initialized. BERT's layers contain useful pre-trained representations, but the classifier hasn't learned to use them yet. After fine-tuning on labeled sentiment data, these probabilities would reflect meaningful predictions.
Training Loop for Classification
Fine-tuning requires careful attention to learning rates, batch sizes, and training duration. Here's a complete training loop:
from torch.optim import AdamW
from torch.utils.data import TensorDataset
from transformers import get_linear_schedule_with_warmup
def create_classification_dataset(texts, labels, tokenizer, max_length=128):
"""Tokenize texts and create a dataset."""
encodings = tokenizer(
texts,
truncation=True,
padding="max_length",
max_length=max_length,
return_tensors="pt",
)
dataset = TensorDataset(
encodings["input_ids"],
encodings["attention_mask"],
torch.tensor(labels),
)
return dataset
def train_classifier(
model,
train_dataloader,
val_dataloader,
num_epochs=3,
learning_rate=2e-5,
warmup_ratio=0.1,
device="cpu",
):
"""Fine-tune BERT for classification."""
model = model.to(device)
model.train()
# Optimizer with different learning rates is optional but common
optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
# Learning rate scheduler with warmup
total_steps = len(train_dataloader) * num_epochs
warmup_steps = int(total_steps * warmup_ratio)
scheduler = get_linear_schedule_with_warmup(
optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
)
loss_fn = nn.CrossEntropyLoss()
history = {"train_loss": [], "val_loss": [], "val_accuracy": []}
for epoch in range(num_epochs):
# Training
model.train()
epoch_loss = 0
for batch in train_dataloader:
input_ids, attention_mask, labels = [b.to(device) for b in batch]
optimizer.zero_grad()
logits = model(input_ids, attention_mask)
loss = loss_fn(logits, labels)
loss.backward()
# Gradient clipping prevents exploding gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step()
epoch_loss += loss.item()
avg_train_loss = epoch_loss / len(train_dataloader)
history["train_loss"].append(avg_train_loss)
# Validation
model.eval()
val_loss = 0
correct = 0
total = 0
with torch.no_grad():
for batch in val_dataloader:
input_ids, attention_mask, labels = [
b.to(device) for b in batch
]
logits = model(input_ids, attention_mask)
loss = loss_fn(logits, labels)
val_loss += loss.item()
preds = torch.argmax(logits, dim=-1)
correct += (preds == labels).sum().item()
total += labels.size(0)
avg_val_loss = val_loss / len(val_dataloader)
val_accuracy = correct / total
history["val_loss"].append(avg_val_loss)
history["val_accuracy"].append(val_accuracy)
return historyfrom torch.utils.data import DataLoader
# Create a small demo dataset
train_texts = [
"I loved this movie, it was incredible!",
"Best film I've seen all year.",
"Absolutely wonderful storytelling.",
"A masterpiece of cinema.",
"Brilliantly acted and directed.",
"Terrible waste of time.",
"I hated every minute of it.",
"Worst movie I've ever seen.",
"Complete garbage, avoid at all costs.",
"Boring and predictable plot.",
]
train_labels = [1, 1, 1, 1, 1, 0, 0, 0, 0, 0] # 1=positive, 0=negative
val_texts = [
"Really enjoyed watching this.",
"Not worth the ticket price.",
]
val_labels = [1, 0]
# Create datasets
train_dataset = create_classification_dataset(
train_texts, train_labels, tokenizer
)
val_dataset = create_classification_dataset(val_texts, val_labels, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2)
# Train (in practice, you'd use more data and epochs)
model_fresh = BertForClassification(num_classes=2)
history = train_classifier(model_fresh, train_loader, val_loader, num_epochs=2)Training Complete ---------------------------------------- Epoch 1: Train Loss=0.7341, Val Loss=0.6740, Val Acc=50.00% Epoch 2: Train Loss=0.5325, Val Loss=0.6292, Val Acc=50.00%
The decreasing training loss indicates the model is learning from the examples. Even with only 10 training examples and 2 epochs, the model begins adapting to the sentiment classification task. Real fine-tuning uses thousands of examples and typically runs for 3-4 epochs, achieving much stronger validation accuracy.
Multi-class Classification
The same architecture handles multi-class classification by changing the number of output classes:
# Topic classification with 4 categories
topic_model = BertForClassification(num_classes=4)
# Example categories: Sports, Politics, Technology, Entertainment
category_names = ["Sports", "Politics", "Technology", "Entertainment"]
test_texts = [
"The team won the championship game.",
"New legislation passed in Congress.",
"Apple releases new smartphone model.",
"The movie grossed millions at box office.",
]topic_model.eval()
predictions = []
for text in test_texts:
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
logits = topic_model(**inputs)
probs = F.softmax(logits, dim=-1)
pred = torch.argmax(probs, dim=-1).item()
predictions.append((text, category_names[pred]))Multi-class Classification (Untrained) -------------------------------------------------- 'The team won the championship game....' Predicted: Sports 'New legislation passed in Congress....' Predicted: Sports 'Apple releases new smartphone model....' Predicted: Politics 'The movie grossed millions at box office...' Predicted: Sports
The predictions are random because the model is untrained. The architecture remains identical to binary classification; only the output dimension changes from 2 to 4 classes. The softmax over 4 classes produces a probability distribution over topics. After fine-tuning on topic-labeled data, the model would correctly classify each text into its appropriate category.