Self-attention with trainable weights

Earlier in this chapter, a context vector was built for one input token at a time. That approach works, but it means repeating the same steps for every position in a sequence. This article shows how the attention weights you already computed can be applied to all the inputs in one operation, producing a full set of context vectors together.

The whole computation is a single line. It takes the attention weight matrix named attn_weights, multiplies it by the inputs with the @ operator, and stores the result in all_context_vecs. Before stepping through why that line works, read the tensor it produces and reason backward from the output to the operation that made it. This habit matters for debugging: when a result looks wrong, you need to know which operation produced which part of it.

The key idea is that a single matrix product replaces a loop over tokens. For each row of the attention weights, the product forms a weighted combination of the input vectors. Doing this for every row at once gives all context vectors together. The operands flow into the product, and the result lands in one tensor.

Each row of the output holds one context vector. Row one belongs to the first input token, row two belongs to the second, and so on. The rows are aligned with the rows of the attention weight matrix, because each row of attention weights describes how much one token attends to every input token. Reading the output therefore means reading one row per token.

Shape reasoning lets you check a result before you inspect its values. Suppose attn_weights has shape $n \times n$, one row and one column per input token, and inputs has shape $n \times d$, one row per token with $d$ values each. The product consumes the shared dimension $n$: each row of attn_weights has $n$ entries that line up with the $n$ rows of inputs. The result therefore has one row per token and $d$ columns, giving shape $n \times d$.

The following toy example uses three tokens with two values each. The weights are chosen so the arithmetic can be checked by hand. Each row of attn_weights sums to one, and each row of the result is a weighted mix of the input rows.

Check the second row by hand. Row two of attn_weights is [0, 1, 0], so it selects the second input token alone, giving [3, 4]. Row three is [0.25, 0.25, 0.5], so it blends the three inputs: 0.25 times [1, 2] plus 0.25 times [3, 4] plus 0.5 times [5, 6] gives [3.5, 4.5]. Each printed row can be reproduced from one row of weights and the input rows.

A useful check is to compare one row of the batched result with a context vector computed earlier by a different route. Earlier, the second context vector was computed one token at a time, so its values should match the second row of the batched result. Printing that earlier vector alongside the new tensor gives a direct comparison.

This check is a debugging technique, not only a confirmation. If a single row disagrees, the problem lies in the batched path for that token, so inspect the attention weights in that row. If every row disagrees, the problem is more likely in the shared inputs or in how the weights were built. A mismatch in one row points to a specific row of attn_weights; a mismatch everywhere points to the shared setup.

The mechanism you have just run is a simplified version of attention. The next step adds trainable weights, which lets the model learn from data and improve its performance on specific tasks. The extended mechanism is called self-attention with trainable weights, also known as scaled dot-product attention. It is the self-attention used in the original transformer architecture, in GPT models, and in most other popular large language models.

The batched product itself stays the same: a single matrix multiplication still yields a full set of context vectors. What changes is the origin of the weights. In the trainable version, the weights are learned rather than fixed, so the same product is applied with parameters the model adjusts during training.

The source describes four attention variants built in sequence. The simplified attention mechanism comes first, followed by self-attention with trainable weights. The later stages add causal attention and multi-head attention. The attention step is one stage in a longer pipeline for building a language model, not the whole model.