2.3 The Attention Mechanism From Scratch¶

If you remember one equation from this entire book, make it this one:

Every GPT, every Llama, every Claude, every Gemini — at its computational heart — is this line, stacked and repeated a few hundred times. The Multi-Layer Perceptrons (MLPs) store knowledge; the normalization layers keep gradients sane; the positional encodings tell the model where tokens are. But attention is the only operation in the Transformer that lets one token’s representation depend on another token’s representation. It is the mechanism for mixing information across the sequence. Remove it and you have an expensive per-token MLP that can never relate “it” to the noun three sentences ago.

This chapter builds that one line from the ground up. We will derive it as differentiable soft retrieval — a database lookup that returns a smooth blend of values rather than a single hard match — and we will see why every piece is there: why we project into queries, keys, and values; why we take a dot product; why we divide by \(\sqrt{d_k}\) (a detail that looks cosmetic but actually controls whether your network trains at all); why softmax and not something else; and how masking turns the same machinery into either a bidirectional encoder or a causal language model.

We assume you are comfortable with Linear Algebra for Deep Learning (matrix multiplication, transposes, the dot product as similarity) and with the softmax and cross-entropy from Probability, Statistics & Information Theory. We will write everything twice: once in pure NumPy so you can see the arithmetic with no framework magic, and once in PyTorch with autograd so it is actually usable. This is the crown-jewel chapter; we will take our time.

From Hard Retrieval to Soft Retrieval¶

The cleanest way to understand attention is to start with a thing you already know: a dictionary lookup, or a key–value store like Redis or a Python dict.

A hard key–value store holds a set of (key, value) pairs. You issue a query, the store finds the key that exactly matches your query, and returns the associated value:

store = {"paris": "france", "tokyo": "japan", "lima": "peru"}
store["tokyo"]   # -> "japan"   (exact match on the key "tokyo")

Three things are happening, and naming them precisely is the whole trick:

1. A query (\(q\)) — the thing you are looking up (“tokyo”).
2. A set of keys (\(k_1, \dots, k_n\)) — the addressable handles (“paris”, “tokyo”, “lima”).
3. A set of values (\(v_1, \dots, v_n\)) — the payloads you actually want back.

The lookup compares the query against every key, picks the matching one, and returns its value. In a Python dict the comparison is exact string equality and the selection is “take exactly one.” That makes it non-differentiable: there is no smooth notion of “the query is 70% like this key.” You cannot backpropagate through store["tokyo"].

Attention is what you get when you make every step of that lookup soft and differentiable:

• Keys and queries become vectors, and “match” becomes a continuous similarity score — the dot product \(q \cdot k\). A bigger dot product means a better match.
• Selection becomes a weighted average instead of a hard pick. We convert the scores into a probability distribution with softmax, then return the blend of all values weighted by how much each key matched. If the query is 70% like key 3 and 30% like key 7, we return \(0.7\,v_3 + 0.3\,v_7\).

This “soft retrieval” framing is not a loose analogy bolted on after the fact — it is literally what the math computes, and it pays off immediately. It tells you what each tensor means: a query says “here is what I am looking for,” a key says “here is what I contain, address me with this,” and a value says “if you decide to read me, here is the information I hand back.” It also tells you why attention can implement things like “copy the value from the most similar earlier token” (induction heads), “average over all tokens that mention the subject,” or “ignore everything except the punctuation.” Those are all just different shapes of the weight distribution.

Scaled Dot-Product Attention, One Query at a Time¶

Let us derive the formula for a single query \(q \in \mathbb{R}^{d_k}\) attending over \(n\) key–value pairs. Stack the keys as rows of a matrix \(K \in \mathbb{R}^{n \times d_k}\) and the values as rows of \(V \in \mathbb{R}^{n \times d_v}\). (We allow \(d_v \neq d_k\) in general, though in practice they are usually equal.)

Step 1 — Score every key against the query. The similarity between the query and key \(j\) is the dot product \(s_j = q \cdot k_j = \sum_{c=1}^{d_k} q_c\, k_{j,c}\). Computing all \(n\) scores at once is a single matrix–vector product:

The dot product is the natural similarity here for a reason: it is large and positive when \(q\) and \(k_j\) point in the same direction, near zero when they are orthogonal, and negative when they oppose. It rewards both alignment of direction and magnitude. (Compare to Embeddings & The Input Pipeline, where the same dot-product-as-similarity intuition underlies nearest-neighbor search over embeddings.)

Step 2 — Scale by \(\sqrt{d_k}\). We divide every score by \(\sqrt{d_k}\):

This is the “scaled” in scaled dot-product attention, and we devote the next section to why it matters. For now, take it as: it keeps the scores from blowing up as the head dimension \(d_k\) grows.

Step 3 — Softmax into attention weights. Turn the scaled scores into a probability distribution over the \(n\) positions:

The vector \(a \in \mathbb{R}^n\) is the attention distribution for this query: \(a_j\) is the fraction of “attention” the query pays to position \(j\). Softmax is the right choice because it is the smooth, differentiable relaxation of \(\arg\max\): it concentrates mass on the largest scores (recovering near-hard selection when one score dominates) while staying differentiable everywhere and always producing a valid probability vector. We unpack softmax’s numerical behavior — overflow, the log-sum-exp trick, the temperature induced by the scale — below.

Step 4 — Blend the values. The output is the attention-weighted sum of the value vectors:

That is the entire single-query mechanism. Notice the output lives in value-space (\(\mathbb{R}^{d_v}\)), is a convex combination of the value vectors (it lies inside their convex hull — it can never extrapolate beyond the values it is given), and is differentiable with respect to \(q\), every \(k_j\), and every \(v_j\). Information has flowed from the positions the query cared about into a single output vector.

Batching all queries at once¶

In practice we have not one query but one per position: \(n\) tokens each emit a query, so \(Q \in \mathbb{R}^{n \times d_k}\). We want every query to attend over every key independently. Stacking step 1 over all queries turns the matrix–vector product into a matrix–matrix product:

This \(n \times n\) object is the attention matrix (pre-softmax it is the “score matrix” or “logits”; post-softmax it is the “attention weights”). Row \(i\) holds the scores of query \(i\) against all keys. We softmax along each row (dimension -1, the key axis) so each query’s weights sum to 1, then multiply by \(V\):

Output row \(i\) is query \(i\)’s personalized blend of the values. We have arrived at the headline equation, and now every symbol has a meaning rooted in soft retrieval.

Why Divide by \(\sqrt{d_k}\)? The Variance Argument¶

The scale factor \(1/\sqrt{d_k}\) looks like a footnote. It is not. Get it wrong and softmax saturates, gradients vanish, and the network refuses to learn. Here is the argument from Attention Is All You Need (Vaswani et al., 2017), made fully explicit.

Treat the components of a query \(q\) and a key \(k\) as independent random variables, each with mean 0 and variance 1 — a reasonable approximation at initialization, where the projection weights are small-random and the inputs are normalized. The raw score is their dot product:

Each term \(q_c k_c\) is a product of two independent zero-mean, unit-variance variables, so it has mean \(\mathbb{E}[q_c k_c] = \mathbb{E}[q_c]\mathbb{E}[k_c] = 0\) and variance \(\operatorname{Var}(q_c k_c) = \mathbb{E}[q_c^2]\mathbb{E}[k_c^2] = 1 \cdot 1 = 1\). Summing \(d_k\) independent such terms, the variances add:

So the standard deviation of the raw score grows like \(\sqrt{d_k}\). For a typical head dimension \(d_k = 64\), scores have a standard deviation of \(8\). For \(d_k = 128\), about \(11\). These are large numbers to feed into a softmax.

Why is that bad? Softmax is extremely sensitive to the spread of its inputs. If one logit is, say, 24 (three standard deviations out at \(d_k=64\)) and the rest sit near 0, then \(e^{24} \approx 2.6\times10^{10}\) dwarfs every other term: the softmax output is essentially a one-hot vector. The distribution has saturated — it has effectively become a hard \(\arg\max\). And the gradient of a saturated softmax is nearly zero: the Jacobian of softmax is \(\operatorname{diag}(a) - a a^\top\), and when \(a\) is one-hot, that matrix is all zeros. No gradient flows back to the queries and keys. Training stalls.

Dividing by \(\sqrt{d_k}\) exactly cancels the growth:

The scaled scores have unit variance regardless of head dimension. Softmax receives inputs of a sane, dimension-independent spread, starts life close to uniform (gradients flow to all positions), and learns to sharpen as training progresses. The scale is a variance-preserving normalization that decouples the temperature of the attention distribution from an architectural hyperparameter (\(d_k\)) it has no business being coupled to. This is the same family of reasoning that motivates careful initialization and normalization throughout deep nets; see Neural Networks From Scratch: MLPs & Backprop.

A From-Scratch Implementation in NumPy¶

Let us write the whole thing with nothing but NumPy so there is no framework to hide behind. We will build a numerically stable softmax first (this matters — naïve softmax overflows), then attention on top of it, then verify on a tiny hand-checkable example.

import numpy as np

def softmax(x, axis=-1):
    """Numerically stable softmax along `axis`.

    The math identity softmax(x) = softmax(x - c) for any constant c lets us
    subtract the per-row max BEFORE exponentiating. This guarantees the largest
    exponent is exp(0) = 1, so we never overflow to +inf. Without this, a logit
    of, say, 1000 would make exp(1000) = inf and the whole row becomes NaN.
    """
    x_max = np.max(x, axis=axis, keepdims=True)      # largest logit per row
    e = np.exp(x - x_max)                             # in (0, 1], no overflow
    return e / np.sum(e, axis=axis, keepdims=True)    # normalize -> sums to 1


def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled dot-product attention, fully from scratch.

    Shapes (single, unbatched):
        Q : (n_q, d_k)   one query per output position
        K : (n_k, d_k)   one key per source position
        V : (n_k, d_v)   one value per source position
        mask : (n_q, n_k) or None. Entries that are True (or 1) are positions
               we are ALLOWED to attend to; False/0 positions are forbidden
               and get score -inf before softmax (so weight -> exactly 0).

    Returns:
        out     : (n_q, d_v)  the attended output, one vector per query
        weights : (n_q, n_k)  the attention distribution (rows sum to 1)
    """
    d_k = Q.shape[-1]

    # Step 1: raw similarity scores. scores[i, j] = q_i · k_j
    scores = Q @ K.T                                  # (n_q, n_k)

    # Step 2: scale so the variance of scores is ~1 regardless of d_k
    scores = scores / np.sqrt(d_k)

    # Step 3 (optional): masking. Set forbidden scores to -inf so that
    # exp(-inf) = 0 and those positions receive exactly zero attention weight.
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)

    # Step 4: softmax over the KEY axis (last axis) -> attention weights
    weights = softmax(scores, axis=-1)                # (n_q, n_k)

    # Step 5: weighted sum of values
    out = weights @ V                                 # (n_q, d_v)
    return out, weights

Now a tiny worked sanity check we can verify by hand. We make the keys axis-aligned (one-hot-ish) and the values easy-to-read tags, then issue a query that points almost exactly at the second key.

np.random.seed(0)

# Three keys, deliberately pointing along distinct axes, in d_k = 4.
K = np.array([
    [1.0, 0.0, 0.0, 0.0],   # key 0
    [0.0, 1.0, 0.0, 0.0],   # key 1
    [0.0, 0.0, 1.0, 0.0],   # key 2
])
# Values are easy to recognize: value j is the number (j+1) repeated.
V = np.array([
    [10.0, 10.0],           # value 0
    [20.0, 20.0],           # value 1
    [30.0, 30.0],           # value 2
])
# One query that aligns strongly with key 1 (the second axis).
Q = np.array([[0.1, 3.0, 0.1, 0.0]])

out, w = scaled_dot_product_attention(Q, K, V)
print("attention weights:", np.round(w, 3))   # ~ [[0.16 0.681 0.16]]
print("output:", np.round(out, 2))            # ~ [[20. 20.]] -> mostly value 1
print("rows sum to 1:", np.allclose(w.sum(axis=-1), 1.0))  # True

The query points mostly along axis 1, so it scores highest against key 1, softmax concentrates ~0.68 of its mass there, and the output is pulled toward value 1 ([20, 20]) — a soft retrieval of “the value whose key best matches my query.” Crank the query’s second component up to, say, 30 and the weight on key 1 approaches 1.0 and the output approaches exactly [20, 20]: hard retrieval recovered in the limit. That single experiment is the intuition for everything attention does.

Causal & Padding Masks: Same Machinery, Different Information Flow¶

The mask argument we slipped into the implementation is what makes attention flexible enough to serve as both a bidirectional encoder and an autoregressive decoder. A mask is just an additive bias on the score matrix — set forbidden entries to \(-\infty\) before softmax so they receive exactly zero weight afterward. Two masks matter in practice.

The causal (autoregressive) mask¶

A language model predicts token \(t+1\) from tokens \(1, \dots, t\). During training we feed the whole sequence at once for efficiency, but we must forbid each position from “peeking” at future tokens — otherwise the model could trivially copy the answer and learn nothing. We enforce this by allowing query \(i\) to attend only to keys \(j \le i\): a lower-triangular mask.

This single triangular mask is the entire reason a decoder-only Transformer can be trained with one forward pass yet still behave autoregressively at inference. It also makes the KV cache possible: because token \(i\)’s output never depends on future keys/values, those keys and values are fixed once computed and can be cached and reused at each decode step — the foundation of fast generation, covered in The Anatomy of LLM Inference: Prefill, Decode & The KV Cache. For the variants and trade-offs of where causal masking is or isn’t applied (prefix-LM, encoder–decoder), see Architecture Variants: Encoder-Decoder, Decoder-Only & Prefix-LM.

The padding mask¶

To batch sequences of different lengths together, we pad short sequences with a special <pad> token up to a common length. Those pad positions are meaningless and must never contribute to or receive attention. A padding mask zeroes out every score whose key is a pad token, for every query. In a real model you combine the causal and padding masks by taking the logical AND (a position is attendable only if it is both non-future and non-pad).

Causal attention in NumPy¶

def causal_mask(n):
    """Lower-triangular boolean mask: entry (i, j) is True iff j <= i.
    True = allowed to attend. Shape (n, n)."""
    return np.tril(np.ones((n, n), dtype=bool))       # tril keeps j <= i

# Demonstrate that future tokens get exactly zero weight.
n, d_k, d_v = 4, 8, 8
np.random.seed(1)
Q = np.random.randn(n, d_k)
K = np.random.randn(n, d_k)
V = np.random.randn(n, d_v)

out, w = scaled_dot_product_attention(Q, K, V, mask=causal_mask(n))
# Upper triangle of the weight matrix (strictly future) must be all zeros:
print(np.round(w, 3))
print("no future leakage:", np.allclose(np.triu(w, k=1), 0.0))  # True
# Row 0 attends only to position 0, so its weight there must be exactly 1:
print("row 0 weight:", np.round(w[0], 3))   # [1. 0. 0. 0.]

The Same Thing in PyTorch, With Autograd and Batching¶

The NumPy version teaches the mechanism; the PyTorch version is what you would actually train. The two differences that matter in real code are (1) a batch and head dimension so we process (batch, heads, seq, d_k) tensors at once, and (2) autograd, so gradients flow back through the whole computation automatically. The arithmetic is byte-for-byte the same.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0,
                                 is_causal=False):
    """From-scratch SDPA matching torch.nn.functional.scaled_dot_product_attention.

    Tensor shapes (the canonical Transformer layout):
        q : (B, H, Lq, d_k)   B=batch, H=heads, Lq=query length
        k : (B, H, Lk, d_k)
        v : (B, H, Lk, d_v)
        attn_mask : broadcastable to (B, H, Lq, Lk). Either a boolean mask
                    (True = keep) or an additive float mask (-inf = forbid).
    Returns:
        out : (B, H, Lq, d_v)
    """
    d_k = q.size(-1)

    # (B,H,Lq,d_k) @ (B,H,d_k,Lk) -> (B,H,Lq,Lk). transpose(-2,-1) makes Kᵀ.
    scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)

    # Build / apply the causal mask if requested.
    if is_causal:
        Lq, Lk = q.size(-2), k.size(-2)
        # Lower-triangular: position i may see j <= i. Forbidden -> -inf.
        causal = torch.ones(Lq, Lk, dtype=torch.bool, device=q.device).tril()
        scores = scores.masked_fill(~causal, float("-inf"))

    # Apply an explicit mask (padding, or arbitrary attention bias).
    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            scores = scores.masked_fill(~attn_mask, float("-inf"))
        else:
            scores = scores + attn_mask          # additive bias (e.g. ALiBi)

    # Softmax over the KEY axis (last dim). Cast to fp32 for a stable softmax
    # even when q/k/v are bf16 — a standard trick in production kernels.
    weights = torch.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)

    if dropout_p > 0.0:
        weights = F.dropout(weights, p=dropout_p)

    return torch.matmul(weights, v)              # (B,H,Lq,d_v)

A complete, trainable self-attention module¶

Now we wrap the kernel in a module that creates the queries, keys, and values from a single input via learned projections — i.e. genuine self-attention. This is single-head; Multi-Head Attention, MQA, GQA & MLA generalizes it to many heads (and to the memory-saving GQA/MLA variants that modern models use).

import torch.nn as nn

class SelfAttention(nn.Module):
    """Single-head self-attention: x -> (Q,K,V) -> attention -> output projection."""

    def __init__(self, d_model, d_k=None, causal=False, dropout=0.0):
        super().__init__()
        d_k = d_k or d_model
        self.causal = causal
        self.dropout_p = dropout
        # One learned linear map per role. No bias is the modern default.
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_k, bias=False)
        self.W_o = nn.Linear(d_k, d_model, bias=False)   # mix attended info back

    def forward(self, x, attn_mask=None):
        # x : (B, L, d_model). Project into the three roles.
        q = self.W_q(x)                          # (B, L, d_k): "what I seek"
        k = self.W_k(x)                          # (B, L, d_k): "how to address me"
        v = self.W_v(x)                          # (B, L, d_k): "what I return"

        # Add a singleton head dim so we can reuse the (B,H,L,d) kernel with H=1.
        q, k, v = (t.unsqueeze(1) for t in (q, k, v))    # (B, 1, L, d_k)

        attended = scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=self.causal,
        )                                         # (B, 1, L, d_k)

        attended = attended.squeeze(1)            # (B, L, d_k)
        return self.W_o(attended)                 # (B, L, d_model)


# Smoke test: shapes, gradient flow, and causality.
torch.manual_seed(0)
B, L, d_model = 2, 5, 16
x = torch.randn(B, L, d_model, requires_grad=True)
attn = SelfAttention(d_model, causal=True)

y = attn(x)
print("output shape:", y.shape)                  # (2, 5, 16)

y.sum().backward()                               # autograd through the whole op
print("grad flows to input:", x.grad is not None and x.grad.abs().sum() > 0)

Two things are worth emphasizing about this module. First, the output projection \(W_O\) is not decoration: after the attention step each position holds a blend of value vectors, and \(W_O\) is what lets the model recombine and reshape that blend before it re-enters the residual stream. (With multiple heads, \(W_O\) also mixes information across heads.) Second, notice that we never wrote a backward pass. Because every operation — matmul, softmax, masked_fill, the linear layers — is a differentiable primitive with a registered gradient, PyTorch’s autograd assembles the full backward pass for us; see Automatic Differentiation & PyTorch Internals.

Validating against PyTorch’s built-in kernel¶

A good from-scratch implementation should agree with the library to floating-point tolerance. PyTorch ships a fused, memory-efficient scaled_dot_product_attention (it dispatches to a FlashAttention-style kernel under the hood). Let us check our hand-rolled version matches it.

B, H, L, d = 2, 4, 7, 32
q = torch.randn(B, H, L, d)
k = torch.randn(B, H, L, d)
v = torch.randn(B, H, L, d)

ours = scaled_dot_product_attention(q, k, v, is_causal=True)
ref  = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # PyTorch builtin

print("max abs diff:", (ours - ref).abs().max().item())   # ~1e-6, fp32 noise
assert torch.allclose(ours, ref, atol=1e-5)

The built-in does the exact same math — it just never materializes the full \(L \times L\) score matrix in slow high-bandwidth memory, computing softmax incrementally instead. That IO-aware reformulation is the heart of FlashAttention; we devote a whole chapter to it in FlashAttention I: IO-Awareness & The Online Softmax. For everyday work, prefer the fused builtin; reach for the from-scratch version to understand, debug, or modify the mechanism.

Which kernel actually runs: the open-source attention layer¶

F.scaled_dot_product_attention is a dispatcher, not one kernel. At call time PyTorch chooses among several backends — a FlashAttention backend, a memory-efficient (xFormers-derived) backend, a cuDNN fused backend, and a pure-PyTorch MATH fallback that does materialize the \(n^2\) matrix — based on dtype, head dimension, mask type, and hardware. The surprise in practice is silent fallback: pass an arbitrary float attn_mask, an unsupported head dim, or fp32 inputs, and you quietly drop to MATH and pay the quadratic memory you thought you had avoided. Pin and inspect the choice rather than hoping:

import torch
from torch.nn.attention import sdpa_kernel, SDPBackend

# Restrict the dispatcher to an explicit preference list. If none of the listed
# backends can service the call, SDPA raises instead of silently falling back.
prefer = ([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]
          if torch.cuda.is_available() else [SDPBackend.MATH])

with sdpa_kernel(prefer):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)

Three more pieces of the real ecosystem sit around this call, and you will meet all of them again later in the book:

• torch.nn.attention.flex_attention (PyTorch ≥ 2.5) compiles a user-written score_mod — an arbitrary function of (score, b, h, q_idx, kv_idx) — and a block_mask predicate into a FlashAttention-class kernel. It is how you get ALiBi biases, Gemma-style logit soft-capping, sliding-window attention, or document-packed block-causal masks without leaving the fused fast path (compile it; uncompiled, it materializes scores).
• Dao-AILab/flash-attention is the reference CUDA implementation. Beyond flash_attn_func, its flash_attn_varlen_func takes a cu_seqlens array of cumulative document lengths and treats a packed buffer as a ragged batch — no mask, no padding waste. That is what Megatron-LM and TRL’s packing path call; see Chat Templates, Data Formatting & Sequence Packing.
• HuggingFace Transformers exposes the same choice as a one-line flag: AutoModelForCausalLM.from_pretrained(..., attn_implementation="sdpa"), with "eager" (the readable, mask-materializing path — the one to use when you want to see attention weights), "flash_attention_2", and "flex_attention" as the alternatives.

For serving, the equivalent layer is a paged kernel that reads KV from non-contiguous blocks — vLLM’s PagedAttention and SGLang’s RadixAttention, covered in PagedAttention & KV-Cache Memory Management.

Cost, Complexity & What the Gradients Look Like¶

Two practical facts about attention dominate every systems decision downstream: its quadratic cost and the shape of its gradient.

The quadratic bottleneck¶

For a sequence of length \(n\), head dimension \(d\), the score matrix \(QK^\top\) is \(n \times n\) and costs \(\mathcal{O}(n^2 d)\) multiply-adds; the \(AV\) product costs another \(\mathcal{O}(n^2 d)\). Both compute and the memory to store the attention matrix scale as \(\mathcal{O}(n^2)\) in the sequence length. Double the context and you quadruple the attention cost.

The softmax Jacobian and why values get clean gradients¶

Where do gradients go when we backprop through attention? The output is \(o_i = \sum_j a_{ij} v_j\). By the product rule, gradient flows along two routes: through the values \(v_j\) (weighted by the attention \(a_{ij}\)) and through the weights \(a_{ij}\) (and from there through softmax back into the scores, queries, and keys).

The value path is beautifully simple: \(\partial o_i / \partial v_j = a_{ij}\). A value vector receives gradient in proportion to how much it was attended to — attended values get strong learning signal, ignored values get little. The query/key path runs through the softmax Jacobian. For a single row with weights \(a\), the Jacobian of softmax is

This is where the \(\sqrt{d_k}\) scaling pays off a second time. When \(a\) is near-uniform (early training, thanks to scaling) the Jacobian has healthy non-zero entries and gradients propagate to queries and keys. When \(a\) is saturated (one-hot), \(J \to 0\) and the query/key gradients vanish — the failure mode the scale factor exists to prevent. The scale is not just numerically tidy at the forward pass; it keeps the backward pass alive.

The backward pass, written out¶

Autograd will assemble this for you, but you should be able to write it by hand: it is four matrix products, it is what a fused kernel actually implements, and you cannot read the FlashAttention backward pass without it. Let \(S = QK^\top/\sqrt{d_k}\), \(A = \operatorname{softmax}(S)\) row-wise, \(O = AV\), and let \(dO = \partial \mathcal{L}/\partial O\) be the gradient arriving from above. Then

where \(\odot\) is elementwise product and \(D \in \mathbb{R}^{n \times 1}\) is the row-wise sum \(D_i = \sum_j A_{ij} (dA)_{ij}\), broadcast across the row. The \(dS\) line is just the softmax Jacobian \(\operatorname{diag}(a) - aa^\top\) applied row by row, written without ever forming the \(n \times n \times n\) object. The useful identity is that \(D_i = \sum_j A_{ij}(dO_i \cdot v_j) = dO_i \cdot o_i\): the correction term is a row-wise dot product of the output with its own incoming gradient, so it can be computed from \(dO\) and \(O\) alone, in one pass, without \(dA\). That is precisely the quantity FlashAttention’s backward kernel precomputes and calls \(D\) — the trick that lets it recompute \(A\) on the fly instead of storing it; see FlashAttention I: IO-Awareness & The Online Softmax.

torch.manual_seed(0)
L, d = 6, 8
Q = torch.randn(L, d, requires_grad=True)
K = torch.randn(L, d, requires_grad=True)
V = torch.randn(L, d, requires_grad=True)

S = (Q @ K.T) / d**0.5
A = torch.softmax(S, dim=-1)
O = A @ V
dO = torch.randn(L, d)          # pretend gradient arriving from the layer above
O.backward(dO)                  # autograd fills Q.grad, K.grad, V.grad

with torch.no_grad():           # the same three gradients, by hand
    dV = A.T @ dO
    dA = dO @ V.T
    D  = (dO * O).sum(dim=-1, keepdim=True)   # == rowsum(dA * A); FlashAttention's D
    dS = A * (dA - D)
    dQ = (dS @ K) / d**0.5
    dK = (dS.T @ Q) / d**0.5

for name, manual, auto in [("dQ", dQ, Q.grad), ("dK", dK, K.grad), ("dV", dV, V.grad)]:
    print(name, "max abs err vs autograd:", (manual - auto).abs().max().item())  # ~1e-7

Two things fall straight out of these formulas. First, \(dV = A^\top dO\) confirms the earlier claim quantitatively: value \(j\) collects \(\sum_i A_{ij}\, dO_i\), i.e. gradient in proportion to how much each query attended to it. Second, notice that the \(1/\sqrt{d_k}\) appears in \(dQ\) and \(dK\) but not in \(dV\): the scale factor damps the query/key gradient path only, so it also sets the relative learning rate of \(W_Q, W_K\) against \(W_V\) — another reason the “obvious” alternative of dividing by \(d_k\) is not benign.

Putting It All Together¶

We started with a dict lookup and ended with a trainable, batched, masked, autograd-ready attention layer — and at no point did we wave our hands. Every component earned its place: queries/keys/values are the three roles a token plays in a soft retrieval; the dot product is the similarity; \(\sqrt{d_k}\) keeps that similarity’s variance dimension-independent so softmax stays soft and gradients stay alive; softmax is the differentiable relaxation of “pick the best match” into “blend by match quality”; the attention matrix is the explicit, inspectable record of who-attends-to-whom; and masking is a single additive bias that reshapes information flow into causal, padding-aware, or arbitrary patterns.

From here the path forks in three directions, all of which build directly on this chapter. Up the stack: replicate this kernel across many heads, inject position information (since bare attention is permutation-equivariant — it has no inherent notion of order), wrap it in a Transformer block with norms and residuals and an MLP, and stack the blocks into a full GPT. Down to the metal: make this kernel fast and memory-frugal with FlashAttention and serve it with a KV cache. Sideways: question whether the \(n^2\) cost is fundamental at all, in SSMs, Mamba & linear attention. And if you want to see this exact kernel inside a model you will actually train end to end — with GQA, QK-norm, RoPE, a document-packed mask and a KV cache wrapped around it — that is The Stack-100M Architecture in the capstone. Whichever way you go, the equation at the top of this page is the thing you are scaling, accelerating, or replacing.

Further reading¶

• Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin — Attention Is All You Need (2017). The original Transformer paper; introduces scaled dot-product and multi-head attention and the \(\sqrt{d_k}\) scaling.
• Bahdanau, Cho, Bengio — Neural Machine Translation by Jointly Learning to Align and Translate (2014). The additive-attention precursor that first framed attention as differentiable soft alignment.
• Dao, Fu, Ermon, Rudra, Ré — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022). How to compute exactly this operation without materializing the \(n^2\) matrix.
• Alammar — The Illustrated Transformer (blog). A widely used visual walkthrough of Q/K/V and the attention computation.
• Karpathy — nanoGPT (code repository) and Let’s build GPT (video). A minimal, readable from-scratch implementation of causal self-attention inside a full GPT.
• Elhage et al. (Anthropic) — A Mathematical Framework for Transformer Circuits (2021). Reads attention as soft retrieval/information-movement and introduces QK and OV circuits and induction heads.

Exercises¶

1. In the soft-retrieval framing, a token’s representation is projected into three roles — query, key, and value — via learned matrices \(W_Q, W_K, W_V\). Explain in one or two sentences what each of the three roles “says,” and then explain why we need three separate projections rather than, say, using the raw input \(x_i\) directly as query, key, and value all at once.

2. The chapter states that an attention output is a convex combination of the value vectors, so it “can never extrapolate beyond the values it is given.” Suppose a single query attends over three values \(v_1 = [0, 0]\), \(v_2 = [10, 0]\), \(v_3 = [0, 10]\) with attention weights \(a = [0.5, 0.3, 0.2]\). Compute the output vector. Then argue briefly why no choice of softmax weights could ever produce the output \([12, 0]\).

3. (Quantitative — the scale factor.) Consider a query scoring against three keys, producing raw (un-scaled) dot-product scores \(s = [12, 6, 0]\) at head dimension \(d_k = 36\).

(a) Compute the softmax attention weights on the raw scores.

(b) Now apply the \(1/\sqrt{d_k}\) scaling and recompute the weights on the scaled scores.

© Comment on what the comparison shows about gradient flow, referencing the softmax Jacobian.

4. (Conceptual — masking.) The mask argument sets forbidden scores to \(-\infty\) before softmax. (a) Why \(-\infty\) specifically, rather than \(0\)? (b) The chapter warns that in float16/bfloat16 you should use a large finite negative number instead of true \(-\infty\). Describe the exact failure that true \(-\infty\) can cause, and when it arises.

5. (Implementation.) Temperature-controlled attention. Extend the NumPy scaled_dot_product_attention from the chapter to accept a temperature parameter \(\tau > 0\) that divides the scaled scores before softmax (so the effective scores are \(\frac{QK^\top}{\tau\sqrt{d_k}}\)). Then, using the chapter’s hand-check keys/values and query, describe qualitatively what happens to the attention weights as \(\tau \to 0^+\) and as \(\tau \to \infty\).

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None, temperature=1.0):
    """Scaled dot-product attention with an explicit softmax temperature.

    temperature (tau > 0) divides the scaled scores before softmax:
        effective scores = (Q @ K.T) / (sqrt(d_k) * tau)
    tau = 1.0 recovers the standard chapter behavior.
    """
    assert temperature > 0, "temperature must be positive"
    d_k = Q.shape[-1]

    scores = Q @ K.T                      # (n_q, n_k) raw similarities
    scores = scores / np.sqrt(d_k)        # standard sqrt(d_k) scaling
    scores = scores / temperature         # <-- the one new line

    if mask is not None:
        scores = np.where(mask, scores, -np.inf)

    weights = softmax(scores, axis=-1)    # reuse the chapter's stable softmax
    out = weights @ V
    return out, weights

for tau in (0.1, 1.0, 10.0):
    out, w = scaled_dot_product_attention(Q, K, V, temperature=tau)
    print(tau, np.round(w, 3), np.round(out, 2))