Attention & Transformers

01Intuition

Attention is a soft dictionary lookup.

A Python dictionary matches one key exactly. d["cat"] hands back a single value and ignores everything else it holds.

Attention asks a fuzzier question. Not "which key matches?" but "how similar is my query to every key?" It then returns a blend of all the values, weighted by that similarity. Three steps:

• Score. Compare the query against every key. One similarity number per key.
• Normalize. Push those scores through a softmax so they sum to one.
• Blend. Take the weighted average of every value, using those weights.

Nothing is discarded, and every token contributes by a different amount.

Here is the same idea in numbers. Three words each carry a two-number key that says how animal-like or object-like the word is. The word it sends out a query asking for something animal-like. The vectors are hand-picked to make the point; a trained model learns them.

import numpy as np

words = ["cat", "mat", "it"]
K = np.array([[1.0, 0.0],        # "cat": an animal
              [0.0, 1.0],        # "mat": an object
              [0.5, 0.5]])       # "it": not sure yet
V = K.copy()                     # here each value is just its key
q = np.array([2.0, 0.5])         # what "it" is looking for: animal-like

scores = K @ q / np.sqrt(2)      # how well each word matches the query
w = np.exp(scores) / np.exp(scores).sum()   # softmax: weights sum to 1
out = w @ V                      # a weighted average of the values

for word, s, weight in zip(words, scores, w):
    print(f"{word:>4}: score {s:5.2f}  weight {weight:5.2f}")
print("new vector for 'it':", np.round(out, 2))

That example is Jay Alammar's, from The Illustrated Transformer. Worth twenty minutes if diagrams land faster for you than equations.

Every self-attention layer does this for every token against every other token, all at once. Queries, keys, values, multiple heads and the $\sqrt{d_k}$ scaling are machinery that makes this one idea learnable and numerically well-behaved.

02Timeline

03Architecture: how the data flows

Every token's embedding is linearly projected three different ways: into a Query, a Key, and a Value vector. Three learned lenses on the same token.

• Query. "What am I looking for?"
• Key. "Here's what I contain."
• Value. "Here's what I'll give you if you attend to me."

Now wrap that in a block: a residual connection, a normalization layer, and a position-wise feed-forward network. Repeat that 12–96+ times and you have a Transformer stack:

That block, stacked, is the whole architecture. Here is how the original paper drew the full encoder-decoder picture: an encoder stack on the left, feeding a decoder stack on the right through cross-attention.

Reference figureThe full encoder-decoder Transformer stack. Source: Jay Alammar, The Illustrated Transformer.

Shapes, concretely:

• Input token ids: (batch, seq_len)
• After the embedding table: (batch, seq_len, d_model)
• After a decoder-only LM head: (batch, seq_len, vocab_size) logits, one distribution per position predicting the next token

Every attention and feed-forward layer preserves that middle shape. Attention mixes information across the seq_len axis but never changes it. Only the final LM head projects d_model → vocab_size.

Encoder and decoder stacks differ in exactly one line.

• Encoder (used in BERT). It is bidirectional, so every position sees the whole input at once.
• Decoder (used in GPT). It is causal, so each position sees only itself and earlier positions.

The decoder gets there by masking out future positions before the softmax, setting their scores to $-\infty$. Same mechanism, one masking triangle.

import numpy as np
rng = np.random.default_rng(0)

T, d_k = 4, 8                    # 4 tokens, 8-dim queries and keys
Q = rng.normal(size=(T, d_k))
K = rng.normal(size=(T, d_k))
V = rng.normal(size=(T, 3))      # values need not share d_k

scores = Q @ K.T / np.sqrt(d_k)  # (T, T): everyone against everyone
future = np.triu(np.ones((T, T), dtype=bool), k=1)
scores[future] = -np.inf         # token t may not read past t

m = scores.max(axis=1, keepdims=True)
W = np.exp(scores - m)
W /= W.sum(axis=1, keepdims=True)     # softmax, one row at a time
Z = W @ V

print("Q", Q.shape, "K.T", K.T.shape, "scores", scores.shape)
print("W", W.shape, "V", V.shape, "Z", Z.shape, "\n")
print(np.round(scores, 2), "\n")
print(np.round(W, 3))
print("row sums:", W.sum(axis=1))
print("Z[0] is exactly V[0]:", np.allclose(Z[0], V[0]))

04The equations

• Q, K, V matrices of shape (seq_len, d_k) — every row is one token's query/key/value vector, produced by three separate learned linear layers on the input embeddings.
• QKᵀ a (seq_len, seq_len) matrix of raw similarity scores — every token compared against every other token via dot product.
• √d_k scaling factor: without it, dot products of high-dimensional vectors grow large in magnitude, which pushes softmax into a region with near-zero gradients almost everywhere except one spike. Dividing by $\sqrt{d_k}$ keeps the variance of the scores roughly constant regardless of dimension, so training stays stable.
• softmax applied row-wise, turning each token's raw scores into a probability distribution over all tokens (including itself) that sums to 1.
• · V the weighted sum — multiplying those probabilities back into the value vectors gives, for each token, a blend of everyone else's information, proportioned by relevance.

import numpy as np
rng = np.random.default_rng(0)

def dots(d_k, n=200_000):        # n independent q.k scores
    q = rng.normal(size=(n, d_k))
    k = rng.normal(size=(n, d_k))
    return (q * k).sum(1)

for d_k in (8, 64, 512):
    s = dots(d_k)
    print(f"d_k={d_k:4d}  var(q.k)={s.var():7.1f}"
          f"   var after /sqrt(d_k)={s.var()/d_k:.2f}")

def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()

d_k = 512                        # one query against six keys
s = rng.normal(size=(6, d_k)) @ rng.normal(size=d_k)
p_raw, p_sc = softmax(s), softmax(s / np.sqrt(d_k))
print("\nraw scores:      ", np.round(s, 1))
print("unscaled softmax:", np.round(p_raw, 4))
print("scaled softmax:  ", np.round(p_sc, 4))
print(f"smallest weight: {p_raw.min():.1e} unscaled,"
      f" {p_sc.min():.1e} scaled")

Multi-head attention

Running one attention function gives the model exactly one notion of "relevance." In practice you want several. One head might learn to track subject-verb agreement, another coreference, another local syntax. Multi-head attention runs $h$ independent, smaller attention operations in parallel and concatenates the results:

If $d_{model}=512$ and $h=8$, each head works in $d_k = 64$ dimensions. The total compute is the same as one 512-dim head, and it buys 8 independent subspaces to specialize in, where a single head has to average everything together.

import torch
torch.manual_seed(0)

T, d_model, H = 5, 12, 3         # 5 tokens, 3 heads of width 4
d_h = d_model // H
x = torch.randn(T, d_model)      # one projection's output, e.g. Q

heads = x.view(T, H, d_h).transpose(0, 1)   # (H, T, d_h)
print("x", tuple(x.shape), "-> heads", tuple(heads.shape))

# No arithmetic happened: the heads share x's memory.
print("same storage:", heads.data_ptr() == x.data_ptr())
print("head 1, token 0:", [round(v, 3) for v in heads[1, 0].tolist()])
print("x[0, 4:8]      :", [round(v, 3) for v in x[0, 4:8].tolist()])

# Each head gets its own (T, T) attention map, for the same FLOPs.
A = torch.softmax(heads @ heads.transpose(-2, -1) / d_h**0.5, -1)
print("attention maps:", tuple(A.shape))

back = heads.transpose(0, 1).reshape(T, d_model)   # the Concat step
assert torch.equal(back, x)
print("concat recovers x bit-for-bit:", torch.equal(back, x))

Reference figureEach head gets its own learned Q/K/V projection matrices, run in parallel on the same input. Source: Jay Alammar, The Illustrated Transformer.

Positional information

Attention itself is permutation-invariant. Shuffle the input tokens and you get the same set of outputs, shuffled the same way, so the layer has no sense of order. Order has to be injected, and there are two standard ways to do it:

• Sinusoidal (original paper). Fixed position vectors added to the input embeddings, one frequency per dimension.
• RoPE, Rotary Position Embedding (modern default). Rotate each query/key vector by an angle proportional to its position.

RoPE's payoff is subtle but large. The dot product between a query at position $m$ and a key at position $n$ ends up depending on the relative distance $m-n$ and not on the absolute positions. This relative-distance property makes RoPE extrapolate more gracefully to longer sequences than fixed sinusoidal or learned absolute embeddings.

05Why this architecture

The obvious alternatives were "keep the RNN" and "use a CNN." Attention wins on the two axes that matter most for training at scale:

• Parallelism. How much of the computation can run at the same time.
• Path length. How many steps separate any two related tokens.

The $O(1)$ path length is the main prize. In an RNN, information from token 1 has to survive $n{-}1$ sequential updates before it can influence token $n$. This is where vanishing gradients bite (see 07). In self-attention, token 1 and token $n$ are always one dot product apart, regardless of distance.

06Complexity, failure modes, and what breaks

Self-attention costs $O(n^2 \cdot d)$ time and, naively, $O(n^2)$ memory to materialize the full attention matrix, where $n$ is sequence length. Double the context window and compute roughly quadruples. That single fact drives a huge amount of LLM systems engineering. It is why 23 (Efficient AI & Systems) exists as its own section.

What shows up in production

• KV cache growth. During autoregressive generation you don't recompute attention from scratch each step. You cache every past token's K and V vectors. Cache size grows linearly with tokens generated × layers × heads × head_dim. For long conversations this becomes the dominant memory cost on the serving GPU, often larger than the model weights themselves.
• FlashAttention doesn't reduce the $O(n^2)$ FLOPs. It restructures the computation to avoid ever writing the full $n \times n$ attention matrix to slow GPU memory, computing softmax in fused tiles instead. Same math, dramatically less memory traffic. Memory traffic is usually the real bottleneck on modern GPUs.
• Grouped-query / multi-query attention shrinks the KV cache by having multiple query heads share one set of K/V heads. A deliberate quality-for-memory tradeoff, used in almost every modern serving-optimized LLM.
• Attention sinks and long-context dilution. Models reliably show a strong bias toward attending heavily to the first few tokens, regardless of relevance. Needle-in-a-haystack evaluations show retrieval accuracy dropping for facts placed in the middle of very long contexts. A model that technically can attend anywhere does not mean it attends usefully everywhere.

07Build this

You can read the attention equation twenty times and still not believe it works. Building the smallest thing that uses it fixes that in an afternoon.

Where this runs in production

Every long-running chat assistant relies on this directly. Causal self-attention lets today's question attend straight back to something you said 50 turns ago, in one hop, with a weight the model learned to assign based on relevance. It is not filtered through a shrinking, decaying RNN memory state.

Cross-attention does the other half of the work: query from one sequence, key and value from another. It is the mechanism a retrieval-augmented system or a translation decoder uses to let generated output consult retrieved documents or the source sentence, at every single generation step. See 11 for how this composes with a retrieval pipeline.

08Interview questions

09Go deeper

●Now write it yourself

Reading the derivation and being able to produce it are different skills. These are Deep-ML problems that exercise what this page covers — each one is checked against real test cases, not multiple choice.

Matched to this page from Deep-ML's catalogue of 1,380 problems. More at deep-ml.com, and Where to practise covers the other platforms and what each one trains.

10Related pages

Other techniques for this problem

A scoped slice of the full Technique Map — every technique this page covers, grouped by what it solves.