Attention weights for a single input token

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 of the inputs in one operation, producing a full set of context vectors together.

The central idea is that a single matrix product replaces a loop over tokens. Instead of asking, for each token in turn, how much it should borrow from every other token, you store all the attention weights in one matrix and all the input vectors in another, then multiply them. The result contains every context vector at once.

The whole computation is one line. It takes the attention weight matrix named attn_weights, multiplies it by the inputs using the @ operator, and stores the result in all_context_vecs. A second line prints that result so you can read it.

The method the source recommends is to run the line first, read the printed tensor, and only then reason backward about the operation that produced it. When a result looks wrong, you need to know which operation produced which part of it, and reading the output before stepping through the execution gives you a concrete picture to test against.

Each row of the output tensor holds one context vector. Row number one belongs to the first input token, row number two belongs to the second input token, and so on. The rows line up with the rows of the attention weight matrix, because each row of attention weights describes how much one token attends to every input token.

For a single row, the product forms a weighted combination of the input vectors. Each weight in that row scales one input vector, and the scaled vectors are summed. Written for a token $i$ with weights $\alpha_{i1}, \alpha_{i2}, \ldots, \alpha_{in}$ and input vectors $x_1, x_2, \ldots, x_n$, the context vector for that row is

The diagram shows the bookkeeping. Row $i$ of the weights is combined with the whole input matrix, and the result lands in row $i$ of the output. Nothing in the multiplication moves a row to a different position, so the index of a weight row is also the index of the context vector it produces.

The shape rule for matrix multiplication is the same one you use for any product. The number of columns in the first operand must match the number of rows in the second. The result takes the row count of the first operand and the column count of the second. Because the weights are square, the inner dimension matches the input token count automatically.

The source's debugging technique is to compare one row of the batched result with a context vector computed earlier by a different route. The earlier vector was built one token at a time, so if the batched and per-token paths agree, their values should match row by row.

The mechanism you have just run is a simplified version of attention. Its weights are fixed by the computation you already did. The next step is to add trainable weights, which let the model learn from data and improve its performance on specific tasks.

This extended mechanism is self-attention, also called scaled dot-product attention. The source notes that it is the self-attention used in the original transformer architecture, in GPT models, and in most other popular large language models. The simplified version and the trainable version share the same underlying idea: attention weights are applied to inputs through a matrix product to produce context vectors for all tokens together.