Multi-Head Attention › Implementation
Language AI HandbookAll chapters Multi-Head Attention › ImplementationOne section · Sections in this chapter6 / 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
With the mathematics clear, let's translate these ideas into working code. We'll build multi-head attention from scratch using NumPy, making every operation explicit so you can see exactly how the formulas map to implementation. This hands-on approach will solidify your understanding and prepare you to work with production implementations in PyTorch or TensorFlow.
Our implementation needs two building blocks: first, the scaled dot-product attention we covered in the previous chapter; second, the multi-head wrapper that orchestrates multiple attention operations in parallel.
Building Block 1: Scaled Dot-Product Attention
We start with a function that computes attention for tensors that already include a head dimension. This allows us to process all heads simultaneously through batched matrix multiplication:
import numpy as np
def softmax(x, axis=-1):
"""Numerically stable softmax."""
exp_x = np.exp(x - x.max(axis=axis, keepdims=True))
return exp_x / exp_x.sum(axis=axis, keepdims=True)
def scaled_dot_product_attention(Q, K, V, mask=None):
"""
Compute scaled dot-product attention.
Args:
Q: Queries, shape (batch, heads, seq_len, d_k)
K: Keys, shape (batch, heads, seq_len, d_k)
V: Values, shape (batch, heads, seq_len, d_v)
mask: Optional mask, shape broadcastable to (batch, heads, seq_len, seq_len)
Returns:
output: Attention output, shape (batch, heads, seq_len, d_v)
weights: Attention weights, shape (batch, heads, seq_len, seq_len)
"""
d_k = Q.shape[-1]
# Compute attention scores
scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_k)
# Apply mask if provided
if mask is not None:
scores = np.where(mask, scores, -1e9)
# Softmax to get attention weights
weights = softmax(scores, axis=-1)
# Weighted sum of values
output = weights @ V
return output, weightsNotice the tensor shape: (batch, heads, seq_len, d_k). By placing the head dimension second, we can compute attention scores for all heads at once. The matrix multiplication Q @ K.transpose(...) produces a (batch, heads, seq_len, seq_len) tensor containing separate attention matrices, one per head. This is the key to efficient parallel computation.
Building Block 2: The Multi-Head Attention Module
Now we build the main module. The MultiHeadAttention class encapsulates all the projection matrices and orchestrates the four-step process we described earlier.
A key implementation detail: instead of storing separate , , matrices for each head, we use a single combined projection matrix W_qkv of shape (d_model, 3 * d_model). This computes all QKV representations for all heads in one matrix multiplication. The result is then split and reshaped to separate the heads. This approach is more efficient on GPUs because it maximizes parallelism and minimizes memory transfers.
class MultiHeadAttention:
"""
Multi-head attention implementation.
Args:
d_model: Model dimension
num_heads: Number of attention heads
"""
def __init__(self, d_model, num_heads, seed=42):
assert d_model % num_heads == 0, (
"d_model must be divisible by num_heads"
)
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
# Initialize projection matrices with small random values
np.random.seed(seed)
scale = np.sqrt(2.0 / d_model)
# Combined QKV projections: (d_model, 3 * d_model)
self.W_qkv = np.random.randn(d_model, 3 * d_model) * scale
# Output projection: (d_model, d_model)
self.W_o = np.random.randn(d_model, d_model) * scale
def split_heads(self, x):
"""
Split the last dimension into (num_heads, d_k).
Args:
x: Input tensor, shape (batch, seq_len, d_model)
Returns:
Reshaped tensor, shape (batch, num_heads, seq_len, d_k)
"""
batch, seq_len, _ = x.shape
x = x.reshape(batch, seq_len, self.num_heads, self.d_k)
return x.transpose(0, 2, 1, 3)
def combine_heads(self, x):
"""
Inverse of split_heads.
Args:
x: Input tensor, shape (batch, num_heads, seq_len, d_k)
Returns:
Reshaped tensor, shape (batch, seq_len, d_model)
"""
batch, _, seq_len, _ = x.shape
x = x.transpose(0, 2, 1, 3)
return x.reshape(batch, seq_len, self.d_model)
def forward(self, x, mask=None):
"""
Compute multi-head self-attention.
Args:
x: Input tensor, shape (batch, seq_len, d_model)
mask: Optional attention mask
Returns:
output: Attention output, shape (batch, seq_len, d_model)
weights: Attention weights per head, shape (batch, num_heads, seq_len, seq_len)
"""
batch, seq_len, _ = x.shape
# Step 1: Project to Q, K, V (combined for efficiency)
qkv = x @ self.W_qkv # (batch, seq_len, 3 * d_model)
# Split into Q, K, V
Q, K, V = np.split(qkv, 3, axis=-1)
# Step 2: Split into multiple heads
Q = self.split_heads(Q) # (batch, num_heads, seq_len, d_k)
K = self.split_heads(K)
V = self.split_heads(V)
# Step 3: Compute attention for all heads in parallel
attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
# Step 4: Combine heads
combined = self.combine_heads(attn_output) # (batch, seq_len, d_model)
# Step 5: Output projection
output = combined @ self.W_o
return output, attn_weightsThe split_heads and combine_heads methods handle the tensor reshaping that makes parallel head computation possible. split_heads takes a tensor of shape (batch, seq_len, d_model) and reorganizes it to (batch, num_heads, seq_len, d_k), separating the head dimension. combine_heads reverses this transformation after attention is computed.
Testing the Implementation
Let's verify our implementation produces the expected output shapes and that the attention weights are properly normalized:
# Create a small example
batch_size = 1
seq_len = 6
d_model = 64
num_heads = 4
# Random input sequence
np.random.seed(123)
x = np.random.randn(batch_size, seq_len, d_model)
# Create multi-head attention module
mha = MultiHeadAttention(d_model, num_heads)
# Forward pass
output, attention_weights = mha.forward(x)Input shape: (1, 6, 64) Output shape: (1, 6, 64) Attention weights shape: (1, 4, 6, 6) Head dimension (d_k): 16 Number of heads: 4
The output tensor has shape (1, 6, 64), matching the input exactly. This dimensional consistency is essential for residual connections in transformer architectures. The attention weights have shape (1, 4, 6, 6), representing a separate attention matrix for each of the 4 heads.
Attention weights sum per query position (head 1): [1. 1. 1. 1. 1. 1.]
Each row sums to exactly 1.0, confirming that softmax normalization is working correctly. This is expected: for each query position, the attention weights across all key positions must form a valid probability distribution.
How Multi-Head Attention Transforms Representations
Let's visualize how multi-head attention changes the token representations. We'll use PCA to project the 64-dimensional embeddings down to 2D, allowing us to see how the representations shift after passing through the attention layer.
from sklearn.decomposition import PCA
# Get the input and output representations
input_embeds = x[0] # Shape: (6, 64)
output_embeds = output[0] # Shape: (6, 64)
# Combine for PCA fitting
combined = np.vstack([input_embeds, output_embeds])
# Fit PCA on combined data
pca = PCA(n_components=2)
combined_2d = pca.fit_transform(combined)
# Split back
input_2d = combined_2d[:6]
output_2d = combined_2d[6:]
# Token labels for this example
tokens_example = [f"Pos {i}" for i in range(6)]
The visualization shows how multi-head attention repositions each token's representation in embedding space. The arrows indicate the direction and magnitude of change. Tokens that attended strongly to semantically related positions move toward each other, while tokens that focused on different aspects of the input shift in different directions. This is the contextual enrichment in action: each position now carries information gathered from across the sequence.