Multi-head attention runs several independently learned attention mechanisms in parallel, letting a model examine different relationships and representation subspaces within the same sequence. section 3.4.2 section 3.6
Each head applies causal self-attention with its own query, key, and value projections. The resulting context vectors are concatenated along the embedding dimension. section 3.6.1 section 3.7
A practical implementation follows this data flow:
- Create one CausalAttention module per head, keeping each module’s learned projections separate. section 3.6.1
- Run every head on the same input sequence and concatenate the outputs with torch.cat(..., dim=-1). section 3.6.1
- For two heads producing two features per token, obtain a four-dimensional token representation; generally, the concatenated width is num_heads × d_out. section 3.6.1
- Use batched matrix multiplications for a more compact and computationally efficient implementation than separately stacking full modules. section 3.7
Two heads over a batch
With d_in=3, d_out=2, two heads, and six tokens per input, MultiHeadAttentionWrapper returns context_vecs with shape [2, 6, 4]: two texts, six tokens, and four concatenated features per token. section 3.6.1
Listing 3.4 A wrapper class to implement multi-head attention
class MultiHeadAttentionWrapper(nn.Module):
def __init__(self, d_in, d_out, context_length,
dropout, num_heads, qkv_bias=False):
super().__init__()
self.heads = nn.ModuleList(
[CausalAttention(
d_in, d_out, context_length, dropout, qkv_bias
)
for _ in range(num_heads)]
)
def forward(self, x):
return torch.cat([head(x) for head in self.heads], dim=-1)
The wrapper makes the architecture explicit: ModuleList constructs independent causal heads, and the forward pass concatenates their context vectors. section 3.6.1
Inside a GPT-style transformer block, masked multi-head attention analyzes relationships between sequence positions, while a feed-forward network transforms each position independently. Layer normalization, dropout, and residual connections surround these operations. section 4.5
