Attention Step by Step

Attention lets each query mix value vectors according to its similarity with keys. For key width $d_k$, scaled dot-product attention is

$\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V.$

The mask $M$ is applied before softmax: disallowed key positions receive negative infinity and therefore zero probability.

Prerequisites

Know batched matrix multiplication, softmax, and the token/padding distinction from Text Data and Tokenization.

Learning goals

• trace Q, K, V, score, weight, and output shapes.
• apply a key padding mask before softmax and verify padded keys receive zero weight.
• distinguish a key mask from a query mask and from a causal mask.

Worked example

Here (batch, query, key) scores have shape (1, 4, 4). The fourth key is padding. masked_fill with -torch.inf removes it from every query's softmax denominator. We separately zero the padded query output.

import math
import torch

queries = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.0, 0.0]]])
keys = queries.clone()
values = torch.tensor([[[10.0, 0.0], [0.0, 10.0], [5.0, 5.0], [99.0, 99.0]]])
key_padding_mask = torch.tensor([[False, False, False, True]])

scores = queries @ keys.transpose(-2, -1) / math.sqrt(keys.shape[-1])
scores = scores.masked_fill(key_padding_mask[:, None, :], -torch.inf)
weights = torch.softmax(scores, dim=-1)
output = weights @ values
output = output.masked_fill(key_padding_mask[:, :, None], 0.0)

print(f"Q={tuple(queries.shape)} K={tuple(keys.shape)} V={tuple(values.shape)}")
print(f"scores={tuple(scores.shape)} weights={tuple(weights.shape)} output={tuple(output.shape)}")
print(f"first_weights={[round(v, 3) for v in weights[0, 0].tolist()]}")
print(f"padded_key_max={weights[..., -1].max():.1f} padded_query={output[0, -1].tolist()}")

Output:

Q=(1, 4, 2) K=(1, 4, 2) V=(1, 4, 2)
scores=(1, 4, 4) weights=(1, 4, 4) output=(1, 4, 2)
first_weights=[0.401, 0.198, 0.401, 0.0]
padded_key_max=0.0 padded_query=[0.0, 0.0]

A key padding mask blocks padding as a source. It does not automatically erase padded query rows. A causal mask is different: query position $t$ cannot attend to later keys. Multi-head attention repeats projections in several subspaces, concatenates head outputs, and projects them again.

Common mistakes

• Applying a binary mask after softmax and forgetting to renormalize.
• Reversing boolean-mask meaning between libraries or APIs.
• Masking keys but then pooling padded query outputs.
• Treating attention weights as a faithful causal explanation of a prediction.
• Forgetting the quadratic score matrix cost in sequence length.

Practice

Add a causal mask to the score tensor. For query index 1, verify that key indices 2 and 3 receive zero weight. State both mask shapes before writing code.

Next steps

Transformer Architecture places multi-head attention inside residual, normalization, and feed-forward blocks.

Primary sources

Checked 2026-10-16: Attention Is All You Need, PyTorch scaled_dot_product_attention, and PyTorch MultiheadAttention.