Transformers and Self-Attention

The recurrent network reads a sentence one word at a time, carrying a hidden state forward. That design has two costs that gating only mitigates. The first is information loss: a dependency between two distant words must be copied through every hidden state in between, degrading at each step; even the LSTM's cell only slows the decay. The second is serial computation: the hidden state at position $t$ needs the hidden state at $t-1$, so the sequence cannot be processed in parallel — training time grows with sentence length.¹

Self-attention discards the recurrence. Instead of passing information down a chain, it lets every position look directly at every other position and pull in a weighted mixture of them, in one step. A dependency between word 1 and word 50 is a single comparison rather than a chain of 50 copies. And because the output at each position is computed independently of the others, the whole sequence is processed at once, so the computation parallelizes across a GPU. The transformer is a stack of layers built almost entirely from this one operation.

The two costs are really one geometric fact: in a recurrent network the number of steps a signal travels between two positions is their distance in the sequence, so a dependency between token $1$ and token $50$ passes through $49$ hidden states, and the training gradient that would tie them passes through $49$ multiplications. In self-attention that path length is $1$ for every pair — token $1$ and token $50$ are one dot product apart, no matter the gap. Shorter paths mean less information decay on the forward pass and less gradient decay on the backward pass, and they are why attention learns long-range structure that recurrence struggles with.

Attention as a soft lookup

Attention already appeared, in miniature, at the end of the recurrent chapter: the decoder of an encoder-decoder model got to look back at all the encoder's hidden states, weighting each by relevance. Self-attention applies that same idea inside a single sequence — a word attends to the other words of its own sentence rather than to a separate source text.²

Start with the simplest version. Each input is a vector ${\mathbf{x}}_i$, and we want an output ${\mathbf{y}}_i$ that summarizes how ${\mathbf{x}}_i$ relates to the words around it. Relatedness is measured by a dot product: ${\mathbf{x}}_i\cdot {\mathbf{x}}_j$ is large when the two vectors point the same way. Compute one such score for every earlier position, normalize the scores with a softmax into weights ${\alpha }_{ij}$ that sum to one, and read off the weighted average of the inputs:

This is a soft dictionary lookup: the query ${\mathbf{x}}_i$ is compared against a set of entries, the comparison scores become a probability distribution, and the answer is the distribution-weighted blend of the entries. Nothing is recurrent, and every ${\mathbf{y}}_i$ is a separate sum, so all of them can be computed simultaneously.³

Queries, keys, and values

Using the same vector ${\mathbf{x}}_i$ for all three jobs — the thing doing the comparing, the thing being compared against, and the thing being summed — is wasteful, because those are three genuinely different roles. Transformers give each role its own learned projection of the input.⁴

The score is now a dot product between the current word's query and each earlier word's key, and the output is the weighted sum of the corresponding values rather than of the raw inputs:

The three matrices ${\mathbf{W}}^Q,{\mathbf{W}}^K,{\mathbf{W}}^V$ are the only learned parameters of the attention layer; they are shared across all positions and trained end-to-end with the rest of the network. Query and key vectors have dimension $d_k$, values have dimension $d_v$, and the model as a whole works in dimension $d_{model}$; a single-head layer takes $d_k=d_v=d_{model}$, though the multi-head version below will pull them apart.⁴

Why the $\sqrt{d_k}$ scaling

A dot product of two $d_k$-dimensional vectors sums $d_k$ products. If the components are roughly independent with unit variance, the dot product has variance proportional to $d_k$, so as the dimension grows the scores spread out into large positive and negative values. Feed large values into a softmax and it saturates: almost all the probability collapses onto the single largest score, and the gradient with respect to the others goes to zero, which stalls learning.⁵

The fix is to divide the score by $\sqrt{d_k}$ before the softmax, rescaling the variance back to a constant regardless of dimension. This is scaled dot-product attention:

Intuitively: a softmax turns a gap between two scores into a gap between two probabilities, but the size of the probability gap depends on the scale of the scores. Left unscaled, that scale grows with dimension, so at large $d_k$ even an ordinary gap saturates the softmax. Dividing by $\sqrt{d_k}$ holds the scale fixed, so the same-sized gap always produces the same-sized probability gap.

To see why the constant matters, make the argument quantitative. Suppose two keys produce raw dot products with the query that differ by one standard deviation of the score distribution. Because that distribution has standard deviation $\sqrt{d_k}$, a one-sigma gap is a raw difference of $\sqrt{d_k}$: the winning logit is $\sqrt{d_k}$, the loser is $0$. Run those through the softmax and the winner's probability climbs sharply with dimension — at $d_k=4$ it wins with probability $0.88$, at $d_k=64$ with $0.9997$, at $d_k=1024$ it is $1.0000$ to four decimals. The distribution has collapsed onto one word, and the gradient that would nudge the others has vanished. Divide by $\sqrt{d_k}$ first and the same one-sigma gap becomes a difference of exactly $1$ at every dimension, so the winner's probability holds steady at $0.73$ regardless of $d_k$. That is the whole point of the scaling: it fixes the softmax's operating point so the layer learns the same way whether $d_k$ is $4$ or $1024$.

A numeric trace of one attention step

For example, take three tokens and work in $d_k=d_v=2$, computing the output ${\mathbf{y}}_3$ for the third position, which (being causal) attends to positions $1$, $2$, and $3$. Suppose the projections have already produced these query, key, and value vectors:

Step 1 — score. Dot the query against each key: ${\mathbf{q}}_3\cdot {\mathbf{k}}_1=1$, ${\mathbf{q}}_3\cdot {\mathbf{k}}_2=1$, ${\mathbf{q}}_3\cdot {\mathbf{k}}_3=2$.

Step 2 — scale. Divide by $\sqrt{d_k}=\sqrt{2}\approx 1.414$: $0.707,0.707,1.414$.

Step 3 — softmax. Exponentiate and normalize. The weights come out ${\alpha }_{3,1}=0.248$, ${\alpha }_{3,2}=0.248$, ${\alpha }_{3,3}=0.503$, summing to $1$. Position $3$ attends most to itself (its key matched the query best) and splits the rest evenly between positions $1$ and $2$.

Step 4 — weighted sum of values.

The output is a blend of the three value vectors, tilted toward ${\mathbf{v}}_3$. Two details. First, the scaling changed the answer: with the raw scores $1,1,2$ the softmax would give weights $0.212,0.212,0.576$, a sharper peak on position $3$; the $\sqrt{2}$ divisor softened the distribution, as intended. Second, nothing here touched positions $4$ or beyond; had this been an encoder (bidirectional) layer, the query would also have scored against future keys, but the causal mask forbids it.

The whole layer as one matrix product

Because every output is independent, the layer has an efficient batched form. Pack the $N$ input vectors as rows of a matrix $\mathbf{X}\in {\mathbb{R}}^{N\times d}$ and project them all at once:

The product $\mathbf{Q}{\mathbf{K}}^{\top }$ is an $N\times N$ matrix holding every query-key score at once; scale it, softmax each row, and multiply by $\mathbf{V}$. The entire self-attention computation for a sequence of $N$ tokens collapses to a single expression:

This is the equation the entire architecture is built around. Note its cost: forming $\mathbf{Q}{\mathbf{K}}^{\top }$ requires $N^2$ dot products, so attention is quadratic in the sequence length. That quadratic cost is why transformer inputs are usually capped at a page or a paragraph, and why more efficient attention variants are an active research area.⁶

Causal (masked) self-attention

For a language model the constraint is absolute: predicting the next word may use only the words before it. If position $i$ were allowed to attend to position $i+1$, the model would see the answer it is supposed to predict. So a language-model attention layer must be causal — each query attends only to keys at its own position or earlier, exactly the $j\le i$ condition written into the sums above.⁷

In the matrix form the full $\mathbf{Q}{\mathbf{K}}^{\top }$ scores every query against every key, including future ones. Those forbidden entries sit in the upper triangle of the matrix. To forbid them, set the upper triangle to $-\infty$ before the softmax; $e^{-\infty }=0$, so every future key receives exactly zero weight and contributes nothing to the output.

The same layer, without the mask, is bidirectional: every position attends to the whole sequence in both directions. Bidirectional self-attention is what an encoder uses, where the full input is available at once; causal self-attention is what a decoder and any left-to-right language model use. It is the identical operation with one triangular mask flipped on or off.⁷

With the attention operation in hand — a query scored against keys, softmaxed into weights, mixing values, optionally masked to be causal — we can assemble the full network. That is the job of the next part, which stacks attention into multiple heads, wraps it in the transformer block, restores word order, and arrives at the architectures modern NLP is built on. This continues in The Transformer Architecture.