Multi Head Self Attention — Module 10.2

How This Lesson Fits the Module & Volume

Module 10.1 introduced multi-head attention and self-attention separately. Here we fuse them as used inside encoder and decoder blocks: multi-head self-attention (MHSA)—Q, K, V from the same sequence, split across heads.

This is the workhorse mixer of the Transformer stack and the substrate for masked variants and ViT patch mixing.

Learning Objectives

By the end of this lesson, students should be able to:

Define MHSA as multi-head attention where Q, K, V come from the same input.
Relate head dimension d_k = d_model / h to the projection shapes.
Implement MHSA from linear projections + scaled dot-product attention in PyTorch.
Explain why multiple heads capture different relation types in parallel.
Connect MHSA to Module 10.1 QKV and attention-score lectures.
State how MHSA plugs into residual + LayerNorm + FFN blocks.

Definition

Multi-head self-attention (MHSA) projects an input sequence X into queries, keys, and values, splits them into h heads, runs scaled dot-product attention per head, concatenates the heads, and applies an output projection—letting each position mix information from all (allowed) positions in multiple representation subspaces.

From Single-Head Self-Attn to Multi-Head

Project

X → Q, K, V.

Split heads

Reshape to h parallel views.

Attend

Softmax(QK^T/√d_k) V.

Merge

Concat + W^O.

Symbol	Meaning
X ∈ R^L×d_model	Input token representations
h	Number of heads
d_k = d_model / h	Per-head key/query dimension
W^Q, W^K, W^V, W^O	Learned projection matrices

MHSA vs MHA vs Cross-Attention

Self-attention

Q, K, V from same X.
Builds context within a sequence.
Encoder & decoder (masked).

Cross-attention

Q from target; K/V from source.
Links decoder to encoder.
See Module 10.1 cross-attn.

Multi-head

Orthogonal to self vs cross.
Parallel subspaces.
Used in both settings.

Code: MHSA from Scratch (Educational)

import math
import torch
from torch import nn
import torch.nn.functional as F

class MHSA(nn.Module):
    def __init__(self, d_model=64, nhead=4, dropout=0.0):
        super().__init__()
        assert d_model % nhead == 0
        self.h = nhead
        self.d_k = d_model // nhead
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        # x: (B, L, d_model)
        B, L, _ = x.shape
        qkv = self.qkv(x).reshape(B, L, 3, self.h, self.d_k).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # (B, h, L, d_k)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if attn_mask is not None:
            scores = scores + attn_mask  # broadcastable (L, L) or (B, 1, L, L)
        weights = self.drop(scores.softmax(dim=-1))
        ctx = (weights @ v).transpose(1, 2).reshape(B, L, -1)
        return self.out(ctx)

mhsa = MHSA()
x = torch.randn(2, 8, 64)
print(mhsa(x).shape)  # torch.Size([2, 8, 64])
# Production alternative: nn.MultiheadAttention(..., batch_first=True)

Strengths and Tradeoffs

Strengths

Parallel relation detectors (syntax, coreference, etc.).
Same API for text tokens and ViT patches.
Composes cleanly with residuals and FFN.

Tradeoffs

O(L² · d) cost per layer.
Head count must divide d_model cleanly.
Without positions, MHSA is order-agnostic.

Common Misconception

“More heads always mean better models.” Past a point, tiny d_k per head hurts; papers often find a sweet spot. Capacity also comes from depth and d_model—not heads alone.

Knowledge Check

Short Answer: What makes attention “self” attention? Answer: Q, K, and V are derived from the same sequence.
True/False: Multi-head means running several attentions in different subspaces then concatenating. Answer: True.
Multiple Choice: If d_model=64 and h=8, then d_k is: (a) 64, (b) 8, (c) 512. Answer: (b).
Short Answer: Write the scaled dot-product formula. Answer: softmax(QK^T/√d_k) V.
True/False: Cross-attention is a type of self-attention. Answer: False—K/V come from another sequence.
Multiple Choice: W^O is applied: (a) before splitting heads, (b) after concatenating heads, (c) only on values. Answer: (b).
Short Answer: Which Module 10.1 lectures define Q, K, V? Answer: Query, Key, and Value.
Short Answer: Why divide by √d_k? Answer: To keep dot-product magnitudes stable before softmax.
Multiple Choice: MHSA alone encodes absolute order: (a) always, (b) never without positions/masks structure, (c) only with dropout. Answer: (b).
True/False: nn.MultiheadAttention can implement MHSA when query=key=value. Answer: True.

Key Takeaways

MHSA = multi-head attention with QKV from the same input.
Heads provide parallel subspaces; merge via concat + W^O.
Built from Module 10.1 scaled dot-product attention.
Used in every encoder/decoder block (optionally masked).
Next: Masked Attention.

Trainer’s Guide

Hands-on idea: Compare your MHSA output to nn.MultiheadAttention on the same random init (or close shapes) and discuss numerical differences.

Discussion prompt: If one head learns positional patterns and another learns lexical links, how does that motivate multi-head design?

Recap: MHSA is the Transformer’s multi-subspace self-mixer. Continue with Masked Attention.