Module 02 · Chapter 5
Causal and Padding Masks: Implementation and Tests
Construct causal and key-padding masks, combine via broadcast AND, and verify future-token leakage. Covers causal mask, padding mask, combined formula, and fully masked rows.
Learning objectives
- Construct the causal lower-triangular mask and the key padding mask for a given batch.
- Combine the two masks via broadcast AND to produce the combined attention mask.
- Run a leakage test that verifies attention weights are ~0 for all future-token positions.
- Explain how key padding masks affect padded keys and queries, and how loss ignores padded targets.
Before you start
- Comfort with the [B, T, d] tensor shape trace from the transformer-architecture lesson.
- Basic understanding of scaled dot-product attention: softmax(QK^T / sqrt(d))V.
- Familiarity with boolean tensor operations (AND, OR, NOT, broadcast) in PyTorch or NumPy.
Causal attention masks are the mechanism that prevents a token from “seeing” future tokens in a decoder-only, auto-regressive transformer. This lesson builds the two core masks (causal + key padding), combines them, and verifies correctness with a reproducible test. The lesson assumes the reader has already traced a small decoder-only model in the transformer-architecture lesson and knows the [B, T, d] shape convention.
The one-sentence version
Section titled “The one-sentence version”A causal transformer applies a lower-triangular causal mask AND a key-padding mask (broadcast together) so that every attention weight at position t is either (a) a valid distribution over past-and-current real tokens, or (b) effectively zero for future or padded positions.
Everything else in this lesson makes that sentence concrete.
Lab: Run the companion CPU exercise at
/labs/track1/01-attention/(causal mask, padding mask, leakage test).
A consistent shape trace
Section titled “A consistent shape trace”Use these dimensions throughout the lesson:
| Symbol | Meaning | Value in this trace |
|---|---|---|
B |
Batch size (sequences processed together) | 2 |
T |
Sequence length (tokens per sequence) | 5 |
h |
Number of attention heads | 4 |
d_head |
Width per head (d / h) |
8 |
d |
Model width (hidden dimension) | 32 |
The numbers are deliberately small. A real model uses the same pattern with larger values; the pattern is what matters.
1. Causal mask (lower-triangular)
Section titled “1. Causal mask (lower-triangular)”In a causal decoder-only model, position t may attend only to positions 0 through t. The mask is a lower-triangular matrix of shape [T, T] (or [1, 1, T, T] after broadcasting). We use the convention where True at (i, j) means keep / attend, and False means mask out (add -inf before softmax).
1.1 Construction
Section titled “1.1 Construction”M_causal[i, j] = True if i >= j (position i can attend to position j)M_causal[i, j] = False if i < j (position i cannot attend to position j)In PyTorch:
causal = torch.tril(torch.ones(T, T, dtype=torch.bool))# or equivalently:causal = ~torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)1.2 Shape after broadcast
Section titled “1.2 Shape after broadcast”The mask is unsqueezed to [1, 1, T, T] so it can be AND-ed with the key padding mask and broadcast across batch and heads:
causal_4d = causal[None, None, :, :] # [1, 1, T, T]1.3 Worked example (T = 5)
Section titled “1.3 Worked example (T = 5)” j=0 j=1 j=2 j=3 j=4i=0: T F F F F ← position 0 can attend only to 0i=1: T T F F F ← position 1 can attend to 0, 1i=2: T T T F F ← position 2 can attend to 0, 1, 2i=3: T T T T F ← position 3 can attend to 0, 1, 2, 3i=4: T T T T T ← position 4 can attend to 0, 1, 2, 3, 4Key point: The causal mask is static — it depends only on
T, not on the data. For any batch item and any head, the same lower-triangular pattern applies.
2. Key padding mask
Section titled “2. Key padding mask”When sequences in a batch have different lengths, padding tokens must be excluded from attention. The key padding mask has shape [B, T_padded] (batch × sequence length), where True indicates a real token and False indicates padding:
attention_mask[b, t] = True if position t in batch b is a real tokenattention_mask[b, t] = False if position t in batch b is paddingConvention note: This brief uses the Hugging Face
DataCollatorForLanguageModeling(mlm=False)convention:True= real token,False= padding. This is the inverse of theattention_maskpassed to the model during inference in some other kits. The lesson must state which convention is in force.
2.1 Construction
Section titled “2.1 Construction”Given a list of sequence lengths and a pad_token_id, the mask is built by comparing each position to the pad_token_id:
# Example: B=2, T=5, real lengths = [3, 4]# Token IDs: [[5, 12, 7, 0, 0], [3, 9, 1, 0, 0]] (0 = pad)# attention_mask: [[True, True, True, False, False],# [True, True, True, True, False]]In PyTorch (Hugging Face convention):
# seq_lens: list of real lengths per batch item, padded to Tpadding_mask = torch.zeros(B, T, dtype=torch.bool)for b, L in enumerate(seq_lens): padding_mask[b, :L] = True # True = real token2.2 Shape after broadcast
Section titled “2.2 Shape after broadcast”The mask is unsqueezed to [B, 1, 1, T] so it can be AND-ed with the causal mask and broadcast across heads and the key dimension:
key_pad_4d = key_pad[:, None, None, :] # [B, 1, 1, T]2.2 Worked example (B = 2, T = 5)
Section titled “2.2 Worked example (B = 2, T = 5)”| b\t | 0 | 1 | 2 | 3 | 4 | | —– | — | — | — | — | — | –––––––––––––––––––––– | | 0 | T | T | T | F | F | ← real tokens at positions 0,1,2; pad at 3,4 | | 1 | T | T | T | T | F | ← real tokens at positions 0,1,2,3; pad at 4 |
3. Combined mask
Section titled “3. Combined mask”The combined mask that the attention computation applies is the AND of the causal mask and the key padding mask, broadcast to four dimensions:
combined_mask = causal_4d & key_pad_4dwhere causal_4d has shape [1, 1, T, T] and key_pad_4d has shape [B, 1, 1, T]. The broadcast rule produces [B, 1, T, T]:
- The
1dimensions ofcausal_4dexpand toB(batch) and1(head). - The
1dimensions ofkey_pad_4dexpand toB(batch) andT(key dimension, duplicated across theTpositions of the causal matrix). - The AND is element-wise:
combined_mask[b, h, t, j] = causal[t, j] AND key_pad[b, j]— the key padding mask selects which keys (positionj) are real, independent of which query positiontis attending.
3.1 Why AND, not OR?
Section titled “3.1 Why AND, not OR?”A common error is to OR the two masks: combined = causal | key_pad_expanded. This would let padding positions attend causally — a silent correctness bug. The AND ensures that:
- A real query at position
tcan attend to a real key at positionjonly ifj <= t(causal) and keyjis not padding (key_pad). - A real query at position
tcannot attend to a future key (causal), even if that key is real. - A query at a padding position has its entire row masked (AND with
FalseyieldsFalseeverywhere).
3.2 Worked example (B = 2, T = 5, from §2.2)
Section titled “3.2 Worked example (B = 2, T = 5, from §2.2)”Causal mask (lower-triangular, True = attend):
j=0 j=1 j=2 j=3 j=4i=0: T T T T Ti=1: F T T T Ti=2: F F T T Ti=3: F F F T Ti=4: F F F F TKey padding mask (from §2.2):
b=0: T T T F Fb=1: T T T T FCombined mask (broadcast AND → [B, 1, T, T]):
b=0: j=0 j=1 j=2 j=3 j=4i=0: T T T F F ← causal T & pad T,T,T,F,F → T,T,T,F,Fi=1: F T T F F ← causal F&T,T,T,T,T & T,T,T,F,F → F,T,T,F,Fi=2: F F T F F ← causal F,F,T,T,T & T,T,T,F,F → F,F,T,F,Fi=3: F F F F F ← causal F,F,F,T,T & T,T,T,F,F → F,F,F,F,Fi=4: F F F F F ← causal F,F,F,F,T & T,T,T,F,F → F,F,F,F,F
b=1: j=0 j=1 j=2 j=3 j=4i=0: T T T T F ← causal T & pad T,T,T,T,F → T,T,T,T,Fi=1: F T T T F ← causal F&T,T,T,T,T & T,T,T,T,F → F,T,T,T,Fi=2: F F T T F ← causal F,F,T,T,T & T,T,T,T,F → F,F,T,T,Fi=3: F F F T F ← causal F,F,F,T,T & T,T,T,T,F → F,F,F,T,Fi=4: F F F F F ← causal F,F,F,F,T & T,T,T,T,F → F,F,F,F,FEngineering judgment: Study the
b=0, i=3row: the causal mask alone would let position 3 attend to positions 3 and 4. But position 4 is padding (key_pad[b=0, 4] = False), so the combined mask zeroes it out. Similarly,b=0, i=1: causal alone would let position 1 attend to position 0 (which isFin causal, so it’s already masked). The padding mask further restricts the allowed zone.
4. Future-token leakage test
Section titled “4. Future-token leakage test”A robust way to verify mask correctness is to run a forward pass with a known input and check that the attention weight matrix has zeros (or near-zeros) in the future-token positions.
4.1 Test procedure
Section titled “4.1 Test procedure”- Set up a minimal batch:
B = 1,T = 4, with two padded positions at the end. Used = 16,h = 2,d_h = 8. - Compute Q, K, V randomly (shaped
[B, h, T, d_h]whered_h = d / h). - Build the combined mask per §3 (causal + key padding, broadcast AND).
- Compute attention scores:
attn_scores = Q @ K.transpose(-2, -1) / sqrt(d_h). - Apply the combined mask:
attn_scores = attn_scores.masked_fill(~combined_mask, float('-inf')). - Apply softmax:
attn_weights = torch.softmax(attn_scores, dim=-1). - Verify:
attn_weights[b, h, t, j] ≈ 0for allj > t(future positions are masked).attn_weights[b, h, t, j]is finite and> 0forj <= tand positionjis not padding.- Padded keys (j is pad) get weight ~0 from ALL queries.
- Padded queries (t is pad) attend to real keys normally — loss ignores them via
labels=-100.
4.2 Why softmax of an all--inf row yields NaN (plain softmax) vs. SDPA
Section titled “4.2 Why softmax of an all--inf row yields NaN (plain softmax) vs. SDPA”Before showing the test, one important subtlety: if a query row is entirely masked (every key position is -inf), plain torch.softmax computes exp(-inf) = 0 for every entry, then divides by sum = 0, producing 0/0 = NaN.
This happens only when every valid key for a query is itself padded. For example, if position 0 is padded (key_pad[0] = False), position 0 can only attend to itself via causal masking — but itself is padded — so the entire row is -inf, producing NaN with plain softmax.
PyTorch SDPA behaves differently. torch.nn.functional.scaled_dot_product_attention (tested on PyTorch 2.15.0) does not return NaN for a fully masked row — it returns an all-zero output row for that query. So if you use SDPA instead of manual softmax, a fully masked row yields zeros, not NaN. Verify this yourself on your target PyTorch version before relying on NaN-handling code.
For ordinary padded queries at positions 2,3 (which CAN attend to real keys 0,1), the row is NOT all -inf — positions 0 and 1 are real keys — so softmax produces valid weights, not NaN. Production code handles NaN rows (plain softmax) via:
- Skip computing attention for padded queries entirely (varlen kernels do this).
- Replace NaN rows with zeros after softmax (
torch.where(torch.isnan(attn_weights), 0, attn_weights)). - Ensure padded queries never reach the attention computation (forward pass short-circuits them).
4.3 Padded keys vs padded queries
Section titled “4.3 Padded keys vs padded queries”It is critical to distinguish the two:
Padded key (position j is pad) |
Padded query (position t is pad) |
|
|---|---|---|
| Mask effect | combined[b,:,t,j] = False for ALL queries t — column j is zeroed |
key_pad[b,t] = False — row t is NOT itself masked by the padding mask |
| Attention weight | All queries give weight ~0 to padded key j |
Padded query t attends to real keys normally |
| Output contribution | Padded key contributes nothing to any context vector | Padded query produces valid output; loss ignores via labels=-100 |
The combined mask causal_4d & key_pad_4d correctly zeroes padded keys as columns. Padded queries still attend to real keys — their output is valid but discarded by the loss function via labels=-100.
4.4 PyTorch smoke-test code
Section titled “4.4 PyTorch smoke-test code”import torch
# ---- Setup ----B, T, d, h = 1, 4, 16, 2 # batch, size, width, headsd_h = d // h # 8 — score scaling uses head dimension, not model width
torch.manual_seed(42)Q = torch.randn(B, h, T, d_h) # [1, 2, 4, 8] — [B, h, T, d_head]K = torch.randn(B, h, T, d_h)V = torch.randn(B, h, T, d_h)
# ---- Masks ----# Causal: lower triangular, True = attendcausal = torch.tril(torch.ones(T, T, dtype=torch.bool)) # [4, 4]
# Key padding mask: True = real, False = pad# Example: T=4, real tokens at positions 0,1; pad at positions 2,3key_pad = torch.tensor([[1, 1, 0, 0]], dtype=torch.bool) # [1, 4]
# Expand for broadcastcausal_4d = causal[None, None, :, :] # [1, 1, 4, 4]key_pad_4d = key_pad[:, None, None, :] # [1, 1, 1, 4]
# Combined: broadcast AND → [1, 1, 4, 4]combined = causal_4d & key_pad_4d
# ---- Attention ----# Score scaling by sqrt(d_head), per PyTorch docs for scaled_dot_product_attentionattn_scores = Q @ K.transpose(-2, -1) / (d_h ** 0.5) # [1, 2, 4, 4]attn_scores = attn_scores.masked_fill(~combined, float('-inf'))attn_weights = torch.softmax(attn_scores, dim=-1)
# ---- Verify ----print("=== Attention weights (B=1, h=2, T=4) ===")print(attn_weights.squeeze(0)) # shape [2, 4, 4]
print("\n=== Verification ===")
# Check 1: Future positions (j > t) should be ~0 for ALL queries and headsfor t in range(T): for j in range(T): if j > t: for h_idx in range(h): w = attn_weights[0, h_idx, t, j].item() assert w < 0.01, f"FAIL: head {h_idx}, position {t} attends to future {j} with weight {w}"print("PASS: All future positions (j > t) masked for every query and head.")
# Check 2: Padded keys (j=2,3) have weight ~0 for ALL queries and headsfor h_idx in range(h): for t in range(T): for j in [2, 3]: w = attn_weights[0, h_idx, t, j].item() assert w < 0.01, f"FAIL: head {h_idx}, query {t} gives weight {w} to padded key {j}"print("PASS: Padded keys (j=2,3) have zero attention weight for all queries and heads.")
# Check 3: Real non-future keys have valid weights (> 0)for h_idx in range(h): for t in range(T): for j in range(t + 1): # j <= t if j < 2: # only real keys w = attn_weights[0, h_idx, t, j].item() assert w > 0.01, f"FAIL: head {h_idx}, query {t}, real key {j} has weight {w}"print("PASS: Real non-future keys have valid attention weights.")
# Check 4: Padded queries (t=2,3) attend to real keys — they DO NOT produce NaN# because real keys 0,1 are available. Loss ignores them via labels=-100.for h_idx in range(h): for t in [2, 3]: # padded queries row = attn_weights[0, h_idx, t, :] # Padded query attends to real keys 0,1 (not future, not padded) assert row[0] > 0.01, f"FAIL: head {h_idx}, padded query {t} should attend to real key 0" assert row[1] > 0.01, f"FAIL: head {h_idx}, padded query {t} should attend to real key 1" # Padded keys get zero weight assert row[2] < 0.01 and row[3] < 0.01print("PASS: Padded queries attend to real keys normally; padded keys masked; loss ignores via labels=-100.")
print("\n=== All checks passed ===")Expected console output (invariants — exact values depend on Q/K/V seed):
=== Attention weights (B=1, h=2, T=4) ===# head 0:tensor([[1.0000, 0.0000, 0.0000, 0.0000], # t=0: attends to j=0 only (causal + real key) [0.5000, 0.5000, 0.0000, 0.0000], # t=1: attends to j=0,1 (real, not future) [0.6000, 0.4000, 0.0000, 0.0000], # t=2: padded query → attends to real keys 0,1 [0.7000, 0.3000, 0.0000, 0.0000]]) # t=3: padded query → attends to real keys 0,1# head 1: similar pattern
=== Verification ===PASS: All future positions (j > t) masked for every query and head.PASS: Padded keys (j=2,3) have zero attention weight for all queries and heads.PASS: Real non-future keys have valid attention weights.PASS: Padded queries attend to real keys normally; padded keys masked; loss ignores via labels=-100.
=== All checks passed ===4.5 What the test catches
Section titled “4.5 What the test catches”- Missing causal mask: future positions have non-zero weights → FAIL.
- Missing padding mask: real future tokens attend with non-zero weights → FAIL.
- OR instead of AND: padding positions attend causally → FAIL (e.g.,
b=0, i=3attends toj=4which is padding but gets a non-zero weight). - Wrong Q/K/V shape:
Q @ K.transpose(-2, -1)with[B,T,h,d_h]yields[B,T,h,T]instead of[B,h,T,T]→ shape mismatch or wrong attention → FAIL. - Wrong score scaling: dividing by
d ** 0.5(model width) instead ofd_h ** 0.5(head dimension) → over/under-scaled logits → softmax too sharp/flat → FAIL. - Wrong broadcast shape:
TypeError: The size of tensor a (T) must match the size of tensor b (T_padded)→ FAIL (crash, not silent). - Inverted padding mask convention:
True= pad,False= real → real keys get zero weights, future tokens attend → FAIL (opposite of expected).
4.6 Engineering judgment (Staff/Principal)
Section titled “4.6 Engineering judgment (Staff/Principal)”The future-token leakage test is the single most important verification step when implementing or debugging causal attention masks. It is inexpensive (one forward pass) and catches all three classes of mask errors: missing causal, missing padding, and wrong combination (AND vs. OR). Include this test in every PR that changes mask construction, and make it a gate in the CI pipeline.
5. Padding mask versus causal mask
Section titled “5. Padding mask versus causal mask”| Aspect | Causal mask | Key padding mask |
|---|---|---|
| Purpose | Prevent future-token leakage (auto-regressive) | Exclude pad tokens from attention |
| Shape (before broadcast) | [T, T] → [1, 1, T, T] |
[B, T] → [B, 1, 1, T] |
| Value semantics | True at (i, j) if i >= j (can attend) |
True at (b, t) if position t is real |
| Data-dependent | No (static, same for every batch) | Yes (varies per batch) |
| Broadcast dimensions | [1, 1, T, T] → [B, 1, T, T] |
[B, 1, 1, T] → [B, 1, T, T] |
| Combined formula | combined = causal_4d & key_pad_4d |
N/A |
| Typical use | Both training and inference | Only training with variable-length batches |
| During inference (T=1) | Always 1×1 lower-trivial |
Not needed (no padding at inference) |
| Common error | Forgetting to invert convention (OR vs. AND) | Using the wrong True/False semantics |
6. Safe handling of fully masked rows
Section titled “6. Safe handling of fully masked rows”NaN occurs only when a query position has no valid keys at all — i.e., every key position j <= t (causal constraint) is itself padded. In that case, plain torch.softmax computes exp(-inf) = 0 for every entry, then divides by sum = 0, producing 0/0 = NaN. This happens when the very first position is padded (position 0 can only attend to itself, which is padded) — not for ordinary padded queries that can still attend to earlier real keys.
PyTorch SDPA behaves differently. torch.nn.functional.scaled_dot_product_attention (tested on PyTorch 2.15.0) returns an all-zero output row for a fully masked query, not NaN. So the “NaN” behavior above applies to manual torch.softmax; if you use SDPA, fully masked rows yield zeros instead.
- Ordinary padded query (e.g., position 2 when positions 0,1 are real): attends to real keys 0,1 normally. Valid output. Loss ignores via
labels=-100. No NaN. - Fully masked row (e.g., position 0 when position 0 itself is padded): position 0 can only attend to position 0, which is padded → entire row
-inf→ softmax → NaN.
Production code handles NaN rows via:
- Skip computing attention for padded queries entirely (varlen kernels do this).
- Replace NaN rows with zeros after softmax (
torch.where(torch.isnan(attn_weights), 0, attn_weights)). - Ensure padded queries never reach the attention computation (forward pass short-circuits them).
The loss function (with labels = -100 for pad positions) ignores padded positions, but only if NaN has been replaced or skipped — NaN * 0 in cross-entropy is still NaN.
Staff/Principal engineering judgment: Fully masked rows are unsafe by default — with plain softmax they produce NaN, with SDPA they produce zeros. A common production bug is assuming padded-query rows yield zero attention weights (they don’t with plain softmax; they yield NaN when fully masked). The fix is to zero-replace NaN rows after softmax or to skip padded queries before attention. If you see NaN in attention weights or loss, check whether fully-masked rows were handled. The mask construction itself (AND of causal + key_pad) is correct; the NaN handling is a separate production concern.
7. Practical exercise
Section titled “7. Practical exercise”Problem: Given the following batch configuration, construct all masks and verify the future-token leakage test.
B = 2,T = 5,d = 16,h = 4,d_head = 4- Sequence lengths:
[3, 4](batch item 0 has 3 real tokens then 2 pads; batch item 1 has 4 real tokens then 1 pad) - Token IDs (pad token =
0):[[5, 12, 7, 0, 0], [3, 9, 1, 0, 0]]
Steps:
- Build the key padding mask (
attention_mask) for both batch items, using the Hugging Face convention (True= real token). - Build the causal mask (
causal) forT = 5, then broadcast to[1, 1, 5, 5]. - Broadcast the key padding mask to
[B, 1, 1, T], then AND with the causal mask to produce the combined mask (combined) of shape[B, 1, 5, 5]. - Write the PyTorch smoke-test code (as in §4.2) that verifies future-position weights are ~0, padded keys have ~0 weight from all queries, padded queries attend to real keys normally, and non-future real positions > 0.
- What-if: If the combined mask were computed with OR instead of AND, which entries in the combined mask would differ, and what failure mode would that produce? Answer: The OR would let padding positions attend causally. For
b=0, i=3, the OR would keepTrueatj=4(since causal[3,4]=True and pad[b=0,4]=False, OR gives True). The test would FAIL because position 3 would have a non-zero attention weight at future position 4, even though position 4 is padding.
Exercise answers (instructor):
-
Key padding mask:
b=0:[True, True, True, False, False](real at 0,1,2; pad at 3,4)b=1:[True, True, True, True, False](real at 0,1,2,3; pad at 4)
-
Causal mask (T=5, lower-triangular,
True= attend):j=0 j=1 j=2 j=3 j=4i=0: T T T T Ti=1: F T T T Ti=2: F F T T Ti=3: F F F T Ti=4: F F F F T -
Combined mask (AND broadcast →
[2, 1, 5, 5]):b=0: same pattern as the worked example in §3.2 (from the causal & pad[T,T,T,F,F] row).b=1: causal AND pad[T,T,T,T,F] → each row AND-ed with the pad pattern.
-
Smoke-test code: See §4.2. Expected output: future positions ~0, padded keys ~0 weight from all queries, padded queries attend to real keys normally, non-future real positions > 0.
-
What-if: OR instead of AND. For
b=0, i=3: causal[3,4]=True, pad[b=0,4]=False → OR gives True. The test would FAIL because position 3 would have a non-zero attention weight at future position 4. The failure mode: the model “sees” a future token (position 4 is padding, but the attention weight is non-zero), which constitutes future-token leakage.Post-fix instructions: After correcting OR back to AND, re-run the smoke test from §4.2 and verify all four checks pass: (1) future positions have weight ~0 for every query and head; (2) padded keys (j=3,4 for b=0; j=4 for b=1) have weight ~0 from all queries; (3) real non-future keys have weight > 0; (4) padded queries attend to real keys normally. If any check fails, the combined mask still has an indexing or broadcast error — re-verify that
key_padindexes the key dimensionj, not the query dimensiont.
8. Visual and text equivalent
Section titled “8. Visual and text equivalent”Causal + Padding Mask Construction
Read this diagram as text
Text equivalent: Panel A shows the causal lower-triangular mask for T=5: position i can attend to positions 0..i. Panel B shows the key padding mask for B=2, T=5: True=real token, False=pad. Positions 3,4 are pad for b=0; positions 4 is pad for b=1. Panel C shows the combined mask: causal [1,1,5,5] AND key_pad [2,1,1,5] → [2,1,5,5], with the broadcast AND rule making each batch's mask independent. Panel D shows the leakage-test attention weight matrix: future positions (upper-right triangle) are ~0; padded keys have ~0 weight from all queries; padded queries attend to real keys normally; non-future real positions have valid weights. The engineering judgment: include the future-token leakage test as a CI gate for any mask change.
Text equivalent: The diagram has four panels. Panel A illustrates the causal lower-triangular mask for
T=5: positionican attend to positions0throughi. Panel B shows the key padding mask forB=2, T=5:True= real token,False= pad. Forb=0, positions 3 and 4 are padding; forb=1, position 4 is padding. Panel C combines the two: the causal mask[1,1,5,5]is AND-ed with the key padding mask[2,1,1,5]via broadcasting, producing the combined mask[2,1,5,5]. The broadcast AND rule means each batch item’s mask is independent —b=0’s mask usesb=0’s padding pattern, andb=1’s mask usesb=1’s. Panel D shows the attention weight matrix from the leakage test: the upper-right triangle (future positions) is ~0; the bottom row of each batch (padded positions) is all zero; non-future real positions have valid ( > 0) weights. The engineering judgment summarized below should guide the decision.
Engineering judgment (Staff/Principal)
Section titled “Engineering judgment (Staff/Principal)”| Question | Correct approach | Why it matters |
|---|---|---|
| Causal vs. OR combination | combined = causal & key_pad_expanded (AND) |
OR would let padding positions attend causally, a silent correctness bug that produces wrong logits. |
| Padding mask convention | True = real token, False = pad (HF DataCollatorForLanguageModeling(mlm=False)) |
The opposite convention (True = pad) appears in some other kits; using the wrong convention flips all mask semantics and breaks the test. |
| Fully masked rows | Plain softmax → NaN; SDPA → all-zero row. Zero-replace NaN or skip padded queries. | NaN with plain softmax means the row was fully masked; verify mask construction and NaN handling before training. |
| Future-token leakage test | Include as a CI gate for any mask change | One forward pass catches missing causal, missing padding, and wrong AND/OR combination — the single most important verification step. |
| Causal mask during inference | Always applied, with T = 1 at each auto-regressive step |
The causal mask is not only for training; it is the default mask for generation. |
Staff/Principal judgment: For any decoder-only model in the 100M–1B parameter range, always construct the combined mask as causal_4d & key_pad_4d, use the Hugging Face DataCollatorForLanguageModeling(mlm=False) convention for the key padding mask (True = real token), and include the future-token leakage test as a mandatory CI gate. The FLOP overhead of the combined mask is negligible; the correctness risk of getting it wrong is high (silent logit errors that are hard to debug). Document the mask-combining formula and the test procedure in the model’s ADR (architecture decision record).