Calculate attention, then change what it can see

Two positions offer different information. Your current query matches the first more strongly, but both may contribute. Attention turns those match scores into a mixture. What changes if one position is not allowed to be read?

Before you begin: Be comfortable with vector dot products and reading short Python code.

Give each quantity a role

A query is a learned vector used to score potential sources. A key is a learned vector compared with the query. A value is the vector contributed to the output. In self-attention, all three are computed from representations in the same sequence, usually through different learned projections.

For one head, let the query width be d_k. A query q and key k_j receive the score dot(q, k_j) / sqrt(d_k). The dot product multiplies corresponding coordinates and sums them. Scaling by the square root helps control score magnitude under the assumptions motivating the original design.

Softmax converts the allowed scores into nonnegative weights that sum to one. The output is the weighted sum of the value vectors. The score decides the weight; it is not itself the information being returned.

Work one row by hand

Use q = [1, 0], keys [1, 0] and [0, 1], and values [10, 0] and [0, 20]. The unscaled scores are one and zero. Dividing by the square root of two gives approximately 0.707 and zero.

Softmax assigns weights of about 0.670 and 0.330. The output is approximately [6.698, 6.605]. Both sources contribute. These vectors are hand-chosen for arithmetic; they are not observed attention weights from a trained language model.

Now mask the second position. Its contribution becomes zero and the first weight becomes one, giving [10, 0]. The query and values did not change. The allowed information changed.

Run the calculation

This complete Python 3 program uses only the standard library. It computes one attention row and rejects an all-masked row, whose softmax would not be meaningful here.

# Runnable: Python 3, standard library.
import math

def attend(query, keys, values, allowed):
    assert len(keys) == len(values) == len(allowed)
    assert query and keys and any(allowed)
    assert all(len(k) == len(query) for k in keys)
    width = len(values[0])
    assert width and all(len(v) == width for v in values)
    scores = [sum(x*y for x, y in zip(query, key))
              / math.sqrt(len(query)) for key in keys]
    peak = max(s for s, ok in zip(scores, allowed) if ok)
    numerators = [math.exp(s-peak) if ok else 0.0
                  for s, ok in zip(scores, allowed)]
    weights = [x / sum(numerators) for x in numerators]
    output = [sum(w*v[j] for w, v in zip(weights, values))
              for j in range(width)]
    return weights, output

keys = [[1, 0], [0, 1]]
values = [[10, 0], [0, 20]]
weights, output = attend([1, 0], keys, values, [True, True])
assert abs(sum(weights)-1) < 1e-12
assert attend([1, 0], keys, values, [True, False])[1] == [10, 0]
print([round(x, 3) for x in weights])
print([round(x, 3) for x in output])

Subtracting the largest allowed score before exponentiation improves numerical stability and leaves softmax unchanged. This teaching function uses assertions for a small fixed experiment; a reusable library would validate inputs with explicit errors and handle batching and tensor dtypes.

Extend the row to a sequence

Stacking query rows gives Q; stacking key rows gives K; stacking value rows gives V. The compact expression is softmax(QKᵀ / sqrt(d_k) + mask)V. The transpose aligns key coordinates for pairwise dot products. Softmax is applied across source positions for each query row.

A causal mask allows a position to read itself and earlier positions while blocking later ones. A padding mask excludes artificial padding tokens. They solve different problems and may be combined. If the mask is wrong during training, the model may read the answer it is supposed to predict.

Explore a change in information

Keep the query and keys fixed, but change the second value to [100, 100]. Will the weights change? What about the output?

Next, we will compare where attention reads from in encoder-only, decoder-only, and encoder-decoder systems.

Source

Attention Is All You Need defines scaled dot-product attention. The program above demonstrates the operation without training a language model.

Continue to the next lesson.