Attention Masking › Mask Shapes and Broadcasting
Language AI HandbookAll chapters Attention Masking › Mask Shapes and BroadcastingOne section · Sections in this chapter7 / 13 ↓
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
Attention masks can have different shapes depending on the use case. Understanding these shapes and how they broadcast determines whether the implementation works as intended.
The attention score matrix has shape (batch, num_heads, seq_len, seq_len) in a full transformer with multi-head attention. Masks can be provided in more compact forms and rely on NumPy or PyTorch's broadcasting rules to expand to the full shape automatically. The four common mask shapes are:
- Full shape
(batch, num_heads, seq_len, seq_len): Fully specified mask for every batch item and head, giving maximum control but consuming the most memory - Per-batch
(batch, 1, seq_len, seq_len): Same mask across all heads for each batch item; the head dimension broadcasts tonum_heads - Per-batch compact
(batch, 1, 1, seq_len): Key-only masking where the query dimension also broadcasts; this is the natural shape for padding masks, since padding depends only on which keys are padding, not on which query is asking - Global
(1, 1, seq_len, seq_len): Same mask for the entire batch; this is the natural shape for causal masks, since the causal constraint is the same for every sequence
Broadcasting allows us to store the mask in its most natural compact form and let the framework handle expansion. This is both memory-efficient and computationally clean: you define the mask at the level of abstraction that matches its semantics, and broadcasting handles the rest.
# Different mask shapes and their purposes
batch_size = 2
num_heads = 4
seq_len = 8
# Causal mask: same for all batches and heads
# Shape: (1, 1, seq_len, seq_len)
causal_mask = create_causal_mask(seq_len)[np.newaxis, np.newaxis, :, :]
# Padding mask: varies per batch, same across heads
# Original shape: (batch, seq_len)
# Expanded shape: (batch, 1, 1, seq_len)
padding_indicator = np.array(
[
[True, True, True, True, True, False, False, False], # 5 real tokens
[True, True, True, False, False, False, False, False], # 3 real tokens
]
)
padding_mask_expanded = padding_indicator[:, np.newaxis, np.newaxis, :]
attention_padding = np.where(padding_mask_expanded, 0.0, -1e9)
# Simulated attention scores
scores = np.random.randn(batch_size, num_heads, seq_len, seq_len)
# Broadcasting handles dimension matching automatically
masked_scores = scores + causal_mask + attention_paddingCausal mask shape: (1, 1, 8, 8) Padding mask shape: (2, 1, 1, 8) Scores shape: (2, 4, 8, 8) Masked scores shape: (2, 4, 8, 8) Broadcasting expands smaller dimensions automatically: Causal (1,1,8,8) + Padding (2,1,1,8) + Scores (2,4,8,8) = (2,4,8,8)
Broadcasting rules make mask combination elegant. The causal mask (1, 1, 8, 8) applies identically to all batch items and heads. The padding mask (2, 1, 1, 8) applies per batch item across all heads, masking keys based on which positions are padding. The final result (2, 4, 8, 8) has the correct combined mask for each batch item and attention head. Neither mask needs to be explicitly replicated to match the full score tensor shape; broadcasting handles that automatically and without allocating additional memory.
Memory-Efficient Masking
For very long sequences, storing full (seq_len, seq_len) masks becomes expensive. A 10,000-token sequence requires 100 million entries per mask. For a batch of 32 sequences with 4 attention heads, the full mask tensor would require 32 * 4 * 10,000 * 10,000 * 4 bytes = 51.2 GB, which is completely infeasible. Several strategies reduce this cost.
Lazy mask generation: Instead of precomputing the full mask, generate it during the attention computation.
def apply_causal_mask_lazy(scores):
"""
Apply causal mask without pre-allocating a full mask tensor.
Uses numpy's triu to directly set upper triangle to -inf.
"""
seq_len = scores.shape[-1]
# Create mask in-place using views
masked = scores.copy()
# Get upper triangle indices
row_indices, col_indices = np.triu_indices(seq_len, k=1)
# Set upper triangle to -inf
masked[..., row_indices, col_indices] = -1e9
return masked
# Compare memory: full mask vs lazy
seq_len = 1000
full_mask_memory = seq_len * seq_len * 4 # 4 bytes per float32Full mask memory for seq_len=1000: 4,000,000 bytes (4.0 MB) Lazy approach: generates mask on-the-fly, no pre-allocation needed
A 1000-token sequence requires 4 MB just for the mask tensor. The lazy approach avoids this allocation by setting mask values directly on the score matrix during computation. For very long sequences, this savings becomes critical.
Boolean masks: Store masks as boolean arrays (1 byte per element) rather than float32 (4 bytes), converting to float only when needed.
# Boolean mask storage
seq_len = 1000
bool_mask = np.tril(np.ones((seq_len, seq_len), dtype=bool))
float_mask = np.tril(np.ones((seq_len, seq_len), dtype=np.float32))
bool_memory = bool_mask.nbytes
float_memory = float_mask.nbytesBoolean mask: 1,000,000 bytes (1.0 MB) Float32 mask: 4,000,000 bytes (4.0 MB) Memory savings: 75%
For sequences of length 1000, boolean masks use 75% less memory. This adds up when processing large batches or very long documents. Modern deep learning frameworks like PyTorch support boolean masks natively and perform the float conversion internally only when needed, so you get the memory savings without manual conversion code.