Theory

Reading time: ~25 minutes (~1500 words) · Prerequisites: Lesson 20 (Tokenisation & Embeddings), basic linear algebra (matrix multiplication, dot products)

1. The Transformer at a Glance

The transformer, introduced by Vaswani et al. in Attention Is All You Need (2017), is the architectural foundation of every major language model today. It replaced recurrent architectures with a mechanism called self-attention that allows every token to relate to every other token in a single operation — enabling massive parallelism during training and capturing long-range dependencies that RNNs struggled with.

A transformer-based language model processes input through three stages: (1) embed tokens and add positional information, (2) pass through a deep stack of identical transformer blocks, and (3) project the final representations back to vocabulary size and sample the next token. This section examines each component in detail.

2. Scaled Dot-Product Attention

Self-attention is the core operation. Given an input sequence of n tokens, each represented as a d_model-dimensional vector, we compute three matrices by multiplying the input by learned weight matrices:

Q = X @ W_Q    # Queries:  [n x d_k]
K = X @ W_K    # Keys:     [n x d_k]
V = X @ W_V    # Values:   [n x d_v]

The attention output is then:

Attention(Q, K, V) = softmax( Q @ K^T / sqrt(d_k) ) @ V

The dot product Q @ K^T produces an [n x n] matrix of raw attention scores. Each row represents how much one token should attend to every other token. The scaling factor 1/sqrt(d_k) is essential: without it, for large d_k, the dot products grow in magnitude, pushing the softmax into regions where its gradient is extremely small, making training unstable. After softmax normalisation, we get a probability distribution over positions, which weights the Value vectors to produce the output.

In decoder-only models (GPT, Claude, Llama), a causal mask is applied before the softmax: positions are prevented from attending to future tokens by setting those scores to negative infinity. This enforces the autoregressive property — each token can only depend on tokens that came before it.

3. Multi-Head Attention

Rather than performing a single attention computation, the transformer runs h parallel attention heads. The model dimension is split so that each head operates on a d_k = d_model / h dimensional subspace:

head_i = Attention(X @ W_Q_i, X @ W_K_i, X @ W_V_i)
MultiHead(X) = Concat(head_1, ..., head_h) @ W_O

Each head can learn to attend to different types of relationships. Empirical analysis of trained models shows that different heads specialise: some track syntactic dependencies (subject-verb agreement), others capture positional patterns (attending to the previous token), and others handle semantic relationships (coreference resolution). The concatenated outputs are projected back to d_model dimensions by W_O.

Grouped Query Attention

Standard multi-head attention (MHA) gives each head its own Q, K, and V projections. Multi-query attention (MQA) shares a single K and V across all heads, dramatically reducing KV cache memory. Grouped query attention (GQA) is the middle ground: query heads are divided into groups, and each group shares one set of K, V heads. Llama 2 70B uses GQA with 64 query heads and 8 KV heads, reducing KV cache by 8x with negligible quality loss.

4. Positional Encoding

Self-attention treats its input as a set, not a sequence — if you shuffle the input tokens, the attention scores change only in their arrangement, not their values. Positional encodings break this symmetry by injecting position information.

Sinusoidal Encoding (Original Transformer)

The original paper used fixed sinusoidal functions at different frequencies:

PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

These are added directly to the token embeddings. The different frequencies allow the model to learn to attend to relative positions: PE(pos+k) can be expressed as a linear function of PE(pos) for any fixed offset k.

RoPE (Rotary Position Embedding)

RoPE, used by Llama, Mistral, and most modern open-source models, encodes position by rotating the Query and Key vectors in 2D subspaces. Pairs of dimensions are rotated by an angle proportional to the position. The key insight is that the dot product between rotated Q and K vectors depends only on their relative position, not absolute position. This makes RoPE naturally suited to capturing relative relationships and supports length extrapolation through frequency scaling techniques like NTK-aware interpolation.

ALiBi (Attention with Linear Biases)

ALiBi takes a different approach entirely: instead of modifying the embeddings, it adds a linear penalty to the attention scores based on the distance between tokens:

attention_score(i, j) = q_i . k_j - m * |i - j|

where m is a head-specific slope (set to geometric series, not learned). ALiBi requires zero additional parameters and offers excellent length extrapolation — models trained on short sequences generalise to much longer ones at inference time. It was used in BLOOM (176B) and MPT models.

5. Feed-Forward Network and Layer Normalisation

After the attention sub-layer, each transformer block applies a position-wise feed-forward network (FFN) independently to each token position. In the original transformer:

FFN(x) = W_2 * ReLU(W_1 * x + b_1) + b_2

The inner dimension d_ff is typically 4x the model dimension. Modern models replace ReLU with GELU or SwiGLU. SwiGLU, used by Llama and PaLM, introduces a gating mechanism: FFN(x) = W_2 * (SiLU(W_1 * x) * W_3 * x). The FFN is where the majority of the model's parameters and "knowledge" reside — attention computes which tokens to combine, and the FFN transforms what is combined.

Each sub-layer is wrapped with a residual connection (output = sublayer(x) + x) and layer normalisation. Modern models use pre-norm (normalise before the sub-layer) rather than the original post-norm, as it leads to more stable training. Many recent models use RMSNorm, which normalises by the root mean square without centering by the mean, offering equivalent quality with fewer operations.

6. Architecture Families

The original transformer has both an encoder (bidirectional attention) and a decoder (causal attention with cross-attention to encoder outputs). Modern practice has diverged into three families:

Encoder-only (BERT, RoBERTa): each token attends to all others. Excellent for classification, NER, and embedding extraction. Not suited for generation.
Decoder-only (GPT, Claude, Llama): causal masking. The dominant architecture for general-purpose language models. Every task is framed as text generation.
Encoder-decoder (T5, BART, Flan-T5): the encoder processes the full input bidirectionally, and the decoder generates output while attending to both previous decoder tokens and encoder representations via cross-attention. Natural fit for tasks with distinct input and output (translation, summarisation).

Decoder-only models have become dominant because they offer maximum flexibility — any task can be expressed as a completion — and they scale more predictably. The tradeoff is that they cannot process input bidirectionally, which is why encoder models still dominate embedding and retrieval workloads.

7. Attention Complexity and Flash Attention

Standard attention computes and stores the full [n x n] attention matrix, giving O(n²) time and memory complexity. For a 128K-token context, this matrix has over 16 billion entries — the memory cost alone can exceed the model weights.

Flash Attention (Dao et al., 2022) solves this by exploiting the GPU memory hierarchy. Instead of materialising the full attention matrix in slow GPU HBM, it computes attention in tiles that fit in fast on-chip SRAM, using an online softmax algorithm. The result is mathematically identical to standard attention but uses O(n) memory and runs 2-4x faster due to reduced memory reads/writes. Flash Attention 2 further optimises the tiling strategy and is now the default backend in PyTorch 2.0+.

8. KV Cache for Inference

During autoregressive generation, the model produces tokens one at a time. Without optimisation, generating the t-th token would require recomputing attention over all t-1 previous tokens. The KV cache stores the Key and Value projections of all previously processed tokens, so each generation step only needs to compute Q for the new token and look it up against the cached K and V matrices.

This reduces per-step complexity from O(n) matrix multiplications to O(1) per layer (ignoring the attention lookup itself). The tradeoff is memory: KV cache for a single request on Llama 2 70B at 4K tokens requires approximately 2.5 GB. For serving systems handling many concurrent requests, KV cache memory often becomes the primary bottleneck, which is why techniques like GQA, quantised KV cache, and paged attention (vLLM) are critical in production.

Key Takeaway

The transformer is a modular architecture: self-attention learns which tokens to combine, the FFN transforms what is combined, positional encodings provide sequence order, and residual connections and normalisation enable training hundreds of layers deep. Understanding each component gives you the vocabulary to read model papers, debug fine-tuning runs, and make informed decisions about architecture and serving trade-offs.