The Attention Mechanism Step by Step — A Computational Example

Introduction

The attention mechanism is the heart of the Transformer architecture and the foundation of every modern large language model. The theory behind it can feel abstract, which is why the best way to understand it is to follow a concrete computational example.

In this post we walk through the full set of calculations on a three-word sentence — from embeddings all the way to the next-token probability distribution — showing every matrix along the way.

Note: for clarity we omit positional encoding and the final output projection \(W_O\). Both are essential in real Transformers: positional encoding encodes token order, and \(W_O\) projects the concatenated outputs of multiple heads back into the model's space.

Example: "cat chases mouse"

Consider a sentence made of three tokens, cat chases mouse, and a vocabulary of size 5:

• token 0: "cat"
• token 1: "chases"
• token 2: "mouse"
• token 3: "quickly"
• token 4: "sleeping"

Step 1: Token embeddings

We represent each token as a vector of dimension \(d_{model} = 2\) (in real models this is typically 512, 768 or more). The entire vocabulary therefore fits into a \(2 \times 5\) matrix whose columns are the embeddings of successive tokens:

Our sentence consists of tokens 0, 1 and 2, so its embedding matrix \(E\) is formed by taking the first three columns of \(W_{vocab}\) and writing them as rows:

where:

• \(E[0] = [1.0,\ 0.0]\) — embedding for "cat"
• \(E[1] = [0.2,\ 1.0]\) — embedding for "chases"
• \(E[2] = [0.8,\ 0.0]\) — embedding for "mouse"

Interpreting the dimensions

We do not directly control what a single dimension means, but we can try to read it post hoc. The classic observation is that embeddings for pairs such as king and queen, or uncle and aunt, are separated by a similar vector — the model has encoded gender along a particular direction in the space.

In our example one might speculate that:

• the first dimension corresponds to "animality" (cat = 1.0, mouse = 0.8, chases = 0.2),
• the second corresponds to "action/movement" (chases = 1.0, the rest = 0.0).

These numbers were chosen to fit the story, though. In models with hundreds of dimensions, interpretation is much harder and rarely unambiguous.

Step 2: Weight matrices — Query, Key, Value

We define three weight matrices that transform embeddings into Query, Key and Value representations. The key and value dimensions are conventionally written \(d_k\) and \(d_v\). With a single head, \(d_k = d_v = d_{model}\); in multi-head attention each head operates on \(d_k = d_v = d_{model} / h\), where \(h\) is the number of heads.

We use identity matrices to keep the arithmetic readable. In real models these are learned parameters, and they are precisely what determines what each head looks for in the sequence.

Step 3: Computing Q, K, V

We multiply the embeddings by the respective weight matrices. With identity matrices, all three representations are identical to \(E\):

Step 4: Attention scores

Attention is computed with the formula which can be found in original attention article:
:

where \(d_k = 2\).

The product \(QK^T\):

Entry \((i, j)\) is the dot product of token \(i\)'s query with token \(j\)'s key — a measure of how well they match.

Scaling by \(\sqrt{d_k} = \sqrt{2} \approx 1.414\):

Step 5: Masking and softmax

In a decoder-only architecture (GPT, Claude) we apply a causal mask, which cuts off access to the future — each token sees only itself and earlier positions. This property is what makes autoregressive generation possible.

An encoder-only architecture (BERT) applies no mask: every token sees the entire sequence (bidirectional attention).

Adding the mask (before softmax):

Note: the notation \(a + (-\infty)\) is mathematically informal but is a standard implementation convention. In floating-point arithmetic, -inf is a concrete value for which \(\exp(-\infty) = 0\), zeroing out masked positions after the softmax.

Softmax is applied independently to each row:

Each row shows how much that token attends to the available positions:

• row 0 ("cat") — sees only itself, so all the weight goes to the single available token;
• row 1 ("chases") — sees "cat" and itself, attending more to itself (0.65) than to "cat" (0.35);
• row 2 ("mouse") — sees all three positions: most to "cat" (0.40), then itself (0.35), least to "chases" (0.25).

Step 6: Output — contextual representations

The final step of attention itself is multiplying the weights by the Value matrix:

The resulting matrix holds the contextual representation of each token — every row is a weighted combination of Value vectors, with the weights coming from the attention scores.

What changed?

• "cat" — unchanged at \([1.0,\ 0.0]\), since the mask means it sees only itself;
• "chases" — from \([0.2,\ 1.0]\) to \([0.48,\ 0.65]\): "animality" rose (influence of "cat"), "action" fell;
• "mouse" — from \([0.8,\ 0.0]\) to \([0.73,\ 0.25]\): an "action" component appeared (influence of "chases").

Each token has absorbed information about its context. "Mouse" now carries a trace of being chased, and that information will matter when predicting the next token. The example is simplified, but exactly the same mechanism — aggregating context through weighted sums — runs inside production models.

Note: in the full architecture the output is additionally multiplied by a projection matrix \(W_O \in \mathbb{R}^{d_v \times d_{model}}\), which we omit here.

Step 7: Feed-forward network (FFN)

The attention block is followed by a feed-forward network applied independently to each position. In the original architecture (Vaswani et al., 2017) it is two linear layers with a ReLU activation:

For simplicity we use a minimal version: one linear layer, no activation, no bias.

Note: we omit LayerNorm and residual connections, which stabilize training in a real Transformer.

Step 8: Predicting the next token

The next-token distribution is read from the representation of the last position:

This is where attention has gathered the context of the whole sentence. It is worth stressing, though, that the other positions do not become redundant: their keys and values are still needed in subsequent layers, and during training every position predicts its own next token. Reading from the last row is the generation-time picture.

Projecting to logits (unembedding)

To go from a representation to scores over the whole vocabulary, we reuse \(W_{vocab}\). Using the same matrix for both embedding and unembedding is called weight tying; it is common in real models, saves parameters and usually improves quality.

From the logits we compute probabilities:

According to our model the most likely next token is "chases", which gives the sentence:

cat chases mouse chases

The result is nonsense — and that is a good starting point for asking why.

Why this does (not) work

Because the weights did not come from training. I picked them by hand so that the arithmetic could be followed on paper. I could have tuned them toward a sensible answer, but the example would then imply something false: that two dimensions and a dozen parameters are enough to model language. They are not. Real models have billions of parameters, and it is only their joint fitting on an enormous corpus that makes predictions meaningful.

What does work regardless of the weight values is the structure of the computation. Attention lets every token look at the available positions and decide which of them matter for its own representation.

In a decoder-only architecture with a causal mask (as in our example):

• each token builds its representation only from itself and earlier positions,
• "mouse" can incorporate context from "cat" and "chases", but not the reverse,
• this enables autoregressive generation: the model predicts the next token knowing only the previous ones.

In an encoder-only architecture (no mask):

• each token sees the entire sequence,
• "cat" can attend to "mouse" and vice versa,
• the model learns bidirectional relationships.

This flexibility lets Transformers capture the structure and semantics of text without hand-written grammatical rules.

Key properties

Scaling by \(\sqrt{d_k}\)

Dividing by \(\sqrt{d_k}\) keeps dot products from growing too quickly as dimensionality increases. With large values the softmax enters its saturated regime — the distribution becomes almost one-hot and gradients vanish, which makes learning harder.

Softmax

Softmax normalizes the scores so that attention weights in each row sum to 1.0. It also sharpens differences: larger values become even more dominant after the transformation.

Q, K, V matrices

Splitting into Query, Key and Value gives the model freedom:

• Query — "what am I looking for?"
• Key — "what can I be found by?"
• Value — "what information do I pass on?"

A token can search by one set of features (Q), be retrieved by another (K) and transmit yet another (V). This separation permits asymmetric relations — A attending to B does not imply B attending to A.

Multi-head attention

In practice, modern Transformers use multi-head attention:

1. embeddings are projected into \(h\) different subspaces (heads),
2. each head has its own \(W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}\) leading to \(d_k = d_v = d_{model} / h\),
3. attention is computed in parallel in every head,
4. the results are concatenated and projected through \(W_O\).

where:

This lets the model learn several kinds of relationship at once — one head may track syntactic dependencies, another semantic ones, another positional ones.

Summary

Attention is a mathematically economical way of modelling dependencies between elements of a sequence. The key steps are:

1. embeddings — representing tokens as vectors (+ positional encoding in the full architecture),
2. Q, K, V transformations — projections into query, key and value spaces,
3. attention scores — computing \(QK^T/\sqrt{d_k}\),
4. masking — a causal mask (decoder) or none (encoder),
5. softmax — normalization into a probability distribution,
6. weighted sum of Values — contextual representations,
7. output projection — multiplication by \(W_O\) (in the full architecture).

Repeated across many layers and many heads, this mechanism forms the architecture that changed natural language processing.

Useful links

• Attention is All You Need (original paper)
• The Illustrated Transformer
• Transformer: A Novel Neural Network Architecture for Language Understanding