Causal attention and multi-head attention

Earlier, 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 have already computed can be applied to all the inputs in a single operation, producing a full set of context vectors together.

The whole computation is one line. It takes the attention weight matrix named attn_weights, multiplies it by the inputs using the @ operator, stores the result in all_context_vecs, and prints it. Before reading further, run that line and look at the printed tensor, then predict what it should look like. The rest of this article uses that output to reason backward about the operation that produced it.

Before trusting the one-line version, trace what it does on a small case. Suppose there are three tokens, each represented by a two-number input vector, and the attention weights for each token are three numbers that sum to 1. The values below are invented for teaching purposes; they are not from a real model.

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 of the result 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.

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 all rows at once gives all context vectors together. The diagram below shows how each weights row produces exactly one output row.

You can predict the output shape before running the program. Let the attention weights be a $T \times T$ matrix, one row and one column per token, and let the inputs be a $T \times d$ matrix, one row per token with $d$ values each. The inner dimensions match, the shared $T$ is consumed by the product, and what remains is the outer pair.

In the 3 by 3 times 3 by 2 example, the middle 3 is the inner dimension and it cancels. The first matrix contributes its row count, 3, and the second contributes its column count, 2, so the result is 3 by 2. In symbols: $(3 \times 3)(3 \times 2) = (3 \times 2)$.

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. Its second row in the batched result should match that vector. Printing the earlier vector alongside the new tensor gives a direct comparison.

Compare values position by position, starting at the first value. The first position where the two vectors disagree tells you which token and which weight produced the error. In the faulty row, the first value agrees and the second does not, so the problem lies in the weights row for that token or in the product that formed its entries.

The pattern of disagreement also points to the cause. This check is a debugging technique, not just a confirmation. If a single row disagrees, the problem is in the batched path for that token, and you should 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.

Causal attention is the later stage that restricts which tokens each position can draw from. Each token may use itself and the tokens before it, but not tokens that come later in the sequence. In the attention weights matrix, this means the entries that would connect a token to a future position are excluded, so each row only blends the allowed inputs.

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. This extended mechanism is self-attention, also called 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 source sequence of attention variants is built step by step. The simplified attention mechanism comes first, followed by self-attention with trainable weights. The later stages add causal attention and multi-head attention. Multi-head attention runs several attention computations in parallel and combines their results into one context vector per token; this article does not walk through its internal steps, but it builds directly on the single matrix product described above.