Simplified Self-Attention with Dot Products

Many learners first meet the word weight in a neural network course and come away with one meaning: a number that training adjusts until the model works. That meaning is reasonable, but it causes a real confusion once self-attention enters the picture. The word appears in two places that play very different roles. One place is the set of weight matrices that a self-attention layer learns. The other is the attention weights that the mechanism computes while it runs.

The rest of this article follows the path from input tokens to attention weights. Keep the two meanings of weight separate as you read, because each value in the computation is either fixed by training or produced fresh for each query.

Weight parameters are the values of a neural network that get optimized during training. They are the fundamental, learned coefficients that define the network's connections. Once training has set them, they describe how the network connects its inputs to its outputs. In a self-attention layer, the weight matrices are these learned parameters.

Attention weights are a separate quantity. They determine the extent to which a context vector depends on the different parts of the input. In other words, they express how strongly the network focuses on each input element when it builds the context vector for one query. Because they are computed fresh for each query, they are dynamic and context-specific rather than fixed coefficients.

Suppose the goal is only to compute a single context vector for one input token. It is still necessary to have key and value vectors for every input element, because all of the keys take part in computing the attention weights with respect to the query. Keys and values for the full set of inputs can be obtained together by matrix multiplication with the key and value weight matrices. This projection maps each input token from its original embedding space into the space of keys and values. The source example projects six input tokens from a three-dimensional space onto a two-dimensional embedding space, so each token ends up with a two-dimensional key and a two-dimensional value.

The second step is to compute attention scores. An attention score is computed as a dot product between a query vector and a key vector. This is the same kind of dot-product computation used in simplified self-attention, with one new aspect: the query and the key are not the raw input elements. They are obtained by transforming the inputs through their respective weight matrices. The result is a raw, unnormalized number that measures how well the query matches that key.

The source computes the score for the second input token first. The query comes from that second token and is paired with the key at index 1. Because the second key is keys[1], an index-based mistake here would pair the query with the wrong key and produce a different score.

Rather than repeating that dot product for every key, the same result can be obtained for all keys at once by multiplying the query by the transposed key matrix. The second entry of the resulting row matches the single score computed for index 1, which is a useful quick check when you write this code yourself.

Attention scores are not yet attention weights. The next step converts the scores by scaling them and then applying the softmax function. Scaling divides each score by the square root of the embedding dimension of the keys. Taking the square root is the same as raising to the power 0.5, so the relationship can be written as follows.

Softmax is what converts the scaled scores into the weights. Its output values are positive and add up to one across the dimension it is applied to. Larger weights mean the context vector leans more heavily on those tokens. Each attention weight routes a fraction of its input embedding into the output context vector for that query, so the context vector is a weighted combination of the input embeddings.

When a computed score or weight does not match the expected value, the source recommends checking four things in a fixed order. Working through them in sequence narrows down the stage where the divergence starts.