BERT Fine-tuning › Question Answering Fine-tuning
Language AI HandbookAll chapters BERT Fine-tuning › Question Answering Fine-tuningOne section · Sections in this chapter5 / 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
Extractive question answering finds answer spans within a context passage. Given a question and context, the model predicts which tokens constitute the answer.
Span Prediction Architecture
Question answering presents a fundamentally different challenge from classification or sequence labeling. The answer to a question isn't a single label but a contiguous span of text within the context. For "Where does John work?", the answer "Google" is a substring of the passage. We need to predict both where this substring starts and where it ends.
Consider the alternatives. We could treat this as token classification, labeling each token as "in answer" or "not in answer." But this approach has a flaw: it doesn't guarantee a contiguous span. The model might label disconnected tokens as part of the answer. We could add constraints, but there's a simpler approach.
Instead of classifying each token, we compute two scores: how likely each position is to be the answer's start, and how likely each position is to be the answer's end. For each token position :
where:
- : the start score for position (higher means more likely to be the answer start)
- : the end score for position (higher means more likely to be the answer end)
- : the learned weight vector for start prediction (768 dimensions)
- : the learned weight vector for end prediction (768 dimensions)
- : the hidden state at position from BERT's final layer
Notice that these are dot products, not full linear layers. Each score is a simple weighted sum of the hidden state dimensions. We use separate weight vectors and because the features that indicate "this is where an answer starts" differ from those indicating "this is where an answer ends."
The predicted answer span is the substring from position to position . We find the token with the highest start score and the token with the highest end score, then extract everything between them (inclusive). During inference, we add a constraint: the end position must not precede the start position.
class BertForQuestionAnswering(nn.Module):
"""BERT with span prediction heads for extractive QA."""
def __init__(self, model_name: str = "bert-base-uncased"):
super().__init__()
self.bert = BertModel.from_pretrained(model_name)
# Separate heads for start and end positions
self.start_classifier = nn.Linear(self.bert.config.hidden_size, 1)
self.end_classifier = nn.Linear(self.bert.config.hidden_size, 1)
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,
)
sequence_output = outputs.last_hidden_state
# Compute start and end logits for each position
start_logits = self.start_classifier(sequence_output).squeeze(-1)
end_logits = self.end_classifier(sequence_output).squeeze(-1)
return start_logits, end_logits
def predict_answer(
self, input_ids, attention_mask, token_type_ids, tokenizer
):
"""Extract the predicted answer span."""
self.eval()
with torch.no_grad():
start_logits, end_logits = self.forward(
input_ids, attention_mask, token_type_ids
)
# Get best start and end positions
start_idx = torch.argmax(start_logits, dim=-1).item()
end_idx = torch.argmax(end_logits, dim=-1).item()
# Ensure end >= start
if end_idx < start_idx:
end_idx = start_idx
# Extract answer tokens
answer_tokens = input_ids[0, start_idx : end_idx + 1]
answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
return answer, start_idx, end_idxqa_model = BertForQuestionAnswering()
# Format: [CLS] question [SEP] context [SEP]
question = "Where does John work?"
context = "John Smith is a software engineer at Google in California."
# Tokenize with both segments
encoding = tokenizer(
question,
context,
return_tensors="pt",
padding=True,
truncation=True,
max_length=256,
)qa_model.eval()
with torch.no_grad():
start_logits, end_logits = qa_model(**encoding)
tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"][0])
# Find predicted span
start_probs = F.softmax(start_logits, dim=-1)
end_probs = F.softmax(end_logits, dim=-1)
start_idx = torch.argmax(start_probs, dim=-1).item()
end_idx = torch.argmax(end_probs, dim=-1).item()Question Answering -------------------------------------------------- Question: Where does John work? Context: John Smith is a software engineer at Google in California. Predicted start position: 17 (token: '.') Predicted end position: 7 (token: 'john')
The predicted positions are random because the span prediction heads haven't been trained. In a fine-tuned model, the start position would point to "Google" and the end position would also point to "Google" (or to "California" if the question asked about location). The model learns to identify answer boundaries by training on question-context-answer triplets from datasets like SQuAD.
The SQuAD Format
The Stanford Question Answering Dataset (SQuAD) is the standard benchmark for extractive QA. Each example contains a context paragraph, a question, and the answer's start position in the context.
def prepare_qa_features(
question, context, answer_text, answer_start, tokenizer, max_length=384
):
"""
Prepare features for QA training.
Returns input tensors and the token positions of the answer.
"""
# Tokenize question and context together
encoding = tokenizer(
question,
context,
max_length=max_length,
truncation="only_second", # Only truncate context, not question
padding="max_length",
return_offsets_mapping=True, # Map tokens to character positions
return_tensors="pt",
)
# Find answer token positions
offset_mapping = encoding.pop("offset_mapping")[0]
# Character positions of answer
answer_end = answer_start + len(answer_text)
# Find token positions
start_token = None
end_token = None
for idx, (start, end) in enumerate(offset_mapping):
# Skip special tokens and question tokens (segment 0)
if encoding["token_type_ids"][0, idx] == 0:
continue
if start <= answer_start < end:
start_token = idx
if start < answer_end <= end:
end_token = idx
break
return {
"input_ids": encoding["input_ids"],
"attention_mask": encoding["attention_mask"],
"token_type_ids": encoding["token_type_ids"],
"start_position": start_token,
"end_position": end_token,
}from transformers import BertTokenizerFast
# Use fast tokenizer for offset mapping support
qa_tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
# Example from SQuAD-style data
qa_question = "What color is the sky?"
qa_context = (
"The sky is blue during the day. At sunset, it turns orange and red."
)
answer_text = "blue"
answer_start = qa_context.index(answer_text)
features = prepare_qa_features(
qa_question, qa_context, answer_text, answer_start, qa_tokenizer
)
# Verify the answer extraction
qa_tokens = qa_tokenizer.convert_ids_to_tokens(features["input_ids"][0])
if features["start_position"] and features["end_position"]:
answer_tokens = qa_tokens[
features["start_position"] : features["end_position"] + 1
]QA Feature Preparation ---------------------------------------- Question: What color is the sky? Context: The sky is blue during the day. At sunset, it turns orange and red. Answer: 'blue' (starts at char 11) Answer token positions: start=11, end=11 Extracted tokens: ['blue']
The answer "blue" was correctly located at token positions within the context. During training, these start and end positions serve as labels: the model learns to maximize the probability of the correct positions. The key challenge is aligning character-level answer positions to token-level positions. The offset_mapping from the tokenizer provides this alignment, mapping each token to its character span in the original text.
Unanswerable Questions
SQuAD 2.0 introduced unanswerable questions: the answer might not exist in the context. To handle this, the model can predict the [CLS] position as both start and end, indicating "no answer."
