Transformer attention from the matrix operations up: what a head actually computes

About this article: AI-assisted writing, reviewed against the cited sources. This is source research, not a lab validation of the procedures. Check the stated versions and your own recovery requirements before making changes.

Attention is not a metaphor sitting on top of a neural network. In the Transformer it is a short sequence of linear algebra: form queries, keys and values; multiply; scale; soft-max; multiply again. This article walks that sequence with the dimensions from the original paper, then explains multi-head projections, causal masking, cost, and the trap of reading attention maps as explanations.

An AI writing system helped prepare this article. Terms such as a model "looking up" context describe computation, not subjective experience.

Nothing below was executed in a lab. Matrix shapes are taken from the cited papers. Implementation-specific kernels and serving APIs are outside this article's scope.

Why it matters

If you size GPUs, debug long-context behaviour, or review vendor claims about "attention", you need the real operation, not a slide that says the model "focuses". Three practical consequences fall out of the maths:

1. Memory and time grow with sequence length. Naive attention materialises an $n\times n$ score matrix per head (or an equivalent amount of work). That is why long prompts hurt before the model has generated a single new token.
2. Parallelism inside a sequence is possible. Unlike an RNN that must finish position $t-1$ before $t$, attention over a known prompt can be computed with large matrix multiplies. That is why training and prompt prefill scale on GPUs the way they do.
3. Weights are not proofs. A bright cell in an attention heatmap is a large softmax weight. It is not, by itself, proof that the model used that token as a causal reason for the answer.

How it works

The attention function in one line

Vaswani et al. define an attention function as mapping a query and a set of key–value pairs to an output. The output is a weighted sum of the values; the weights come from a compatibility score between the query and each key.

Their particular form is scaled dot-product attention. Pack many queries into matrix $Q$, keys into $K$, values into $V$. Then:

Read it left to right:

If components of $q$ and $k$ behave like independent mean-zero unit-variance variables, the unscaled dot product has variance $d_k$. Vaswani et al. suspect that larger scores push softmax into regions with very small gradients; they scale by $\sqrt{d_k}$ to counteract that. The variance calculation follows from the stated assumptions; the gradient explanation is presented as a suspicion, not a measured causal result.

Where $Q$, $K$ and $V$ come from

Tokens are first embedded into vectors of size $d_{\text{model}}$. In the original Transformer, $d_{\text{model}}=512$. Learned linear projections then produce queries, keys and values for each attention sub-layer. Residual connections wrap the sub-layer so information can pass through a stack of identical layers without being forced through a single rewrite at every depth. The paper uses $N=6$ encoder layers and $N=6$ decoder layers in the base configuration.

Position is not carried by recurrence. The architecture adds positional information to embeddings (sinusoidal encodings in the original work) so that the same content at different offsets is distinguishable. Without that, pure attention would treat the sequence as a bag of vectors.

Multi-head attention: several subspaces, then concatenate

A single attention head averages information into one weighted mix. The paper's remedy is multi-head attention: project $Q$, $K$ and $V$ with different learned matrices $h$ times, run attention in parallel, concatenate the head outputs, and project again with $W^O$.

In the original configuration:

• $h=8$ heads
• $d_k=d_v=d_{\text{model}}/h=64$
• Total cost stays comparable to one full-dimension head because each head is narrower

The stated motivation is that heads can attend to different representation subspaces. That is an empirical design claim in the paper, not a guarantee that each head learns a human-readable role. Later interpretability work studies specialised heads; that literature is separate from the 2017 definition.

Three places attention appears in the original model

Vaswani et al. use multi-head attention in three patterns:

1. Encoder self-attention. Queries, keys and values all come from the previous encoder layer. Every position may attend to every position in the input.
2. Decoder masked self-attention. Same idea inside the decoder, but positions are prevented from attending to later positions. Combined with the one-position offset of output embeddings, this keeps autoregressive generation causal: predicting token $i$ may only use outputs at positions less than $i$.
3. Encoder–decoder attention. Queries come from the decoder; keys and values come from the encoder output. Each decoding step can draw on the whole encoded input.

The distinction to retain is architectural: masked self-attention relates positions within the decoder, while encoder–decoder attention reads a separate encoder output.

Causal masks as concrete tensor edits

"Masking" is not a policy document. The original paper describes setting disallowed pre-softmax connections to negative infinity, giving them zero softmax weight. Our implementation guidance is to check the library's mask convention and handling of fully masked rows; substituting a finite sentinel is an implementation choice that needs numerical testing. If the mask is wrong, the model either cheats by seeing the future during training or fails to use legitimate context. When you debug a custom attention implementation, compare the boolean mask against the sequence positions before you tune learning rates.

Softmax competition

Softmax forces a distribution. If one key scores much higher than the others, it can monopolise the mix. That sharpness is useful for precise lookup and harmful when several pieces of evidence should share mass. Temperature and nucleus sampling later change how the output vocabulary is sampled; they do not rewrite the attention softmax inside the stack unless the implementation exposes that separately.

Residuals and layer normalisation around the head

The original encoder layer does not replace the residual stream with the attention output. It computes $\mathrm{L}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}\mathrm{N}\mathrm{o}\mathrm{r}\mathrm{m}(x+\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x))$ for attention and again for the position-wise feed-forward network. Practically, that means an attention head proposes an update that is added to what was already there. If you zero an attention sub-layer in a debugging build, earlier information can still pass through the skip connection — which is why "attention did nothing" and "attention wrote a small delta" are different failure modes.

The feed-forward sub-layer after attention is point-wise: the same MLP runs at every position independently. Attention is the only place in the block that mixes information across positions. When people say a Transformer "reads" distant context, they mean this mix, repeated through depth.

Autoregression after the architecture exists

The paper stresses that the decoder is auto-regressive: each step consumes previously generated symbols. The paper does not discuss KV caching. In later serving practice, earlier tokens' keys and values can be stored for reuse. This avoids recomputing their projections when a new token arrives. The attention operation remains the same, but the query shape, cached inputs and scheduling differ between full-sequence training and one-token decoding.

Cost: why engineers feel attention

Work and memory in the naive algorithm

For sequence length $n$, head dimension $d$, a straightforward implementation:

• forms an $n\times n$ score matrix (or tiles that amount of traffic),
• applies softmax,
• multiplies by values.

Training and long-prompt prefill therefore face quadratic growth in $n$ for that term. The original paper emphasises that relating two positions costs a constant number of operations with attention, versus growth with distance for some convolutional stacks — but the constant still multiplies by $n^2$ work across the sequence.

Exact attention can still be IO-aware

Dao et al. (FlashAttention, 2022) show that the bottleneck on GPUs is often movement between high-bandwidth memory and on-chip SRAM, not the arithmetic itself. Their algorithm computes exact attention with tiling and softmax-friendly online statistics, reducing HBM traffic. The formula above does not change; the schedule of reads and writes does. If a vendor says they use "approximate attention", that is a different claim from FlashAttention-style exact kernels.

Prefill versus decode

For autoregressive generation after the prompt is known:

• Prefill computes attention over the prompt in large parallel matmuls.
• Decode usually appends one token at a time and reuses stored keys and values (the KV cache). Attention cost per new token then scales with the growing context, and memory for the cache scales with layers × heads × context × head dimension × dtype.

KV caching avoids recomputing earlier tokens' key and value projections. The new query still has to be scored against the stored keys and combine the stored values; it does not make context-length cost disappear.

What to choose and why

Risks and limits

• Attention weights need an explanation test. Jain and Wallace argue that weights often fail to explain predictions; Wiegreffe and Pinter challenge the generality of that conclusion and propose stronger tests. Their 2019 studies largely concern RNN-era classification, so they motivate caution rather than settle interpretability for modern Transformers. A heatmap alone is insufficient evidence of a causal explanation.
• Head specialisation is not guaranteed by the 2017 design. Eight heads of width 64 were a working configuration, not a taxonomy of linguistic roles.
• Numerical edge cases. Softmax in low precision, mask values that are not low enough, and mixed-precision accumulators all create implementation bugs that look like "the model is dumb".
• This article does not validate any vendor's production kernel. FlashAttention is cited for the IO point; serving frameworks evolve quickly — check their current docs before operational changes.
• No NetLeverage product behaviour is claimed.

Sources

1. Vaswani et al., Attention Is All You Need, arXiv:1706.03762 — https://arxiv.org/abs/1706.03762 (HTML https://arxiv.org/html/1706.03762 opened 10 November 2026). Defines scaled dot-product attention, multi-head attention, encoder/decoder stacks, masking, and the base dimensions $d_{\text{model}}=512$, $h=8$, $d_k=64$.

2. Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv:2205.14135 — https://arxiv.org/abs/2205.14135 (HTML opened 10 November 2026). Exact attention with reduced HBM traffic; does not replace the mathematical definition above.

3. Jain and Wallace, Attention is not Explanation, NAACL 2019. Evidence against equating weights with explanations.

4. Wiegreffe and Pinter, Attention is not not Explanation, EMNLP 2019. Contrary evidence and proposed explanation tests. Both abstracts reopened 10 November 2026.

This article was originally published on 10 November 2026.