§3.3Multi-Headed Attention

At the end of the last section we hit a wall: with query = key = value, a token attends mostly to itself. The escape is to give the layer room to specialize — and then to run several such specialists in parallel. That is multi-headed attention multi-head attention Running several attention heads in parallel, each with its own Q/K/V projections, then concatenating and linearly projecting their outputs. defined in ch. 3 — open in glossary .

One head, three projections

A single attention head attention head One independent set of query/key/value projections; BERT uses 12 heads of dimension 768/12 = 64, each free to focus on a different relation. defined in ch. 3 — open in glossary applies three independent learned linear projections to each embedding, producing distinct query query The learned projection of a token embedding that does the 'asking' — matched against every key to decide how much to attend. defined in ch. 3 — open in glossary , key key The learned projection of a token embedding that is 'matched against' by queries via a similarity (dot product). defined in ch. 3 — open in glossary , and value value The learned projection of a token embedding that carries the content mixed into the output according to the attention weights. defined in ch. 3 — open in glossary vectors instead of reusing the embedding three times:

class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.q = nn.Linear(embed_dim, head_dim)
        self.k = nn.Linear(embed_dim, head_dim)
        self.v = nn.Linear(embed_dim, head_dim)

    def forward(self, hidden_state):
        attn_outputs = scaled_dot_product_attention(
            self.q(hidden_state), self.k(hidden_state), self.v(hidden_state))
        return attn_outputs

Each nn.Linear maps the 768-dim embedding down to head_dim dimensions. Now the query and key of a token can differ, so a token can look for complementary information rather than a copy of itself.

Why more than one head?

The softmax of a single head tends to focus on one aspect of similarity — often the current token itself. But language has many kinds of relationships at once, so we run several heads in parallel and let each learn its own:

Try it. The lab below has two hand-built heads for “time flies like an arrow”. The semantic head routes “flies” to “arrow”; the verb-centric head makes every token attend to the verb “flies”. Same sentence, completely different attention — that is the value of multiple heads.

Concatenate the heads

The full multi-head layer multi-head attention Running several attention heads in parallel, each with its own Q/K/V projections, then concatenating and linearly projecting their outputs. defined in ch. 3 — open in glossary concatenates every head’s output and passes the result through one more linear layer:

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        embed_dim = config.hidden_size          # 768
        num_heads = config.num_attention_heads  # 12
        head_dim = embed_dim // num_heads        # 64
        self.heads = nn.ModuleList(
            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]
        )
        self.output_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, hidden_state):
        x = torch.cat([h(hidden_state) for h in self.heads], dim=-1)
        x = self.output_linear(x)
        return x

In matrix notation, with $h$ heads:

• $W_i^Q, W_i^K, W_i^V$ — the $i$-th head’s projection matrices (its own learnable parameters).
• $\text{head}_i$ — the §3.2 scaled dot-product attention run inside head $i$.
• $W^O$ — the final output projection (output_linear) mixing the heads.

Follow the shapes through the whole layer — note how the hidden dimension splits into heads and then re-merges:

Passing BERT’s config through this layer confirms the shape round-trips:

multihead_attn = MultiHeadAttention(config)
attn_output = multihead_attn(inputs_embeds)
attn_output.size()          # torch.Size([1, 5, 768])

Reading real heads

With a pretrained checkpoint you can visualize genuine learned attention (the book uses BertViz’s head_view). Feeding it two sentences — “time flies like an arrow” and “fruit flies like a banana” — and inspecting head 8 shows two things:

• Attention is strongest between words in the same sentence: the model has learned that the [SEP]-separated sentences are distinct.
• For the ambiguous word “flies”, the head attends to “arrow” in the first sentence and to “fruit” and “banana” in the second — exactly the context needed to tell the verb from the noun.

That single head has learned a slice of syntax and word-sense disambiguation, purely from data. Multiply by twelve, stack across layers, and you get the rich contextual representations that make transformers work.

With attention fully assembled, the encoder layer still needs its second sublayer and the glue that holds a deep stack together: the feed-forward network, layer normalization, and positional embeddings.