Attention Scores and Weights

A self-attention class can look correct line by line and still produce a result you did not expect. The fastest way to catch that problem is to form a prediction first and then check it against what the code actually does. This article builds that habit using a compact self-attention class written as a subclass of nn.Module.

Start with the input. The input is a matrix with one row per token, and each row is an embedding vector. Before reading any internals, decide what the output should look like. Ask how many rows the output has, and what determines how many columns it has. Those two questions are the foundation of every check that follows.

Many learners meet the word 'weight' for the first time in a neural network course and assume it always means the same thing: a number that training adjusts until the model works. That assumption is reasonable, but once attention enters the picture it becomes a source of confusion. 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 clean way to keep the two apart is to ask when each value is determined. Weight parameters are settled by training and then reused. Attention weights are produced during a forward computation, so the same weight matrix can yield different attention weights for different queries. The diagram below places the two side by side.

The class SelfAttention_v1 inherits from nn.Module. The source describes this base class as a fundamental building block of PyTorch models that supplies the functionality needed to create and manage model layers. Inheriting from it is what lets the class hold trainable parameters in a way PyTorch can track.

Inside the constructor, three trainable weight matrices are created: W_query, W_key, and W_value. Each one is wrapped in nn.Parameter and sized so that it maps the input dimension d_in to the output dimension d_out. The three matrices start from random values, and the same input vectors are multiplied by each one to produce three different views of the data.

The constructor only sets up the starting values. As the model is exposed to more data during training, it adjusts these trainable weights. Learning changes them later, which is why they are called trainable. Nothing in the forward method changes the matrices themselves; it only reads them.

The forward method is where the input actually flows through the mechanism. It takes the input matrix x, projects it into keys, queries, and values, computes a score for every pair of tokens, converts those scores into weights, and then uses the weights to combine the values. Read the method as a pipeline with five stages: projection, scores, scaled softmax weights, and context vectors.

Read each arrow as one matrix operation. The input X is multiplied by three weight matrices to produce queries, keys, and values. Queries and keys are combined into the attention score matrix. Scores are scaled and normalized into attention weights, and those weights are then applied to the values to produce the context vectors Z. The numbers shown in the shape labels are the ones that matter here; the source's figure also notes that its attention weight values are placed for visual consistency rather than being real weights, so the structure is what you should memorize.

Notice what this trace reveals. The number of rows in the output comes from the number of input tokens, and the number of columns comes from d_out, the output dimension chosen for the weight matrices. Context vectors have one row per input token and d_out columns, which makes the output shape a useful first check against what you expected.

The second stage computes attention scores. An attention score is a dot product between a query vector and a key vector, and it is a raw, unnormalized number. 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.

Index positions matter here. 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. Python starts indexing at 0, so the second key is keys[1], not keys[2]. The score is the dot product of the two vectors, giving one number that measures how well that query matches that key.

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 entry at index 1 of the resulting row matches the single score computed earlier, which is a useful quick check when you write this code yourself. The diagram below shows how the matrix is laid out: each row belongs to one query and each column to one key.

The highlighted cell is the score for the second query, which in zero-based Python indexing is row 1, paired with key 1. Because there are six queries and six keys, the full score matrix has one row per query and one column per key.

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

The source code performs the same steps. It reads the key dimension from the shape of the keys, divides the scores by its square root, and applies softmax along the last dimension. Softmax is what turns the scaled scores into the weights. Its outputs are positive and sum to one across the dimension it is applied to, so the attention weights for one query form a distribution over the input tokens. Larger weights mean the context vector leans more heavily on those tokens.

When a result looks wrong, the cause often sits in a step that seems routine. The following mistakes are the ones most likely to produce a value that differs from what you expected.

If the rows do not match, the mismatch tells you which part of the pipeline to examine. Because the forward method is a fixed sequence of matrix operations, you can check its output at each stage rather than only at the end.

The source offers a simple verification. In the output matrix, the second row should match the context vector computed for the second input token in the single-token computation. Matching that row is a strong check, because it confirms that the batched matrix version and the single-token computation agree.

When a score or weight differs from the expected value, work through a fixed checklist in this order: the query source, the key index (remembering that Python indexes from 0), the scaling step, and the softmax axis. Stop at the first checkpoint that fails. That step is where the value first diverges.