Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Attention Masking › Efficient Masking Implementation

Michael BrenndoerferUpdated May 27, 20255 min read

Language AI HandbookAll chapters

Part 13 of 67Self-Attention

Chapter 93Attention Masking › Efficient Masking Implementation

Reading section 8 of 13 · Efficient Masking Implementation · 5 minSection 8 of 13 · 5 min

Section 8 of 13 · 5 min read

Attention Masking › Efficient Masking ImplementationOne section · Sections in this chapter8 / 13 ↓

Sections in this chapter13

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

With our understanding of padding and causal masks in place, we can now build a complete masked attention function. The key insight is that masking fits directly into the standard attention formula through a single addition operation before the softmax. No additional branches, no special-casing, no changes to the backpropagation graph: just one extra tensor addition that changes everything about which positions contribute.

The Masked Attention FormulaLink Copied

Recall that standard scaled dot-product attention computes:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

where:

$Q$ : the query matrix, shape $(n, d_k)$ , where each row is a query vector for one position
$K$ : the key matrix, shape $(n, d_k)$ , where each row is a key vector for one position
$V$ : the value matrix, shape $(n, d_v)$ , where each row is a value vector for one position
$d_k$ : the dimension of the query and key vectors, used in the scaling factor $1/\sqrt{d_k}$

To add masking, we simply insert the mask matrix $M$ before the softmax:

\text{Attention}(Q, K, V, M) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V

where:

$Q$ : the query matrix of shape $(n, d_k)$ , representing what each position is "looking for"
$K$ : the key matrix of shape $(n, d_k)$ , representing what each position "offers"
$V$ : the value matrix of shape $(n, d_v)$ , containing the information to aggregate
$M$ : the mask matrix of shape $(n, n)$ , with 0 for allowed and $-\infty$ for blocked positions
$d_k$ : the dimension of queries and keys, used for scaling to prevent vanishing gradients
$n$ : the sequence length

The formula unfolds in three steps:

Compute raw scores: $QK^T$ produces an $(n, n)$ matrix where entry $(i, j)$ measures the similarity between query $i$ and key $j$ .
Scale and mask: Divide by $\sqrt{d_k}$ to stabilize gradients, then add the mask $M$ . Masked positions receive $-\infty$ , which the softmax will convert to near-zero weights.
Normalize and aggregate: Softmax converts scores to weights summing to 1 (per row), then these weights select and combine values from $V$ .

This formulation is elegant because the mask $M$ does not change the computational structure. We still compute all pairwise scores, but the masking happens through simple addition before softmax. The exponential function in softmax then "erases" the masked positions. The gradient with respect to masked positions is also effectively zero, so backpropagation naturally ignores them during parameter updates.

Implementing Masked AttentionLink Copied

Let's translate the formula into code. The implementation follows the three-step structure exactly:

In[27]:

Code

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Compute scaled dot-product attention with optional masking.

    Args:
        query: shape (batch, seq_len, d_k) or (batch, heads, seq_len, d_k)
        key: same shape as query
        value: same shape as query
        mask: broadcastable mask, 0 for attend, -inf for block

    Returns:
        output: weighted values, same shape as value
        attention_weights: softmax weights
    """
    d_k = query.shape[-1]

    # Step 1: Compute raw similarity scores
    # (batch, ..., seq_len, d_k) @ (batch, ..., d_k, seq_len)
    # -> (batch, ..., seq_len, seq_len)
    scores = query @ key.swapaxes(-2, -1) / np.sqrt(d_k)

    # Step 2: Apply mask before softmax
    if mask is not None:
        scores = scores + mask

    # Step 3: Normalize to weights and aggregate values
    attention_weights = softmax(scores, axis=-1)
    output = attention_weights @ value

    return output, attention_weights

Notice how the mask application is just a single line: scores = scores + mask. This simplicity is intentional. The mask contains 0 for allowed positions (no effect on scores) and large negative values for blocked positions (which softmax converts to near-zero weights). No branching, no special cases, just addition. The mask is a first-class citizen of the attention formula, not a post-hoc patch.

Testing the ImplementationLink Copied

Let's verify that our implementation correctly handles both padding and causal masks:

In[28]:

Code

# Test setup
batch_size = 2
seq_len = 6
d_model = 8

# Random query, key, value
Q = np.random.randn(batch_size, seq_len, d_model)
K = np.random.randn(batch_size, seq_len, d_model)
V = np.random.randn(batch_size, seq_len, d_model)

# Padding: batch 0 has 6 tokens, batch 1 has 4 tokens
padding_indicator = np.array(
    [
        [True, True, True, True, True, True],  # All real
        [True, True, True, True, False, False],  # 4 real + 2 padding
    ]
)

# Create combined mask (causal + padding)
causal = create_causal_mask(seq_len)[np.newaxis, :, :]  # (1, seq_len, seq_len)
padding = np.where(
    padding_indicator[:, np.newaxis, :],  # (batch, 1, seq_len)
    0.0,
    -1e9,
)
combined_mask = causal + padding

# Run attention
output, weights = scaled_dot_product_attention(Q, K, V, mask=combined_mask)

Out[29]:

Console

Output shape: (2, 6, 8)
Attention weights shape: (2, 6, 6)

Batch 1 attention weights (4 real tokens + 2 padding):
Last 2 columns should be ~0 (padding), upper triangle should be ~0 (causal)
[[1.    0.    0.    0.    0.    0.   ]
 [0.286 0.714 0.    0.    0.    0.   ]
 [0.046 0.626 0.328 0.    0.    0.   ]
 [0.158 0.31  0.169 0.363 0.    0.   ]
 [0.552 0.061 0.124 0.262 0.    0.   ]
 [0.369 0.195 0.336 0.1   0.    0.   ]]

The attention weight matrix confirms that both masks work together correctly. Looking at the output for batch 1 (which has 4 real tokens and 2 padding tokens):

Columns 4 and 5 are nearly zero: The padding mask successfully blocks attention to padding positions. No real token wastes attention on meaningless padding.
Upper triangle is nearly zero: The causal mask blocks attention to future positions. Position 0 cannot see positions 1-5, position 1 cannot see positions 2-5, and so on.
Lower-left region has non-zero weights: The intersection of "past positions" and "real tokens" receives all the attention. These are exactly the positions each query should attend to.

The combined mask creates a triangular pattern truncated by the padding boundary. This is precisely what a causal language model needs when processing variable-length batched sequences. The two masks, one enforcing temporal ordering and one enforcing content validity, combine into a single constraint.