Multi-Head Attention

Scaled Dot-Product Attention

The building block of MHA is Scaled Dot-Product Attention. Given a matrix of queries $Q \in \mathbb{R}^{T_q \times d_k}$, keys $K \in \mathbb{R}^{T_k \times d_k}$, and values $V \in \mathbb{R}^{T_k \times d_v}$:

Explanation: Each query vector dot-products with every key vector to produce a raw compatibility score, which is scaled to prevent vanishing gradients in the softmax, normalized to a probability distribution, and used to form a weighted sum of the value vectors. The $\sqrt{d_k}$ scaling factor is critical. Without it, dot products grow with dimensionality, pushing the softmax into regions of near-zero gradient:

Explanation: Assuming query and key components are i.i.d. with zero mean and unit variance, dividing by $\sqrt{d_k}$ normalizes the dot-product variance to 1 regardless of dimensionality, stabilizing training.

Projecting Into Multiple Subspaces

Multi-Head Attention applies $h$ independent learned linear projections to $(Q, K, V)$ before computing attention:

Explanation: Each head $i$ learns its own projection matrices $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$, $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ that carve out a different relational subspace of the input.

Concatenation and Output Projection

The $h$ head outputs are concatenated and linearly projected back to $d_{\text{model}}$:

Explanation: Concatenation stacks the $h$ value-weighted summaries into a vector of dimension $h \cdot d_v$, and the output matrix $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$ fuses the multi-perspective information back into the model's working dimension.

Parameter Count

In the standard setting $d_k = d_v = d_{\text{model}} / h$, so total parameters are:

Explanation: Three projection matrices $(W^Q, W^K, W^V)$ each of size $d_{\text{model}} \times d_{\text{model}}$ plus the output projection $W^O$ of the same size sum to $4 d_{\text{model}}^2$ parameters, independent of the number of heads $h$ — confirming that heads are free in terms of parameter count but not in terms of per-head capacity.

PyTorch Implementation

```python import torch import torch.nn as nn import torch.nn.functional as F import math class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1): super().__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_model = d_model self.h = num_heads self.d_k = d_model // num_heads # per-head dimension

# Fused QKV projection for efficiency self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) self.dropout = nn.Dropout(dropout) def forward( self, query: torch.Tensor, # (B, T_q, d_model) key: torch.Tensor, # (B, T_k, d_model) value: torch.Tensor, # (B, T_k, d_model) mask: torch.Tensor | None = None, # (B, 1, T_q, T_k) boolean ) -> tuple[torch.Tensor, torch.Tensor]: B, T_q, _ = query.shape T_k = key.shape[1]

# --- Project Q, K, V ---

# For self-attention query==key==value; for cross-attention they differ. Q = self.qkv_proj(query)[:, :, :self.d_model] K = self.qkv_proj(key) [:, :, self.d_model:2self.d_model] V = self.qkv_proj(value)[:, :, 2self.d_model:]

# --- Reshape to (B, h, T, d_k) --- def split_heads(x, T): return x.view(B, T, self.h, self.d_k).transpose(1, 2) Q = split_heads(Q, T_q) # (B, h, T_q, d_k) K = split_heads(K, T_k) # (B, h, T_k, d_k) V = split_heads(V, T_k) # (B, h, T_k, d_k)

# --- Scaled dot-product attention --- scale = math.sqrt(self.d_k) scores = torch.matmul(Q, K.transpose(-2, -1)) / scale # (B, h, T_q, T_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = F.softmax(scores, dim=-1) # (B, h, T_q, T_k) attn_weights = self.dropout(attn_weights)

# --- Weighted sum of values --- context = torch.matmul(attn_weights, V) # (B, h, T_q, d_k)

# --- Merge heads --- context = context.transpose(1, 2).contiguous().view(B, T_q, self.d_model) output = self.out_proj(context) # (B, T_q, d_model) return output, attn_weights # return weights for interpretability

# --- Quick sanity check --- if __name__ == "__main__": B, T, d = 2, 16, 512 mha = MultiHeadAttention(d_model=d, num_heads=8) x = torch.randn(B, T, d) out, attn = mha(x, x, x) # self-attention print(out.shape) # torch.Size([2, 16, 512]) print(attn.shape) # torch.Size([2, 8, 16, 16]) ``` > Note on `torch.nn.MultiheadAttention`: PyTorch ships a fused, optimized implementation. The custom version above matches its semantics exactly and is useful for learning or when you need custom masking logic. For production use `F.scaled_dot_product_attention` (PyTorch ≥ 2.0), which dispatches to FlashAttention kernels automatically.