Attention Mechanism

1. Why a Token Needs Context

Consider the sentence “Maya lent Ravi her notebook.” The word “her” cannot identify a person on its own. Interpreting it requires information from the surrounding sentence. A model that processed every word independently would have no way to connect the pronoun with someone already mentioned.

Rather than copying one earlier word, attention combines numerical information from several positions. The contribution from each position can be different. The representation at “her” can therefore depend on Maya, Ravi, and the action connecting them.

This operation is called self-attention when the receiving positions and the positions supplying information belong to the same sequence. “Self” refers to the sequence attending to itself, not to each word looking only at its own position.

2. From Text to Vectors

Before attention starts, a tokenizer splits the text into tokens. We will treat each word in our sentence as one token to keep the example readable; an actual tokenizer may split it differently.

Each token has an integer ID. An embedding layer uses that ID to look up a learned vector, an ordered list of numbers. These numbers are the starting representation of the token. Individual coordinates usually do not correspond to tidy labels such as “person” or “ownership.”

The Transformer updates a vector at every position as the text passes through its layers. These intermediate vectors are called hidden states. A hidden state in a later layer can already contain information gathered from earlier layers, so attention is not always comparing the original word embeddings.

To compute over a whole sequence, we stack its vectors into a matrix called $X$. A matrix is simply a rectangular arrangement of numbers: here, one row per token and one column per feature. If there are five tokens and each vector has 64 numbers, $X$ has shape $5 \times 64$. The width of the model's hidden states is usually written as $d_{\text{model}}$.

Word order also needs to be represented. Transformers include positional information, either in the input vectors or in the attention computation. We will return to its role once the query and key vectors are introduced.

3. Queries, Keys, and Values

Attention has two jobs: decide how strongly positions should be connected, and collect the information carried along those connections. It creates three vectors from each input hidden state to handle these jobs. They are the query, key, and value.

Each vector is produced by a linear projection: multiply the input vector by a learned weight matrix. For an input row $x_i$ at position $i$, the three projections are:

The matrices $W_Q$, $W_K$, and $W_V$ are learned during training and reused at every token position in that head and layer. The resulting query, key, and value vectors change with the input. Bias terms are omitted here to keep the notation focused on the main operation.

Keeping these projections separate lets the model learn one representation for deciding relevance and another for carrying information. At “her,” its query is compared with the keys of allowed positions; the resulting weights will be applied to their values. No human assigns a rule saying that this query must select a particular name.

4. Turning Scores into Attention Weights

To compare a query with a key, standard Transformer attention uses their dot product. Multiply matching coordinates and add the results. For example, the query $[2, 1]$ and key $[2, 0]$ give $2 \times 2 + 1 \times 0 = 4$. A larger score gives that position more influence after normalization, relative to the other allowed positions.

Let $d_k$ be the number of coordinates in each query and key. The dot product is divided by $\sqrt{d_k}$. Adding products across a wider vector can produce larger score magnitudes; this scaling helps keep their range manageable. Our two-coordinate example produces $4 / \sqrt{2} \approx 2.83$.

These scores are not yet mixing weights: they can be negative, and they do not sum to one. Softmax converts the allowed scores into positive weights whose total is one. It exponentiates each score and divides by the sum of those exponentials. Higher scores receive larger weights, while other positions can still contribute.

An attention weight describes how much a value contributes to this head's mixture. A weight of 0.6 is not a 60% probability that a word is correct, nor a confidence score for the final answer. It belongs to an intermediate computation inside one layer.

5. Why a Language Model Hides Future Tokens

A model generating text from left to right predicts the next token using what is already available. In our example, the hidden state at “her” will eventually help predict “notebook.” If that state could read “notebook” during training, it would have access to the answer it was supposed to predict.

A causal mask prevents this. Position $i$ can attend to itself and all earlier positions, but not to any position after it. Future scores are replaced by negative infinity before softmax, giving them exactly zero weight. The remaining weights are normalized over the allowed positions.

The diagonal remains available because the current input token is already known. The prediction at that position is for the next token. This one-position shift is what allows a causal model to use “her” while still hiding “notebook.”

An encoder reading an already supplied sentence can use bidirectional self-attention instead, allowing both earlier and later positions. This difference in information access is part of the encoder and decoder distinction described in the Hugging Face Transformer course. Separate masks can also exclude padding tokens used to make examples the same length in a batch.

6. Combining the Values

Once the weights are available, each value vector is multiplied by its weight and the results are added coordinate by coordinate. This produces one output vector for the receiving position. A value with twice the weight contributes twice as much of that particular vector to the sum.

For example, mixing the values $[2, 0]$ and $[0, 2]$ with weights 0.75 and 0.25 gives $[1.5, 0.5]$. The output keeps both contributions; attention does not have to select a single winning token.

The first position is a useful boundary case. Under a causal mask, “Maya” has only its own value available, so its sole weight is one. Later positions have more values to combine. Changing the receiving token in the example changes the query and the set of allowed sources, which changes the output.

Attention weights alone do not fully explain a model's decision. The value vectors, the other heads, and the rest of the network also affect what reaches the final prediction. The visual shows the arithmetic of one head, not a complete account of how a trained model interprets a pronoun.

7. Computing All Positions Together

The same operations can be carried out for every query at once. Stack the query vectors into $Q$, keys into $K$, and values into $V$. Their projections are $Q = XW_Q$, $K = XW_K$, and $V = XW_V$.

Transposing $K$, written $K^\top$, swaps its rows and columns. Multiplying $QK^\top$ then computes every query-key dot product. Each row belongs to one query position, and each column belongs to one source position. The complete single-head operation is:

The mask matrix $M$ adds zero at allowed positions and negative infinity at blocked positions. Softmax runs across columns within each row. Multiplying its result by $V$ performs the weighted sums. This is the scaled dot-product attention formulation introduced in Attention Is All You Need, with the mask shown explicitly.

Attention preserves the number of positions. Five input positions produce five output vectors for this head. Query and key widths must match for the dot product; the value width can be different, although using the same width is common.

8. Why Transformers Use Multiple Heads

One attention head produces one set of mixing weights at each position. A token may need several kinds of context, and those need not favor the same source positions. Multi-head attention runs several attention computations in parallel, each with its own learned projections and its own set of weights.

Concatenation places the head outputs side by side into a longer vector. A learned output matrix, $W_O$, mixes those features and maps them back to the model's hidden-state width. For example, eight heads producing 64 coordinates each give 512 coordinates before this projection.

Heads are not assigned fixed roles such as “grammar” or “pronouns” by the architecture. Their behavior develops through training and may be distributed or redundant. The standard multi-head setup described here gives each head its own query, key, and value projections; variants such as grouped-query attention share some of the key and value projections.

Positional information influences these comparisons too. With Rotary Position Embedding (RoPE), query and key vectors are rotated according to their token positions before scoring, making their dot products sensitive to relative position. Other Transformers use different positional encodings. The causal mask specifies which connections are allowed, while positional encoding helps represent order and distance.

9. Where Attention Fits in a Transformer Block

A Transformer is built from repeated blocks. In a common decoder-only, pre-normalization design, a block contains an attention sublayer followed by a feed-forward sublayer. Both contribute updates to the stream of hidden states moving through the model, often called the residual stream.

Before attention, a normalization operation such as LayerNorm or RMSNorm controls the scale of each token's features. Attention uses those normalized vectors to produce its update. A residual connection adds the update to the block's incoming hidden states, preserving a direct path for information through the network.

Here, $X$ is the block input, $H$ is the state after the attention update, and $Y$ is the block output. The multi-head operation includes the output projection. The exact placement of normalization varies across architectures; these equations describe this particular pre-normalization arrangement.

The feed-forward network (FFN) applies learned linear transformations and a nonlinear activation or gate independently at every position. Attention mixes information across positions; the FFN transforms the features at each position after that mixing. Its weights are shared across positions within the layer, but it does not directly read other positions in this step.

After many blocks, the model applies its final normalization and output projection to produce a score for each possible next token in the vocabulary. A separate softmax turns those vocabulary scores into next-token probabilities. That final softmax serves a different purpose from the attention softmax, which distributes weights across positions in the input.

10. Attention During Training and Generation

During training, all tokens in an example are already available. The model computes the query-key comparisons for all positions in parallel, while the causal mask stops later tokens from influencing earlier predictions. The next-token loss sends gradients through the value mixtures, attention weights, and projection matrices. This is how the model learns useful connections without being given a separate label for every attention weight.

At generation time, the prompt is processed first, a phase called prefill. The final prompt position produces the distribution for the first new token. Once that token is selected, it becomes another input position, and the model computes the distribution for the following token. New tokens depend on earlier generated tokens, so this part proceeds sequentially.

A KV cache stores the keys and values already computed for past positions in each attention layer. A new position supplies a new query and can reuse those stored keys and values instead of recomputing the earlier prefix. It still needs to compare its query with the available keys and combine their values. The cache saves repeated work; it does not remove attention over the context.

For a full sequence of length $n$, dense attention has roughly $n^2$ query-key pairs per head, or $n(n+1)/2$ allowed pairs with a causal mask. Doubling the length therefore gives about four times as many pairs. A single cached decoding step instead computes one new query row over the growing context.

A straightforward implementation stores the full score and weight matrices. FlashAttention reorganizes the computation into tiles so that these large matrices do not need to be stored in GPU main memory. It computes the same attention operation, up to floating-point differences, while reducing memory traffic. Dense attention's pairwise arithmetic still grows quadratically during full-sequence processing.

11. A Small Implementation

The PyTorch function below implements one causal attention head for one sequence. It starts after the query, key, and value projections, assumes equal sequence lengths with at least one token, and leaves out batching, padding, positional encoding, and dropout. Keeping those details separate makes the attention calculation visible.

import math
import torch

def causal_attention(q, k, v):
    # One sequence and one head, with no padding or KV cache.
    # q, k: [tokens, key_width]; v: [tokens, value_width]
    scores = (q @ k.T) / math.sqrt(q.shape[-1])

    positions = torch.arange(q.shape[0], device=q.device)
    future = positions[None, :] > positions[:, None]
    scores = scores.masked_fill(future, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights

The mask compares each column's position with the row's position, blocking columns that lie in the future. The final dimension in softmax is the key-position axis, so every query gets its own normalized row. The returned output has one row per input token and the same feature width as the values.

Production implementations generally use optimized operations such as PyTorch's scaled_dot_product_attention. Its causal option matches this triangular mask when query and key sequence lengths are equal. Cached decoding can have different query and key lengths and needs a mask aligned with the positions in the cache, so this small function is not a complete cached decoder.

Following one row through this calculation is enough to connect the notation with the operation: a query is scored against allowed keys, the scores become weights, and the weights combine values into an output vector. The surrounding projections, heads, and Transformer blocks build on that computation to produce context-dependent token representations.