Attention: How Each Token Decides Which Other Tokens Matter

In AI Foundations #10, machaMochaLatte followed text through tokenization:

text -> tokens -> token IDs -> embeddings -> vectors

Now I finally have the question I kept postponing.

Suppose the model has a vector for every token. How does the vector for one token use information from the other tokens around it?

That is the job of attention.

The same word can need different information

Consider these two sentences:

The bank approved my mortgage.

We sat on the bank of the river.

The token bank starts from a learned embedding associated with that token. But its useful meaning in the sentence depends on context.

In the first sentence, mortgage is highly informative. In the second, river is.

A model therefore needs a mechanism that can ask something like:

For this token, which other positions contain useful information, and how much should each one contribute?

Attention turns that question into matrix operations.

Three versions of each token: query, key, value

This terminology looked unnecessarily mysterious to me at first.

For each input vector, the attention layer creates three projected vectors:

Q = query
K = key
V = value

They are produced by learned weight matrices. In simplified form:

Q = XWq
K = XWk
V = XWv

where X contains the token representations and Wq, Wk, and Wv are learned parameters.

I find it easier to think about their jobs rather than their names.

A query represents what the current position is looking for.

A key represents what another position can be matched on.

A value contains the information that position can contribute if it receives attention.

The query and key decide how much to attend. The value supplies what information gets mixed in.

Query meets key

For one token to decide how relevant another token is, attention compares its query with the other token’s key.

The standard Transformer uses a dot product.

A larger compatible dot product generally means a larger attention score.

For a tiny imaginary example, suppose the token bank produces scores against nearby tokens like this:

The        0.2
bank       0.5
approved   1.1
my         0.3
mortgage   2.4

These are not probabilities yet. They are raw compatibility scores.

The important idea is that mortgage can receive a larger score because the learned query/key projections make it useful to the current position.

Why divide by the square root of dimension?

The Transformer paper defines scaled dot-product attention as:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

The sqrt(d_k) part puzzled me.

As the key/query dimension grows, dot products can grow in magnitude. Very large values can push softmax into extremely sharp regions, which makes optimization less well behaved.

Dividing by the square root of the key dimension keeps the score scale under better control.

This is why the operation is called scaled dot-product attention.

Softmax turns scores into weights

After computing the scaled scores, softmax converts them into non-negative weights that sum to 1 across the allowed positions.

Our imaginary scores might become something like:

The         0.05
bank        0.07
approved    0.13
my          0.06
mortgage    0.69

These numbers are only an illustration, not values from a real model.

Now attention can take a weighted combination of the value vectors.

Conceptually:

new representation
= 0.05 × value(The)
+ 0.07 × value(bank)
+ 0.13 × value(approved)
+ 0.06 × value(my)
+ 0.69 × value(mortgage)

The result is no longer just the original embedding for bank. It now contains information gathered from the sentence.

That is the key transition:

static token representation
-> attention over context
-> context-dependent representation

Attention does not literally retrieve words

The analogy can become misleading if I take it too literally.

The model is not running a database search for the word mortgage.

Queries, keys, and values are vectors learned during training. Their useful dimensions are not manually labeled with human concepts such as financial_word = true.

Training adjusts the projection matrices through the same machinery we already studied:

forward pass
-> loss
-> backpropagation
-> gradients
-> parameter updates

So the model learns attention patterns because they help reduce training loss.

The matrix view

Doing this token by token would be cumbersome. The elegant part is that the comparisons can be expressed with matrices.

If Q contains all query vectors and K contains all key vectors, then:

QK^T

produces a matrix of query-key scores.

Each row answers:

For this query position, how compatible is every key position?

After scaling, masking, and softmax, we get an attention-weight matrix. Multiplying it by V produces the weighted combinations for all positions.

This is one reason our earlier lessons on tensors and matrix multiplication mattered. Attention sounds linguistic when explained with sentences, but the implementation is tensor algebra.

Why masking matters for a language model

During next-token prediction, a causal language model must not look into the future.

If the training sequence is:

The cat sat on the mat

when computing the representation at sat, the model should not be allowed to use on the mat as future information for that prediction step.

A causal mask blocks attention to positions that come later in the sequence.

Conceptually the allowed pattern looks like:

position 1 -> 1
position 2 -> 1,2
position 3 -> 1,2,3
position 4 -> 1,2,3,4

The masked scores are excluded before the final attention weighting.

This small detail is fundamental to autoregressive language modeling.

One attention pattern is not enough

A sentence can contain many relationships at once.

One position may benefit from grammatical information. Another interaction may track a name, a previous object, punctuation, or something much harder to describe cleanly.

Transformers therefore use multiple attention heads.

Each head has its own learned projections and can form a different pattern of query-key compatibility.

Their outputs are combined so the layer can use several kinds of relationships in parallel.

I would not interpret each head as having one simple human-readable job. Some heads can show recognizable patterns, but the network is free to distribute computation across them.

Attention is not the whole Transformer

This distinction matters because I used to treat “attention” and “Transformer” as synonyms.

They are not.

Attention is a major component inside a Transformer block. A Transformer also contains other operations, including feed-forward/MLP computation, normalization, residual connections, and positional information.

So our curriculum has now reached:

text
-> tokenizer
-> token IDs
-> embeddings
-> token vectors
-> attention
-> context-dependent vectors

But we still have not assembled the full machine.

That is the next lesson.

The mental model I am keeping

If I have to remember attention without the equations, I use this:

query: what am I looking for?
key: what can I match against?
value: what information do I receive?

Then:

query × keys
-> relevance scores
-> scale + mask + softmax
-> attention weights
-> weighted values
-> context-dependent output

The equation is compact. The consequence is enormous: every token representation can be rebuilt using information from other relevant positions.

Next, we can finally put attention together with the other pieces and ask what a Transformer block actually does.