What does a padding mask zero out, and why is it needed when batching variable-length sequences?

Flashcard·Easy·4.0 · 0·~30s

Attempt it

A padding mask zeros out attention to the special pad tokens that are appended to shorter sequences so an entire batch can pack into a single rectangular tensor. Without the mask, attention would treat pad positions as real tokens and let their keys/values pollute the weighted sum, dragging every output toward a meaningless average. Mechanically, the mask is added to the pre-softmax score matrix with -inf at pad columns. After softmax, those positions contribute exactly 0 to the context vector. The mask is orthogonal to the causal mask: padding mask hides batching artifacts, causal mask hides future tokens.

Your answer

**Real sequences come in different lengths.** A batch of training examples might have sequences of length 12, 47, 8, and 33. GPUs want one rectangular tensor per batch, so every sequence gets padded with a special token (usually `[PAD]` or token ID 0) to match the length of the longest sequence in the batch. Those pad tokens are not real content. They are placeholders that exist only to make the tensor rectangular. If attention treats them as legitimate keys, then every query token at a real position will read from the pad columns and mix their (meaningless) values into its output. The model's representation gets contaminated by content that does not exist. **The padding mask** fixes this. It is a 1D boolean (or 0/1) tensor marking which positions are pad. Before applying softmax over the attention scores, the mask is converted to an additive bias: 0 at real positions, -inf at pad positions. The score matrix gets the bias added column-wise, so every pad column becomes -inf for every query row. After softmax, e^(-inf) = 0 contributes exactly nothing to the normalization, and the pad positions get probability 0. The attention output for every real query is a weighted sum over only the real keys. The padding mask is **orthogonal** to the causal mask. The causal mask blocks future tokens; the padding mask blocks fake tokens. Both can apply at the same time, and at training of causal LMs on variable-length batches, both usually do.

**The padding mask zeros out attention to pad tokens** that get appended to short sequences when batching. The mechanism is the same trick used by every attention mask: add -inf to the pre-softmax scores at the masked positions, and softmax converts that to a probability of 0 after normalization. **Why it is needed.** Variable-length sequences need to share a rectangular batch tensor for GPU efficiency. Padding short sequences to the batch's max length is the standard solution. Without a padding mask, attention treats pad tokens as legitimate keys: their projected K and V vectors influence the softmax-weighted sum at every real query position. The model's output drifts toward the padding embedding, and quality collapses. **The mechanism, step by step.** 1. Compute scores: `scores = Q @ K^T / sqrt(d_head)`, shape `(B, n_heads, T, T)`. 2. Build a mask: `mask = padding_mask[None, None, None, :]` of shape `(B, 1, 1, T)`, with 0 at real columns and -inf at pad columns. 3. Add the mask: `scores = scores + mask`. 4. Softmax along the last axis: `weights = softmax(scores, dim=-1)`. Pad columns now have weight exactly 0. 5. Apply weights to V: pad positions contribute nothing. **Compatible with the causal mask.** Both masks are typically broadcast and summed before softmax. A position that is either future (causal-blocked) or padding (padding-blocked) becomes -inf, and the model sees only real, non-future keys. At inference, padding masks are less common because batches are typically formed from queued requests of similar length (or single-example inference). At training, especially on diverse-length corpora, the padding mask is non-negotiable.

Quick reference

TL;DR

A padding mask zeros out attention to pad tokens added during batching. It adds -inf to pre-softmax scores at pad columns so e^(-inf) = 0 contributes nothing after normalization.

Memory aid

Easy to grasp

Picture a teacher running a group quiz where every team must turn in an answer sheet with exactly 100 boxes filled. Some teams only have 30 real answers, so they fill the remaining 70 boxes with the word 'BLANK' to make the sheet look the same size as everyone else's. When the teacher grades by averaging answers, they need to know to skip every box marked BLANK; otherwise the average of every team is dragged toward 'BLANK' and nothing means anything. The padding mask is exactly the rule 'skip BLANK boxes'. Without it, attention happily averages in the pad tokens and the model's output becomes garbage.

Key concepts

Padding mask Causal mask Variable-length batching

Detailed answer & concept explanation

~7 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.

Quick scan · 2–4 min

Everything important, quickly.

Explain why padding exists, walk through the additive -inf trick at the score matrix, show the softmax converting -inf to 0, distinguish padding from causal masks, and close with the FlashAttention-2 cu_seqlens replacement and continuous-batching alternative at inference.

Code example

python

import torch
import torch.nn.functional as F

B, T, n_heads, d_head = 2, 8, 4, 16
# Sequence 0 has 5 real tokens, sequence 1 has 8 real tokens.
real_lengths = torch.tensor([5, 8])

Q = torch.randn(B, n_heads, T, d_head)
K = torch.randn(B, n_heads, T, d_head)
V = torch.randn(B, n_heads, T, d_head)

# Build padding mask: True at real positions, False at pad.
idx = torch.arange(T)[None, :]  # (1, T)
is_real = idx < real_lengths[:, None]  # (B, T)

# Convert to additive: 0 at real, -inf at pad.
additive = torch.where(is_real, 0.0, float('-inf'))  # (B, T)
additive = additive[:, None, None, :]  # (B, 1, 1, T) -- broadcast over heads and queries

scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)  # (B, n_heads, T, T)
scores = scores + additive  # pad columns become -inf
weights = F.softmax(scores, dim=-1)  # pad columns now exactly 0
output = weights @ V  # (B, n_heads, T, d_head), no contribution from pad keys

Real-world examples

Real products, models, and research that use this idea.

PyTorch's torch.nn.MultiheadAttention takes a key_padding_mask of shape (B, T) where True means 'this is padding'.
Hugging Face Transformers uses attention_mask of shape (B, T) where 1 means 'real token' (opposite polarity).
FlashAttention-2's cu_seqlens API replaces explicit padding masks with cumulative example lengths, avoiding (B, T, T) mask materialization.
vLLM's continuous batching avoids padding entirely at inference by dynamically composing batches at every decode step.
Llama 4 Maverick, Mistral, and Qwen 3.5 pretraining stacks use packed sequences with cu_seqlens rather than padding.

Interviewer follow-ups

What an interviewer would ask next. Try answering before peeking at the approach.

QWhy the additive -inf trick instead of multiplying the weights by zero after softmax?

QHow does the padding mask interact with the causal mask in causal-LM training?

QWhat is the memory cost of an explicit padding mask at long context?

Don't say this

Red flags and common mistakes that signal junior thinking. Click to expand.

Most common mistake

Confusing padding mask with causal mask. Padding hides fake batching tokens; causal hides real future tokens. Both can coexist in the same attention call.

60-second cheat sheet

The night-before-the-interview bullets. Scan these on the way to the call.

References

Primary sources. Skim if you want the original framing.