# How Attention Works in Transformer Architecture Explained

> Learn how attention mechanism works in transformer architecture, from query-key-value intuition and formulas to multi-head attention in a worked example.

- Author: Swarnava Dutta (https://swarnava.dev)
- Published: 2026-07-28
- Tags: Attention Mechanism Formula, Scaled Dot Product Attention, Query Key Value Attention
- Reading time: 10 min (2306 words)
- Canonical: https://swarnava.dev/blogs/transformer-attention-mechanism-explained

---

![Illustration of how attention mechanism works in transformer architecture: A wide spotlight rail: one uniquely-shaped](/images/blogs/transformer-attention-mechanism-explained-hero.jpg)

I once spent an embarrassing chunk of a weekend staring at a shape mismatch error - `RuntimeError: mat1 and mat2 shapes cannot be multiplied` - because I'd transposed my key matrix in the wrong dimension while hand-rolling attention for a toy model. The fix was a one-line `.transpose(-2,-1)`, but the real cost was that I'd been treating attention as a black box formula I'd memorized rather than something I understood. That's the trap most people fall into with how the attention mechanism works in transformer architecture: they can recite softmax(QK^T/√d_k)V but couldn't tell you why that particular arrangement of matrix multiplies does anything useful.

The intuition that finally made it click for me is simpler than the notation suggests. Attention is differentiable information retrieval - every token asks a question, scores everything else in the sequence, and blends the answers by relevance. Once you see it that way, the query-key-value split stops being arbitrary and starts being obvious.

This piece walks through that intuition, the formula itself, and a small worked example you can trace by hand before you ever touch a framework.

## How the Attention Mechanism Works in Transformer Architecture

At its core, attention is a learned, context-dependent lookup: for every token, the model decides which other tokens matter and by how much, then blends their representations accordingly. Nothing about that lookup is fixed in advance - the weighting itself is learned during training, which is what separates it from a static averaging or convolution window.

The sequence is always the same four moves:

- Project each token into a query, a key, and a value.
- Score compatibility between queries and keys.
- Normalize those scores into a probability-like distribution.
- Use the distribution to combine the values into a new representation.

Attention isn't a standalone neural network - it's a differentiable operation built from a handful of learned linear projections and a softmax, dropped inside a larger architecture that trains it end to end.

### Where Attention Fits Inside a Transformer Block

A transformer block wraps self-attention in residual connections and layer normalization, then hands the output to a position-independent feed-forward network. Embeddings plus positional information go in first, since attention alone has no sense of token order. Encoder blocks use unmasked self-attention, decoder blocks mask future tokens, and encoder-decoder setups add cross-attention where queries come from the decoder and keys/values from the encoder. Attention mixes information *across* tokens; the feed-forward layer transforms each token *independently* - that division of labor is what makes the block work.

## Query-Key-Value Attention: What Q, K, and V Represent

The search analogy is the fastest way in. Every token emits a **query** - "what am I looking for?" - while every token also emits a **key** - "here's what I contain." The **value** is the actual payload handed over once a match is found, the way a search engine returns a document's content rather than its index entry.

None of these roles are hand-assigned. An embedding doesn't "know" it's a key any more than a raw pixel knows it's an edge - a learned matrix projects it there. Three weight matrices, W_Q, W_K, and W_V, each turn the same input embedding into a different vector, and training shapes those matrices until the projections behave usefully.

For a batch of sequences shaped `(batch, seq_len, d_model)`, each projection typically maps to `(batch, seq_len, d_k)`. Queries and keys must share that last dimension, d_k, because the compatibility score between them is a dot product - mismatched dimensions are exactly the shape error I mentioned earlier.

### Queries, Keys, and Values Come From the Same Sequence

Concretely: Q = XW_Q, K = XW_K, V = XW_V, where X is the input sequence embedding matrix. When Q, K, and V all derive from the *same* X, that's self-attention - every token looks at every other token in its own sequence. Cross-attention breaks that symmetry: queries come from one sequence (say, a decoder), while keys and values come from another (the encoder output), letting one sequence attend over a different one entirely.

## Scaled Dot-Product Attention Formula, Step by Step

The full equation is Attention(Q,K,V) = softmax(QKᵀ/√d_k + M)V, and every symbol maps to a concrete tensor shape. Q is `(seq_len, d_k)`, K is `(seq_len, d_k)`, so QKᵀ produces `(seq_len, seq_len)` - one compatibility score for every query-key pair. That score matrix is the heart of it: row i tells you how much token i's query resonates with every other token's key.

Dividing by √d_k rescales those raw scores, M optionally masks some of them, and softmax turns each row into a probability distribution summing to 1. Multiplying that `(seq_len, seq_len)` weight matrix by V, shaped `(seq_len, d_v)`, produces the output - `(seq_len, d_v)` - a contextualized blend of every token's value, weighted by relevance to the query.

![Flow diagram showing queries and keys multiplied, scaled, masked, softmaxed, then combined with values to form contextual output](/images/blogs/transformer-attention-mechanism-explained-diagram-1.jpg "The scaled dot-product attention pipeline")

### Why Divide Attention Scores by the Square Root of d_k?

As d_k grows, dot products between random vectors tend to grow in variance too, since you're summing more independent products. Left unscaled, some logits blow up, softmax saturates toward one-hot outputs, and gradients through the attention weights shrink toward nothing. Dividing by √d_k tames the magnitude without touching the relative ranking - the largest dot product stays largest, softmax just gets a fairer, less saturated input to work with.

### How Padding and Causal Masks Change the Scores

Masking happens *before* softmax, not after - add a large negative number (often -1e9) to forbidden positions so their post-softmax weight rounds to zero. Padding masks blank out pad tokens so real words don't attend to filler; causal masks additionally zero out any position ahead of the current token, forcing decoders to predict without peeking. Get the axis wrong, or mask post-softmax, and you'll silently leak future tokens into training - a bug that trains fine and only shows up as suspiciously good validation loss.

## A Worked Attention Mechanism Example With Small Matrices

Take a three-token sequence, each embedded into a two-dimensional space after projection, so Q, K, and V are all 3×2 matrices. Say token 1's query is [1, 0], and the keys for tokens 1-3 are [1, 0], [0, 1], and [1, 1]. The dot products QKᵀ for that row come out to 1, 0, and 1 - token 1 resembles itself and token 3 more than token 2.

Scale by √d_k (here √2 ≈ 1.41), giving roughly 0.71, 0, 0.71, then softmax that row into weights near 0.42, 0.16, 0.42. Multiply those weights against the corresponding value vectors and sum - token 1's new representation becomes a blend, dominated by tokens 1 and 3, barely touching token 2.

Nothing here is permanent word importance - it's routing specific to this query. Swap token 1's query to [0, 1] instead, keep every key and value fixed, and the weights invert: token 2 now dominates. Same sequence, same keys, same values - different question, different answer.

## From One Attention Head to Multi-Head Self-Attention

One attention head learns one notion of compatibility. Multi-head self-attention runs several of those heads in parallel, each with its own W_Q, W_K, W_V, so each head can specialize on a different relationship in the same sequence.

Trace it through the worked example: head 1 might learn the [1,0]-style query that binds token 1 to token 3, while head 2, with different weight matrices, could bind token 1 more strongly to token 2 instead. Each head produces its own contextualized output, then you concatenate all of them along the feature axis and pass the result through a final output projection, W_O, back down to d_model.

![Comparison diagram of single-head attention versus multi-head attention where separate heads run in parallel then concatenate and project](/images/blogs/transformer-attention-mechanism-explained-diagram-2.jpg "Single head vs multi-head attention")

Standard practice splits d_model across heads evenly - 512 dimensions across 8 heads gives d_k = 64 per head - so total compute stays roughly comparable to one large head. Heads don't reliably map onto human labels like "subject-verb" or "coreference," even when visualizations tempt you to describe them that way.

### Why Multiple Heads Help - and What They Do Not Guarantee

Running subspaces in parallel lets the model route syntax, position, and semantic similarity simultaneously instead of squeezing everything through one score matrix. In practice, some heads specialize sharply while others turn out redundant - pruning studies have shown many heads can be removed at inference with little accuracy loss [[1]](#ref-1). More heads isn't automatically better; past a point you're paying compute for redundancy, not capability, and architecture search - not head count alone - decides real gains, a lesson worth remembering before you go [comparing transformers to state space models](/blogs/state-space-models-vs-transformers).

## Attention Mechanism vs Self-Attention, Cross-Attention, and RNNs

"Attention" is the general operation: a query attends to some set of keys and values, wherever those come from. Self-attention is the special case where all three - Q, K, and V - come from the same sequence, which is what lets every token look at every other token in its own input.

Cross-attention breaks that symmetry on purpose. The decoder generates queries from its own partial output, while keys and values come from the encoder's final representations, so translation or captioning models can ask "what in the source matters for this word" without folding the source back into their own self-attention pass.

Compared to an RNN, the difference is structural, not incremental. A recurrent network updates one hidden state sequentially, token by token, so information from token 1 has to survive dozens of update steps to reach token 50. Self-attention connects any two tokens in one hop, which is why it handles long-range dependencies so much more gracefully - at the cost of a score matrix that grows quadratically with sequence length, the exact trade-off projects like [FlashAttention](/blogs/is-flash-attention-stable-production-guide) exist to soften.

## Why the Attention Mechanism Works - and Where Intuition Can Mislead

Attention works because it's content-addressable, differentiable routing - every projection matrix and score gets nudged by backpropagation toward whatever weighting minimizes the training loss, not toward some hand-designed notion of relevance. No one tells W_Q to detect subject-verb pairs; gradient descent discovers that shape because it happens to help predict the next token.

That dynamic weighting is what produces genuinely contextual representations. The word "bank" gets a different output vector next to "river" than next to "loan," because its query pulls different keys into focus each time - same embedding in, different blend out, depending entirely on context.

But attention alone doesn't supply everything the block needs. Positional encodings inject order, since raw attention is permutation-invariant; residual connections preserve earlier signal so deep stacks don't collapse; feed-forward layers add per-token nonlinearity attention itself can't. Training data, not the mechanism, teaches reasoning.

Two warnings worth keeping close: attention weights are not reliable explanations of model decisions - high weight doesn't guarantee causal importance [[2]](#ref-2). And the mechanism still scales quadratically, can spread weight too diffusely over long context, and stays brittle to masking mistakes.

## Implementing Transformer Attention Without Shape or Masking Bugs

Every self-attention bug I've hit traces back to shapes, so this is where I'd start if you're wiring one up yourself:

```
Q, K, V = X @ Wq, X @ Wk, X @ Wv         # (batch, seq, d_model)
Q, K, V = split_heads(Q, K, V)          # (batch, heads, seq, d_k)
scores = Q @ K.transpose(-2, -1) / sqrt(d_k)   # (batch, heads, seq, seq)
scores = scores.masked_fill(mask == 0, -1e9)
weights = softmax(scores, dim=-1)
weights = dropout(weights)
out = weights @ V                       # (batch, heads, seq, d_k)
out = concat_heads(out)                 # (batch, seq, d_model)
out = out @ Wo
```

Frequent breakage points: forgetting `transpose(-2,-1)` on K, softmaxing over the wrong axis (it must be the last, key, dimension), masking after softmax instead of before, and dropping the final W_o projection entirely.

Sanity checks worth automating: assert each attention row sums to 1, feed a shuffled sequence through a causal mask to confirm future positions stay invisible, and diff your output against PyTorch's `scaled_dot_product_attention` on identical inputs.

## FAQ

### What is an attention mechanism?

It's a learned mechanism that lets a model weigh how much each element in a sequence should influence the representation of another element. Instead of treating every input as equally relevant, the model computes a relevance score for each pair and blends information accordingly. That blend - not a fixed rule - gets learned from data during training.

### Why does the attention mechanism work?

Because the weighting is differentiable, gradient descent can shape the query, key, and value projections toward whatever routing minimizes prediction error. Nobody hand-codes "verbs attend to subjects" - the model discovers useful patterns because they lower loss. The result is context-sensitive representations: the same word gets a different output vector depending on what surrounds it.

### Is an attention mechanism a neural network?

Not on its own - it's a differentiable operation, a handful of matrix multiplies and a softmax, not a standalone network with its own training objective. It only becomes useful embedded inside a larger architecture, like a transformer block, that trains those projection matrices end to end alongside everything else.

### How does the attention mechanism work in transformer architecture?

Each token gets projected into a query, key, and value; queries and keys are scored via dot product, scaled by √d_k, masked if needed, and passed through softmax; the resulting weights blend the values into a new, contextualized representation. Transformers run several of these heads in parallel, each specializing in different relationships, then concatenate and reproject the results.

### What is the difference between attention and self-attention?

Attention is the general mechanism - a query attends over some set of keys and values, which can come from anywhere. Self-attention is the specific case where the query, key, and value all come from the same sequence, letting every token look at every other token in its own input.


## References

1. [Single Headed Attention RNN: Stop Thinking With Your Head](https://arxiv.org/abs/1911.11423v2) - Stephen Merity (2019)
2. [Toward Faithful Explanations in Acoustic Anomaly Detection](https://arxiv.org/abs/2601.12660v1) - Maab Elrashid, Anthony Deschênes, Cem Subakan et al. (2026)