Multi-head
and the causal mask

1 The problem with a single head

The attention from chapter 03 works, but it has a serious limitation: each row of the attention map sums to 1. That means attention is a fixed budget. If a token spends 70% looking at its syntactic antecedent, only 30% is left for everything else.

But a word takes part in several relations at once. Take "mordeu" (bit) in "o cachorro grande mordeu o carteiro ontem" (the big dog bit the mailman yesterday):

• It needs to look at "cachorro" (dog) — who is the subject?
• It needs to look at "carteiro" (mailman) — who is the object?
• It needs to look at "ontem" (yesterday) — when did this happen?
• It needs to look at the preceding "o" (the) — agreement, local structure.

With a single head, these four needs compete for the same 100% budget. The model is forced to average across relations that are qualitatively different — and an average of different things is usually a bad thing.

2 How it works: divide, don't multiply

The detail that makes the idea elegant: multi-head does not cost more. Instead of twelve heads of 768 dimensions each, the model splits the existing 768 dimensions into twelve slices of 64.

The full procedure:

The matrix W^O is essential and frequently overlooked. Without it, the twelve heads would sit in sealed compartments, each writing into its own 64 dimensions and never talking to each other. It is W^O (called c_proj in GPT-2) that lets information from head 3 influence the dimensions that head 7 uses.

The dimensions, in practice

GPT-2 small, 5-token sentence, 12 heads:

h                    (5, 768)
  |
c_attn (768 -> 2304) (5, 2304)          a single matmul
  split into 3       Q,K,V: (5, 768) each
  reshape into heads Q,K,V: (12, 5, 64) <- 12 heads in parallel
  |
attention per head   (12, 5, 64)         12 independent 5x5 maps
  |
concat (join)        (5, 768)            regroups the slices
c_proj (768 -> 768)  (5, 768)            mixes the heads
  |
output               (5, 768)            same shape as the input

What the heads actually learn

Nobody assigns roles. The specializations emerge from training. Interpretability studies on GPT-2 and BERT have consistently found heads such as:

3 The causal mask — the line that defines everything

We now reach the most important point of the entire course. Pay close attention to this section, because everything in Part III depends on it.

The problem

GPT-2 is trained to predict the next word. We give it "O Brasil é" and it must produce "gigante". But, to be efficient, training does not run one word at a time — it processes the whole sentence at once and demands all the predictions simultaneously:

Input:    O      Brasil  é       gigante  pela
Target:   Brasil é       gigante pela     própria
          ^      ^       ^       ^        ^
          5 predictions demanded at the same time, in a single forward pass.

And here is the catastrophe: the attention from chapter 03 lets each token see all the others — including the ones that come after it. While trying to predict that "gigante" comes after "é", the token "é" is looking directly at "gigante".

The model would learn the trivial task of copying the answer. When it came time to generate real text, where the future does not exist, it would have no idea what to do. This is label leakage — the cardinal sin of supervised learning.

The solution: forbid it explicitly

Before the softmax, add a mask matrix that puts −∞ in every future position:

Why −∞ and not zero?

That is the right question, and the answer is beautiful. Zeroing the score does not help — a score of 0 still becomes e⁰ = 1 in the softmax, a perfectly respectable weight. And zeroing the weight after the softmax would break the sum-to-1.

The −∞ works because:

In implementations a very negative number is used (-1e9, or torch.finfo(dtype).min) to avoid NaN in floating-point arithmetic.

The triangular shape

Attention map AFTER the mask (5 tokens). "." = forbidden.

           O   Brasil  e   gigante  pela
   O     [0.99   .      .     .       .   ]  <- sees only itself
   Brasil[0.41  0.59    .     .       .   ]
   e     [0.12  0.55   0.33   .       .   ]
   gigante[0.08 0.31   0.19  0.42     .   ]
   pela  [0.05  0.22   0.11  0.38    0.24 ]  <- sees everything up to here

Lower triangular. Each row sums to 1. The future does not exist.

The practical consequence: the KV cache

The causal mask brings a free gift. Since token 5 never influences token 3, the K and V vectors of earlier positions never change when a new token is appended.

So, when generating text, there is no need to recompute the whole sentence for every word. You store K and V for everything that has already passed (the KV cache) and compute only the new token. Generation drops from O(n²) per token to O(n).

This is what makes autoregressive generation viable in practice — and RoBERTa, being bidirectional, cannot do this, because in it appending a token changes the representation of all the others.

4 The complete code

import torch, math
import torch.nn as nn
import torch.nn.functional as F

class AtencaoMultiCabeca(nn.Module):
    def __init__(self, d_model=768, n_cabecas=12, causal=True):
        super().__init__()
        self.h = n_cabecas
        self.d_k = d_model // n_cabecas      # 64
        self.causal = causal                  # <<< GPT=True, RoBERTa=False
        self.c_attn = nn.Linear(d_model, 3 * d_model)  # Q,K,V fused
        self.c_proj = nn.Linear(d_model, d_model)      # W^O

    def forward(self, x):
        B, n, d = x.shape

        # 1) a single matmul produces Q, K and V at once
        q, k, v = self.c_attn(x).split(d, dim=2)

        # 2) (B, n, 768) -> (B, 12, n, 64): separates the heads
        q = q.view(B, n, self.h, self.d_k).transpose(1, 2)
        k = k.view(B, n, self.h, self.d_k).transpose(1, 2)
        v = v.view(B, n, self.h, self.d_k).transpose(1, 2)

        # 3) scaled scores -> (B, 12, n, n)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 4) ===== THE LINE THAT DEFINES THE ARCHITECTURE =====
        if self.causal:
            futuro = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
            att = att.masked_fill(futuro, float("-inf"))
        # ==================================================

        att = F.softmax(att, dim=-1)
        y = att @ v                                    # (B, 12, n, 64)

        # 5) regroups the heads and mixes
        y = y.transpose(1, 2).contiguous().view(B, n, d)
        return self.c_proj(y)

Read step 4 again. causal=True → GPT-2. causal=False → RoBERTa. Everything else is identical.

5 In the lab

.venv\Scripts\python.exe lab\04_multihead_mask.py

1. Prints the triangular mask in ASCII, so you can see its shape.
2. Runs the same input with causal=True and False, and shows the two attention matrices side by side. The visual difference is immediate.
3. Proves the leakage: changes a token at the end of the sentence and measures how much the representation at the start changed. In causal mode: zero. In bidirectional mode: quite a lot. This is the test that demonstrates causality instead of asserting it.
4. Portrait of the 12 heads of layer 0 in the real GPT-2: for each head, computes the average distance it looks at and the entropy of its distribution. You will identify the positional heads (distance ≈ 1) and the "sink" heads (extremely low entropy, everything on token 0).
5. Measures the gain from the KV cache by generating 50 tokens with and without it.

Attention moves information between positions. But it is essentially a weighted average — a linear operation.
Chapter 05 shows the missing pieces: where non-linear thinking comes in, and how 48 layers get stacked without training blowing up.