Query, key, and value are the three inputs to attention. Queries express what to match, keys provide the features used for matching, and values provide the content to aggregate. Attention compares queries with keys, assigns relationship scores, and uses those scores to combine values. section 9.2.1
Attention uses the three inputs in a short sequence of operations:
- Compares each query with keys using a dot product, then scales the scores by the square root of the key dimension. section 9.2.1
- Converts the scaled scores into attention weights with softmax, optionally applying a mask or dropout. section 9.2.1
- Multiplies the attention weights by the values to produce the attention output. section 9.2.1
Multihead attention for “How are you?”
An embedding with shape (1, 6, 256) passes through three neural networks to produce Q, K, and V with the same shape. Multihead attention splits each into eight heads of shape (1, 6, 32), calculates attention independently in each head, and joins the results back into an attention vector of shape (1, 6, 256). Different heads can capture different aspects of meaning, such as the distinct senses of “bank.” section 9.2.1
Listing 9.1 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 implementation makes the data flow explicit: calculate scaled query-key scores, apply an optional mask, normalize with softmax, and aggregate the values. section 9.2.1
