The attention mechanism, worked out by hand

Attention is usually introduced with a formula and a promise that it will make sense later.

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:

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:

Q = X W_q      K = X W_k      V = X W_v

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?

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.

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

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:

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.

d_k = 64   ->  raw scores std ~ 8    ->  softmax saturates
               after / sqrt(64) = 8  ->  std ~ 1, healthy gradients

Step 3 — Softmax into attention weights

weights = softmax(scaled, axis=-1)      # each row sums to 1

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

output = weights @ V                    # (4, d_k)

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.

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.

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.

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.

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.

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.

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
• Embeddings explained mathematically
• CNNs from scratch in NumPy