Attention & Transformers

A recurrent network reads a sequence one step at a time, squeezing the entire past into a single hidden state $h_t$. That bottleneck is the defect: information from position $1$ reaches position $n$ only after $n$ sequential hops, and the gradient connecting them decays along the way: the very problem LSTMs and GRUs only partially patch. Attention removes the hops entirely. Instead of routing information through a chain, every position reads directly from every other position in one step, weighting what it reads by learned, content-dependent relevance.¹

Queries, keys, and values

The vocabulary is borrowed from retrieval. A query is what a position is looking for; a key advertises what each position offers; a value is the content actually returned. All three are linear projections of the inputs, so the mechanism is fully learned: three weight matrices determine what each token searches for, what it exposes for matching, and what content it returns.

Compatibility is measured by a dot product: $q\cdot k_j$ is large when the query and key point the same way. Normalizing the scores with a softmax turns them into a probability distribution over positions, and the output is the expectation of the values under that distribution,

Scaled dot-product attention

Stack the $n$ queries into the rows of $Q$, the $m$ keys into the rows of $K$, and the $m$ values into the rows of $V$. Keeping the query count $n$ and the key/value count $m$ separate matters: in self-attention $n=m$, but in cross-attention a decoder of length $n$ reads an encoder of length $m\ne n$. Queries and keys must share a width $d_k$ so their dot product is defined; values carry an independent width $d_v$. The shapes chain as follows.

Every query-key score is a single matrix product $Q{K}^{\top }\in {\mathbb{R}}^{n\times m}$: entry $(i,j)$ is the dot product $q_i\cdot k_j$ of query $i$ with key $j$. Divide by $\sqrt{d_k}$, run a softmax along each row so that row $i$ becomes a probability distribution over the $m$ keys, and right-multiply by $V$. The inner dimension $m$ cancels in $AV$, so an $n\times m$ weight matrix times an $m\times d_v$ value matrix returns an $n\times d_v$ output — one $d_v$-vector per query, exactly the input query count. The whole operation collapses to one expression.

The only unexplained piece is the $1/\sqrt{d_k}$ factor. It is essential: without it the softmax saturates and gradients vanish.²

Why the $\sqrt{d_k}$ scaling

Treat the entries of a query $q$ and a key $k$ as independent random variables with mean $0$ and variance $1$. The dot product is a sum of $d_k$ such products, and we can compute the variance of that sum exactly.

So the raw scores have standard deviation $\sqrt{d_k}$, which grows with the model dimension. Feed scores of magnitude $\sim \sqrt{d_k}$ into a softmax and, for $d_k$ in the hundreds, the largest entry dwarfs the rest: the softmax saturates onto one position, approaching a one-hot vector. In that regime its Jacobian is nearly zero, and almost no gradient flows back to the queries and keys. Dividing by $\sqrt{d_k}$ rescales the scores to unit variance, $\mathrm{Var}\left((q\cdot k)/\sqrt{d_k}\right)=1$, keeping the softmax in its sensitive, high-gradient region regardless of dimension.

A worked example

Push numbers through the formula. Take three tokens with $d_k=d_v=2$, and, to keep the arithmetic legible, let the queries and keys coincide, so this is self-attention:

Step 1 — raw scores $S=Q{K}^{\top }$. Each entry is $q_i\cdot k_j$. Row $1$ is query $(1,0)$ dotted against the three keys: $(1,0)\cdot (1,0)=1$, $(1,0)\cdot (0,1)=0$, $(1,0)\cdot (1,1)=1$. Filling in all nine,

Step 2 — scale by $\sqrt{d_k}=\sqrt{2}\approx 1.414$. Every score shrinks by the same factor, pulling the logits toward zero before the softmax:

Step 3 — row-wise softmax. Work row $1$ explicitly. Exponentiate: $e^{0.707}\approx 2.028$, $e^0=1$, $e^{0.707}\approx 2.028$, summing to $5.056$. Dividing gives $2.028/5.056\approx 0.401$ for the first and third keys and $1/5.056\approx 0.198$ for the second. Doing the same to every row,

and each row sums to $1$, as a distribution over the three keys must. Notice the third token — whose query $(1,1)$ aligns with every key — spreads its weight most evenly but still favours itself (key $3$, weight $0.503$), because $(1,1)\cdot (1,1)=2$ is the largest score in that row.

Step 4 — blend the values, $O=AV$. Row $1$ of the output is $0.401v_1+0.198v_2+0.401v_3=0.401(1,0)+0.198(0,1)+0.401(0.5,0.5)=(0.602,0.398)$. All three outputs:

Each output row is a point inside the triangle spanned by the three value vectors — a convex combination, never an extrapolation beyond them. That is the geometric signature of attention: the output can only ever be a weighted average of the values it was given.

Self-attention

When the queries, keys, and values are all linear projections of the same sequence $X\in {\mathbb{R}}^{n\times d}$, attention becomes self-attention: each position attends to every other position in its own input, including itself.

Its central object is the attention matrix

a row-stochastic matrix (${\sum }_j{A}_{ij}=1$) that is a soft, content-based, all-pairs routing table: row $i$ says where token $i$ gathers its information. Unlike a convolution's fixed local stencil or a recurrence's strictly leftward chain, every entry is non-zero and learned from content: the word it can route directly to the noun it refers to, however far back.

Permutation equivariance

Self-attention has no notion of order. Permuting the rows of $X$ permutes the rows of the output identically and changes nothing else; the mechanism treats its input as a set, not a sequence.

Equivariance renders position invisible: dog bites man and man bites dog produce the same set of outputs, merely reordered. Order must therefore be injected explicitly, the job of positional encoding below.

Multi-head attention

A single attention function forces every position to mix all of its relationships into one set of weights. Multi-head attention runs $h$ attention functions in parallel, each with its own projections, so different heads can specialize: one tracking syntactic agreement, another long-range coreference, another local adjacency.³

The shapes are worth tracing once. Fix $d=d_{\text{model}}$ and $h$ heads, and set the per-head width $d_k=d_v=d/h$ (so $h=8$ heads of a $d=512$ model each work in $64$ dimensions). For head $i$ the projections are $W_Q^{(i)},W_K^{(i)},W_V^{(i)}\in {\mathbb{R}}^{d\times d/h}$, so from an input $X\in {\mathbb{R}}^{n\times d}$,

each head returns an $n\times (d/h)$ block. Concatenating the $h$ blocks along the width rebuilds a full $n\times d$ matrix, which $W_O\in {\mathbb{R}}^{d\times d}$ mixes back into the model width:

Because $h\cdot (d/h)=d$, the total projection and attention cost is the same as one full-width head; heads split the width rather than add to it. The expressive gain is that the model holds $h$ distinct routing tables at once instead of averaging them into one — one head can lock onto the previous token while another tracks long-range coreference, and $W_O$ learns how to combine their outputs.

Positional encoding

Because self-attention is permutation-equivariant, the model cannot tell position $1$ from position $5$ unless we tell it. The fix is to add a position-dependent vector $\mathrm{PE}(p)$ to each input embedding before the first layer, so that order rides along in the representation itself.⁴

The geometric spread of frequencies is the key property: low-index dimensions oscillate quickly to distinguish neighbors, high-index dimensions oscillate slowly to distinguish distant regions. The encoding also admits relative shifts, since $\mathrm{PE}(p+\Delta )$ is a fixed linear function of $\mathrm{PE}(p)$, so the model can learn to attend by offset, not just absolute index.

Learned positional embeddings — a trainable lookup table indexed by position — are the common alternative; they fit the data more tightly but, unlike the sinusoids, do not extrapolate to sequences longer than those seen in training.

The Transformer block

Self-attention mixes information across positions; it does no nonlinear processing within a position. The Transformer block pairs the two: a multi-head self-attention sublayer followed by a position-wise feed-forward network, each sublayer wrapped in a residual connection and layer normalization.⁵

Each sublayer is wrapped identically, $\mathrm{out}=\mathrm{LayerNorm}(x+\mathrm{Sublayer}(x))$. The residual skip lets gradients reach early layers undiminished — the same mechanism that powers residual networks — and LayerNorm holds each position's activations at a stable scale across depth.

Encoder and decoder

The original Transformer is an encoder-decoder. The encoder stacks identical blocks with unrestricted self-attention — every position sees every other. The decoder adds two twists that make it a valid autoregressive generator.

• Masked self-attention. To generate token $t$ using only tokens $<t$, the decoder sets $S_{ij}=-\infty$ for $j>i$ before the softmax (the masking branch of the algorithm above), forcing $A_{ij}=0$ on the future. Generation stays causal.
• Cross-attention. A second attention sublayer takes its queries from the decoder but its keys and values from the encoder's output, letting each generated token read the entire source sequence — the channel through which a translation model reads the sentence it is translating.

Cross-attention is the setting where the $n\ne m$ shapes matter. In an encoder of length $m$ and a decoder generating length $n$, the decoder supplies $Q\in {\mathbb{R}}^{n\times d_k}$ while the encoder supplies $K,V\in {\mathbb{R}}^{m\times d_k}$, so the score matrix is a rectangular $n\times m$ lookup — every output position against every source position. Self-attention is the special case $Q=K=V$ from one sequence, giving the square $n\times n$ matrix; cross-attention keeps queries and keys from different sequences.

The causal mask

The mask is a single additive matrix $M\in {\mathbb{R}}^{n\times n}$ applied to the scaled scores before the softmax: $M_{ij}=0$ for $j\le i$ and $M_{ij}=-\infty$ for $j>i$. Adding $-\infty$ to a logit sends $e^{-\infty }=0$, so the softmax assigns those future positions weight exactly $0$ — the reason for the $-\infty$ rather than merely a large negative number. For four tokens,

leaving a lower-triangular weight matrix: token $i$ attends only to tokens $1\dots i$. The mask costs nothing at training time (one addition), yet it is what lets the decoder train on a whole target sequence in parallel while still predicting each token from its prefix alone.

Recurrence vs. convolution vs. self-attention

Why did attention displace both recurrence and convolution for sequences? Compare the three on the costs that matter: how far information must travel to connect two positions (maximum path length, which bounds how easily long-range gradients flow), how much of the work is inherently sequential, and the per-layer compute.

Two columns decide it. Self-attention's maximum path length is $O(1)$: any two positions interact directly, so the gradient between them never has to pass through a long chain. Its sequential depth is $O(1)$, so the whole layer parallelizes across positions on a GPU, unlike the strictly stepwise recurrence. The price is the $O(n^{2}d)$ term: the $n\times n$ attention matrix is quadratic in sequence length, cheap when $n<d$ (the usual regime) but the dominant cost for long sequences, the bottleneck that later efficient-attention variants attack.

Where attention came from

Attention postdates Goodfellow's 2016 text, so the primary sources carry the detail Goodfellow cannot.

• Where attention came from. The dot-product form here is the streamlined version. The idea started as additive attention (Bahdanau et al., ICLR 2015), which scored a query against each key with a small MLP, $\mathrm{score}(q,k)=v^{\top }\tanh (W_{q}q+W_{k}k)$, to let a translation decoder read every encoder state instead of one fixed vector. Vaswani et al., Attention Is All You Need (NeurIPS 2017), replaced the MLP with the cheap scaled dot product, added the $1/\sqrt{d_k}$ correction derived above, and removed recurrence entirely.
• Reducing the $O(n^2)$ cost. The quadratic score matrix is the target of a whole line of work. Two directions matter. Sparse and linear approximations — Longformer (Beltagy et al., 2020) with sliding-window attention, and the Performer/linear-attention family — reduce the asymptotic cost to $O(n\log n)$ or $O(n)$ by never forming the full matrix. Exact but IO-aware kernels — FlashAttention (Dao et al., NeurIPS 2022) — keep the exact softmax but tile the computation so the $n\times n$ matrix never leaves fast on-chip memory, cutting the memory footprint from $O(n^2)$ to $O(n)$ with no change to the result. The second approach is now the default in production Transformers.
• Relative and rotary position. The additive sinusoidal encoding is only the first of several schemes. Shaw et al. (2018) added learned relative-position biases inside the scores; RoPE (Su et al., 2021) rotates the query and key vectors by an angle proportional to position, so a dot product depends only on the relative offset — the standard choice in modern decoder-only LLMs.

The full architecture built from this mechanism — the stacked encoder–decoder, causal masking, and the parameter/compute accounting — is the subject of the next lesson.

Takeaways

• Attention is a differentiable soft lookup: a query scores a set of keys, a softmax turns the scores into weights, and the output is the weighted blend of the values, $\mathrm{Attn}(Q,K,V)=\mathrm{softmax}(Q{K}^{\top }/\sqrt{d_k})V$.
• The $\sqrt{d_k}$ scaling cancels the dot product's variance growth ($\mathrm{Var}(q\cdot k)=d_k$), keeping the softmax out of its saturated, zero-gradient regime.
• Self-attention ($Q,K,V$ from one sequence) yields an all-pairs routing matrix $A$; being permutation-equivariant, it needs positional encoding — sinusoidal or learned — to know order.
• Multi-head attention runs $h$ parallel attentions in $d/h$-dimensional subspaces, capturing several relationship patterns at once for the cost of one.
• The Transformer block = multi-head self-attention + position-wise FFN, each in a residual + LayerNorm wrapper; decoders add a causal mask and cross-attention.
• Against recurrence and convolution, self-attention wins on path length ($O(1)$) and parallelism ($O(1)$ sequential ops), paying an $O(n^{2}d)$ per-layer cost — the trade that made it the dominant sequence architecture.