Inside Transformers, Part 4: Multi-head and masked attention

Part 3 derived scaled dot-product attention for one head. A Transformer uses several heads in parallel, then uses a mask and a second attention block on the decoder side.

Multi-head attention

Each head has its own query, key, and value projections. If the model width is d_model and there are h heads, a common configuration gives each head a smaller width d_k = d_v = d_model / h. The base Transformer used d_model = 512 and h = 8, so each head used width 64. Eight heads are a detail of that original base configuration, not a rule for every Transformer.

For head r, the attention output is

head_r = Attention(Q_r, K_r, V_r)

The model concatenates the head outputs along their feature dimension and applies another learned matrix W_O:

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_O

With eight width-64 heads, concatenation returns to width 512. The output projection mixes information from the heads before the residual connection.

Cross-attention in the decoder

The encoder and decoder use self-attention, but the decoder also has a cross-attention block. The decoder’s current hidden states provide the queries. The encoder’s output sequence is projected separately to keys and values. This lets each target position select information from any permitted source position.

For an English-to-French translation, the encoder reads the English sentence. A decoder position producing a French token uses its query to score the encoder keys, then mixes the encoder values. The query does not come from the encoder, and the keys and values do not come from the decoder’s target prefix.

Causal masking

The decoder receives a shifted target sequence during training. For example, it might receive <BOS> Je suis and learn to predict Je suis ... one position to the right. At decoder position i, masked self-attention may read positions j <= i in that shifted input. It must not read a target token that belongs to a later position.

The mask is applied to attention scores before softmax. A boolean upper-triangular mask expresses the rule directly. If a score tensor is used, a dtype-compatible fill value can represent excluded entries:

import torch

future = torch.triu(
    torch.ones(seq_len, seq_len, dtype=torch.bool, device=scores.device),
    diagonal=1,
)
scores = scores.masked_fill(future, torch.finfo(scores.dtype).min)
weights = torch.softmax(scores, dim=-1)

The future matrix marks keys with j > i. Those entries receive a very negative value before softmax, so the causal self-attention weights for them become zero. Using torch.finfo(scores.dtype).min keeps the fill value compatible with the score tensor’s floating-point type. Frameworks that expose a boolean mask or an is_causal option can express the same rule without manually choosing a constant.

What attention patterns can tell us

Different heads can produce different score patterns, and some analyses find heads that correlate with local syntax or longer-range relations. That observation is useful for investigating a trained model, but it does not prove that a head represents one human-defined concept. Attention weights alone are not definitive explanations of a model’s decision.

Part 5 covers the residual additions and layer normalization that wrap these attention blocks in the original post-normalization Transformer.