Multi-Head Attention › Practical Considerations
Language AI HandbookAll chapters Multi-Head Attention › Practical ConsiderationsOne section · Sections in this chapter10 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
When implementing multi-head attention in production systems, several practical considerations come into play.
Choosing the Number of Heads
The original transformer paper used 8 heads with a 512-dimensional model, giving each head 64 dimensions. Modern large models often use more heads, but the per-head dimension typically stays between 64 and 128.
# Real model configurations
model_configs = {
"Transformer (orig.)": {"d_model": 512, "num_heads": 8},
"BERT-base": {"d_model": 768, "num_heads": 12},
"GPT-2 Small": {"d_model": 768, "num_heads": 12},
"BERT-large": {"d_model": 1024, "num_heads": 16},
"GPT-2 Medium": {"d_model": 1024, "num_heads": 16},
"GPT-2 Large": {"d_model": 1280, "num_heads": 20},
"GPT-3 (175B)": {"d_model": 12288, "num_heads": 96},
}
# Compute d_k for each
for name, config in model_configs.items():
config["d_k"] = config["d_model"] // config["num_heads"]
The visualization reveals a consistent design pattern: as model dimension increases, the number of heads scales proportionally to maintain per-head dimensions between 64 and 128. This suggests an empirically validated "sweet spot" for head expressiveness. Too few dimensions per head limit each head's ability to capture complex patterns; too many heads with too few dimensions fragment the representation unhelpfully.
Computational Efficiency
Although multi-head attention computes separate attention operations, modern implementations parallelize this efficiently. The key insight is that we can stack all heads into a single tensor and use batched matrix multiplication.
# Efficient batched implementation
def efficient_multihead_attention(x, W_qkv, W_o, num_heads):
"""
Efficient multi-head attention using batched operations.
All heads are computed in parallel using a single matrix multiplication,
then reshaped to separate the head dimension.
"""
batch, seq_len, d_model = x.shape
d_k = d_model // num_heads
# Single projection for all heads: (batch, seq, 3*d_model)
qkv = x @ W_qkv
# Reshape to (batch, seq, 3, num_heads, d_k)
qkv = qkv.reshape(batch, seq_len, 3, num_heads, d_k)
# Transpose to (3, batch, num_heads, seq, d_k) and unpack
qkv = qkv.transpose(2, 0, 3, 1, 4)
Q, K, V = qkv[0], qkv[1], qkv[2]
# Batched attention: (batch, num_heads, seq, seq)
scores = (Q @ K.transpose(0, 1, 3, 2)) / np.sqrt(d_k)
weights = softmax(scores)
# (batch, num_heads, seq, d_k)
attn_output = weights @ V
# Reshape back: (batch, seq, d_model)
output = attn_output.transpose(0, 2, 1, 3).reshape(batch, seq_len, d_model)
return output @ W_oThis batched approach processes all heads with the same matrix operations, using GPU parallelism. The computational cost is essentially equivalent to single-head attention with some reshaping overhead.
Head Pruning
Research has shown that not all attention heads are equally important. Some heads can be removed after training with minimal impact on performance. This head pruning technique reduces computation and memory for inference.
Studies on BERT found that many heads could be pruned without significant accuracy loss, suggesting some redundancy in the learned representations. Michel et al. (2019) found that on the WMT 2014 English-to-German translation task, up to 20 out of 48 heads could be eliminated at test time with less than 0.25 BLEU-point loss. On some GLUE tasks, models with as few as one attention head per layer performed competitively with the full 12-head BERT-base. This has practical implications for deploying models in resource-constrained environments: by identifying and pruning redundant heads before serving a model, you can reduce memory footprint and inference latency without significant quality regression.
The practical takeaway is that you should treat head count as a tunable hyperparameter rather than a fixed design choice. If you are deploying a model to an edge device or a latency-sensitive production system, structured pruning of attention heads is often one of the first compression techniques worth trying.