Dot-product attention assembles every piece covered so far in this category โ Q/K/V and attention scores โ into the complete mechanism: scores, normalized via softmax into weights, used to combine values into a final output.
Formula
Reading left to right: compute every query's raw score against every key (\(\mathbf{Q}\mathbf{K}^\top\), a matrix of scores), convert each query's row of scores into a valid probability distribution via softmax (see Softmax Function), then use those probabilities as weights for a weighted sum of the value vectors.
Numerical Example โ Continuing From Attention Score
Using the scores from Attention Score: \([1.0, 0.0, 0.7]\). Applying softmax:
With values \(\mathbf{v}_1=[1,0]\), \(\mathbf{v}_2=[0,1]\), \(\mathbf{v}_3=[0.5,0.5]\):
The output is a weighted blend of every value, dominated by \(\mathbf{v}_1\) (since key 1 scored highest against the query) but with genuine contributions from every other value too โ exactly the "soft," differentiable selection mechanism attention provides, as opposed to a hard, all-or-nothing choice.
Code
import torch
import torch.nn.functional as F
Q = torch.tensor([[1.0, 0.0]]) # one query
K = torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.7, 0.7]]) # three keys
V = torch.tensor([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]) # three values
scores = Q @ K.T # (1, 3) -- raw scores
weights = F.softmax(scores, dim=-1) # normalize into a probability distribution
output = weights @ V # weighted sum of values
print(weights) # tensor([[0.4741, 0.1743, 0.3516]]) -- matches the hand-worked example
print(output) # tensor([[0.6483, 0.3517]])
Why This Is a "Soft," Differentiable Lookup
Unlike a hard lookup (picking exactly one best-matching key/value), this weighted-sum approach is fully differentiable โ gradients can flow back through every step (the scores, the softmax, the weighted sum) via backpropagation, letting the entire attention mechanism, including the Q/K/V projection weights, be learned end-to-end alongside the rest of the network, exactly like any other layer.
Common Mistakes
- Forgetting the softmax step and using raw scores directly as weights โ without normalization, the weighted sum wouldn't represent a valid convex combination of the values, and could produce nonsensical (unbounded, or even negative-weighted) outputs.
- Applying softmax across the wrong dimension for a full batch of queries โ softmax must normalize each query's row of scores independently (across the key dimension), not across queries or across the batch.
Interview Relevance
Q: "Why is dot-product attention described as a 'soft' lookup, and why does that matter for training?" Instead of selecting exactly one best-matching key/value pair, attention computes a weighted combination of all values, weighted by softmax-normalized similarity scores โ every value contributes something, proportional to relevance. This makes the entire operation smooth and differentiable, so gradients can flow through it during backpropagation, allowing the whole mechanism (including the learned Q/K/V projections) to be trained end-to-end via standard gradient descent.
Practice Question
If one key's score is vastly larger than all the others (e.g. 100 vs single-digit scores for the rest), what will softmax do to the resulting attention weights, and what will the output approximately equal?