Simple Self-Attention Without Trainable Weights

This article starts with code rather than theory. SelfAttention_v2 is an improved implementation of self-attention in PyTorch. Instead of managing its three projection matrices by hand, it builds them from PyTorch's Linear layers. The approach is to read the class once, predict what a call to it returns, run it, and then use the explanation to make sense of the printed result.

The constructor takes d_in and d_out, which are the input and output dimensions, and an optional qkv_bias flag that defaults to False. From these arguments it creates three Linear layers, one each for queries, keys, and values. Each layer maps an input of size d_in to an output of size d_out. Because all three layers share the same dimensions, the queries, keys, and values all have width d_out.

The forward method runs in a fixed order. It first projects the input x three times, producing keys, queries, and values. It then computes raw attention scores by multiplying the queries with the transpose of the keys. Those scores are divided by the square root of the key dimension, which is read from keys.shape[-1], and a softmax is applied along the last dimension. The result is the attention weights. Finally, the attention weights are multiplied by the values, and the product is the context vector that the method returns.

The raw scores are produced by multiplying queries with the transpose of keys. Each query row is compared against each key column, so the result is a square matrix. If the input has T tokens, the score matrix is T by T. The entry in row i and column j measures how strongly token i's query matches token j's key. Each row therefore describes one token's view of every token in the input, including itself.

Here $d_k$ is the key dimension, read from keys.shape[-1]. The division by the square root of $d_k$ scales the raw scores before the softmax. The softmax is taken along the last dimension, so each row of weights is normalized on its own and each row describes how one token distributes its attention across all tokens.

The final step multiplies the attention weights by the values. Each row of the attention weights picks out a mixture of the value vectors. The row for token i gives one weight to every value vector, and the context vector for token i is the weighted sum of those value vectors. Value vectors that receive larger weights contribute more to that token's context vector. One context vector is produced per input token.

The earlier version handled its projections by hand, meaning the weight matrices were managed directly in the class. The Linear-layer version uses nn.Linear modules instead. Each module holds its own weight matrix that maps d_in inputs to d_out outputs, and it can optionally hold a bias term controlled by qkv_bias. The class no longer tracks these matrices itself. It creates three modules and calls them on x, and the modules handle the stored weights.