Self-attention is the operation that lets every token in a sequence look at every other token and pull in the information it needs. It is the core of the Transformer: a single matrix-multiply-heavy step that replaces the sequential state of an RNN with direct, content-based connections between all positions.
This page derives scaled dot-product attention, explains masks and multiple heads, implements it in PyTorch, and covers the quadratic cost and the techniques used to tame it.
Given input rows X ∈ Rn×d, three learned projections produce queries, keys and values: Q = XWQ, K = XWK, V = XWV. Then
Attention(Q, K, V) = softmax(QKT / √dk + M) V
The dot product qi·kj measures how relevant token j is to token i; softmax turns each row of scores into weights summing to 1; the output for token i is the weighted average of value vectors. The √dk scaling keeps the variance of the logits near 1 — without it, large dot products saturate softmax into near one-hot distributions with vanishing gradients.
The additive mask M holds 0 for allowed pairs and −∞ for forbidden ones, which become exactly zero weight after softmax. A causal mask (upper triangle) prevents a token from attending to the future; it is what makes decoder-only models valid next-token predictors. A padding mask hides pad tokens in batched variable-length input. Encoder-only models use only the padding mask, so attention is bidirectional.
Cross-attention is the same formula with queries from one sequence and keys/values from another — the decoder-to-encoder link in encoder-decoder models.
Instead of one attention over d dimensions, the model runs h heads over d/h dimensions each, concatenates them, and mixes with an output projection WO. Different heads can specialize (previous-token, syntactic dependency, coreference) at no extra FLOPs compared with a single wide head.
Decoder inference variants shrink the KV cache by sharing keys and values: multi-query attention (one K/V head for all query heads) and grouped-query attention (a K/V head per group of query heads), now common in large LLMs.
import torch, torch.nn as nn, torch.nn.functional as F, math
class MultiHeadSelfAttention(nn.Module):
def __init__(self, d_model=512, n_heads=8, causal=False):
super().__init__()
assert d_model % n_heads == 0
self.h, self.dk, self.causal = n_heads, d_model // n_heads, causal
self.qkv = nn.Linear(d_model, 3 * d_model)
self.out = nn.Linear(d_model, d_model)
def forward(self, x, pad_mask=None): # x: (B, T, d_model)
B, T, _ = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
# (B, T, d) -> (B, h, T, dk)
q, k, v = (t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v))
scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk) # (B, h, T, T)
if self.causal:
causal = torch.ones(T, T, dtype=torch.bool, device=x.device).triu(1)
scores = scores.masked_fill(causal, float("-inf"))
if pad_mask is not None: # (B, T) True = padding
scores = scores.masked_fill(pad_mask[:, None, None, :], float("-inf"))
attn = F.softmax(scores, dim=-1)
y = (attn @ v).transpose(1, 2).reshape(B, T, -1)
return self.out(y), attn
x = torch.randn(2, 10, 512)
y, attn = MultiHeadSelfAttention(causal=True)(x)
print(y.shape, attn.shape) # (2, 10, 512) (2, 8, 10, 10)
# In practice use the fused kernel (FlashAttention-backed when available):
# F.scaled_dot_product_attention(q, k, v, is_causal=True)The explicit version is useful for understanding and for extracting attention maps; production code should call F.scaled_dot_product_attention, which dispatches to memory-efficient fused kernels.
The score matrix is n × n per head, so time is O(n²d) and naive memory is O(n²). Doubling the context quadruples attention cost. Approaches:
O(n) rather than recomputing the prefix.Take three tokens with dk = 2. Let the query of token 1 be q = [1, 0], the keys k1 = [1, 0], k2 = [0, 1], k3 = [1, 1], and the values v1 = [1, 0], v2 = [0, 2], v3 = [1, 1].
Token 1 draws mostly on tokens 1 and 3, whose keys align with its query, and less on token 2. With a causal mask, token 1 could only see itself and its output would simply be v1. Scale up the vectors by 10 and the scores become 7.07, 0, 7.07, softmax gives nearly 0.5, 0, 0.5, and the gradient through the near-zero weight almost vanishes, which is the saturation that the √dk factor is meant to prevent as dimensions grow.
Self-attention is permutation-equivariant: shuffle the input tokens and the outputs shuffle the same way. Order must therefore be injected. The original Transformer added fixed sinusoidal encodings to the embeddings; BERT and GPT-2 used learned absolute position embeddings, which do not extrapolate beyond the trained length. Most current LLMs use rotary position embeddings (RoPE, Su et al., 2021), which rotate pairs of dimensions in q and k by an angle proportional to the position, so the dot product qi·kj depends on the relative offset i − j. Context-extension methods (position interpolation, NTK-aware scaling, YaRN) rescale those rotation frequencies. ALiBi (Press et al., 2021) instead adds a distance-proportional penalty directly to the scores. Encoder models for images (see Vision Transformers) use 2D variants of the same ideas.
During generation, each new token needs the keys and values of every previous token in every layer. Storing them is the KV cache, whose size is 2 × layers × kv_heads × head_dim × tokens × bytes. For a model shaped like Llama 2 7B (32 layers, 32 heads of dimension 128, 16-bit values), that is 2 × 32 × 32 × 128 × 2 bytes = 512 KiB per token, or 2 GiB for one 4,096-token sequence; multiply by the batch size and the cache, not the weights, becomes the memory limit for serving.
Multi-query attention (Shazeer, 2019) shares a single key/value head across all query heads, cutting the cache by the number of heads at some quality cost. Grouped-query attention (Ainslie et al., 2023) is the middle ground used by many open LLMs: query heads are split into groups that each share one K/V head, recovering most of MHA's quality with a several-fold smaller cache. Serving systems add paged cache management (vLLM's PagedAttention), cache quantisation, and prefix sharing across requests; see vLLM and quantization.
| Setting | Attention pattern | Example models |
|---|---|---|
| Text understanding | Bidirectional self-attention | BERT, RoBERTa, sentence embedders |
| Text generation | Causal self-attention with KV cache | GPT-style decoders |
| Translation, summarisation | Encoder self-, decoder causal self-, plus cross-attention | T5, BART |
| Images | Self-attention over patches, sometimes windowed | ViT, Swin |
| Graphs | Attention restricted to neighbours | GAT |
Compared with recurrence (LSTMs), attention connects any two positions in one step and parallelises across the sequence during training, at the price of quadratic cost in sequence length. Convolutions are cheaper but see only a local window per layer. State-space models such as Mamba scale linearly and carry a fixed-size state, which is attractive for very long sequences; hybrid architectures interleave a few attention layers with them, because exact retrieval of a specific earlier token remains a strength of attention.
Further reading: Bahdanau, Cho & Bengio, "Neural Machine Translation by Jointly Learning to Align and Translate" (2014); Vaswani et al., "Attention Is All You Need" (2017, arXiv:1706.03762); Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (2019); Ainslie et al., "GQA" (2023); Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding" (2021); Dao et al., "FlashAttention" (2022); Jain & Wallace, "Attention is not Explanation" (2019).