Attention from first principles

Attention is the operation that lets a token look at other tokens and decide, per query, which ones matter. It is the reason transformers work, and it is the single most important thing to be able to derive rather than invoke, because every later chapter (the KV cache and its memory, GQA, the roofline of inference) is an argument about the cost of this specific computation. So I am going to build scaled dot-product attention from the ground up: what queries, keys, and values are, why the dot product is the right similarity, why we divide by $\sqrt{d}$ and not by $d$ or nothing, how the softmax turns scores into weights (and what its Jacobian is, because that is what backprop needs), how causal masking makes it autoregressive, and how multi-head attention buys representational capacity for almost free. Then I will implement it in raw PyTorch and check it, number for number, against a real Hugging Face layer.

Theory

The problem attention solves

A token embedding on its own is context-free: the vector for "bank" is the same whether the sentence is about a river or a loan. To build a contextual representation, each token needs to pull in information from the other tokens that are relevant to it, and "relevant" has to be learned and content-dependent, not fixed by position. Attention is a differentiable, content-addressed lookup that does exactly this. Every token emits a query (what am I looking for?), every token exposes a key (what do I offer as a match?) and a value (what information do I carry?), and each token's output is a weighted average of everyone's values, where the weights come from how well its query matches each key. It is a soft, learnable dictionary lookup: instead of retrieving the one value whose key matches exactly, you retrieve a blend, weighted by match quality.

Queries, keys, values as learned projections

Start from the input to an attention layer: a sequence of $L$ token vectors stacked into a matrix $X\in {\mathbb{R}}^{L\times d_{\text{model}}}$, one row per position. Attention does not use $X$ directly; it first projects it into three separate spaces with three learned weight matrices:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V,\text{\hspace{1em}(3.1)}$$

with $W_Q,W_K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$ and $W_V\in {\mathbb{R}}^{d_{\text{model}}\times d_v}$, giving $Q,K\in {\mathbb{R}}^{L\times d_k}$ and $V\in {\mathbb{R}}^{L\times d_v}$. The point of three separate projections is that the same token plays three different roles: the query space is tuned for asking, the key space for being found, the value space for what actually gets copied. These are the only learned parameters in attention (plus an output projection I will add shortly); everything else is fixed arithmetic.

Scaled dot-product attention, derived

The similarity between query $i$ and key $j$ is their dot product $q_i\cdot k_j={\sum }_{m=1}^{d_k}{q}_{im}{k}_{jm}$, which is large and positive when the two vectors point the same way and near zero when they are orthogonal. Stacking all pairs gives the score matrix $Q{K}^{\top }\in {\mathbb{R}}^{L\times L}$, entry $(i,j)$ being how much token $i$'s query matches token $j$'s key. We convert each query's row of scores into a probability distribution over the $L$ keys with a softmax, then use those probabilities to average the values. The whole operation is one line:

$$\mathrm{Attn}(Q,K,V)=\underset{A\in {\mathbb{R}}^{L\times L}}{\underbrace{\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)}}V.\text{\hspace{1em}(3.2)}$$

Equation (3.2) is the whole mechanism, and the row-stochastic matrix $A$ (each row sums to 1, all entries nonnegative) is the attention matrix: $A_{ij}$ is the fraction of its output that token $i$ pulls from token $j$'s value. The output is $AV\in {\mathbb{R}}^{L\times d_v}$, each row a convex combination of value vectors. The only piece not yet justified is that $\sqrt{d_k}$ in the denominator, and it is not cosmetic.

The softmax and its Jacobian

The softmax turns a row of scores $z\in {\mathbb{R}}^L$ into a probability vector $a\in {\mathbb{R}}^L$ via

$$a_i=\frac{e^{z_i}}{{\sum }_{j=1}^L{e}^{z_j}},\text{\hspace{1em}(3.6)}$$

which is nonnegative and sums to $1$ by construction. To train through attention, backprop needs the derivative of these probabilities with respect to the scores, and the softmax Jacobian is clean and worth deriving once, because it explains both why saturated softmaxes kill gradients and how the attention weights get their gradient.

Causal masking makes it autoregressive

For a language model, token $i$ must predict token $i+1$ using only tokens $1$ through $i$; letting it attend to future tokens would leak the answer. Attention enforces this with a causal mask applied to the scores before the softmax: set $S_{ij}=-\infty$ for all $j>i$, so that after the softmax $A_{ij}=0$ there. Since $e^{-\infty }=0$, the future tokens get exactly zero weight and drop out of every convex combination, while the surviving weights still renormalize to sum to $1$ over the allowed positions:

$$S_{ij}=\frac{q_i\cdot k_j}{\sqrt{d_k}}+M_{ij},M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}\text{\hspace{1em}(3.10)}$$

The mask is a fixed lower-triangular pattern, not learned. In practice frameworks add a large negative number (like the dtype's most-negative finite value) rather than literal $-\infty$ to avoid NaNs, but the effect is the same: future is unreachable. This lower-triangular structure is why the KV cache works at inference (each new token only ever attends backward, so past keys and values can be stored and reused), which is the entire subject of the KV-cache chapter in Part II.

Multi-head attention

A single attention head produces one set of weights, one "way of looking." But a token often needs to attend to several things at once for different reasons: the subject it agrees with, the clause it modifies, the earlier mention it refers to. Multi-head attention runs $h$ independent attention operations in parallel, each with its own small projections, then concatenates and mixes the results. Split the model dimension into $h$ heads of size $d_k=d_{\text{model}}/h$ each, and for head $t$:

$${\mathrm{head}}_t=\mathrm{Attn}(X{W}_Q^{(t)},X{W}_K^{(t)},X{W}_V^{(t)}),t=1,\dots ,h,\text{\hspace{1em}(3.11)}$$

then concatenate the $h$ outputs along the feature axis and apply a final output projection $W_O\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{model}}}$:

$$\mathrm{MHA}(X)=[{\mathrm{head}}_1\Vert {\mathrm{head}}_2\Vert \cdots \Vert {\mathrm{head}}_h]W_O.\text{\hspace{1em}(3.12)}$$

The elegance is in the bookkeeping: because each head is size $d_{\text{model}}/h$, the total parameter count of the $h$ heads equals that of a single full-width attention, so multi-head attention buys $h$ different views at essentially the same parameter and compute cost as one wide head. It is capacity for almost free, which is why every model uses it. In implementation the per-head projections of equation (3.11) are fused into single $d_{\text{model}}\times d_{\text{model}}$ matrices $W_Q,W_K,W_V$ and then reshaped into $h$ heads, so on disk you see one big projection per role, not $h$ small ones. When I get to GQA in the next chapter, the whole trick will be to keep $h$ separate query heads but share keys and values across groups of them, purely to shrink the KV cache; everything about why that is even coherent comes from equations (3.11)–(3.12).

Tooling

The tool that embodies attention is, at the bottom, torch.nn.functional.scaled_dot_product_attention (SDPA), the fused primitive that PyTorch exposes and that Hugging Face models call under the hood. It implements exactly equation (3.2) with the causal mask of (3.10), but it does not build the $L\times L$ matrix $A$ in memory when it can avoid it: on supported hardware it dispatches to a FlashAttention-style kernel that computes the softmax-weighted sum in tiles, streaming over the sequence and never materializing the full attention matrix. That matters enormously for the 16GB budget, because the naive $A\in {\mathbb{R}}^{L\times L}$ scales quadratically with sequence length and is the first thing to blow up at long context; the fused kernel keeps attention's memory linear in $L$ even though its compute stays quadratic. I will unpack that roofline in Part II, but the tooling point for now is that SDPA is the single call that a hand-written attention and a production model both ultimately reduce to, which is why I can verify one against the other.

On the Hugging Face side, a model's attention module (for example Qwen3Attention) is a thin wrapper: it holds the fused q_proj, k_proj, v_proj, and o_proj linear layers of equations (3.1) and (3.12), reshapes into heads, applies rotary position embeddings (next chapter) to $Q$ and $K$, and then calls SDPA. Because the wrapper is thin, I can pull the exact projection weights out of a real layer, run the arithmetic myself with the same weights, and expect a numerical match to within floating-point tolerance. That is the whole design of the lab.

Lab

The goal is to implement scaled dot-product multi-head attention in raw PyTorch, straight from equations (3.1)–(3.12), and verify it numerically against a real Qwen3 attention layer by feeding both the same input and the same weights and checking the outputs agree to floating-point tolerance. The artifact is a saved report of the max absolute difference, which is the evidence that "attention" is exactly the arithmetic derived above and nothing hidden.

uv init labs/attention
cd labs/attention
uv add torch transformers

"""Raw-PyTorch scaled dot-product attention, verified against a HF layer.

Artifact: artifacts/attention_check.json with the max abs difference.
"""
import json
import math
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

def manual_attention(x, layer, cfg):
    """Reproduce one attention layer's output from raw ops (eqs 3.1-3.12).

    x: (batch, seq, d_model). Uses the layer's own projection weights.
    Note: we skip RoPE here and disable it on the reference too, so this
    isolates the scaled-dot-product-attention math of this chapter. We do
    apply Qwen3's per-head q_norm/k_norm, because the reference layer applies
    them to Q and K before attention and the outputs will not match without it.
    """
    b, s, _ = x.shape
    n_q = cfg.num_attention_heads
    n_kv = cfg.num_key_value_heads
    d_head = cfg.head_dim if hasattr(cfg, "head_dim") else cfg.hidden_size // n_q

    # Qwen3 applies a per-head RMSNorm (q_norm/k_norm) over head_dim to Q and K
    # before attention; replicate it so the manual output matches the reference.
    q = layer.q_norm(layer.q_proj(x).view(b, s, n_q, d_head)).transpose(1, 2)   # (b, n_q, s, d)
    k = layer.k_norm(layer.k_proj(x).view(b, s, n_kv, d_head)).transpose(1, 2)  # (b, n_kv, s, d)
    v = layer.v_proj(x).view(b, s, n_kv, d_head).transpose(1, 2)

    # Expand KV heads to match query heads (GQA; trivial when n_kv == n_q).
    rep = n_q // n_kv
    k = k.repeat_interleave(rep, dim=1)
    v = v.repeat_interleave(rep, dim=1)

    scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_head)         # eq 3.2/3.5
    causal = torch.triu(torch.full((s, s), float("-inf")), diagonal=1)
    scores = scores + causal                                       # eq 3.10
    attn = torch.softmax(scores, dim=-1)                           # eq 3.6
    out = attn @ v                                                 # eq 3.2
    out = out.transpose(1, 2).contiguous().view(b, s, -1)
    return layer.o_proj(out)                                       # eq 3.12

def main() -> None:
    torch.manual_seed(0)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.float32, attn_implementation="eager"
    )
    model.eval()
    cfg = model.config
    layer = model.model.layers[0].self_attn

    b, s = 1, 12
    x = torch.randn(b, s, cfg.hidden_size)

    with torch.no_grad():
        mine = manual_attention(x, layer, cfg)

        # Reference: call the same layer but neutralize RoPE by passing
        # position_embeddings of (ones cos, zeros sin) so no rotation happens,
        # and pass the same causal mask the manual path applies so both are
        # autoregressive (q_norm/k_norm live inside the layer already).
        d_head = cfg.head_dim if hasattr(cfg, "head_dim") else cfg.hidden_size // cfg.num_attention_heads
        cos = torch.ones(b, s, d_head)
        sin = torch.zeros(b, s, d_head)
        mask = torch.triu(torch.full((s, s), float("-inf")), diagonal=1)
        ref = layer(x, position_embeddings=(cos, sin),
                    attention_mask=mask.view(1, 1, s, s))[0]

    max_abs = (mine - ref).abs().max().item()
    mean_abs = (mine - ref).abs().mean().item()
    report = {
        "model": MODEL,
        "seq_len": s,
        "max_abs_diff": max_abs,
        "mean_abs_diff": mean_abs,
        "match": max_abs < 1e-4,
    }
    (OUT / "attention_check.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))
    print(f"Artifact: {(OUT / 'attention_check.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python verify_attention.py

What you should see. The script writes artifacts/attention_check.json reporting a max absolute difference between my hand-rolled attention and the real Qwen3 layer that is down at the floating-point-noise floor, on the order of $10^{-5}$ to $10^{-6}$ in FP32, with "match": true. That tiny residual is not error in the derivation; it is the accumulated rounding of doing the same real-valued arithmetic in a slightly different operation order (equation 1.8 from the tensors chapter, made visible). The takeaway is concrete and load-bearing for the rest of the book: a production attention layer is exactly equations (3.1) through (3.12), same weights, same softmax, same $\sqrt{d_k}$, and once you have watched your own code reproduce it to five decimal places, attention stops being a black box and becomes a thing you can reason about the cost of. Record your exact max-diff (measured on the baseline machine, record value, date, driver), since it depends on dtype and the transformers version you pinned.