Attention

Problem

A language model has to decide what "it" refers to in "The cat sat because it was tired", and a translation model has to decide which source word to look at while producing each target word. The answer depends on the whole sentence, and the distance between the words varies. Earlier designs squeezed a sentence through a recurrent state one step at a time (see the recurrent networks that came before), so information from the start had to survive every step. Convolutions see only a fixed window.

What we want is a layer where each position can look directly at any other position, decide how much each matters for the current job, and take a weighted blend of what it finds, with the weights computed from the content rather than fixed by position. That layer is attention.

Intuition

Think of a library lookup. Each position writes a query ("what am I looking for?"), a key ("what do I contain, for the purpose of being found?") and a value ("what will I hand over if you pick me?"). To compute the output at position i, take its query, compare it with every key by dot product (a big dot product means a good match, see vectors and dot products), turn the scores into weights that sum to 1 with a softmax, and return the weighted average of the values. A position whose key matches strongly contributes most of its value.

Nothing here is fixed in advance: the queries, keys and values are produced from the input by learned matrices, so training decides what counts as a match. Multi-head attention runs several such lookups in parallel, each in its own smaller subspace, because one blended average can express only one kind of relationship at a time (who is the subject, which noun does this pronoun point to, what came just before).

The same four tokens, two heads, two different patterns. Both heads still put weight on the last column, which is padding. Masks (below) remove that.

Mechanics

Queries, keys and values

Let X be the input, shape [L, E]: L positions, each an E-dimensional vector (the output of an embedding layer, or of the previous layer). Three learned matrices turn it into queries, keys and values:

Q = X W_Q^T   K = X W_K^T   V = X W_V^T    shapes: [L, E] @ [E, E] → [L, E]

The transpose is there because PyTorch stores a linear layer's weight as [out, in] and computes x @ weight.T. For self-attention all three come from the same X. For cross-attention the queries come from one sequence and the keys and values from another (below). Queries and keys must share a width (they get dotted); values may differ.

Scaled dot-product attention

Attention(Q, K, V) = softmax( Q K^T / √d_k + M ) V    shapes: [L, d_k] @ [d_k, S] → [L, S]; [L, S] @ [S, d_v] → [L, d_v]

1. Scores. Q K^T is an [L, S] matrix: entry (i, j) is the dot product of query i with key j.
2. Scale. Divide by √d_k. A dot product of two d_k-vectors with independent unit-variance entries has variance d_k, so without scaling the scores grow with the width, the softmax saturates toward a one-hot row, and its gradient (which is tiny where the softmax is saturated) stalls training. Dividing by √d_k brings the variance back to 1.
3. Mask. Add M: 0 where attention is allowed and −∞ where it is not, so the softmax gives those positions exactly zero weight.
4. Softmax along the key axis (each row). Row i becomes a probability distribution over the S positions.
5. Mix. Multiply by V: output row i is the weights-weighted average of the value rows.

The score matrix is L × S, so memory and compute grow with the square of the sequence length; that cost drives most of the engineering at the end of this page.

Multi-head attention

With E = 4 and h = 2 heads, each head works in d_k = E / h = 2 dimensions. The layer projects once to full width and then reshapes so that each head sees its own slice of columns:

[N, L, E] → view [N, L, h, E/h] → transpose [N, h, L, E/h]   attend per head   → [N, L, h·E/h] → @ W_O^T → [N, L, E]

The head split costs nothing: it is a view of the same memory. The heads attend independently (the head dimension acts as a batch dimension), their outputs are concatenated back to width E, and a final linear layer W_O mixes the heads. Total compute is about the same as one head at full width, but the model gets h different weight patterns. In F.multi_head_attention_forward the code spells the split out as q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1) and, on the way out, attn_output.permute(2, 0, 1, 3).reshape(bsz * tgt_len, embed_dim), followed by the output projection.

Green and amber are the two heads' column slices. Heads never see each other's columns until the concatenation, and the output projection is the only place they interact.

Masks: padding, causal, boolean and additive

The mask M in the formula says which key positions a query may use. Two situations need it. Padding: sequences in a batch have different lengths and are padded to the longest; padded positions must be ignored. Causal: when a model predicts the next token, position i must not see positions after i, or training would simply copy the answer. A causal mask is a lower-triangular pattern; a padding mask blocks whole columns.

A mask can be given two ways. An additive (float) mask is added to the scores: 0 means allowed, −∞ means blocked (finite negative values act as a soft penalty, which is how relative-position biases such as ALiBi work). A boolean mask is converted to the additive form before use. The trap is the direction of True, which differs between the two main APIs:

PyTorch's own docstring warns about exactly this and tells you to invert a boolean mask (~mask) when moving from MultiheadAttention to scaled_dot_product_attention. Inside MultiheadAttention, _canonical_mask performs the conversion (torch.zeros_like(mask, dtype=...).masked_fill_(mask, float("-inf"))), so a boolean True becomes −∞. A key_padding_mask of shape [N, S] is reshaped to [N·h, 1, S] and added to attn_mask, so the two combine by addition.

Two details matter in practice. If a whole row is masked (a padded query position that is also not allowed to see anything), softmax has no allowed entry and the result is NaN; implementations avoid this by keeping at least one allowed entry per row or by zeroing such rows afterwards. And is_causal=True is a hint that lets fused kernels skip building the mask; for a non-square [L, S] mask it means the upper-left aligned triangle, which differs from the lower-right alignment you want when a cache supplies S > L keys (see CausalBias below).

Masks are added to the scores, so they stack by addition: −∞ plus anything is −∞, and two 0s stay 0.

Cross-attention

When Q comes from one sequence (length L) and K, V from another (length S), the score matrix is [L, S] and the same formula applies. An encoder–decoder translation model uses this in every decoder layer: queries are the target tokens produced so far, keys and values are the encoder's output for the source sentence, so each target word picks the source words it needs. The projection code reflects it: in _in_projection_packed the "encoder–decoder" branch (keys and values are the same tensor but the query is not) splits the packed weight into W_Q and a stacked W_KV and projects the source once. A padding mask in this setting is over the source (S) axis.

Grouped-query attention

In ordinary multi-head attention every head has its own K and V. In grouped-query attention (GQA) there are H_q query heads but only H_kv key/value heads, and each K/V head serves a group of H_q / H_kv query heads. H_kv = H_q is ordinary multi-head attention; H_kv = 1 is multi-query attention. The saving is memory during generation: the keys and values of all past tokens are cached, and the cache is proportional to H_kv. For a model shaped like the example in PyTorch's docstring (32 query heads, 8 key/value heads, head width 128), with 32 layers and 2-byte values, the cache costs 131,072 bytes per token against 524,288 with 32 key/value heads, a 4× saving. F.scaled_dot_product_attention(..., enable_gqa=True) accepts K and V with fewer heads; its reference code does key.repeat_interleave(query.size(-3) // key.size(-3), -3), so with 8 query heads over 2 K/V heads, query heads 0 to 7 use K/V heads 0, 0, 0, 0, 1, 1, 1, 1, and each K/V head is shared by four query heads. It requires H_q divisible by H_kv.

nn.MultiheadAttention in PyTorch

The module holds one packed parameter in_proj_weight of shape [3E, E] (rows 0..E−1 are W_Q, the next E are W_K, the last E are W_V), an in_proj_bias of length 3E, and an out_proj linear layer [E, E]. Packing lets self-attention compute Q, K and V with one matrix multiplication and then split. If the key and value widths kdim and vdim differ from E, it stores three separate matrices instead (q_proj_weight, k_proj_weight, v_proj_weight) and in_proj_weight is None. All of them start from Xavier-uniform initialisation with zero biases. The parameter count is 4E² + 4E; for E = 512 that is 1,050,624.

Two code paths follow from the arguments. With need_weights=True (the default) the forward pass computes the scaled scores with baddbmm, a softmax, dropout and a bmm explicitly, because it must return the weights (averaged over heads unless average_attn_weights=False). With need_weights=False it calls F.scaled_dot_product_attention and returns None for the weights, which allows fused kernels. Separately, when the module is in eval mode, without gradients, with self-attention on batched batch_first=True input, an even number of heads and a few other conditions, forward takes a fast path to a fused inference kernel (and supports nested tensors for padding). The code records the first failing condition in why_not_fast_path; every condition must hold, so for example training mode or requires_grad inputs always use the ordinary path.

F.scaled_dot_product_attention and its backends

F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False) is the primitive. Its docstring gives the mathematical reference (the "Efficient implementation equivalent to the following" block) and the function dispatches to one of several kernels that compute the same thing:

By default PyTorch picks the best backend whose input constraints are met (dtype, head width, mask type, device). To force or exclude backends, use the context manager with sdpa_kernel(SDPBackend.FLASH_ATTENTION): ... (or a list, with set_priority=True to order them); it restores the previous state on exit. The fused kernels give slightly different floating-point results from each other, and dropout_p is applied whenever it is above 0, so a module should pass dropout_p=self.p if self.training else 0.0.

CausalBias and FlexAttention

torch.nn.attention.bias provides causal_upper_left(L, S) and causal_lower_right(L, S), mask objects that fused kernels recognise without materialising the matrix. For L = 3 and S = 4 the allowed pattern is [1000, 1100, 1110] for upper-left and [1100, 1110, 1111] for lower-right, where the lower-right form is what you want when the query is the last L positions of a sequence whose cache has S entries. The two coincide when L = S.

FlexAttention (torch.nn.attention.flex_attention, a prototype feature) generalises the idea. Instead of a mask tensor you write small Python functions and a compiler fuses them into the attention kernel:

• score_mod(score, b, h, q_idx, kv_idx) returns a modified score for one (query, key) pair, enough for relative-position biases, soft-capping or ALiBi.
• mask_mod(b, h, q_idx, kv_idx) returns a boolean, True meaning the pair is allowed (the SDPA direction). create_block_mask(mask_mod, B, H, Q_LEN, KV_LEN) evaluates it per block and produces a BlockMask, so blocks that are entirely masked (causal upper triangle, sliding-window outside, other documents in a packed batch) are skipped, not computed and discarded. and_masks and or_masks combine them.

The call is flex_attention(query, key, value, score_mod=None, block_mask=None, scale=None, enable_gqa=False) on 4-D tensors [B, H, L, E], and it needs torch.compile to be fast.

Worked example

Implementation

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

N, L, E, H = 1, 4, 4, 2
x = torch.randn(N, L, E)

# the reference formula, one head
def attention(q, k, v, mask=None):                    # q: [.., L, d], k, v: [.., S, d]
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))   # [.., L, S]
    if mask is not None:
        scores = scores + mask                        # additive: 0 allowed, -inf blocked
    return torch.softmax(scores, dim=-1) @ v

# the module: packed projection, heads, output projection
mha = nn.MultiheadAttention(E, H, batch_first=True)
mha.in_proj_weight.shape                              # [12, 4]: W_Q, W_K, W_V stacked
out, weights = mha(x, x, x)                           # out [1, 4, 4], weights [1, 4, 4] (mean over heads)
out, per_head = mha(x, x, x, average_attn_weights=False)   # weights [1, 2, 4, 4]

# masks, with the two opposite boolean conventions
padding = torch.tensor([[False, False, False, True]])         # MHA: True = ignore this key
causal_blocked = torch.triu(torch.ones(L, L, dtype=torch.bool), diagonal=1)   # MHA: True = not allowed
out, _ = mha(x, x, x, key_padding_mask=padding, attn_mask=causal_blocked)

q = k = v = x.view(N, L, H, E // H).transpose(1, 2)           # [N, H, L, E/H]
allowed = ~causal_blocked                                     # SDPA: True = allowed
y = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)             # [N, H, L, E/H]
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)                # same causal mask, no tensor
y = y.transpose(1, 2).reshape(N, L, E)                                     # merge heads

# grouped-query attention: 8 query heads share 2 key/value heads
q = torch.randn(1, 8, L, 16)
k = v = torch.randn(1, 2, L, 16)
F.scaled_dot_product_attention(q, k, v, enable_gqa=True)

# choose a backend explicitly
from torch.nn.attention import SDPBackend, sdpa_kernel
with sdpa_kernel(SDPBackend.MATH):
    F.scaled_dot_product_attention(q, k, v, enable_gqa=True)

The same computation in plain Python, with the projection, head split and mask spelled out:

import math

def softmax(row):
    top = max(row)
    exps = [math.exp(s - top) for s in row]
    return [e / sum(exps) for e in exps]

def attend(Q, K, V, bias=None):                       # Q [L, d], K [S, d], V [S, dv]
    out = []
    for i, q in enumerate(Q):
        scores = [sum(a * b for a, b in zip(q, k)) / math.sqrt(len(q))
                  + (bias[i][j] if bias else 0.0) for j, k in enumerate(K)]
        w = softmax(scores)
        out.append([sum(w[j] * V[j][c] for j in range(len(V))) for c in range(len(V[0]))])
    return out

def multihead(X, W_in, W_out, heads, bias=None):      # W_in [3E, E], W_out [E, E]
    E = len(X[0])
    proj = lambda W: [[sum(a * b for a, b in zip(x, row)) for row in W] for x in X]
    Q, K, V = proj(W_in[:E]), proj(W_in[E:2 * E]), proj(W_in[2 * E:])
    d = E // heads
    cols = lambda M, h: [row[h * d:(h + 1) * d] for row in M]
    merged = [sum(rows, []) for rows in zip(*[attend(cols(Q, h), cols(K, h), cols(V, h), bias)
                                              for h in range(heads)])]
    return [[sum(a * b for a, b in zip(m, row)) for row in W_out] for m in merged]

Tradeoffs

• Quadratic cost. The score matrix has L × S entries per head. Doubling the context quadruples the work and, in a naive implementation, the memory for the scores. Fused kernels (flash, memory-efficient) remove the memory term by never storing the matrix, but the arithmetic is still quadratic. Sliding-window masks, block-sparse patterns (FlexAttention) and linear-time alternatives such as state-space models exist to bend this.
• No built-in order. Attention treats its input as a set: shuffle the rows and the outputs shuffle the same way. Order enters only through positional information added to the embeddings, or through a causal mask.
• Heads versus width. More heads give more distinct patterns but each has fewer dimensions, so very narrow heads (below about 32 to 64) get less expressive. GQA and multi-query trade a little quality for much less cache memory at inference.
• Generation cache. At inference, keys and values of earlier tokens do not change, so they are cached instead of recomputed; the cache grows linearly with context and dominates memory for long sequences. The cache is why decoding speed is memory-bound.
• Pitfalls. Reversing the boolean mask direction between MultiheadAttention and SDPA silently attends to padding and masks the real tokens. Fully masked rows produce NaN. dropout_p in SDPA applies even in eval mode unless you pass 0. need_weights=True forces the slower explicit path, so request weights only when you need them. is_causal=True with a non-square mask is upper-left aligned. The head count must divide the embedding width.
• When not to use it. For short, fixed-size inputs (small images, tabular rows) a convolution or plain linear layer is cheaper and often as good; attention pays off when long-range, content-dependent mixing matters.

Check yourself

Connections

References

• Vaswani et al. Attention Is All You Need. NeurIPS, 2017 (arXiv 1706.03762).
• Bahdanau, Cho and Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR, 2015 (arXiv 1409.0473; the first attention mechanism).
• Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. 2019 (arXiv 1911.02150; multi-query attention).
• Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP, 2023 (arXiv 2305.13245).
• Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022 (arXiv 2205.14135); Dao. FlashAttention-2. 2023 (arXiv 2307.08691).
• Dong et al. Flex Attention: A Programming Model for Generating Optimized Attention Kernels. 2024 (arXiv 2412.05496).
• PyTorch v2.15.0: MultiheadAttention, scaled_dot_product_attention, torch/nn/attention/bias.py and flex_attention.