Part 3: Attention, context, and the geometry of relevance¶
This chapter is the heart of the Transformer. The earlier pieces gave us tokens, embeddings, and position. Attention is the part that turns a static token into a context-dependent meaning.
Foundation connection: Dot products, outer products, and softmax provide the numerical backbone. This chapter turns those primitives into the actual routing mechanism used in transformers.
1. Why static embeddings are not enough¶
A token like “bank” can mean a financial institution, the side of a river, or a place where a plane turns. The embedding lookup gives the same initial vector to the same token ID no matter what sentence it appears in.
That is useful, but it is not enough. A language model needs a representation that changes with the surrounding context.
The same is true for “mole,” “tower,” or “charge.” The meaning depends on nearby words, topic, and legal or scientific context. A single fixed embedding is just a starting point.
Attention solves this by letting each token ask:
- Which nearby tokens matter to me here?
- How strongly should I rely on each one?
- What information should be mixed into my current representation?
This is the step from a token identity to a context-aware meaning.
2. The geometric picture¶
Think of the embedding space as a semantic geometry. A token embedding is a point in a high-dimensional space, but the embedding alone does not contain the full sentence context.
The attention block computes a new representation by pulling information from relevant positions. The resulting vector moves in a different direction depending on the surrounding words.
So the role of attention is not to invent meaning from nowhere. It is to update a token’s representation using the evidence in the sequence.
That is why attention is often described as a learned routing mechanism: the model learns which context signals are relevant and how strongly to mix them.
A useful intuition is to think of attention as a content-based lookup system. A token asks a question, compares that question against all candidate tokens, and then retrieves the information from the most relevant ones. In a dictionary, you search for a word, find the matching entries, and read the corresponding definitions. In attention, the query plays the role of the search key, the keys act like indexes or labels, and the values are the actual content that gets retrieved.
This is the basic principle behind attention: instead of treating each token as isolated, the model lets one token “look up” the information it needs from the surrounding context. The matching is not hard-coded; it is learned by the model through the projection matrices and the training objective.
That is the bridge from the intuition to the formal Transformer notation: a token asks “which other tokens are relevant to me?” and then receives a weighted mixture of their value vectors.
3. Query, key, and value¶
After embedding, every token gets a vector \(e_i\). The model then projects it into three different spaces:
These are the usual Transformer ingredients:
- \(q_i\) is the query: “what information am I looking for?”
- \(k_i\) is the key: “what information does this token carry?”
- \(v_i\) is the value: “what content should be passed along if this token is relevant?”
The model learns the matrices \(W^Q\), \(W^K\), and \(W^V\) from training data. They are not handcrafted features. They are learned projections that align token representations with useful comparison patterns.
The equations are simple, but the meaning is important. The same token is projected into three different roles, and the model learns which role is useful for which context. The query tells us where to look; the key tells us what each token offers; the value tells us what information is actually passed forward when a match is found.
4. Compare each token to all keys¶
For a token at position \(i\), we compare its query against every key in the sequence by dot product:
A large positive value means the two positions look relevant to each other under the current learned metric. A small or negative value means they are not strongly aligned.
If we stack all queries and keys into matrices \(Q\) and \(K\), then the full score matrix is:
This is the matrix of all pairwise similarities. Each entry \(S_{ij}\) measures how much query \(i\) matches key \(j\) after the learned projection. A larger value means “this token is a stronger candidate source for context.”
The scaling by \(\sqrt{d_k}\) is a numerical stabilization trick. It keeps the score magnitudes from growing too large as the key dimension increases. Without this scaling, some heads can become overly sharp, and the attention pattern can collapse too aggressively during training.
This matrix is the heart of attention: it records pairwise relevance between every query position and every key position.
5. Softmax turns scores into weights¶
Raw dot products are not probabilities. We convert them into attention weights with softmax:
This does two things at once:
- it turns the raw similarity scores into positive values
- it normalizes them so the row sums to 1
So each row of \(A\) becomes a probability distribution over the context positions. The model is not picking a single token. It is saying: “here is how much I want to rely on each token in the context.”
A tiny example makes this concrete. Suppose the score row for one query is:
Applying softmax gives something like:
The token with the largest score, \(2.9\), gets the largest weight, but the other scores are not discarded. They still contribute in smaller amounts. This is exactly why the attention pattern is soft rather than hard: the model learns a weighted mixture, not a single winner.
This is why the operation is called attention: the model pays more attention to some positions and less to others, while still retaining a broad context.
6. Values carry the information¶
Once the weights are known, the model mixes the value vectors:
This is the actual information transfer step.
The process is:
- compute how relevant each context token is
- normalize that into attention weights
- use those weights to combine value vectors
- produce the updated representation for the current token
The weighting is the important idea. The attention pattern decides which context content enters the token state; the values are the content being passed along.
The output is then used in the layer update:
This is the residual connection. The original hidden state \(h\) is not thrown away. Instead, the attention result \(O\) is added back to it, so the model keeps the original token information while also incorporating context. In a full Transformer block the usual pattern is:
This is why attention is a contextual refinement, not a replacement: the model keeps the token’s identity and adds the context-aware correction.
7. The canonical attention equation¶
The full attention layer is written as:
Read this carefully in order:
- compute the score matrix \(QK^\top\)
- scale it to keep magnitudes stable
- apply softmax row by row so each query gets a valid probability distribution
- use those weights to mix the value vectors
The output is a weighted sum of values, so the model passes information from the context positions that were most relevant to the current token.
This is the core Transformer primitive.
It is mathematically simple, but the behavior is rich because the model learns \(W^Q\), \(W^K\), and \(W^V\) from data. The same formula can learn different relationship patterns depending on the task, training objective, and layer depth.
8. A full worked example with exact dimensions¶
A real Transformer may have:
- vocabulary size \(V \approx 50{,}000\)
- embedding dimension \(d = 128\)
- one attention head with head dimension \(d_k = d_v = 64\)
This is far too large to draw exactly, so we show the shape and the logic with a small example and a sentence-level matrix view.
8.1 One token: \(q\), \(k\), and \(v\) as vectors¶
Take the token “bank.” Its embedding is a vector:
After positional encoding, it is still a \(128\)-dimensional vector:
Now the model projects this single token into three different spaces:
with
For the current word “bank,” the model compares its query against the keys of nearby words:
This gives a score vector:
Then the model rescales and normalizes this score vector:
Suppose this gives:
The updated representation for “bank” is then:
This is the one-word worked example. At this point, the shapes are clear:
- token embedding: \(128\)
- query/key/value: \(64\)
- score vector against the sentence: \(T\)
- attention weight vector: \(T\)
- output vector: \(64\)
8.2 Full sentence: \(X\), \(Q\), \(K\), \(V\), \(S\), \(A\), \(O\)¶
Now take a sentence of length \(T=6\). The token embeddings form a matrix:
Add position:
Project the whole sentence:
Then compute all pairwise similarities at once:
This matrix is the full attention score table. Each row is one query token, and each column is one key token.
Now apply row-wise softmax:
Finally, mix values:
The row for the token “bank” is exactly the one-word computation in matrix form:
This is the exact end-to-end path:
The one-word case is simply one row of the sentence-level computation. The full sentence is not different in principle; it is just the same operation applied to all tokens at once.
This is the real computational picture behind attention: the model does not view tokens independently. It turns the whole sequence into a set of query, key, and value matrices, builds a score grid, converts it to attention weights, and updates each token with a weighted mix of the relevant values.
The figure is intentionally simplified. A real transformer has \(50{,}000\) vocabulary entries and \(128\)-dimensional vectors, but we only show a small set of words in a 3D projection so the geometry is visible.
- gray points: the base word embeddings
- red point: the current query for “bank”
- blue points: the context keys being compared against it
- gold point: the final attention output after weighting the values
- bars on the right: the softmax row that tells us how much each word contributes
This is the heart of attention: the model does not just look at a token’s identity. It updates that token using the context that best matches its current query.
9. Self-attention and cross-attention¶
A single consistent example helps here. Take the sentence:
the bank was crowded near the river
and imagine the model is updating the word “bank.”
Self-attention¶
In self-attention, the current token asks about other tokens in the same sequence. The same hidden state matrix supplies everything:
where \(H \in \mathbb{R}^{T \times d}\) is the hidden representation of the whole sentence. The query for “bank” is compared against all keys in the same sentence, and the values from the relevant positions are mixed back in.
This is the standard language-model setup. The model learns which nearby words matter for the current token, and it updates the representation based on that context.
Cross-attention¶
In cross-attention, the math stays the same, but the sources change.
The decoder asks a question about the encoder output, not about its own previous tokens. A good worked example is machine translation:
- source English sentence: “The bank was crowded near the river.”
- target French sentence: “Le bord était bondé.”
Suppose the French decoder is generating the word for “bank” and needs to decide which English words matter. The decoder creates a query from its own representation:
The encoder provides the keys and values for the source sentence:
Now the score matrix is:
followed by
This is the crucial point for cross-attention: the decoder chooses where to look, but the encoder provides the keys and values that carry the content. The values do not come from nowhere; they are the transformed source-side hidden states from the encoder.
So the difference is not a different formula. The difference is which sequence provides the query and which sequence supplies the key/value content.
- self-attention: \(Q, K, V\) all come from the same sentence
- cross-attention: \(Q\) comes from the decoder, while \(K\) and \(V\) come from the encoder
In both cases, the final step is the same:
The current token keeps its original representation, then receives a context-aware update.
One consistent geometric picture¶
The same picture works for both settings:
- the current token asks a question with \(Q\)
- candidate tokens provide keys with \(K\)
- the selected content is delivered through \(V\)
- the model mixes the relevant values into a new, context-aware representation
This is what makes self-attention useful for language modeling and cross-attention useful for translation, summarization, and other conditioning tasks.
10. Why masked attention is required¶
For autoregressive generation, a token cannot attend to future tokens. Otherwise the model would look ahead and cheat during training.
This is handled with a causal mask.
For a sequence of length \(T\), we define:
Then the effective scores become:
After softmax, positions with \(-\infty\) become zero weight. The token can only attend to the current prefix, never the future.
This is what makes next-token prediction possible while still allowing parallel training across positions.
A causal-mask heatmap¶
The causal mask is easy to visualize as a matrix of allowed and forbidden connections.
The lower triangle is allowed; the upper triangle is blocked. This ensures each token only reads the prefix that is legally available at its position.
11. Why a single head is not enough¶
One attention head learns one relevance pattern. That is useful, but language contains many different kinds of relationships at once:
- pronoun to antecedent
- noun to adjective
- syntax to semantics
- local pattern to long-range context
- entity tracking across a sentence
A single head is usually too narrow to capture all of this. That is why transformers use multi-head attention.
12. Multi-head attention¶
Suppose there are \(h\) heads and model width \(d\). Each head uses its own learned projections:
Then each head computes its own attention output:
The results are concatenated and projected back into the model width:
This lets the model view the same sequence through multiple relational lenses at once. One head might emphasize syntax, another local similarity, another long-range topic continuity.
A useful mental picture is to imagine the same sentence being viewed under several filters at once:
- one head tracks local phrase structure
- one head tracks entity mentions
- one head tracks topic continuity
- one head tracks disagreement or sentiment edges
These are not hand-coded rules. They are learned patterns that emerge from the training objective.
Different heads do not all learn the same pattern. The same sentence is viewed under several learned filters, and each filter can specialize in a different relationship.
A simple schematic is:
flowchart LR
X[sequence tokens] --> H1[head 1: local syntax]
X --> H2[head 2: entity links]
X --> H3[head 3: topic continuity]
H1 --> C[concatenate]
H2 --> C
H3 --> C
C --> O[projected multi-head output]
The exact pattern is not a literal grammar; it is a learned attention structure. Multi-head attention is the model’s way of having several distinct relevance views at once.
The important point is not that each head is interpretable in a human sense. It is that multiple learned projections give the model more flexibility than a single attention pattern.
13. Why attention became the dominant idea¶
The transformer architecture replaced recurrence because it gave the model direct access to context without hiding everything inside a single compressed state.
In a recurrent model, information has to pass through a hidden state over time. In attention, each token can directly compare itself to any position in the sequence.
This makes long-range dependencies much easier to model. It also makes the architecture naturally parallelizable during training.
That is a major reason the transformer took over language modeling.
14. The computational cost¶
The naive self-attention score matrix has size \(T \times T\) for a sequence of length \(T\).
That makes the cost roughly:
for sequence length \(T\) and hidden width \(d\). This is the famous quadratic cost of attention.
The issue is not just time. It is also memory. The model must store and read a large attention table and the corresponding key and value tensors.
This is one reason later systems research cares so much about efficient attention, caching, and sparse approximations.
15. The practical view in generation¶
At inference time, language models generate one token at a time.
Without optimization, each newly generated token would recompute attention against the full sequence again. That is wasteful because previous keys and values do not change.
The key-value cache stores previous key and value vectors so that only the newest token needs a full attention computation. This is the standard efficiency trick behind generation systems.
This is one of the places where the neat mathematical idea meets large-scale engineering.
16. The basic numeric recipe¶
The attention computation is simple when broken into steps:
- compute query-key compatibility scores
- convert them to a probability distribution with softmax
- multiply by the value vectors
- obtain the updated context-aware representation
This is the basic numeric recipe behind the whole layer. The visual examples above show what each step means in practice: scores tell us which positions are relevant, weights turn those scores into a distribution, and values carry the actual information that gets mixed into the output.
17. The central intuition to remember¶
If you remember only one idea from this chapter, remember this:
Attention is a learned way to route information from the context into each token’s representation.
It is the place where a token stops being “just itself” and becomes a contextual state shaped by the sentence around it.
That is why attention sits at the center of the Transformer. It is not a small add-on. It is the mechanism that allows a model to build meaning from sequence structure.
The next chapter puts this mechanism into a full Transformer block with residual connections, layer normalization, and feed-forward layers.