--- title: "The attention mechanism, worked out by hand" description: "Queries, keys and values with real numbers. Why the scaling factor is sqrt(d_k), what multi-head attention buys, and how causal masking works." url: https://articles.sythra.ai/articles/attention-mechanism-visualized slug: attention-mechanism-visualized author: "vaibhavkothari" author_url: https://articles.sythra.ai/writers/vaibhavkothari date_published: 2026-08-29T09:26:53.619Z date_modified: 2026-08-29T10:22:15.966Z topics: ["Machine Learning", "Ai", "Python", "Deep Learning"] reading_time_minutes: 9 publisher: "Sythra" publisher_url: https://sythra.ai access: free language: en --- # The attention mechanism, worked out by hand > Queries, keys and values with real numbers. Why the scaling factor is sqrt(d_k), what multi-head attention buys, and how causal masking works. Source: https://articles.sythra.ai/articles/attention-mechanism-visualized · Author: vaibhavkothari · Published: 2026-08-29 · Reading time: 9 min · Topics: Machine Learning, Ai, Python, Deep Learning Attention is usually introduced with a formula and a promise that it will make sense later. ```text Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V ``` Later rarely arrives. So let us do it with actual numbers on a four-word sentence, small enough to check by hand, and let the intuition come from the arithmetic rather than before it. ## The problem attention solves Take the sentence: **"The animal crossed the road because it was tired."** What does *it* refer to? A model processing tokens independently cannot know. It needs, when building a representation of *it*, to **look back and pull in information from "animal"**. That is attention: for each position, produce a weighted mixture of information from all positions, where the model learns the weights from the content itself. ## Three roles: query, key, value Every token gets projected into three different vectors, each with a distinct job: | Vector | Role | Analogy | |---|---|---| | **Query** (q) | What this token is looking for | Your search box | | **Key** (k) | What this token offers as a match target | A document's title | | **Value** (v) | What this token actually contributes | The document's contents | The separation is the clever part. What a token advertises (key) is not the same as what it delivers (value), and neither is the same as what it wants (query). Three learned projections: ```text Q = X W_q K = X W_k V = X W_v ``` ```python import numpy as np # 4 tokens, model dimension 8, attention dimension 4 X = np.random.randn(4, 8) # token embeddings W_q = np.random.randn(8, 4) W_k = np.random.randn(8, 4) W_v = np.random.randn(8, 4) Q = X @ W_q # (4, 4) K = X @ W_k # (4, 4) V = X @ W_v # (4, 4) ``` ## Step 1 — Scores: every query against every key The raw score is a dot product: how well does query *i* match key *j*? ```matrix # Q (4 tokens × 2 dims, simplified) [1.0] 0.5 0.2 1.4 0.9 0.1 0.3 0.8 × # K^T (2 dims × 4 tokens) [1.1] 0.3 0.8 0.2 [0.4] 1.2 0.1 0.9 = # scores (4 × 4) [1.30] 0.90 0.85 0.65 1.78 1.74 0.30 1.30 1.03 0.39 0.73 0.27 0.65 1.05 0.32 0.78 ``` The highlighted cell is `q_1 · k_1 = 1.0(1.1) + 0.5(0.4) = 1.30` — how much token 1 attends to token 1. Read the matrix by rows: **row `i` is how much token `i` attends to every other token.** Row 2 has its largest values in columns 1 and 2, so token 2 is mostly interested in tokens 1 and 2. ```python scores = Q @ K.T # (4, 4) ``` Note the shape. This matrix is `(seq_len, seq_len)` — every token against every token. That is where the famous quadratic cost of transformers comes from, and it is why context length is expensive. ## Step 2 — Scale by sqrt(d_k), and why ```python d_k = Q.shape[-1] scaled = scores / np.sqrt(d_k) ``` This is the step everyone glosses over, and there is a concrete reason for it. A dot product of two `d_k`-dimensional vectors with roughly unit-variance components sums `d_k` products. Variance adds, so the dot product has variance about `d_k` and standard deviation about `sqrt(d_k)`. As `d_k` grows, scores spread further apart. Feed a widely spread vector to softmax and it saturates — one value gets essentially all the probability: ```python def softmax(x, axis=-1): x = x - x.max(axis=axis, keepdims=True) # numerical stability e = np.exp(x) return e / e.sum(axis=axis, keepdims=True) small = np.array([1.0, 2.0, 3.0]) large = small * 8 # what d_k=64 would produce print(softmax(small)) # [0.09 0.24 0.67] soft, informative print(softmax(large)) # [0.00 0.00 1.00] saturated, one-hot ``` A saturated softmax is a disaster for training, not just for expressiveness: its gradients are almost exactly zero, so the layer stops learning. Dividing by `sqrt(d_k)` normalises the variance back to roughly 1 and keeps the distribution in the responsive region. ```text d_k = 64 -> raw scores std ~ 8 -> softmax saturates after / sqrt(64) = 8 -> std ~ 1, healthy gradients ``` ## Step 3 — Softmax into attention weights ```python weights = softmax(scaled, axis=-1) # each row sums to 1 ``` ```matrix # attention weights (rows sum to 1) [0.31] 0.26 0.25 0.18 0.29 0.28 0.13 0.30 0.35 0.20 0.28 0.17 0.22 0.31 0.19 0.28 ``` Now each row is a probability distribution: *how much of each token's value should flow into this position?* Row 1 says token 1 draws 31% from itself, 26% from token 2, and so on. This matrix is what people mean by "attention maps". It is directly interpretable — and it is why you can visualise which words a model looked at. ## Step 4 — Weighted sum of values ```python output = weights @ V # (4, d_k) ``` ```matrix # weights (4 × 4) [0.31] [0.26] [0.25] [0.18] 0.29 0.28 0.13 0.30 0.35 0.20 0.28 0.17 0.22 0.31 0.19 0.28 × # V (4 × 2) [2.0] 1.0 [0.5] 3.0 [1.5] 0.5 [0.0] 2.0 = # output (4 × 2) [1.13] 1.55 0.95 1.79 1.19 1.31 0.97 1.86 ``` Output row 1: `0.31(2.0) + 0.26(0.5) + 0.25(1.5) + 0.18(0.0) = 1.13`. That is the whole mechanism. Every output position is a content-weighted blend of every input position's value vector. ```python def attention(Q, K, V, mask=None): scores = Q @ K.T / np.sqrt(Q.shape[-1]) if mask is not None: scores = np.where(mask, scores, -np.inf) weights = softmax(scores, axis=-1) return weights @ V, weights ``` Nine lines, and it is the core of every transformer. ## Causal masking: no peeking at the future For language modelling, position 3 must not see positions 4 and 5 — otherwise the model trivially cheats at next-token prediction. The fix is to set the forbidden scores to `-inf` **before** softmax, so `exp(-inf) = 0` and they receive exactly zero weight. ```matrix # causal mask (1 = allowed) [1] 0 0 0 1 1 0 0 1 1 1 0 1 1 1 1 ``` ```python n = 4 mask = np.tril(np.ones((n, n), dtype=bool)) # lower triangular out, w = attention(Q, K, V, mask=mask) print(w.round(2)) # [[1. 0. 0. 0. ] token 1 sees only itself # [0.51 0.49 0. 0. ] token 2 sees 1, 2 # [0.38 0.22 0.40 0. ] token 3 sees 1, 2, 3 # [0.24 0.33 0.20 0.23]] token 4 sees everything ``` Masking before softmax rather than after matters: if you zeroed the weights afterwards, the rows would no longer sum to 1 and the output would be scaled inconsistently across positions. ## Multi-head attention: several questions at once One attention head produces one weighting. But a token might need to track several relationships simultaneously — the syntactic subject, the coreferent, the topic. Multi-head attention runs `h` independent attentions in parallel on **slices** of the dimensions, then concatenates. ```text d_model = 512, h = 8 -> each head works in d_k = 64 dims head 1: Q1 K1 V1 -> (seq, 64) --+ head 2: Q2 K2 V2 -> (seq, 64) | ... +--> concat (seq, 512) --> W_o --> (seq, 512) head 8: Q8 K8 V8 -> (seq, 64) --+ ``` The crucial point: **total compute is roughly unchanged.** You are not running 8× the work — you split 512 dimensions into 8 slices of 64. Multi-head attention buys diversity of attention patterns for approximately free. ```python def multi_head_attention(X, W_q, W_k, W_v, W_o, h=8): n, d_model = X.shape d_k = d_model // h Q = (X @ W_q).reshape(n, h, d_k).transpose(1, 0, 2) # (h, n, d_k) K = (X @ W_k).reshape(n, h, d_k).transpose(1, 0, 2) V = (X @ W_v).reshape(n, h, d_k).transpose(1, 0, 2) scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k) # (h, n, n) weights = softmax(scores, axis=-1) heads = weights @ V # (h, n, d_k) merged = heads.transpose(1, 0, 2).reshape(n, d_model) # concat heads return merged @ W_o, weights ``` Trained models really do develop specialised heads — some track positional offsets, some track syntactic dependencies, some match repeated tokens. That specialisation is emergent, not designed. ## The cost, precisely The `(seq_len, seq_len)` score matrix is the whole story of transformer scaling. | Sequence length | Score matrix cells | Relative cost | |---|---|---| | 512 | 262,144 | 1× | | 2,048 | 4,194,304 | 16× | | 8,192 | 67,108,864 | 256× | | 32,768 | 1,073,741,824 | 4,096× | Doubling context quadruples attention cost and memory. Every long-context technique — FlashAttention, sliding windows, sparse patterns, linear attention — is an attempt to avoid materialising that matrix. FlashAttention in particular does not change the maths at all; it computes the same result in tiles that fit in fast on-chip memory, which is why it is a pure speed and memory win with identical outputs. ```remember # Remember this Query -> what this token is looking for Key -> what a token advertises to be matched against Value -> what it actually contributes when matched --- Divide by `sqrt(d_k)` or softmax saturates and gradients vanish. The score matrix is `(seq, seq)` — that is the quadratic cost. ``` ## Frequently asked questions ### Why divide by the square root of d_k? Dot products of `d_k`-dimensional vectors have variance proportional to `d_k`, so scores spread out as dimension grows. Wide-spread scores saturate the softmax into a one-hot distribution with near-zero gradients, and the layer stops learning. Dividing by `sqrt(d_k)` restores unit variance. ### What is the difference between query, key and value? Query is what a token is looking for, key is what a token advertises to be matched against, value is what it actually contributes when matched. Keeping them separate lets a token be a strong match target for one kind of query while contributing something quite different. ### Why multiple attention heads? One head produces one attention distribution per token. Multiple heads let a token track several relationships at once — syntax, coreference, position. Because the heads split the existing dimensions rather than duplicating them, the extra capability is close to free. ### Is attention the same as the attention maps I see visualised? Yes — the visualisations are the post-softmax weight matrix. Row `i` shows how much token `i` drew from every other token. Treat them as informative rather than as a complete explanation: attention weights show where information flowed, not what was computed with it. ### Why is long context expensive? The score matrix is `(seq_len, seq_len)`, so cost and memory grow quadratically. Going from 2k to 8k tokens is 16× more attention work. FlashAttention computes the same result without materialising the full matrix, which removes the memory wall but not the quadratic compute. ## Next reading on Sythra Articles - [Build a transformer block from scratch in PyTorch](https://articles.sythra.ai/articles/transformer-from-scratch-pytorch) - [Embeddings explained mathematically](https://articles.sythra.ai/articles/embeddings-explained-mathematically) - [CNNs from scratch in NumPy](https://articles.sythra.ai/articles/cnn-from-scratch-numpy)