TechByteByByte

Multi-Head Attention

Understand why a single attention computation isn't enough, how multiple attention heads run in parallel on different learned subspaces, and how their outputs are concatenated and projected — with a fully verified two-head example.

#Transformers#Multi-Head Attention#Attention#AI#LLM

Begin with the central question

Why ask several attention questions in parallel instead of relying on one attention pattern?

Essential words

An attention head has its own learned projections. Concatenation joins head outputs. An output projection mixes the joined information back into model width.

What You Will Understand

Why Transformers don’t just run attention once — they run it several times in parallel, each with its own learned projections, then combine the results. You’ll compute two full attention heads by hand, on the same input, and see them produce genuinely different attention patterns.

input -> several attention heads -> concatenate -> output projection

The problem this module solves

Module 5’s scaled dot-product attention computes one specific kind of relevance — determined entirely by one set of Wq, Wk, Wv matrices. But language has many kinds of relationships simultaneously — a token might need to relate to a nearby token for grammatical agreement, and also to a distant token for topical coherence.

A single attention computation, with one shared projection, has to compress all of this into one relevance scoring — multi-head attention exists to let the model represent several different kinds of relationships at once.

Build the intuition

instead of one person trying to evaluate a sentence from every possible angle at once, imagine several specialists each reading the same sentence, each asking a different kind of question — then combining their separate assessments into one final judgment. Each attention head is one of these “specialists,” working from its own learned projection of the same input.

4. Real-World Analogy

Think of a hiring committee reviewing the same resume. One member focuses purely on technical skills, another on communication style, another on culture fit — each looking at the same document but through a different lens, producing a different assessment. The final hiring decision combines all these separate, parallel assessments.

Multi-head attention works the same way: the same input X is projected through several different, independently learned Wq/Wk/Wv sets, each head producing its own “assessment” of token relationships.

Analogy: The Hiring Committee Panel of Specialists Think of multi-head attention as a specialized evaluation committee reading a single applicant’s profile (the input sequence):

  • The Recruiter Panel:
    • Recruiter 1 (Head 1): Focuses solely on Java coding years. They filter out communication or degree lines, projecting the resume into a smaller technical subspace.
    • Recruiter 2 (Head 2): Focuses solely on leadership experience and communication cues (subspace 2).
    • Recruiter 3 (Head 3): Focuses solely on location compatibility (subspace 3).
  • Each recruiter writes their own independent score sheet (attention map).
  • The Decision (Merge): The department head collects the sheets, staples them side-by-side (Concatenation), and runs a final combined review (Output projection Wo) to hire the candidate.
  • The model gets multifaceted insight without any single head getting overwhelmed by details.

📊 Visual Flowchart: Multi-Head Splitting, Parallel Attention, and Projection Merge

Here is how input spaces are split into heads, processed in parallel, and consolidated:

graph TD
    X["Input X<br>(Seq Length x d_model)"] --> Head1Split["1. Head 1 Projection<br>(Q1, K1, V1 with dimension d_head)"]
    X --> Head2Split["1. Head 2 Projection<br>(Q2, K2, V2 with dimension d_head)"]




    subgraph ParallelHeads ["Independent Parallel Execution"]
        Head1Split --> Head1Attn["2. Scaled Dot-Product Attention 1<br>(Output: Seq Length x d_head)"]
        Head2Split --> Head2Attn["2. Scaled Dot-Product Attention 2<br>(Output: Seq Length x d_head)"]
    end




    Head1Attn --> Concat["3. Concatenation along features<br>(Restores shape: Seq Length x d_model)"]
    Head2Attn --> Concat




    Concat --> OutProj["4. Final Output Projection: Wo<br>(Intermixes features across heads)"]
    OutProj --> MultiHeadOut["Multi-Head Attention Output<br>(Seq Length x d_model)"]

5. Core Concept

Input X
  ↓
Q/K/V projections -- but with SEVERAL SEPARATE sets, one per head
  ↓
Head 1: attention(X @ Wq_1, X @ Wk_1, X @ Wv_1)
Head 2: attention(X @ Wq_2, X @ Wk_2, X @ Wv_2)
...
Head N: attention(X @ Wq_N, X @ Wk_N, X @ Wv_N)
  ↓
Concatenate all heads' outputs
  ↓
Output projection (Wo)
  ↓
Multi-head attention output

Each head typically works in a smaller subspace: if the full model dimension is d_model and there are num_heads heads, each head’s Wq/Wk/Wv project down to d_head = d_model / num_heads — so the concatenated output of all heads has dimension num_heads × d_head = d_model again, ready for the output projection.

⚠️ Do not claim any specific head “always” learns a specific, human-nameable relationship (like “head 3 always learns grammar”). Real trained models sometimes show heads with roughly interpretable, specialized behavior in research analysis — but this isn’t guaranteed or universal across all heads, layers, or models. The architectural capacity for heads to specialize differently is the reliable claim; exactly what each one learns is an empirical, model-specific question.

6. How It Works — Step by Step

1. Choose the number of heads (a hyperparameter, e.g., 8, 16, or
   more in real models)
2. Each head gets its OWN, independently initialized and learned
   Wq, Wk, Wv matrices, projecting from d_model down to d_head
3. Run scaled dot-product attention (Module 5) INDEPENDENTLY and
   IN PARALLEL for every head, using each head's own Q, K, V
4. CONCATENATE every head's output along the feature dimension --
   this restores the original d_model dimensionality
5. Apply a final OUTPUT PROJECTION (Wo) to the concatenated result
   -- this lets the model learn how to best COMBINE information
   from all the heads

7. Mathematical Intuition

Read the mathematics as a story

Instead of forcing one attention calculation to capture every relationship, several smaller heads learn separate projections. Their outputs are joined and mixed.

X -> head 1
  -> head 2 -> concatenate -> output projection
  -> head 3

If d_model = 512 and num_heads = 8, each head works with d_head = 512 / 8 = 64. Each head’s attention computation (Module 5) is identical in form to single-head attention — same formula, same scaling by √d_head — just operating on a smaller slice of the total representation space, with its own independently learned parameters.

8. Small Worked Example

Walk through the example

  1. Split model width across two heads. 2. Compute Q/K/V and attention independently per head. 3. Concatenate both outputs. 4. Project back to model width.

With 3 tokens (“cat”, “sat”, “mat”) and d_model=4, splitting into 2 heads gives each head d_head=2. Each head computes its own full attention pipeline — its own Q, K, V, its own scores, its own softmax weights — completely independently of the other head, using different learned projection matrices on the same input.

9. Python / NumPy Example

What the code will demonstrate

This small NumPy example makes Multi-Head Attention visible with inspectable numbers and shapes. Read it in three passes: identify each input, follow the transformation line by line, and connect the printed output to the diagram above. The arrays are intentionally tiny teaching values; unless the text explicitly says otherwise, they are not weights or measurements from GPT, Gemini, or another trained model.

# The arrays are intentionally small so each transformation can be inspected.
# Printed values illustrate the mechanism; they are not trained-model measurements.
import numpy as np




def softmax(x, axis=-1):
    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)




np.random.seed(3)




tokens = ["cat", "sat", "mat"]
X = np.array([
    [1.0, 0.0, 1.0, 0.0],
    [0.0, 1.0, 0.0, 1.0],
    [1.0, 1.0, 0.0, 0.0],
])




d_model = 4
num_heads = 2
d_head = d_model // num_heads
print(f"d_model={d_model}, num_heads={num_heads}, d_head={d_head}")




def single_head_attention(X, Wq, Wk, Wv):
    Q = X @ Wq
    K = X @ Wk
    V = X @ Wv
    scores = Q @ K.T / np.sqrt(Wq.shape[1])
    weights = softmax(scores, axis=-1)
    output = weights @ V
    return output, weights




# Head 1: its own, independently learned projection matrices
Wq1 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)
Wk1 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)
Wv1 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)




# Head 2: DIFFERENT, independently learned projection matrices
Wq2 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)
Wk2 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)
Wv2 = np.round(np.random.randn(d_model, d_head) * 0.5, 2)




head1_output, head1_weights = single_head_attention(X, Wq1, Wk1, Wv1)
head2_output, head2_weights = single_head_attention(X, Wq2, Wk2, Wv2)




print("\nHead 1 attention weights:\n", np.round(head1_weights, 4))
print("\nHead 2 attention weights:\n", np.round(head2_weights, 4))




# Concatenate heads
concatenated = np.concatenate([head1_output, head2_output], axis=-1)
print("\nConcatenated output (shape", concatenated.shape, "):\n", np.round(concatenated, 4))




# Output projection: back to d_model
Wo = np.round(np.random.randn(d_model, d_model) * 0.5, 2)
multi_head_output = concatenated @ Wo
print("\nFinal multi-head attention output (shape", multi_head_output.shape, "):\n", np.round(multi_head_output, 4))




print("\n--- What 'cat' (row 0) attends to, per head ---")
for i, tok in enumerate(tokens):
    print(f"  Head 1 -> '{tok}': {head1_weights[0, i]:.4f}   |   Head 2 -> '{tok}': {head2_weights[0, i]:.4f}")

Expected Output:

d_model=4, num_heads=2, d_head=2




Head 1 attention weights:
 [[0.472  0.2676 0.2603]
 [0.2635 0.3618 0.3746]
 [0.4565 0.2735 0.27  ]]




Head 2 attention weights:
 [[0.3627 0.1943 0.443 ]
 [0.3326 0.4491 0.2182]
 [0.3124 0.3025 0.3851]]




Concatenated output (shape (3, 4)):
 [[-0.0194 -0.8694 -0.9616 -0.3999]
 [ 0.3008 -0.8794 -0.8637 -0.4963]
 [ 0.0033 -0.8711 -0.9156 -0.4626]]




Final multi-head attention output (shape (3, 4)):
 [[-1.0682  0.1421 -0.3248  1.0458]
 [-1.4414 -0.0608 -0.1886  1.0093]
 [-1.1078  0.0533 -0.322   0.9565]]




--- What 'cat' (row 0) attends to, per head ---
  Head 1 -> 'cat': 0.4720   |   Head 2 -> 'cat': 0.3627
  Head 1 -> 'sat': 0.2676   |   Head 2 -> 'sat': 0.2603
  Head 1 -> 'mat': 0.2603   |   Head 2 -> 'mat': 0.4430

10. How It Works

  • Head 1 and Head 2 produce genuinely different attention weight patterns for the same token “cat” — Head 1 attends most to “cat” itself (0.472), while Head 2 attends most to “mat” (0.443) — concrete proof that different heads, using different learned projections, capture different relevance patterns from the identical input.
  • Each head’s output has shape (3, 2) — 3 tokens, d_head=2 dimensions. Concatenating both heads restores (3, 4) — the original d_model.
  • The final output projection (Wo) mixes information across both heads’ concatenated outputs, producing the final (3, 4) multi-head attention result — this is what actually feeds into the rest of the Transformer block (Module 9).

11. How Is This Used in Modern AI?

Where this concept lives

Follow the concept at three levels: inside the model, where the computation happens; inside the AI product, where that computation supports a visible feature; and inside production, where engineers measure speed, memory, quality, and failure cases. The details below connect those levels.

🤖 How Is This Used in Modern AI?

Real LLMs use far more heads than this toy example — commonly 8 to 96+ per layer, depending on model size — each operating on its own learned subspace, at every single Transformer block. Multi-head attention is not an optional refinement; it’s a standard, structural component of essentially every modern Transformer-based model.

12. How Is This Used in LLMs?

Trace one model call

User text → tokens → Transformer computation → output-token probabilities
          this topic affects one part of that computation

An LLM does not apply this idea as a separate magic step. It uses it as part of the repeated numerical pipeline that transforms token vectors and produces the next-token probabilities.

Every attention computation inside a real LLM is multi-head, not single-head. When you read about a model’s architecture (e.g., “32 attention heads”), this is precisely what’s being described — 32 parallel attention computations per layer, each on its own subspace, concatenated and projected exactly as demonstrated above.

Real systems you can recognize

The original Transformer used multi-head attention so different representation subspaces could be attended to in parallel. Modern models may use multi-head, grouped-query, or multi-query variants; the public API name alone does not reveal every implementation detail.

13. How Is This Used in Agentic AI?

Trace one agent step

Goal + history + tool results
              ↓
     LLM processes the context
              ↓
Suggested answer or tool call
              ↓
Agent runtime validates and executes it

This distinction matters: the Transformer helps produce the proposal, while the surrounding agent software controls tools, permissions, retries, memory, and execution.

Direct relevance to Agentic AI: Moderate, indirectly. Multi-head attention is part of how an agent’s underlying LLM builds rich, multi-faceted contextual understanding of a conversation or retrieved document — the ability to simultaneously track multiple kinds of relationships (who’s speaking, what tool was called, what the user originally asked) benefits directly from having multiple independent attention “views” over the same context, rather than one shared, compressed view.

When this knowledge is useful

Use Multi-Head Attention when you need to explain, implement, debug, evaluate, or optimize the corresponding part of a Transformer pipeline. It is also useful when a model API behaves unexpectedly and you need to trace the behavior back to tokens, tensor shapes, attention visibility, training, or inference mechanics.

When it is not enough

Understanding this mechanism does not by itself prove that a complete model or application is accurate, safe, fast, or cost-effective. Production decisions still require representative evaluation data, latency and memory measurements, model-specific documentation, and tests of the surrounding retrieval or agent code.

14. Common Beginner Mistakes

⚠️ Mistake

Incorrect idea: assuming multi-head attention means running the FULL attention computation multiple times at full dimension.

Why it is incorrect: Each head typically works in a smaller subspace (d_head = d_model / num_heads), not the full d_model — this keeps total compute roughly comparable to a single, full-dimension attention computation, while gaining the benefit of multiple independent “views.”

⚠️ Mistake

Incorrect idea: claiming specific heads always learn specific, human-interpretable concepts.

Why it is incorrect: As Section 5 cautions, this is an empirical, model-specific observation at best — not a guaranteed architectural property.

⚠️ Mistake

Incorrect idea: forgetting the output projection step.

Why it is incorrect: Concatenating heads alone isn’t the final output — the learned output projection (Wo) is what lets the model combine the heads’ separate information meaningfully, rather than just stacking it side by side.

15. Important Distinctions

Single-Head AttentionMulti-Head Attention
One set of Wq/Wk/Wv, full d_model dimensionMultiple sets, each a smaller d_head subspace
One attention pattern per computationMultiple, independently-learned attention patterns computed in parallel
ConcatenationOutput Projection
Simply stacks all heads’ outputs side by sideA LEARNED transformation combining information across the concatenated heads

16. Production / Engineering Considerations

  • Number of heads is a real architectural hyperparameter — more heads mean more distinct relationship “views” but each in a smaller subspace; this trade-off is typically tuned empirically per model scale.
  • Grouped-query attention (GQA) and multi-query attention (MQA), covered fully in Module 17, are modern efficiency variants that share key/value projections across multiple heads specifically to reduce memory usage during inference — worth knowing this module’s standard “one Wk/Wv per head” setup is the baseline these variants optimize.

17. Interview Questions

Beginner

Q: Why does a Transformer use multiple attention heads instead of just one?

Ans: A single attention computation, using one shared set of learned projections, has to compress every kind of token relationship into one relevance scoring. Multiple heads let the model compute several independent attention patterns in parallel, each using its own learned projections — allowing it to represent different kinds of relationships simultaneously rather than being limited to one.

Intermediate

Q: What happens to the dimensionality of Q, K, and V when going from single-head to multi-head attention with the same total d_model?

Ans: Instead of one set of projections mapping the full d_model dimension, each head typically projects down to a smaller d_head = d_model / num_heads. All heads run in parallel on this smaller dimension, and their outputs are concatenated back to the original d_model before a final output projection — keeping total computation roughly comparable to a single full-dimension attention pass.

Advanced

Q: Why is the output projection step (Wo) necessary after concatenating multi-head attention’s outputs?

Ans: Simply concatenating each head’s output places their information side by side without any interaction — the model has no way to combine or weigh information across heads without an additional learned transformation.

The output projection is a learned linear layer applied to the concatenated result, allowing the model to learn how to best integrate the different heads’ independently-computed information into one unified representation, rather than leaving them as disconnected segments.

Scenario

Q: You’re told a Transformer model has “no meaningful difference” between using 1 head or 8 heads at the same total dimension, and its performance is identical either way. Would this be expected? Why or why not?

Ans: This would be unexpected and worth investigating — with properly initialized and trained heads, using multiple heads should generally let the model capture a richer set of relationships than a single head constrained to the same total dimension, since each head can specialize its own learned projections differently.

Identical performance might suggest a bug (e.g., heads accidentally sharing the same weights instead of being independently initialized) or, less likely, that this particular task’s relationships are simple enough that additional heads provide no practical benefit — but the architectural expectation is that multiple heads provide genuine representational advantage, verified in this module by the two heads producing meaningfully different attention patterns on identical input.

Architecture

Q: How does multi-head attention’s parallel structure relate to Module 1’s point about attention being parallelizable?

Ans: Just as individual attention score computations have no step-to-step dependency (Module 1), different heads’ computations are also entirely independent of each other — Head 1’s Q/K/V and attention output don’t depend on Head 2’s in any way.

This means multi-head attention, like single-head attention, can be computed efficiently in parallel on GPU hardware, preserving the core parallelization advantage that motivated moving away from recurrent architectures in the first place.

Engineering

Q: Why might a production LLM serving system care about the number of attention heads specifically, beyond just model quality?

Ans: The number of heads (and their key/value dimensions specifically) directly affects the size of the KV cache during inference (Module 16) — more heads with full, independent key/value projections mean more memory consumed per cached token.

This is precisely why efficiency variants like grouped-query and multi-query attention (Module 17), which reduce the number of independent key/value projections while keeping multiple query heads, have become important for reducing serving costs at scale.

18. What You Should Remember

  • Multi-head attention runs several independent attention computations in parallel, each on its own learned subspace (d_head = d_model / num_heads) — verified directly: two heads produced genuinely different attention weight patterns on identical input.
  • Heads are concatenated, then passed through a learned output projection to combine their information.
  • Don’t claim specific heads always learn specific human-nameable concepts — the architectural capacity for specialization is guaranteed; what’s actually learned is empirical.

19. How This Helps Me Build AI Systems

Every time you see “N attention heads” in a model’s architecture description, you now know exactly what’s being computed — N parallel, independently-learned relevance patterns over the same input, combined through concatenation and a learned output projection. This directly sets up Module 16’s KV cache discussion, where the number of heads becomes a genuine, practical memory and cost consideration.

Next: Module 8 — Positional Information — sinusoidal encoding, learned embeddings, and why modern LLMs increasingly use RoPE.

Author
TechByteByByte Editorial Team
Reviewed by
TechByteByByte Admin
Published
Last reviewed