Skip to main content
Abhik Sarkar's avatar

Multi-Head Attention

How multi-head attention runs scaled dot-product attention in parallel across several representation subspaces to build context-aware token embeddings.

One query, several attention calculations

Self-attention updates a token by combining value vectors from the same sequence. Multi-head attention runs several such calculations in parallel. Each head has its own learned query, key, and value projections, so the same token can produce a different mixture in each head.

Follow one query through those mixtures and into the combined output. The figure keeps both head outputs visible while you inspect either head. Selecting a head changes what you examine; it does not remove the other head from the computation.

Project the complete input in every head

For input states X \in \mathbb{R}^{n \times d_{\text{model}}}, each head projects the full model-width vector:

Q_i = XW_i^Q,\qquad K_i = XW_i^K,\qquad V_i = XW_i^V

The query and key projection matrices, W_i^Q and W_i^K, have shape d_{\text{model}} \times d_k; the value projection W_i^V has shape d_{\text{model}} \times d_v. The resulting Q_i and K_i have shape n \times d_k, while V_i has shape n \times d_v. A common choice is d_k = d_v = d_{\text{model}} / h for h heads. Packed implementations can combine these projections into larger matrices and then reshape the projected result into heads. They do not first give each head a disjoint slice of the raw input.

The example has four input tokens, model width 4, and two heads of width 2. The fixed vectors and matrices are invented teaching values, not a trained language model. Every token is visible to every query, including later positions; a causal decoder would apply a mask before softmax.

Start with the like query. Use Next or Calculation stage to follow projections, scores, weights, value mixing, concatenation, and the output projection. Switch Inspect head to compare the two calculations, then switch Query token and watch both head outputs change. Reset example restores the same values every time.

Normalize within each head, then mix values

For the selected token, let q_i be its projected query in head i. That head computes:

z_i = \operatorname{softmax}\!\left(\frac{q_iK_i^\top}{\sqrt{d_k}}\right)V_i

The softmax normalizes across allowed key positions within this head. Its weights sum to one independently of the other heads. The resulting z_i is a weighted sum of this head’s value vectors, with d_v components. The query and key vectors determine the weights; the value vectors provide what is combined.

Different projections can produce different distributions, but a head is not assigned a guaranteed role such as “grammar” or “meaning.” A large attention weight in this toy example also does not establish a semantic relationship. The same calculation for one head is explained in scaled dot-product attention.

Concatenate first, project once

The head outputs are placed next to one another in a fixed order, then mixed by the learned output matrix:

y = \operatorname{Concat}(z_0,\ldots,z_{h-1})W^O,\qquad W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}

The figure numbers heads from zero. Its two 2-component outputs concatenate into a 4-component vector, and a 4 × 4 output matrix returns a 4-component attention output. Inspect the fixed matrices to see how each output component receives contributions from both heads.

This is the attention branch output for one token. A full Transformer block also includes its residual connection, normalization, and a feed-forward layer. In practice, the same head calculations run for all query tokens together.

At fixed model width with d_k = d_v = d_{\text{model}}/h, the dense attention arithmetic across heads remains O(n^2 d_{\text{model}}). Separate projections and the output projection also cost arithmetic. More heads are not a free improvement: a materialized attention matrix has h n^2 entries, and narrower head vectors change what each head can represent. Attention Is All You Need, §3.2.2 defines these projections, dimensions, and concatenation.

Key Concepts & Further Exploration

  • Scaled Dot-Product Attention: The core calculation involving Q, K, V, scaling, and softmax.
  • Multi-Head: Separate learned projections let each head form its own query, keys, values, and context vector. Concatenation preserves the individual outputs before W_O mixes them.
  • Contextualization: The attention output combines information from the allowed sequence positions for each query. The surrounding Transformer block uses that output to update its hidden state.
  • Efficiency: Multi-query attention shares one K/V projection across all heads to shrink the KV cache at inference.
Transformers & LLMs
Positional Embeddings in Vision Transformers

Move the same image patches with and without positional addition to see how learned absolute positions change a Vision Transformer’s attention inputs.

Transformers & LLMs
Self-Attention in Vision Transformers

Follow one image patch through Q/K/V projection, scaled scores, row-wise softmax, value mixing, and the residual update inside a Vision Transformer.

Transformers & LLMs
CLS Token in Vision Transformers

Trace how a learned CLS row joins image patches, gathers evidence through self-attention, and becomes the image-level classification readout.

Transformers & LLMs
Hierarchical Attention in Vision Transformers

Trace how local windows, shifted cross-window exchange, and patch merging turn one high-resolution token grid into a multi-scale vision hierarchy.

Deep Learning
Adaptive Tiling: Efficient Visual Token Generation

Learn adaptive tiling in vision transformers: dynamically partition images based on visual complexity to reduce token counts while preserving detail.

Deep Learning
Convolution Operation: The Foundation of CNNs

Interactive guide to convolution in CNNs: visualize sliding windows, kernels, stride, padding, and feature detection with step-by-step demos.

If you found this explanation helpful, consider sharing it with others.

Mastodon