Transformers & Attention Mechanisms

🎯 The Attention Problem

Before Transformers, the state-of-the-art for sequence-to-sequence tasks (like machine translation) was an RNN encoder-decoder architecture. The encoder read the entire source sentence and compressed it into a single fixed-size hidden state vector — the "thought vector" — which the decoder would then use to generate the output one word at a time. This bottleneck was the fatal flaw.

🔑 Scaled Dot-Product Attention

The Transformer paper (Vaswani et al., 2017) formalised attention using three components — Queries, Keys, and Values — and introduced scaled dot-product as the scoring function. This formulation is elegant, parallelisable, and extraordinarily effective.

# Scaled Dot-Product Attention — the core formula
Attention(Q, K, V) = softmax(Q · Kᵀ / √dₖ) · V

# Shapes (batch omitted for clarity):
#   Q (Queries): shape (seq_len_q, dₖ)   — "what am I looking for?"
#   K (Keys):    shape (seq_len_k, dₖ)   — "what do I contain?"
#   V (Values):  shape (seq_len_v, dᵥ)   — "what do I output if selected?"
#   Output:      shape (seq_len_q, dᵥ)   — weighted combination of Values

# Step-by-step walkthrough for a 3-token sequence ["The", "cat", "sat"]:

# Step 1: Project input embeddings into Q, K, V spaces (via learned Wq, Wk, Wv matrices)
Q = X · Wq    # (3 × dₖ)
K = X · Wk    # (3 × dₖ)
V = X · Wv    # (3 × dᵥ)

# Step 2: Compute raw attention scores — how much each query attends to each key
scores = Q · Kᵀ           # (3 × 3) matrix of dot products

# Step 3: Scale by √dₖ to prevent softmax saturation in high dimensions
scores_scaled = scores / √dₖ

# Step 4: Convert scores to probabilities
weights = softmax(scores_scaled, dim=-1)   # each row sums to 1.0
# weights[i][j] = how much token i should attend to token j

# Step 5: Weighted sum of values
output = weights · V      # (3 × dᵥ)
# Each output token is a blended combination of all value vectors,
# proportional to attention weights.

# Example attention weights for "cat" might look like:
# "cat" attends to: "The"→0.05, "cat"→0.85, "sat"→0.10
# → output for "cat" is mostly the value of "cat" itself (self-attention)
      

🔀 Multi-Head Attention

Single-head attention can only attend to information from one "perspective" at a time. Multi-head attention runs h attention operations in parallel, each with its own learned Q/K/V projections, then concatenates and projects the results. Different heads learn to specialise in different types of relationships.

# Multi-Head Attention — h parallel attention heads
MultiHead(Q, K, V) = Concat(head₁, head₂, …, headₕ) · Wₒ

# Each head i uses separate projection matrices:
headᵢ = Attention(Q · Wᵢq, K · Wᵢk, V · Wᵢv)

# Typical dimensions (GPT-3 large example):
#   d_model = 1024    (embedding dimension)
#   h = 16            (number of heads)
#   dₖ = dᵥ = d_model / h = 64   (per-head dimension)
#
# Each head works in a 64-dimensional subspace.
# Concatenating 16 × 64-dim outputs = 1024-dim, then project with Wₒ.

# Total parameters per MHA layer:
#   16 heads × (Wq + Wk + Wv): 3 × 1024 × 1024 = 3.1M
#   Output projection Wₒ: 1024 × 1024 = 1.0M
#   Total: ~4.1M parameters per MHA layer
      

🏗️ The Full Transformer Architecture

The Transformer is built from two stacks: an encoder (processes the input) and a decoder (generates the output). Many modern models use only one stack: BERT uses encoders only; GPT uses decoders only. T5 uses both.

# ── ENCODER (one layer, repeated N times, e.g. N=12 in BERT-base) ──

Input tokens  →  Token Embeddings + Positional Encoding
                          ↓
               ┌─────────────────────────┐
               │  Multi-Head Self-Attention│  ← every token attends to every other token
               └──────────┬──────────────┘
                          │ + residual (add input to attention output)
                    Layer Normalisation
                          ↓
               ┌─────────────────────────┐
               │   Feed-Forward Network   │  ← 2-layer MLP applied independently to each position
               │ (expand to 4×d → GELU → │
               │  project back to d)      │
               └──────────┬──────────────┘
                          │ + residual
                    Layer Normalisation
                          ↓  (becomes input to next encoder layer)

# ── DECODER (one layer, also repeated N times) ──

Target tokens (shifted right) →  Embeddings + Positional Encoding
                          ↓
              Masked Multi-Head Self-Attention  ← causal mask: can only see past positions
                          ↓ + residual + LayerNorm
              Multi-Head Cross-Attention        ← Q from decoder, K/V from encoder output
                          ↓ + residual + LayerNorm
              Feed-Forward Network
                          ↓ + residual + LayerNorm
                     Linear + Softmax → Output probabilities

# ── POSITIONAL ENCODING ──
# Since Transformers process all positions in parallel, they have no inherent
# sense of order. Positional encodings add position information to embeddings.

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

# Learned positional embeddings (BERT, GPT-2) are an alternative.
      

🌐 Transformer Variants & Models

Since 2017, the Transformer has spawned hundreds of variants. The most important distinction is the architecture type: encoder-only (good at understanding), decoder-only (good at generation), or encoder-decoder (good at transformation tasks like translation or summarisation).