Lecture 8 — Beyond Next-Token Prediction

Project connection. Project Step 8 is the capstone. Its scale/fine-tune/LoRA, preference-optimization, and mechanistic-interpretability tracks correspond to Sections 1—2, 3, and 4 of this lecture.

Chapter overview. Pretraining estimates a conditional distribution of text. Fine-tuning changes the data distribution, LoRA restricts the rank of the parameter update, and preference optimization changes the objective itself. The KL-regularized reward objective has an exact Gibbs-form optimum, from which direct preference optimization follows. Mechanistic interpretability asks a different question: which algorithms are already encoded in the trained weights?

0. Base models and policies

Let $x$ denote a prompt and $y$ a finite response sequence. A conditional distribution $\pi(y\mid x)$ is called a policy in reinforcement-learning terminology. Let $\pi_{\mathrm{ref}}$ denote a pretrained or instruction- tuned reference policy.

A base language model is trained by next-token cross-entropy. This objective rewards accurate continuation of its training texts; instruction following or human preference enters only if such behavior is present in the data or in a later objective.

1. Fine-tuning

Let $\theta_0$ be pretrained parameters and let $\mathcal D'$ be a target dataset. Write $\mathcal L_{\mathcal D'}(\theta)$ for its average next-token cross-entropy.

Definition 1.1 (fine-tuning). Fine-tuning initializes at $\theta_0$ and continues gradient-based minimization of $\mathcal L_{\mathcal D'}(\theta)$. Instruction tuning is fine-tuning on prompt—response pairs, usually masking the prompt positions so that the loss is evaluated only on response tokens.

Remark 1.2 (transfer). Pretraining supplies features and linguistic regularities before target-domain optimization begins. Fine-tuning can therefore adapt with less target data than training from a random initialization. The magnitude of this benefit is empirical and depends on domain similarity.

Remark 1.3 (catastrophic forgetting). A narrow target loss contains no term preserving performance on the pretraining distribution. Degradation away from the target domain is called catastrophic forgetting. Mixing old data into the target set or constraining the update can reduce it.

2. LoRA and low-rank updates

The rank of a matrix is the dimension of its column space. If $B\in\mathbb R^{d\times r}$ and $A\in\mathbb R^{r\times k}$, then $\operatorname{rank}(BA)\leq r$.

Definition 2.1 (LoRA). Let $W_0\in\mathbb R^{d\times k}$ be a frozen pretrained matrix. Low-rank adaptation (LoRA) parametrizes the adapted matrix as

$$W=W_0+BA, \qquad B\in\mathbb R^{d\times r}, \quad A\in\mathbb R^{r\times k}, \quad 1\leq r<\min\{d,k\},$$

Here $r$ is the adapter rank and is chosen much smaller than $\min\{d,k\}$. Only $A$ and $B$ are trained. A common initialization takes $A$ random and $B=0$, so the initial adapted matrix is exactly $W_0$.

The factorization introduces $r(d+k)$ trainable scalars rather than $dk$. For $d=k=128$ and $r=4$, the counts are 1,024 and 16,384, respectively. After training, $BA$ may be added to $W_0$, so the merged matrix has the same inference shape as the original. Many task-specific adapters can also share one frozen base model, storing only their small factor pairs.

Remark 2.2 (the low-rank hypothesis). LoRA does not assume that the pretrained matrix $W_0$ has low rank. It assumes that the task-specific update $\Delta W=W-W_0$ can be approximated at low rank. A sweep over $r$ tests this hypothesis: saturation of target performance estimates the effective rank needed for that task and layer set.

Geometric aside 2.3. The matrices of rank at most $r$ form a determinantal variety. Its rank-exactly-$r$ part has dimension $r(d+k-r)$, whereas the factor pair $(B,A)$ has $r(d+k)$ coordinates. The surplus $r^2$ reflects the invariance $(B,A)\mapsto(BR,R^{-1}A)$ for invertible $R\in\mathbb R^{r\times r}$. The variety is singular—not locally a smooth manifold of the expected dimension—on lower-rank parts; optimizing the factorization replaces this constrained matrix geometry by an unconstrained but non-identifiable parametrization.

3. KL-regularized preference optimization

Fix a prompt $x$. Let $\mathcal Y_x$ be a finite response set, or a finite truncation of the sequence space. Assume $\pi_{\mathrm{ref}}(y\mid x)>0$ for every $y\in\mathcal Y_x$. Let $r(x,y)\in\mathbb R$ be a finite reward, and let $\beta>0$ be a regularization parameter. For distributions $p,q$ on $\mathcal Y_x$, recall

$$D_{\mathrm{KL}}(p\|q) =\sum_{y\in\mathcal Y_x}p(y)\log\frac{p(y)}{q(y)}.$$

Reinforcement learning from human feedback (RLHF) uses human preference data to construct or optimize a reward-guided policy. In this lecture the reward may also be programmatic, as in Project Track B.

Definition 3.1 (KL-regularized RLHF objective). For a prompt distribution $\mathcal D$, define

$$J(\pi) =\mathbb E_{x\sim\mathcal D}\left[ \mathbb E_{y\sim\pi(\cdot\mid x)}[r(x,y)] -\beta D_{\mathrm{KL}}\bigl( \pi(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x) \bigr) \right].$$

The first term rewards preferred responses. The second penalizes departure from the reference policy.

Theorem 3.2 (closed-form optimum). For every prompt $x$, the unique maximizer of the conditional objective is

$$\pi^*(y\mid x) =\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x) \exp\!\left(\frac{r(x,y)}{\beta}\right),$$

where

$$Z(x)=\sum_{y\in\mathcal Y_x} \pi_{\mathrm{ref}}(y\mid x)e^{r(x,y)/\beta}.$$

Consequently, maximizing each conditional also maximizes $J(\pi)$.

Proof. Fix $x$ and abbreviate conditional arguments. The proposed $\pi^*$ is a strictly positive normalized distribution. For any $\pi$,

$$\begin{aligned} D_{\mathrm{KL}}(\pi\|\pi^*) &=\sum_y\pi(y)\log\frac{\pi(y)}{\pi_{\mathrm{ref}}(y)} -\frac1\beta\sum_y\pi(y)r(y)+\log Z. \end{aligned}$$

Rearranging gives

$$\sum_y\pi(y)r(y) -\beta D_{\mathrm{KL}}(\pi\|\pi_{\mathrm{ref}}) =\beta\left(\log Z-D_{\mathrm{KL}}(\pi\|\pi^*)\right).$$

Gibbs’ inequality makes the right side at most $\beta\log Z$, with equality exactly when $\pi=\pi^*$. $\square$

Statistical-mechanics aside 3.3. The normalizer $Z$ is a partition function, and the optimal regularized value is $\beta\log Z$. If energy is defined as $E(y)=-r(y)$ and $\beta$ is read as temperature, then $F=-\beta\log Z$ is the Helmholtz free energy; the optimized objective is $-F$.

Remark 3.4 (exponential tilting). The optimum multiplies the reference measure by $e^{r/\beta}$ and renormalizes. This is the same Gibbs tilting seen in attention and temperature sampling. As $\beta\to\infty$, $\pi^*\to\pi_{\mathrm{ref}}$. As $\beta\to0^+$, mass concentrates on the reward-maximizing responses that have reference support; among tied maximizers, their relative reference probabilities are retained.

Remark 3.5 (reward hacking and alignment tax). Small $\beta$ permits a large distributional move and can exploit defects in a learned or programmatic reward. Large $\beta$ preserves the reference policy but limits reward improvement. The resulting reward—KL frontier makes this tradeoff measurable. Loss of useful base-model behavior under alignment is often called an alignment tax. In best-of-$n$ decoding, one draws $n$ candidate responses and returns the one with largest reward; it is a simple way to expose both the reward gain and possible reward hacking.

The partition function $Z(x)$ is generally intractable because the response space is enormous. Classical RLHF learns a reward from preference pairs and uses Proximal Policy Optimization (PPO), a clipped policy-gradient method. A policy gradient differentiates expected reward with respect to policy parameters. A rollout is a response sampled from a policy; on-policy means rollout data come from the policy currently being optimized. Direct preference optimization avoids explicit evaluation of $Z(x)$.

Define the logistic sigmoid

$$\sigma(t)=\frac1{1+e^{-t}}.$$

In the Bradley—Terry preference model, the probability that response $y_w$ is preferred to $y_l$ is

$$\mathbb P(y_w\succ y_l\mid x) =\sigma\bigl(r(x,y_w)-r(x,y_l)\bigr).$$

Theorem 3.6 (DPO reduction). Theorem 3.2 can be inverted as

$$r(x,y) =\beta\log\frac{\pi^*(y\mid x)} {\pi_{\mathrm{ref}}(y\mid x)} +\beta\log Z(x).$$

Substituting this expression into the Bradley—Terry likelihood and modeling $\pi^*$ by $\pi_\theta$ yields the direct preference optimization loss

$$\mathcal L_{\mathrm{DPO}}(\theta) =-\mathbb E\left[ \log\sigma\!\left( \beta\log\frac{\pi_\theta(y_w\mid x)} {\pi_{\mathrm{ref}}(y_w\mid x)} -\beta\log\frac{\pi_\theta(y_l\mid x)} {\pi_{\mathrm{ref}}(y_l\mid x)} \right)\right].$$

The expectation is over observed preference triples $(x,y_w,y_l)$. The loss can be minimized by ordinary gradient methods without an explicit reward model or on-policy rollouts.

Proof. Taking logarithms in Theorem 3.2 gives the inverted reward formula. In the reward difference $r(x,y_w)-r(x,y_l)$, the term $\beta\log Z(x)$ is identical and cancels. The remaining difference is exactly the argument of the displayed sigmoid. Negative log-likelihood gives the DPO loss. $\square$

Remark 3.7 (why the reduction works). The intractable partition function depends on $x$ but not on the candidate response $y$. Any preference model depending only on reward differences eliminates such an additive prompt-dependent constant.

4. Mechanistic interpretability

Mechanistic interpretability seeks a causal, component-level account of the computation implemented by trained weights. Lecture 5 writes head $h$ as

$$A^h(X)XW_{OV}^h, \qquad W_{QK}^h=W_Q^h(W_K^h)^\top, \quad W_{OV}^h=W_V^hW_O^h.$$

The QK circuit controls the pattern read from the residual stream, and the OV circuit controls the update written to it. In an attention-only transformer, expanding residual additions expresses the output exactly as a sum over paths through these head components; MLP nonlinearities make the analogous decomposition more involved.

Definition 4.1 (induction behavior). On a sequence containing a repeated pattern, an induction head attends from the current token to the position after an earlier occurrence of that token and copies information useful for predicting the repeated continuation. Schematically, on [A][B] ... [A], it uses the earlier transition [A] -> [B] to predict [B] after the second [A].

Mechanistic remark 4.2. A standard two-layer account composes a previous-token head with an induction head. The earlier head writes information about the preceding token into the residual stream; the later head’s QK circuit matches the current token against that information and its OV circuit copies the corresponding successor. Ablating the candidate head and measuring loss on repeated sequences tests causality rather than mere correlation. Ablation means removing or zeroing a component while holding the rest of the model fixed.

Empirical aside 4.2. Induction-like heads have been observed to appear abruptly during training, near changes in in-context learning: improved use of examples or patterns in the prompt without any parameter update. Step 8 Track C asks students to search for this behavior in the course model, compare fresh and repeated halves, and confirm a candidate by ablation. The experiment may reveal a partial or noisy circuit rather than a textbook-perfect head.

A feature direction is a direction in the residual space whose coefficient represents some property of the input. Features are sparse when only a small fraction are active on a typical input.

Definition 4.3 (superposition). Superposition is the representation of more features than residual dimensions by assigning features to non-orthogonal directions. Interference remains manageable when the features are sufficiently sparse and their directions have small pairwise inner products. A coordinate responding to several unrelated features is called polysemantic.

Linear-algebra aside 4.5. High-dimensional Euclidean spaces contain large collections of nearly orthogonal vectors. A Johnson—Lindenstrauss bound says that $N$ points can have pairwise distances approximately preserved in dimension $d=O(\varepsilon^{-2}\log N)$. Read in reverse, a $d$-dimensional space can support exponentially many points at fixed distortion. This does not prove that a trained network uses a particular feature code, but it explains the geometric capacity that superposition exploits.

Remark 4.4 (sparse autoencoders). If meaningful features are directions rather than coordinate axes, individual neurons need not be interpretable. Sparse autoencoders seek an overcomplete dictionary, one containing more feature atoms than ambient dimensions, with sparse activation coefficients, turning feature recovery into dictionary learning, the recovery of atoms and their coefficients. Its sparse-recovery aspect is closely related to compressed sensing, the recovery of sparse signals from a smaller collection of linear measurements.

5. Open questions

Research aside 5.1. The following questions delimit what the preceding mathematics does not yet explain.

1. What implicit regularization—optimizer preferences not written as explicit penalty terms—makes stochastic optimization favor solutions that generalize?
2. Why do empirical scaling curves follow power laws over broad regimes?
3. Which forward-pass algorithms underlie in-context learning beyond induction behavior?
4. How many features are stored in superposition, and when can their dictionary be recovered?
5. How can optimization against a learned reward be related formally to human intent rather than only to the proxy reward?

Closing aside 5.2. The course model is small enough for every matrix to be inspected yet large enough to display attention structure, tokenization artifacts, overfitting, scale-dependent behavior, and sometimes in-context copying. The capstone asks for one precise question, one controlled experiment, and one connection to a result from the lectures.

Summary

Fine-tuning continues the pretraining objective on a new distribution. LoRA restricts task-specific matrix updates to a low-rank factorization. KL-regularized reward maximization has the exact Gibbs optimum $\pi^*\propto\pi_{\mathrm{ref}}e^{r/\beta}$, exposing the tradeoff between reward and deviation from the reference. DPO substitutes this optimum into a reward-difference preference model, where the intractable partition function cancels. Mechanistic interpretability studies trained attention through QK and OV circuits; induction behavior supplies a causal test case, while superposition explains why feature directions need not align with individual neurons.

Exercises (paired with Step 8)

A star marks a Project Step 8 track task.

1. ★ Prove Theorem 3.2 and its $\beta\to\infty$ and $\beta\to0^+$ limits, treating tied reward maximizers.
2. ★ Derive Theorem 3.6 in full and identify exactly which assumption on the preference model causes $\log Z(x)$ to cancel.
3. Interpret $\beta(\log Z-D_{\mathrm{KL}}(\pi\|\pi^*))$ as a free-energy identity and make the statistical-mechanics analogy precise.
4. ★ Apply LoRA to attention projections, sweep $r\in\{1,2,4,8,16\}$, and plot target loss against trainable parameter count.
5. Define a programmatic reward, optimize it by best-of-$n$ or a preference loss, and plot reward against KL divergence. Exhibit a reward exploit.
6. ★ Search for an induction head on repeated random sequences, compare loss on fresh and repeated halves, and confirm a candidate by ablation.
7. Use a Johnson—Lindenstrauss or spherical-packing bound—a bound on the number of well-separated points on a unit sphere—to estimate how many pairwise nearly orthogonal unit vectors can fit in $\mathbb R^{128}$ at a chosen tolerance.
8. Express a two-head induction circuit as a composition of an earlier OV write with a later QK read, and identify the rank constraints.

Pointers

Ouyang et al., InstructGPT (2022); Rafailov et al., DPO (2023); Hu et al., LoRA (2021); Elhage et al., A Mathematical Framework for Transformer Circuits (2021); Olsson et al., In-context Learning and Induction Heads (2022); Elhage et al., Toy Models of Superposition (2022). See the resources page and Project Step 8.