The Mathematics Behind Self-Attention

Every sequence model faces the same problem: the meaning of a word depends on other words, and the model must decide which ones. In “the keys on the piano”, the word “keys” means something because of “piano”, five positions away. A recurrent network answers this by passing a hidden state along the sequence, which forces information from “piano” to survive a relay race of matrix multiplications before it reaches “keys”, and the gradient mathematics of that relay, as the CNNs and RNNs article showed, is an exponential that erases long-range influence.

Self-attention abolishes the relay. Every position looks at every other position directly, in one step, and decides for itself how much each one matters. The entire mechanism is four small ideas composed: similarity measured by dot products, three learned projections that split a word’s roles, a scaling factor with a precise variance justification, and a softmax that turns scores into a weighted average. This article derives each piece and shows why the composition, and not any single part, is the innovation that made transformers possible.

The goal, stated as mathematics

Start with what the mechanism must produce. The input is a sequence of n tokens, each represented by an embedding vector of dimension d, stacked into a matrix:

Each row is one token’s vector, and at this point every row is context-free: the vector for “keys” is the same whether the sentence is about pianos or prisons. The output should be a new matrix of the same shape in which each row has been updated to reflect its context:

The most natural way to build a context-aware vector is as a weighted average of the vectors around it:

where the weight $w_{ij}$ says how much token j matters to token i. If “piano” matters greatly to “keys”, its weight should be large, and the updated “keys” vector will lean toward “piano”, disambiguating it. The whole design question is now: where do the weights come from? They cannot be fixed parameters, because which words matter to which changes with every sentence. The weights must be computed from the content of the sequence itself, which is what “self” in self-attention means.

Dot products measure relevance

The simplest measure of how much two vectors have in common is their dot product:

The dot product is large and positive when the vectors point in similar directions, near zero when they are unrelated, and negative when they oppose, which is exactly the behaviour a relevance score needs. Embeddings are trained so that related words occupy related directions, so a dot product between embeddings is a learned similarity measure available for free.

A first draft of attention is therefore: score every pair of tokens with a dot product, and weight each token’s update by the scores. But this draft has a flaw worth catching before fixing, because the fix is the most important design decision in the mechanism. Under raw dot products, the score between “keys” and “piano” is forced to equal the score between “piano” and “keys”, since $x_i \cdot x_j = x_j \cdot x_i$, and the relevance relationship is genuinely not symmetric: “piano” may be essential for interpreting “keys” while “keys” contributes little to interpreting “piano”. Worse, a token’s relevance to others and the content it offers others are tangled in a single vector.

Queries, keys, and values: three roles, three projections

The fix is to give each token three different vectors for its three different jobs, produced by three learned linear maps:

with weight matrices $W_Q, W_K \in \mathbb{R}^{d \times d_k}$ and $W_V \in \mathbb{R}^{d \times d_v}$ learned during training. The names describe the jobs. The query $q_i$ is what token i is looking for, its question to the rest of the sequence. The key $k_j$ is what token j advertises about itself, the label by which others find it. The value $v_j$ is what token j actually hands over once selected, the content that flows into other tokens’ updates.

The relevance score becomes a query-key dot product:

and the asymmetry problem dissolves: $q_i \cdot k_j$ and $q_j \cdot k_i$ involve different matrices and are free to differ, so “keys attends strongly to piano” no longer forces the reverse. Separating value from key resolves the other entanglement, letting a token be easy to find for one reason and contribute content of a different kind. The three projections are the only learned parameters in the whole mechanism, which is worth pausing on: attention itself has no weights, only the maps that prepare its inputs.

The retrieval analogy is exact enough to be the right mental model. A search engine matches your query against page titles, the keys, and returns page contents, the values. Self-attention runs that search n times simultaneously, every token querying every other, with the match scores deciding how much of each value to blend.

The scaling factor, derived rather than asserted

All the pairwise scores assemble into one matrix product:

Row i holds token i’s scores against every token. Before these become weights, they pass through a softmax, and here a subtle numerical problem arises that the famous $\sqrt{d_k}$ exists to solve.

Consider the size of a single score. Suppose, as is roughly true at initialisation, the components of $q_i$ and $k_j$ are independent with mean 0 and variance 1. The dot product is a sum of $d_k$ such component products, and the variance of a sum of independent terms adds:

So the typical magnitude of a score grows like $\sqrt{d_k}$, and with $d_k = 64$ the scores are routinely 8 times larger than the unit-scale inputs. Feeding large scores into a softmax is destructive: the exponential amplifies differences, so a softmax over large inputs saturates, placing essentially all its mass on the single largest score. Two things break. The attention becomes prematurely hard, each token attending to exactly one other, and the gradient of a saturated softmax is near zero everywhere, so the mechanism stops learning, the same vanishing-gradient arithmetic that plagued RNNs sneaking back in through a different door.

Dividing by $\sqrt{d_k}$ restores unit variance:

which keeps the softmax in its responsive regime regardless of how large the key dimension grows. The scaling is not a stylistic constant; it is a variance correction, and the derivation shows exactly why it is a square root and why it involves $d_k$ and nothing else.

Softmax: scores into a probability distribution

The scaled scores in each row are converted into weights by the softmax:

Softmax does three jobs at once. It makes every weight positive. It makes each row sum to exactly 1, so the weights form a probability distribution over the sequence and the subsequent combination is a genuine weighted average, keeping the output on the same scale as the values regardless of sequence length. And it is differentiable, unlike a hard argmax selection, so the relevance decisions themselves receive gradients and improve with training. The exponential’s amplification is a feature at the right scale: it sharpens moderate differences in score into decisive differences in weight, letting the model commit to the tokens that matter while still passing gradient to the ones that nearly mattered.

The full mechanism in one line

Composing the pieces gives the equation that defines the transformer era:

Read it inside out, because each symbol is now earned. $QK^{\top}$ is every token questioning every other, $n^2$ dot products in one matrix multiply. The division by $\sqrt{d_k}$ keeps the scores at unit variance so the softmax stays trainable. The softmax turns each row into a probability distribution over the sequence. And the final multiplication by V computes, for every token simultaneously, the weighted average of value vectors that becomes its context-aware representation:

which is precisely the goal we wrote down at the start, with the weights now computed from content, asymmetric, scaled, and normalised. The “keys” row places heavy weight on “piano”, its updated vector absorbs piano-ness, and the ambiguity is resolved in a single layer with no relay race.

What the matrix form buys: parallelism

The RNN comparison deserves one more pass, because the practical revolution is computational as much as representational. A recurrent network’s hidden state at position t depends on the state at t-1, so the sequence must be processed in order, one step at a time, and no hardware parallelism can help across the length of the sequence. The attention equation contains no such dependency: $QK^{\top}$ and the softmax and the final product are dense matrix operations over the whole sequence at once, which is exactly the workload GPUs are built for. Training on long sequences went from inherently serial to embarrassingly parallel, and that, as much as the gradient story, is why transformers displaced recurrence.

The price is written in the same expression. The score matrix has $n^2$ entries, so memory and compute grow quadratically with sequence length:

Doubling the context multiplies the attention cost by four. This quadratic term is the central constraint of the LLM era, the reason long-context models are an engineering achievement rather than a parameter change, and the motivation for the efficient-attention variants covered later in the roadmap.

Two properties the equation forces on the architecture

Two consequences follow directly from the mathematics, and both shape everything downstream.

First, permutation blindness. Nothing in $QK^{\top}$ knows where a token sits: shuffle the input rows and the output rows shuffle identically, with the same values. Attention treats the sequence as a bag of vectors, so word order, which language depends on, is invisible to it. The cure is to inject position into the embeddings themselves, adding a positional encoding to each row of X before the projections, so that “piano” at position 5 and “piano” at position 2 present different vectors and order re-enters through the content channel. The design of those encodings is its own topic in the roadmap.

Second, seeing the future. In a model trained to predict the next token, position i must not attend to positions after i, or the training task collapses into copying. The fix is the causal mask: before the softmax, set every score $s_{ij}$ with j > i to negative infinity, so the exponential sends those weights to exactly zero:

Each token then attends only to its past, the weighted average runs over positions 1 through i, and the mechanism becomes the autoregressive engine that generates text one token at a time.

Where this leads

One attention layer computes one pattern of relevance, one set of weighted averages under one set of projections. A sentence needs several patterns at once: syntactic relationships, coreference, semantic similarity, positional neighbours. Running several attention operations in parallel with separate learned projections, then concatenating the results, is multi-head attention, and stacking such layers with feed-forward networks between them is the transformer. The Deep Dive into Attention article takes up that construction, along with the geometry of what the heads learn.

What this article should leave behind is the anatomy of the single mechanism. A weighted average whose weights are computed from content. Dot products as the similarity measure, split across queries and keys so relevance can be asymmetric, with values separating what a token offers from how it is found. A $\sqrt{d_k}$ that is a variance correction, not a convention. A softmax that makes the weights a trainable probability distribution. And a matrix form whose parallelism built the modern era and whose $n^2$ bills it. Every headline about context windows, attention variants, and inference cost is downstream of this one equation, and you can now read it term by term.

See you soon.