Softmax converts a vector of raw scores into nonnegative values that sum to one. In attention, these values become attention weights, expressing how strongly each token attends to every token. section 2.1.1
Scaled dot-product attention applies softmax in a short sequence:
- Compute similarities by taking dot products between query and key vectors. section 2.1.1
- Divide the scores by the square root of the key-vector dimension to stabilize training before applying softmax. section 2.1.1
- Use the resulting attention weights to calculate a weighted combination of the value vectors. section 3.3.3
Softmax in self-attention
For the four-token sentence “Wie geht es dir,” the input embedding has shape 4 × 256. Softmax turns each token’s attention scores into weights whose total is 100%, allowing the resulting weighted values to produce contextually enriched token representations. section 2.1.1
Figure 2.3 A diagram of the self-attention mechanism
Visualizes how softmax fits between query-key scoring and value aggregation in self-attention. 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)
Shows a PyTorch attention head applying softmax to scaled attention scores. section 3.3.3
