MultiHeadAttention integrates multi-head attention in one class. It projects inputs into queries, keys, and values, reshapes each projected tensor into multiple heads, computes attention, and combines the head outputs. section 3.6.2
Its implementation follows this data flow:
- Check that d_out is divisible by num_heads, then derive head_dim. section 3.6.2
- Apply linear projections for queries, keys, and values before reshaping and transposing them into separate heads. section 3.6.2
- Compute attention for the heads, concatenate their results, and pass the combined representation through out_proj. section 3.6.2
Two heads with a small batch
With d_out=2 and num_heads=2, a batch produces context vectors shaped [2, 6, 2]. Separate head calculations reproduce the corresponding attention-score matrices. section 3.6.2
Listing 3.5 An efficient multi-head attention class
class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out,
context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert (d_out % num_heads == 0), \
"d_out must be divisible by num_heads"
self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads #1
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out) #2
self.dropout = nn.Dropout(dropout)
self.register_buffer(
"mask",
torch.triu(torch.ones(context_length, context_length),
diagonal=1)
)
def forward(self, x):
b, num_tokens, d_in = x.shape
keys = self.W_key(x) #3
queries = self.W_query(x) #3
values = self.W_value(x) #3
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) #4
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(
b, num_tokens, self.num_heads, self.head_dim
)
keys = keys.transpose(1, 2) #5
queries = queries.transpose(1, 2) #5
values = values.transpose(1, 2) #5
attn_scores = queries @ keys.transpose(2, 3) #6
mask_bool = self.mask.bool()[:num_tokens, :num_tokens] #7
attn_scores.masked_fill_(mask_bool, -torch.inf) #8
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)
context_vec = (attn_weights @ values).transpose(1, 2) #9
#10
context_vec = context_vec.contiguous().view(
b, num_tokens, self.d_out
)
context_vec = self.out_proj(context_vec) #11
return context_vec
Inspect the PyTorch class skeleton, including projections, head_dim, dropout, causal mask registration, and output projection. section 3.6.2
