Large Language Models

Before any neural computation can happen, raw text must become a sequence of discrete integers. The choice of how to split text into tokens is more consequential than it first appears — it affects vocabulary size, sequence length, out-of-vocabulary handling, multilingual coverage, and arithmetic performance.

Character-level tokenization is the simplest approach: each Unicode code point is a token. Vocabulary size is tiny (~256 for bytes, ~140,000 for all Unicode), sequences are long, and the model must learn spelling from scratch. RNNs used this but it was computationally expensive for transformers.

Word-level tokenization splits on whitespace and punctuation. Vocabulary must be enormous to cover morphological variants ('run', 'runs', 'running', 'ran'), and rare or unknown words become out-of-vocabulary tokens — a hard failure.

Byte Pair Encoding (BPE), introduced for NMT and adopted by GPT-2 onward, is the dominant approach. The algorithm:

1. Start with a vocabulary of individual bytes (or characters).
2. Count all adjacent token pairs in the corpus.
3. Merge the most frequent pair into a new token.
4. Repeat until the vocabulary reaches size V (typically 32K–100K for LLMs).

The result is a vocabulary of frequent subword units. Common words ('the', 'of') become single tokens; rare words are split into morphemes or even bytes. GPT-4 uses ~100,277 tokens via tiktoken. The key insight is that BPE is corpus-adaptive: the vocabulary reflects actual frequency distributions in pretraining data.

SentencePiece (used by LLaMA, T5) adds language-agnostic preprocessing and can operate on raw Unicode without pre-tokenization, making it better for non-Latin scripts.

Tokenization has surprising failure modes. 'strawberry' contains three r's but GPT-4 famously miscounts them — because 'strawberry' is a single token and character-level access requires the model to have learned the internal structure of token strings. Numbers like '127385' may be split into ['127', '385'] or ['1', '27', '385'] depending on the corpus — a known source of arithmetic errors. These are not bugs but consequences of the tokenization boundary.

# Minimal BPE training loop (simplified)
from collections import Counter

def get_pairs(vocab):
    """Count adjacent pairs across all word-frequency entries."""
    pairs = Counter()
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i+1])] += freq
    return pairs

def merge_vocab(pair, vocab):
    """Merge the most frequent pair into a single token."""
    merged = {}
    bigram = ' '.join(pair)
    replacement = ''.join(pair)
    for word, freq in vocab.items():
        new_word = word.replace(bigram, replacement)
        merged[new_word] = freq
    return merged

# vocab: {' l o w </w>': 5, ' l o w e r </w>': 2, ...}
vocab = {' l o w </w>': 5, ' l o w e r </w>': 2,
         ' n e w e r </w>': 6, ' w i d e r </w>': 3}

for _ in range(10):  # 10 merge operations
    pairs = get_pairs(vocab)
    if not pairs:
        break
    best = max(pairs, key=pairs.get)
    vocab = merge_vocab(best, vocab)
    print(f'Merged: {best}')
# Merged: ('e', 'r') -> 'er', then ('er', '</w>') -> 'er</w>', etc.

The central idea of modern LLMs is deceptively simple: given all previous tokens, predict the next one. This is autoregressive language modeling, and the training objective is cross-entropy loss over the true next token:

$$\mathcal{L} = -\sum_{t=1}^{T} \log P_{\theta}(x_t \mid x_1, \ldots, x_{t-1})$$

What makes this powerful is that it is self-supervised: the labels are the data itself, shifted by one position. You need no human annotation. You can use every token in every document as a training example. This is why pretraining can exploit internet-scale corpora — Common Crawl (petabytes of web text), Books3, GitHub, arXiv, Wikipedia — without any labeling cost.

Why does predicting the next word require understanding? Consider the sentence: 'The surgeon who operated on the patient washed ___ hands.' To predict 'her' or 'his' correctly, the model must resolve a coreference across a complex relative clause. To predict 'The capital of France is ___', the model must store and retrieve a world fact. To predict the next line of Python code, it must track variable scope, types, and intent. The task is deceptively comprehensive — it is a proxy for understanding.

The pretraining corpus is massive and heterogeneous. GPT-3 trained on ~300B tokens; LLaMA-2 on 2 trillion tokens; modern frontier models on ~10-15 trillion tokens. The corpus is typically filtered (remove boilerplate, deduplication, quality scoring via a fastText classifier trained on Wikipedia vs. random web), weighted (upsampling books and code), and shuffled at document granularity.

Pretraining is computationally staggering. GPT-3 (175B parameters) required approximately 3.14×10²³ FLOPs — equivalent to roughly 3,640 petaflop-days. This is done on clusters of thousands of H100 GPUs using a combination of data parallelism (each GPU processes a different microbatch, gradients are averaged via AllReduce), tensor parallelism (split individual weight matrices across GPUs), and pipeline parallelism (split model layers across GPUs, with microbatching to hide bubble overhead). Mixed-precision training (BF16 forward, FP32 master weights) and gradient checkpointing are standard.

# Illustrative pretraining loop (single GPU, educational)
import torch, torch.nn as nn

def pretrain_step(model, batch_tokens, optimizer):
    """
    batch_tokens: LongTensor [B, T+1]
    Returns scalar loss.
    """
    # Inputs are tokens 0..T-1, targets are tokens 1..T
    x = batch_tokens[:, :-1]   # [B, T]
    y = batch_tokens[:, 1:]    # [B, T]

    logits = model(x)           # [B, T, vocab_size]
    loss = nn.functional.cross_entropy(
        logits.reshape(-1, logits.size(-1)),  # [B*T, V]
        y.reshape(-1)                          # [B*T]
    )
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # stability
    optimizer.step()
    return loss.item()

# Perplexity = exp(average cross-entropy loss)
# A well-trained GPT-2 achieves ~29 ppl on Penn Treebank
# GPT-4 achieves single-digit ppl on many benchmarks

The GPT family is built on the decoder-only transformer, a direct descendant of the original 'Attention Is All You Need' architecture but stripped of the encoder and cross-attention blocks. Understanding this architecture precisely is essential.

Input representation. Token indices are embedded via a learned matrix $W_E \in \mathbb{R}^{V \times d}$ into dense vectors of dimension $d$ (768 for GPT-2 small, 12288 for GPT-4). Positional information is added either via learned absolute positional embeddings (GPT-2, GPT-3) or via Rotary Position Embedding (RoPE) (LLaMA, GPT-NeoX), which encodes position by rotating query and key vectors, giving better extrapolation to sequence lengths beyond training.

Causal self-attention. Each transformer block begins with multi-head self-attention. With $h$ heads, each head computes: $$\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$ where $Q = XW_Q$, $K = XW_K$, $V = XW_V$. The crucial constraint for autoregressive generation is causal masking: position $t$ may only attend to positions $\leq t$. This is implemented by adding $-\infty$ to the upper triangle of $QK^T/\sqrt{d_k}$ before the softmax, so those positions receive zero attention weight. Causal masking is what allows training on all positions simultaneously (teacher forcing) while maintaining the proper left-to-right dependency.

Feed-forward network. After attention, each token passes through a two-layer MLP with a nonlinearity: $$\text{FFN}(x) = W_2 \cdot \text{GELU}(W_1 x + b_1) + b_2$$ The FFN hidden dimension is typically $4d$. Modern variants use SwiGLU (LLaMA): $\text{SwiGLU}(x, W, V, b, c) = \text{Swish}(xW+b) \odot (xV+c)$, which improves performance. The FFN is where most factual knowledge is believed to reside — ablation studies show that patching the FFN layers changes factual outputs more than attention layers.

Layer normalization. GPT-2 and later use Pre-LayerNorm (norm before attention and FFN, rather than after), which stabilizes training at large scale. LLaMA replaces LayerNorm with RMSNorm (no mean centering, just scale by RMS), which is faster.

The residual stream. Every block adds its output to the residual stream: $$x \leftarrow x + \text{Attention}(\text{LN}(x)); \quad x \leftarrow x + \text{FFN}(\text{LN}(x))$$ The residual stream picture (from Anthropic mechanistic interpretability research) is a powerful mental model: think of $x \in \mathbb{R}^d$ as a shared workspace that each attention head and FFN block reads from and writes to in parallel.

Output. After all $L$ blocks, a final LayerNorm and a linear projection $W_U \in \mathbb{R}^{d \times V}$ (often tied to $W_E^T$) produce logits over the vocabulary. Softmax gives a probability distribution for the next token.

GPT-2 (1.5B) uses $d=1600$, $h=25$, $L=48$. GPT-3 (175B) uses $d=12288$, $h=96$, $L=96$.

# Minimal GPT block in PyTorch
import torch, torch.nn as nn, math

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, max_seq_len=2048):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        # Causal mask: lower-triangular ones
        self.register_buffer('mask',
            torch.tril(torch.ones(max_seq_len, max_seq_len)))

    def forward(self, x):           # x: [B, T, D]
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)
        # Split into heads: [B, H, T, d_k]
        def split_heads(t):
            return t.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        q, k, v = split_heads(q), split_heads(k), split_heads(v)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
        scores = scores.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
        attn = scores.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
        return self.proj(out)

class GPTBlock(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model)
        )
    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # residual + attention
        x = x + self.ffn(self.ln2(x))    # residual + FFN
        return x

One of the most striking empirical discoveries in deep learning is that the test loss of language models follows precise power laws in model size, dataset size, and compute. This was formalized in the Kaplan et al. (2020) paper (OpenAI) and refined in Hoffmann et al. (2022) (DeepMind, 'Chinchilla').

The Kaplan scaling laws. For a fixed compute budget $C$ (FLOPs), model size $N$ (parameters), and dataset size $D$ (tokens), the loss approximately satisfies: $$L(N) \approx \left(\frac{N_c}{N}\right)^{\alpha_N}, \quad L(D) \approx \left(\frac{D_c}{D}\right)^{\alpha_D}, \quad L(C) \approx \left(\frac{C_c}{C}\right)^{\alpha_C}$$ where $\alpha_N \approx 0.076$, $\alpha_D \approx 0.095$, $\alpha_C \approx 0.050$ are measured empirically. These are independent of architecture details (within a class), optimizer, and data domain — a surprising universality.

Kaplan et al. concluded: given a fixed compute budget, you should scale model size much faster than data size — train a large model on fewer tokens (roughly $D \propto N^{0.74}$).

The Chinchilla correction. Hoffmann et al. (2022) reran scaling experiments more carefully and found a different optimal allocation. Their key result: to minimize loss for a given compute budget $C \approx 6ND$ (the approximation that one forward+backward pass costs roughly 6 FLOPs per parameter per token), the optimal balance is: $$N_{opt} \propto C^{0.5}, \quad D_{opt} \propto C^{0.5}$$ That is, model size and token count should scale equally. The practical implication: GPT-3 (175B params, 300B tokens) was undertrained by a factor of ~5 in data. Chinchilla (70B params, 1.4T tokens) matches GPT-3's performance at 4x fewer parameters, making it cheaper at inference time.

This finding reshaped the field. LLaMA-1/2/3 use the Chinchilla-optimal recipe but deliberately overtrain (well past compute-optimal) to minimize inference cost — compute-optimal training is different from deployment-optimal training.

Emergent abilities. Some capabilities appear to 'emerge' discontinuously as scale increases — BIG-Bench tasks that score near zero below a threshold and then jump sharply. The interpretation is contested: some researchers argue emergence is an artifact of evaluation metrics (thresholded vs. continuous). Regardless, scaling laws give practitioners a powerful tool: you can run small-scale ablations, fit power laws, and extrapolate to predict whether a 10x larger model is worth training.

The irreducible loss. Even with infinite compute, loss has a floor $L_\infty > 0$ corresponding to genuine ambiguity in language (different valid continuations). The actual parameterization often used is: $$L(N, D) = E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}}$$ where $E$ is the irreducible entropy of language.

# Chinchilla compute-optimal allocation
# Given a compute budget C (in FLOPs), find optimal N (params) and D (tokens)
# C ≈ 6 * N * D  (approximate, ignoring attention FLOPs)

import numpy as np

def chinchilla_optimal(C_flops):
    """
    Hoffmann et al. 2022: N_opt ~ 0.2 * C^0.5, D_opt ~ 8.3 * C^0.5
    (constants fitted from their Table A3)
    """
    N_opt = 0.2 * (C_flops ** 0.5)
    D_opt = C_flops / (6 * N_opt)  # from C = 6*N*D
    return N_opt, D_opt

# GPT-3 training compute: ~3.14e23 FLOPs
C_gpt3 = 3.14e23
N_opt, D_opt = chinchilla_optimal(C_gpt3)
print(f'Chinchilla-optimal for GPT-3 compute budget:')
print(f'  N = {N_opt/1e9:.1f}B parameters')
print(f'  D = {D_opt/1e9:.0f}B tokens')
# N ≈ 3.5B params, D ≈ 148B tokens
# vs GPT-3 actual: 175B params, 300B tokens  <- undertrained on data

# Power-law loss extrapolation
def kaplan_loss(N, A=406.4, alpha=0.34):
    """Simplified single-variable power law (illustrative)."""
    return A * N**(-alpha)

for n_params in [1e8, 1e9, 1e10, 1e11]:
    print(f'{n_params/1e9:.2f}B params -> predicted loss: {kaplan_loss(n_params):.3f}')

A pretrained LLM is a remarkable artifact, but it is not yet an assistant. Its output is a probability distribution over continuations of arbitrary text — it will complete 'The capital of France is' with 'Paris', but it will also complete 'How do I make a bomb? Here is a step-by-step' in an equally helpful tone. Alignment is the process of shaping the model's behavior to be helpful, harmless, and honest.

Supervised Fine-Tuning (SFT). The simplest approach: collect a dataset of (instruction, ideal response) pairs, and continue training the pretrained model on these with the same cross-entropy objective, but only computing loss on the response tokens (not the instruction). This teaches the model the format and style of helpfulness. OpenAI used SFT as the first stage of InstructGPT; LLaMA-2-Chat began with SFT on ~27,500 curated instruction-response pairs.

Parameter-Efficient Fine-Tuning (PEFT). Full fine-tuning of a 70B model requires as much GPU memory as pretraining — often infeasible. Several PEFT methods achieve near-full-fine-tune quality at a fraction of the cost:

• LoRA (Low-Rank Adaptation): Freeze all pretrained weights. For each target weight matrix $W_0 \in \mathbb{R}^{d \times k}$, inject a low-rank decomposition $\Delta W = BA$ where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$, $r \ll \min(d,k)$. During forward pass: $h = W_0 x + \frac{\alpha}{r} BAx$. Only $A$ and $B$ are trained, reducing trainable parameters by 10,000x. Typical $r \in {4, 8, 16, 64}$.

• QLoRA: LoRA on top of a 4-bit quantized base model (NF4 quantization with double quantization), enabling fine-tuning of 65B models on a single 48GB GPU.

• Prefix Tuning / Prompt Tuning: Prepend trainable 'soft prompt' vectors to the input sequence. Only these vectors are trained. Works for simple tasks but underperforms LoRA for complex instruction following.

Catastrophic forgetting is a real risk in fine-tuning: the model loses pretrained capabilities as it adapts to a narrow task. Mitigations include mixing fine-tuning data with pretraining data (replays), low learning rates, and PEFT methods that leave most weights frozen.

# LoRA applied to a linear layer (core idea)
import torch, torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base_layer: nn.Linear, rank: int, alpha: float = 1.0):
        super().__init__()
        d_out, d_in = base_layer.weight.shape
        self.base = base_layer
        self.base.weight.requires_grad_(False)  # freeze base weights
        # Low-rank factors: A initialized with Kaiming, B with zeros
        self.lora_A = nn.Parameter(torch.empty(rank, d_in))
        self.lora_B = nn.Parameter(torch.zeros(d_out, rank))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.scale = alpha / rank

    def forward(self, x):
        # Base path (frozen) + low-rank update (trainable)
        return self.base(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scale

# Result: for a 4096x4096 attention weight matrix (16M params),
# LoRA rank-8 adds only 2*4096*8 = 65,536 trainable params — a 250x reduction.

Supervised fine-tuning teaches the model format, but it cannot directly optimize for nuanced human preferences — helpfulness, tone, refusal of harmful requests — which are hard to specify as a fixed dataset. RLHF, introduced by OpenAI and formalized in the InstructGPT paper (Ouyang et al., 2022), addresses this.

The pipeline has three stages:

Stage 1: SFT. Fine-tune the pretrained model on a small set of high-quality (prompt, response) demonstrations to get an initial policy $\pi_{SFT}$.

Stage 2: Reward Model (RM) training. For each prompt, sample multiple responses from $\pi_{SFT}$ and have human labelers rank them by quality. Given a comparison $(y_w \succ y_l \mid x)$ (response $y_w$ preferred over $y_l$ for prompt $x$), train a scalar reward model $r_\phi$ to maximize: $$\mathcal{L}{RM} = -\mathbb{E}{(x,y_w,y_l)}\left[\log \sigma(r_\phi(x, y_w) - r_\phi(x, y_l))\right]$$ This is a Bradley-Terry pairwise preference model. The RM takes a (prompt, response) pair and outputs a scalar reward.

Stage 3: RL fine-tuning with PPO. Use the RM as a reward signal to fine-tune $\pi_{SFT}$ via Proximal Policy Optimization. The RL objective: $$\max_{\pi_\theta} \mathbb{E}{x \sim D,, y \sim \pi\theta(\cdot|x)}\left[r_\phi(x, y)\right] - \beta \cdot \text{KL}\left[\pi_\theta(\cdot|x) |, \pi_{\text{ref}}(\cdot|x)\right]$$ The KL penalty prevents the policy from straying too far from the SFT reference model (which would cause reward hacking: generating gibberish that scores highly on a flawed RM). $\beta$ is a hyperparameter controlling the tradeoff.

PPO-clip updates the policy in small steps, clipping the probability ratio to prevent large policy updates that destabilize training: $$L^{CLIP}(\theta) = \mathbb{E}_t\left[\min\left(r_t(\theta)\hat{A}_t,, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t\right)\right]$$

RLHF challenges. (1) Reward hacking: the policy finds ways to score highly on the RM that don't correspond to actual quality (e.g., very long responses, sycophancy). (2) Distribution shift: the RM is trained on $\pi_{SFT}$ outputs but evaluated on $\pi_\theta$ outputs that can diverge significantly. (3) Human labeler disagreement: preferences are noisy and culturally variable.

Direct Preference Optimization (DPO), proposed by Rafailov et al. (2023), eliminates the separate RM and RL step. It shows that the RLHF objective has a closed-form optimal policy, leading to a supervised loss directly on preference pairs: $$\mathcal{L}{DPO} = -\mathbb{E}{(x,y_w,y_l)}\left[\log \sigma\left(\beta \log \frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)}\right)\right]$$ DPO is simpler to implement, stable to train, and competitive with PPO-based RLHF.

# DPO loss (simplified, illustrative)
import torch, torch.nn.functional as F

def dpo_loss(policy_logps_w, policy_logps_l,
            ref_logps_w, ref_logps_l, beta=0.1):
    """
    policy_logps_w: log P_theta(y_w | x)  -- sum of token log-probs
    policy_logps_l: log P_theta(y_l | x)
    ref_logps_w:    log P_ref(y_w | x)
    ref_logps_l:    log P_ref(y_l | x)
    beta:           KL penalty coefficient
    """
    # Log-ratio of policy to reference for each response
    ratio_w = policy_logps_w - ref_logps_w  # log [pi/ref] for winner
    ratio_l = policy_logps_l - ref_logps_l  # log [pi/ref] for loser

    # Prefer winner: maximize (ratio_w - ratio_l)
    logits = beta * (ratio_w - ratio_l)
    loss = -F.logsigmoid(logits).mean()
    return loss

# Interpretation:
# - If policy assigns relatively higher prob to y_w vs ref -> ratio_w increases
# - If policy assigns relatively lower prob to y_l vs ref  -> ratio_l decreases
# - Both effects increase logits -> decrease loss -> correct direction

Perhaps the most surprising capability of large LLMs is in-context learning (ICL): the ability to perform a new task when given just a few examples (shots) in the prompt, with no weight updates. GPT-3 demonstrated this dramatically — a model that was never trained on translation can translate French to English after seeing three examples in the prompt.

The phenomenon. Given a prompt of the form:

Input: "chat" -> Output: "cat"
Input: "chien" -> Output: "dog"
Input: "livre" -> Output:

GPT-3 produces 'book' with high probability, having inferred the translation task from context alone. This is called few-shot prompting (1-shot, 5-shot, etc.). With no examples at all — just a natural language task description — it is zero-shot prompting.

Why does ICL work? This is an active research question with several competing hypotheses:

1. Meta-learning view (Brown et al., 2020): During pretraining, the model sees many different implicit tasks (different styles, domains, formats). It learns an implicit algorithm for inferring and executing tasks from patterns. ICL is meta-learning at inference time.

2. Bayesian posterior update view (Xie et al., 2021): Pretraining instills a prior over latent tasks. Given few-shot examples, the model implicitly performs Bayesian inference over which task is intended, then executes the most probable one. This explains why more shots generally help.

3. Gradient descent in the residual stream (Akyürek et al., 2022; Von Oswald et al., 2023): Linear self-attention heads can implement gradient descent steps. ICL may literally be executing a form of implicit gradient descent in the forward pass activations, without any weight updates.

Chain-of-Thought (CoT) prompting (Wei et al., 2022) extends ICL to reasoning. By showing examples with intermediate reasoning steps ('Let's think step by step...'), the model generates reasoning chains before the final answer, dramatically improving performance on math and logic tasks. CoT works because generating the intermediate steps forces the model to allocate more computation (more tokens) to the problem, and because it was pretrained on text that contained such reasoning chains.

Prompt sensitivity is a known failure mode: ICL performance depends strongly on the order of examples, the exact wording of the task description, and even the formatting. A prompt that works for GPT-4 may fail on LLaMA-3. This makes ICL less reliable than fine-tuning for production systems.

Context length and retrieval. Modern LLMs support 128K–1M token context windows. Long context enables 'many-shot' prompting — providing hundreds of examples, full documents, or entire codebases in context. But 'lost in the middle' experiments show that models attend more to the beginning and end of long contexts, with degraded performance in the middle.

# Few-shot prompting structure (template, not an API call)
few_shot_prompt = """
Classify the sentiment of each review as POSITIVE or NEGATIVE.

Review: The food was absolutely delicious and the service was impeccable.
Sentiment: POSITIVE

Review: I waited 45 minutes and the order was wrong.
Sentiment: NEGATIVE

Review: Decent enough for the price, nothing special.
Sentiment: NEGATIVE

Review: Best pizza I've had in years, will definitely return!
Sentiment:"""
# A well-trained LLM continues: POSITIVE

# Chain-of-Thought template (math)
cot_prompt = """
Q: Roger has 5 tennis balls. He buys 2 more cans of 3 balls each.
How many tennis balls does he have?
A: Roger starts with 5. 2 cans x 3 balls = 6 new balls. 5 + 6 = 11. Answer: 11

Q: There are 15 trees. After workers plant trees, there are 21 trees.
How many trees did the workers plant?
A:"""
# CoT causes the model to generate: 21 - 15 = 6. Answer: 6

The trained model defines a conditional distribution $P(x_t | x_{<t})$, but how we sample from it dramatically affects output quality, diversity, and coherence. This is the decoding strategy.

Greedy decoding. At each step, pick the most probable token: $x_t = \arg\max P(x_t | x_{<t})$. Fast, deterministic, but produces repetitive, mode-seeking output. A famous failure: greedy decoding on GPT-2 famously produces looping text (the model gets stuck in probability-maximizing cycles).

Beam search. Maintain the top-$k$ most probable partial sequences ('beams') at each step, expanding each by all vocabulary tokens and keeping the top-$k$. Returns the globally most probable sequence up to length $T$. Better than greedy but still prone to degeneracy (all beams converging to near-identical, bland outputs). Widely used in translation but less so for open-ended generation.

Temperature sampling. Divide logits by temperature $\tau$ before softmax: $P_\tau(x) \propto \exp(z_x / \tau)$. $\tau \to 0$ approaches greedy; $\tau \to \infty$ approaches uniform. Typical creative writing uses $\tau \in [0.7, 1.2]$; factual tasks use $\tau < 0.5$.

Top-k sampling. Restrict sampling to the $k$ most probable tokens (zero out the rest, renormalize). Prevents sampling extremely unlikely tokens. $k = 50$ is a common default.

Nucleus (Top-p) sampling (Holtzman et al., 2020). Instead of a fixed $k$, include the smallest set of tokens whose cumulative probability exceeds $p$. $p = 0.9$ means: take the top tokens that together account for 90% of the probability mass. Adaptive: when the model is confident (peaked distribution), fewer tokens are included; when uncertain (flat distribution), more are included. This is the dominant approach in LLM APIs.

Repetition penalties. To combat the tendency to repeat phrases, many systems apply a multiplicative penalty to the logits of tokens that have appeared recently. Min-P sampling (2024) is a newer variant that sets a dynamic floor on token probabilities, avoiding both truncation of coherent completions and sampling of incoherent ones.

Speculative decoding. A technique for fast inference: use a small 'draft' model to generate $k$ tokens speculatively, then verify all $k$ with one forward pass of the large model (using the fact that verification is parallelizable). Achieves 2-3x speedup with no change in output distribution.

import torch

def top_p_sample(logits, p=0.9, temperature=1.0):
    """
    logits: [vocab_size] unnormalized log-probabilities
    Returns: sampled token index
    """
    # Apply temperature
    logits = logits / temperature

    # Sort tokens by descending probability
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)

    # Remove tokens beyond the nucleus (cumprob > p)
    # Shift right so that the token pushing over p threshold is included
    sorted_remove = cumulative_probs - torch.softmax(sorted_logits, dim=-1) > p
    sorted_logits[sorted_remove] = float('-inf')

    # Scatter back to original indexing and sample
    logits_filtered = logits.scatter(0, sorted_indices, sorted_logits)
    token = torch.multinomial(torch.softmax(logits_filtered, dim=-1), 1)
    return token.item()

Generating a 1000-token response with a 7B-parameter model is not trivially fast. Understanding the inference bottleneck and the KV cache optimization is essential for anyone deploying LLMs.

The naive cost. Generating token $t$ requires computing attention over all $t-1$ preceding tokens. Without caching, generating $T$ tokens requires $O(T^2)$ attention computations total — quadratic in sequence length.

The KV cache. In causal self-attention, the key ($K$) and value ($V$) vectors for all previous tokens are the same at every generation step — because they are computed from the same frozen prefix and the model weights don't change. We can compute and store them once, then at each new step only compute the $Q$, $K$, $V$ for the single new token and append $K, V$ to the cache.

Concretely: the KV cache stores $L \times 2 \times T \times d_k$ floating-point tensors (2 for K and V, L layers). For a 7B model (32 layers, $d_k = 128$, FP16), caching 4096 tokens requires approximately $32 \times 2 \times 4096 \times 128 \times 2\text{ bytes} = 67$MB — manageable. For a 70B model with 128K context, this grows to ~320GB, which exceeds GPU memory and motivates paged attention (vLLM) and sliding window attention (Mistral).

Prefill vs. decode. Generation has two phases:

• Prefill: process the entire prompt in parallel (all tokens at once, like training). Compute-bound (GPU utilization ~100%).
• Decode: generate one token at a time, each step doing a small matmul against the KV cache. Memory-bandwidth-bound (the bottleneck is reading model weights from HBM, not arithmetic). This is why LLMs are fast at processing long prompts but slow at generating long outputs.

Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce KV cache size by sharing K and V heads across multiple Q heads (GQA groups Q heads, MQA uses a single K/V head for all Q heads). LLaMA-3, Mistral, and Gemma use GQA. This shrinks the KV cache and reduces memory bandwidth, enabling longer contexts and higher batch sizes.

# KV cache simulation (single layer, single head, illustrative)
import torch

class CachedAttention:
    def __init__(self):
        self.k_cache = []  # list of key vectors
        self.v_cache = []  # list of value vectors

    def prefill(self, Q, K, V):  # [T, d_k]
        """Process entire prompt at once, populate cache."""
        self.k_cache = list(K.unbind(0))  # store each token's K
        self.v_cache = list(V.unbind(0))
        # Standard attention over full sequence
        return self._attend(Q, K, V)

    def decode_one(self, q, k, v):  # [d_k] vectors for new token
        """Generate one new token using cached K, V."""
        self.k_cache.append(k)
        self.v_cache.append(v)
        K_full = torch.stack(self.k_cache)  # [T+1, d_k]
        V_full = torch.stack(self.v_cache)  # [T+1, d_k]
        # Q is just the new token: [1, d_k]
        return self._attend(q.unsqueeze(0), K_full, V_full)

    def _attend(self, Q, K, V):
        import math
        d_k = Q.shape[-1]
        scores = (Q @ K.T) / math.sqrt(d_k)
        weights = scores.softmax(dim=-1)
        return weights @ V

# Key insight: decode_one does O(T) work per step (matmul with growing K,V)
# but avoids recomputing K,V for previous tokens -> 2x-3x speedup in practice

We can describe how to train an LLM; the deeper question is what it learns. Mechanistic interpretability is the young field that tries to reverse-engineer the algorithms implemented by a neural network's weights.

Circuits and features. Collins, Elhage, and colleagues at Anthropic pioneered the 'circuits' framework. The key claim: within a transformer, small subnetworks (circuits) — a few attention heads and MLP neurons connected via the residual stream — implement specific algorithms. Examples discovered:

• Induction heads (a 2-head circuit in GPT-2): given a token $A$ followed later by $B$, then $A$ again, predict $B$. This implements 'copy the completion of the previous occurrence of this prefix' — a form of in-context pattern completion. Forming around layer 1-2, induction heads are believed to underlie much of in-context learning.

• Indirect object identification (Wang et al., 2022): in 'John gave Mary the book. Then John told ___', the correct answer is 'Mary'. A specific circuit of ~26 attention heads was identified that implements the full algorithm (name mover heads, duplicate token heads, S-inhibition heads).

• Superposition hypothesis: models represent far more 'features' than they have neurons by embedding them as near-orthogonal directions in activation space, exploiting the fact that most features are rarely active simultaneously (sparsity). Sparse Autoencoders (SAEs) are used to disentangle these superposed features.

Factual knowledge in MLPs. Geva et al. (2021) found that MLP layers act as 'key-value memories': each row of $W_1$ is a 'key' that activates for certain input patterns; the corresponding column of $W_2$ is a 'value' that distributes over the residual stream, promoting certain output tokens. ROME and MEMIT (Meng et al.) use this insight to edit specific facts directly in model weights.

Limitations. Mechanistic interpretability has made impressive progress on small models (GPT-2) and specific circuits, but scaling it to frontier-sized models remains an open challenge. Most model behaviors emerge from the collective action of thousands of circuits, not clean modular components.