Phase 1.1 · OptionalOptionalModule 9 of 21

Multi Head Attention

A single attention head can only capture one type of relationship at a time. Multiple heads let the model simultaneously attend to different representation subspaces — one head for subject-verb agreement, another for coreference, another for local context.

A newsroom with multiple reporters covering the same event from different angles — politics, economics, human interest. The editor (concatenation + projection) combines all perspectives into one comprehensive story.

Visual Workflows

Start here — study each diagram, then use + / − to zoom if needed.

100%

Loading diagram...

Scroll inside the frame · use + / − to zoom

Key Takeaways

1.Multi-head attention runs several attention operations in parallel — each 'head' learns to focus on different types of relationships (syntax, semantics, position) and their outputs are concatenated.
2.Instead of one attention with d_model dimensions, use h heads each with d_k = d_model/h dimensions.
3.Each head: Q_i = XW_Q_i, K_i = XW_K_i, V_i = XW_V_i → head_i = Attention(Q_i, K_i, V_i).
4.Concatenate all heads: [head_1; head_2; ...; head_h] → multiply by W_O (output projection) to get d_model output.
5.Typical: 8-16 heads for base models, 96+ for large LLMs.

Real Example

Scenario

In 'The cat sat on the mat because it was comfortable', Head 3 might attend 'it'→'mat' (location), Head 7 might attend 'sat'→'cat' (subject-verb), Head 12 might capture local bigram patterns. Combined, the model understands the full sentence.

What you would do

In Transformer & ML Foundations, apply Multi Head Attention to this scenario: In 'The cat sat on the mat because it was comfortable', Head 3 might attend 'it'→'mat' (location), Head 7 might attend 'sat'→'cat' (subject-verb), Head 12 might capture local bigram patterns. Identify the inputs, run the technique, validate the output, and note one thing you would monitor in production.

Practice Task

Open the Code Walkthrough below and run it locally. Change one parameter related to Multi Head Attention (e.g. model, temperature, top_k, or tool name), observe the difference in output, and write 2–3 sentences explaining what changed.

Code Walkthrough

Highlighted lines show where Multi Head Attention happens in the code.

Multi Head Attention

1import torch  # import dependencies2import torch.nn as nn  # import dependencies3
4class MultiHeadAttention(nn.Module):  # define a data structure or component5    def __init__(self, d_model: int, num_heads: int):  # define a reusable function6        super().__init__()7        assert d_model % num_heads == 0  # key line for Multi Head Attention8        self.d_k = d_model // num_heads  # key line for Multi Head Attention9        self.num_heads = num_heads  # key line for Multi Head Attention10        self.W_q = nn.Linear(d_model, d_model)11        self.W_k = nn.Linear(d_model, d_model)12        self.W_v = nn.Linear(d_model, d_model)13        self.W_o = nn.Linear(d_model, d_model)14
15    def forward(self, x):  # define a reusable function16        B, S, D = x.shape17        Q = self.W_q(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)  # key line for Multi Head Attention18        K = self.W_k(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)  # key line for Multi Head Attention19        V = self.W_v(x).view(B, S, self.num_heads, self.d_k).transpose(1, 2)  # key line for Multi Head Attention20        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)21        weights = torch.softmax(scores, dim=-1)22        out = torch.matmul(weights, V)23        out = out.transpose(1, 2).contiguous().view(B, S, D)24        return self.W_o(out)  # return the result

Cheat Sheet

Quick recap

quick ref

•d_k = d_model / h
•head_i = Attn(Q_i, K_i, V_i)
•out = Concat(heads) @ W_O
•Typical: 8-16 heads (base models)
•view + transpose for head splitting
•W_O mixes cross-head information

Common Mistakes

✕d_model not divisible by num_heads — reshape fails
✕Forgetting output projection W_O after concatenation
✕Assuming each head has a fixed linguistic role — they're learned
✕Confusing num_heads with num_layers