Attention Explained: 4 Simple Steps Behind Modern AI

โšก TL;DR: Attention lets every token look at every other token and decide which ones matter. It works by comparing a query vector against key vectors with dot products, turning those scores into weights, and taking a weighted average of value vectors.

Attention is the mechanism that made modern language models possible, and it is far less intimidating than its reputation. If you understand a dot product and a softmax, you already have both halves.

The problem attention solves

Consider: “The animal didn’t cross the street because it was tired.”

What does “it” refer to? The animal โ€” obviously, to you. But to a model processing tokens in order, “it” is just a token. Something has to connect it to “animal” seven positions back.

Attention heatmap over a sentence showing each token attending to earlier tokens, with the word it attending strongly to animal
Each row shows where one token looks. The row for “it” puts most of its weight on “animal”.

Older architectures processed text strictly left to right, carrying a single summary forward โ€” and the further back a word was, the more diluted it became. Attention removes the distance penalty entirely: every token can reach any other token in one step.

Query, key and value

Each token produces three vectors from its embedding, by multiplying it against three learned weight matrices.

Query
What this token is looking for. “it” asks: which noun am I standing in for?
Key
What this token offers. “animal” advertises: I am a noun that could be referred to.
Value
What this token actually contributes once it has been selected.

๐Ÿงช The library analogy

You walk into a library with a question in mind โ€” that is your query.

Every book has a title on its spine โ€” those are the keys. You scan the spines and judge which look relevant.

The values are the contents of the books. You do not read all of them equally: you read mostly from the best match, a little from the second, and none from the irrelevant ones.

Attention does exactly this, except the “how relevant” judgement is a dot product and you always take a blend rather than picking one book.

The formula, one piece at a time

$$\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{T}}{\sqrt{d_k}}\right)V$$

It looks dense; it is four steps.

  1. QKแต€ โ€” take the dot product of every query with every key. A large dot product means the vectors point in similar directions, so that key is relevant to that query. This produces a grid of scores, one per token pair โ€” exactly the heatmap above.
  2. รท โˆšdk โ€” divide by the square root of the key dimension. With long vectors the dot products grow large, which pushes softmax into a near-one-hot output and kills the gradients. This scaling keeps them in a workable range.
  3. softmax โ€” convert each row of scores into weights that are positive and sum to 1. See the softmax function.
  4. ร— V โ€” take the weighted average of the value vectors. Tokens with high weight contribute most.
โœ… The output: a new vector for each token, blended from whichever other tokens it judged relevant. The vector for “it” now carries information from “animal” โ€” which is precisely how a contextual embedding gets its context.

Multi-head attention

One set of query, key and value matrices can only learn one kind of relationship. Real language has many running at once โ€” grammatical agreement, what refers to what, topic, tone.

So the whole mechanism is run several times in parallel with different learned matrices. Each copy is a head, and the results are concatenated. One head might track subject-verb agreement while another tracks pronoun reference. Nobody assigns these roles; they emerge during training.

Self-attention and causal masking

Self-attention means queries, keys and values all come from the same sequence โ€” the text is looking at itself.

In a text-generating model there is a constraint: predicting token 5 must not use token 6, because at that moment token 6 does not exist yet. Causal masking enforces this by setting the scores for all future positions to negative infinity before the softmax, so their weights become zero.

That is why the heatmap above is triangular โ€” every token sees itself and everything before it, and nothing ahead.

The cost, and why context windows are limited

โš ๏ธ Attention scales quadratically

Every token compares against every other token, so n tokens require nยฒ comparisons. Double the input and the work quadruples.

1,000 tokens โ†’ a million comparisons. 100,000 tokens โ†’ ten billion. This is the single biggest reason long context windows are expensive and were slow to arrive, and why a great deal of research goes into approximations like FlashAttention and sparse attention.

๐Ÿ”‘ Key Takeaways

  • Attention lets every token pull information from any other token in one step.
  • Query asks, key advertises, value delivers โ€” relevance is a dot product between query and key.
  • Softmax turns raw scores into weights that sum to 1; the output is a weighted average of values.
  • Multiple heads run in parallel so different relationships can be tracked at once.
  • Cost grows with the square of sequence length, which is what limits context windows.

Further reading

Jay Alammar’s Illustrated Transformer walks through the same mechanism with diagrams at every step, and Attention Is All You Need is the original paper โ€” surprisingly readable once you have the intuition above.

Where to go next

Scroll to Top