Attention as content-dependent weighted retrieval
How to study this lesson
Read one section, trace the worked example, and try the knowledge check. Bookmark saves a shortcut; notes appear in My notebook. Mark lesson complete is your own assessment and does not mark its lab passed.
By the end, you can
- Calculate scaled dot-product attention.
- Explain the roles of queries, keys, and values.
- Apply causal masking and stable softmax.
Queries ask, keys address, values contribute
Attention produces an output by weighting a collection of value vectors according to how well their associated keys match a query. The query and key coordinates determine compatibility; the value coordinates determine what information is combined. In a transformer, learned projections create queries, keys, and values from hidden representations. These names describe computational roles, not literal database fields.
Imagine a query [1,0], keys [1,0] and [0,1], and values [10,0] and [0,20]. The query aligns more strongly with the first key, so the weighted output leans toward the first value. The output need not equal either stored value. Attention is a mixture, and its ability to blend information is as important as its ability to emphasize one position.
From compatibility scores to a mixture
Scaled dot-product attention computes each score as q dot k divided by the square root of the key dimension. Softmax exponentiates the scores and divides by their sum, producing nonnegative weights that sum to one. The final output is the weighted sum of values. With scores [1,0], the weights are approximately [0.731,0.269]. Values [10,0] and [0,20] then yield approximately [7.311,5.379].
The division by the square root of dimension controls score scale under assumptions about coordinate variance. It is not a universal normalization of semantic similarity. Changing score scale changes softmax sharpness: multiplying all scores by a large positive constant concentrates weight on the largest score, while shrinking them moves toward a uniform mixture.
Stable softmax preserves the distribution
Directly computing exp(1000) can overflow even though the final normalized probabilities are well behaved. Subtract the maximum score before exponentiating. Softmax is unchanged because the common exponential factor cancels between numerator and denominator. Scores [1000,999] therefore become [0,-1] for exponentiation, preserving their relative preference without enormous intermediate values.
Subtracting the maximum solves one numerical problem; it does not validate all inputs. Nonfinite values, empty lists, and a row with no allowed keys need explicit handling. In the lab, an all-masked row is rejected because no probability distribution exists over an empty allowed set. Some tensor kernels choose a different convention, so read their behavior before using them as a reference for edge cases.
Masks express information availability
A causal language model predicting the next token must not use later tokens from the training sequence. For position t, a causal mask permits keys only from positions at or before t. Disallowed positions receive zero weight after masked normalization. Masking is a statement about what information is available, not just a performance optimization.
Suppose the sequence has three positions. The first row can attend only to the first key, so its allowed weight is exactly one. The second can mix the first two; the third can mix all three. If future keys are accidentally allowed during training, loss can look impressively low because the model sees information unavailable at generation time. This is a particularly direct form of leakage.
Multiple heads and interpretation limits
Several attention heads apply different learned projections and produce different mixtures. Their outputs are combined through another learned projection. This lets a layer represent several compatibility patterns at once. The heads are not guaranteed to become neatly named specialists, and an attention weight alone is not a complete explanation of why a final prediction occurred: values, other heads, residual paths, and later layers also matter.
The example computes one query against two keys, once with both allowed and once with only the first. It is a transparent attention calculation, not a transformer model or a production attention kernel. Trace the output as a weighted sum before considering faster implementations. That understanding will make the next module's cache layout and grouped-query attention much easier to reason about.
Work through the code
The query is scaled so the attention scores are exactly [1,0]. The first run blends both values; the second masks out the second value. This compact demo assumes valid shapes and at least one allowed key. The lab turns those assumptions into an explicit contract.
import math
def attend(query, keys, values, allowed):
scores = [sum(q * k for q, k in zip(query, key)) / math.sqrt(len(query))
for key in keys]
maximum = max(score for score, use in zip(scores, allowed) if use)
masses = [math.exp(score - maximum) if use else 0.0
for score, use in zip(scores, allowed)]
total = sum(masses)
weights = [mass / total for mass in masses]
context = [sum(weight * value[j] for weight, value in zip(weights, values))
for j in range(len(values[0]))]
return weights, context
query = [math.sqrt(2), 0]
keys = [[1, 0], [0, 1]]
values = [[10, 0], [0, 20]]
for mask in ([True, True], [True, False]):
weights, context = attend(query, keys, values, mask)
print("weights:", [round(x, 3) for x in weights])
print("context:", [round(x, 3) for x in context])
weights: [0.731, 0.269] context: [7.311, 5.379] weights: [1.0, 0.0] context: [10.0, 0.0]
Run Python snippets locally with the prerequisites named above. The module coding lab runs directly in your browser.
Pause and reason
If two allowed keys have identical scores and values [2,6] and [8,0], what is the attention output? What changes if you add 100 to both scores?
Check your understanding
Why must a decoder training example mask future token positions?
Your notes
Explain the mechanism in your own words. Add a failure you want to test.