Self-Attention Mechanism

Read a sentence and you constantly decide which other words matter for understanding a given word. A pronoun such as 'it' is meaningless until you connect it to something earlier in the text. A word near the end of a long passage may depend on a phrase introduced many words before it. Language is full of these links, and they do not always sit next to each other.

Transformers and large language models are built around a component called the self-attention mechanism, which is designed to handle exactly this problem. Its job is to let the model weigh the importance of different words or tokens in a sequence relative to each other. Because the weighing happens between the words of the input itself, the mechanism is described as self-attention.

This article follows one concrete implementation, a compact class called SelfAttention_v1, and traces it from its input to its output. Before you read the internals, you will practise forming a prediction about what the output should look like. The source notes that a self-attention class can look correct line by line and still produce a result you did not expect, and that the fastest way to catch this is to predict first and then check against what the code actually does.

The input is a matrix with one row per token, and each row is an embedding vector. Suppose the input has 6 tokens and each embedding has 3 values, so the input is 6 rows by 3 columns. The three weight matrices each map the input dimension d_in to the output dimension d_out. Use d_in equal to 3 and d_out equal to 2.

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 is 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 sketch above is an illustrative reconstruction consistent with the description in the source, which explains each stage in prose rather than printing the full class. The line that divides by the square root of the key dimension is the scaling step that the source calls scaled softmax.

The constructor only sets up the starting values. The source explains that 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. This distinction matters when you debug: the matrices you see at construction time are not the ones the model ends up with.

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: receive the input, project it into queries, keys, and values, score every pair of tokens, turn the scores into weights, and combine the values using those weights.

The diagram traces the same flow the forward method performs, but as data flow instead of 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 the weights are then applied to the values to produce the context vectors. Notice that the scores and weights sit between the projections and the final output. When something goes wrong, you can check the output of each stage rather than only the final result.

Take a single token, say token 2. Its query is compared against the key of every token in the sequence, including itself. Each comparison yields a score, and the scores become weights. The context vector for token 2 is then a weighted blend of the value vectors of all tokens, so tokens that score highly contribute more to it. Each token therefore gets its own context vector, shaped by its surroundings.

The example above shows the kind of judgement self-attention is meant to make: deciding which parts of the input bear on which other parts. It does not show how the model computes those judgements in detail. The source introduces the mechanism at this conceptual level and says nothing more specific about the arithmetic inside it.

Shape tracing is the most reliable way to find a bug before reading any numbers. The example below is an illustrative case 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. Follow each intermediate tensor from the input to the final output.

The source offers a simple way to verify the implementation. In the output matrix, the second row should match the context vector computed for the second input token in the 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 its output at each stage rather than only at the end. Compare the shapes first, then the projections, then the scores, then the weights, then the context vectors. The first stage that disagrees with your prediction is where the unexpected result begins.

Self-attention is what allows a model to relate every position in an input to other positions, rather than reading each word in isolation. The source identifies two closely linked benefits. First, the mechanism enables the model to capture long-range dependencies, meaning relationships between words that sit far apart. Second, it captures contextual relationships within the input data, so that a token's role is shaped by its surroundings.

The reason this works is that a token's query is compared directly with the key of every other token in one step. A distant token is therefore no harder to reach than a neighbouring one. This is the contrast with a model that passes information only from one position to the next, where a distant clue would have to travel through every intermediate position.

The source connects these capabilities to output quality. By capturing these relationships, the mechanism enhances the model's ability to generate output that is coherent and contextually relevant. The reason self-attention matters is not only that it sees many words at once, but that it uses those words to judge what each one means in its setting.

The self-attention mechanism appears within a larger structure. A transformer is a deep learning model architecture for language translation. In the simplified depiction the source gives, it is made of two parts: an encoder and a decoder. In the original transformer architecture, the encoder processes the complete input text and produces an embedding representation, and the decoder uses that representation to generate output one word at a time.

Consider one translation. The input text is the sentence 'This is an example', and the partially translated sentence is 'Das ist ein'. The final stage is where the decoder has to generate only the last word, 'Beispiel', to complete the translation. The encoder has already read the whole input, so the decoder's job at this point is to produce the remaining output word using the representation it was handed.

Keeping these jobs separate makes it easier to see where the attention-based weighing of words does its work. The encoder is where the input is read as a whole, and the decoder depends on the representation that results. The source's introductory treatment places self-attention within this two-part structure without describing the internal arithmetic of either part in detail.

When you reach the detailed treatment, keep the structure from this article in mind. The shapes you predicted, the five stages of the forward method, and the row-by-row check will still apply. The later material supplies the mathematics behind each stage.