LLM Atlas/Part 06

Cluster A: Quadratic Attention O(n²)

The O(n²) lineage: sparse patterns, sliding windows, FlashAttention. Keep exact attention, attack the constant factor instead.

Published: 2 August 2026
Reading time: 10 min read
Figures: 4 figures
Equations: 9 equations

Encoder-Only: BERT & Masked Language Modeling

BERT (Devlin et al., 2018) uses Masked Language Modeling (MLM): randomly replace ~15% of input tokens with a [MASK] placeholder, and train the model to predict the original tokens using bidirectional context — both left and right.

Bidirectionality is the key differentiator from GPT: each masked token's prediction can use information from any other position in the sequence. This makes BERT excellent for understanding tasks (classification, extraction, question answering) but unsuitable for autoregressive generation.

A special [CLS] token prepended to every input accumulates a sentence-level representation via its final hidden state — useful for classification.

Key descendants: RoBERTa dropped Next Sentence Prediction and added dynamic masking — matched BERT with better efficiency. ALBERT used cross-layer parameter sharing to reduce memory. DistilBERT used knowledge distillation to produce a 40%-smaller, 60%-faster model retaining 97% of BERT's performance.

MLM loss

L_{MLM} = - i \in masked \sum lo g P (x_{i} ∣ x_{context})

Encoder-Only: BERT & Masked Language Modeling

Decoder-Only: GPT & Causal Language Modeling

GPT-1 (Radford et al., 2018) uses Causal Language Modeling (CLM): predict the next token given all previous tokens. A causal mask sets all future-position attention scores to −∞ before softmax, making those weights exactly 0.

Autoregressive generation: feed a prompt, sample the next token from the output probability distribution, append it to the context, repeat. This loop generates sequences of arbitrary length.

GPT-3 (175B, 2020) demonstrated in-context learning: present a few examples of a task in the prompt, and the model solves new instances without any weight updates. This is qualitatively different from fine-tuning — the 'learning' happens inside the forward pass, via attention computations over the context.

Decoder-only became dominant for generation because: the CLM objective covers any text, scales cleanly with compute, and eliminates the need for paired input-output training data.

CLM objective

L = - t \sum lo g P (x_{t} ∣ x_{< t})

Decoder-Only: GPT & Causal Language Modeling

python

class TransformerBlock(nn.Module):
    def __init__(self):
        self.attn = MultiHeadAttention()
        self.ffn = FeedForward()
        self.ln1 = RMSNorm()
        self.ln2 = RMSNorm()

    def forward(self, x):
        # Pre-normalization architecture
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

Encoder-Decoder: T5 & Span Corruption

T5 (Raffel et al., 2019) unified all NLP tasks as text-to-text: the model always takes a text string in and produces a text string out. 'Classify sentiment: I love this film' → 'positive'. 'Translate to French: Hello' → 'Bonjour'. Even regression becomes text: 'STS-B score: ...' → '3.8'.

Pre-training uses span corruption: randomly mask contiguous spans of tokens, have the decoder reconstruct only the masked spans. This preserves more context in the encoder than BERT-style per-token masking.

Cross-attention bridges encoder and decoder: at each decoder layer, the decoder representations are queries while the encoder's final representations are keys and values. This lets the decoder dynamically 'read' any part of the encoded input while generating output.

BART (Lewis et al., 2019) is a denoising autoencoder — various corruptions (deletion, permutation, text infilling) applied to the input, decoder reconstructs the original. Especially strong for abstractive summarization.

Cross-attention

Attn (Q_{dec}, K_{enc}, V_{enc}) = softmax (\frac{Q _{dec} K _{enc}^{⊤}}{d _{k}}) V_{enc}

Encoder-Decoder: T5 & Span Corruption

Mixture of Experts (MoE)

Core intuition: instead of routing every token through every parameter, selectively activate a small subset of specialized 'expert' sub-networks. This decouples total parameters (knowledge capacity) from active compute (FLOPs per forward pass).

Each MoE layer replaces the FFN with N expert FFNs. A gating network computes routing probabilities for each expert. Sparse MoE activates only the top-k experts per token — the others contribute nothing.

Load balancing is critical: without constraints, the router degenerates — it collapses onto a few favorite experts and ignores the rest. The auxiliary loss penalizes imbalance by minimizing Σ f_i · P_i across experts, where f_i is the fraction of tokens routed to expert i and P_i is the mean routing probability.

DeepSeek-V3 (2024/25): 671B total parameters / 37B active per token. Uses 256 routed + 1 shared expert per layer, 8 activated per token. An auxiliary-loss-free load balancer uses per-expert bias terms instead of the standard loss term. Trained on 14.8T tokens in 2.788M H800 GPU-hours — a remarkably efficient frontier model.

MoE output

y = i \in top-k \sum g_{i} (x) \cdot E_{i} (x)

Gating

g (x) = softmax (TopK (x W_{gate}))

Load-balance loss

L_{aux} = α N i \sum f_{i} P_{i}

Mixture of Experts (MoE)

python

def sparse_moe(x, router, experts, top_k=2):
    # router outputs logits for each expert
    routing_logits = router(x)
    routing_probs = torch.softmax(routing_logits, dim=-1)
    
    # Select top_k experts
    top_probs, top_indices = torch.topk(routing_probs, top_k)
    top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True) # Normalize
    
    # Combine expert outputs
    out = sum(prob * experts[idx](x) for prob, idx in zip(top_probs, top_indices))
    return out

GPT-4 Architecture Notes (2023)

While OpenAI never published the exact architecture, widespread leaks and community consensus paint GPT-4 not as a single dense model, but a massive Mixture of Experts.

Estimates suggest ~1.8 Trillion total parameters, split across 16 experts of ~111B parameters each. During inference, it routes each token to 2 experts, meaning active compute is 'only' ~280B parameters per token. This architectural leap allowed a massive increase in knowledge capacity without an unmanageable increase in inference latency.

Supervised Fine-Tuning (SFT)

Pre-training creates a 'document completer' that might respond to a question by asking more questions. SFT converts it into an assistant.

SFT trains the model on tens of thousands of high-quality, human-written prompt-response pairs. The objective is exactly the same (next-token prediction cross-entropy), but the data distribution is entirely conversational. This teaches the model the format of interaction.

python

def compute_sft_loss(logits, targets, ignore_index=-100):
    # Only compute loss on target tokens, ignore prompt tokens
    loss_fct = nn.CrossEntropyLoss(ignore_index=ignore_index)
    # Shift logits and targets
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = targets[..., 1:].contiguous()
    return loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

RLHF: Reinforcement Learning from Human Feedback

SFT is expensive (requires writing perfect answers) and limited by the human's skill. RLHF (Ouyang et al., 2022) scales alignment by having humans grade answers instead.

Step 1: Train a Reward Model (RM) on human preference data (e.g., 'Response A is better than B'). The RM learns to score text quality. Step 2: Use PPO (Proximal Policy Optimization) to fine-tune the LLM to maximize the RM's score. A KL-divergence penalty ensures the LLM doesn't drift too far from the original SFT model (which would cause it to output 'reward-hacking' gibberish).

RLHF Objective

θ max E [R (x, y)] - β D_{K L} (π_{θ} ∣∣ π_{SFT})

DPO: Direct Preference Optimization (2023)

RLHF is notoriously unstable because PPO involves multiple models acting simultaneously. DPO (Rafailov et al., 2023) eliminates the Reward Model and the RL loop entirely.

Mathematical insight: the RLHF objective can be solved exactly for the optimal policy. DPO reparameterizes the reward in terms of the policy itself. You can train the LLM directly on the human preference pairs using a simple binary cross-entropy loss. It is much more stable, requires less memory, and is now the industry standard (used in LLaMA 3).

DPO Loss

L_{D P O} = - lo g σ (β lo g \frac{π _{θ} ( y _{w} )}{π _{ref} ( y _{w} )} - β lo g \frac{π _{θ} ( y _{l} )}{π _{ref} ( y _{l} )})

Constitutional AI (Anthropic, 2022)

Relying on humans to label harmlessness is difficult because it exposes them to toxic content, and human values are subjective. Constitutional AI replaces human raters with an AI supervisor.

Give the model a 'Constitution' (a list of principles like 'choose the response that is least racist').
Have the model generate responses to toxic prompts, then ask it to critique and revise its own responses based on the constitution.
Train the model on its own revised safe responses (RLAIF: RL from AI Feedback). Claude is built on this principle.

Knowledge Distillation

How do you get a smart 8B model? Train it to mimic a 70B model.

Instead of training the small model (student) on hard labels (one-hot distributions), train it to match the exact output probabilities (soft labels) of a large model (teacher). The soft labels contain 'dark knowledge' — for example, that a dog is 10% likely to be confused with a cat, but 0% likely to be confused with a car. This richer signal allows small models to punch far above their weight class.

Distillation Loss

L = α L_{C E} (y, p_{student}) + (1 - α) T^{2} D_{K L} (p_{teacher}^{τ} ∣∣ p_{student}^{τ})