Multi-Head Attention

Imagine a team reading a contract together. One person looks for dates. One looks for names. One looks for risks. Then they combine their notes. That's much better than one person trying to do it all. Multi-head attention works the same way.

• A team reads a contract
• Each one has a different focus
• Then they combine notes

Different heads, different focus

In the same sentence, one attention head might connect it back to cat, working out who's being talked about. Another might link tired to sat, connecting cause and effect. A third might track grammar. We split each word's vector into pieces, give one piece to each head, run attention separately, and then join everything together. Nobody tells the heads what to focus on. They work it out during training.

Multi-head from scratch

Here's the full thing. One linear layer makes queries, keys, and values in a single step. Split reshapes them, so that each head gets its own slice. In this case, sixty four numbers become eight heads of eight each. The attention line is the same formula as before, applied to all heads at once. Then we glue the heads back together, and mix them with one last linear layer. Run it. Same shape in, same shape out. This module is used, almost exactly like this, inside the biggest language models on Earth.

import torch
from torch import nn

class MultiHead(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads, self.size = heads, dim // heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim)

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        split = lambda t: t.view(B, T, self.heads, self.size).transpose(1, 2)
        q, k, v = split(q), split(k), split(v)
        w = (q @ k.transpose(-2, -1) / self.size ** 0.5).softmax(dim=-1)
        y = (w @ v).transpose(1, 2).reshape(B, T, D)
        return self.out(y)

layer = MultiHead(dim=64, heads=8)
print("In: ", [2, 10, 64])
print("Out:", list(layer(torch.rand(2, 10, 64)).shape))
print("Each head looks at", 64 // 8, "numbers")

In:  [2, 10, 64]
Out: [2, 10, 64]
Each head looks at 8 numbers

Summary