Attention

TL;DR — Attention lets every position in a sequence look at every other position and pull in whatever is relevant, in a single matrix multiply. Each token emits a query, a key, and a value; the query–key dot product scores relevance, softmax turns scores into weights, and the weights take a weighted average of the values. Dividing by $\sqrt{d_k}$ keeps those scores at unit variance so softmax does not saturate. Multi-head attention runs several of these in parallel subspaces, so one layer can track several kinds of relationship at once. The cost is $O(n^2)$ time and memory in sequence length — which is why long context is expensive.

1. Why this exists

Before attention, sequences were handled by recurrent networks, which read a sequence one step at a time and carry a hidden state forward:

$$h_t=f(h_{t-1},x_t)$$

That design has three problems, and attention is essentially a response to all three.

It cannot be parallelised across time. Computing $h_t$ requires $h_{t-1}$, which requires $h_{t-2}$, and so on. A 1000-token sequence is 1000 sequential steps, and no amount of GPU width helps.

Long-range dependencies decay. The signal linking token 1 to token 1000 has to survive 999 multiplications. Gradients flowing back through that chain vanish or explode, and in practice the start of a long sequence is simply forgotten.

Everything is squeezed through a fixed-size state. The hidden state has the same width at step 5 and step 5000, so all history must be compressed into one vector of fixed dimension.

Attention attacks the first two directly. Between any two positions $i$ and $j$, the path length is one — they interact directly, no matter how far apart they are. And because attention is expressed as matrix products over the whole sequence at once, there is no sequential dependency at all within a layer.

Where this matters for you. If someone asks “why did transformers replace RNNs?”, the expected answer is not “better accuracy”. It is: parallelism over the sequence dimension (which is what made training on web-scale data feasible) and constant path length between distant tokens.

2. How it works

2.1 The lookup analogy

Attention is a soft dictionary lookup. Imagine a dictionary where each entry has a key and a value. A normal lookup takes a query, finds the entry whose key matches, and returns its value. Attention does the same thing, except instead of returning one exact match it returns a weighted blend of every value, weighted by how well each key matches the query.

• Query ($Q$) — what this token is looking for.
• Key ($K$) — what each token advertises about itself.
• Value ($V$) — the content each token actually contributes.

All three are learned linear projections of the same input, which is why this is called self-attention:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$

where $X\in {\mathbb{R}}^{n\times d_{model}}$ is the sequence of $n$ token vectors, and $W_Q,W_K,W_V\in {\mathbb{R}}^{d_{model}\times d_k}$ are the learned projections.

2.2 Scaled dot-product attention

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

Read it in three steps:

1. Score. $Q{K}^{\top }\in {\mathbb{R}}^{n\times n}$ — entry $(i,j)$ is the dot product of query $i$ with key $j$. A large positive value means “token $j$ is relevant to token $i$”.
2. Normalise. Softmax over each row turns raw scores into a probability distribution: how much attention token $i$ pays to each token $j$. Rows sum to 1.
3. Mix. Multiplying by $V$ produces, for each row, a weighted average of all value vectors — the token’s new representation.

flowchart LR
    X["X (n × d_model)"] --> WQ["× W_Q"]
    X --> WK["× W_K"]
    X --> WV["× W_V"]
    WQ --> Q["Q"]
    WK --> K["K"]
    WV --> V["V"]
    Q --> S["QKᵀ / √d_k"]
    K --> S
    S --> SM["softmax (row-wise)"]
    SM --> A["attention weights (n × n)"]
    V --> O["× V"]
    A --> O
    O --> OUT["output (n × d_k)"]

2.3 Why divide by $\sqrt{d_k}$

This is the single most-asked question about attention, and the answer is a variance argument — not a heuristic.

Assume the entries of $q$ and $k$ are independent, zero-mean, unit-variance. Their dot product is

$$q\cdot k=\sum _{i=1}^{d_k}{q}_i{k}_i$$

Each term $q_i{k}_i$ has mean $0$ and variance $\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]-(\mathbb{E}[q_i]\mathbb{E}[k_i]{)}^2=1$. The $d_k$ terms are independent, so the sum has variance $d_k$ and standard deviation $\sqrt{d_k}$.

So the magnitude of the scores grows with the head dimension. Feed scores with standard deviation $\sqrt{d_k}$ into a softmax and, for large $d_k$, the distribution becomes nearly one-hot: one weight near 1, the rest near 0. In that regime the softmax is saturated, its local gradient is close to zero, and the attention weights stop moving. Dividing by $\sqrt{d_k}$ rescales the scores back to unit variance, keeping the softmax in the region where it has a useful gradient.

This is also why the scaling is not $\sqrt{d_{model}}$. It is the size of the per-head projection $d_k$ that governs the variance of the dot product, so that is what the score is divided by.

2.4 Masking

A decoder must not look at future tokens — otherwise, predicting token $t$ would be trivial, since the answer is sitting in the input. The fix is to add $-\infty$ to the scores in the forbidden positions before the softmax:

$$S_{ij}\leftarrow \{\begin{array}{ll}{S}_{ij}& j\le i\\ -\infty & j>i\end{array}$$

Since $e^{-\infty }=0$, those positions get exactly zero attention weight and the rows still sum to 1. The result is a causal (lower-triangular) attention pattern:

$$\left(\begin{array}{cccc}■& & & \\ ■& ■& & \\ ■& ■& ■& \\ ■& ■& ■& ■\end{array}\right)$$

Padding masks are applied the same way, but to columns: padding tokens must not be attended to.

2.5 Multi-head attention

One attention operation computes a single weighted average. That is limiting: a token often needs to relate to several different things at once — the subject of the sentence, the most recent noun, a matching bracket. Multi-head attention runs $h$ attention operations in parallel on different learned projections, then concatenates and projects the results:

$$\mathrm{MultiHead}(X)=\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h)W_O$$ $$\text{where }{\mathrm{head}}_i=\mathrm{Attention}(X{W}_Q^{(i)},X{W}_K^{(i)},X{W}_V^{(i)})$$

The key design choice: each head works in a subspace of width $d_k=d_{model}/h$, so the total compute is roughly the same as a single head of full width. Multi-head is not a wider model; it is the same width partitioned into $h$ independent relationship-trackers.

flowchart TB
    X["X"] --> H1["head 1: Attention(d_k = d_model/h)"]
    X --> H2["head 2: Attention(d_k)"]
    X --> HD["⋯ head h"]
    H1 --> C["Concat (n × d_model)"]
    H2 --> C
    HD --> C
    C --> WO["× W_O"]
    WO --> OUT["output (n × d_model)"]

2.6 What it costs

Let $n$ be sequence length and $d$ the model width. For the standard implementation:

The quadratic memory term is the important one: it is what makes naive attention unusable at long context, and it is precisely what FlashAttention (see Modern Architecture Deltas) eliminates by never materialising the full matrix.

3. Implementation

A correct, readable from-scratch implementation:

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


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, T, _ = x.shape

        # (B, T, d_model) -> (B, n_heads, T, d_k)
        # .view splits d_model into n_heads x d_k; .transpose moves heads
        # next to the batch dimension so the matmuls below are batched.
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # (B, n_heads, T, T)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)

        if mask is not None:
            # mask: True where attention is ALLOWED.
            # Use the dtype's most-negative finite value, not -inf: a fully
            # masked row would be softmax over all -inf, i.e. 0/0 -> NaN.
            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)

        attn = self.dropout(F.softmax(scores, dim=-1))
        out = attn @ V                                   # (B, n_heads, T, d_k)

        # merge heads: (B, n_heads, T, d_k) -> (B, T, d_model)
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_o(out)


def causal_mask(T, device):
    """Lower-triangular boolean mask: position i may attend to j <= i."""
    return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))

Four details in that code are exactly what interviewers probe:

• .transpose(1, 2) moves the head dimension next to the batch dimension so that Q @ K.transpose(-2, -1) is a batched matrix multiply over (B, n_heads) — the heads are computed in parallel rather than in a Python loop.
• / math.sqrt(self.d_k) is applied to the scores, not to $Q$ or $K$ individually. Mathematically the two are equivalent; numerically, scaling the scores is what the derivation describes.
• Masking applies after scaling and before the softmax — applying it after the softmax would give a distribution that no longer sums to 1. Note the code uses torch.finfo(dtype).min rather than -inf; see the pitfalls below for why.
• .contiguous() before .view() — after transpose, the tensor is not contiguous in memory, and view requires it to be. Omitting this raises a runtime error (or silently produces a wrong reshape with older APIs).

The version you should actually ship

Since PyTorch 2.0, F.scaled_dot_product_attention dispatches to a fused kernel (FlashAttention / memory-efficient attention) and never materialises the $n\times n$ matrix:

out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

It is faster and uses less memory than the explicit version above, at the same result. Write the explicit version on a whiteboard; ship this one.

4. Numbers and trade-offs

The KV cache is the real memory cost at inference

During autoregressive decoding you generate one token at a time. Naively you would recompute keys and values for the entire prefix on every step. Instead you cache them — the KV cache — and the cache size is:

$$\text{bytes}=2\times n_{layers}\times n_{heads}\times d_{head}\times \text{seq\_len}\times \text{bytes\_per\_element}$$

The factor of 2 is for keys and values. For LLaMA-2 7B (32 layers, 32 heads, $d_{head}=128$) in fp16, that is 512 KiB per token (binary units throughout: 1 KiB = 1024 B):

$$2\times 32\times 32\times 128\times 2\text{ B}=512\text{ KiB/token}$$

At 4096 tokens that is 2 GiB per sequence — frequently more than the model weights once you batch. This is why long context is expensive, and why PagedAttention and continuous batching (see Serving) exist.

Head dimension

Per-head width $d_k$ is usually kept in the 64–128 range even as models grow, with head count increasing instead. Keeping $d_k$ fixed means the per-head compute stays constant and the $\sqrt{d_k}$ scaling stays in a well-behaved range, while the number of independent relationship-trackers grows with capacity.

Which axis you are optimising

Attention is a three-way trade-off and it is worth naming which axis a given decision moves:

5. Pitfalls

Forgetting the $\sqrt{d_k}$ scale. The model still trains — just worse, and the damage is invisible until you compare against a properly scaled baseline. Symptom: attention weights saturate to one-hot, gradients through the softmax collapse.

Masking with 0 instead of -inf. masked_fill(mask, 0) leaves masked positions at score 0, which softmax interprets as “moderately relevant” — not “forbidden”. The model silently leaks future information and reports a suspiciously good training loss that does not transfer to generation.

Using float('-inf'), or a value fp16 cannot represent. A fully-masked row (every position masked) becomes softmax over all $-\infty$, which is 0/0 → NaN, and NaN propagates through the whole model. The tempting fix is a large finite negative such as -1e9 — and in bf16 that works, because bf16 has fp32’s exponent range. In fp16 it does not: the largest finite magnitude is 65,504, so -1e9 is not representable and masked_fill rounds it to -inf, producing exactly the NaN you were trying to avoid. Use a representable value such as torch.finfo(torch.float16).min ($=-65504$) — which is what the code above does — or fall back to fp32 for the mask fill.

Forgetting .contiguous() before .view(). After transpose the tensor’s memory layout no longer matches its logical shape; view will raise, or in older code paths produce a scrambled reshape.

Confusing self-attention with cross-attention. In self-attention, $Q$, $K$, and $V$ all come from the same sequence. In cross-attention, $Q$ comes from the decoder and $K$, $V$ come from the encoder — that is how a seq2seq model reads its input. Getting this backwards is a common slip when describing an encoder–decoder stack.

Assuming nn.MultiheadAttention is a drop-in for the above. PyTorch’s built-in expects inputs in (seq_len, batch, d_model) by default (batch_first=False) and handles masking with a different convention. It is not a drop-in replacement, and mixing the two conventions is a classic source of silently wrong results.

6. Using it in a project

You almost never implement attention. In practice you control four things:

Context length. This is your dominant cost lever, because both attention time and KV-cache memory scale with it. Before increasing it, ask what actually needs to be in context — often the answer is “retrieved chunks”, which is the RAG pattern rather than a longer window.

Attention implementation. Prefer the fused kernel (F.scaled_dot_product_attention, or the equivalent in your serving stack). If your serving framework supports GQA, using a GQA checkpoint can cut KV-cache memory several-fold for a small quality cost — the single highest-leverage serving decision.

Head configuration, when fine-tuning or training small. If you are training a model from scratch, $d_k$ of 64–128 is the well-trodden range. If you are fine-tuning, head structure is fixed by the checkpoint; do not fight it.

Precision. Attention score computation is where fp16 overflows most often. If you see NaN losses early in training, this is the first place to look.

When not to touch any of this. For a 7B model at short context, attention is not your bottleneck — retrieval quality, prompt design, or the output parser usually is. Optimising attention before measuring is a classic premature optimisation.

7. How you would prove it worked

Attention has no product metric of its own; what you prove is that a design choice about attention matters. Three concrete experiments:

The scaling ablation. Train two identical small transformers on the same data, one with the $\sqrt{d_k}$ divisor and one without. Plot the loss curves. (example) The unscaled run typically shows a flat or barely-decreasing loss and much higher attention entropy variance, because the softmax has saturated. This is the experiment that turns “because the paper says so” into “because I watched it break”.

Attention entropy as a diagnostic. Log the entropy of the attention distribution per layer during training. A healthy model starts fairly uniform and sharpens; a model whose entropy collapses to near-zero at layer 0 early in training is usually saturated, and often indicates a scaling or learning-rate problem.

Long-context degradation. Evaluate the same task at increasing sequence lengths and plot quality against $n$. (example) Quality is typically flat up to the training length and then falls off sharply — the “lost in the middle” effect is partly an attention-pattern phenomenon. This is the measurement that justifies investing in retrieval rather than a bigger window.

Reporting template (example — replace with your own numbers). “We ablated the attention kernel: switching from the eager implementation to FlashAttention at 8k context cut peak memory by 40% (example) at identical output quality, and raised throughput 2.1× (example). We kept the eager path as a fallback for older GPUs.”

8. Interview drill

“Why divide by $\sqrt{d_k}$?” Strong answer: the dot products of independent unit-variance vectors have variance $d_k$; unnormalised, the softmax saturates and its gradient vanishes; dividing by $\sqrt{d_k}$ restores unit variance. Common wrong answer: “to keep the numbers small so it doesn’t overflow.” That is a consequence, not the reason — and it misses the gradient argument entirely.

“Why multiple heads instead of one big one?” Strong answer: the same total width split into subspaces lets the layer track several independent relationships simultaneously; it is not more compute. Follow-up: “how do the heads get combined?” → concatenated along the feature dimension, then projected by $W_O$. Heads are not averaged.

“Why is attention $O(n^2)$?” Because the score matrix is $n\times n$ — every token against every token. Follow-up: “what do you do about it?” → FlashAttention to remove the memory term; GQA to shrink the KV cache; sliding window or sparse attention to reduce the work; or retrieval, to avoid needing a long context in the first place.

“How does the causal mask work?” Strong answer: add $-\infty$ to scores above the diagonal before the softmax, so future positions receive exactly zero weight. Trap: saying “multiply by a triangular matrix afterwards” — that is a different operation and does not produce a normalised distribution.

“Self-attention vs cross-attention?” Strong answer: in self-attention $Q$, $K$, $V$ all come from the same sequence; in cross-attention $Q$ comes from the decoder and $K$, $V$ from the encoder. Follow-up: “which does a decoder-only model like GPT use?” → only self-attention; there is no separate encoder to cross-attend to.

“What is the difference between your implementation and nn.MultiheadAttention?” Strong answer: tensor layout conventions differ (batch_first), and the built-in handles masking differently; they are not interchangeable, and the most common bug is mixing the two conventions.

Next: Positional Encoding — attention as described here is permutation-invariant, so it cannot tell “dog bites man” from “man bites dog” without help.