An attention score measures how relevant a query is to a key. Scaled dot-product attention computes it by multiplying Q by the transpose of K, then dividing by the square root of the key dimension. section 2.1.1
Trace an attention score through the mechanism:
- Project input embeddings into query, key, and value vectors. section 3.3.3
- Calculate QKᵀ; larger dot products indicate stronger query-key similarity. section 2.1.1
- Scale the dot product by √dₖ to prevent large magnitudes when embeddings have high dimensionality. section 2.1.1
- Apply softmax to turn scores into attention weights, then combine the value vectors using those weights. section 3.3.3
Recommendation and contextualization
For a query such as “generative AI” or “PyTorch,” book titles and descriptions act as keys: stronger query-key similarity produces higher-ranked values. For “Wie geht es dir,” attention transforms four 256-dimensional token embeddings into context-sensitive representations with the same dimensions. section 2.1.1
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 AttentionHead implementation, including query, key, and value projections, scaled matrix multiplication, and softmax. section 3.3.3
