Queries, Keys, and Values

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.

The input to the class is a matrix with one row per token, and each row is an embedding vector. Before you read the internals, decide what the output should look like. Ask two questions: how many rows does the output have, and what determines how many columns it has?

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. In shape terms, each matrix has d_in rows and d_out columns, so multiplying the input matrix by one of them produces a matrix with d_out columns.

The three matrices start from random values, and the same input vectors are multiplied by each one. That produces three different views of the same data: queries from W_query, keys from W_key, and values from W_value. The three matrices are separate parameters, so each view can be shaped differently during training even though they all start from the same input.

The constructor only sets up starting values. The source explains that as the model is exposed to more data during training, it adjusts these trainable weights. This is why the matrices are called trainable: the constructor creates them, and learning changes them later.

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, where each stage's output is the next stage's input.

The diagram follows the same sequence the forward method performs, written as a data flow rather than as code. The input matrix X is multiplied by three weight matrices to produce queries, keys, and values. Queries and keys are combined into the attention weight matrix, and those weights are then applied to the values to produce the context vectors Z.

Each of the three stages below corresponds to a group of lines in the forward method. Each is explained in its own subsection so you can locate the line that produces each kind of output.

Shape tracing is the most reliable way to find a bug before you read any numbers. The example below was written for this article, not taken from the source. It uses six input tokens, an input dimension of 3, and an output dimension of 2. Trace each intermediate tensor and record which dimension changes at each step.

Reading the shapes this way gives you a check you can perform without running any numbers. If your output has 6 rows and 2 columns, the row count and the output dimension are correct. If it has 3 columns, a projection probably used d_in where it should have used d_out. If it has 6 rows but a wrong column count, the context stage is the first place to look.

The three projections play different roles even though they all come from the same input. The queries and keys are compared with each other to produce scores and then weights. The values are not used when computing the weights. They are only used in the final stage, where the weights combine them into context vectors.

This dependency explains a useful debugging fact. A mistake in how queries or keys are built will change the attention weights, and therefore change every context vector. A mistake that only affects the values changes the context vectors directly, but leaves the weights alone.

Suppose the context vectors look wrong. The source offers a verification method: in the output matrix, the second row should match the context vector computed for the second input token in the previous single-token computation. Matching a row like this is a strong check, because it confirms that the batched matrix version and the single-token computation agree.

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 output at each stage rather than only at the end. Work backward from the context vectors, or forward from the projections, until you find the first stage where the values diverge from what you predicted.