Attention Mechanism Deep Dive

§ 01 · HookA Promise Kept

Here is the whole thing in one line, from the 2017 paper:

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

By the end of this article, every symbol in that line will make sense, and you'll know why each piece is there. We'll build up to it from scratch, introducing each part only after you've seen the problem it solves.

This matters beyond theory. Engineers who understand attention can reason about why a long prompt gets expensive, what a model card means by "RoPE" or "FlashAttention", and why a model sometimes "loses" information buried in the middle of a long context.

§ 02 · The Problem, PreciselyThe Problem Attention Solves

After tokenisation, each token ID is swapped for a vector from the embedding table (Article #02, §03). That vector is static: it's a single row of a table, the same every time the token appears, whatever surrounds it.

So "bank" enters the Transformer with exactly the same vector in "the river bank eroded" and "the savings bank collapsed". (In BERT-base that's 768 numbers; in the small model we'll use below, 384.) Yet the two uses mean completely different things. The input vector can't know which one it is.

Attention's job is to update each token's vector by letting it gather information from the rest of the sentence. What comes out of an attention layer isn't the original static vector; it's a new, context-aware one that has "looked at" the whole sentence and blended in whatever was relevant.

§ 03 · The Core MechanismQueries, Keys, and Values

This is the heart of the article, so take it slowly. We'll build it in three stages.

Stage 1: the search story, generalised

Article #02 used a video-search example. Here's the general version: every token is both a searcher and a search result, at the same time.

As a searcher, each token asks, "What information do I need from the rest of this sentence?" That question is its Query. As a result, each token offers two things: a Key, a label other tokens match their questions against, and a Value, the actual content handed over when there's a match. The better a query matches a key, the more of that token's value gets retrieved. Every token searches all the others and is searched by all the others, simultaneously.

Stage 2: three learned projections

Each token's vector x is multiplied by three different learned matrices:

Q = x · W_Q    # what am I looking for?
K = x · W_K    # what do I advertise?
V = x · W_V    # what do I deliver?

W_Q, W_K, and W_V are ordinary parameters, learned by gradient descent (Article #01, §05) along with everything else. They're deliberately separate, because "what I'm looking for" and "what I have to offer" are different things. A pronoun like "it" is looking for a noun; the noun "trophy" is offering itself as one. Separate matrices let the model learn those roles independently.

The projected vectors are usually smaller than the original. In the 2017 paper, each token vector has 512 numbers and there are 8 attention heads, so each head's Q, K, and V have 512 ÷ 8 = 64 numbers.

What does "multiply by a matrix" actually mean? If you haven't done matrix maths since school, this phrase can feel like a black box. It isn't. Each number in the output is just a dot product: take the token's vector, pair it up with one column of the matrix, multiply matching positions, and add. Do that once per column and you have the new vector. Figure 3 shows it with small numbers you can check in your head.

Two things to notice. First, the sizes have to line up: the 4 numbers across the vector must match the 4 rows down the matrix. Second, the number of columns decides how long the result is. That's how a matrix can shrink a 512-number vector into a 64-number query: it simply has 64 columns. W_K and W_V work the same way, just with different numbers inside, which is why the same token comes out as three different vectors.

Stage 3: the four steps, each with a reason

Step 1, dot product. To decide how much token i should attend to token j, compute score(i, j) = Q_i · K_j. A dot product multiplies matching positions of two vectors and adds them up (Article #02, §04); it's large when the two vectors point the same way. A high score means i's question is well matched by j's label.

Step 2, scale by √d_k. Here's why this step exists. A dot product of two 64-number vectors adds up 64 products, so as vectors get longer, the scores naturally get bigger. Big scores push the next step, softmax, into an extreme where one token gets almost all the weight and the rest get almost none. In that state the gradients become tiny and learning stalls, the same vanishing-gradient problem from Article #02, §02. Dividing by √d_k (√64 = 8 for 64-number keys) keeps the scores in a healthy range.

Step 3, softmax. Turn each token's row of scores into weights between 0 and 1 that add up to 1 (Article #01, §08). Each weight says how much of token j's information token i should take in. Softmax works on one row at a time, so each token gets its own set of weights.

Now we can see what "the extreme" from Step 2 looks like. Take three scores, 1, 2 and 3. Softmax turns them into 0.09, 0.24 and 0.67: the best match wins, but the others still count. Multiply the same scores by 8 and softmax gives 0.00, 0.00 and 1.00. One token takes everything. It's like a microphone turned up too loud, where the loudest voice drowns out everyone else in the room. And because the losing tokens are stuck at zero, small changes to them make no difference, so the model gets almost no signal about how to improve them.

Why divide by √d_k and not some other number? We measured it (bottom of Figure 4). For random vectors, the typical dot product between two 4-number vectors is about 2, between two 64-number vectors about 8, and between two 512-number vectors about 22.6. Those are exactly √4, √64 and √512. So dividing by √d_k cancels out the growth and brings scores back to a similar size, however long the vectors are.

Step 4, weighted sum of Values. Multiply each Value vector by its weight and add them all up: output_i = Σ weight_ij × V_j. The result is a blend of every token's Value, dominated by the most relevant ones.

A blend is easier to picture with something familiar. Say it's Friday movie night, and your mood splits 50% towards a heist thriller, 30% towards a family comedy and 20% towards a cricket drama. Describe each film with three numbers: action (out of 10), comedy (out of 10) and running time in minutes.

Each column of the last row is the same recipe: 0.5 × 9 + 0.3 × 2 + 0.2 × 5 = 6.1 for action, and so on. The result isn't any one of the three films. It's a new description that leans towards the one you cared about most. In attention, the weights come from softmax, the films are the tokens' Value vectors, and the blended mood is the token's new, context-aware vector. When "it" puts most of its weight on "trophy", the new vector for "it" ends up carrying a good share of what "trophy" means.

Attention from scratch in NumPy

It's short enough to write yourself. The weights below are random (nothing is trained), so the pattern of who looks at whom means nothing. The point is to watch the mechanics.

import numpy as np
np.set_printoptions(precision=2, suppress=True)

def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)          # subtract the max for numerical safety
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, causal=False):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)                 # steps 1 and 2: dot products, then scale
    if causal:                                      # decoder models: hide the future
        future = np.triu(np.ones_like(scores), k=1).astype(bool)
        scores = np.where(future, -np.inf, scores)
    weights = softmax(scores)                       # step 3: each row sums to 1
    return weights @ V, weights                     # step 4: blend the Values

rng = np.random.default_rng(0)
tokens = ["the", "cat", "sat", "down"]
X = rng.normal(size=(4, 8))                         # 4 token embeddings, 8 numbers each
W_Q, W_K, W_V = (rng.normal(size=(8, 4)) for _ in range(3))

Q, K, V = X @ W_Q, X @ W_K, X @ W_V                 # three projections, 4 numbers each
out, w = attention(Q, K, V)
print(w)
print("row sums:", w.sum(axis=1))
print("shapes:", X.shape, Q.shape, w.shape, out.shape)
print("new vector for 'cat':", out[1])

# [[0.38 0.13 0.46 0.03]
#  [0.01 0.   0.13 0.86]
#  [0.91 0.   0.08 0.  ]
#  [0.01 0.25 0.08 0.66]]
# row sums: [1. 1. 1. 1.]
# shapes: (4, 8) (4, 4) (4, 4) (4, 4)
# new vector for 'cat': [ 0.94 -0.4  -0.03  0.59]

Each row of w is one token's attention: how much it takes from each of the four tokens. Every row sums to 1. The shapes tell the story too: 4 tokens of 8 numbers go in, each is projected down to 4 numbers, the weights form a 4 × 4 grid (every token against every token), and out comes one new 4-number vector per token. Notice that the new vector for "cat" has negative numbers in it. That's fine: it's a blend of Values, not a probability. Figure 5 follows the "cat" row through all four steps.

The paper's diagram, decoded

If you open the 2017 paper, you'll find this recipe drawn as a small stack of boxes called "Scaled Dot-Product Attention". It looks intimidating at first, but you now know every box in it. Figure 6 redraws it top to bottom, with the shapes from our NumPy example written beside each step.

The diagram answers a question many learners ask: why do Q and K travel together while V takes a separate route? Because Q and K have one job, working out the weights, and V has a different job, supplying the content. "MatMul" is simply the paper's short name for matrix multiplication.

§ 04 · VariantsSelf-Attention, Causal Masking, and Cross-Attention

The same mechanism comes in three variants, and mixing them up causes real confusion when you read model cards and papers.

Self-attention (bidirectional)

Q, K, and V all come from the same sequence, and every token can look at every other token, before and after it. This is what encoder models like BERT use, for understanding tasks such as classification, named entity recognition (finding and labelling names of people, places, and organisations in text), and embeddings.

Causal (masked) self-attention

Q, K, and V still come from the same sequence, but each token may only look at the tokens before it. This is what decoder-only LLMs use: the GPT, Llama, and Mistral families, and, as far as their makers have described them, most of today's chat models.

Cross-attention

The queries come from one sequence and the keys and values from another. In an encoder-decoder model like T5, the decoder's tokens ask questions of the encoder's output. Image generators such as Stable Diffusion use cross-attention too, letting the picture being generated "consult" the words of your prompt.

Why causal masking is necessary

During training, the model sees the whole sequence at once, because that's what makes training parallel (Article #02). Without a mask, token 5 could look at token 6 while being trained to predict token 6. It would be copying the answer.

The fix is to set the scores for all future positions to −∞ before the softmax. Since e^−∞ is 0, those tokens get exactly zero weight, as if they didn't exist yet. You can see this in the NumPy code: pass causal=True and every weight above the diagonal becomes 0. It's the same rule as your phone's predictive text (Article #01, §08): it can only use what you've typed so far, never the word you haven't typed yet.

At generation time, the mask is almost automatic. The model produces one token at a time, left to right, so future tokens genuinely don't exist yet.

§ 05 · Multiple RelationshipsMulti-Head Attention

One attention operation tends to capture one kind of relationship. But a sentence is a web of relationships at once: who did what, what "it" refers to, which words sit next to each other, which words share a topic. So Transformers run several attention operations side by side, each with its own W_Q, W_K, and W_V, then join the results. This is the mechanical version of Article #02's "reading the same novel several times, each time looking for something different".

MultiHead(Q, K, V) = Concat(head₁, …, head_h) · W_O
    where head_i = Attention(Q · W_Q⁽ⁱ⁾, K · W_K⁽ⁱ⁾, V · W_V⁽ⁱ⁾)

Each head works in a smaller space. With 512-number token vectors and 8 heads, each head gets 64-number queries, keys, and values. So adding heads doesn't multiply the work; the same total is split into several parallel, specialised jobs whose outputs are concatenated and mixed by one final matrix, W_O.

"Concatenated" just means placed side by side. Each of the 8 heads hands back a 64-number vector for every token. Lay the 8 vectors end to end and you get 8 × 64 = 512 numbers, the same size the token started with. W_O then lets the model mix what the different heads found. In a sentence like "I ate the cake and it was delicious", one head might link "it" to "cake", another might link "ate" to "cake", and a third might carry the positive feeling of "delicious" over to "cake". Joining the heads puts all of that into one vector.

What heads actually learn

Nobody tells a head what to specialise in. Researchers who study what trained models have learned inside (an area nicknamed "BERTology") have found heads that track recognisable patterns:

• Neighbour heads that look at the previous or next word.
• Syntax heads where verbs attend to their subjects, or objects to their verbs.
• Coreference-like heads, where a pronoun such as "it" attends to an earlier noun. (Coreference means working out which earlier thing a word like "it" or "she" refers to.)
• Separator heads that park most of their attention on special tokens like [SEP] when they have nothing better to do.

You don't have to take this on trust. Figure 8 shows three heads from a real model, all reading the same sentence.

§ 06 · Hands-on CodeSee Real Attention with bertviz

bertviz is an open-source tool that draws a Transformer's actual attention weights as interactive coloured lines inside a Jupyter notebook. You can click through layers and heads and see the patterns from §05 for yourself. We'll use all-MiniLM-L6-v2, the small BERT-style model from Article #04 (6 layers, 12 heads each), because it downloads quickly. The same code works with bert-base-uncased (12 layers, 12 heads).

$ pip install bertviz transformers torch

import torch
from bertviz import head_view
from transformers import AutoTokenizer, AutoModel

name = "sentence-transformers/all-MiniLM-L6-v2"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModel.from_pretrained(name, output_attentions=True,
                                  attn_implementation="eager")   # needed to return attention weights
model.eval()

sentence = "The trophy didn't fit in the suitcase because it was too big."
inputs = tokenizer(sentence, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

attention = outputs.attentions            # one tensor per layer: (batch, heads, seq_len, seq_len)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

print(len(attention), attention[0].shape[1], len(tokens))
# 6 12 17     6 layers, 12 heads, 17 tokens

print(tokens)
# ['[CLS]', 'the', 'trophy', 'didn', "'", 't', 'fit', 'in', 'the', 'suitcase',
#  'because', 'it', 'was', 'too', 'big', '.', '[SEP]']

head_view(attention, tokens)             # interactive view, renders in Jupyter

Notice there are 17 tokens, not 12 words: BERT-style tokenisers add [CLS] and [SEP], and split "didn't" into didn, ', and t (Article #03 in action). In the interactive view:

• Click between heads and watch the patterns change completely.
• Find a head where every token points one step back: a neighbour head.
• Find a head where "it" sends a thick line to "trophy".
• Earlier layers tend to show more local, syntax-like patterns; later layers tend to show longer-range ones.

Find the "it" head automatically, then draw it

Instead of hunting by eye, let the code search all heads for the one where "it" attends most strongly to "trophy":

import matplotlib.pyplot as plt

it, trophy = tokens.index("it"), tokens.index("trophy")
layer, head = max(
    ((l, h) for l in range(len(attention)) for h in range(attention[0].shape[1])),
    key=lambda lh: attention[lh[0]][0, lh[1], it, trophy].item(),
)
weights = attention[layer][0, head].numpy()     # rows = token looking, columns = token looked at
print(f"layer {layer+1}, head {head+1}: it → trophy = {weights[it, trophy]:.2f}")
# layer 4, head 2: it → trophy = 0.58

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(weights, cmap="Oranges", vmin=0, vmax=1)
ax.set_xticks(range(len(tokens)), tokens, rotation=60, ha="right", fontsize=9)
ax.set_yticks(range(len(tokens)), tokens, fontsize=9)
ax.set_title(f"Attention weights: layer {layer+1}, head {head+1}")
plt.colorbar(im, ax=ax, label="attention weight")
plt.tight_layout()
plt.savefig("attention_heatmap.png", dpi=150)
plt.show()

§ 07 · PositionPositional Encoding Revisited

Attention has a built-in blind spot: it's permutation-invariant. Scores come from dot products between pairs of tokens, and nothing in that calculation knows which token came first. Shuffle the sentence and each token ends up with the same attention weights for the same partners. "Dog bites man" and "man bites dog" would look identical. Word order carries meaning, so that would be a disaster.

Article #02, §03 introduced the fix: add position information. Here are the three main ways it's done.

Sinusoidal (the original 2017 paper)

Each position gets a fixed pattern built from sine and cosine waves of different speeds, added to the token's embedding. Nothing is learned, and a vector can be computed for any position, however large. In practice, though, models trained this way still get worse on sequences much longer than the ones they were trained on.

Learned position table (GPT-2)

A simple lookup table with one learned vector per position, added to the token's embedding. It's easy and works well, but it has a hard limit: GPT-2's table has 1,024 rows, so there's simply no vector for position 1,025.

RoPE, Rotary Position Embedding (Llama, Mistral, and most new LLMs)

Instead of adding a vector, RoPE rotates the query and key vectors by an angle that depends on each token's position, just before the dot product. Picture each vector as a clock hand, turned further for tokens later in the sentence.

The clever part: when you take the dot product of two rotated vectors, what matters is the difference between their angles, so attention ends up depending on how far apart two tokens are rather than on their absolute positions. It's like a stopwatch: you can tell five minutes passed by comparing the two hand positions, without knowing what time you started.

RoPE has no lookup table to run out of, and it has become the standard for decoder-only LLMs. It doesn't magically work at any length, though. Models still get worse beyond the length they were trained on unless they're extended with techniques such as position interpolation or YaRN. That's how many models are stretched to much longer context windows after their main training.

§ 08 · Engineering RealityAttention Cost and the Long-Context Problem

Standard attention compares every token with every other token. For n tokens, that's roughly n × n comparisons, written O(n²). Double the length and you quadruple the work and memory for attention.

Think of a room where everyone must shake hands with everyone else. Ten people make 45 handshakes. Double it to twenty people and you don't get 90, you get 190, because every newcomer has to greet everyone, including the other newcomers. Article #03 used exactly this reasoning to explain why character-level tokens would be so expensive.

With context windows now reaching hundreds of thousands of tokens, and a million or more in some models, this is a serious engineering problem. Here's what the field has done about it.

FlashAttention (Dao et al., 2022)

FlashAttention computes exactly the same result, still O(n²), but reorganises the work so the GPU spends far less time moving data between its small, fast on-chip memory and its larger, slower main memory. That data movement turned out to be the real bottleneck. Think of a chef who doesn't chop any faster, but stops walking to a distant pantry for every ingredient and keeps what the next few steps need within arm's reach. It typically runs attention several times faster and uses much less memory, and it's now the default in most frameworks. FlashAttention-2 (2023) and FlashAttention-3 (2024, tuned for NVIDIA's H100 GPUs) push the same idea further.

Sliding-window attention (for example Mistral 7B)

Each token attends only to a fixed window of recent tokens, so the cost grows with n × window instead of n². It works because most attention weight is local anyway, and information from further back can still travel forward through the stacked layers.

Sparse attention (Longformer, BigBird)

Only certain pairs are computed: local windows plus a few "global" tokens that everything can see. This was designed for document-length inputs.

§ 09 · SummaryThe Complete Mental Model