Computing Attention Weights for All Input Tokens

Self-attention starts with a simple question: how strongly does each input token relate to every other input token? If a sequence contains six input vectors, every vector must be compared against every other vector. That produces a grid with six rows and six columns, so there are 36 comparisons in total. This article follows that grid through three stages: raw dot-product scores, a matrix-multiplication shortcut that produces the same grid, and a normalization step that turns each row into proportions called attention weights.

The first version of the computation allocates an empty six by six tensor. Two nested loops visit every pair of input vectors. The outer loop index i selects a row, the inner loop index j selects a column, and each cell receives the dot product of the two vectors that index it. The dot product of vector i and vector j is the score that the cell at row i, column j stores.

Read the diagram as a map from loop position to memory of the grid. Cells on the diagonal, such as scores[0, 0] and scores[1, 1], pair a vector with itself. Off-diagonal cells, such as scores[0, 1] and scores[1, 0], compare two different tokens. Because the loop visits all pairs, the finished grid contains 36 entries, one per cell.

The loop version is correct, but for loops are generally slow. The same grid can be produced with a single matrix multiplication: 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 shape check explains why the result is square. The input has T rows and d columns, and the transpose has d rows and T columns. The shared inner dimension d is summed away by the multiplication, leaving a T by T matrix. Each entry is indexed by a pair of tokens, which is exactly the token-to-token grid the loops produced.

The second stage normalizes each row of the score grid so that its values behave like proportions. The normalized result is the attention weight tensor. Attention weights are the softmax-normalized scores, so each row of weights describes how one token distributes its attention across all tokens. In the course pipeline, these weights are the output of stage two, and the context vectors are formed from them in stage three.

The rule the weights must satisfy is that each row sums to 1 across its columns. In symbols, for a row i of the weight tensor, the entries satisfy $\sum_{j} w_{ij} = 1$, where $w_{ij}$ is the weight at row i and column j.

The before and after values come from the same course output. Before normalization, both cells hold 0.9544 because the dot product is symmetric. After normalization, the cell at row 0, column 1 holds 0.2006 and the cell at row 1, column 0 holds 0.1385. Each row was normalized by a different total, so the mirror symmetry of the scores is lost in the weights. Symmetry is a property of the raw scores, not of the weights.

The dim argument decides which axis softmax runs across. With dim=-1, softmax works along the last dimension. For a two-dimensional score tensor with rows and columns, that means each row is normalized across its columns, and the values inside one row sum to 1. Changing dim changes the direction of normalization, which changes what each row means. The choice of dim is therefore part of whether the result is correct, not a cosmetic setting.

The computation has checkpoints. Stage one produces raw dot-product scores with shape T x T. Stage two applies softmax along dim=-1 to produce weights whose rows sum to 1. Stage three forms context vectors from those weights. Knowing the expected output of each stage tells you where a result went wrong.