The Heart of Attention: Letting Patches Talk to Each Other

The intuition: every patch asks a question

In the previous guide we turned an image into a sequence of patch tokens — each token is a vector that summarises a small square of the picture. But a single patch, on its own, is almost blind. A patch showing a brown furry texture could be part of a dog, a horse, or a carpet; it cannot know which until it compares notes with the patches around it. The whole point of self-attention is to let every token do exactly that: enrich itself by pulling in information from the other tokens.

Picture a noisy cocktail party. You are standing in a crowd, and silently you ask one question: "Who here is relevant to me?" You don't listen to everyone equally — you tune in to the few people whose conversation matters to you, and you tune out the rest. Self-attention does the same for a token. Each token broadcasts a question, every other token offers an answer about how relevant it is, and the token then listens more to the relevant ones and less to the irrelevant ones.

Concretely, the output for each token is a weighted average of all the tokens. The weights are numbers between 0 and 1 that add up to 1, and they encode "how relevant." A weight near 1 means "listen closely to this token"; a weight near 0 means "basically ignore it." So a furry-texture patch can end up blending in information from a nearby ear patch and a snout patch, and thereby start to "realise" it is part of a dog. The token doesn't move or change position — it just rewrites its own contents as a relevance-weighted mixture of everyone else's contents.

Query, Key, Value: three roles for every token

To turn the cocktail-party intuition into mechanism, each token plays three roles at once. Think of searching a library. You walk in with a Query — a description of what you're looking for ("books about edges and corners"). Every book on the shelf carries a Key — a short label or index card that says what it's about, so it can be matched against queries. And every book also has a Value — its actual content, the thing you carry home once you've decided it's relevant. This is the query–key–value decomposition, and it is the gearbox inside self-attention.

Three linear projections turn the token embeddings X into queries, keys, and values.

Let's unpack every symbol. X is the matrix of token embeddings, with shape N×D: there are N tokens (rows), and each token is a D-dimensional vector (one row of X). W_Q, W_K, W_V are three separate learned projection matrices, each of shape D×d_k. "Learned" means their numbers are adjusted by training — they start random and gradient descent shapes them. Multiplying X (N×D) by W_Q (D×d_k) gives Q with shape N×d_k: each row of Q is one token's query vector. Likewise K and V are N×d_k, one key and one value per token. Concretely, in a standard ViT-Base you might have N = 197 tokens (196 patches + 1 class token), D = 768, and d_k = 64; then a single token's embedding (a 1×768 row) times W_Q (768×64) produces that token's 1×64 query. The matrices are the only thing learned here — they decide how a patch phrases its question, how it advertises itself, and what it offers up.

Scaled dot-product attention: the master formula

Now we assemble the engine. Step one is measuring relevance. How do we score how well a query matches a key? With a dot product. The dot product of two vectors is big and positive when they point the same way, near zero when they're perpendicular, and negative when they point opposite ways — think of two arrows: aligned arrows give a large number, crossed arrows give a small one. So computing the dot product of token i's query with token j's key gives a single number: "how relevant is j to i?" Doing this for every pair (i, j) at once is exactly the matrix product Q·Kᵀ, an N×N grid of all pairwise scores. This is the core of scaled dot-product attention, built on the query–key–value vectors we just made.

Scaled dot-product attention: score, scale, normalise, then blend.

Read the formula from the inside out. QKᵀ is the N×N matrix of all pairwise similarity scores; its entry in row i, column j is the dot product of token i's query with token j's key — "raw relevance of j to i." We then divide every score by √d_k, a fixed number (the square root of the key dimension); this is the scaling that keeps the numbers in a sensible range — the next section is devoted to why. Next, softmax is applied to each row independently: it exponentiates the scores and divides by their sum, turning a row of arbitrary real numbers into a row of positive weights that sum to exactly 1 — these are the attention weights, the "how much I listen to each token" probabilities. Finally we multiply that N×N weight matrix by V (N×d_k): each output row becomes a weighted average of all the value vectors, weighted by that token's attention. The result has shape N×d_k — one enriched vector per token. The softmax step itself is just:

Softmax turns a row of scores s into positive weights that sum to 1.

Let's run real numbers through it for one token, with d_k = 2 and three tokens. Suppose token 1's query is q₁ = [1, 0], and the three keys are k₁ = [1, 0], k₂ = [0, 1], k₃ = [1, 1]. The raw scores q₁·kⱼ are 1, 0, and 1. Divide by √2 ≈ 1.414 to get 0.707, 0, 0.707. Exponentiate: e^0.707 ≈ 2.03, e^0 = 1, e^0.707 ≈ 2.03, summing to ≈ 5.06; so the attention weights are 0.40, 0.20, 0.40 (and indeed they sum to 1). Now take the values v₁ = [1, 0], v₂ = [0, 1], v₃ = [1, 1]. The output for token 1 is 0.40·[1,0] + 0.20·[0,1] + 0.40·[1,1] = [0.80, 0.60]. Notice what happened: token 1 paid most attention to tokens 1 and 3 (its keys aligned with q₁) and little to token 2, and its new vector is a blend dominated by those values. That is one full pass of attention, by hand.

import numpy as np

def softmax(x, axis=-1):
    # subtract the row max for numerical stability (no change to result)
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def self_attention(X, W_Q, W_K, W_V):
    # X: (N, D) token embeddings; W_*: (D, d_k) learned projections
    Q = X @ W_Q                         # (N, d_k) what each token is looking for
    K = X @ W_K                         # (N, d_k) how each token advertises itself
    V = X @ W_V                         # (N, d_k) what each token will hand over
    d_k = Q.shape[-1]
    scores  = (Q @ K.T) / np.sqrt(d_k)  # (N, N) scaled pairwise similarity
    weights = softmax(scores, axis=-1)  # (N, N) each row sums to 1
    out     = weights @ V               # (N, d_k) relevance-weighted blend
    return out, weights

Why divide by √d_k?

The √d_k looks like a tiny detail, but leaving it out quietly breaks training. Here is the problem. A dot product is a sum of d_k product-terms. The more dimensions you sum over, the larger that sum tends to grow in magnitude. When the scores fed into softmax become very large, softmax saturates: it pushes almost all the weight onto the single biggest score and leaves the rest near zero — the distribution becomes "peaky." A peaky softmax is bad for two reasons: the token effectively listens to only one neighbour (throwing away the soft, blended averaging that makes attention powerful), and, worse, the gradient through a saturated softmax is almost zero, so the network can barely learn. Scaling the scores down keeps softmax in a gentle, responsive region.

See it with numbers. Suppose three raw, well-behaved scores are [2, 1, 0]. Softmax gives roughly [0.67, 0.24, 0.09] — a healthy spread where every token still gets heard. Now imagine the dimension is large, so the same pattern of scores comes out ten times bigger: [20, 10, 0]. Softmax of that is about [0.99995, 0.000045, 0.000000002] — essentially all the weight collapses onto the first token, and the other two vanish. The blend has degenerated into a hard pick. Dividing [20, 10, 0] by √d_k (here √100 = 10) brings us right back to [2, 1, 0] and restores the healthy spread. Same relevance ordering, but a usable, trainable distribution.

If query and key components have variance 1, the raw dot product has variance ≈ d_k.

Here is the reasoning in plain words. Suppose, as a clean assumption, that the components of a query q and a key k are independent random numbers with mean 0 and variance 1. The dot product q·k = Σ qᵢkᵢ is a sum of d_k terms. Each term qᵢkᵢ is a product of two independent mean-0, variance-1 numbers, which itself has variance 1. When you add up independent quantities, their variances add, so the total variance is about d_k — it grows linearly with the number of dimensions. Variance is the square of the typical size, so the typical magnitude of q·k is about √d_k. That's the culprit: bigger d_k ⇒ bigger scores. To cancel it, we divide by exactly √d_k; dividing a quantity by c divides its variance by c², so variance d_k becomes (√d_k)²·1/d_k = 1. The scores are pulled back to a typical size of 1 regardless of dimension. For a real number: with d_k = 64, raw dot products run around √64 = 8, and dividing by 8 brings them back to ≈ 1.

Seeing attention: reading attention maps

Remember the N×N weight matrix from the softmax step — every row is one token's set of "how much I listen to each other token" weights. Because those weights are just numbers between 0 and 1, we can draw them. Pick a token (a row), reshape its N weights back onto the grid of patch positions, and colour each patch by its weight: bright where attention is high, dark where it's low. That picture is an attention map — a literal visualisation of what a token, under self-attention, is paying attention to.

The most revealing row to look at is the class token's — the special token (from the previous guide) whose final vector is fed to the classifier. When you visualise the class token's attention, the bright regions tend to fall on the actual object: the dog's body, the bird against the sky, the cup on the table. This is the satisfying payoff of all the math — it suggests the network has, on its own, learned to gather evidence from the foreground object and to downweight the background, and we get a window into where it is "looking" when it makes a decision.

You now have the complete heart of the Transformer: tokens form queries, keys, and values; queries and keys are compared (scaled by √d_k) and softmaxed into weights; values are blended into a richer output; and those weights can be read back as a picture of what each token attends to. In the next guide we'll wrap many copies of this mechanism into multi-head attention — letting the model ask several different questions at once — and stack the result with positions and depth to build the full encoder.