Computing Context Vectors with Attention Weights

Attention weights tell us how strongly each input token should influence a chosen query token. On their own, though, they are only numbers that sum to 1.0. The step that turns them into something a model can use is the context vector: a single vector that blends the whole input sequence according to those weights. This article walks through the full pipeline, from raw comparisons between input vectors, to normalized weights, to the blended context vector, and it checks at each stage what the output should look like.

Self-attention begins by measuring how strongly each input token relates to every other input token. For a sequence of six input vectors, that means comparing every vector against every other vector, which produces a grid with six rows and six columns. The first version of this computation fills that grid one cell at a time with nested loops. The same grid can later be produced with a single matrix multiplication, and then the raw grid is turned into weights whose rows behave like proportions.

Read the stages as a pipeline with a checkpoint at each stage. Each stage produces a tensor, and each tensor has an expected shape and expected row behavior. Knowing what each stage should produce is what lets you find the exact place where a result goes wrong.

The loop-based computation starts by allocating an empty six by six tensor. Two nested loops visit every pair of input vectors, and each cell receives the dot product of the two vectors that index it. The outer loop picks a row i and the inner loop picks a column j, so the cell at position (i, j) holds the dot product of input i with input j. Every cell gets exactly one dot product, so the grid has thirty six entries.

Compare that prediction with the source output. The scores are symmetric across the diagonal, so the cell at row 0, column 1 and the cell at row 1, column 0 both hold 0.9544. The diagonal, where each vector is paired with itself, is worth noticing as well, because it is the one place where a vector meets its own copy.

The loop version is correct, but for loops are generally slow. The same grid of scores can be produced with a single matrix multiplication, written inputs @ inputs.T. The input matrix multiplied by its own transpose contains, in each cell, the dot product of the corresponding row vectors. The transpose turns the columns of the second copy into rows, so the product lines up every vector against every other vector in one operation.

The source notes that the matrix version produces the same results as the loop version, and the values can be checked visually. This is the verification habit to carry forward: when you change how a computation is written, compare the new output against the old one, element by element or as a whole grid.

The second stage normalizes the score grid so that the values in each row behave like proportions. The result of this step is called the attention weight tensor, and attention weights are the normalized version of the attention scores. Each raw score is passed through the softmax function, which makes each value positive and divides it by the total of its row, so that the row sums to 1.

In this expression s_{ij} is the raw score at row i and column j, and a_{ij} is the resulting attention weight. The denominator sums over the columns k of the same row i, which is exactly why each row ends up summing to 1.

The dim argument decides which axis the normalization runs across. In PyTorch, dim=-1 tells softmax to work along the last dimension. For a two-dimensional score tensor shaped as rows and columns, this means each row is normalized across its columns. Changing dim changes which direction is normalized, and that changes what the rows mean. This is why the choice of dim is part of the correctness of the result, not a cosmetic setting.

The source output shows why the normalized weights are not symmetric. In the score grid, the cells at row 0, column 1 and row 1, column 0 both hold 0.9544. In the weight grid, the same two positions hold 0.2006 and 0.1385, which differ. Each row is normalized by its own total, so each row changes by a different amount, and the mirror symmetry of the scores is lost in the weights. Scores are symmetric, but the normalized weights generally are not.

After normalization, the expected property is simple: every row of the weight tensor should sum to 1. Checking this is the fastest way to catch a normalization step that was applied along the wrong axis. The check can be done for a single row, by adding the values of one row, or for the whole tensor, by summing along the last dimension and confirming that every total is 1.

A context vector is the weighted sum of all input vectors, where each input is multiplied by its attention weight for one chosen query. Read this definition as a recipe with three ingredients. You need the full list of input vectors, the attention weight attached to each of them for one chosen query, and a running total that the weighted vectors are added into. Nothing in the recipe depends on the order of the inputs except that each weight must stay paired with its own input vector.

In this notation z^{2} is the context vector for the second token, x^{i} is the i-th input vector, and a^{2}_{i} is the attention weight that the second token assigns to input i. The sum runs over all T input tokens, so every input contributes to the result, even the ones that receive a very small weight.

The implementation follows the formula directly. It starts from a zero vector with the same shape as the query, then walks through the inputs one at a time. On each pass it scales the current input vector by its matching weight and adds the result into the running vector. Because the loop uses enumerate, the index i selects the right weight for the input x_i at the same position.

The zero start matters. The running vector has to begin empty so that the first addition does not carry leftover values from anywhere else. After the loop finishes, the running vector holds the complete sum, and printing it shows the context vector. The source's printed output for its query is three values, which the trace above is designed to let you reproduce by reasoning about the pass-by-pass totals.

Errors in this computation can be located by checking the query index, the normalization of the weights, and the pairing of each weight with its input. The following mistakes are the most common ways the pipeline fails.

Computing one context vector answers a single question: how should the second token's representation be blended from the whole sequence? A model needs an answer for every token. The source handles this by repeating the same three steps for each query rather than inventing a new method. Attention weights are first computed for every input token as a query, producing a full table with one row per query. Each row of that table is then used in the same weighted sum to produce that query's context vector.

The source notes that the change needed to compute all context vectors, instead of only one, is a modification to the code rather than to the underlying steps. The three steps remain the same: take the query's row of weights, scale each input by its matching weight, and add the results into a running total. A matrix formulation organizes these repeated passes into one pass, but each output row is still a weighted sum of all input vectors for its query.