A key is an attention vector that represents information for comparison with a query. The query–key similarity determines how strongly one sequence element relates to another; the resulting attention weights select information from value vectors. section 2.1.1 section 3.3.3
A key participates in self-attention through a sequence of transformations:
- Project each input embedding X through a learned linear layer: K = X * W_K. section 3.3.3
- Compare each query with keys using a dot product; larger values indicate stronger similarity. section 3.3.3
- Scale the scores by the square root of the key-vector dimension, then apply softmax to obtain attention weights. section 3.3.3
- Use the weights to combine value vectors, producing the attention output rather than returning the keys themselves. section 3.3.3
Keys across attention heads
With 48-dimensional token representations split across four heads, each head uses 12-dimensional query, key, and value vectors. Different heads can focus on different input aspects, such as texture or edges, while their attention outputs are combined. section 3.3.3
Listing 3.4 Calculating self-attention in each attention head
class AttentionHead(nn.Module):
def __init__(self, hidden_size, attention_head_size, bias=True):
super().__init__()
self.hidden_size = hidden_size
self.attention_head_size = attention_head_size
self.query = nn.Linear(hidden_size, attention_head_size, bias=bias)
self.key = nn.Linear(hidden_size, attention_head_size, bias=bias)
self.value = nn.Linear(hidden_size, attention_head_size, bias=bias)
def forward(self, x):
query = self.query(x)
key = self.key(x)
value = self.value(x) #1
attention_scores = torch.matmul(query, key.transpose(-1, -2))
attention_scores = attention_scores / math.sqrt(\
self.attention_head_size)
attention_probs = nn.functional.softmax(attention_scores,
dim=-1) #2
attention_output = torch.matmul(attention_probs, #3
value) #3
return (attention_output, attention_probs)
Inspect the incremental PyTorch implementation that creates key vectors, computes query–key scores, scales them, and normalizes them with softmax. section 3.3.3
