BERT Fine-tuning › Sequence Labeling Fine-tuning
Language AI HandbookAll chapters BERT Fine-tuning › Sequence Labeling Fine-tuningOne section · Sections in this chapter4 / 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
Named Entity Recognition (NER), part-of-speech tagging, and similar tasks require a prediction for each token rather than a single sequence-level prediction. We classify every token position instead of using only the [CLS] token.
Token Classification Architecture
Classification uses a single representation, the [CLS] token, to make one prediction for the entire sequence. But what if we need a prediction for every token? In Named Entity Recognition, each word gets a label: "John" is a person, "Google" is an organization, "works" is outside any entity. The model must make as many predictions as there are tokens.
The solution extends naturally from classification. Instead of applying our linear layer only to [CLS], we apply the same linear transformation to every token's representation. For each position in the sequence:
where:
- : the logits for position with dimension (number of labels)
- : the hidden state at position from BERT's final layer
- : the shared weight matrix of shape applied to all positions
- : the shared bias vector
An important design choice: we use the same and for every position. This weight sharing makes sense because the meaning of labels doesn't change across positions: a "person" label means the same thing whether we're classifying the first token or the tenth. Sharing weights also keeps the parameter count manageable; without sharing, we'd need separate parameters for each possible position.
The output is a matrix of logits with shape (sequence_length, num_labels). We apply softmax independently to each row to get per-token probability distributions, then take the argmax to get the predicted label for each position:
class BertForTokenClassification(nn.Module):
"""BERT with a token classification head for sequence labeling."""
def __init__(
self,
num_labels: 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_labels)
self.num_labels = num_labels
def forward(self, input_ids, attention_mask=None, token_type_ids=None):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
)
# Get all token representations
sequence_output = outputs.last_hidden_state # (batch, seq_len, hidden)
# Apply dropout and classify each token
sequence_output = self.dropout(sequence_output)
logits = self.classifier(
sequence_output
) # (batch, seq_len, num_labels)
return logits# NER with IOB tagging scheme
# Labels: O, B-PER, I-PER, B-ORG, I-ORG, B-LOC, I-LOC
ner_labels = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
ner_model = BertForTokenClassification(num_labels=len(ner_labels))
# Example sentence
text = "John Smith works at Google in California."
inputs = tokenizer(text, return_tensors="pt")ner_model.eval()
with torch.no_grad():
logits = ner_model(**inputs)
preds = torch.argmax(logits, dim=-1)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
token_predictions = [
(token, ner_labels[pred.item()]) for token, pred in zip(tokens, preds[0])
]Token Classification (NER) ---------------------------------------- Token Predicted Label ---------------------------------------- [CLS] B-ORG john I-ORG smith I-PER works B-ORG at B-PER google B-PER in B-PER california B-PER . B-ORG [SEP] B-PER
The predicted labels are random because the classifier weights haven't been trained. In a trained NER model, "John" and "Smith" would receive B-PER and I-PER tags, "Google" would receive B-ORG, and "California" would receive B-LOC. The O (Outside) tag would correctly mark non-entity tokens like "works" and "at".
Handling WordPiece Tokenization
A subtle challenge arises from WordPiece tokenization. When a word splits into multiple subwords, which subword's prediction should we use?
def align_predictions_to_words(tokens, predictions, tokenizer):
"""
Align subword predictions back to original words.
Uses the first subword's prediction for each word.
"""
word_predictions = []
current_word = ""
current_pred = None
for token, pred in zip(tokens, predictions):
# Skip special tokens
if token in ["[CLS]", "[SEP]", "[PAD]"]:
continue
if token.startswith("##"):
# Continuation of previous word
current_word += token[2:]
else:
# New word: save previous word's prediction
if current_word:
word_predictions.append((current_word, current_pred))
current_word = token
current_pred = pred
# Don't forget the last word
if current_word:
word_predictions.append((current_word, current_pred))
return word_predictions# Example with multi-token words
text_subwords = "The unbelievable transformation surprised everyone."
inputs_sw = tokenizer(text_subwords, return_tensors="pt")
tokens_sw = tokenizer.convert_ids_to_tokens(inputs_sw["input_ids"][0])ner_model.eval()
with torch.no_grad():
logits_sw = ner_model(**inputs_sw)
preds_sw = torch.argmax(logits_sw, dim=-1)[0]
word_preds = align_predictions_to_words(
tokens_sw, [ner_labels[p.item()] for p in preds_sw], tokenizer
)WordPiece Tokenization: Original: 'The unbelievable transformation surprised everyone.' Tokens: ['[CLS]', 'the', 'unbelievable', 'transformation', 'surprised', 'everyone', '.', '[SEP]'] Aligned Predictions: the: B-PER unbelievable: B-PER transformation: B-ORG surprised: B-PER everyone: O .: B-PER
The word "unbelievable" tokenizes into multiple subwords (likely "un", "##believ", "##able"). The alignment function takes the first subword's prediction as the label for the entire word. This is necessary because NER labels apply to words, not subwords. Alternative strategies include averaging logits across subwords or using the last subword's prediction.
The IOB Tagging Scheme
NER commonly uses Inside-Outside-Beginning (IOB) tagging. "B-" marks the beginning of an entity, "I-" continues it, and "O" marks non-entity tokens.
The IOB scheme allows multi-word entities: "John Smith" spans two tokens with B-PER and I-PER, while "New York" spans B-LOC and I-LOC. This is why sequence labeling is more complex than simple classification: the model must learn to produce coherent tag sequences.