A query is an attention vector representing the information whose relevance to sequence elements is being evaluated. Attention compares each query with keys, assigns relationship scores, and combines values according to the resulting weights. section 2.1.1 section 3.3.3
The mechanism begins with token or feature embeddings. Learned linear layers transform an input sequence X into Q, K, and V. A dot product between Q and K measures similarity; scaling and softmax turn those scores into attention weights, which weight the value vectors. section 2.1.1 section 3.3.3
In multi-head attention, queries support the following sequence of operations:
- Project each input through a learned query layer to produce Q, alongside corresponding key and value projections. section 3.3.3
- Split the projected vectors across heads so different heads can focus on different aspects of the input. section 3.3.3
- Compare queries with keys using scaled dot products, apply softmax, and use the weights to aggregate values. section 3.3.3
- Concatenate the per-head attention outputs into one attention tensor for the surrounding model. section 7.3.1
From translation to vision
For “Wie geht es dir,” token embeddings and positional encodings enter attention so the model can relate words while generating an English translation. In a vision transformer, four query heads divide a 48-value representation into 12-value projections, allowing different heads to examine aspects such as texture or edges. section 2.1.1 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 PyTorch AttentionHead implementation, including query projection, query-key scoring, scaling, softmax, and value aggregation. section 3.3.3
