Multi-head attention runs attention through several heads rather than one set of query, key, and value vectors. Each head can attend to different aspects of an input, giving the combined result a broader view of context. This is useful when a word has more than one possible meaning. section 9.1.1
Eight heads for a short sentence
For “How are you?”, query, key, and value each have shape (1, 6, 256): one sentence, six tokens including beginning- and end-of-sequence tokens, and 256 values per token. Split each into eight heads of 32 values, calculate attention in each head, then join the results into a (1, 6, 256) attention vector. section 9.2.1
To implement the operation:
- Project the inputs into query, key, and value; require the model dimension to be divisible by the number of heads. section 9.2.1
- Reshape and transpose the projected tensors to separate heads. Within each head, scale query–key scores, optionally mask them, apply softmax, and use the resulting weights to combine values. section 9.2.1
- Join the head outputs back into one representation. section 9.2.1
Listing 9.2 Calculating multihead attention
from copy import deepcopy
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
super().__init__()
assert d_model % h == 0
self.d_k = d_model // h
self.h = h
self.linears = nn.ModuleList([deepcopy(
nn.Linear(d_model, d_model)) for i in range(4)])
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
if mask is not None:
mask = mask.unsqueeze(1)
nbatches = query.size(0)
query, key, value = [l(x).view(nbatches, -1, self.h,
self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))] #1
x, self.attn = attention(
query, key, value, mask=mask, dropout=self.dropout) #2
x = x.transpose(1, 2).contiguous().view(
nbatches, -1, self.h * self.d_k) #3
output = self.linears[-1](x) #4
return output
The PyTorch `MultiHeadedAttention` listing shows the head-count argument `h`, the per-head dimension `d_k`, and the projection and reshape before calling `attention()`. section 9.2.1
In a Transformer decoder, masked multi-head self-attention prevents a position from attending to later tokens. A separate cross-attention sublayer then relates the decoder state to the encoder output. The same parallel-head idea also appears in a decoder-only GPT model for text generation and in attention over image features in a denoising U-Net. section 9.3.1 section 12.1.1 section 15.3.1
