1 Why RNNs Were Replaced

To appreciate why the Transformer was such a breakthrough, you need to understand what it replaced. Recurrent Neural Networks (RNNs) and LSTMs were the dominant architecture for sequence processing through the mid-2010s. They had two serious problems that became crippling at scale.

Problem 1: Sequential Processing — Parallelism Is Impossible

An RNN processes tokens one at a time, left to right. To compute the hidden state at step t, you must first have the hidden state at step t-1. This is a fundamental sequential dependency — you cannot compute step 5 until you've finished steps 1, 2, 3, and 4. On a modern GPU with thousands of CUDA cores that are all idle and waiting, this is a catastrophic inefficiency.

For a sequence of 1,000 tokens, an RNN requires exactly 1,000 sequential matrix multiplications. A Transformer computes attention for ALL token pairs simultaneously — a single parallel matrix operation. This difference in training speed is not 2× or 10×. It is measured in orders of magnitude. Training GPT-3 with RNNs would take years; with Transformers it takes weeks.

Problem 2: Long-Range Dependencies Still Struggle

LSTMs improved over vanilla RNNs with their gating mechanisms, but the core problem remained: information from the first token in a sequence must travel through every subsequent hidden state to reach the last. After 500 steps, the signal from token 1 is hopelessly diluted by everything that came after. This is why LSTMs still struggled with very long documents, long-horizon dependencies in code ("this variable was declared 800 lines above"), and cross-sentence reasoning.

The Transformer's attention mechanism has a fundamentally different property: every token can directly attend to every other token in a single step, regardless of their distance. Token 1 and token 1,000 have exactly the same "distance" in attention space. Long-range dependencies are not just easier — they are structurally free.

import torch
import time

# Illustrate sequential vs parallel processing
seq_len = 512
hidden_dim = 768
batch_size = 8

# Simulate RNN: must loop through time steps sequentially
rnn = torch.nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
x = torch.randn(batch_size, seq_len, hidden_dim)

start = time.time()
for _ in range(10):
    out, _ = rnn(x)  # internally sequential — can't parallelize over time
rnn_time = (time.time() - start) / 10

# Simulate Transformer self-attention: one parallel matrix op
attention = torch.nn.MultiheadAttention(hidden_dim, num_heads=12, batch_first=True)

start = time.time()
for _ in range(10):
    out, _ = attention(x, x, x)  # ALL positions attended in parallel
attn_time = (time.time() - start) / 10

print(f"LSTM forward (seq={seq_len}):       {rnn_time*1000:.1f} ms")
print(f"Self-Attention forward (seq={seq_len}): {attn_time*1000:.1f} ms")
print(f"Speedup: {rnn_time/attn_time:.1f}x")

2 Positional Encoding: Injecting Order Information

Here's the crucial problem that emerges when you replace sequential RNNs with parallel attention: attention is permutation-invariant. If you shuffle the tokens in the input, self-attention produces the exact same attention scores (just rearranged). The phrases "cat sat on mat" and "mat on sat cat" would produce identical representations for each word — because attention only computes pairwise similarities between token embeddings, which don't change when you reorder the tokens.

But word order is meaning. "The dog bit the man" and "The man bit the dog" have opposite meanings. We must explicitly tell the Transformer where each token sits in the sequence.

Sinusoidal Positional Encoding

The original Transformer paper's elegant solution: add a unique vector to each token embedding that encodes its position. This vector is computed using sine and cosine functions at different frequencies:

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

Intuition: think of it like a clock or a binary counter, but smooth. Each dimension of the encoding oscillates at a different frequency — some dimensions oscillate fast (differentiating nearby positions), others oscillate slowly (differentiating distant positions). Together, they give each position a unique "fingerprint" that is smooth enough for the model to learn relative position patterns.

Key properties: (1) each position has a unique encoding; (2) nearby positions have similar encodings; (3) the encoding generalises to sequence lengths longer than seen during training; (4) the model can learn to use relative positions because PE(pos+k) is a linear function of PE(pos).

import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Build the PE matrix once: shape (max_len, d_model)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()  # (max_len, 1)
        # Denominator: 10000^(2i/d_model)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)  # even dims: sin
        pe[:, 1::2] = torch.cos(position * div_term)  # odd dims: cos

        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model) — batch dimension
        self.register_buffer('pe', pe)  # not a parameter, but moves with model.to(device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

# Visualize the PE matrix
d_model = 128
pe_module = SinusoidalPositionalEncoding(d_model, max_len=50)
pe_matrix = pe_module.pe[0].numpy()  # shape: (50, 128)

print(f"PE matrix shape: {pe_matrix.shape}")
print(f"Position 0, first 8 dims: {pe_matrix[0, :8].round(3)}")
print(f"Position 1, first 8 dims: {pe_matrix[1, :8].round(3)}")
print(f"Position 10, first 8 dims: {pe_matrix[10, :8].round(3)}")

# Show that positions are unique
pos0 = pe_matrix[0]
pos1 = pe_matrix[1]
pos10 = pe_matrix[10]
similarity_01  = np.dot(pos0, pos1)  / (np.linalg.norm(pos0) * np.linalg.norm(pos1))
similarity_010 = np.dot(pos0, pos10) / (np.linalg.norm(pos0) * np.linalg.norm(pos10))
print(f"\nCosine similarity pos 0 vs pos 1:  {similarity_01:.3f}  (nearby = similar)")
print(f"Cosine similarity pos 0 vs pos 10: {similarity_010:.3f} (distant = less similar)")

Numbers on their own are hard to build intuition from. Let's see the positional encoding matrix instead. The heatmap below plots PE(pos, dim) for every position (rows) and embedding dimension (columns) — and the line chart shows a handful of individual dimensions as waves across position. Drag the slider to change the model dimension d_model and watch how the frequency spectrum stretches or compresses:

Modern Alternatives: Learned and Rotary Positional Embeddings

The original sinusoidal encoding was later found to be improvable. Modern models use two main alternatives:

Learned positional embeddings (used in BERT, GPT-2): instead of a formula, treat positions 0 through max_len as a vocabulary of their own. Learn a separate embedding for each position via backpropagation. Slightly better in practice, but cannot generalize to longer sequences than seen during training.

Rotary Position Embeddings (RoPE) (used in Llama, Mistral, Falcon): rather than adding a positional vector, rotate the query and key vectors in complex space by an angle proportional to the position. The clever property: the dot product Q_pos_i · K_pos_j naturally encodes only the RELATIVE distance (i − j), not absolute positions. This makes RoPE naturally length-generalisable and gives better performance on long contexts.

# BERT-style learned positional embeddings (conceptually simple)
class LearnedPositionalEmbedding(nn.Module):
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        # One learnable embedding per position — just like word embeddings
        self.pos_embedding = nn.Embedding(max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device)  # [0, 1, 2, ..., seq_len-1]
        return x + self.pos_embedding(positions)  # broadcast over batch

# Both approaches can be combined: token emb + position emb + token type emb (BERT)
class BERTEmbedding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, max_len: int):
        super().__init__()
        self.token_emb    = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.position_emb = nn.Embedding(max_len, d_model)
        self.token_type   = nn.Embedding(2, d_model)  # sentence A=0, sentence B=1
        self.layer_norm   = nn.LayerNorm(d_model)
        self.dropout      = nn.Dropout(0.1)

    def forward(self, token_ids, token_type_ids=None):
        seq_len = token_ids.size(1)
        positions = torch.arange(seq_len, device=token_ids.device)
        x = self.token_emb(token_ids) + self.position_emb(positions)
        if token_type_ids is not None:
            x = x + self.token_type(token_type_ids)
        return self.dropout(self.layer_norm(x))

3 The Transformer Encoder Block

Think of the Transformer encoder as a stack of identical "processing floors." Each floor takes a sequence of vectors, enriches them with context from every other position (via attention), then applies a simple position-wise neural network (the FFN) to further transform each vector independently. The output is the same sequence of vectors — same shape as the input — but now each vector contains information from the entire sequence.

One encoder layer has exactly three components, applied in sequence:

Component 1: Multi-Head Self-Attention (covered in Lesson 52)

Each token queries all other tokens, asking "which positions are most relevant to understanding me?" The output for each position is a weighted sum of all value vectors. The "multi-head" part runs this process independently h times with different learned projections, then concatenates the results — letting the model attend to different aspects of the input simultaneously (syntax, semantics, co-reference, etc.).

Component 2: Position-Wise Feed-Forward Network (FFN)

After attention mixes information across positions, a two-layer MLP is applied independently to each position's vector. This is where the model can process the newly-gathered contextual information and combine features non-linearly. The inner dimension is typically 4× the model dimension (3072 in BERT-base with d_model=768). This sounds large but is where most of the model's "memory" lives — ablation studies show the FFN stores factual knowledge.

Component 3: Add & Norm (Residual + LayerNorm)

After each of the above components, two operations are applied: (1) add the block's input back to its output (residual connection), and (2) apply LayerNorm. The residual connection is the gradient highway from Lesson 40 — it allows gradients to flow directly to early layers without passing through every attention and FFN layer. LayerNorm normalises across the feature dimension (not batch), making it independent of batch size and suitable for sequences of variable length.

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

class TransformerEncoderBlock(nn.Module):
    """One Transformer encoder layer: attention + FFN, each with add-and-norm."""
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        # Multi-head self-attention
        self.self_attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=num_heads,
            dropout=dropout, batch_first=True
        )
        # Position-wise feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        # Layer normalisations (post-norm, original paper style)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, src_key_padding_mask=None) -> torch.Tensor:
        # x shape: (batch, seq_len, d_model)

        # --- Sub-layer 1: Multi-Head Self-Attention ---
        attn_out, attn_weights = self.self_attn(
            query=x, key=x, value=x,
            key_padding_mask=src_key_padding_mask
        )
        x = self.norm1(x + self.dropout(attn_out))  # Add & Norm

        # --- Sub-layer 2: Feed-Forward Network ---
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))   # Add & Norm

        return x

# Stack N encoder blocks to build a full encoder
class TransformerEncoder(nn.Module):
    def __init__(self, vocab_size, d_model=256, num_heads=8, d_ff=1024, num_layers=4, dropout=0.1):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_enc   = SinusoidalPositionalEncoding(d_model, dropout=dropout)
        self.layers    = nn.ModuleList([
            TransformerEncoderBlock(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, token_ids, padding_mask=None):
        x = self.token_emb(token_ids)  # (B, T) → (B, T, d_model)
        x = self.pos_enc(x)
        for layer in self.layers:
            x = layer(x, src_key_padding_mask=padding_mask)
        return self.norm(x)

# Test it
vocab_size = 10000
encoder = TransformerEncoder(vocab_size, d_model=256, num_heads=8, num_layers=4)
token_ids = torch.randint(1, vocab_size, (4, 50))  # batch=4, seq_len=50
output = encoder(token_ids)
print(f"Input shape:  {token_ids.shape}")     # (4, 50)
print(f"Output shape: {output.shape}")        # (4, 50, 256)
total_params = sum(p.numel() for p in encoder.parameters())
print(f"Total parameters: {total_params:,}") # ~12M

4 The Transformer Decoder Block

The decoder is used in sequence-to-sequence tasks: you have an encoded source sequence (e.g., French text) and you want to generate a target sequence (e.g., English text) one token at a time. The decoder needs to simultaneously: (1) attend to the tokens it has generated so far, and (2) attend to the encoder's output to "read" the source.

A decoder block has THREE sub-layers instead of the encoder's two:

Sub-layer 1: Masked Multi-Head Self-Attention

The decoder attends to the tokens it has already generated. Critically, it uses a causal mask (also called an autoregressive mask) that prevents each position from attending to future positions. When generating token at position 5, the model may only see tokens 1–5, not 6–N. This mask is a triangular matrix of −∞ values added to the attention scores before the softmax — after softmax, those positions have probability 0.

Sub-layer 2: Cross-Attention (Encoder-Decoder Attention)

This is the "reading comprehension" layer. The Query vectors come from the decoder (what am I looking for?), but the Key and Value vectors come from the encoder's output (what information is available in the source?). This is how translation works: when generating each English word, the decoder attends back to the relevant French words in the encoder output.

Sub-layer 3: Feed-Forward Network

Same position-wise MLP as in the encoder.

class TransformerDecoderBlock(nn.Module):
    """One Transformer decoder layer: masked self-attn + cross-attn + FFN."""
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        # Masked self-attention (causal)
        self.self_attn  = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
        # Cross-attention: Q from decoder, K/V from encoder
        self.cross_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
        # Position-wise FFN
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def make_causal_mask(self, seq_len: int, device) -> torch.Tensor:
        """Upper-triangular mask with -inf: prevents attending to future positions."""
        mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
        return mask.masked_fill(mask == 1, float('-inf'))

    def forward(self, tgt, memory, tgt_key_padding_mask=None, memory_key_padding_mask=None):
        # tgt: (B, T_tgt, d_model) — decoder input
        # memory: (B, T_src, d_model) — encoder output
        T = tgt.size(1)
        causal_mask = self.make_causal_mask(T, tgt.device)  # (T, T)

        # Sub-layer 1: Masked self-attention
        sa_out, _ = self.self_attn(tgt, tgt, tgt,
                                   attn_mask=causal_mask,
                                   key_padding_mask=tgt_key_padding_mask)
        x = self.norm1(tgt + self.dropout(sa_out))

        # Sub-layer 2: Cross-attention (Q from decoder, K/V from encoder output)
        ca_out, ca_weights = self.cross_attn(x, memory, memory,
                                             key_padding_mask=memory_key_padding_mask)
        x = self.norm2(x + self.dropout(ca_out))

        # Sub-layer 3: Feed-forward
        x = self.norm3(x + self.dropout(self.ffn(x)))
        return x, ca_weights  # return cross-attn weights for visualization

# Demonstrate encoder → decoder flow
d_model = 256
enc_out  = torch.randn(4, 30, d_model)   # encoder output: 4 sequences, 30 src tokens
dec_input = torch.randn(4, 20, d_model)  # decoder input: 4 seqs, 20 tgt tokens so far

decoder_block = TransformerDecoderBlock(d_model=256, num_heads=8, d_ff=1024)
dec_out, attn_w = decoder_block(dec_input, enc_out)
print(f"Decoder output:      {dec_out.shape}")   # (4, 20, 256)
print(f"Cross-attn weights:  {attn_w.shape}")    # (4, 20, 30) — each tgt attends to src

5 Encoder-Only vs Decoder-Only vs Encoder-Decoder

The original Transformer had both encoder and decoder. But researchers quickly discovered you could use just one part, depending on your task. This split led to three distinct families that dominate NLP today:

Encoder-Only: BERT Family

Encoder-only models process the full sequence with bidirectional attention — every token attends to every other token freely. This gives the richest possible contextual representations. The tradeoff: they cannot generate text autoregressively. Use them for understanding tasks where you need to classify, compare, or extract from existing text.

• Text classification: sentiment analysis, spam detection, topic labeling
• Named entity recognition: extract names, dates, organisations
• Question answering: BERT reads a passage and finds the answer span
• Semantic similarity: sentence embeddings, duplicate detection

Key models: BERT, RoBERTa, ALBERT, DeBERTa, ModernBERT.

Decoder-Only: GPT Family

Decoder-only models use the causal (autoregressive) mask. They naturally generate text by predicting one token at a time from left to right. Each position can only see previous positions — so the representations are less rich for classification than BERT's, but the architecture is perfect for generation.

• Text generation: story writing, code completion, chat
• Summarization: generate a summary as output
• Translation: prompt "Translate to English: ..."
• In-context learning: few-shot classification via prompting

Key models: GPT-2/3/4, Llama, Mistral, Falcon, Gemma.

Encoder-Decoder: T5, BART Family

The full architecture with both encoder and decoder. The encoder reads the input with full bidirectional attention; the decoder generates the output autoregressively while cross-attending to the encoder. Best for tasks where both understanding the input deeply AND generating fluent output are required.

• Machine translation: encode French, decode English
• Abstractive summarization: encode article, decode summary
• Structured output: encode a question, decode a structured answer

Key models: T5, BART, MarianMT, PEGASUS.

6 How Transformers Are Pre-trained

An untrained Transformer is just random weights. The magic comes from pre-training on enormous text corpora — hundreds of billions of words — which teaches the model the statistical structure of language. The pre-training objective differs between the two main families.

BERT Pre-training: Masked Language Modeling (MLM)

BERT's training procedure: take a sentence, randomly mask 15% of its tokens (replace with a [MASK] token), then train the model to predict the original masked tokens. Because BERT can see all tokens on both sides of the masked positions, it learns deep bidirectional context. After pre-training, BERT "knows" that "The [MASK] sat on the mat" should be "cat" — it has learned syntax, semantics, commonsense, and world knowledge from billions of such examples.

The 15% masking splits as: 80% replaced with [MASK], 10% replaced with a random word (forces the model to maintain a representation for every token), 10% unchanged (helps the representation stay grounded). BERT also had a secondary objective, Next Sentence Prediction (NSP) — predict if sentence B follows sentence A — though later work showed NSP adds little value.

GPT Pre-training: Causal Language Modeling (CLM)

GPT's training is conceptually simpler: given a document, for every position t, predict the next token using only the tokens before it. The entire document is the training signal — every token is a prediction target. This scales beautifully: more data = more training signal, no special processing needed. The simplicity is the point: if you just train long enough on enough text, the model learns everything.

T5: Text-to-Text Transfer Transformer

T5 (Raffel et al., 2020) takes a radical approach: cast EVERY task as a text-to-text problem. Translation input: "translate English to German: The cat sat". Summarization: "summarize: [article text]". Classification: "mnli premise: ... hypothesis: ... entailment or contradiction?". This unified format allows training on many tasks simultaneously with the same loss function.

from transformers import AutoTokenizer

# BERT tokenizer: MLM masking example
bert_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
text = "The cat sat on the mat."
tokens = bert_tokenizer.tokenize(text)
print("BERT tokens:", tokens)
# ['the', 'cat', 'sat', 'on', 'the', 'mat', '.']

# Encode with special tokens
encoded = bert_tokenizer(text, return_tensors='pt')
print("Input IDs:    ", encoded['input_ids'])
print("Attention mask:", encoded['attention_mask'])
# [CLS] token=101 is prepended, [SEP] token=102 is appended

# What MLM does during training: mask 15% of tokens
import torch
input_ids = encoded['input_ids'].clone()
labels = input_ids.clone()

# Create a random mask (exclude [CLS] and [SEP])
mask_prob = 0.15
rand = torch.rand(input_ids.shape)
# Mask only non-special tokens (not [CLS]=101, [SEP]=102, [PAD]=0)
maskable = (input_ids != 101) & (input_ids != 102) & (input_ids != 0)
mask_indices = (rand < mask_prob) & maskable

input_ids[mask_indices] = 103   # 103 = [MASK] token in BERT
labels[~mask_indices] = -100    # -100 = ignore in cross-entropy loss

print("\nOriginal IDs:", encoded['input_ids'].tolist())
print("Masked IDs:  ", input_ids.tolist())
print("Labels:      ", labels.tolist())  # -100 for unmasked, original ID for masked

7 Scaling Laws: Why Bigger Is Better

One of the most profound empirical findings in modern ML is that neural language model performance follows clean, predictable scaling laws. Kaplan et al. (2020) at OpenAI found that model performance (measured by test loss) follows a power law with three variables: model size (N parameters), training data size (D tokens), and compute budget (C FLOPs). Double any one of them and you get a reliable, predictable improvement.

The key insight: these scaling laws hold over many orders of magnitude. The same relationship that describes the performance jump from a 10M parameter model to a 100M parameter model also describes the jump from 1B to 10B parameters. This predictability is what allowed OpenAI and others to confidently invest in training ever-larger models — they could forecast the return before spending the compute.

The Chinchilla Correction

Hoffmann et al. (2022) from DeepMind trained a 70B parameter model (Chinchilla) on 1.4T tokens and showed it matched or outperformed the 280B parameter Gopher model. Their finding: for a given compute budget, previous models (GPT-3, Gopher) were heavily over-parameterized and under-trained. The optimal compute allocation is approximately 20 tokens per parameter. A 7B model should train on ~140B tokens, not the 300B+ tokens used by LLaMA 1 (which is actually more than Chinchilla-optimal — an intentional choice to trade training compute for inference efficiency).

import numpy as np

# Simplified Chinchilla scaling law: compute optimal tokens for a given model size
def chinchilla_optimal_tokens(num_params: float) -> float:
    """Chinchilla: optimal training tokens ≈ 20 × number of parameters."""
    return 20 * num_params

models = {
    'BERT-base':   110e6,
    'GPT-2 Small': 117e6,
    'GPT-2 XL':    1.5e9,
    'Llama-3 8B':  8e9,
    'Llama-3 70B': 70e9,
    'GPT-3':       175e9,
}

print(f"{'Model':<15} {'Params':>12} {'Optimal Tokens':>18} {'Optimal Data (GB)':>20}")
print("-" * 68)
for name, params in models.items():
    optimal_tokens = chinchilla_optimal_tokens(params)
    # rough estimate: 1 token ≈ 4 bytes
    data_gb = optimal_tokens * 4 / 1e9
    print(f"{name:<15} {params:>12.1e} {optimal_tokens:>18.1e} {data_gb:>20.1f} GB")

8 Attention Complexity and Efficient Transformers

Standard self-attention has O(n²) time and memory complexity in sequence length n. For each of the n tokens, you compute attention scores against every other n token — producing an n×n attention matrix. For n=512 (BERT's limit), that's 262,144 attention scores per head, per layer. For n=4,096 (a reasonable document length), that's 16,777,216 scores — and the memory required grows quadratically.

This quadratic complexity is why BERT has a 512-token limit and original GPT-2 has a 1024-token limit. For long documents, genomes, or high-resolution images, this becomes computationally prohibitive.

FlashAttention: Same Math, Much Less Memory

Dao et al. (2022) observed that the bottleneck was not the number of operations (FLOPs) but memory bandwidth — repeatedly reading and writing the large attention matrix to GPU VRAM. FlashAttention fuses the softmax and matrix multiplication steps and processes the attention in tiles that fit in the GPU's fast on-chip SRAM (like L1 cache). The result: exactly the same mathematical output as standard attention, but using O(n) memory instead of O(n²). This enabled context lengths of 32K, 100K, and eventually 1M+ tokens.

Sparse and Linear Attention Variants

Longformer: each token attends only to its local window (e.g., 512 surrounding tokens) plus a few global tokens (like [CLS]) that attend to everything. Complexity: O(n). Great for processing long documents with mostly local dependencies.

BigBird: random attention (each token attends to a random subset) + local attention + global tokens. Proved that O(n) sparse attention can approximate O(n²) full attention for many tasks.

Linear attention: approximate the softmax with a kernel function that factorises Q·K^T into separate operations — reducing complexity to O(n) but with some accuracy loss on tasks requiring precise attention.

import torch
import time

def standard_attention(Q, K, V, mask=None):
    """Standard O(n²) scaled dot-product attention."""
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)  # (B, H, T, T)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, V), attn

# Memory scaling demonstration
device = 'cpu'
d_k = 64
batch, heads = 1, 8

print(f"{'Seq Len':<12} {'Attention Matrix Size':<25} {'Memory (MB)':<15}")
print("-" * 52)
for seq_len in [128, 256, 512, 1024, 2048, 4096]:
    Q = torch.randn(batch, heads, seq_len, d_k)
    K = torch.randn(batch, heads, seq_len, d_k)
    V = torch.randn(batch, heads, seq_len, d_k)
    scores = torch.matmul(Q, K.transpose(-2, -1))
    mem_mb = scores.nelement() * scores.element_size() / 1e6
    print(f"{seq_len:<12} {f'{seq_len}×{seq_len}':^25} {mem_mb:>10.1f} MB")

9 Vision Transformer (ViT): Attention Beyond Text

Everything so far treated the Transformer as an NLP architecture — but nothing about self-attention (Lesson 52) is actually specific to text. The Vision Transformer (ViT) (Dosovitskiy et al., 2020) applies the exact same encoder block from Section 3 to images, with one clever preprocessing trick to turn pixels into a sequence of tokens.

Patch Embeddings: An Image Is Worth 16×16 Words

A Transformer needs a sequence of vectors, and an image is a grid of pixels — ViT bridges the two by slicing the image into fixed-size patches (typically 16×16 pixels), flattening each patch, and linearly projecting it into the same d_model-dimensional space a word embedding (Lesson 49) would occupy. A 224×224 image cut into 16×16 patches becomes a sequence of 196 "visual tokens" — from there, it's a completely ordinary Transformer encoder.

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, d_model=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2   # 224/16=14 -> 14*14=196 patches
        # A single strided convolution does BOTH the patch-slicing and the
        # linear projection in one operation: each 16x16xC patch becomes one
        # d_model-dimensional vector.
        self.projection = nn.Conv2d(in_channels, d_model, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: (batch, channels, height, width) -> (batch, d_model, 14, 14)
        x = self.projection(x)
        # Flatten spatial dims into a sequence: (batch, d_model, 196) -> (batch, 196, d_model)
        x = x.flatten(2).transpose(1, 2)
        return x   # exactly the shape a text Transformer encoder expects

patch_embed = PatchEmbedding()
image_batch = torch.randn(4, 3, 224, 224)   # 4 images, 3 channels (RGB), 224x224
patches = patch_embed(image_batch)
print(f"Image batch shape:  {image_batch.shape}")
print(f"Patch sequence shape: {patches.shape}")   # (4, 196, 768) -- a "sentence" of 196 "words"

The [CLS] Token and Position Embeddings — Borrowed Directly from BERT

ViT reuses two ideas straight from Lesson 54's BERT: a learnable [CLS] token is prepended to the patch sequence, and after the encoder stack, that token's final representation is used as the whole image's summary vector for classification — exactly how BERT's [CLS] token summarizes a whole sentence. And because self-attention has no inherent sense of position (Section 2), ViT adds a learned positional embedding to every patch, encoding where in the image each patch came from.

import torch
import torch.nn as nn

class ViT(nn.Module):
    def __init__(self, img_size=224, patch_size=16, d_model=768, n_heads=12, n_layers=12, n_classes=1000):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, d_model=d_model)
        n_patches = self.patch_embed.n_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))          # learnable, like BERT's [CLS]
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, d_model))  # +1 for the CLS token

        encoder_layer = nn.TransformerEncoderLayer(d_model, n_heads, dim_feedforward=4*d_model, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, n_layers)      # same encoder block as Section 3
        self.classifier = nn.Linear(d_model, n_classes)

    def forward(self, x):
        batch_size = x.size(0)
        patches = self.patch_embed(x)                                     # (batch, n_patches, d_model)
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        tokens = torch.cat([cls_tokens, patches], dim=1)                  # prepend [CLS]
        tokens = tokens + self.pos_embed                                   # add positional info

        encoded = self.encoder(tokens)
        cls_output = encoded[:, 0]                                         # [CLS] token's final state
        return self.classifier(cls_output)

vit = ViT()
logits = vit(image_batch)
print(f"ViT output shape: {logits.shape}")   # (4, 1000) -- classification logits per image