Multi-head attention runs attention through several sets of query, key, and value vectors instead of just one. Each head can attend to a different aspect of the input, giving a transformer a broader, more contextual representation. For example, 256-dimensional vectors can be divided into eight heads of 32 dimensions each. section 2.1.1
To calculate attention within a head:
- Project the input embeddings into query, key, and value vectors using learned linear layers. section 3.3.3
- Multiply queries by transposed keys, then divide the scores by the square root of the key dimension. section 2.3.1
- Apply softmax to obtain attention weights, then multiply those weights by the values to produce the head’s attention output. A mask can exclude positions before softmax. section 2.3.1
Listing 2.4 Calculating attention based on query, key, and value
def attention(query, key, value, mask=None, dropout=None):
d_k = query.size(-1)
scores = torch.matmul(query,
key.transpose(-2, -1)) / math.sqrt(d_k) #1
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9) #2
p_attn = nn.functional.softmax(scores, dim=-1) #3
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn #4
The PyTorch attention function shows the score, optional mask, softmax, and weighted-value operations used within attention heads. section 2.3.1
Where the heads sit matters. An encoder layer pairs multi-head self-attention with a feed-forward network, layer normalization, and residual connections. A decoder layer first uses masked self-attention so a position cannot attend to later positions; it then uses cross-attention over the encoder output. In that cross-attention step, decoder representations supply queries, while encoder output supplies keys and values. section 2.1.2 section 2.3.2
Listing 2.6 Creating a decoder layer
class DecoderLayer(nn.Module):
def __init__(self, size, self_attn, src_attn,
feed_forward, dropout):
super().__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = nn.ModuleList([deepcopy(
SublayerConnection(size, dropout)) for i in range(3)])
def forward(self, x, memory, src_mask, tgt_mask):
x = self.sublayer[0](x, lambda x:
self.self_attn(x, x, x, tgt_mask)) #1
x = self.sublayer[1](x, lambda x:
self.src_attn(x, memory, memory, src_mask)) #2
output = self.sublayer[2](x, self.feed_forward) #3
return output
The decoder-layer listing traces masked self-attention, encoder–decoder cross-attention, and the feed-forward sublayer in execution order. section 2.3.2
Different relationships in text and images
In “Why is the river so rich? Because it has two banks,” separate heads can attend to different interpretations of “bank.” In an image transformer, four heads divide a 48-dimensional embedding into 12-dimensional heads; they can attend to aspects such as texture or edges. The same mechanism thus works on sequences of words or image patches. section 2.1.1 section 3.3.3
