Attention and the Transformer

The transformer (Vaswani et al., 2017, “Attention Is All You Need”) replaced recurrence with attention. It underlies BERT, GPT, Claude, Gemini, Llama, vision transformers and more.

The idea of attention

To understand a word, look at the other words that matter for it. In “The animal didn’t cross the street because it was too tired”, the representation of “it” should draw mostly from “animal”.

Each token produces three vectors from its embedding $\mathbf{x}$ via learned matrices:

• Query $\mathbf{q} = W_Q\mathbf{x}$: what this token is looking for.
• Key $\mathbf{k} = W_K\mathbf{x}$: what this token offers to be matched against.
• Value $\mathbf{v} = W_V\mathbf{x}$: the information it passes on if attended to.

Scaled dot-product attention

For all tokens at once (rows of $Q$, $K$, $V$):

1. $QK^\top$: an $n \times n$ matrix of scores, how well each query matches each key.
2. Divide by $\sqrt{d_k}$ so that the scores don’t grow with dimension and saturate the softmax.
3. Softmax each row into weights that sum to 1.
4. Multiply by $V$: each output is a weighted average of the values.

flowchart LR
  X[Token embeddings] --> Q[Q = XW_Q] & K[K = XW_K] & V[V = XW_V]
  Q & K --> S["scores = QKᵀ/√dₖ"] --> M["(mask)"] --> SM[softmax] --> O["weights × V"]
  V --> O

Every token can directly use information from every other token, regardless of distance, and the whole computation is matrix multiplications that parallelise across the sequence. Training no longer has to step through time like an RNN.

Cost: the score matrix is $n \times n$, so time and memory are $O(n^2)$ in sequence length. Long-context models use tricks like FlashAttention (memory-efficient exact attention), sliding windows and sparse attention.

Multi-head attention

One attention pattern can’t capture everything (syntax, coreference, nearby words). Multi-head attention runs $h$ attentions in parallel with smaller dimensions ($d_k = d_\text{model}/h$), concatenates the results and mixes them:

Positional information

Attention by itself ignores order: shuffle the tokens and the outputs shuffle the same way. Position must be injected:

• Sinusoidal encodings (original paper): fixed sin/cos waves added to embeddings.
• Learned position embeddings (BERT, GPT-2).
• Rotary embeddings (RoPE): rotate queries and keys by position, so attention depends on relative positions. Used by most modern LLMs.

The transformer block

flowchart TB
  IN[x] --> LN1[LayerNorm] --> MHA[Multi-head self-attention] --> ADD1(("+"))
  IN --> ADD1
  ADD1 --> LN2[LayerNorm] --> FFN["Feed-forward:<br/>Linear → GELU → Linear"] --> ADD2(("+"))
  ADD1 --> ADD2
  ADD2 --> OUT[to next block]

Each block has self-attention (tokens exchange information) and a position-wise feed-forward network (each token is processed independently, typically 4× wider), each wrapped with a residual connection and layer normalisation (pre-norm, as drawn, is standard now). Models stack dozens of these blocks.

Three families

The causal mask sets future scores to $-\infty$ before the softmax, so position $t$ only sees positions $\le t$. Cross-attention takes queries from the decoder and keys and values from the encoder.