How transformers actually attend — a weighted lookup you can calculate

An attention diagram is only useful if you can name what travels along its arrows. “Query,” “key” and “value” are not three different kinds of token; they are different roles played by projected vectors. Start with one query, two keys and two values, and compute the result before scaling up to a whole sequence.

A token is represented by a vector

A token is an identifier produced by a tokenizer. An embedding lookup maps it to an initial vector, and later layers produce contextual hidden vectors. Their dimensionality is an architecture choice, not necessarily 768 or 1024. Position information can be added to representations or enter the attention computation through other mechanisms, such as rotary transformations of queries and keys; it is not always simply a field already stored inside each token embedding.

Attention transforms representations by combining information from permitted positions. Calling the output a “better vector” presumes training success. The operation itself only defines a learned computation; it does not guarantee semantic understanding or improved task quality.

Attention, in one sentence

For each query, score the eligible keys, normalize those scores across keys, and blend the corresponding values.

“Eligible” matters. Bidirectional self-attention, causal self-attention, cross-attention and sparse masks do not all let every position look at every other position. The mask is part of the computation, not an annotation added after the output.

What Q, K, V actually mean

For a self-attention input $X\in\mathbb R^{T\times d_{model}}$, use learned projections:

With $W_Q,W_K\in\mathbb R^{d_{model}\times d_k}$ and $W_V\in\mathbb R^{d_{model}\times d_v}$, Q and K have shape (T,d_k) and V (T,d_v). The scores are

M is zero for allowed positions and negative infinity for excluded ones in the mathematical formulation. P is (T,T) and O (T,d_v). Normalize one row across its eligible keys, not one scalar pair in isolation. A row with no eligible keys needs an explicit policy; the formula is undefined there.

• Q determines how the current position scores potential information sources.
• K provides the vectors against which those queries are compared.
• V provides the content to be combined, not the matching score itself.

The lookup analogy is useful but limited: these are learned vectors, not human-readable database keys. In cross-attention, Q can come from a decoder with length $T_q$, while K and V come from an encoder with length $T_k$. Then P is (T_q,T_k) and O (T_q,d_v). The old blanket statement that all three always come from the same token vector was only a self-attention statement.

A numeric head: why a blend is not a hard selection

Take already-scaled logits (log 3, 0). Their exponentials are (3,1), so weights are (0.75,0.25). With values (2,0) and (0,4), the output is

The result is neither original value vector. If both logits are equal, the weights become (0.5,0.5) and output (1,2). Changing V while holding Q and K fixed changes the content output without changing these weights. Changing Q or K can change the weights.

These weights describe contribution in this particular computation. They are not a complete causal explanation of a network prediction: output projections, residual connections, other heads and later layers also matter.

Why divide by the square root?

The Transformer paper motivates scaling by assuming independent query/key components with mean zero and variance one. Under those simplifying assumptions, the dot product $\sum_{r=1}^{d_k}q_rk_r$ has variance $d_k$, and division by $\sqrt{d_k}$ returns its variance to one. Large score differences can make softmax saturate and reduce some gradients.

This is an initialization/statistical motivation, not a hard bound on learned logits. Correlations, learned norms and training can violate the assumptions; the division does not prevent arbitrary values from becoming large. Compute a numerically stable softmax by subtracting each finite row's maximum before exponentiating. That preserves the normalized weights because the same multiplicative exponential factor cancels from numerator and denominator.

A bounded implementation with failure checks

This program first verifies the two-value calculation, then uses a three-position self-attention example with identity projections. It tests row sums, causal exclusion, and the fact that changing a future input cannot change an earlier causal output in this isolated layer.

import numpy as np
 
def softmax_rows(scores):
    # Reject the undefined all-masked case instead of returning NaNs.
    if not np.all(np.isfinite(scores).any(axis=-1)):
        raise ValueError('each query needs a finite eligible key')
    shifted = scores - scores.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)
 
P = softmax_rows(np.array([[np.log(3.), 0.]]))
V = np.array([[2., 0.], [0., 4.]])
assert np.allclose(P, [[.75, .25]])
assert np.allclose(P @ V, [[1.5, 1.]])
 
X = np.array([[1., 0.], [0., 1.], [1., 1.]])
def attend(x, causal=False):
    q, k, v = x, x, x  # fixed identity projections for this toy only
    scores = q @ k.T / np.sqrt(q.shape[-1])
    if causal:
        allowed = np.tril(np.ones(scores.shape, dtype=bool))
        scores = np.where(allowed, scores, -np.inf)
    weights = softmax_rows(scores)
    return weights @ v, weights
 
full, weights = attend(X)
assert np.allclose(weights.sum(axis=-1), 1)
assert np.allclose(weights[2], [.24825508, .24825508, .50348984], atol=1e-8)
causal, masked_weights = attend(X, causal=True)
assert np.all(masked_weights[np.triu_indices(3, 1)] == 0)
assert np.allclose(causal[0], X[0])
changed = X.copy(); changed[2] = [20., -7.]
changed_out, _ = attend(changed, causal=True)
assert np.allclose(changed_out[:2], causal[:2])
# Without positional information/masks, reordering rows reorders outputs.
perm = np.array([2, 0, 1])
reordered, _ = attend(X[perm])
assert np.allclose(reordered, full[perm])
try:
    softmax_rows(np.array([[-np.inf, -np.inf]]))
except ValueError:
    pass
else:
    raise AssertionError('all-masked row must be rejected')
print('two-value blend:', (P @ V).tolist())
print('three-token weights:', weights.tolist())
print('causal independence, permutation and mask checks passed')

The permutation check explains why a model needs some way to represent order: unmasked self-attention without positional information is permutation-equivariant, not sequence-order aware. Causal masks themselves introduce an order constraint, so the same unrestricted permutation test should not be expected to hold for the causal computation.

This teaching function accepts finite logits plus negative infinity masks; it is not a general validator for arbitrary NaN or positive-infinity inputs. It also has no learned weights, dropout, padding integration, backward pass, cache or GPU kernel. The checks verify the stated mechanism rather than a complete transformer implementation.

Multi-head: several projections, then one output

For head h, form separate projected Q, K and V, compute an output of shape (T,d_v), then concatenate heads to (T,H*d_v) and multiply by an output projection to the model width. Implementations often combine the projections in larger matrix multiplications before reshaping; they need not run a Python loop per head.

Different projections allow different relationships to be represented. A head may exhibit a pattern such as local matching or long-distance dependence, but no training rule assigns “head 3 must represent subject–verb agreement.” Head count and dimensions vary by architecture; an 8–64-head range is not a definition. Some architectures also share key/value heads across query heads, changing parameter and cache accounting.

For a shape exercise, let T=3, H=2 and each head's value width be 4. Concatenation has shape (3,8). An output projection of shape (8,6) gives (3,6). Each head's normalized weights sum to one before any attention dropout, but the final projected result is not necessarily a convex combination of the original input vectors.

Context length: opportunity, cost and limits

For full attention with fixed head dimensions, there are $T_qT_k$ query–key pairs; self-attention has $T^2$. Doubling T makes the full score grid four times as large. A dense causal triangle has $T(T+1)/2$ eligible pairs, so the exact small-T count is not literally quadrupled, although its asymptotic growth is quadratic.

That statement concerns pairwise attention, not every part of an LLM. Projection and feed-forward work, memory traffic and kernel efficiency also contribute. During cached autoregressive decoding, one new query attends over accumulated keys; its current attention step is linear in the available key count for fixed dimensions, while cache size also grows with sequence length. Do not copy the full-prefill $T^2$ statement onto every decoding step.

FlashAttention changes IO and which intermediates are materialized; it does not turn full exact pairwise attention into a linear-arithmetic algorithm. Longer context can expose more relevant evidence, but distractors, position effects, task design and training determine whether it helps. A context-window limit is not a measured guarantee of retrieval or reasoning quality.

The earlier “learned groupBy” analogy is dropped: groupBy normally partitions items using discrete keys and aggregates within those groups. Attention instead computes query-dependent soft weights over eligible values. Weighted lookup is a closer first mental model, still subject to the limits above.

Exercises and solutions

1. With logits (0,0) and values (2,0), (0,4), calculate the output.
2. With one query and seven keys, each of width 8, and values of width 5, give score and output shapes.
3. A causal mask is applied after softmax without renormalizing. What goes wrong?
4. Why is dividing by √d not a numerical-overflow guarantee?
5. What part of the toy implementation establishes that future values cannot influence an earlier causal query?

Where to go next

Use the NumPy shape lesson to review batched matmul and axis semantics. Then study the full self-attention derivation for gradients and a larger failure-mode discussion. Attention can be used in recommendation models, but it is not synonymous with every ranking or retrieval algorithm.

Sources and execution boundary

• Attention Is All You Need, v7, §3.2: scaled dot-product attention, multi-head projections and encoder–decoder uses.
• FlashAttention, v2: IO-aware exact attention and reduced materialization.

The exact small NumPy program is executed locally. No transformer training, learned head interpretation, long-context quality benchmark or GPU speed measurement is claimed.