Sequence Masking (causal and padding masks)

Why masking exists (motivation)

In attention, every query position $q$ produces a distribution over key positions $k$:

The softmax turns each row of scores $S_{q,:}$ into weights that sum to 1. That “sum to 1” is the core reason masking matters: if your sequence contains invalid tokens (padding) or forbidden tokens (future positions during autoregressive decoding), softmax will still allocate some probability mass to them unless you explicitly prevent it.

Masking is the standard way to encode structural constraints into attention: “these keys may be attended to, those keys must not.”

Definition

An attention mask is a matrix $M$ with entries:

• •$M_{q,k} = 1$ means attention from query position $q$ to key position $k$ is allowed.
• •$M_{q,k} = 0$ means it is blocked.

Shape-wise, the conceptual mask is often $(L_q \times L_k)$, where:

• •$L_q$ = number of query positions
• •$L_k$ = number of key positions

In self-attention, typically $L_q = L_k = L$.

How a mask actually affects softmax

Masks are applied to the attention logits (scores) before the softmax. The most common implementation is additive masking:

Then

Because $\exp(-\infty)=0$, masked positions get zero probability.

In real code, you don’t literally use $-\infty$; you use a large negative constant like $-10^4$ or $-10^9$ (depending on dtype and stability). The concept is “so negative that softmax assigns ~0.”

The two mask families

1. 1)Padding mask (data validity): blocks keys that are padding tokens.
2. 2)Causal mask (information flow): blocks keys that are to the right of the query (future positions), enforcing autoregressive behavior.

They solve different problems and are often combined.

A first mental model: “mask edits the graph”

Think of attention as a directed weighted graph from queries to keys. Masking removes edges before normalization.

• •Padding mask removes edges that point into padded (fake) nodes.
• •Causal mask removes edges that point forward in time.

That’s it: masking is not a learned component; it’s a rule baked into the computation.

Why causal masking is needed

In autoregressive language modeling, at position $t$ the model must predict token $x_t$ using only past context $x_{<t}$.

Self-attention without constraints would let position $t$ attend to future tokens $x_{>t}$ during training (because they are present in the same sequence). That creates label leakage: the model can “cheat” by looking at the answer.

Causal masking prevents this by allowing attention only to the prefix.

The causal rule

For a sequence length $L$ with 0-indexed positions:

• •Query position $q$ may attend to key positions $k \le q$.
• •Keys with $k > q$ are blocked.

So the mask is lower-triangular.

Inline diagram: the (L×L) causal triangle

Let $L=6$ and write 1=allow, 0=block. Rows are queries $q$, columns are keys $k$.

Causal mask M_causal (L=6):

      k: 0 1 2 3 4 5
q=0      1 0 0 0 0 0
q=1      1 1 0 0 0 0
q=2      1 1 1 0 0 0
q=3      1 1 1 1 0 0
q=4      1 1 1 1 1 0
q=5      1 1 1 1 1 1

This is the “causal triangle.” It encodes “no looking ahead.”

Construction in math

A clean definition is:

where $\mathbf{1}[\cdot]$ is the indicator function.

If you use an additive mask $B$ that adds $-\infty$ where blocked:

Then you compute $S' = S + B$.

Self-attention vs cross-attention

• •Decoder self-attention in an autoregressive Transformer uses a causal mask.
• •Encoder self-attention typically does not use a causal mask (it’s bidirectional).
• •Decoder cross-attention (decoder queries attending to encoder keys) typically does not use a causal mask, because encoder outputs are not “future tokens” of the decoder timeline; they’re a separate source sequence.

Subtlety: “causal” depends on the task

Causal masking is a task constraint, not a Transformer requirement.

• •Autoregressive generation: causal = required.
• •Masked language modeling (BERT-like): no causal mask; instead you hide tokens in input.
• •Prefix-LM (sometimes used in instruction tuning): allow full attention inside prefix, causal afterwards (a block-structured mask).

Practical detail: training vs inference

• •Training: you usually compute attention for all positions in parallel, so the causal mask is essential to ensure the parallel computation matches the autoregressive factorization.
• •Inference: if you generate token-by-token with KV caching, causal structure is implicit (the query is only for the current step and keys are only past). But many implementations still keep a causal mask for correctness and simplicity.

Why padding masking is needed

Real batches contain sequences with different lengths. To form a rectangular tensor, shorter sequences are padded to a common length $L$.

Padding tokens are not real data. If attention can attend to them, several things go wrong:

• •Probability mass gets wasted on “nothing.”
• •Values at padded positions can pollute representations.
• •Gradients can flow through padding in ways that destabilize training.

Padding masks ensure padded keys do not receive attention.

Key idea: padding mask is usually about keys

In attention, the output at query $q$ is a weighted sum over values at keys $k$.

If key $k$ is padding, then value vₖ is meaningless; we must prevent it from contributing.

So most padding masks are applied along the key dimension: for each query, block the same set of padded key positions.

From per-token validity to a (Lq×Lk) mask

Often you start with a 1D “valid tokens” vector for each sequence:

• •$p_k = 1$ if position $k$ is a real token
• •$p_k = 0$ if position $k$ is padding

This is shape $(L_k)$ (or $(B, L_k)$ for a batch).

To turn that into an attention mask you broadcast across queries:

So each row is identical: any query cannot attend to padded keys.

Inline diagram: combined (causal ∧ padding) mask over a padded sequence

Suppose we have a batch element with length 4 padded to $L=6$:

Positions: 0 1 2 3 4 5

Tokens: A B C D <pad> <pad>

Validity: 1 1 1 1 0 0

Padding-only mask (1=allow, 0=block) applied on keys:

M_pad (allow real keys only):

      k: 0 1 2 3 4 5
q=0      1 1 1 1 0 0
q=1      1 1 1 1 0 0
q=2      1 1 1 1 0 0
q=3      1 1 1 1 0 0
q=4      1 1 1 1 0 0
q=5      1 1 1 1 0 0

Causal mask (from earlier):

M_causal:

      k: 0 1 2 3 4 5
q=0      1 0 0 0 0 0
q=1      1 1 0 0 0 0
q=2      1 1 1 0 0 0
q=3      1 1 1 1 0 0
q=4      1 1 1 1 1 0
q=5      1 1 1 1 1 1

Combined mask uses logical AND (allow only if both allow):

Result:

M_combined = M_causal AND M_pad:

      k: 0 1 2 3 4 5
q=0      1 0 0 0 0 0
q=1      1 1 0 0 0 0
q=2      1 1 1 0 0 0
q=3      1 1 1 1 0 0
q=4      1 1 1 1 0 0
q=5      1 1 1 1 0 0

Notice how:

• •Causality blocks “future” keys.
• •Padding blocks keys 4 and 5 for every query.
• •Even though the causal mask would allow (q=5, k=5), padding still blocks it.

Do we also mask padded queries?

There are two related concerns:

1. 1)Keys/values padding: must be masked so they don’t contribute.
2. 2)Queries padding: if query positions are padding, their outputs are meaningless. Some systems:

• •compute them anyway but ignore them in the loss, or
• •explicitly zero them out afterward, or
• •prevent attention computation for them via masking.

In practice, padding keys is non-negotiable; padding queries is optional but can save compute and avoid odd numerical artifacts.

Combining masks: additive form is usually easiest

If your implementation uses additive masks $B$ (0 for allow, −∞ for block), then combining is simple:

Because:

• •if either term is −∞, the sum is −∞ (blocked)
• •if both are 0, the sum is 0 (allowed)

This “sum of additive masks” is a very common pattern in Transformer codebases.

Where the mask goes in scaled dot-product attention

For one batch element and one head, attention typically follows:

1) Compute scores:

2) Apply mask to get masked logits $S'$.

3) Softmax over keys:

4) Weighted sum of values:

Masking must happen before softmax. Masking after softmax is not equivalent because the distribution would no longer be normalized correctly.

Shape/broadcasting schematic (most common source of bugs)

In multi-head attention with batching, you commonly work with these shapes:

• •$Q$: (B, H, Lq, d)
• •$K$: (B, H, Lk, d)
• •$V$: (B, H, Lk, d)
• •Scores $S = QK^\top$: (B, H, Lq, Lk)

You want a mask that can broadcast to (B, H, Lq, Lk).

Here are typical mask shapes and how they broadcast:

Target scores S:      (B, H, Lq, Lk)

Causal mask:          (1, 1, Lq, Lk)   or (Lq, Lk)
Padding key mask:     (B, 1, 1,  Lk)   from (B, Lk)
Combined mask:        (B, 1, Lq, Lk)   after broadcast AND/add

Key principle:

• •Causal depends on (Lq, Lk) only (same for every batch element).
• •Padding depends on which tokens are real for each batch element (varies with B), but is usually identical across heads and queries.

Table: common mask representations

Numerical stability: why “−∞” must be chosen carefully

Softmax is typically computed using a stabilized form:

If masked positions are very negative, $\exp(\text{very negative}) \approx 0$, which is what you want.

But two pitfalls:

1. 1)If all positions in a row are masked, you get softmax over all −∞, which can produce NaNs.
2. 2)With low-precision dtypes (fp16/bf16), extremely negative constants can underflow in unexpected ways.

Typical engineering fixes:

• •Ensure every query has at least one unmasked key (e.g., for causal masks, diagonal is allowed).
• •If some queries correspond to padding, avoid computing loss on them and/or force a safe attention pattern.
• •Use library attention kernels that handle masking robustly.

Worked integration: masked attention logits

Let $S$ be scores and $M$ be a binary mask.

Convert $M$ to an additive mask:

Then apply:

Then:

Finally output:

Connection to Transformers

Masks are a “plumbing” idea that becomes crucial in:

• •Decoder-only Transformers (GPT-like): causal mask is central.
• •Encoder-decoder Transformers (T5-like):
• •encoder self-attention uses padding mask,
• •decoder self-attention uses causal + padding,
• •decoder cross-attention uses padding mask for encoder keys.

Once you can build and broadcast these masks correctly, multi-head attention and full Transformers become much easier to implement without silent correctness bugs.