Lesson 107: Extractive QA — SQuAD, BERT, and Open-Domain Retrieval

Comparison matrix

From SQuAD: from reading comprehension to token indices: refill the example column from what you know. The rest of the table is as it appeared.

Comparison matrix

From BERT input for SQuAD: the concatenated sequence: every row here is a choice with a cost. Fill the segment ID column, then say which row you would actually pick and what you give up for it.

Discussion prompt

BERT produces a hidden vector h_i ∈ ℝ^{d} for every input token i. Two separate nn.Linear(d, 1) layers score each position as a candidate start and a candidate end.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Predict first

Trace a batch of B=2 sequences, each T=16 tokens, through BERT hidden size d=64. Predict tensor shapes at each step before reading on.

Commit before you compute: what does Forward pass: shapes and loss — verified come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: start_logits and end_logits: both torch.Size([2, 16])

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Linear(64,1) applied to each token's 64-d vector → scalar per token → squeeze the trailing dim → (B=2, T=16).

Discussion prompt

Work backwards. The example finished here:

start_logits and end_logits: both torch.Size([2, 16])

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Predict first

A student writes this, and it looks reasonable:

Decode naively: start = argmax(start_logits), end = argmax(end_logits) — return the span (start, end) regardless of order.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: This is a common failure mode: the two heads are trained independently and can produce an inverted span (end < start), which is syntactically impossible.

Enumerate all valid (i, j) pairs where i ≤ j, compute start_logits[i] + end_logits[j], pick the argmax.

Why: This is a common failure mode: the two heads are trained independently and can produce an inverted span (end < start), which is syntactically impossible.

Predict first

Predict the EM and F1 for each pair before reading the table. Remember: normalize first (lowercase, split on whitespace), then compute bag-of-word overlap.

Commit before you compute: what does EM and F1 — four concrete examples come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: "albert einstein" vs "einstein": EM=0, P=0.500, R=1.000, F1=0.667

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Common = {'einstein'} (1 token).

Discussion prompt

Work backwards. The example finished here:

"albert einstein" vs "einstein": EM=0, P=0.500, R=1.000, F1=0.667

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Predict first

A student writes this, and it looks reasonable:

Compute F1 by counting characters in common between predicted and gold strings.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: Character counting gives an inflated or deflated number depending on string length, and ignores token boundaries entirely.

Normalize both strings (lowercase, strip articles/punct), split on whitespace, take bag-of-word intersection, compute token-level precision and recall.

Why: Character counting gives an inflated or deflated number depending on string length, and ignores token boundaries entirely. 'states' shares characters with 'states' but 'the' contributes characters not in gold.

Discussion prompt

The rule this trap just fixed:

P = 2/3 = 0.667 (2 of 3 pred tokens are correct), R = 2/4 = 0.500 (2 of 4 gold tokens recovered). F1 = 2·0.667·0.5/(0.667+0.5) = 0.571. This matches the official SQuAD script.

Now break it on purpose. Build a case that violates it and follow the consequences until something visibly fails. Where does the failure first show up — and would you have noticed it if you had not been looking?

Hint: The dangerous rules are the ones whose violation still produces an answer. If yours fails loudly, try to find one that fails quietly.

Predict first

Five candidate passage embeddings, one query embedding — all unit-normalized (seed=7). Which passage ranks first? Compute the dot products mentally before revealing.

Commit before you compute: what does DPR retrieval — dot-product ranking verified come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Passage 3 ranks first with score 0.2467

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Dot product of two unit vectors equals cosine similarity.

Comparison matrix

From DPR retrieval — dot-product ranking verified: refill the passage column from what you know. The rest of the table is as it appeared.

Discussion prompt

For USAAIO: extractive is architecturally cleaner and the dominant SQuAD paradigm. Generative models (FiD, RAG) are used for open-ended questions where the answer is not a verbatim passage substring.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Discussion prompt

Run The QA system playbook with this step confiscated:

Train: loss = (cross_entropy(start_logits, gold_start) + cross_entropy(end_logits, gold_end)) / 2.

Is it still possible? If it is, say what takes its place and what it costs you. If it is not, say exactly what that step was providing that nothing else does.

Hint: A step you can drop for free was never load-bearing. If you cannot drop it, name the thing that goes wrong the moment it is gone.

Discussion prompt

The QA system playbook works on the cases you have just seen. Push it to the edge: what is the most degenerate input it still handles — empty, zero, one item, everything equal — and what is the first case where it stops being true? Name the case, not just "it breaks".

Hint: Try the smallest legal input, then the largest, then the one where two things collide. Methods are specified at their edges; the middle takes care of itself.

Eliminate the wrong options

BERT hidden size d=768, two Linear(768, 1) heads. How many new parameters do the QA heads add?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. 1,538 (768×2 weights + 2 biases)
• B. 1,536 (768×2 weights only, no biases)
• C. 768 (one shared linear layer)
• D. 589,824 (768×768 per head)

Survives elimination: A

Why: Each Linear(768, 1) has 768 weight parameters + 1 bias = 769 parameters. Two independent heads: 769 × 2 = 1,538. This is why fine-tuning BERT for QA is so lightweight — the pretrained encoder (110 M params) does the heavy lifting; the QA head adds 0.001%.

Check your understanding

BERT hidden size d=768, two Linear(768, 1) heads. How many new parameters do the QA heads add?

• A. 1,538 (768×2 weights + 2 biases) (correct)
• B. 1,536 (768×2 weights only, no biases)
• C. 768 (one shared linear layer)
• D. 589,824 (768×768 per head)

Answer: A

Why: Each Linear(768, 1) has 768 weight parameters + 1 bias = 769 parameters. Two independent heads: 769 × 2 = 1,538. This is why fine-tuning BERT for QA is so lightweight — the pretrained encoder (110 M params) does the heavy lifting; the QA head adds 0.001%.

Why B tempts people

Forgets the bias term; PyTorch nn.Linear has bias=True by default, adding 1 scalar per output unit.

Why C tempts people

BERT-SQuAD uses two separate heads — one predicts start, one predicts end — not a single shared layer. Sharing would conflate the two tasks.

Why D tempts people

Confuses a square weight matrix (768×768) with the correct (768×1) output projection; the QA head projects to a single logit, not to d dimensions.

Predict first

pred = 'New York City', gold = 'New York'. What are EM and F1 (token-level)?

Answer it in your own words, now, with nothing to choose from. The options are on the next slide — and picking the right one off a list is an easier skill than producing it.

Correct: EM = 0, F1 = 0.800

Why: Normalized pred tokens = {new, york, city} (3 tokens). Normalized gold tokens = {new, york} (2 tokens). Common = {new, york} (2 tokens). P = 2/3 ≈ 0.667, R = 2/2 = 1.000. F1 = 2·(2/3)·1 / (2/3+1) = (4/3)/(5/3) = 4/5 = 0.800. EM = 0 because the token lists are not identical.

Check your understanding

pred = 'New York City', gold = 'New York'. What are EM and F1 (token-level)?

• A. EM = 0, F1 = 0.800 (correct)
• B. EM = 0, F1 = 0.667
• C. EM = 1, F1 = 1.000
• D. EM = 0, F1 = 1.000

Answer: A

Why: Normalized pred tokens = {new, york, city} (3 tokens). Normalized gold tokens = {new, york} (2 tokens). Common = {new, york} (2 tokens). P = 2/3 ≈ 0.667, R = 2/2 = 1.000. F1 = 2·(2/3)·1 / (2/3+1) = (4/3)/(5/3) = 4/5 = 0.800. EM = 0 because the token lists are not identical.

Why B tempts people

Computes P = 2/3 correctly but uses F1 = 2·P·R/(P+R) = 2·(2/3)·(2/3)/((2/3)+(2/3)) = 0.667, incorrectly taking R = 2/3 instead of R = 2/2 = 1. Recall is common/gold, not common/pred.

Why C tempts people

Incorrectly concludes EM=1 because 'new york' is a substring of 'new york city'; EM requires exact string equality after normalization, not substring match.

Why D tempts people

Sets F1=1 because recall R=1 (both gold tokens are in the prediction), but ignores that precision < 1 since 'city' is a spurious extra token; harmonic mean of P<1 and R=1 is always <1.

Eliminate the wrong options

BERT-SQuAD is trained with loss = (CE(start_logits, gold_start) + CE(end_logits, gold_end)) / 2. What does this training signal teach the model to do, and why does the independence assumption make this factorization valid?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. It teaches the model to maximize P(start=i) and P(end=j) independently, which is valid because both heads share BERT's context vectors that jointly encode passage semantics
• B. It teaches the model to predict the joint span distribution P(start=i, end=j) exactly, which requires a T×T output matrix
• C. It teaches start position only; the end head is supervised by the start head's output, making them dependent
• D. It maximizes the log-probability of the correct character offsets, not token indices, so the loss is computed in character space

Survives elimination: A

Why: CE(start_logits, gold_start) trains the start head to assign maximum probability to the correct start token index — i.e., maximize log P_s(gold_start). Similarly for end. The two heads share the same BERT encoder output h_i, so the encoder learns a single representation that makes both position predictions accurate. The implicit coupling through shared h_i is why factoring the span probability as P_s · P_e is a good approximation despite the independence assumption.

Check your understanding

BERT-SQuAD is trained with loss = (CE(start_logits, gold_start) + CE(end_logits, gold_end)) / 2. What does this training signal teach the model to do, and why does the independence assumption make this factorization valid?

• A. It teaches the model to maximize P(start=i) and P(end=j) independently, which is valid because both heads share BERT's context vectors that jointly encode passage semantics (correct)
• B. It teaches the model to predict the joint span distribution P(start=i, end=j) exactly, which requires a T×T output matrix
• C. It teaches start position only; the end head is supervised by the start head's output, making them dependent
• D. It maximizes the log-probability of the correct character offsets, not token indices, so the loss is computed in character space

Answer: A

Why: CE(start_logits, gold_start) trains the start head to assign maximum probability to the correct start token index — i.e., maximize log P_s(gold_start). Similarly for end. The two heads share the same BERT encoder output h_i, so the encoder learns a single representation that makes both position predictions accurate. The implicit coupling through shared h_i is why factoring the span probability as P_s · P_e is a good approximation despite the independence assumption.

Why B tempts people

A joint distribution over all (start, end) pairs would require a T² output and a joint softmax — far more parameters and intractable for T=512. The factored formulation with two T-way softmaxes is the efficiency trick that makes BERT-SQuAD practical.

Why C tempts people

The end head is supervised directly by the gold end label, not by the start head's output. The two heads are trained in parallel with independent losses; there is no sequential dependency at training time.

Why D tempts people

SQuAD labels are stored as character offsets in the dataset, but they are converted to token indices before training. The cross-entropy loss operates over the discrete token vocabulary of positions 0…T-1, not in character space.

Sort into buckets

Sort these by output shape, from memory, without looking back at Your turn — Milestone 1: model definition and…. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

(B, T, 64)

tok_embed Embedding(500,64); pos_embed Embedding(32,64); seg_embed Embedding(2,64); MHSA (per Block): qkv+out; FFN (per Block): 64→128→64

—

LayerNorm (2 per block + 1 final); TOTAL

(B, T, 2)

QA head Linear(64,2)

g1

output shape is "(B, T, 64)" for tok_embed Embedding(500,64), pos_embed Embedding(32,64), seg_embed Embedding(2,64), MHSA (per Block): qkv+out, FFN (per Block): 64→128→64 — that is what the table on "Your turn — Milestone 1: model…" records, and it is the single property separating this group from the rest.

g2

output shape is "—" for LayerNorm (2 per block + 1 final), TOTAL — that is what the table on "Your turn — Milestone 1: model…" records, and it is the single property separating this group from the rest.

g3

output shape is "(B, T, 2)" for QA head Linear(64,2) — that is what the table on "Your turn — Milestone 1: model…" records, and it is the single property separating this group from the rest.

Predict first

The table runs: 0 | ~3.0 | uniform over 20 classes: log(20) ≈ 3.0 · 50 | ~2.0 | model narrowing to correct region · 100 | ~1.0 | substantially learned start/end offsets · 150 | ~0.5 | near-convergence on synthetic task

In Your turn — Milestone 2: training loop with synthetic data, given the rows so far: what is the next one — the row where step is 200?

Correct: 200 | ~0.2 | overfitting to synthetic random seeds

Why: The relationship between the columns, not the individual numbers, is what generates the next row. Using fixed random seed ensures reproducibility.

Comparison matrix

From Your turn — Milestone 2: training loop with synthetic data: every row here is a choice with a cost. Fill the expected loss column, then say which row you would actually pick and what you give up for it.

Fill in the blanks

From Your turn — Milestone 3: EM/F1 evaluation and full program — one line has had its right-hand side removed. Put it back.

def best_valid_span(s_lg, e_lg):
"""Return (start, end) maximizing s[i]+e[j] subject to i<=j."""
T = s_lg.shape[0]
best, bi, bj = float('-inf'), 0, 0
for i in range(T):
for j in range(i, T):
v = s_lg[i].item() + e_lg[j].item()
if v > best:
best, bi, bj = v, i, j
return bi, bj

def token_f1(pred_ids, gold_ids):
ps, gs = set(pred_ids), set(gold_ids)
common = ps & gs
if not common: return 0.0
p = len(common)/len(ps); r = len(common)/len(gs)
return 2pr/(p+r)

model.eval()
total_em, total_f1, n = 0, 0.0, 0
for seed in range(20): # 20 held-out examples
ids, segs, gs, ge = make_batch(B=1, seed=seed+1000)
with torch.no_grad():
s_lg, e_lg = model(ids, segs)
ps, pe = best_valid_span(s_lg[0], e_lg[0])
pred_ids = ids[0, ps:pe+1].tolist()
gold_ids = ids[0, gs[0]:ge[0]+1].tolist()
em = int(pred_ids == gold_ids)
f1 = token_f1(pred_ids, gold_ids)
total_em += em; total_f1 += f1; n += 1
print(f'EM=___ F1=___ (n=___)')

Why: pred_ids is what everything below it consumes, so the wrong expression here fails later and somewhere else. In a real system you would map token indices to character offsets via the tokenizer's offset_mapping.

Comparison matrix

From Your turn — Milestone 3: EM/F1 evaluation and full program: refill the what to check column from what you know. The rest of the table is as it appeared.

Draw it

One page, no notation unless you need it: draw how these connect — SQuAD — the span-prediction task · Two heads on BERT — start & end prediction · EM and F1 — evaluating span predictions · Open-domain QA — retrieve then read. Put an arrow wherever one of them is what makes another possible, and label the arrow with why.