Attention and Transformers¶

This is the most important document in the entire roadmap. Everything downstream — BERT, GPT, ViT, diffusion, AlphaFold — is an application of what is derived here. Read it slowly. Implement every equation before moving on.

1. The Problem Attention Solves¶

An RNN processes a sequence one token at a time. At step t, it takes the current input x_t and the previous hidden state h_{t-1}, and produces a new hidden state h_t. The entire history of the sequence up to position t is compressed into h_t.

This creates two structural problems:

Problem 1 — Sequential dependency. Computing h_t requires h_{t-1}, which requires h_{t-2}, and so on. This is an O(n) sequential dependency chain. You cannot parallelize across time steps on a GPU. For a sequence of length 512, you have 512 sequential operations. Modern GPUs have thousands of CUDA cores that sit idle waiting for the chain to resolve.

Problem 2 — Information bottleneck. The hidden state h_t is a fixed-size vector (say, 512 dimensions). Compressing all of “what happened in the first 500 tokens of this document” into 512 numbers is lossy. Long-range dependencies suffer — the model literally cannot hold enough information to connect a pronoun at position 400 to its antecedent at position 12.

Bahdanau et al. (2014) proposed a partial solution: at each decoding step, the decoder could “look back” at all encoder hidden states and compute a weighted sum — this was the original attention mechanism. It helped, but recurrence remained.

Vaswani et al. (2017) made the decisive move: remove recurrence entirely. If attention lets every position look at every other position, you don’t need sequential hidden states at all. The transformer processes all tokens simultaneously, with every token able to directly attend to every other token in a single layer. Parallelism is fully restored. The information bottleneck is eliminated.

2. Query, Key, Value — The Library Intuition¶

Before the math, build the correct mental model.

Imagine you’re in a library looking up a topic. You have a query — the question you’re trying to answer. The library has books with keys (titles, subject headings on the spine) and values (the actual content inside). You scan all the keys, compute how relevant each book’s key is to your query, and then retrieve a weighted mix of the corresponding values — spending more time on books whose keys matched your query well.

That is exactly what self-attention computes:

Q (Query): a linear projection of the current token — “what am I looking for?”
K (Key): a linear projection of every token — “what do I contain?”
V (Value): a linear projection of every token — “what do I contribute if selected?”

The attention output for a given token is a weighted sum of all Values, where the weights are determined by how well that token’s Query matches each token’s Key.

3. Scaled Dot-Product Attention — Full Derivation¶

For a sequence of n tokens, each of dimension d_model, we compute:

Q = X W_Q       # shape: (n, d_k)
K = X W_K       # shape: (n, d_k)
V = X W_V       # shape: (n, d_v)

where W_Q, W_K ∈ ℝ^(d_model × d_k) and W_V ∈ ℝ^(d_model × d_v) are learned projection matrices.

The attention formula:

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

Let’s unpack each component:

Step 1 — Similarity scores: S = Q K^T
Shape: (n, n). Entry S_{ij} = dot product between query of token i and key of token j. High value = token i finds token j relevant.

Step 2 — Scale: S_scaled = S / √d_k
Why divide by √d_k? With large d_k, dot products grow in magnitude proportional to d_k (each of d_k terms contributes ~1 in expectation if Q, K are standard normal). If d_k = 64, typical dot products have variance 64, so standard deviation ≈ 8. The softmax input ranges over something like [-40, +40]. Softmax saturates in this regime — most weight goes to one extreme position and gradients of the other positions vanish. Dividing by √d_k restores variance to 1, keeping softmax in its stable, gradient-friendly operating range. This is not a heuristic; it is a variance normalization.

Step 3 — Attention weights: A = softmax(S_scaled)
Applied row-wise. Each row i is a probability distribution over all n tokens — “what fraction of its attention does token i allocate to each other token?” Shape: (n, n).

Step 4 — Weighted value sum: Output = A V
Shape: (n, d_v). Each output token is a weighted combination of all value vectors. Tokens with high attention weight contribute more to the output.

In one equation:

Attention(Q, K, V) = softmax( Q K^T / √d_k ) · V

Computational cost: O(n² · d_k) for the Q K^T multiplication. This is why long sequences are expensive — the attention matrix scales quadratically with sequence length. For n=4096, that’s 16M attention weights per head per layer.

4. Multi-Head Attention¶

A single attention operation learns one type of relationship. But a sentence has multiple simultaneous structures: syntactic dependencies, semantic co-reference, short-range n-gram patterns, long-range discourse structure. You want the model to track all of these simultaneously.

Multi-head attention runs h independent attention operations in parallel, each with its own projections:

head_i = Attention(Q W_Q_i, K W_K_i, V W_V_i)

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_O

where each W_Q_i ∈ ℝ^(d_model × d_k), and typically d_k = d_model / h (so the total computation is comparable to single-head attention at full dimension).

Why this works: Different heads empirically learn to specialize. In BERT, analyses show some heads track subject-verb agreement, others track coreference, others attend predominantly to adjacent tokens (acting like a local window). This wasn’t designed in — it emerges from training. The concatenation and final projection W_O allow the model to integrate all these relationship types into a single output representation.

For BERT-base: d_model = 768, h = 12 heads, d_k = d_v = 64.

5. Positional Encoding¶

Here is a critical property of the attention formula: it is permutation-invariant. If you shuffle all the tokens in your input sequence, the attention computation produces the same output (just shuffled). There is no inherent notion of “token 3 comes before token 4.”

This is a problem. Language is ordered. “The dog bit the man” ≠ “The man bit the dog.”

The fix: inject positional information as an additive signal to the token embeddings before feeding them to the transformer.

Sinusoidal Encoding (Original Transformer)¶

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

where pos is the position index and i is the dimension index. This creates a unique encoding for each position using sine/cosine waves of different frequencies. The encoding is added (not concatenated) to the token embedding.

Why sinusoidal? The original paper argues that for any fixed offset k, PE(pos+k) can be expressed as a linear function of PE(pos). This means the model can learn to attend to relative positions. Also, it generalizes to sequence lengths longer than those seen during training.

Learned Positional Encoding (GPT-2 style)¶

Simply learn an embedding table of shape (max_seq_len, d_model). Each position gets its own trainable vector. Simpler, often works as well, does not generalize beyond max_seq_len.

RoPE — Rotary Position Embedding (2024 Standard)¶

RoPE (Su et al., 2021, arXiv:2104.09864) encodes position by rotating the query and key vectors in 2D planes. The key property: the dot product Q_m · K_n depends only on the relative offset (m - n), not absolute positions. This gives the model true relative position awareness without explicit relative position bias.

Used in: LLaMA, Mistral, Gemma, Qwen, and virtually every open-source LLM released after 2022. RoPE is the current default. If you are implementing a transformer from scratch in 2025, use RoPE.

6. The Full Transformer Layer¶

A single transformer encoder layer applies these operations in sequence:

# Pre-norm (modern, more stable)
x = x + MultiHeadAttention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

LayerNorm: Normalizes across the feature dimension for each token independently. Stabilizes training. Applied before the sublayer in the modern “pre-norm” variant.

Multi-Head Attention sublayer: As derived above.

Residual connection: x = x + sublayer(x). Critical for gradient flow in deep networks. Without residuals, gradients vanish over 12+ layers.

FFN (Feed-Forward Network):

FFN(x) = max(0, x W_1 + b_1) W_2 + b_2

Two linear layers with a ReLU (or in modern variants, GELU, SwiGLU) in between. Dimension typically expands to 4 × d_model then contracts back: W_1 ∈ ℝ^(d_model × 4·d_model). The FFN is applied independently to each token — it is not a sequence operation. Approximately 2/3 of transformer parameters live in the FFN layers, not the attention.

Pre-norm vs. post-norm:
The original paper used post-norm: x = LayerNorm(x + sublayer(x)). This is numerically less stable for deep networks because the LayerNorm gate is after the residual, allowing gradient magnitudes to grow unchecked through the residual stream. Pre-norm applies LayerNorm before: x = x + sublayer(LayerNorm(x)). The residual stream remains unnormalized, which empirically allows training of much deeper networks without gradient explosion. All modern large transformers use pre-norm.

7. Encoder vs. Decoder¶

Encoder (BERT-style): all tokens attend to all tokens — bidirectional. Token at position 5 can see tokens at positions 1, 2, 3, 4, 6, 7, … The model sees the full context when building each token’s representation. Appropriate for tasks where you have the complete input (classification, NER, extractive QA).

Decoder (GPT-style): autoregressive — each token can only attend to previous tokens. Implemented via a causal mask: set S_{ij} = -∞ for j > i before softmax, which forces A_{ij} = 0 for future positions. This is what makes GPT generative: it can produce token t+1 given only tokens 1..t. The constraint is architectural, not just a training choice.

Encoder-Decoder (original transformer for machine translation): encoder processes the source sequence bidirectionally; decoder generates the target sequence autoregressively with a cross-attention sublayer that attends to the encoder’s output. The decoder has three sublayers: self-attention (causal), cross-attention (to encoder), FFN.

8. Implementation — Attention From Scratch¶

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (batch, heads, seq_len, d_k)
    K: (batch, heads, seq_len, d_k)
    V: (batch, heads, seq_len, d_v)
    mask: (batch, 1, seq_len, seq_len) — optional causal or padding mask
    Returns: (batch, heads, seq_len, d_v)
    """
    d_k = Q.size(-1)
    
    # Similarity scores: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply mask if provided (e.g., causal mask for decoder)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # Attention weights
    attn_weights = F.softmax(scores, dim=-1)
    
    # Weighted value sum
    output = torch.matmul(attn_weights, V)
    return output, attn_weights


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # dimension per head
        
        # Projection matrices
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)
    
    def split_heads(self, x):
        """(batch, seq_len, d_model) → (batch, heads, seq_len, d_k)"""
        batch_size, seq_len, _ = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, x, mask=None):
        batch_size = x.size(0)
        
        Q = self.split_heads(self.W_Q(x))  # (batch, heads, seq, d_k)
        K = self.split_heads(self.W_K(x))
        V = self.split_heads(self.W_V(x))
        
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # Concatenate heads: (batch, seq_len, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)
        
        return self.W_O(attn_output)


class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        
        self.attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # Pre-norm attention sublayer
        x = x + self.dropout(self.attn(self.norm1(x), mask))
        # Pre-norm FFN sublayer
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x


# ── Verification: compare to torch.nn.MultiheadAttention ──────────────────────

def verify_equivalence():
    torch.manual_seed(42)
    batch, seq_len, d_model, num_heads = 2, 10, 64, 4
    
    x = torch.randn(batch, seq_len, d_model)
    
    # Our implementation
    our_mha = MultiHeadAttention(d_model, num_heads)
    
    # PyTorch implementation — copy weights to ensure fair comparison
    torch_mha = nn.MultiheadAttention(d_model, num_heads, bias=False, batch_first=True)
    
    # Copy weights from our model to PyTorch's
    # PyTorch packs Q, K, V projections into in_proj_weight
    in_proj = torch.cat([
        our_mha.W_Q.weight,
        our_mha.W_K.weight,
        our_mha.W_V.weight
    ], dim=0)
    torch_mha.in_proj_weight.data.copy_(in_proj)
    torch_mha.out_proj.weight.data.copy_(our_mha.W_O.weight)
    
    our_output = our_mha(x)
    torch_output, _ = torch_mha(x, x, x, need_weights=False)
    
    max_diff = (our_output - torch_output).abs().max().item()
    print(f"Max absolute difference: {max_diff:.2e}")
    assert max_diff < 1e-5, f"Outputs diverge: {max_diff}"
    print("PASS — implementations are numerically equivalent.")

if __name__ == "__main__":
    verify_equivalence()

Expected output: Max absolute difference: < 1e-5 followed by PASS.

9. Complexity Summary¶

Operation	Time	Space
Attention score matrix	O(n² · d_k)	O(n²)
FFN per token	O(n · d_model · d_ff)	O(n · d_ff)
Full transformer layer	O(n² · d_model + n · d_model²)	O(n² + n · d_model)

The O(n²) attention cost is why long-context models (>100K tokens) require approximations: FlashAttention rewrites the attention kernel to minimize memory bandwidth (same O(n²) computation, but IO-optimal). Linear attention methods approximate the softmax to get O(n). Know which you’re using and what you’re trading away.

What Most People Get Wrong¶

Attention is not “what the model is paying attention to” in a human-interpretable sense. It is a similarity-weighted retrieval operation. When you visualize attention weights and claim “the model focused on the subject when processing the verb,” you are post-hoc rationalizing a linear algebra operation. The weights reflect softmax-normalized dot-product similarities in a learned projection space. They have no guaranteed semantic meaning.

More concretely: multiple attention heads can collectively route information through paths that look uninterpretable when inspected head-by-head. Anthropic’s mechanistic interpretability research has documented “induction heads,” “copy heads,” and other functional circuits that only make sense when analyzed as coordinated multi-head operations, not single-head attention maps.

Use attention visualization for debugging gradient flow and diagnosing degenerate patterns (e.g., all weight on [CLS], attention collapse). Do not use it as explanation.