Attention Mechanism Step by Step - A Computational Example

Introduction

The attention mechanism is the heart of the Transformer architecture and the foundation of all modern large language models (LLMs). Although the theory behind attention can seem abstract, the best way to understand it is to walk through a concrete computational example step by step.

In this post, we'll go through a complete computational example of the attention mechanism on a simple sentence, showing all matrices and numerical calculations.

Note: For clarity, we omit positional encoding and the final output projection \(W_O\) in this example. In real Transformers, both are essential — positional encoding encodes token order, and \(W_O\) projects the concatenated outputs from multiple heads into the model's embedding space.

Example: "cat chases mouse"

Consider a simple sentence consisting of three words (tokens): cat chases mouse

And the following 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 an embedding vector of dimension \(d_{model} = 2\) (in real models this is typically 512, 768 or more).

Each token has the following embeddings:

where columns correspond to successive tokens in the vocabulary.

Embeddings:

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:

Although we don't directly control what each dimension means, we can try to discover it post hoc. Online you can find many examples where embeddings for tokens like uncle and aunt are shifted by the same constant as king and queen, meaning the model encoded gender information in a specific direction in the embedding space.

In our example, one might speculate:
- First dimension — "animality" (cat=1.0, mouse=0.8, chases=0.2)
- Second dimension — "action/movement" (chases=1.0, rest=0.0)

That said, this is just an example constructed to fit the narrative. In real 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.

In standard notation, the key and value dimensions are denoted \(d_k\) and \(d_v\). With a single head (single-head attention), we have \(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.

Query weight matrix:

Key weight matrix:

Value weight matrix:

In this simple example we use identity matrices, but in real models these are learned parameters.

Step 3: Computing Q, K, V

We multiply the embeddings by the respective weight matrices:

Query:

Key:

Value:

Step 4: Attention Scores

We compute attention scores using the formula:

where \(d_k = 2\) (key dimension).

Computing \(QK^T\):

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

Step 5: Masking and Softmax

In a decoder-only architecture (e.g. GPT, Claude) we apply a causal mask, which prevents tokens from attending to future tokens — they can only attend to themselves and earlier positions. This enables autoregressive text generation.

In contrast, an encoder-only architecture (e.g. BERT) applies no mask — every token sees the entire sequence (bidirectional attention).

After masking (before softmax):

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

Applying softmax row-wise (\(-\infty\) values become 0):

Each row shows how much that token "attends" to the available tokens:
- Row 0 (cat): sees only itself (the only available token), so weight = 1.0
- Row 1 (chases): sees "cat" and itself; attends more to itself (0.65) than to "cat" (0.35)
- Row 2 (mouse): sees all three tokens; attends most to "cat" (0.40), then to itself (0.35), least to "chases" (0.25)

Step 6: Output (Contextual Representations)

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

The final output matrix contains contextual representations for each token — each vector is a weighted combination of Value vectors, where the weights depend on attention scores.

What changed?

• "cat" — unchanged \([1.0, 0.0]\), because masking means it only sees itself
• "chases" — changed from \([0.2, 1.0]\) to \([0.48, 0.65]\): "animality" increased (influence of "cat"), "action" decreased
• "mouse" — changed 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 knows it is being chased — the information encoded in the "action" dimension will be crucial for predicting the next token.

This is a simplified example, but exactly the same mechanism of context aggregation through weighted sums operates in widely used models.

Note: In the full Transformer architecture, the output is then 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)

After the attention block comes a feed-forward network (FFN), applied independently to each position. In the original Transformer architecture (Vaswani et al., 2017) it consists of two linear layers with a ReLU activation:

For simplicity, we use a minimal FFN with a single linear layer, no activation, and no bias:

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

Step 8: Predicting the Next Token

To predict the next token, we need the representation of the last token:

It's worth noting that at this stage all other token representations are no longer needed. After the attention step, all information they carry should already be encoded in the last token.

Projecting to logits (unembedding)

From the logits we can compute probabilities:

Results:

According to our simple model, the most probable next token is "chases", making the sentence:

Cat chases mouse chases

Which is complete nonsense?

Why This Does (Not) Work

Because the weights were chosen arbitrarily, not through training. I originally intended to pick weights that would give a sensible result, then decided it would be more valuable to remind you again that real models have millions of weights — so this simple example has no reason to work correctly. (Honestly, the first motivation was laziness; the educational narrative came after.)

If we do believe the mechanism works, it's worth noting that this is thanks to attention, which allows each token to "look at" the available tokens and decide which are most important for its representation.

In a decoder-only architecture with a causal mask (as in our example):
- Each token builds its representation only from itself and previous tokens
- "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 between tokens

This flexibility allows Transformers to understand the structure and semantics of text without the need to define explicit grammatical rules.

Key Properties

Scaling by \(\sqrt{d_k}\)

Dividing by \(\sqrt{d_k}\) prevents dot products from growing too large as dimensionality increases. Without this scaling, softmax could produce very extreme values (close to 0 or 1), which would hinder learning through vanishing gradients.

Softmax

The softmax function normalizes scores so that the attention weights for each token sum to 1.0. It also sharpens differences between values — larger values become even more dominant after softmax.

Q, K, V Matrices

Separating into Query, Key, and Value gives the model flexibility:
- Query: "what am I looking for?"
- Key: "what do I offer as a matching key?"
- Value: "what information do I pass on?"

A token can search for specific features (Q), be retrieved by different features (K), and transmit yet other information (V). This separation allows asymmetric relationships between tokens.

Multi-Head Attention

In practice, modern Transformers use multi-head attention, where:

1. Embeddings are projected into \(h\) different subspaces (heads)
2. Each head has its own matrices \(W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}\) with dimensions leading to \(d_k = d_v = d_{model} / h\)
3. Attention is computed in parallel for each head
4. Results are concatenated and projected through \(W_O\)

where:

This allows the model to learn different types of relationships simultaneously — one head may track syntactic relations, another semantic ones, another positional ones.

Summary

The attention mechanism is a mathematically elegant way to model 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 — optional causal mask (decoder) or no mask (encoder)
5. Softmax — normalization to probabilities
6. Weighted sum of Values — contextual representations
7. Output projection — multiplication by \(W_O\) (in the full architecture)

This mechanism, repeated many times across many layers and heads, creates the powerful Transformer architecture that has revolutionized natural language processing.

Useful Links

• Attention is All You Need (original paper)
• The Illustrated Transformer
• Visualizing Attention in Transformer Models