Attention, from the ground up.

We have a vector for every input position. Each vector starts with token and position information, but it cannot yet incorporate the rest of the sequence. Attention gives a position a way to gather information from other positions, with weights that depend on the current input.

The central operation is a weighted mixture. We will build that mixture without learned projections, calculate the weights by hand, then put the matrix equation and code around it. If you arrived here directly, the course starts with tokens and vectors .

1. Start with values, not with the formula

Suppose three source positions offer these value vectors:

If the receiving position assigns weights \(a=[0.2,0.3,0.5]\), its output is:

That is all the final multiplication does. Nonnegative weights summing to one make this a convex combination of the values. It usually does not select a single vector exactly. The output is a new mixture, not a copied token ID.

The remaining question is where the weights come from. A fixed average would use the same weights regardless of the query. Attention computes them from a compatibility function, so different receiving positions can gather different mixtures.

Earlier neural translation attention used a decoder state to score encoder annotations and construct a weighted context vector. Bahdanau et al. ↗ , §3, describe that learned alignment mechanism. Dot-product attention is a particular scoring rule, not the definition of every possible attention mechanism.

2. Queries ask, keys match, values contribute

Consider “The robot picked up the cup because it was empty.” Understanding a useful representation of “it” may require information from another position. A fixed average would combine every source equally; an input-dependent mixture can favor different sources in different contexts. This motivates attention, but it does not prove that a particular trained head performs coreference or that one weight explains the whole model.

The receiver needs a way to compare its current features with available sources. Giving matching and carried content separate projections lets the model use one feature space to choose a mixture and another to represent what is mixed. These projections are learned through the task loss, as in the foundations learning step , with many parameters instead of one.

A query is the vector used by a receiving position to score potential sources. A key is a source’s vector for matching. A value is the vector that source contributes to the mixture.

The words “ask” and “match” are analogies. These are learned numbers, not explicit natural-language questions, dictionary keys, or human-interpretable labels. A key is not a discrete identifier: several keys can receive substantial weight.

For self-attention, we obtain the three sets of vectors from the same input states X through different learned matrices:

Separate projections let the model learn different features for compatibility and for the information being transferred. Q and K need equal feature width because we take their dot products. V may have a different feature width; it only needs one value per source key.

The projections are parameters shared across positions. The resulting Q, K, V, and attention weights depend on the input and are recalculated for each sequence. Do not confuse the learned matrices \(W_Q,W_K,W_V\) with those temporary activations.

3. Calculate one query all the way through

Choose \(q=[1,0]\) and keys \(k_1=[1,0]\), \(k_2=[0,1]\), \(k_3=[1,1]\). Keep the three value vectors from the first section.

The unscaled dot products are [1, 0, 1]. The key dimension is two, so the scaled scores are approximately [0.7071, 0, 0.7071].

Softmax exponentiates the scores, then divides by their sum:

The weights are approximately [0.4011, 0.1978, 0.4011]. The output is therefore approximately [1.2033, 0.7967]. The middle source still contributes. Having the smallest score does not mean its weight is zero.

If we change only the values, we change the mixture while keeping the weights fixed. If we change a key, we can change its score and, through the shared denominator, every weight in that query’s row. If we change the query, we change the compatibility calculation for that receiving position.

4. Make every query a row

Stack T queries into Q and S source keys into K. We deliberately use different letters for the two lengths: cross-attention can have different numbers of receiving and source positions.

Q @ K.T computes all pairwise scores. Softmax acts on the last axis, the S source keys. Then (T, S) @ (S, dᵥ) gives (T, dᵥ). The output length follows the queries, not the keys.

5. Why divide by the square root?

A mean is an average; variance is the average squared distance from that mean. For features [−1,1], the mean is zero and the variance is one. Variance describes scale here, not semantic certainty or the probability of a word. “Independent” in the following argument means the random coordinates do not influence one another; it is a simplifying initialization assumption, not a description of every trained activation.

Under a simple initialization model, assume the query and key coordinates are independent, have mean zero, and have variance one. Each coordinate product then has variance one, so a sum over dₖ independent products has variance dₖ. Dividing by \(\sqrt{d_k}\) brings that variance back to one.

This is a scale argument under assumptions, not a promise about every trained query and key. Without scaling, increasing the feature dimension can produce larger logits and very peaked softmax distributions. That can make gradients through softmax small in saturated regions. The derivation appears with the definition in §3.2.1 of Attention Is All You Need ↗ .

A dot product is not automatically cosine similarity. Cosine similarity explicitly divides by both vector norms. Standard scaled dot-product attention does not generally normalize Q and K that way.

6. Change the inputs and watch the result

This experiment adds a fourth key/value pair to our three-position calculation. The labels identify positions; these hand-chosen vectors do not represent learned word semantics. Select a query, adjust the extra temperature control, and optionally hide future positions.

Lowering temperature makes the largest visible scores more dominant. Raising it makes the distribution flatter. A causal mask disallows positions entirely. Changing temperature cannot recover a masked source because its weight remains zero.

Try selecting the first query and enabling the causal mask. Only its own position remains visible, so its weight must be one and its output must equal its own value, [1, 0]. This is a structural consequence of the mask, not a learned behavior.

7. The same operation in two libraries

These functions accept leading batch and head axes as well as simpler two-dimensional matrices. The last two axes always mean position and feature. allowed uses a single course-wide convention: True means a source may participate.

def attention(q: FloatArray, k: FloatArray, v: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
        raise ValueError("Q/K feature widths and K/V source lengths must match")
    scores = (q @ np.swapaxes(k, -2, -1)) / np.sqrt(q.shape[-1])
    if allowed is not None:
        visible = np.broadcast_to(np.asarray(allowed, dtype=bool), scores.shape)
        if not visible.any(axis=-1).all():
            raise ValueError("Every query needs at least one visible key")
        scores = np.where(visible, scores, -np.inf)
    shifted = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(shifted)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ v, weights

def attention(q: Tensor, k: Tensor, v: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
        raise ValueError("Q/K feature widths and K/V source lengths must match")
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    if allowed is not None:
        visible = torch.broadcast_to(allowed.to(device=q.device, dtype=torch.bool), scores.shape)
        if not visible.any(dim=-1).all():
            raise ValueError("Every query needs at least one visible key")
        scores = scores.masked_fill(~visible, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

NumPy subtracts the largest score in each row before exponentiation. This preserves softmax because multiplying every exponential in a row by the same factor cancels between numerator and denominator. It prevents overflow when the original scores are large. PyTorch’s softmax performs the stable calculation internally.

We reject a row with no visible keys. Otherwise, all its logits would be negative infinity and subtracting their maximum would give an undefined result. Our later decoder always has an initial BOS token visible to each valid query.

Run the three-position calculation using the complete NumPy file:

import numpy as np
from numpy_core import attention

q = np.array([[1., 0.]])
k = np.array([[1., 0.], [0., 1.], [1., 1.]])
v = np.array([[1., 0.], [0., 2.], [2., 1.]])
output, weights = attention(q, k, v)

np.testing.assert_allclose(weights.sum(-1), 1)
np.testing.assert_allclose(output, [[1.203336, 0.796664]], atol=1e-5)
print(np.round(weights, 4))
print(np.round(output, 4))

8. What this operation does not tell us

The weights are distributions over source positions, not over the vocabulary. An attention weight of 0.4 is not a 40% probability that a word is the correct next token. A later readout layer makes vocabulary predictions.

A high attention weight is not a complete explanation of a model’s decision. The associated value vector, output projection, residual path, later layers, and prediction head all affect the result. Treat a heatmap as one intermediate calculation to inspect rather than a direct report of the model’s reasoning.

Finally, random or hand-chosen projections demonstrate an operation. Useful context-dependent behavior has to be learned from a training objective. The last lesson will connect output scores to such an objective.

The next missing piece

We can now mix information across positions, but we have not specified which positions are allowed to communicate. In the next lesson, that choice separates bidirectional self-attention, causal self-attention, and cross-attention. The weighted-mixture operation stays the same.