Multi-Head Attention

Run $h$ attention operations in parallel on $d/h$-dimensional slices of the residual stream so different heads can attend to different relations, then concatenate and project.

Equation

$\mathrm{MHA}(x)=\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h)W^O,{\mathrm{head}}_i=\mathrm{Attn}(x{W}_i^Q,x{W}_i^K,x{W}_i^V)$

Head dim $d_h=d/h$, so total FLOPs match a single-head attention of width $d$.

From scratch

import torch, torch.nn as nn
from einops import rearrange
 
class MHA(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.h, self.dh = n_heads, d // n_heads
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d, bias=False)
 
    def forward(self, x):                       # x: (b, n, d)
        qkv = self.qkv(x)                       # (b, n, 3d)
        q, k, v = rearrange(qkv, 'b n (three h dh) -> three b h n dh',
                            three=3, h=self.h)  # each (b, h, n, dh)
        scores = torch.einsum('bhid,bhjd->bhij', q, k) / self.dh**0.5
        n = x.shape[1]
        mask = torch.ones(n, n, dtype=torch.bool, device=x.device).triu(1)
        scores = scores.masked_fill(mask, float('-inf'))
        out = torch.einsum('bhij,bhjd->bhid', scores.softmax(-1), v)
        return self.proj(rearrange(out, 'b h n dh -> b n (h dh)'))

Notes

• The fused qkv projection is one GEMM instead of three, which is why real implementations do it that way.
• The three h dh split order must match how weights were saved; (h dh) vs (dh h) is a classic checkpoint-conversion bug.
• Params: $4d^2$ (Q, K, V, O) ignoring biases.
• Grouped-Query Attention keeps $h$ query heads but fewer K/V heads to shrink the KV Cache.

Related: Scaled Dot-Product Attention, Self-Attention, Transformer Block