Named Entity Recognition › BERT-Based NER
Language AI HandbookAll chapters Named Entity Recognition › BERT-Based NEROne section · Sections in this chapter9 / 17 ↓
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
Modern NER is dominated by transformer-based approaches. Let's understand exactly how BERT is adapted for NER and implement a complete fine-tuning pipeline.
How Tokenization Interacts with NER
One subtlety in applying BERT to NER is that BERT uses subword tokenization (WordPiece), which can split words into multiple tokens. The word "Nadella" might become ["Na", "##della"]. When BERT assigns a label to each subword, we need a strategy for aggregating or ignoring subword predictions.
The standard approach is to label only the first subword of each original word and ignore the labels for continuation subwords during evaluation. This keeps the subword details inside the model while maintaining word-level predictions at output time.
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
# Example with entity that gets split into subwords
example_words = ["Nadella", "leads", "Microsoft", "in", "Redmond", "."]
example_labels = ["B-PER", "O", "B-ORG", "O", "B-LOC", "O"]
# Tokenize word by word, tracking alignment
encoding = tokenizer(
example_words,
is_split_into_words=True,
return_offsets_mapping=True,
return_tensors=None,
)
subword_tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"])
word_ids = encoding.word_ids()Subword alignment for NER: Subword Token Word ID Label (first subword only) ------------------------------------------------------- [CLS] None [special] Na 0 B-PER ##dell 0 I-PER ##a 0 I-PER leads 1 O Microsoft 2 B-ORG in 3 O Red 4 B-LOC ##mond 4 I-LOC . 5 O [SEP] None [special]
The special tokens ([CLS] and [SEP]) get word_id None and are excluded from NER evaluation. For continuation subwords (marked with ## in WordPiece), the standard practice during training is to assign them a special ignore index so they contribute no gradient.
Fine-Tuning Architecture
The BERT-for-NER architecture is straightforward. The model takes a tokenized sentence, runs it through BERT's transformer layers to produce a contextualized embedding for each token, and then passes those embeddings through a linear layer that maps to the number of possible BIO tags.
import torch.nn as nn
from transformers import AutoModel
class BertForNER(nn.Module):
"""Minimal BERT-based NER model with token classification head."""
def __init__(self, model_name, num_labels):
super().__init__()
self.bert = AutoModel.from_pretrained(model_name)
hidden_size = self.bert.config.hidden_size
self.dropout = nn.Dropout(0.1)
self.classifier = nn.Linear(hidden_size, num_labels)
def forward(self, input_ids, attention_mask, token_type_ids=None):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
)
# outputs.last_hidden_state: (batch, seq_len, hidden_size)
sequence_output = self.dropout(outputs.last_hidden_state)
logits = self.classifier(sequence_output)
# logits: (batch, seq_len, num_labels)
return logits
# Entity labels for a CoNLL-style 4-class tagset
label_list = [
"O",
"B-PER",
"I-PER",
"B-ORG",
"I-ORG",
"B-LOC",
"I-LOC",
"B-MISC",
"I-MISC",
]
num_labels = len(label_list)
label2id = {label: i for i, label in enumerate(label_list)}
id2label = {i: label for i, label in enumerate(label_list)}NER label set (9 labels): 0: O 1: B-PER 2: I-PER 3: B-ORG 4: I-ORG 5: B-LOC 6: I-LOC 7: B-MISC 8: I-MISC
The architecture is simple, but its power comes from the BERT encoder. The 12 transformer layers in BERT-base have already learned, through pre-training on billions of words, to produce representations that encode:
- Co-reference context: whether "he" or "she" in context refers to a named person
- Syntactic structure: whether a token is the head of a noun phrase or a modifier
- World knowledge: associations between particular tokens and entity types
These learned representations mean the linear classification head only needs to learn a relatively simple mapping from representation to label, rather than learning everything from scratch.
Training Loop
In practice, fine-tuning BERT for NER uses the cross-entropy loss on the first subword of each original word, ignoring special tokens and continuation subwords:
def compute_ner_loss(logits, labels, attention_mask, ignore_index=-100):
"""
Compute cross-entropy loss for NER.
logits: (batch, seq_len, num_labels)
labels: (batch, seq_len) -- use ignore_index for subwords/specials
attention_mask: (batch, seq_len)
"""
loss_fn = nn.CrossEntropyLoss(ignore_index=ignore_index)
# Flatten batch and sequence dimensions
active_mask = attention_mask.view(-1) == 1
active_logits = logits.view(-1, logits.size(-1))[active_mask]
active_labels = labels.view(-1)[active_mask]
loss = loss_fn(active_logits, active_labels)
return loss
# Demonstrate with a tiny synthetic batch
batch_size, seq_len = 2, 8
logits_demo = torch.randn(batch_size, seq_len, num_labels)
labels_demo = torch.tensor(
[
[
label2id["B-PER"],
label2id["I-PER"],
label2id["O"],
label2id["B-ORG"],
label2id["O"],
label2id["B-LOC"],
label2id["O"],
-100,
],
[
label2id["O"],
label2id["B-ORG"],
label2id["I-ORG"],
label2id["O"],
-100,
-100,
-100,
-100,
],
]
)
mask_demo = (labels_demo != -100).long()
loss = compute_ner_loss(logits_demo, labels_demo, mask_demo)Demo NER loss: 2.5542 Active tokens in batch: 11 / 16 Labels shape: torch.Size([2, 8]) Logits shape: torch.Size([2, 8, 9])
The loss ignores the padding positions (where labels are -100) and computes cross-entropy only on the tokens where we have meaningful labels. After a few thousand gradient steps on annotated NER data, the linear classifier learns to combine BERT's contextual representations into accurate entity type predictions.