Multi-Head Attention

The self-attention mechanism we have looked at so far is complete and LLM-ready, but is a special case called single-head attention.

Modern LLMs instead use multi-head self-attention. This sounds fancy, but all it really means is that we take the exact same mechanism we have just learned about, and run multiple instances of it in parallel. Conceptually this looks like the following:

Throughout this page we will show multi-head attention with just two heads, but it works exactly the same for any number of heads.

The reason we do this is because it allows the LLM to learn multiple different types of attention at the same time. For example, in our previous toy example we showed self-attention learning how to associate pronouns with their referents (she was referring to dog). But there could be many other types of relationships that self-attention can discover — for example tracking which verb belongs to which subject, linking an adjective to the noun it describes, or simply focusing on the words physically nearest to each token. (The interactive demo on the overview page shows three such heads: pronoun resolution, syntax, and local attention.)

Simple implementation

Given our previous implementation of single-head attention, we can write a multi-head attention mechanism as follows: keep a list of independent attention heads, run each one on the input, and concatenate their results side by side.

import torch
import torch.nn as nn


class MultiHeadAttentionWrapper(nn.Module):

    def __init__(self, input_dim, head_dim, num_heads):
        super().__init__()
        self.heads = nn.ModuleList(
            [SelfAttention(input_dim, head_dim) for _ in range(num_heads)]
        )

    def forward(self, x):
        # run every head, then join their results side by side
        return torch.cat([head(x) for head in self.heads], dim=-1)

We also add another linear layer after joining the two result vectors (self.out_proj = nn.Linear(num_heads * head_dim, num_heads * head_dim)).

The reason for this last layer is that the heads work completely independently, so just concatenating two sets of enriched input vectors together, by itself, is not very meaningful to an LLM. Adding a final linear layer allows the LLM to take the combined learnings from its multiple attention heads, and form a representation that encapsulates the combined learnings in a meaningful way.

Our full multi-head attention class (extended to handle any number of heads) now looks as follows:

class MultiHeadAttentionWrapper(nn.Module):

    def __init__(self, input_dim, head_dim, num_heads):
        super().__init__()
        self.heads = nn.ModuleList(
            [SelfAttention(input_dim, head_dim) for _ in range(num_heads)]
        )
        self.out_proj = nn.Linear(num_heads * head_dim, num_heads * head_dim, bias=False)

    def forward(self, x):
        combined = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.out_proj(combined)

Full implementation

The multi-head attention implementation above is fully functional and would work fine, but running all of our heads in series and concatenating the results is not very efficient for running at scale. We'll now show a method which produces the same results, but uses PyTorch's functionality much more efficiently.

We'll show the calculation in depth, showing the full results for a worked example at each step. The purpose of this is so that you can see the shape of the various matrices at each step, and see that we really are achieving the exact same results as in the simple version in our Simple implementation above.

For simplicity, we'll use two heads, each with a head dimension of 3. In a real transformer, a multi-head attention block keeps its output the same width as its input — so the blocks can be stacked, each one adding its result back onto what came before. That means the embedding dimension is exactly num_heads × head_dim: the heads partition the embedding rather than widening it. With two heads of dimension 3, that gives an embedding dimension of 6. (The single-head example on the previous page used an embedding dimension of 2 just to keep the vectors small enough to read; here we use the realistic setup.)

num_tokens = 5
num_heads  = 2
head_dim   = 3
embed_dim  = num_heads * head_dim   # = 6

We'll reuse the start of our sentence from the previous page, "the dog fetched the ball", giving us five input tokens, each embedded into 6 dimensions:

Step 1 — project to queries, keys and values

Exactly as before, we project the input to get queries, keys and values — but now each projection outputs num_heads × head_dim = 6 numbers per token:

queries = self.W_query(x)   # each is shape (5, 6)
keys    = self.W_key(x)
values  = self.W_value(x)

The trick is that those 6 numbers are really two heads stitched together: the first three columns (blue) belong to head 1, the last three (teal) to head 2.

Step 2 — split into heads

We make that grouping explicit with a reshape, then move the head axis to the front so the two heads can be processed as an independent batch:

queries = queries.view(num_tokens, num_heads, head_dim).transpose(0, 1)
keys    = keys.view(num_tokens, num_heads, head_dim).transpose(0, 1)
values  = values.view(num_tokens, num_heads, head_dim).transpose(0, 1)
# each is now shape (num_heads, num_tokens, head_dim) = (2, 5, 3)

No numbers change here — .view(...) just regroups the 6 columns above into (2 heads × 3 dims), and .transpose(...) reorders the axes so that head 1's slice (the blue columns) and head 2's slice (the teal columns) each peel apart into their own (5, 3) block:

From here, every head is just the single-head attention we already know, run independently on its own block.

Step 3 — attention within each head

Each head now runs the familiar scaled dot-product attention — scores, causal mask, softmax, then a weighted sum of its values — on its own (5, 3) query, key and value blocks. Because the two heads hold different numbers, they attend differently:

attention_scores = queries @ keys.transpose(-2, -1)   # shape (2, 5, 5)

mask = torch.triu(torch.ones(num_tokens, num_tokens), diagonal=1).bool()
attention_scores = attention_scores.masked_fill(mask, float('-inf'))

attention_weights = torch.softmax(attention_scores / head_dim ** 0.5, dim=-1)

context = attention_weights @ values   # shape (2, 5, 3)

Running that on our two heads gives:

Step 4 — recombine the heads

Finally we stitch the heads back together — transpose the head axis back and flatten to one vector per token — and pass the result through the output projection:

context = context.transpose(0, 1).reshape(num_tokens, num_heads * head_dim)
output  = self.out_proj(context)   # shape (5, 6)

This leaves us with our final self-attention class:

import torch
import torch.nn as nn


class MultiHeadAttention(nn.Module):

    def __init__(self, input_dim, head_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        out_dim = num_heads * head_dim

        self.W_query = nn.Linear(input_dim, out_dim, bias=False)
        self.W_key   = nn.Linear(input_dim, out_dim, bias=False)
        self.W_value = nn.Linear(input_dim, out_dim, bias=False)
        self.out_proj = nn.Linear(out_dim, out_dim, bias=False)

    def forward(self, x):
        num_tokens = x.shape[0]

        queries = self.W_query(x)
        keys    = self.W_key(x)
        values  = self.W_value(x)

        # split each token's vector into (num_heads, head_dim), then
        # move the head axis to the front
        queries = queries.view(num_tokens, self.num_heads, self.head_dim).transpose(0, 1)
        keys    = keys.view(num_tokens, self.num_heads, self.head_dim).transpose(0, 1)
        values  = values.view(num_tokens, self.num_heads, self.head_dim).transpose(0, 1)

        attention_scores = queries @ keys.transpose(-2, -1)

        mask = torch.triu(torch.ones(num_tokens, num_tokens), diagonal=1).bool()
        attention_scores = attention_scores.masked_fill(mask, float('-inf'))

        attention_weights = torch.softmax(
            attention_scores / self.head_dim ** 0.5, dim=-1
        )

        context = attention_weights @ values
        context = context.transpose(0, 1).reshape(num_tokens, self.num_heads * self.head_dim)

        return self.out_proj(context)

A note on batches

One final note — most LLM implementations allow you to input batches of inputs at the same time, instead of a single input. The nice thing is that the magic of PyTorch means the code works out almost exactly the same: the input gains a leading batch dimension, the reshape becomes .view(batch, num_tokens, num_heads, head_dim) and the head axis is moved with .transpose(1, 2), but every other line stays the same. A full batched implementation will follow.

What's next

Self-attention has now shared information between the tokens. The next stage in a transformer block processes each enriched token independently:

• Feed-forward layer