Multi-head attention: many ways of looking at one sentence

Transformer on-ramp · Part 6 of 8

If one attention head gives a sentence one point of view, why not give it several? That is the whole idea of multi-head attention, and it turns out to be almost free.

Several ways of looking at the same sentence

One attention operation gives each token one blended view of the sentence. But a word relates to others in several ways at once: grammar, who “it” refers to, which words are nearby. Multi-head attention runs several attentions in parallel, each on its own smaller projection of the embedding, and lets each one specialise. Then it glues the results together and mixes them with one more learned projection.

Is it more expensive?

Barely. If the model width is d_model and there are h heads, each head works in a space of width d_k = d_model / h. Narrower heads, more of them, and the total number of numbers is the same as for one wide head. In the companion notebook, one wide head and four heads both use 4,096 weights (for a model width of 32). You get several viewpoints for about the same cost.

The practical trick

In real code the per-head projections are not h separate matrices. One big matrix does all the projecting, and its output is simply reshaped into h heads. Each head’s query, key and value are therefore still computed from the whole embedding. (Splitting the embedding first and giving each head a slice would blind it to most of the input.)

import torch

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, h):
        super().__init__()
        self.h, self.dk = h, d_model // h
        self.qkv = torch.nn.Linear(d_model, 3 * d_model)   # one matrix makes Q, K and V
        self.out = torch.nn.Linear(d_model, d_model)       # W_O

    def forward(self, x, causal=True):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        split = lambda t: t.view(B, T, self.h, self.dk).transpose(1, 2)   # -> (B, h, T, dk)
        q, k, v = split(q), split(k), split(v)
        mask = torch.tril(torch.ones(T, T)) if causal else None
        o, w = scaled_dot_product_attention(q, k, v, mask)   # from the previous post
        o = o.transpose(1, 2).reshape(B, T, D)               # concatenate the heads
        return self.out(o), w

x = torch.randn(2, 6, 32)
y, w = MultiHeadAttention(d_model=32, h=4)(x)
print(y.shape, w.shape)    # torch.Size([2, 6, 32]) torch.Size([2, 4, 6, 6])

The output has the same shape as the input, so layers can be stacked. The weights tensor has one 6×6 table per head: four separate opinions about the same six tokens.

Open in Colab Shape walk-through, the parameter count, and a check against PyTorch. Free in your browser.

Takeaways

• Several smaller attention heads run in parallel, each free to specialise.
• d_k = d_model / h, so total cost is about the same as one big head.
• In code: one big projection, reshaped into heads; heads are concatenated and mixed by W_O.
• Output has the same shape as the input, so the layer can be stacked.

Next in the series (Part 7): Positional encoding: teaching attention about word order