Scaled dot-product attention (SDPA) lets each element in a sequence weigh information from every element, including itself. In self-attention, learned projections turn input embeddings into queries, keys, and values: queries and keys determine relevance, while values supply the information combined into the output. Think of a library search: a query is compared with book descriptions, and stronger matches receive more attention. section 9.1.1
From a phrase to attention weights
For “How are you?”, four token embeddings of width 256 form an input of shape (4, 256). Three learned 256 × 256 weight matrices produce Q, K, and V, each also shaped (4, 256). Positional encoding is added to the word embeddings before these projections so the input carries token-position information. section 9.1.1
For each query, the calculation proceeds in three steps:
- Take dot products with the keys and divide the scores by the square root of the key dimension. section 9.2.1
- Apply softmax to turn the scores into attention weights. An optional mask excludes positions before softmax; dropout can then be applied to the weights. section 9.2.1
- Multiply the attention weights by the values to produce the attention output. section 9.2.1
Listing 9.1 Calculating attention based on query, key, and value
def attention(query, key, value, mask=None, dropout=None):
d_k = query.size(-1)
scores = torch.matmul(query,
key.transpose(-2, -1)) / math.sqrt(d_k) #1
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9) #2
p_attn = nn.functional.softmax(scores, dim=-1) #3
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn #4
A PyTorch attention function shows the score calculation, optional mask and dropout, softmax, and returned output and weights together. section 9.2.1
A Transformer uses this calculation in multihead attention rather than relying on one set of attention vectors. With eight heads, 256-dimensional queries, keys, and values are divided into 32-dimensional heads. Each head calculates attention, and their outputs are joined into a 256-dimensional vector. This gives different heads room to attend to different aspects of a phrase, including differing meanings of a word such as “bank.” section 9.2.1
The same calculation on image pixels
In a denoising U-Net, SDPA is applied to image features rather than text tokens. For an input shaped (1, 128, 64, 64), each feature channel is flattened into a sequence of 4,096 pixel positions. Projected queries, keys, and values are split into four heads; attention helps the model emphasize relevant image features during denoising. section 15.3.1 section 15.1.2
