Scaled dot-product attention (SDPA) lets each element in a sequence weigh information from other elements, including itself. In a transformer, this is a common way to compute self-attention: learned query and key vectors determine where to attend, while value vectors supply the information combined in the output. section 2.1.1 section 3.3.3
Compute attention in four stages:
- Project input embeddings X through learned layers to obtain queries Q, keys K, and values V. section 2.1.1
- Calculate query–key dot products and divide the scores by √dₖ, where dₖ is the key-vector dimension. section 3.3.3
- Apply softmax to the scores so each token’s attention weights across the sequence sum to one. section 3.3.3
- Multiply the weights by V to produce the attention output. section 3.3.3
From image patches to attention heads
A vision transformer can turn a 32 × 32 image into 64 patch tokens, each represented by 48 values. Splitting attention into four heads gives each head 12-dimensional queries, keys, and values. Each head computes SDPA, allowing the model to capture different aspects of the patches, such as texture or edges. section 3.3 section 3.3.3
The same calculation also serves denoising: flatten an image into a sequence of pixels, then use SDPA in a U-Net to emphasize relevant spatial regions. There, attention helps address redundant features from convolutional blocks and skip connections while the model predicts noise. section 7.3 section 7.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)
The attention-head listing connects learned query, key, and value projections to matrix multiplication, scaling, and softmax in PyTorch. section 3.3.3
