38 Interpretability and Explainability of Generative Models

In 2016 a neural network trained to detect wolves in photographs was found, on examination, to be largely relying on the presence of snow in the background rather than wolf morphology (Ribeiro et al., 2016). The model's test accuracy was high. Its reasoning was wrong. The same year, a recidivism prediction algorithm used in US courts was shown to exhibit systematic racial bias in its scores, despite never being given race as an explicit input variable (Angwin et al., 2016). The model's statistical calibration was arguably defensible. Its outputs carried life-altering consequences. And in 2023, GPT-4 passed the bar exam at the 90th percentile, solved international olympiad problems in chemistry, and wrote code that compiled and passed unit tests at rates exceeding many professional programmers. Yet researchers could not reliably explain which part of the network's billions of parameters was responsible for any particular answer.

These three incidents share a common thread: our ability to deploy powerful machine learning systems has dramatically outpaced our ability to understand them. For discriminative models, this gap is uncomfortable. For generative models, it is foundational. When a language model writes a medical report, a legal brief, or a psychological assessment, or when a diffusion model generates a photorealistic face, the question is not merely whether the output is statistically correct. The question is: where does the output come from? What latent variables contributed? Which training examples are implicitly being recalled? Is the model reasoning, pattern-matching, or hallucinating?

Interpretability and explainability are the scientific disciplines that take these questions seriously. This chapter develops their mathematical foundations with particular attention to the features that make generative models qualitatively more challenging than their discriminative counterparts. We proceed from motivating principles through formal definitions and historical context, then develop the modern mechanistic interpretability programme and its application to generative architectures.

Why Interpretability Matters

The Black Box Problem

A function $f:𝒳\to 𝒴$ is a black box with respect to an observer if the observer can query $f$ on inputs and receive outputs but has no access to the internal computation that transforms one into the other. In classical software engineering the black box is a methodological choice: we deliberately hide implementation details behind interfaces to manage complexity. In modern deep learning the black box is an emergent property of scale: a language model with hundreds of billions of parameters implements a function that no human designed explicitly, whose internal representation cannot be read off the weight matrices in any straightforward way, and whose decision boundaries in the input space have never been analytically characterised.

This distinction matters enormously. A software black box can in principle be white-boxed by reading its source code. A deep learning black box resists white-boxing even in principle: the “source code” is a matrix of floating-point numbers that implements no human-interpretable algorithm. Researchers call this the opacity problem or, in the regulatory literature, the right to explanation problem.

Philosophical Dimensions: Consciousness, Explanation, and the Introspection Gap

Before developing formal machinery, it is worth pausing on a question that practitioners often suppress: is interpretability even a well-posed demand? Consider the human brain. It is the most capable cognitive system we know of, and it is spectacularly uninterpretable by its own introspective access. Cognitive psychologists have documented systematic failures of human self-explanation: subjects confabulate reasons for choices made on purely unconscious grounds (Nisbett and Wilson, 1977); patients with split-brain lesions offer verbal explanations for actions initiated by the hemisphere that lacks language (Gazzaniga, 1985); people reliably attribute their own emotional states to external causes unrelated to the actual physiological trigger (Schachter and Singer, 1962).

This raises what we may call the introspection gap: the systematic divergence between an agent's stated reasons for a decision and the actual causal factors that produced it. If the human brain, operating at $\sim$86 billion neurons and $\sim {10}^{14}$ synapses, exhibits a severe introspection gap that does not prevent us from considering humans morally responsible agents, why should we require machine learning systems to close that gap entirely?

The philosophical answer turns on the difference between causal explanation and rationalisation. What we demand of interpretable systems is not that they produce coherent post-hoc narratives (which humans do too, and often wrongly) but that we can identify the actual causal mechanisms by which inputs produce outputs. This is a scientific demand, not an epistemological impossibility. Neuroscience has made genuine progress understanding specific neural circuits, spike-timing codes, and functional modules in biological brains precisely by combining external intervention (lesions, optogenetics, pharmacology) with multi-electrode recordings. Mechanistic interpretability for artificial neural networks is the same programme applied to silicon.

A second philosophical question concerns consciousness and self-awareness. Can a system explain its own decisions in a meaningful sense? When a language model is prompted “explain your reasoning”, the resulting text is generated by the same forward pass mechanism as any other output. It is not a privileged readout of internal state; it is another output token sequence conditioned on the prompt. This is not necessarily deceptive, but it is not the same as genuine introspective access. The distinction matters for regulation, for trust, and for the scientific study of AI systems.

Formal Definitions: Interpretability, Explainability, and Transparency

The terms “interpretability”, “explainability”, and “transparency” are used inconsistently across the literature. We adopt the following formal definitions, which unify the major usages and are compatible with the EU AI Act's framework.

The GenAI Explainability Triangle

Generative models present a richer explainability surface than discriminative models. We organise the challenge using what we call the GenAI Explainability Triangle: three interrelated demands that must all be addressed for a generative system to be considered genuinely explainable.

Vertex 1: Generative Mechanism Transparency. What latent variables, attention patterns, or intermediate representations caused the specific output? For a language model this includes: which training documents are being implicitly recalled, which attention heads are routing information from which positions, and whether the model is “reasoning” or “pattern-matching” (to the extent that distinction is coherent). For a diffusion model it includes: which noise trajectory was taken, which guidance signals shaped the denoising path, and how the latent manifold was traversed.

Vertex 2: User-Centered Interpretability. What explanation is appropriate for a given user with a given need? A data scientist querying a language model needs a different explanation than a radiologist, a judge, or a patient. The explanation must be adapted to the user's domain model, the consequentiality of the decision, and the available interaction modality.

Vertex 3: Evaluation Fidelity. Can we measure whether an explanation is correct? For discriminative models, fidelity metrics like area under the perturbation curve (Samek et al., 2017) are well-established. For generative models, standard fidelity metrics are largely undefined because the ground truth causal mechanism is unknown and the output space is high-dimensional and structured.

Why Generative Models Are Harder to Interpret

Classical XAI was developed primarily for discriminative models: classifiers and regressors that map a fixed input to a scalar or low-dimensional output. Methods like LIME, SHAP, Grad-CAM (Selvaraju et al., 2017), and integrated gradients (Sundararajan et al., 2017) produce saliency maps or feature attributions that highlight which input features most influenced a particular classification.

Generative models introduce three qualitative difficulties that these methods do not address:

Stochastic sampling.

A discriminative model $f:𝒳\to 𝒴$ is typically deterministic given $\theta$. A generative model samples from a distribution $p_{\theta }(𝒚|𝒙)$. Two forward passes with the same input and parameters can produce arbitrarily different outputs. An attribution method that computes gradients of the output with respect to the input is undefined or ambiguous for a stochastic output. Even if we fix the random seed, the attribution reflects that seed's trajectory, not the model's distributional behaviour.

High-dimensional, structured output spaces.

Discriminative model outputs are typically scalar or low-dimensional. Generative outputs are sequences of tokens, images with millions of pixels, audio waveforms, or protein structures. Attribution methods that assign a scalar importance score to each input feature face a fundamental mismatch: the output is not a scalar to differentiate with respect to. Extending gradient-based attributions to sequence outputs requires integrating over output positions, which introduces aggregation choices that interact poorly with long-range dependencies.

Multi-step generation.

Many generative models involve multiple cascaded components: a language model may use retrieval, chain-of-thought, a code interpreter, and post-processing. A diffusion model iterates hundreds of denoising steps through multiple network passes. A VAE encodes, samples, and decodes. Each intermediate step introduces its own representation, its own stochasticity, and its own opportunities for information to be created, discarded, or transformed. Attribution methods designed for single forward passes cannot natively attribute credit across such pipelines.

The following table summarises the qualitative differences in the XAI challenge between discriminative and generative models.

XAI Dimension	Discriminative Models	Generative Models
Output type	Scalar / low-dimensional label	Sequence, image, audio, or structured object
Determinism	Deterministic given parameters	Stochastic; distribution over outputs
Attribution target	Single output scalar	Ambiguous: which token, pixel, or step?
Evaluation of attributions	Pixel-flipping, deletion curves	No standard metrics; output space too complex
Forward pass structure	Single encoder + head	Multi-stage pipeline; retrieval; sampling
Latent space	Compact class embeddings	High-dimensional, entangled; sampling noise
Causal question	Why this class?	Why this generation trajectory?
Human-facing utility	Highlight discriminative pixels	Identify failure modes, bias, copying
Faithfulness measurable?	Yes (deletion / perturbation)	Difficult; ground truth not available
Regulatory touchpoint	GDPR Article 22	EU AI Act high-risk systems

Biological Connections: Neural Coding and Distributed Representations

The question of how information is encoded in neural networks, artificial or biological, has a long history in neuroscience that illuminates the interpretability challenge. Two competing theories of neural coding have shaped the field.

The sparse coding hypothesis.

Barlow (1961) proposed that the nervous system achieves efficient representation by using as few neurons as possible to encode each stimulus. In the extreme, the grandmother cell hypothesis (Gross, 2002) posits that individual neurons represent specific high-level concepts: one neuron fires for your grandmother's face, one fires for Jennifer Aniston, one fires for the Eiffel Tower. This is trivially interpretable: if you can label each neuron, you can read off the model's internal representation.

Evidence for this view comes from the seminal work of Hubel and Wiesel (1959, 1962), who found neurons in primary visual cortex that respond selectively to oriented edges at specific spatial frequencies, and from Quiroga et al. (2005), who found neurons in the human medial temporal lobe responding to single specific individuals or landmarks.

The distributed representation hypothesis.

McCulloch and Pitts (1943) and later Hinton, McClelland, and Rumelhart (1986) argued for a fundamentally different coding scheme: information is represented by patterns of activity across many neurons, no single one of which is individually interpretable. A face is encoded as a high-dimensional activation vector, not as a single neuron's firing rate. This allows combinatorial generalisation: $N$ neurons can represent up to $2^N$ distinct patterns rather than $N$ concepts.

Modern deep neural networks clearly use distributed representations. This makes them powerful and difficult to interpret simultaneously. However, recent work in mechanistic interpretability suggests that superposition (Elhage et al., 2022) allows networks to store more features than dimensions by representing features as near-orthogonal directions in activation space, with small but non-zero interference.

The superposition hypothesis has important implications for interpretability. If a network stores $k\gg n$ features in $n$ dimensions via superposition, then probing individual neurons is insufficient: the features are polysemantic (a single neuron responds to multiple unrelated concepts). Recovering individual features requires finding sparse decompositions of the activation space, which is the goal of sparse autoencoders in mechanistic interpretability (Cunningham et al., 2023; Bricken et al., 2023).

Exercises

The Mechanistic Interpretability Paradigm

Defining the Programme

Mechanistic interpretability is the programme of reverse-engineering trained neural networks into human-understandable computational processes. The goal is not to produce a local explanation for a single prediction (as in SHAP or LIME) but to produce a complete account of the algorithm a network has learned: which neurons compute which features, how those features are combined via attention or MLP operations, and what high-level algorithm the resulting circuit implements.

Three Abstraction Layers

The mechanistic interpretability programme operates at three interrelated levels of abstraction, inspired by Marr's three levels of analysis in computational neuroscience (Marr, 1982).

Layer 1: Neurons as Features.

At the finest grain, individual neurons (or directions in activation space) represent specific features of the input. For vision models these may be edge detectors, colour blobs, or object part detectors. For language models they may represent syntactic roles, entity types, or factual properties. The goal at this layer is to produce a feature dictionary: a mapping from neurons (or sparse directions) to their semantic content.

Methods at this layer include:

• Maximum activation examples: find the training or held-out inputs that maximally activate a neuron, and infer the feature from the pattern.

• Activation patching on concepts: verify that a neuron represents a concept by measuring the downstream effect of setting its activation to its mean.

• Sparse autoencoders (SAEs): learn a sparse, overcomplete dictionary $𝐃\in {\mathrm{ℝ}}^{d\times k}$ such that activations $𝒉\approx 𝐃𝒄$ with ${\Vert 𝒄\Vert }_0\ll k$. Each column of $𝐃$ is a purported monosemantic feature.

Layer 2: Circuits.

A circuit is a sub-graph of the computation graph connecting features to downstream features through specific attention heads, MLP layers, or residual stream connections. Circuits explain specific input-output relationships: “the indirect object identification circuit in GPT-2 (Wang et al., 2022) routes the subject token's representation through three specific attention heads to generate the indirect object at the output position.”

The circuit paradigm was established by Olah et al. (2020) who reverse-engineered a curve-detector circuit in InceptionV1, showing that curve detector neurons receive inputs from oriented edge detector neurons (themselves identified as implementing a Fourier decomposition of local contrast patterns) via a specific set of connections. The insight that the same circuit type (curve from edges) recurs at multiple spatial positions is the notion of circuit universality: similar circuits emerge in different networks trained on different tasks, suggesting convergent algorithmic solutions.

Layer 3: Macroscopic Algorithms.

At the coarsest grain, what high-level algorithm does the network implement? For GPT-2 on the indirect object identification (IOI) task, Wang et al. (2022) showed the algorithm has three phases: (i) token copying, in which name mover heads copy subject representations to the final token position; (ii) subject suppression, in which inhibition heads downweight the current subject relative to alternatives; and (iii) output generation, in which name mover heads promote the correct output. This three-phase description is a macroscopic algorithm that could in principle be implemented differently by a different network, but which GPT-2 happened to implement through the specific circuit identified at Layer 2.

The Causal Intervention Toolkit

Mechanistic interpretability relies on causal interventions: modifying specific activations in the network and observing the downstream effect. This section formalises the four main techniques.

Let $f_{\theta }$ be a transformer with $L$ layers, $H$ attention heads per layer, and hidden dimension $d$. Denote by ${𝒂}_t^{(l,h)}$ the attention head output of head $h$ in layer $l$ at position $t$, and by ${𝒎}_t^{(l)}$ the MLP output at layer $l$, position $t$. The residual stream at position $t$ after layer $l$ is: (Residual)$$\begin{array}{ccc}& {𝒓}_t^{(l)}={𝒓}_t^{(0)}+\sum _{\ell =1}^l(\sum _{h=1}^H{𝒂}_t^{(\ell ,h)}+{𝒎}_t^{(\ell )}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒓}_t^{(0)}$ is the token embedding at position $t$.

Automated Circuit Discovery

Manual circuit analysis (Olah et al., 2020; Wang et al., 2022) requires a researcher to form and test hypotheses about each attention head and MLP component. For GPT-2-small (117M parameters, 12 layers, 12 heads per layer), this process required several months of researcher time per circuit. For modern large language models the brute-force approach is computationally and cognitively intractable.

Two algorithmic approaches have emerged to automate circuit discovery.

ACDC (Automated Circuit DisCovery).

Conmy et al. (2023) introduce ACDC, a greedy top-down algorithm that identifies the minimal sub-graph of the computation graph sufficient to reproduce a target metric. Starting from the full computation graph, ACDC iteratively ablates edges (connections between components) and retains an edge if its ablation causes the metric to change by more than a threshold $\tau$: (ACDC)$$\begin{array}{cccc}& {𝒢}^{(t+1)}& ={𝒢}^{(t)}\setminus \{e\},\text{if }|ℳ({𝒢}^{(t)})-ℳ({𝒢}^{(t)}\setminus \{e\})|<\tau .& \text{<span class="tml-eqn"></span>}\end{array}$$ The output is a minimal circuit ${𝒢}^{\ast }$ such that all retained edges contribute more than $\tau$ to the metric. ACDC is correct by construction for greedy selection but may miss non-minimal circuits that interact synergistically.

Edge Attribution Patching (EAP).

Syed et al. (2023) and Marks et al. (2024) propose EAP as a gradient-based approximation to path patching that scales to large models. For each edge $e=(u\to v)$ in the computation graph, EAP computes an attribution score: (EAP)$$\begin{array}{ccc}& {\alpha }_e={({\widehat{𝒓}}_v-{𝒓}_v^{\ast })}^{\top }\frac{\partial ℳ}{\partial {𝒓}_v}{|}_{{𝒓}_v={𝒓}_v^{\ast }},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\widehat{𝒓}}_v$ is the residual stream at $v$ on the corrupted input and ${𝒓}_v^{\ast }$ is the same on the clean input. The gradient $\partial ℳ/\partial {𝒓}_v$ is computed with a single backward pass on the clean input. EAP thus requires only two forward passes (clean and corrupt) and one backward pass, regardless of circuit size.

The key insight is that ${\alpha }_e$ approximates the first-order contribution of restoring the clean activation at $v$ via edge $e$ to the metric. The approximation is exact in the linear limit and remains accurate when activations are small perturbations of their clean values.

A Taxonomy of Mechanistic Interpretability Methods

Method	Layer	Core Technique	Key Reference	Cost
Max activation examples	Neuron	Find inputs maximising neuron firing	Olah et al., 2017	$O(|𝒟|)$
Feature visualisation	Neuron	Gradient ascent in input space	Olah et al., 2017	$O(I\cdot d)$
Probing classifiers	Neuron	Train linear probe on activations	Alain & Bengio, 2017	$O(|𝒟|\cdot d)$
Sparse autoencoders	Neuron	Learn overcomplete sparse basis	Bricken et al., 2023	$O(k\cdot |𝒟|)$
Logit lens	Neuron/Circuit	Project intermediate states to logits	Nostalgebraist, 2020	$O(L\cdot d)$
Activation patching	Circuit	Restore clean activations selectively	Meng et al., 2022	$O(n_c\cdot T_{\mathrm{fwd}})$
Path patching	Circuit	Patch specific computational paths	Wang et al., 2022	$O(E\cdot T_{\mathrm{fwd}})$
Mean ablation	Circuit	Replace with distributional mean	Chan et al., 2022	$O(n_c\cdot T_{\mathrm{fwd}})$
ACDC	Circuit	Greedy edge pruning via ablation	Conmy et al., 2023	$O(E\cdot T_{\mathrm{fwd}})$
EAP	Circuit	First-order gradient attribution	Syed et al., 2023	$O(T_{\mathrm{fwd}}+T_{\mathrm{bwd}})$
Causal tracing	Algorithm	Corrupt + restore; locate facts	Meng et al., 2022	$O(L\cdot n_{\mathrm{tok}})$
ROME / MEMIT	Algorithm	Rank-one model editing	Meng et al., 2022/2023	$O(d^2)$
Attention analysis	Algorithm	Visualise attention weight patterns	Clark et al., 2019	$O(LH{T}^2)$

Chapter Roadmap

The remainder of this chapter develops the mechanistic interpretability programme for generative models specifically. Figure fig:interp:roadmap shows the organisation.

Exercises

Classical Explainability Methods

Before we can understand why interpretability for generative models demands fundamentally new tools, we must master the classical toolkit. The methods developed between 2016 and 2020 - SHAP, LIME, Grad-CAM, Integrated Gradients, and attention visualisation - represent the mature engineering of a particular idea: that a model's prediction on a given input can be explained by attributing credit or blame to the individual input features that drove it. These are attribution methods, and they are extraordinarily successful at what they do. They are also, as we shall argue in Why Classical XAI Fails for Generative Models, fundamentally inadequate for generative models.

The story of attribution methods begins with a question that seems almost embarrassingly simple. Given a neural network $f$ that classifies an input $𝒙$ and produces a prediction $f(𝒙)=\hat{y}$, which features of $𝒙$ were responsible? For a radiologist looking at a chest X-ray classified as showing pneumonia, this question is not academic: it determines whether the model can be trusted, whether it has learned the right visual patterns, and whether the radiologist's intuition aligns with the machine's. For a judge evaluating whether an algorithmic decision was fair, knowing which features drove the decision is a legal requirement in many jurisdictions.

Yet the question conceals a thicket of subtlety. What does it mean for a feature to be “responsible” for a prediction? Features are not independent; removing or perturbing one changes the context for all others. The model $f$ is a nonlinear function of thousands or millions of parameters, not a transparent formula one can inspect directly. And the very notion of feature importance is ambiguous: do we mean the feature's causal effect on the output, its correlation with the output, or some measure of counterfactual sensitivity?

The four methods we study in this section each answer this question differently, and their different answers reveal complementary facets of model behaviour.

SHAP: Shapley Values from Game Theory

The most principled approach to attribution comes not from machine learning but from cooperative game theory. In 1953, Lloyd Shapley solved the problem of fairly distributing the proceeds of a cooperative game among its players, and his solution - now bearing his name - has become the mathematical foundation of modern feature attribution.

Cooperative Games and Fair Division

A cooperative game is a pair $(𝒩,v)$ where $𝒩=\{1,\dots ,d\}$ is a set of players (here, feature indices) and $v:2^{𝒩}\to \mathrm{ℝ}$ is a characteristic function that assigns a real-valued payoff to every subset (coalition) $S\subseteq 𝒩$. For attribution purposes, $v(S)$ represents how well a model performs when only the features in $S$ are available. Typically, $$v(S)={𝔼}_{{𝒙}_{\bar{S}}}[f({𝒙}_S,{𝒙}_{\bar{S}})]-f(\mathrm{𝟎}),$$ where $\bar{S}=𝒩\setminus S$ and the expectation marginalises over all features not in $S$, while features in $S$ retain their observed values.

The Shapley value of player $i$ measures that player's average marginal contribution across all possible orderings in which players join the coalition.

The Shapley value is the unique allocation satisfying four axioms that formalise the notion of fairness.

KernelSHAP: Tractable Approximation

Direct computation of Shapley values requires evaluating $v$ on all $2^d$ subsets, which is computationally infeasible for high-dimensional inputs. Lundberg and Lee (2017) introduced SHAP (SHapley Additive exPlanations), a practical framework with the KernelSHAP algorithm as its model-agnostic backbone.

The key observation is that Shapley values can be cast as the solution to a weighted least-squares problem. Define a coalition vector $𝐳\in {\{\mathrm{0,1}\}}^d$ indicating which features are present, and let the SHAP kernel weight for coalition $S$ (encoded as $𝐳$) be (Kernel)$$\begin{array}{ccc}& w(S)=\frac{d-1}{\left(\genfrac{}{}{0px}{}{d}{|S|}\right)|S|(d-|S|)}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Then the SHAP attributions $𝝓=({\varphi }_0,{\varphi }_1,\dots ,{\varphi }_d)$ are obtained by solving (Kernelshap)$$\begin{array}{ccc}& 𝝓=\underset{{𝝓}^{\prime }}{\mathrm{arg\hspace{0.17em}min}}\sum _{S\subseteq 𝒩}w(S)[v(S)-{\varphi }_0-{𝝓}^{\prime }\cdot {𝐳}_S{]}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐳}_S\in {\{\mathrm{0,1}\}}^d$ is the binary indicator of $S$. In practice, one samples a manageable number of coalitions (typically $2d+2^{11}$), evaluates $v$ on each (possibly by querying $f$ with missing features replaced by background samples), and solves the weighted least-squares problem.

LIME: Local Interpretable Model-Agnostic Explanations

While SHAP seeks a globally consistent attribution grounded in game theory, LIME (Ribeiro, Singh, and Guestrin, 2016) takes a pragmatically local approach: fit a simple interpretable model to the behaviour of the complex model in the neighbourhood of a specific input.

Local Surrogate Framework

Given an input $𝒙\in 𝒳$ and a black-box classifier $f:𝒳\to \mathrm{ℝ}$, LIME constructs a local surrogate $g\in 𝒢$ (typically a sparse linear model) that approximates $f$ near $𝒙$. The explanation is the solution to (LIME)$$\begin{array}{ccc}& \xi (𝒙)=\underset{g\in 𝒢}{\mathrm{arg\hspace{0.17em}min}}ℒ(f,g,{\pi }_{𝒙})+\mathrm{\Omega }(g),& \text{<span class="tml-eqn"></span>}\end{array}$$ where:

• $ℒ(f,g,{\pi }_{𝒙})={\sum }_{𝒛\in 𝒵}{\pi }_{𝒙}(𝒛)[f(𝒛)-g(𝒛){]}^2$ is the local fidelity loss, measuring how well $g$ mimics $f$ on perturbations $𝒛$ of $𝒙$, weighted by proximity ${\pi }_{𝒙}$;

• ${\pi }_{𝒙}(𝒛)=\exp (-D(𝒙,𝒛{)}^2/{\sigma }^2)$ is an exponential kernel assigning higher weight to perturbations closer to $𝒙$;

• $\mathrm{\Omega }(g)$ is a complexity penalty (e.g., ${\ell }_1$ norm of the linear weights) that encourages sparse, human-readable explanations.

The perturbations $𝒛$ are generated by randomly masking a subset of superpixels (for images) or words (for text), replacing masked regions with a background value or held-out tokens. The surrogate $g$ is then learned on the dataset ${\{({𝒛}_k,f({𝒛}_k),{\pi }_{𝒙}({𝒛}_k))\}}_{k=1}^n$ using weighted ridge regression.

Locality and Its Implications

The locality of LIME's explanation is simultaneously its strength and its vulnerability. On the one hand, a local linear model can capture the local decision boundary of an arbitrarily complex classifier, provided the neighbourhood radius $\sigma$ is calibrated correctly. On the other hand, different runs of LIME on the same input can produce qualitatively different explanations if the sampled perturbations happen to fall in different regions of a complex loss landscape. This instability is a documented limitation that has motivated several remediation strategies, including using larger perturbation budgets, fixing the random seed, and employing regularised regression in place of standard ridge.

Grad-CAM: Gradient-Weighted Class Activation Mapping

SHAP and LIME are model-agnostic: they treat the classifier as a black box and probe it by varying inputs. Grad-CAM (Selvaraju et al., 2017) takes the opposite approach, exploiting the internal gradient structure of convolutional neural networks to produce spatial saliency maps that highlight where in the image the model looked when making its prediction.

Class Activation Maps via Gradients

Let $A^k\in {\mathrm{ℝ}}^{u\times v}$ denote the activation map of the $k$-th feature map in the final convolutional layer, and let $y^c$ be the logit (pre-softmax score) for class $c$. Grad-CAM computes the importance weight of the $k$-th feature map for class $c$ by globally average-pooling the gradient of $y^c$ with respect to $A^k$: (Weight)$$\begin{array}{ccc}& {\alpha }_k^c=\frac{1}{uv}\sum _{i=1}^u\sum _{j=1}^v\frac{\partial y^c}{\partial A_{ij}^k}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The Grad-CAM saliency map is then (MAP)$$\begin{array}{ccc}& L_{\text{Grad-CAM}}^c=\mathrm{ReLU}\left(\sum _k{\alpha }_k^c{A}^k\right)\in {\mathrm{ℝ}}^{u\times v},& \text{<span class="tml-eqn"></span>}\end{array}$$ where the ReLU discards negative contributions (features that suppress class $c$) and retains only the positive evidence. The resulting map is upsampled (bilinearly) to the input resolution and overlaid on the image.

The ReLU operation in (MAP) is a deliberate choice. Without it, the map would include features that suppress the predicted class, producing a noisier and harder-to-interpret visualisation. With it, Grad-CAM produces localised, high-contrast heatmaps that are easy for human experts to interpret.

Integrated Gradients: Path-Based Attribution with Axioms

Integrated Gradients (Sundararajan, Taly, and Yan, 2017) sits at a conceptual crossroads between the game-theoretic rigour of SHAP and the gradient efficiency of Grad-CAM. Like SHAP, it is grounded in a set of formal axioms; like Grad-CAM, it exploits the gradient of the network with respect to its input. Unlike Grad-CAM, it is architecture-agnostic: it applies to any differentiable model, not just convolutional networks.

Path Integral and Axiomatic Foundation

The Integrated Gradients attribution of feature $i$ is defined as the integral of the gradient along a straight-line path from a baseline ${𝒙}^{\prime }\in 𝒳$ (representing “no signal”, e.g. an all-black image or an all-padding text sequence) to the actual input $𝒙$: (IG)$$\begin{array}{ccc}& {\mathrm{IG}}_i(𝒙;{𝒙}^{\prime })=(x_i-x_i^{\prime }){\int }_0^1\frac{\partial f({𝒙}^{\prime }+\alpha (𝒙-{𝒙}^{\prime }))}{\partial x_i}\mathrm{d}\alpha .& \text{<span class="tml-eqn"></span>}\end{array}$$ In practice, the integral is approximated by a Riemann sum over $m$ equally-spaced steps along the path: (Approx)$$\begin{array}{ccc}& {\mathrm{IG}}_i(𝒙;{𝒙}^{\prime })\approx (x_i-x_i^{\prime })\cdot \frac{1}{m}\sum _{k=1}^m{\frac{\partial f(𝒙)}{\partial x_i}|}_{𝒙={𝒙}^{\prime }+\frac{k}{m}(𝒙-{𝒙}^{\prime })}.& \text{<span class="tml-eqn"></span>}\end{array}$$ For this right-endpoint Riemann sum the approximation error decreases as $O(1/m)$ (a midpoint or trapezoidal quadrature would improve this to $O(1/m^2)$); values of $m\in [20,300]$ are standard in practice.

The Completeness axiom is particularly important from an auditing perspective: it guarantees that the explanation accounts for the entire prediction, with no attribution budget left unaccounted for. By contrast, vanilla gradient-times-input methods (which multiply the gradient at $𝒙$ by the input value) can violate Completeness because they do not integrate over the path.

Attention Visualisation: Is Attention Explanation?

The rise of Transformer models introduced a new form of “explanation” that is appealing in its simplicity: attention weights. When a text classifier based on BERT produces a prediction, the multi-head self-attention layers produce weight matrices $𝐀\in {[\mathrm{0,1}]}^{n\times n}$ (one per head per layer) that ostensibly indicate how much each token “attends to” every other token. Visualising these weights as heatmaps became a popular way to explain model behaviour, and several major NLP papers in 2018–2019 relied on attention visualisations to support interpretability claims.

The Jain-and-Wallace Challenge

In a landmark 2019 paper, Jain and Wallace challenged the assumption that attention weights constitute valid explanations. Their argument rests on two empirical observations and one theoretical argument.

Empirical finding 1: Attention and feature importance disagree. When LIME and gradient-based attributions are computed alongside attention weights for the same prediction, the resulting attributions are poorly correlated (Spearman $\rho <0.3$ in many cases). High attention weight on a token does not imply that token drives the prediction.

Empirical finding 2: Adversarial attention. It is possible to construct “adversarial attention” distributions that produce the same model output as the original attention distribution while being maximally different in the attention weights. Formally, given a prediction $\hat{y}$ produced with attention $𝐀$, Jain and Wallace find an alternative attention $\widetilde{𝐀}$ such that $f(𝒙;\widetilde{𝐀})=f(𝒙;𝐀)$ but $\mathrm{TVD}(𝐀,\widetilde{𝐀})$ is large. This demonstrates that the mapping from attention to output is many-to-one: attention can be changed dramatically without changing the prediction.

Theoretical argument: Attention is an intermediate computation. Attention weights are computed inside the model, not at the interface between model and input. They capture which token representations are mixed at a given layer, but the final prediction depends on how those mixed representations are subsequently processed by feed-forward sublayers, layer normalisation, and additional attention heads. Attributing the prediction to attention weights at one layer ignores all downstream processing.

What Attention Does Tell Us

Despite these limitations, attention visualisation is not worthless. Attention patterns in pre-trained language models exhibit systematic behaviours that are genuinely informative:

• Syntactic heads. Certain heads in BERT's lower layers reliably attend along syntactic dependency arcs (subject-verb, noun-modifier), even though the model was never trained with syntactic supervision.

• Coreference heads. Some heads track coreference chains, attending from pronouns to their antecedents.

• Positional heads. A subset of heads implement diagonal attention (attending to the immediately preceding or following token), implementing a form of local context integration.

These patterns reveal something real about what the model has learned, even if they do not constitute faithful attribution of any specific prediction. The distinction matters: attention is a useful diagnostic tool for understanding model behaviour at the population level, but a poor explanatory tool for understanding any specific prediction.

Exercises: Classical Explainability Methods

Why Classical XAI Fails for Generative Models

Classical XAI methods are engineered solutions to a specific problem: given a discriminative model $f:𝒳\to 𝒴$ (a classifier or regressor) and a fixed input $𝒙\in 𝒳$, attribute the scalar prediction $f(𝒙)\in 𝒴$ to the input features. The problem is well-posed: the input is fixed, the output is a scalar, and the attribution is a vector in ${\mathrm{ℝ}}^d$. The methods we surveyed in Classical Explainability Methods are mature, principled, and - within their domain - genuinely useful.

The problem is that generative models are not discriminative models. A diffusion model, an autoregressive language model, a VAE, or a vision-language model does not classify an input; it produces outputs. And the production process is qualitatively different from classification in ways that systematically defeat every classical XAI method. This section catalogues these failure modes precisely, building toward the formal definition of what interpretability should mean for generative models.

The detective analogy is helpful here. Classical XAI assumes that we are trying to understand a judge's verdict: given the evidence (input $𝒙$), which pieces of evidence drove the verdict ($f(𝒙)$)? But a generative model is not a judge passing a verdict; it is a novelist writing a story. The novelist does not declare a fact; they generate a narrative that could have gone in many directions. Understanding the novelist requires a fundamentally different set of questions: What algorithms does the novelist employ? What internal representations mediate the creative process? How does the initial prompt constrain the space of possible stories?

The Attribution Problem: What Is the Input?

All classical XAI methods presuppose a fixed, well-defined input $𝒙\in 𝒳$ to which attributions are assigned. For a ResNet classifying an image, $𝒙$ is a pixel tensor; for a BERT model classifying sentiment, $𝒙$ is a token sequence. The attribution vector $𝝓\in {\mathrm{ℝ}}^d$ is indexed by the same dimensions as $𝒙$.

Generative models confound this assumption in several distinct ways.

The condition vs. the generation.

In a conditional generative model such as a text-conditioned image diffusion model, the input comprises both the condition $𝒄$ (the text prompt) and the noise seed $𝒛$ (the initial Gaussian sample that seeds the reverse process). If a generated image contains a spurious artefact, is it attributable to $𝒄$ or $𝒛$? The artefact may arise from the interaction of specific prompt tokens with specific noise configurations in ways that cannot be decomposed.

Unconditional generation.

In unconditional models (vanilla VAEs, vanilla GANs, unconditional diffusion), there is no condition at all. The “input” is pure noise $𝒛\sim p(𝒛)$, which has no semantic meaning. Attributing a generated face's eye colour to specific dimensions of $𝒛$ is technically feasible (one can compute $\partial {\widehat{𝒙}}_i/\partial z_j$) but semantically meaningless without a prior understanding of what $z_j$ represents - and that understanding is precisely what is missing.

The prompt is not the cause.

In large language models, the prompt $𝒄$ is the nominal input, but the model's generation also depends on its entire training history (encoded in $𝜽$), the decoding algorithm (temperature, top-$p$, beam search), and - in stochastic decoding - the random seed. Attributing a specific generated sentence to specific prompt tokens conflates these three qualitatively different sources of influence.

Stochastic Generation: Same Input, Different Outputs

Classical XAI methods implicitly assume that $f(𝒙)$ is a deterministic scalar. LIME builds a local surrogate by querying $f$ on perturbations of $𝒙$ and recording the outputs; SHAP evaluates $v(S)={𝔼}_{\bar{S}}[f({𝒙}_S,{𝒙}_{\bar{S}})]$; Grad-CAM computes $\partial f(𝒙)/\partial A^k$. All of these operations require either a deterministic output or the ability to marginalise over randomness in a controlled way.

Generative models are stochastic by design. A diffusion model with temperature $\tau >0$ produces a different image on every forward pass for the same prompt. An autoregressive model with nucleus sampling generates a different completion on every call. A VAE samples $𝒛\sim q_{\varphi }(𝒛|𝒙)$ and decodes $\widehat{𝒙}\sim p_{\theta }(\widehat{𝒙}|𝒛)$, with two stochastic steps.

This stochasticity creates a fundamental problem for attribution.

Which output do we explain?

Given $n$ stochastic outputs ${\widehat{𝒙}}^{(1)},\dots ,{\widehat{𝒙}}^{(n)}$ from the same prompt, LIME would need to construct a different surrogate for each output (since each has different pixel values to be attributed). Averaging the surrogates loses information about which outputs are typical and which are outliers. Explaining the average output $\overline{𝒙}=(1/n){\sum }_k{\widehat{𝒙}}^{(k)}$ is misleading because the average is itself not a typical output of the generative model (it is blurry, for diffusion models, or incoherent, for language models).

Output variance as information.

The variance across stochastic outputs is itself a meaningful signal: high variance indicates that the model is uncertain or that the prompt under-specifies the generation. Classical XAI methods entirely ignore this variance, providing no tools to explain why the model is uncertain or which prompt tokens are responsible for under-specification.

Multi-Step Processes: The Attribution Chain Problem

Diffusion models generate through $T\in [50,1000]$ denoising steps. Autoregressive language models generate through $N$ sequential token predictions, where $N$ may be hundreds or thousands. In both cases, the final output is the result of a long chain of stochastic transformations, and the causal structure of this chain creates fundamental obstacles for classical attribution.

Diffusion: The T-Step Attribution Puzzle

In a DDPM with $T$ steps, the generation proceeds as a Markov chain: $${𝒙}_T\to {𝒙}_{T-1}\to \cdots \to {𝒙}_1\to {𝒙}_0,$$ where each step applies the learned denoising network ${\mu }_{\theta }({𝒙}_t,t)$. If we wish to attribute a feature of the final image ${𝒙}_0$ (say, the presence of a red object in the upper-left quadrant) to the initial noise ${𝒙}_T$, we must trace the causal path through all $T$ denoising steps.

A single application of Grad-CAM at the final denoising step $t=1\to 0$ provides a saliency map for the last step, but ignores the contributions of steps $t=2,\dots ,T$. The contribution of early steps is mediated through the intermediate states ${𝒙}_{T-1},\dots ,{𝒙}_1$, which are never observed (in standard sampling) and whose dimensions lack semantic meaning.

Integrated Gradients can in principle be applied by treating ${𝒙}_T$ as the input and defining a baseline ${𝒙}_T^{\prime }$ (e.g., the zero vector). But the straight-line interpolation from ${𝒙}_T^{\prime }$ to ${𝒙}_T$ passes through intermediate noise tensors that are not valid inputs to the denoising chain: the denoising network at step $t$ expects ${𝒙}_t\sim q({𝒙}_t|{𝒙}_0)$, not an arbitrary interpolation.

Autoregressive Models: The Token Dependency Graph

In an autoregressive language model, token $n+1$ is generated conditioned on all previous tokens ${𝒄}_{1:n}$, which include both the original prompt and the model's own previous generations. The causal graph of influences is a directed acyclic graph where every prompt token and every previously generated token can influence every future token.

SHAP can compute attributions for any specific output token ${\hat{c}}_{n+1}$ by treating ${𝒄}_{1:n}$ as the features. But the attributions change for every output token, and aggregating them across a multi-token generation requires a choice: sum? max? mean? Each aggregation discards structural information about how attribution evolves as generation proceeds.

Moreover, in teacher-forcing training, the model predicts each token conditioned on the ground-truth context, but in inference, it conditions on its own (possibly erroneous) previous predictions. Attributing a generation error to a specific prompt token may be misleading if the error was actually caused by an earlier generation error that snowballed.

Latent Space Opacity: VAE and GAN Codes

In discriminative models, the input $𝒙$ has a natural human interpretation: pixels have positions, words have meanings. This semantic grounding is what makes feature attribution meaningful: saying “the model attended to the top-right corner” or “the word `not' had negative attribution” conveys actionable information.

Latent variable models fundamentally lack this semantic grounding.

VAE Latent Codes

In a VAE, the encoder $q_{\varphi }(𝒛|𝒙)$ maps an input to a distribution over a $k$-dimensional latent space. The decoder $p_{\theta }(\widehat{𝒙}|𝒛)$ maps latent codes back to data space. If we attribute a generated image feature to latent dimensions $z_j$, what does this mean semantically?

In general, nothing. The latent code $𝒛$ is organised by the ELBO objective to maximise reconstruction quality and keep the posterior close to the prior $𝒩(\mathrm{𝟎},𝐈)$. There is no mechanism that ensures $z_j$ corresponds to any human-interpretable concept (“hair colour”, “facial expression”, “object shape”). In practice, latent dimensions in vanilla VAEs are entangled: $z_1$ may partially encode hair colour, face orientation, and background brightness simultaneously.

Disentanglement methods ($\beta$-VAE, FactorVAE, TC-VAE) seek to encourage axis-aligned latent factors that correspond to independent generative factors. But even successful disentanglement does not guarantee semantic alignment: a disentangled latent code where $z_1$ controls “size of the leftmost object” and $z_2$ controls “hue of the background” may be interpretable by reference to the images but does not correspond to human-level concepts.

GAN Latent Spaces

In a GAN, the generator $G:𝒵\to 𝒳$ maps a noise vector $𝒛\sim p(𝒛)$ to a data point. Unlike VAEs, GANs impose no constraint on the structure of $𝒵$: the latent space can be geometrically complex, with no guarantee of linearity, convexity, or semantic regularity.

GANSpace (Härkönen et al., 2020) and InterfaceGAN (Shen et al., 2020) discovered empirically that GAN latent spaces often contain interpretable directions: a vector $𝐝\in 𝒵$ such that traversing $𝒛\to 𝒛+\lambda 𝐝$ corresponds to a semantically coherent change in the generated image (e.g., ageing, adding glasses, changing hairstyle). But these directions must be discovered by post-hoc linear probing or supervised steering - they are not intrinsic to the latent code's native coordinates.

Cross-Modal Alignment: VLMs and the Fusion Problem

Vision-language models (VLMs) such as CLIP, DALL-E, Stable Diffusion, and GPT-4V add a further layer of complexity: they must fuse information from two qualitatively different modalities - visual tokens and text tokens - into a unified representation that drives generation. The alignment between modalities is learned implicitly through massive-scale contrastive or generative training, and the resulting representations are highly non-transparent.

The Cross-Attention Attribution Problem

In a latent diffusion model conditioned on text, cross-attention layers compute attention between visual tokens (keys and values from the spatial feature map) and text tokens (queries from the CLIP text encoder). The cross-attention weight $a_{ij}$ measures the degree to which spatial position $i$ of the image attends to text token $j$.

Visualising cross-attention maps (as in Prompt-to-Prompt and Attend-and-Excite) provides a form of spatial-textual attribution: “the word `sunset' drove the redness of the sky region.” But this attribution has the same problems as standard attention visualisation (see Attention Visualisation: Is Attention Explanation?) compounded by cross-modal entanglement: text tokens are processed by CLIP, which itself is a complex non-transparent model; visual tokens are extracted by a VAE encoder; and the cross-attention operates in a compressed latent space, not in pixel space.

Modality Attribution Failure

Consider the question: “In a CLIP-guided diffusion model that generated an image of a dog on a skateboard, how much of the `skateboard' visual content is attributable to the text prompt, how much to the image prior learned during training, and how much to the initial noise seed?” This question is natural for a human engineer debugging a generation failure, but it has no well-defined answer in the framework of classical attribution:

• The text prompt provides the semantic intent (skateboard).

• The model's weights encode the visual prior (what skateboards look like in the training distribution).

• The noise seed determines which specific visual configuration is selected from the prior.

• The CLIP text encoder transforms the word “skateboard” into a 512-dimensional vector whose relationship to the skateboard's visual appearance is encoded in the weights, not in any human-readable form.

Attributing the skateboard's visual presence to “the word `skateboard' in the prompt” conflates all of these factors. Classical SHAP over the text tokens would assign attribution to the word “skateboard” simply because removing it changes the output - but it would assign zero attribution to the training data that taught the model what skateboards look like, and zero attribution to the model architecture that determines how text and image representations are aligned.

The Fundamental Shift: From Feature Attribution to Algorithm Understanding

The failure modes catalogued above are not engineering deficiencies that can be patched by better implementations of SHAP or Integrated Gradients. They reflect a fundamental mismatch between the question that classical XAI asks and the question that generative model interpretability requires.

Classical XAI asks: which input features drove this output? This question presupposes a fixed input, a scalar (or low-dimensional) output, and a model that computes a deterministic function of the input. Generative models violate all three presuppositions.

The shift has a precedent in the history of science. Early pharmacologists asked “which molecules in this extract reduce fever?” - an attribution question. Modern molecular pharmacology asks “what receptor does aspirin bind, and how does binding inhibit prostaglandin synthesis?” - a mechanistic question. The mechanistic question is harder to answer but provides qualitatively different understanding: it explains not just what works but why it works, and it enables prediction and design rather than mere post-hoc explanation.

A Taxonomy of Classical XAI Failure Modes

Method	Autoregressive LM	Diffusion Model	VAE/GAN	VLM
SHAP	Partial: per-token	Critical: $T$-step	Critical: latent	Critical:
	attribution only	coalition eval	space opacity	cross-modal
	ignores generation	intractable	undef. input	entanglement
LIME	Partial: prompt	Critical: no	Critical: no	Partial: prompt
	perturbations	stable output	input to	perturbations
	only	to regress on	perturb	only; cross-
				attn ignored
Grad-CAM	N/A: no	Partial: applies	N/A: no	Partial: cross-
	conv. layers	at single step	conv. stack	attn maps are
		only		not Grad-CAM
Integr. Grads	Partial: per-	Critical: path	Partial: input	Critical: no
	token only;	not in data	is latent;	semantic
	aggregation ad hoc	manifold	no semantics	baseline exists
Attention	Misleading:	Misleading:	N/A	Misleading:
Viz.	not faithful	cross-attn		cross-modal
	attribution	partial only		attn conflated

Interpretability Desiderata for Generative Models

Having identified the failure modes of classical XAI, we can now state positively what a satisfactory interpretability method for generative models must achieve. We formulate this as three formal desiderata.

These desiderata generalise and extend the axiomatic foundations of classical XAI. Faithfulness generalises the Sensitivity axiom: both require that explanations accurately reflect causal relationships. Completeness generalises the Efficiency axiom of SHAP: both require that the explanation accounts for the total output. Minimality is new: it has no direct classical counterpart, because attribution methods over a fixed input space are automatically minimal (they produce exactly $d$ attribution values, one per feature). For generative models, where the explanation space is the set of computational circuits in the model, minimality is non-trivial to enforce.

The proposition is not a counsel of despair. It identifies the correct unit of analysis: rather than attributing outputs to input features, interpretability methods for generative models must attribute outputs to model components (circuits, layers, attention heads, neurons, latent directions). This is the programme of mechanistic interpretability, which we develop in the subsequent sections of this chapter.

Exercises: Failures of Classical XAI

The Mathematical Framework of Transformer Circuits

Modern large language models are, at heart, transformers: architectures built from repeated blocks of multi-head self-attention and feedforward networks connected by residual streams. To truly interpret what these models compute, we need more than visualisations of attention weights; we need a rigorous mathematical language for decomposing the computation into atomic, understandable units. This language is the theory of transformer circuits, developed by Anthropic researchers Elhage et al. as a precise mechanistic account of what individual components of a transformer compute and how they interact.

The central insight is beautifully simple: at every layer, each token maintains a vector called its residual stream, which acts as a shared communication bus. Every component-attention heads and MLP layers alike-reads from this bus, performs a transformation, and writes back an additive update. Because all operations are additive, the residual stream at any layer is an exact linear superposition of all previous contributions. This linearity is not an approximation; it is an architectural fact that makes the transformer uniquely amenable to circuit-level analysis.

The Residual Stream as a Communication Bus

We begin with a precise definition.

The key structural consequence of (DEF) is that every component of the network interacts with every other component only through the residual stream. There are no direct skip connections between layer 1's attention heads and layer 3's MLP; the only path is through sequential reads and writes to the shared bus.

Token embedding ${𝐖}_E$.

The embedding matrix ${𝐖}_E\in {\mathrm{ℝ}}^{d_{\text{model}}\times |𝒱|}$ maps a discrete token index $t\in \{1,\dots ,|𝒱|\}$ to a dense vector ${𝐖}_E{𝒆}_t\in {\mathrm{ℝ}}^{d_{\text{model}}}$, where ${𝒆}_t$ is the $t$-th standard basis vector. In practice, the model learns ${𝐖}_E$ end-to-end, and its columns are the token embedding vectors. The initialisation of the residual stream is (INIT)$$\begin{array}{ccc}& {𝒙}_i^{(0)}={𝐖}_E{𝒆}_{t_i}+{𝒑}_i,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒑}_i$ is the positional encoding for position $i$.

Unembedding ${𝐖}_U$.

The unembedding matrix ${𝐖}_U\in {\mathrm{ℝ}}^{|𝒱|\times d_{\text{model}}}$ maps the final residual stream vector back to logits over the vocabulary. In many implementations, ${𝐖}_U={𝐖}_E^{\top }$ (weight tying), although this is not universal in large models. The pre-softmax logit for token $v$ is ${\ell }_v={({𝐖}_U{𝒙}_i^{(L)})}_v$, and the predicted probability is $\hat{p}(v|\text{context})=\text{softmax}{({𝐖}_U{𝒙}_i^{(L)})}_v$.

Attention Heads as Independent Additive Units

Each attention head in a transformer layer is an independent computation that (1) computes a soft routing pattern over source positions, and (2) moves weighted information from those positions to the destination. Crucially, multiple heads in the same layer act independently and additively: their outputs are summed before being added to the residual stream.

Formally, for layer $\ell$ with $H$ heads, the attention sub-layer update for destination token $i$ is (Additive)$$\begin{array}{ccc}& {\mathrm{\Delta }}_i^{(\ell ,\text{attn})}=\sum _{h=1}^H{\mathrm{\Delta }}_i^{(\ell ,h)},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\mathrm{\Delta }}_i^{(\ell ,h)}$ is the contribution from head $h$ in layer $\ell$. The MLP update ${\mathrm{\Delta }}_i^{(\ell ,\text{MLP})}$ is then added separately.

This additive structure has a profound consequence: we can isolate the contribution of any individual head by ablating or patching it, and its removal changes the residual stream by exactly $-{\mathrm{\Delta }}_i^{(\ell ,h)}$. There is no non-linear interaction between different heads in the same layer (though non-linearities in the MLP and in subsequent attention heads can create effective inter-layer non-linearities through the residual stream).

The QK Circuit: Attention Pattern Computation

Every attention head contains four weight matrices: ${𝐖}_Q^h,{𝐖}_K^h,{𝐖}_V^h\in {\mathrm{ℝ}}^{d_{\text{head}}\times d_{\text{model}}}$ and ${𝐖}_O^h\in {\mathrm{ℝ}}^{d_{\text{model}}\times d_{\text{head}}}$. We derive the two fundamental circuits that characterise each head.

Derivation of the QK circuit.

The query and key vectors for head $h$, computed from source residual streams ${𝒙}_j^{(\ell -1)}$ and destination residual stream ${𝒙}_i^{(\ell -1)}$, are (Vectors)$$\begin{array}{ccc}& {𝒒}_i^h={𝐖}_Q^h{𝒙}_i^{(\ell -1)}\in {\mathrm{ℝ}}^{d_{\text{head}}},{𝒌}_j^h={𝐖}_K^h{𝒙}_j^{(\ell -1)}\in {\mathrm{ℝ}}^{d_{\text{head}}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The (pre-softmax) attention score between destination $i$ and source $j$ is the scaled inner product: (Score)$$\begin{array}{ccc}& s_{ij}^h=\frac{{𝒒}_i^{h\top }{𝒌}_j^h}{\sqrt{d_{\text{head}}}}=\frac{{({𝐖}_Q^h{𝒙}_i^{(\ell -1)})}^{\top }({𝐖}_K^h{𝒙}_j^{(\ell -1)})}{\sqrt{d_{\text{head}}}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Expanding the inner product: (Expansion)$$\begin{array}{ccc}& s_{ij}^h=\frac{1}{\sqrt{d_{\text{head}}}}{𝒙}_i^{(\ell -1)\top }\underset{{𝐖}_{QK}^h}{\underbrace{{𝐖}_Q^{h\top }{𝐖}_K^h}}{𝒙}_j^{(\ell -1)},& \text{<span class="tml-eqn"></span>}\end{array}$$ which reveals the QK circuit matrix ${𝐖}_{QK}^h\in {\mathrm{ℝ}}^{d_{\text{model}}\times d_{\text{model}}}$.

The OV Circuit: Information Movement

The OV circuit describes what information is moved once the attention pattern determines where to attend.

Derivation of the OV circuit.

Given attention weights $a_{ij}^h$, the value vector aggregated at destination $i$ from head $h$ is (AGG)$$\begin{array}{ccc}& {𝒗}_i^h=\sum _j{a}_{ij}^h{𝐖}_V^h{𝒙}_j^{(\ell -1)}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The output of head $h$ at destination $i$ is then projected back to the residual stream dimension: (Output)$$\begin{array}{ccc}& {\mathrm{\Delta }}_i^{(\ell ,h)}={𝐖}_O^h{𝒗}_i^h={𝐖}_O^h\sum _j{a}_{ij}^h{𝐖}_V^h{𝒙}_j^{(\ell -1)}=\sum _j{a}_{ij}^h\underset{{𝐖}_{OV}^h}{\underbrace{{𝐖}_O^h{𝐖}_V^h}}{𝒙}_j^{(\ell -1)}.& \text{<span class="tml-eqn"></span>}\end{array}$$

The OV circuit makes precise the intuition that a head “copies information”: if ${𝐖}_{OV}^h\approx 𝐈$, the head copies the source residual stream verbatim (weighted by attention) to the destination. If ${𝐖}_{OV}^h$ maps token type embeddings to prediction directions in ${𝐖}_U^{\top }$, the head directly increments logits for particular next tokens.

TikZ Figure: The Residual Stream Architecture

The following figure illustrates the residual stream as a shared bus with attention heads reading and writing through their QK and OV circuits.

Head Composition across Layers

Single-layer attention heads implement simple operations. The true power of transformers emerges from composition: when the output of one head enters the residual stream and is read by a later head.

Q-Composition, K-Composition, and V-Composition

Let head $A$ be in layer ${\ell }_1$ and head $B$ be in layer ${\ell }_2>{\ell }_1$. Head $B$ reads from the residual stream ${𝒙}_i^{({\ell }_2-1)}$, which contains the contribution ${\mathrm{\Delta }}_i^{({\ell }_1,A)}={𝐖}_{OV}^A{𝒙}_j^{({\ell }_1-1)}$ (averaged over the head $A$ attention pattern). Depending on which weight matrix of head $B$ is most sensitive to this contribution, we distinguish:

1. Q-Composition. Head $A$'s output ${\mathrm{\Delta }}_i^{({\ell }_1,A)}$ influences ${𝒒}_i^B$ via ${𝐖}_Q^B$. The virtual QK weight for the composed circuit is (Compose)$$\begin{array}{ccc}& {𝐖}_{QK}^{A\to B}={({𝐖}_Q^B{𝐖}_{OV}^A)}^{\top }{𝐖}_K^B,& \text{<span class="tml-eqn"></span>}\end{array}$$ which determines how head $A$'s processing of destination token $i$ modulates what head $B$ attends to from source token $j$.

2. K-Composition. Head $A$'s output ${\mathrm{\Delta }}_j^{({\ell }_1,A)}$ (at source position $j$) influences ${𝒌}_j^B$ via ${𝐖}_K^B$. The virtual weight is (Compose)$$\begin{array}{ccc}& {𝐖}_{QK}^{A\to B}={𝐖}_Q^{B\top }({𝐖}_K^B{𝐖}_{OV}^A),& \text{<span class="tml-eqn"></span>}\end{array}$$ which controls how head $A$'s processing of source token $j$ modulates what head $B$ attends to from $j$.

3. V-Composition. Head $A$'s output ${\mathrm{\Delta }}_j^{({\ell }_1,A)}$ (at source position $j$) influences ${𝒗}_j^B$ via ${𝐖}_V^B$. The virtual OV weight for the composed circuit is (Compose)$$\begin{array}{ccc}& {𝐖}_{OV}^{A\to B}={𝐖}_O^B({𝐖}_V^B{𝐖}_{OV}^A),& \text{<span class="tml-eqn"></span>}\end{array}$$ which maps the original source residual stream (before head $A$ processed it) all the way to the output of head $B$.

Induction Heads: A Canonical Example of Composition

The most studied example of two-layer head composition is the induction head circuit. Induction heads perform a simple but powerful operation: given a sequence containing $\dots AB\dots A$, they predict $B$ to follow the second occurrence of $A$ by attending back to the first occurrence. This is achieved through the composition of two heads across layers:

1. A previous-token head in layer 1 (K-composition): attends from position $i$ to position $i-1$, writing information about the previous token into the residual stream at each position.

2. An induction head in layer 2 (uses the layer-1 output via K-composition): attends from position $i$ to position $j$ where the previous-token residual stream at $j$ matches the current token at $i$.

Together they implement the lookup: “find where my current token appeared before, and copy what came after it.” This is a two-layer circuit with measurable virtual weights, and its ablation causes a sharp drop in in-context learning performance.

Theorem: Decomposition of Transformer Output

We now state the fundamental decomposition theorem that makes transformer circuits mathematically tractable.

Example: Information Flow in a 2-Layer Transformer

Exercises

Multi-Head Attention Decomposition

Having established the mathematical framework of transformer circuits in The Mathematical Framework of Transformer Circuits, we now turn to a systematic decomposition of multi-head attention as an information-routing mechanism. The central question is: given that the residual stream is a linear superposition of head contributions, can we interpret each head's contribution to the model's predictions? The answer is yes, through a combination of logit attribution, the logit lens technique, and a careful analysis of direct versus indirect effects through the residual stream.

Attention as Weighted Information Routing

The attention mechanism can be understood as a content-addressable memory lookup. At each destination position $i$, the head broadcasts a query into the residual stream space and retrieves a weighted combination of values from source positions, where the weights are determined by query-key similarity.

More formally, using the circuit notation from The Mathematical Framework of Transformer Circuits, head $h$ in layer $\ell$ implements the routing: (Formal)$$\begin{array}{ccc}& {\mathrm{\Delta }}_i^{(\ell ,h)}=\sum _{j=1}^n{\mathrm{softmax}}_j\left(\frac{{𝒙}_i^{\top }{𝐖}_{QK}^h{𝒙}_j}{\sqrt{d_{\text{head}}}}\right){𝐖}_{OV}^h{𝒙}_j,& \text{<span class="tml-eqn"></span>}\end{array}$$ where we have abbreviated ${𝒙}_j={𝒙}_j^{(\ell -1)}$ for readability. The softmax ${\mathrm{softmax}}_j(\cdot )$ over all source positions $j$ produces a probability distribution ${𝒂}_i^h=(a_{i1}^h,\dots ,a_{in}^h)\in {\mathrm{\Delta }}^{n-1}$ over source positions.

The routing metaphor.

Information at source position $j$ is routed to destination $i$ with weight $a_{ij}^h$. The OV circuit determines what information is moved: only the component of ${𝒙}_j$ in the row space of ${𝐖}_{OV}^h$ can be moved at all. The QK circuit determines where to route: it scores every (destination, source) pair and produces a soft assignment.

Softmax Attention Patterns and Their Interpretation

Attention weights ${𝒂}_i^h\in {\mathrm{\Delta }}^{n-1}$ are produced by softmax, which has several interpretability implications.

Sparsity and entropy.

The entropy of the attention distribution, (Entropy)$$\begin{array}{ccc}& H({𝒂}_i^h)=-\sum _j{a}_{ij}^h\log a_{ij}^h,& \text{<span class="tml-eqn"></span>}\end{array}$$ measures how diffuse the head's attention is. Low entropy (near $0$) indicates the head attends to a single source position-a hard lookup; high entropy (near $\log n$) indicates uniform attention over all positions-a soft average. Interpreting a head's function often begins with measuring this entropy: induction heads are sparse (they have found the right position to attend to), while positional heads may be diffuse (they average nearby positions).

Attention as a simplex-valued function.

The attention pattern $A^h\in {\mathrm{ℝ}}^{n\times n}$ (with $A_{ij}^h=a_{ij}^h$, rows summing to 1) is a stochastic matrix. Its left eigenvectors decompose the pattern into interpretable modes: the dominant left singular vector of $A^h-\frac{1}{n}{\mathrm{𝟏𝟏}}^{\top }$ (the attention pattern with uniform background subtracted) often corresponds to a clear structural pattern such as attending to punctuation, the previous token, the subject of the sentence, and so forth.

Direct versus Indirect Effects through the Residual Stream

When we ask “what is the effect of head $h$ on the final logits?”, we must distinguish two types of contribution.

The Logit Lens: Tracking Predictions Layer by Layer

The logit lens is a diagnostic technique introduced by Nostalgebraist that applies the unembedding matrix to intermediate residual stream states, revealing how the model's prediction evolves across layers. It exploits the decomposition from Theorem 3 to make the intermediate computation interpretable.

The logit lens allows us to observe the model “changing its mind” across layers: an early layer might assign high probability to a syntactically plausible but semantically wrong continuation, which is corrected by later layers. This layer-by-layer evolution of predictions has been used to identify which layers perform which computational subtasks.

Relationship to the residual stream decomposition.

Note that (Decomp)$$\begin{array}{ccc}& {\hat{\ell }}_v^{(\ell )}-{\hat{\ell }}_v^{(\ell -1)}={𝐖}_U(\sum _{h=1}^H{\mathrm{\Delta }}_i^{(\ell ,h)}+{\mathrm{\Delta }}_i^{(\ell ,\text{MLP})})\cdot {𝒆}_v,& \text{<span class="tml-eqn"></span>}\end{array}$$ the change in logit contributed by layer $\ell$. This can be further decomposed head by head, giving a per-head, per-layer logit contribution chart.

Proposition: Additivity of Head Contributions in Logit Space

We now prove the key structural result justifying per-head attribution.

MLP Blocks: The Role of Feedforward Networks

While attention heads route information between positions, MLP blocks process each token's representation independently (they are applied position-wise with no cross-token communication). A standard two-layer MLP sub-block computes (DEF)$$\begin{array}{ccc}& {\mathrm{\Delta }}_i^{(\ell ,\text{MLP})}={𝐖}_{\text{out}}^{(\ell )}\sigma \left({𝐖}_{\text{in}}^{(\ell )}{𝒙}_i^{(\ell )}\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐖}_{\text{in}}^{(\ell )}\in {\mathrm{ℝ}}^{d_{\text{ff}}\times d_{\text{model}}}$, ${𝐖}_{\text{out}}^{(\ell )}\in {\mathrm{ℝ}}^{d_{\text{model}}\times d_{\text{ff}}}$, and $\sigma$ is an element-wise activation (GELU or ReLU in practice).

The non-linearity $\sigma$ is the key difference from attention heads: MLPs are not linear in the residual stream. However, the post-activation vector $𝒉=\sigma ({𝐖}_{\text{in}}^{(\ell )}{𝒙}_i^{(\ell )})\in {\mathrm{ℝ}}^{d_{\text{ff}}}$ can be analysed neuron by neuron.

The “Attention Gathers, MLPs Compose” Hypothesis

Motivated by empirical findings in vision transformers and large language models, a key interpretability hypothesis has emerged:

Empirical evidence for this hypothesis comes from several directions. In vision transformers (ViTs), attention heads in early layers perform local texture gathering, while MLP blocks in the same layers compute responses to features like edges and colour gradients. In large language models, factual knowledge is disproportionately stored in and recalled by MLP layers, while grammatical and positional relationships are implemented in attention layers.

TikZ Figure: Attention versus MLP Division of Labour

Exercises

Induction Heads and In-Context Learning

Among the most consequential discoveries of the mechanistic interpretability programme is the identification of a concrete, reusable computational motif called the induction head. Found independently in nearly every sufficiently large transformer trained on natural text, the induction head implements a form of pattern matching that appears to be the primary driver of in-context learning-the remarkable ability of large language models to adapt their behaviour from a few examples given in the input prompt, without any gradient update.

This section develops the mathematics of induction heads from the ground up. We begin with the simpler previous token heads that form in early layers of two-layer transformers, then show how a subsequent layer of attention heads exploits the key–value information produced by previous token heads to implement the $[𝐀][𝐁]\cdots [𝐀]\to [𝐁]$ completion pattern. We formalise the notion of K-composition by which two attention heads interact across layers, derive a mathematical characterisation of the attention pattern an induction head must implement, and place the discovery in its historical context as a major milestone on the path from black-box language modelling to circuit-level understanding.

Previous Token Heads

In a two-layer transformer, the first layer of attention heads performs computations on the residual stream before the second layer sees it. Empirically, several heads in layer 1 consistently learn a simple but strategically important operation: they attend almost exclusively to the token immediately preceding the query token in the sequence. We call these previous token heads.

Let the input sequence be $t_1,t_2,\dots ,t_T$, where each $t_i$ is a token index. The residual stream at position $i$ after the embedding and positional encoding is (Residual INIT)$$\begin{array}{ccc}& {𝒙}_i^{(0)}={𝐞}_{t_i}+{𝐩}_i,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐞}_{t_i}\in {\mathrm{ℝ}}^d$ is the token embedding and ${𝐩}_i\in {\mathrm{ℝ}}^d$ is the positional embedding. A previous token head with weight matrices ${𝐖}_Q,{𝐖}_K,{𝐖}_V\in {\mathrm{ℝ}}^{d\times d_h}$ computes attention scores (PREV ATTN Score)$$\begin{array}{ccc}& a_{ij}=\frac{{𝒙}_i^{(0)\top }{𝐖}_Q{𝐖}_K^{\top }{𝒙}_j^{(0)}}{\sqrt{d_h}},j\le i,& \text{<span class="tml-eqn"></span>}\end{array}$$ with a causal mask preventing attention to future positions. The empirical observation is that the learned weight product ${𝐖}_Q{𝐖}_K^{\top }$ is dominated by a component that computes similarity between the positional embedding of position $i$ and the positional embedding of position $j=i-1$. More precisely, if ${𝐩}_i$ is a sinusoidal or learned positional embedding, then the product ${𝐩}_i^{\top }{𝐖}_Q{𝐖}_K^{\top }{𝐩}_{i-1}$ is large relative to any other $(i,j)$ pair, producing an approximately one-hot attention pattern $A_{ij}\approx \mathrm{𝟏}[j=i-1]$.

The output of the previous token head at position $i$ is then (PREV Token Output)$$\begin{array}{ccc}& {𝒛}_i=\sum _{j=1}^i{A}_{ij}{𝒙}_j^{(0)}{𝐖}_V{𝐖}_O\approx {𝒙}_{i-1}^{(0)}{𝐖}_V{𝐖}_O,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐖}_O\in {\mathrm{ℝ}}^{d_h\times d}$ is the output projection. This is added to the residual stream, so after layer 1 position $i$ contains information about both $t_i$ and $t_{i-1}$: (Residual After L1)$$\begin{array}{ccc}& {𝒙}_i^{(1)}={𝒙}_i^{(0)}+{𝒛}_i\approx {𝐞}_{t_i}+{𝐩}_i+f({𝐞}_{t_{i-1}},{𝐩}_{i-1}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $f$ denotes the linear map ${𝐖}_V{𝐖}_O$. This “token $+$ successor-context” representation in the residual stream is the key information that induction heads will exploit in layer 2.

Induction Heads: the [A][B]...[A] $\to$ [B] Pattern

An induction head is an attention head in layer 2 that implements a specific lookup procedure: when the current query position holds token $A$ (which has appeared before in the sequence), the induction head attends back to the position after the previous occurrence of token $A$ and copies the token $B$ that followed it.

The key question is: how does position $t$ know to attend to position $s+1$ rather than $s$? The answer lies in K-composition with the previous token head (Previous Token Heads).

K-Composition Between Layers

Recall from sec:interp:residual-stream that the residual stream is a sum of contributions from all previous layers. The attention pattern of an induction head in layer 2 is determined by the inner product of query and key vectors, both of which are computed from the layer 2 residual stream ${𝒙}^{(1)}$. Because ${𝒙}^{(1)}$ contains the output of previous token heads, the induction head's key at position $j$ encodes information about $t_{j-1}$ (the token before $j$). Meanwhile, the query at position $t$ encodes information about $t_t=A$.

More precisely, the attention score of the induction head between query position $t$ and key position $j$ is (Induction Score)$$\begin{array}{ccc}& a_{t,j}=\frac{{({𝒙}_t^{(1)})}^{\top }{𝐖}_Q^{(2)}{({𝐖}_K^{(2)})}^{\top }{𝒙}_j^{(1)}}{\sqrt{d_h}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The layer-2 residual stream decomposes as ${𝒙}_j^{(1)}={𝒙}_j^{(0)}+{𝒛}_j$, where ${𝒛}_j\approx f({𝐞}_{t_{j-1}})$ is the previous token head output. If the induction head's key weight matrix ${𝐖}_K^{(2)}$ is primarily sensitive to the “previous token” component ${𝒛}_j$, then (K COMP Expand)$$\begin{array}{ccc}& {({𝐖}_K^{(2)})}^{\top }{𝒙}_j^{(1)}\approx {({𝐖}_K^{(2)})}^{\top }{𝒛}_j\approx {({𝐖}_K^{(2)})}^{\top }{𝐖}_V^{(1)}{𝐖}_O^{(1)}{𝐞}_{t_{j-1}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the virtual key contributed to position $j$ by the token $t_{j-1}$. The effective key–query inner product (K COMP Inner)$$\begin{array}{ccc}& a_{t,j}\propto {𝐞}_{t_t}^{\top }{𝐖}_Q^{(2)}{({𝐖}_K^{(2)})}^{\top }{𝐖}_V^{(1)}{𝐖}_O^{(1)}{𝐞}_{t_{j-1}}& \text{<span class="tml-eqn"></span>}\end{array}$$ is large when $t_{j-1}$ is the same token as $t_t$ (i.e., when $j-1=s$, the previous occurrence of $A$), so $j=s+1$.

This is K-composition: the key of the layer-2 head is effectively composed with the value–output mapping of the layer-1 head, allowing the second layer to look up “which position follows the previous occurrence of the current token.”

Traced Example: A Concrete Induction Head Run

To build intuition, we trace the induction head on the classic nursery-rhyme continuation. Suppose the input sequence (tokenised) is:

which we index as positions $1,2,\dots ,12$. Token positions: The$=1$, cat$=2$, sat$=3$, on$=4$, the$=5$, mat$=6$, .$=7$. The subsequence repeats: positions $8$–$12$ mirror positions $1$–$5$ (incomplete repetition; the model must predict the continuation at position $13$).

At position $12$ (query token the, second occurrence), the induction head proceeds as follows:

1. Query: The query vector at position $12$ encodes token the (the same token type as position $5$).

2. Key scan: The keys at positions $1$–$12$ contain, in their “previous token” component, information about tokens $0,1,2,\dots ,11$ respectively (using a dummy token $0$ at position $0$). Position $6$ has previous token the (position $5$, token $5$), so the K-composed key at position $6$ encodes token the.

3. Attention spike: Because the query encodes the and the K-composed key at position $6$ also encodes the, the attention logit $a_{\mathrm{12,6}}$ is large. The softmax concentrates attention weight on position $6$.

4. Value read-out: Position $6$ holds token mat (the token that followed the previous occurrence of the at position $5$). The value–output circuit copies the representation of mat into the residual stream at position $12$.

5. Prediction: The model therefore strongly predicts mat as the next token-which is correct.

Why Induction Heads Work on Random Tokens

A striking experimental finding from Olsson et al. (2022) is that induction heads operate equally well on random token sequences-sequences drawn uniformly at random with no semantic structure whatsoever. If the sequence is 7, 42, 13, 7, the induction head at the second 7 correctly predicts 42 (the token that followed the first 7).

This is not a trivial observation. It rules out the hypothesis that induction heads work by learning semantic associations (“the word cat often follows the word the”). Instead, the mechanism is purely structural: it implements a sequence-level pattern-matching algorithm that can be described entirely in terms of token identities and relative positions.

The random-token experiment also provides a quantitative diagnostic for induction head behaviour. Given a sequence of length $2T$ of the form $(r_1,r_2,\dots ,r_T,r_1,r_2,\dots ,r_T)$ with $r_i$ drawn i.i.d. from a uniform distribution over the vocabulary, an ideal induction head achieves cross-entropy loss (Ideal LOSS)$$\begin{array}{ccc}& {ℒ}_{\mathrm{ideal}}=-\log \frac{1}{V}\text{(on the first half, no context)}& \text{<span class="tml-eqn"></span>}\end{array}$$ for the first repetition ($T+1$) and exactly zero loss for positions $T+2,\dots ,2T$ (since the pattern is perfectly predictable from context). In practice, models with strong induction heads achieve losses close to zero on the second half of such sequences, a metric called the induction-head loss used as a circuit-level benchmark.

The Phase Change: Sudden Emergence of In-Context Learning

Perhaps the most dramatic finding of the Olsson et al. (2022) study is the tight correspondence between the formation of induction heads during training and a sudden, non-linear improvement in in-context learning ability. The training curve of a large transformer on natural language exhibits a characteristic phase change:

1. Early training (before phase change): The model learns basic unigram and bigram statistics; the loss decreases smoothly. In-context learning ability (measured, for example, by how much the model benefits from in-context examples) is approximately zero. On random-token sequences, the second-half loss is approximately equal to the first-half loss.

2. The phase change: At a specific training step (which occurs earlier for larger models), the in-context loss drops sharply-over perhaps a factor-of-ten shorter window than the overall loss decrease. The model's ability to use in-context examples jumps discontinuously.

3. After the phase change: Induction heads are firmly established. The second-half loss on random-token sequences approaches zero. In-context learning ability is mature and roughly invariant to further training.

The mechanistic explanation is as follows. Induction heads form via a two-stage process:

1. In stage (a), previous token heads form in layer 1. This is a relatively gradual process.

2. In stage (b), having discovered the “previous token” information available in the residual stream, layer-2 heads rapidly reorganise to exploit K-composition. This is the sudden stage that produces the observed phase change.

The suddenness arises because K-composition is a multiplicative interaction: the gain from the induction-head algorithm scales as ${\lambda }^2$ (the square of the K-composition strength), so once the previous token head is sufficiently strong, the gradient advantage of the induction-head algorithm becomes self-reinforcing.

Exercises

The IOI Circuit - Interpretability in the Wild

The study of induction heads established that mechanistic interpretability can uncover real computational structure in trained transformers. But induction heads are, in some sense, a simple case: they form in two-layer models, they implement a single well-defined algorithm, and their function can be characterised by a handful of matrix products. The natural next question is whether the circuit-finding programme can scale to more complex, naturalistic tasks in full-sized models.

The answer was provided by Wang et al. (2022) in a landmark paper on the Indirect Object Identification (IOI) task in GPT-2 Small. The IOI task requires determining the indirect object of a sentence such as:

The correct continuation is “Mary” (the indirect object) rather than “John” (the subject, who is grammatically barred from being the recipient of an action performed by himself). Wang et al. identified a specific subgraph of 26 attention heads within GPT-2 Small-together with associated MLP inputs-that is causally sufficient to perform this task. This subgraph, which they called the IOI circuit, is the most complete mechanistic account of a real NLP capability in a full-scale language model to date.

The IOI Task: Formal Description

GPT-2 Small (12 layers, 12 heads per layer, 768-dimensional residual stream) achieves ${\mathrm{\Delta }}_{\ell }\approx 4.3$ on the IOI dataset (Appendix A of Wang et al., 2022), indicating that it reliably performs the task.

Head Roles in the IOI Circuit

The 26 attention heads in the IOI circuit play distinct functional roles, determined by their effect on the logit difference ${\mathrm{\Delta }}_{\ell }$. Wang et al. identified four main categories:

Name Mover Heads

Name Mover Heads (NMH) are the final functional component of the IOI circuit: they read the indirect object name from its position in the context and copy it to the final token position (the position at which the model must predict the next token). The three most important NMHs are located in layers 9 and 10 of GPT-2 Small (heads 9.9, 9.6, and 10.0 in layer.head notation).

Mechanistically, NMHs implement a form of V-composition with earlier heads: their value–output circuit copies the name token embedding from the indirect object's position to the output residual stream. The OV matrix of a NMH satisfies ${𝐖}_V^{\mathrm{NMH}}{𝐖}_O^{\mathrm{NMH}}\approx {𝐏}_{\mathrm{copy}}$, where ${𝐏}_{\mathrm{copy}}$ projects the residual stream component corresponding to name tokens onto the logit space.

More precisely, if the residual stream at the $IO$ position contains the representation ${𝒓}_{IO}$ and the NMH attends to that position with weight $A_{\mathrm{final},IO}\approx 1$, the contribution to the final-position logits is (NMH Contribution)$$\begin{array}{ccc}& \mathrm{\Delta }\ell =A_{\mathrm{final},IO}{𝒓}_{IO}^{\top }{𝐖}_V^{\mathrm{NMH}}{𝐖}_O^{\mathrm{NMH}}{𝐖}_U,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐖}_U\in {\mathrm{ℝ}}^{d\times V}$ is the unembedding matrix. When ${𝐖}_O^{\mathrm{NMH}}{𝐖}_U$ is approximately a scaled copy operator on name-token logits, this directly boosts ${\ell }_{IO}$ relative to all other tokens.

S-Inhibition Heads

S-Inhibition Heads (SIH) serve the complementary role of suppressing the subject name $S$ at the final token position. Without SIHs, the model would predict both $S$ and $IO$ with high probability (since both are plausible sentence completions in isolation), leading to errors on sentences where $S$ appears more recently in context.

The SIHs (principally heads 8.6, 8.10, and 7.3) sit in layers 7–8, above the duplicate token heads that detect the subject's repeated occurrence. Their query vectors at the final position attend to earlier positions that contain subject-name tokens; their value circuit then suppresses the corresponding name logit by writing a negative contribution to the residual stream at the final position: (SIH Contribution)$$\begin{array}{ccc}& \mathrm{\Delta }{\ell }_S=A_{\mathrm{final},S_2}{𝒓}_{S_2}^{\top }{𝐖}_V^{\mathrm{SIH}}{𝐖}_O^{\mathrm{SIH}}{𝐖}_U\approx -c\cdot {𝒆}_S,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $S_2$ denotes the second (most recent) occurrence of the subject name, $c>0$ is a positive scalar, and ${𝒆}_S$ is a unit vector selecting the $S$ logit. This suppression mechanism is what allows the model to distinguish the correct indirect object from the subject.

Duplicate Token Heads

Duplicate Token Heads (DTH) detect the fact that a name appears twice in the sentence-a key signal that identifies it as the subject. The DTHs (heads 0.1, 3.0, and others in early layers) attend from the second occurrence of a name token back to the first occurrence of the same name. Their OV circuit writes a “duplicated” signal into the residual stream of the second occurrence.

Formally, the key–query interaction of a DTH satisfies ${𝐖}_Q^{\mathrm{DTH}}{({𝐖}_K^{\mathrm{DTH}})}^{\top }\approx {𝐏}_{\mathrm{name}}$, where ${𝐏}_{\mathrm{name}}$ is a projection onto the subspace encoding proper-noun identity. As a result, position $S_2$ (second occurrence of subject $S$) attends strongly to position $S_1$ (first occurrence), and not to the single occurrence of $IO$.

Previous Token Heads in the IOI Circuit

The IOI circuit also contains previous token heads analogous to those discussed in Previous Token Heads. Here they serve the specific purpose of propagating name information from the name-token position to the token following the name (e.g., from “John” to “and” or from “Mary” to the verb phrase that follows). This enables later heads to attend to the name indirectly via its successor token, which is important for the sentence structures encountered in IOI templates.

Formal Definition of a Circuit

Activation Patching: Verifying Causal Necessity

Demonstrating that a circuit is causally necessary-not merely correlated with the desired behaviour-requires an intervention methodology. The standard technique in the IOI paper (and the broader mechanistic interpretability literature) is activation patching (also called causal tracing or interchange intervention).

The activation patching procedure partitions the causal effect of a node into the portion attributable to the information it processes (rather than its architectural position), making it more interpretable than simple ablation.

In the IOI analysis, activation patching was applied systematically to every attention head in every layer of GPT-2 Small, producing a heat-map of causal effects across the 12$\times$12 = 144 heads. The heads with the largest effects are precisely the NMHs, SIHs, DTHs, and previous token heads identified in the circuit-providing strong evidence that these heads are causally, not merely correlationally, important.

Step-by-Step Circuit Walkthrough

We now trace the IOI circuit on the sentence:

Let $S=$ John (subject, appears at positions 1 and 9) and $IO=$ Mary (indirect object, appears at position 3).

Step 1: Duplicate Token Heads detect the repeated name.

At position 9 (the second John), the DTHs attend strongly to position 1 (the first John), because both positions share the same name embedding. This writes a “duplicate” signal into the residual stream at position 9, flagging it as the second occurrence of a name.

Step 2: Previous Token Heads propagate context.

At positions 2 and 10 (the tokens immediately following each John), the previous token heads write information about John into the residual stream. Similarly, at position 4 (following Mary), the previous token head writes information about Mary. This propagation ensures that later heads have access to name identity from both the token itself and its successor.

Step 3: S-Inhibition Heads suppress the subject.

At the final query position (the token “to,” which precedes the predicted token), the SIHs in layers 7–8 scan backwards. Their attention weights are concentrated on position 9 (the second, most recent, occurrence of John), particularly because position 9 carries the “duplicate” signal from Step 1. The SIHs' OV circuit writes a negative signal to the residual stream that reduces the logit of the token John.

Step 4: Name Mover Heads copy the indirect object.

At the final position, the NMHs in layers 9–10 attend to position 3 (Mary). Their attention is guided by the absence of a duplicate signal at position 3: Mary appears only once, so the NMHs' QK circuit selects it over the (now-suppressed) John. The OV circuit of the NMHs copies the Mary embedding into the final residual stream.

Step 5: Unembedding.

The final residual stream, boosted toward Mary by the NMHs and away from John by the SIHs, is projected through the unembedding matrix ${𝐖}_U$. The result is ${\ell }_{\mathrm{Mary}}>{\ell }_{\mathrm{John}}$, i.e., the model correctly predicts Mary as the continuation.

Warnings and Limitations

Broader Implications for Interpretability Research

The IOI circuit contributes several lasting methodological lessons to the interpretability programme:

Functional decomposition is possible at scale.

Despite GPT-2 Small's 144 attention heads and 12 MLP sublayers, a mere 26 heads are causally sufficient for the IOI task. This suggests that transformers do not distribute all computations diffusely across all components, but concentrate specific functions in identifiable subsets.

Algorithm identification bridges symbolic and neural views.

The four-step algorithm identified in Step-by-Step Circuit Walkthrough (detect duplicate, propagate context, suppress subject, copy object) is a symbolic description of a computation performed by continuous, high-dimensional neural network activations. This bridging of levels is precisely what mechanistic interpretability aspires to achieve for all model behaviours.

Causal methodology is essential.

The combination of activation patching, attention visualisation, and ablation provides converging evidence that is more compelling than any single method alone. Correlation-based analyses (attention visualisation without patching) would have identified the same heads as “interesting,” but could not establish causal necessity. The IOI paper's causal methodology is now standard in circuit-finding work.

Safety relevance.

The IOI task is innocuous, but the methodology it validated is directly applicable to safety-relevant behaviours: deception, refusal, sycophancy, hallucination. If a model's tendency to produce false confident statements is implemented by a specific circuit, that circuit can potentially be identified, patched, or monitored. This is one of the primary motivations for the broader mechanistic interpretability programme.

Exercises

The Superposition Hypothesis

One of the most striking findings in mechanistic interpretability is also one of the most counterintuitive: individual neurons are not individual concepts. Open a dictionary and each entry denotes exactly one word. Look up a neuron in a transformer and you will find it responds to dozens of unrelated stimuli - currency symbols, geographical locations, names of chemical elements, expressions of negation. This is not a flaw in how we measure neurons; it reflects something deep about how neural networks choose to organise information internally.

The superposition hypothesis, developed at Anthropic in a landmark series of papers on toy models of superposition, proposes a precise explanation: neural networks store more features than they have neurons by packing those features into nearly orthogonal directions in activation space. The inevitable price is interference - features “bleed” into one another whenever two features that share a neuron happen to be active simultaneously - but if features are sparse (rarely co-active), the expected interference is small and the gain in representational capacity is large.

This section develops the mathematics of superposition rigorously. We begin with the formal definition (Formal Definition of Superposition), then study a toy model that makes the phenomenon concrete and analytically tractable (The Toy Model of Superposition). We connect the geometry to the celebrated Johnson–Lindenstrauss lemma (The Johnson–Lindenstrauss Lemma and Superposition Capacity), derive the phase diagram governing when a network chooses to superpose a feature (Phase Transitions in Superposition), and close with a striking biological analogy: sparse coding in the mammalian visual cortex (Biological Connection: Sparse Coding in Visual Cortex).

Why Neurons Are Polysemantic

Consider a language model processing a sentence about Paris. The model must track, among other things, the grammatical role of “Paris”, its semantic category (proper noun, city, capital), its cultural associations (Eiffel Tower, haute cuisine, fashion week), and its position in the discourse. In a network with $d$ neurons in a hidden layer, the model must represent all of these features - and thousands more - using vectors in ${\mathrm{ℝ}}^d$.

The naive solution is monosemanticity: assign one neuron to each feature. This is clean, interpretable, and utterly impractical. Modern language models process tokens drawn from vocabularies of tens of thousands, appearing in contexts that invoke hundreds of thousands of distinct semantic properties. No realistic hidden dimension $d$ can accommodate a dedicated neuron per feature.

The fundamental tension is between two desiderata. First, the network wants representational capacity: the ability to track many features simultaneously. Second, it wants interpretability: a clean mapping between neurons and concepts. Superposition is the network's resolution to this tension, sacrificing interpretability for capacity.

To see why this trade-off arises naturally from gradient descent, recall that a neural network's loss function cares only about predictions, not about the interpretability of intermediate representations. If packing two features into one neuron reduces loss (because both features are useful but rarely co-active), gradient descent will discover this packing automatically. There is no mechanism imposing monosemanticity.

Formal Definition of Superposition

We formalise the notion of superposition in terms of the linear representation hypothesis.

The Toy Model of Superposition

To make the superposition hypothesis analytically tractable, Elhage et al. (2022) introduced an elegant toy model: a one-hidden-layer autoencoder that attempts to reconstruct $n$ features using only $d<n$ dimensions.

Model architecture.

Let $𝐟\in {\mathrm{ℝ}}^n$ be a sparse feature vector, where each component $f_i$ is nonzero with probability $p_i$ (the feature's “sparsity”) and, when active, takes value drawn from some distribution (for simplicity, $f_i\in \{0,1\}$ with $\mathrm{Pr}[f_i=1]=p_i$). The model maps $𝐟$ through a linear bottleneck: (TOY Encode)$$\begin{array}{cccc}& 𝒙& =W𝐟+𝐛,& \text{<span class="tml-eqn"></span>}\\ & \widehat{𝐟}& =\mathrm{ReLU}(W^{\top }𝒙+{𝐛}^{\prime }),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $W\in {\mathrm{ℝ}}^{d\times n}$ is the weight matrix (the superposition dictionary), and the ReLU nonlinearity enforces non-negativity of the reconstructed features. The training objective is a weighted mean-squared error: (TOY LOSS)$$\begin{array}{ccc}& ℒ(W)={𝔼}_{𝐟}\left[\sum _{i=1}^n{S}_i\right(f_i-{\hat{f}}_i{)}^2],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $S_i>0$ is the importance of feature $i$ (proportional to how much the model cares about reconstructing it accurately).

What does the model learn?

Two extreme solutions exist. In the monosemantic solution, the model sets $d=n$ (impossible here since $d<n$) or selects the $d$ most important features and represents each in a dedicated dimension, ignoring the rest. In the superposition solution, the model represents all $n$ features using $n$ nearly-orthogonal directions in ${\mathrm{ℝ}}^d$, accepting small reconstruction errors whenever two features are simultaneously active.

The expected reconstruction error for feature $i$ under superposition is: (TOY Interference)$$\begin{array}{ccc}& 𝔼\left[\right(f_i-{\hat{f}}_i{)}^2]\approx \sum _{j\ne i}{p}_j⟨{𝒗}_i,{𝒗}_j{⟩}^2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where the right-hand side accumulates squared inner products weighted by the co-activation probability $p_j$. When features are sparse ($p_j\ll 1$ for all $j$), this error is small even if the ${𝒗}_j$ are only approximately orthogonal.

Optimal geometry.

A natural question is: what arrangement of $n$ unit vectors in ${\mathrm{ℝ}}^d$ minimises the total squared interference ${\sum }_{i\ne j}{⟨{𝒗}_i,{𝒗}_j⟩}^2$? This is the Tammes problem (also called the Grassmannian packing problem). For $n=d+1$, the optimal solution is the regular simplex: vertices of the regular $(d+1)$-simplex inscribed in the unit sphere $S^{d-1}$. For larger $n$, the optimal packings are known for special cases (e.g., $n=2d$ gives the cross-polytope) but are generally NP-hard to compute.

The toy model trained by gradient descent nonetheless discovers near-optimal solutions, suggesting that the loss landscape has favourable geometry for superposition.

The Johnson–Lindenstrauss Lemma and Superposition Capacity

How many features can a $d$-dimensional space accommodate in superposition? The answer is surprisingly large, and its theoretical foundation is the celebrated Johnson–Lindenstrauss (JL) lemma.

The JL lemma (in its original form) states that any set of $m$ points in ${\mathrm{ℝ}}^n$ can be embedded into ${\mathrm{ℝ}}^d$ (with $d=O({ϵ}^{-2}\log m)$) while preserving all pairwise distances to within a factor of $(1\pm ϵ)$. The connection to superposition is direct: if we regard features as points on the unit sphere, a random projection into $d$ dimensions (the “compression” step) preserves their near- orthogonality, exactly the geometric property that makes superposition work.

The implication for neural networks is profound: the weight matrix $W$ of a linear layer can act as a random projection, preserving the approximate geometry of the feature space even when the bottleneck dimension $d\ll n$. Superposition is, in a sense, the neural network's implicit implementation of the Johnson–Lindenstrauss construction.

Phase Transitions in Superposition

The toy model of The Toy Model of Superposition reveals a sharp phase transition: depending on the importance $S_i$ and sparsity $p_i$ of a feature, the model either represents it cleanly (monosemantically) or superimposes it with other features, or ignores it entirely.

Three phases.

Consider a model with $d=1$ (one-dimensional hidden space) and $n=2$ features with importances $S_1\ge S_2$ and sparsities $p_1,p_2$. The model must choose how to allocate its single dimension. The optimal solution depends on the trade-off between:

1. Monosemantic: represent only feature 1, giving loss $S_2𝔼[f_2^2]=S_2{p}_2$.

2. Superposition: represent both features as ${𝒗}_1=1,{𝒗}_2=-1$ (antipodal), giving interference loss $(S_1+S_2)p_1{p}_2$ (the antipodal cross terms cost error only when both features co-activate, probability $\approx p_1{p}_2$).

3. Ignore both: loss $S_1{p}_1+S_2{p}_2$.

Superposition is preferred when the interference loss is lower than the monosemantic loss: (Phase COND)$$\begin{array}{ccc}& (S_1+S_2)p_1{p}_2<S_2{p}_2⟺p_1<\frac{S_2}{S_1+S_2}.& \text{<span class="tml-eqn"></span>}\end{array}$$ This condition is satisfied precisely when feature 1 is sparse ($p_1<S_2/(S_1+S_2)$), so that the interference it causes is rare enough that sharing the dimension still beats paying $S_2{p}_2$ to ignore feature 2.

General phase diagram.

In the general setting with $n$ features, the phase diagram lives in the space of $(S_i,p_i)$ pairs. Three regimes emerge:

• Clean representation (high importance, high density): the model allocates a dedicated dimension, achieving exact reconstruction.

• Superposition (moderate importance, low density): the model packs multiple features into shared directions, accepting small interference.

• Dropped feature (low importance): the model ignores the feature entirely, setting ${𝒗}_i\approx 0$.

Biological Connection: Sparse Coding in Visual Cortex

The superposition hypothesis does not arise in a vacuum: it has a remarkable parallel in neuroscience. In 1996, Bruno Olshausen and David Field published a seminal paper showing that the receptive fields of simple cells in the primary visual cortex (V1) can be derived from a sparse coding principle applied to natural images.

The Olshausen–Field model.

Given a patch $𝐈\in {\mathrm{ℝ}}^m$ of a natural image, Olshausen and Field proposed learning a dictionary $A\in {\mathrm{ℝ}}^{m\times n}$ (with $n>m$, an overcomplete dictionary) and a sparse code $𝐬\in {\mathrm{ℝ}}^n$ such that $𝐈\approx A𝐬$ with $𝐬$ sparse. The learning objective is: (Olshausen Field)$$\begin{array}{ccc}& \underset{A,𝐬}{\min }{\Vert 𝐈-A𝐬\Vert }^2+\lambda \sum _i\log (1+s_i^2/{\sigma }^2),& \text{<span class="tml-eqn"></span>}\end{array}$$ where the second term encourages sparsity in the code. When trained on natural image patches, the columns of $A$ spontaneously develop into oriented, localised, bandpass filters - precisely the Gabor-like receptive fields observed in V1 neurons.

The connection to superposition.

The parallel is striking:

• The dictionary $A$ plays the role of the superposition dictionary $W$, with overcomplete columns ($n>m$) analogous to $n>d$.

• The sparse code $𝐬$ plays the role of the feature activations $𝐟$, with sparsity enabling efficient representation.

• The Gabor-like columns of $A$ correspond to monosemantic feature directions: each column encodes a single oriented edge at a specific spatial frequency.

The brain appears to have independently discovered the same solution to the representational bottleneck that gradient descent finds in artificial neural networks: pack more features than dimensions by exploiting sparsity.

Exercises

Sparse Autoencoders for Feature Discovery

If the superposition hypothesis is correct, then the fundamental challenge of mechanistic interpretability is not merely to find what neurons do, but to disentangle the superimposed features from the directions in which they are packed. The neuron basis is the wrong basis; we need to find the feature basis.

This insight motivates a deceptively simple idea: train a sparse autoencoder (SAE) on the model's activations. The autoencoder is forced to reconstruct the activations using a small number of active dictionary elements at each forward pass. If the training is successful, those dictionary elements will correspond to the monosemantic features that the network has superimposed - the SAE effectively inverts the superposition and recovers the feature basis.

The field of SAE-based interpretability has exploded since Cunningham et al. (2023) and the Anthropic team (Templeton et al., 2024) demonstrated that SAEs trained on transformer activations discover human-interpretable monosemantic features at scale. This section develops the theoretical foundations, practical training considerations, and empirical findings of this methodology.

SAE Architecture

A sparse autoencoder is a two-layer neural network with an expansion factor: the hidden dimension $m$ is much larger than the input dimension $d$ (i.e., $m\gg d$). The key constraint is that the hidden representation must be sparse: only a small fraction of the $m$ hidden units are active for any given input.

Dictionary elements.

The columns of $W_{\mathrm{dec}}$, denoted ${𝐝}_j\in {\mathrm{ℝ}}^d$ for $j=1,\dots ,m$, are the dictionary elements of the SAE. Each dictionary element is a direction in activation space. When the SAE is trained successfully, each ${𝐝}_j$ should correspond to a monosemantic feature: a single, human-interpretable concept.

The encoder computes a non-negative coefficient $h_j$ for each dictionary element, with $h_j>0$ indicating that feature $j$ is active in the current context. The decoded reconstruction $\widehat{𝒙}={\sum }_j{h}_j{𝐝}_j+{𝐛}_{\mathrm{pre}}$ expresses the activation as a sparse linear combination of dictionary elements.

Training Objective

The SAE training objective balances two competing goals: reconstruction fidelity (the decoded $\widehat{𝒙}$ should match the original $𝒙$ closely) and sparsity (only a few dictionary elements should be active at each step).

Normalisation of decoder columns.

A crucial subtlety: the decoder columns ${𝐝}_j$ must be constrained to have unit norm. Without this constraint, the optimiser can trivially reduce the ${\ell }_1$ sparsity penalty by rescaling: make ${𝐝}_j$ very large (absorbing the scale into ${𝐝}_j$) and make $h_j$ very small (reducing ${\Vert 𝐡\Vert }_1$). Unit-norm columns prevent this degeneracy. In practice, this is enforced by projecting $W_{\mathrm{dec}}$ onto the Stiefel manifold after each gradient step: (SAE NORM)$$\begin{array}{ccc}& {({𝐝}_j)}_{\text{new}}\leftarrow \frac{{({𝐝}_j)}_{\text{old}}}{{\Vert {({𝐝}_j)}_{\text{old}}\Vert }_2}.& \text{<span class="tml-eqn"></span>}\end{array}$$

Hyperparameter sensitivity.

The sparsity coefficient $\lambda$ is a critical hyperparameter. Too small: the SAE learns a near-dense representation with high reconstruction fidelity but uninterpretable features (the trivial solution is to learn the identity mapping with $m=d$). Too large: the SAE learns to use very few features, sacrificing reconstruction fidelity to the point where many input dimensions are ignored. The optimal $\lambda$ must be found by balancing reconstruction loss and sparsity, often using metrics like $L_0$ (average number of active features per token) and reconstruction $R^2$.

The Shrinkage Problem

A fundamental limitation of ${\ell }_1$-penalised estimation, long recognised in the statistics literature under the name LASSO shrinkage, is that it systematically underestimates the magnitude of large coefficients.

The bias of LASSO.

Suppose the true feature activation for feature $j$ in context $𝒙$ is $h_j^{\ast }$. Under the ${\ell }_1$ penalty, the optimal estimate ${\hat{h}}_j$ is the soft-thresholding operator applied to the pre- activation ${\tilde{h}}_j$: (SOFT Threshold)$$\begin{array}{ccc}& {\hat{h}}_j={\mathrm{ST}}_{\lambda }({\tilde{h}}_j):=\max (0,{\tilde{h}}_j-\lambda ),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\mathrm{ST}}_{\lambda }$ is the soft-thresholding function with threshold $\lambda$. The bias is $𝔼[{\hat{h}}_j]-h_j^{\ast }=-\lambda$ whenever ${\tilde{h}}_j>\lambda$ (the feature is active). This means:

1. Active features have their magnitudes systematically underestimated by $\lambda$.

2. The SAE must compensate by adjusting the decoder columns ${𝐝}_j$, potentially distorting the feature geometry.

3. Features with true activation $h_j^{\ast }<\lambda$ are entirely suppressed to zero, even if they are genuinely active.

Gated SAEs: Separating Gating from Magnitude

The shrinkage problem arises because a single value ($h_j$) must simultaneously encode two distinct decisions: (1) is feature $j$ active? (a binary gating decision) and (2) how strongly is feature $j$ active? (a continuous magnitude). The ${\ell }_1$ penalty acts on the magnitude, but its effect is felt through the gating decision as well.

Gated SAEs (Rajamanoharan et al., 2024) resolve this by introducing separate pathways for gating and magnitude estimation.

Loss modification.

The gated SAE uses a modified loss that places the sparsity penalty on the gate logits rather than the final activations: (Gated LOSS)$$\begin{array}{ccc}& {ℒ}_{\mathrm{gated}}={\Vert 𝒙-\widehat{𝒙}\Vert }_2^2+\lambda {\Vert \mathrm{ReLU}(\widetilde{𝐡})\Vert }_1,& \text{<span class="tml-eqn"></span>}\end{array}$$ The sparsity penalty acts on $\mathrm{ReLU}(\widetilde{𝐡})$ (the smooth version of the gate) rather than on the magnitude activations, decoupling the gating and magnitude estimation problems.

Empirical improvements.

Gated SAEs consistently outperform standard SAEs on the reconstruction–sparsity Pareto frontier: for the same level of sparsity (average $L_0$), gated SAEs achieve lower reconstruction error, and for the same reconstruction error, they achieve lower $L_0$. The interpretability of discovered features is also qualitatively improved, suggesting that clean gating leads to cleaner feature separation.

Transcoders: Mapping MLP Input to Output

Standard SAEs are trained on residual stream states: the activations $𝒙$ before or after a given layer. An alternative, proposed by Dunefsky et al. (2024), is the transcoder: a sparse model that maps the input to an MLP layer directly to its output.

Advantages of transcoders.

The transcoder provides a fundamentally different perspective from residual-stream SAEs. Whereas a residual-stream SAE decomposes “what information is stored at a given position”, a transcoder decomposes “what computation the MLP performs”. Specifically, the dictionary elements ${𝐝}_j$ of the transcoder are input-output pairs: active feature $j$ contributes $h_j{𝐝}_j$ to the output, where ${𝐝}_j\in {\mathrm{ℝ}}^d$ describes the direction in which that feature updates the residual stream.

This makes transcoders particularly suitable for understanding knowledge storage in MLP layers, where factual associations (e.g., “Paris is the capital of France”) are believed to be stored as superimposed key-value pairs (Meng et al., 2022 - ROME).

SAE Architecture: A Visual Overview

Monosemantic Features: What SAEs Discover

The proof of the SAE methodology lies in what it finds. When SAEs are trained on the internal activations of large language models, the dictionary elements ${𝐝}_j$ exhibit striking monosemanticity: each element responds strongly to a specific, human-interpretable concept and weakly to everything else.

Categories of discovered features.

SAE features can be broadly grouped into three categories:

1. Geographic features. Features that activate strongly for tokens that are part of geographical references - country names, city names, continent names, or demonyms. Different features distinguish different granularities and types: one feature may activate for “European capital city names”, another for “South American river names”, a third for “US state abbreviations”.

2. Temporal features. Features encoding temporal information: specific decades (“1990s”, “2000s”), recurring events (“every year”, “annual”), temporal expressions (“yesterday”, “next month”), and calendar dates. Temporal features often have a strong interaction with domain: “1990s technology” and “1990s music” may activate different features despite sharing the decade marker.

3. Conceptual and relational features. Features for abstract semantic relationships: negation, causality, analogy, contrast, and entailment. Also present are syntactic features (subject-verb agreement, genitive markers), domain markers (medical terminology, legal jargon, code syntax), and emotion/sentiment features.

Validation of monosemanticity.

How do we verify that a discovered feature is genuinely monosemantic? The standard approach is automatic interpretability (Bills et al., 2023): use a separate language model (e.g., GPT-4) to generate a natural-language description of the feature by examining the top-activating tokens and contexts, then evaluate whether the description correctly predicts which new contexts will activate the feature. A feature is considered monosemantic if a single short description achieves high prediction accuracy.

Templeton et al. (2024) applied this methodology at scale to Claude 3 Sonnet, finding that the majority of the 1,000,000+ discovered SAE features received high-quality monosemantic descriptions from the automated evaluator.

Scaling SAEs: Dictionary Size, Dead Features, and Fidelity

Training SAEs at scale introduces new challenges that do not appear in small-scale experiments. Three phenomena dominate the scaling behaviour of SAEs: the dictionary size, the dead feature problem, and the reconstruction fidelity trade-off.

Dictionary size and feature granularity.

As the dictionary size $m$ increases, the SAE discovers progressively more fine-grained features. With $m=512$, an SAE trained on a language model might discover a single “European country” feature. With $m=\mathrm{16,384}$, it discovers separate features for each European country. With $m=\mathrm{131,072}$, it distinguishes not only countries but also contexts (“France as a tourist destination”, “France in the context of wine”, “France in WWII context”).

This suggests a feature hierarchy: broader features at small dictionary sizes, finer-grained features at larger sizes. The optimal dictionary size depends on the application: larger dictionaries provide richer interpretability but are more expensive to train and interpret.

Dead features.

A common failure mode in SAE training is the emergence of dead features: dictionary elements ${𝐝}_j$ that are never (or very rarely) activated across the training distribution. Dead features waste model capacity and can arise from:

• Initialisation failures: random initialisation places ${𝐝}_j$ in a region of the unit sphere that no training example maps near, and the gradient updates never pull it into an active region.

• Gradient starvation: if a feature is rarely active, it receives few gradient updates, leading to slow learning and eventual death.

• Supercession: a similar feature ${𝐝}_{j^{\prime }}$ with slightly better initialisation captures the same concept, causing ${𝐝}_j$ to become redundant.

Dead features can be mitigated through auxiliary losses (Anthropic's “AuxK” loss, which explicitly penalises the magnitude of the top-$k$ activations of dead features), resampling (periodically reinitialising dead features to activations of training examples not well-explained by active features), or $k$-sparse autoencoders (TopK SAEs, which use a fixed-$k$ sparse activation instead of the ReLU-plus-L1 approach).

Reconstruction fidelity.

The quality of an SAE is measured by how well the reconstructed $\widehat{𝒙}$ matches the original $𝒙$. Two metrics are standard:

1. Variance explained ($R^2$): the fraction of total variance in $𝒙$ explained by $\widehat{𝒙}$. (SAE R2)$$\begin{array}{ccc}& R^2=1-\frac{𝔼[{\Vert 𝒙-\widehat{𝒙}\Vert }_2^2]}{𝔼[{\Vert 𝒙\Vert }_2^2]}.& \text{<span class="tml-eqn"></span>}\end{array}$$ A well-trained SAE achieves $R^2>0.99$ on the training distribution.

2. Loss recovered: the fraction of the model's language modelling loss that is preserved when the residual stream is replaced by $\widehat{𝒙}$. If the model's clean cross-entropy loss is ${ℒ}_0$ and the loss with the SAE intervention is ${ℒ}_{\mathrm{SAE}}$, then: (SAE LOSS Recovered)$$\begin{array}{ccc}& \text{Loss recovered}=1-\frac{{ℒ}_{\mathrm{SAE}}-{ℒ}_0}{{ℒ}_{\mathrm{SAE}}-{ℒ}_{\mathrm{zero}}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${ℒ}_{\mathrm{zero}}$ is the loss when the residual stream is zeroed out. Loss recovered $>0.95$ indicates that the SAE captures the vast majority of the causally relevant information.

Discovered Features: A Visual Taxonomy

Example: Training an SAE on GPT-2

To ground the preceding theory, we trace through a concrete end-to-end example: training a sparse autoencoder on the residual stream activations of GPT-2 small (117M parameters) after the 8th layer.

Setup.

GPT-2 small has a hidden dimension of $d=768$. We train an SAE with dictionary size $m=\mathrm{16,384}$ (expansion factor $\approx 21$) and sparsity coefficient $\lambda =4\times {10}^{-5}$. The SAE is trained on 1 billion tokens of text from the OpenWebText corpus using the Adam optimiser with learning rate ${10}^{-4}$.

Training dynamics.

During early training (first $\mathrm{100,000}$ steps), the SAE quickly learns a near-identity mapping: the reconstruction loss drops sharply while the sparsity penalty remains high. By step $\mathrm{500,000}$, a bimodal pattern emerges in the feature activation histogram: most features are inactive (activating fewer than 0.1% of the time), while a smaller number of features activate frequently. The dead feature rate is approximately 30% at this stage.

After applying the AuxK auxiliary loss and running for $1.5$ million total steps, the dead feature rate drops to approximately 5%, $R^2=0.987$, and the average $L_0$ (number of active features per token) stabilises at around 64.

Inspecting top features.

We examine the five highest-activating tokens for selected dictionary elements:

• Feature 2314: activates on “Paris”, “Lyon”, “Marseille”, “Bordeaux”, “Toulouse” - clearly a “French cities” feature.

• Feature 7891: activates on “DNA”, “RNA”, “nucleotide”, “genome”, “chromosome” - a “molecular biology terminology” feature.

• Feature 11203: activates on “however”, “nevertheless”, “nonetheless”, “yet”, “although” - an “adversative/concessive conjunction” feature.

• Feature 3047: activates on “def”, “class”, “import”, “return”, “for” - a “Python syntax keyword” feature.

• Feature 9156: activates on “1990”, “1992”, “1994”, “1996”, “1998” - a “1990s year” feature.

Each of these features is strikingly monosemantic: a single description captures all top-activating tokens. Crucially, the features were not engineered; they emerged spontaneously from training on next-token prediction with a sparsity constraint.

Causal intervention.

To verify that the features are causally relevant (not merely correlational), we perform an activation patching experiment: we replace the feature activation $h_{2314}$ (“French cities”) with its maximum observed value in a forward pass on an unrelated context (e.g., a passage about computer science). The model's completion probability for French-related tokens increases significantly, confirming that feature 2314 causally influences the model's predictions about French cities.

Exercises

Knowledge Storage and Factual Editing

How does a language model know that the Eiffel Tower is in Paris? The question sounds almost trivial, yet answering it rigorously requires us to open the model's weights and trace, neuron by neuron, which parameters encode the association between the subject Eiffel Tower, the relation is located in, and the object Paris. Over the past several years, a branch of mechanistic interpretability has done exactly this, developing tools that localise factual knowledge to specific weight matrices, quantify how strongly each layer contributes to a given completion, and, most strikingly, surgically modify stored facts without retraining the entire model.

The practical motivation is compelling. A language model trained in 2023 encodes the fact that a particular country's prime minister is a particular person. By 2025, that fact may have changed. Rather than retraining the full model, a knowledge-editing algorithm inserts the updated fact at the precise location where the old one is stored. The operation is targeted, fast, and-when done correctly-does not disrupt thousands of unrelated facts that surround the edited one in weight space.

This section develops the mathematical foundations for all of these operations. We begin with the empirical discovery that localised the problem, proceed through the formal framework of MLP key-value memories, and culminate in the ROME and MEMIT algorithms with full mathematical derivations.

The ROME Hypothesis: Facts Live in Middle-Layer MLPs

The story begins with a deceptively simple intervention experiment. Given a factual prompt such as “The Eiffel Tower is located in the city of”, Meng et al. (2022) asked: which activations, if corrupted, would prevent the model from completing the sentence with Paris? Their technique, called causal tracing, works as follows.

Run the model three times on the same prompt:

1. Clean run. Process the prompt normally; record every hidden state ${𝒉}_i^{(l)}$ at layer $l$, token position $i$.

2. Corrupted run. Replace the subject tokens with a corrupted version (e.g. Gaussian noise is added to their embeddings) and run the model again; record all states.

3. Restored run. For each candidate state $(l,i)$, take the corrupted run but patch in the clean state at position $(l,i)$; measure whether the clean prediction is recovered.

Formally, define the average indirect effect (AIE) of restoring state $(l,i)$ as (AIE)$$\begin{array}{ccc}& \mathrm{AIE}(l,i)={𝔼}_{\text{prompts}}[p_{\mathrm{restored}}(\text{correct})-p_{\mathrm{corrupt}}(\text{correct})],& \text{<span class="tml-eqn"></span>}\end{array}$$ where the expectation is over a dataset of factual prompts. A large value of $\mathrm{AIE}(l,i)$ means that restoring state $(l,i)$ from the clean run substantially recovers the correct prediction, i.e., that state $(l,i)$ is causally necessary for the fact.

The empirical finding from causal tracing across GPT-2 XL and GPT-J is striking and reproducible:

• The early attention layers at subject tokens contribute modestly by spreading subject information across the context.

• The middle-layer MLP modules (roughly layers 5–15 in a 28-layer model) at the last subject token exhibit dramatically elevated AIE.

• Restoring a single MLP hidden state at the right location can recover essentially the full probability of the correct object.

MLPs as Key-Value Memory Stores

To understand why facts live in MLP weights, we need to interpret the two-layer MLP formally as an associative memory. Consider the standard transformer MLP block: (MLP Block)$$\begin{array}{ccc}& \mathrm{MLP}(𝒉)={𝐖}_{\mathrm{out}}\sigma \left({𝐖}_{\mathrm{in}}𝒉\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐖}_{\mathrm{in}}\in {\mathrm{ℝ}}^{d_{\mathrm{ff}}\times d}$ and ${𝐖}_{\mathrm{out}}\in {\mathrm{ℝ}}^{d\times d_{\mathrm{ff}}}$ are the weight matrices, $\sigma$ is the activation function (GELU or ReLU), and $𝒉\in {\mathrm{ℝ}}^d$ is the residual-stream vector at the current position.

Rewrite this as (MLP KV Expansion)$$\begin{array}{ccc}& \mathrm{MLP}(𝒉)=\sum _{j=1}^{d_{\mathrm{ff}}}\sigma (⟨{𝒘}_{\mathrm{in},j},𝒉⟩)\cdot {𝒘}_{\mathrm{out},j},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒘}_{\mathrm{in},j}$ is the $j$-th row of ${𝐖}_{\mathrm{in}}$ and ${𝒘}_{\mathrm{out},j}$ is the $j$-th column of ${𝐖}_{\mathrm{out}}$.

The analogy to an associative memory is now transparent. Each “neuron” $j$ stores:

• a key ${𝒌}_j\stackrel{\mathrm{def}}{=}{𝒘}_{\mathrm{in},j}$: a direction in residual-stream space that the neuron “attends to”;

• a value ${𝒗}_j\stackrel{\mathrm{def}}{=}{𝒘}_{\mathrm{out},j}$: a direction in residual-stream space that is written when the neuron fires.

The scalar $\sigma (⟨{𝒌}_j,𝒉⟩)$ is the attention coefficient: it is large when the input $𝒉$ aligns with key ${𝒌}_j$, causing value ${𝒗}_j$ to be added to the residual stream.

The key-value interpretation makes a precise prediction: if we want to change the stored object from $o$ to $o^{\ast }$ (e.g. from Paris to London), we must find the dominant key-value pair in the target MLP layer that encodes the current association and replace its value with a new vector ${𝒗}^{\ast }$ that promotes $o^{\ast }$.

The ROME Algorithm

ROME (Rank-One Model Editing) implements the key-value replacement idea with a rank-one update to a single MLP weight matrix. We present the full derivation.

Setup

Let $l^{\ast }$ be the target MLP layer identified by causal tracing. We focus on the second weight matrix of that MLP, its output (down-projection), $𝐖\stackrel{\mathrm{def}}{=}{𝐖}_{\mathrm{out}}^{(l^{\ast })}\in {\mathrm{ℝ}}^{d\times d_{\mathrm{ff}}}$. The goal is to find an update $\widehat{𝐖}=𝐖+\mathrm{\Delta }𝐖$ such that (ROME Constraint NEW)$$\begin{array}{ccccc}& \widehat{𝐖}{𝒌}^{\ast }& ={𝒗}^{\ast }& \text{(new fact inserted)},& \text{<span class="tml-eqn"></span>}\\ & \widehat{𝐖}𝒌& =𝐖𝒌& \text{for all }𝒌\notin \{{𝒌}^{\ast }\},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒌}^{\ast }\in {\mathrm{ℝ}}^{d_{\mathrm{ff}}}$ is the key vector associated with the subject and ${𝒗}^{\ast }\in {\mathrm{ℝ}}^d$ is the target value vector associated with the new object $o^{\ast }$.

Constraint cannot hold for all vectors simultaneously unless $\mathrm{\Delta }𝐖=0$, so we relax it: we require it to hold in the direction of every key that appears in a set of preservation memoranda ${\{{𝒌}_i\}}_{i=1}^N$, a sample from the space of contexts unrelated to the edited fact.

Computing the Key Vector ${𝒌}^{\ast }$k*

The key vector ${𝒌}^{\ast }$ represents the MLP's intermediate representation when processing the subject $s$ in the prompt. Concretely:

1. Run the transformer forward on the prompt template $𝚙(s,r)$.

2. Extract the hidden state at the last subject token at the input to layer $l^{\ast }$, i.e., ${𝒉}_{i_s}^{(l^{\ast }-1)}$ where $i_s$ is the last token index of $s$.

3. Apply the MLP input projection: $${𝒌}^{\ast }=\sigma \left({𝐖}_{\mathrm{in}}^{(l^{\ast })}{𝒉}_{i_s}^{(l^{\ast }-1)}\right)\in {\mathrm{ℝ}}^{d_{\mathrm{ff}}}.$$

To make ${𝒌}^{\ast }$ robust to paraphrases of $s$, ROME averages over several prompt templates: (Kstar)$$\begin{array}{ccc}& {𝒌}^{\ast }=\frac{1}{|𝒫|}\sum _{p\in 𝒫}\sigma ({𝐖}_{\mathrm{in}}^{(l^{\ast })}{𝒉}_{i_s}^{(l^{\ast }-1)}(p)),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒫$ is a small set of paraphrase templates.

Computing the Value Vector ${𝒗}^{\ast }$v*

The target value vector ${𝒗}^{\ast }$ must, when multiplied by the key, cause the residual stream to subsequently produce the correct token $o^{\ast }$. This is solved by an optimisation problem. Let ${𝒛}^{(l^{\ast })}$ be the output of layer $l^{\ast }$'s MLP; the rest of the network (layers $l^{\ast }+1$ through $L$) maps this to a distribution over tokens. Define the target probability objective: (Vstar Objective)$$\begin{array}{ccc}& {𝒗}^{\ast }={\mathrm{arg\hspace{0.17em}max}}_{𝒗}\log p_{\theta }(o^{\ast }|𝚙(s,r);\text{layer }{l}^{\ast }\text{ output}=𝐖{𝒌}^{\ast }+(𝒗-𝐖{𝒌}^{\ast })).& \text{<span class="tml-eqn"></span>}\end{array}$$ In practice, ${𝒗}^{\ast }$ is computed by gradient descent on the output embedding of layer $l^{\ast }$: one freezes all layers except the “virtual” output of layer $l^{\ast }$ and optimises to maximise the log-probability of $o^{\ast }$ at the prediction head.

Rank-One Update

With ${𝒌}^{\ast }$ and ${𝒗}^{\ast }$ in hand, the update is derived as follows. We seek $\mathrm{\Delta }𝐖$ of minimal Frobenius norm such that: (ROME Constrained)$$\begin{array}{ccc}& (𝐖+\mathrm{\Delta }𝐖){𝒌}^{\ast }={𝒗}^{\ast },(𝐖+\mathrm{\Delta }𝐖)𝐊=𝐖𝐊,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝐊=[{𝒌}_1,\dots ,{𝒌}_N]\in {\mathrm{ℝ}}^{d_{\mathrm{ff}}\times N}$ is the matrix of preservation keys.

The second constraint fixes what the weight does on the column space of $𝐊$. We use the least-squares solution in the $𝐂=𝐊{𝐊}^{\top }$ metric (covariance of the key distribution), which amounts to computing: (ROME Lambda)$$\begin{array}{cccc}& 𝝀& ={𝐂}^{-1}{𝒌}^{\ast },& \text{<span class="tml-eqn"></span>}\\ & \mathrm{\Delta }𝐖& =\frac{({𝒗}^{\ast }-𝐖{𝒌}^{\ast }){𝝀}^{\top }}{{𝝀}^{\top }{𝒌}^{\ast }}=\frac{({𝒗}^{\ast }-𝐖{𝒌}^{\ast }){({𝐂}^{-1}{𝒌}^{\ast })}^{\top }}{{({𝐂}^{-1}{𝒌}^{\ast })}^{\top }{𝒌}^{\ast }}.& \text{<span class="tml-eqn"></span>}\end{array}$$ This is a rank-one matrix: $\mathrm{\Delta }𝐖=𝒖{𝒘}^{\top }$ where $𝒖={𝒗}^{\ast }-𝐖{𝒌}^{\ast }$ (the “value residual”) and $𝒘={𝐂}^{-1}{𝒌}^{\ast }/({({𝐂}^{-1}{𝒌}^{\ast })}^{\top }{𝒌}^{\ast })$.

ROME Pipeline: From Causal Trace to Edit

Example: Editing a Geographic Fact

Evaluation Criteria: Specificity, Generality, Locality

Any knowledge-editing method must be evaluated along three orthogonal axes:

ROME achieves high specificity by the design of the rank-one update: because $\mathrm{\Delta }𝐖$ is rank-one and the preservation key covariance $𝐂$ covers the distribution of typical inputs, the update barely affects activations outside the subject token's direction. Generality is improved by averaging ${𝒌}^{\ast }$ over multiple prompt templates.

MEMIT: Mass-Editing Memory in a Transformer

ROME edits one fact at a time. When thousands of facts must be updated simultaneously (e.g. incorporating a year's worth of news events into a model), the sequential application of rank-one updates is insufficient: each update shifts the covariance $𝐂$, so the approximation degrades, and the locality constraint is increasingly violated.

MEMIT (Mass-Editing Memory in a Transformer, Meng et al. 2022) addresses this by:

1. Spreading edits across multiple layers. Rather than targeting a single layer $l^{\ast }$, MEMIT identifies a range of layers $\{l_1,\dots ,l_R\}$ (typically the same middle layers identified by causal tracing) and distributes the editing responsibility across them.

2. Batch optimisation. Given a batch of $M$ fact tuples ${\{(s_m,r_m,o_m^{\ast })\}}_{m=1}^M$, MEMIT solves a joint least-squares problem for all value residuals simultaneously.

MEMIT Objective

For a batch of $M$ edits, let ${𝒗}_m^{\ast }$ be the target value for edit $m$ and ${𝒌}_m^{\ast }$ the corresponding key. MEMIT distributes each value residual ${𝒓}_m\stackrel{\mathrm{def}}{=}{𝒗}_m^{\ast }-𝐖{𝒌}_m^{\ast }$ equally across $R$ layers: (Memit Residual)$$\begin{array}{ccc}& {𝒓}_m^{(l)}=\frac{1}{R}{𝒓}_m,l\in \{l_1,\dots ,l_R\}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Then for each layer $l$, the weight update solves the constrained least-squares problem: (Memit Update)$$\begin{array}{ccc}& \mathrm{\Delta }{𝐖}^{(l)}={𝐑}^{(l)}({𝐊}^{(l)}{)}^{\top }[{𝐊}^{(l)}({𝐊}^{(l)}{)}^{\top }+\alpha {𝐂}^{(l)}{]}^{-1},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐑}^{(l)}=[{𝒓}_1^{(l)},\dots ,{𝒓}_M^{(l)}]\in {\mathrm{ℝ}}^{d\times M}$, ${𝐊}^{(l)}=[{𝒌}_1^{\ast },\dots ,{𝒌}_M^{\ast }]\in {\mathrm{ℝ}}^{d_{\mathrm{ff}}\times M}$, ${𝐂}^{(l)}$ is the precomputed key covariance at layer $l$, and $\alpha >0$ is a regularisation coefficient controlling how much the update respects the locality constraint.

This is a standard ridge-regression solution and can be computed in closed form. The total number of parameters modified is $R\times d\times d_{\mathrm{ff}}$, but each $\mathrm{\Delta }{𝐖}^{(l)}$ has rank at most $M$, so for $M\ll d_{\mathrm{ff}}$ the updates remain low-rank.

Historical and Broader Context

Exercises

Knowledge Conflicts and Hallucination

A language model carries two kinds of knowledge simultaneously: the parametric knowledge encoded in its weights during training, and the contextual knowledge supplied by the input prompt at inference time. For the most part, these two knowledge sources agree. But sometimes they conflict: the model's weights say the population of a city is one number, while a retrieved document in the context says another. Sometimes neither source is correct. And sometimes the model fabricates a plausible-sounding but entirely false fact, an event that has come to be called hallucination.

This section develops a mechanistic account of all three phenomena. We ask: what circuits inside the transformer are responsible for resolving knowledge conflicts, and when those circuits malfunction, what goes wrong? The mechanistic perspective yields both a rigorous formal definition of hallucination and a set of principled interventions for reducing it.

Context-Memory Conflicts: The Competing Circuits View

Let $𝚌$ be a retrieved context document and $𝚚$ a question whose answer appears in $𝚌$ but contradicts the model's parametric knowledge. For example:

• Context: “According to recent reports, the CEO of Acme Corp is Jane Smith.”

• Question: “Who is the CEO of Acme Corp?”

• Parametric answer: John Doe (from training data).

• Contextual answer: Jane Smith (from context).

Mechanistic interpretability research has identified two competing circuits that process these two information sources:

• The parametric recall circuit: The same middle-layer MLP pathway identified in The ROME Hypothesis: Facts Live in Middle-Layer MLPs that routes subject-relation pairs to stored objects. This circuit fires based on the question tokens alone, ignoring the context.

• The in-context retrieval circuit: An attention-based pathway that reads from context positions using the question tokens as queries. This circuit was first characterised in the study of induction heads (Olsson et al., 2022) and copy suppression (McDougall et al., 2023).

The late-layer attention mechanism acts as a soft resolver: it computes a weighted combination of the two signals. Let ${𝒗}_{\mathrm{param}}$ and ${𝒗}_{\mathrm{ctx}}$ be the value vectors contributed by the parametric and contextual circuits, respectively. The resolved output is approximately: (Conflict Resolve)$$\begin{array}{ccc}& {𝒗}_{\mathrm{out}}\approx \alpha {𝒗}_{\mathrm{ctx}}+(1-\alpha ){𝒗}_{\mathrm{param}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\alpha \in [\mathrm{0,1}]$ is determined by the attention weights and depends on both the strength of the contextual signal and an implicit confidence in the parametric recall. When $\alpha \approx 1$ the model follows the context (desirable for RAG); when $\alpha \approx 0$ it ignores the context (problematic).

Empirically, the value of $\alpha$ is strongly influenced by:

• Salience of the context mention. If the relevant fact in the context is stated early, clearly, and multiple times, $\alpha$ increases.

• Strength of the parametric association. Highly frequent training facts (e.g. capital cities of major countries) produce large $\Vert {𝒗}_{\mathrm{param}}\Vert$, which can dominate even when context disagrees.

• Position in context. The notorious lost-in-the-middle effect (Liu et al., 2023) shows that facts near the beginning or end of a long context are retrieved more reliably than facts buried in the middle, suggesting that attention patterns limit the strength of ${𝒗}_{\mathrm{ctx}}$ for middle positions.

Chain-of-Thought and Factual Circuits

Chain-of-thought (CoT) prompting, which asks the model to reason step by step before giving a final answer, has a complex interaction with factual accuracy. On one hand, CoT enables the model to decompose multi-hop questions and retrieve each sub-fact through separate forward-pass computations. On the other hand, errors introduced in early reasoning steps can propagate and amplify, leading to a phenomenon called reasoning hallucination: the final answer is confidently wrong not because the stored facts are wrong, but because an intermediate reasoning step went astray.

The mechanistic picture is as follows. In a CoT setting, each generated token becomes part of the context for subsequent tokens. The key-value memory retrieval described in MLPs as Key-Value Memory Stores is re-applied at each new token position. If a generated intermediate token incorrectly names a subject (e.g. the model writes “the Louvre is in London” as a reasoning step), then when the model next processes that token as a subject, the parametric recall circuit will look up Louvre in memory and retrieve a different set of associations, compounding the error.

Hallucination as a Mechanistic Phenomenon

The popular press treats hallucination as a mysterious failure mode of AI systems. From the mechanistic interpretability perspective, it is a precisely characterisable event in the transformer's computational graph.

Several distinct mechanisms can cause hallucination:

Type 1: Spurious Co-occurrence

During training, subject $s$ co-occurred with object $o^{\prime }$ in many documents, even though the correct object is $o$. The MLP key vector ${𝒌}^{\ast }$ for $s$ has high cosine similarity with the key of $o^{\prime }$, causing the wrong value to be amplified: (Spurious)$$\begin{array}{ccc}& ⟨{𝒌}^{\ast },{𝒌}_{o^{\prime }}⟩\gg ⟨{𝒌}^{\ast },{𝒌}_o⟩.& \text{<span class="tml-eqn"></span>}\end{array}$$ Example: “The author of [popular novel] is [name of a different popular author who is frequently mentioned in the same reviews]”.

Type 2: Retrieval Head Failure

In the context-memory conflict model , the in-context retrieval circuit has $\alpha \approx 0$: the model ignores a context that contains the correct answer. This happens when the context is long, the relevant span is in the middle, or the question is phrased differently from the context sentence.

Type 3: Out-of-Distribution Subject

The subject $s$ appeared rarely in training data, so its key vector ${𝒌}^{\ast }$ has low norm or is near-orthogonal to all stored value vectors. The model defaults to generating a plausible continuation based on surface patterns (the relation $r$) rather than factual recall. This is why models hallucinate more for obscure entities than famous ones.

Knowledge Evolution During Training

Facts are not learned uniformly during training. The order and frequency of exposure dramatically affects what is stored and how strongly. Several regularities have been identified:

• Frequent facts are learned early. In a curriculum ordered by document date, the model quickly encodes facts that appear in thousands of documents. The MLP key-value pairs for such facts have large $\Vert {𝒌}^{\ast }\Vert$ and their value vectors point strongly toward the correct token.

• Rare facts are learned late and weakly. A fact that appears in only a handful of documents may never be strongly encoded; its value vector remains near the mean of the value distribution, producing uncertain predictions.

• Temporal displacement. If a fact changes over time (e.g. the president of a country), the model will encode the most recent version it saw most frequently, not necessarily the most recent one chronologically. A model trained on data from 2015–2023 may have seen a particular leader's name far more often in 2015–2020 documents than in 2021–2023 ones, causing it to prefer the older answer.

These dynamics can be formalised using the gradient flow perspective. Consider a simplified single-layer MLP with one fact $(s,r,o)$. The cross-entropy loss for this fact at training step $t$ is: (Training LOSS)$$\begin{array}{ccc}& {ℒ}_t=-\log \mathrm{softmax}({𝐖}_{\mathrm{out}}\sigma ({𝐖}_{\mathrm{in}}𝒉){)}_o.& \text{<span class="tml-eqn"></span>}\end{array}$$ The gradient with respect to ${𝐖}_{\mathrm{out}}$ is (Training Gradient)$$\begin{array}{ccc}& {\nabla }_{{𝐖}_{\mathrm{out}}}{ℒ}_t=-(1-p_t(o))\sigma ({𝐖}_{\mathrm{in}}𝒉){𝒆}_o^{\top }+\sum _{o^{\prime }\ne o}{p}_t(o^{\prime })\sigma ({𝐖}_{\mathrm{in}}𝒉){𝒆}_{o^{\prime }}^{\top },& \text{<span class="tml-eqn"></span>}\end{array}$$ where $p_t(o)$ is the current model probability and ${𝒆}_o$ is the one-hot embedding of object token $o$. The update has the form of a rank-one correction: the column ${𝒗}_o={({𝐖}_{\mathrm{out}})}_{:,o}$ is moved toward the current key activation. Thus each gradient step incrementally sharpens the key-value association. Infrequently presented facts receive few such steps and remain weakly encoded.

Catastrophic Forgetting and Its Mechanistic Basis

When a neural network is trained sequentially on two tasks or two corpora, knowledge from the first task is often erased by training on the second. This phenomenon, known as catastrophic forgetting (McCloskey and Cohen, 1989), has a clean mechanistic interpretation in the key-value memory framework.

Suppose the model has encoded a fact $(s_1,r,o_1)$ at layer $l^{\ast }$ via the key-value pair $({𝒌}_1,{𝒗}_1)$. Now we fine-tune the model on a new corpus that presents a different fact $(s_2,r,o_2)$. If $s_2$ is similar to $s_1$, their key vectors ${𝒌}_1\approx {𝒌}_2$ will be close in the MLP's input space. The gradient update for the new fact will shift ${𝐖}_{\mathrm{out}}$ in the direction of ${𝒗}_2$, but since the update direction is ${𝒌}_2\approx {𝒌}_1$, it simultaneously shifts the response to ${𝒌}_1$. Concretely, if the rank-one update for fact 2 is $\mathrm{\Delta }𝐖=({𝒗}_2-𝐖{𝒌}_2){𝒌}_2^{\top }/{\Vert {𝒌}_2\Vert }^2$, then the model's output for fact 1 becomes: (Forgetting)$$\begin{array}{ccc}& (𝐖+\mathrm{\Delta }𝐖){𝒌}_1=𝐖{𝒌}_1+({𝒗}_2-𝐖{𝒌}_2)\frac{{𝒌}_2^{\top }{𝒌}_1}{{\Vert {𝒌}_2\Vert }^2}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The second term is nonzero whenever ${𝒌}_1$ and ${𝒌}_2$ are non-orthogonal, i.e., whenever the two subjects share semantic features (which is the norm rather than the exception in natural language). The original value $𝐖{𝒌}_1\approx {𝒗}_1$ is partially overwritten by ${𝒗}_2$, degrading recall of fact 1.

This perspective immediately suggests remedies:

• Regularisation. Penalise changes to weight dimensions that are active for previously learned facts (analogous to Elastic Weight Consolidation, Kirkpatrick et al., 2017).

• Orthogonal updates. Restrict weight updates to the null space of old key vectors (as in PackNet, Mallya and Lazebnik, 2018, or the OWM method of Zeng et al., 2019).

• Rehearsal. Replay a subset of old facts during fine-tuning to maintain their key-value associations.

Interventions for Reducing Hallucination

The mechanistic account of hallucination points to several targeted interventions:

Activation Steering

Given a direction $𝒅$ in residual-stream space associated with honest/factual behaviour (identified via contrastive activation analysis), one can steer the model toward that direction at inference time by adding $\alpha 𝒅$ to the residual stream at a designated layer. This technique, called representation engineering (Zou et al., 2023), has been shown to reduce hallucination rates on factual QA benchmarks by up to 40% without any weight update.

Formally, if ${𝒉}^{(l)}$ is the residual stream at layer $l$, we apply: (ACT Steering)$$\begin{array}{ccc}& {𝒉}^{(l)}\leftarrow {𝒉}^{(l)}+\alpha \widehat{𝒅},\widehat{𝒅}=\frac{𝒅}{\Vert 𝒅\Vert },& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒅$ is estimated as the mean difference between hidden states on honest vs. hallucinating completions in a calibration set.

Contrastive Decoding

A model that has been lightly fine-tuned on synthetic hallucinations acts as an “amateur” model $f_{{\theta }_a}$ that over-expresses hallucination tendencies. Contrastive decoding generates tokens that maximise: (Contrastive Decoding)$$\begin{array}{ccc}& \log p_{{\theta }_{\mathrm{expert}}}(t|𝚙)-\beta \log p_{{\theta }_{\mathrm{amateur}}}(t|𝚙),& \text{<span class="tml-eqn"></span>}\end{array}$$ with hyperparameter $\beta >0$. The subtraction effectively suppresses tokens that the amateur (hallucinating) model prefers, directing the expert model away from its own hallucination-prone regions.

Uncertainty-Based Abstention

If the model's predictive entropy for a given factual query exceeds a threshold, the model abstains rather than generating a potentially false answer. The predictive entropy is: (PRED Entropy)$$\begin{array}{ccc}& 𝖧[t|𝚙]=-\sum _t{p}_{\theta }(t|𝚙)\log p_{\theta }(t|𝚙).& \text{<span class="tml-eqn"></span>}\end{array}$$ High entropy signals that no single circuit is dominant, which corresponds mechanistically to the case where both $\Vert {𝒗}_{\mathrm{param}}\Vert$ and $\Vert {𝒗}_{\mathrm{ctx}}\Vert$ are small or they roughly cancel. Abstention in this regime avoids presenting a low-confidence guess as a confident fact.

The Forgetting Problem at Scale

For large language models trained on internet-scale data, catastrophic forgetting takes on a new character. During pre-training on a single large corpus, facts encountered early in the training schedule can be partially overwritten by later training steps. In the regime of very large models (100B+ parameters), the MLP layers have sufficient capacity that individual facts are each encoded by a small subset of neurons; the key vectors for unrelated facts tend to be nearly orthogonal simply due to the high dimensionality (Johnson-Lindenstrauss phenomenon). This greatly reduces interference.

However, during fine-tuning (e.g. supervised instruction tuning or RLHF), the model is trained on a much narrower distribution. The fine-tuning gradient steps are concentrated on a small set of task distributions, causing large updates to the MLP layers that are relevant for those tasks. Facts not covered in the fine-tuning distribution can be overwritten. This is sometimes called alignment tax: improving instruction-following or safety behaviour degrades raw factual accuracy.

Concretely, the LIMA paper (Zhou et al., 2023) observed that even very short fine-tuning (1000 examples) can substantially shift the model's factual output distribution. The mechanistic account predicts exactly this: 1000 gradient steps on a narrow distribution produce 1000 rank-one-like updates, each of which potentially interferes with the broader set of encoded facts.

Exercises

VLM Interpretability and Cross-Modal Alignment

The interpretability toolkit developed for pure language models rests on a foundational assumption: that the objects under study are sequences of discrete tokens, each occupying a well-defined position in a shared vocabulary whose statistics can be probed, steered, and ablated. Logit-lens projections presuppose that intermediate representations are already trying to predict a vocabulary item. Activation patching presupposes that a causal model of token-to-token influence captures the dominant information flow. Even probing classifiers, the most assumption-light tool in the box, typically operate on representations that live in a single homogeneous residual stream.

Vision-language models (VLMs) shatter these assumptions. A modern VLM such as LLaVA, InternVL, or GPT-4V processes two fundamentally different information streams simultaneously: a sequence of visual patch embeddings extracted by a vision encoder and a sequence of linguistic tokens processed by a language backbone. The two streams are fused through a projection layer (often a simple linear map or a small MLP) that maps visual features into the token embedding space of the language model. The resulting architecture is not a language model with images bolted on; it is a hybrid system in which the inductive biases of convolutional or attention-based vision encoders interact with the statistical geometry of a pre-trained language model.

The language-centric bias problem is the observation that almost every existing interpretability tool is calibrated for the linguistic modality and fails gracefully-or fails silently-when applied to the visual modality. Logit-lens projections of visual patch embeddings produce vocabulary distributions that are typically incoherent, because visual patch tokens are not expected to predict the next word; they are expected to encode spatial features. Attention head analyses developed for text transformers frequently mis-attribute visual attention patterns to roles that have no linguistic analogue. Even the vocabulary of mechanistic interpretability-circuits, features, subspaces-was developed with text in mind and requires careful reinterpretation when the model processes images.

This section develops the emerging theory of VLM interpretability, focusing on three themes that have produced the most conceptual traction: the cross-modal structure of task representations, the dynamics of modality fusion in the residual stream, and the three-stage evolution of visual representations as they traverse a VLM's layers.

The Language-Centric Bias Problem

To make the language-centric bias precise, consider a VLM $f_{\theta }:{𝒳}_{\text{vis}}\times {𝒳}_{\text{txt}}\to {𝒳}_{\text{txt}}$ that takes a visual input $𝒗\in {𝒳}_{\text{vis}}$ and a textual prompt $𝒕\in {𝒳}_{\text{txt}}$ and produces a textual response. Internally, the model processes a concatenated token sequence (Tokens)$$\begin{array}{ccc}& 𝒔=[\underset{\text{visual patches}}{\underbrace{{𝒗}_1,\dots ,{𝒗}_P}},\underset{\text{text tokens}}{\underbrace{{𝒕}_1,\dots ,{𝒕}_L}}],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $P$ is the number of visual patch tokens and $L$ is the length of the text prompt. Each token $s_i$ is associated with a residual stream ${𝒉}_i^{(\ell )}\in {\mathrm{ℝ}}^d$ at layer $\ell$.

The logit-lens projects ${𝒉}_i^{(\ell )}$ into vocabulary space via the unembedding matrix ${𝐖}_U\in {\mathrm{ℝ}}^{|𝒱|\times d}$: (Logitlens)$$\begin{array}{ccc}& p_i^{(\ell )}=\mathrm{softmax}\left({𝐖}_U{𝒉}_i^{(\ell )}\right).& \text{<span class="tml-eqn"></span>}\end{array}$$ For text tokens at late layers, $p_i^{(\ell )}$ typically concentrates on semantically plausible vocabulary items-the prediction sharpens as the layer index increases. For visual patch tokens, $p_i^{(\ell )}$ remains diffuse and lacks semantic coherence even at late layers, because the visual patches are never the subject of a next-token prediction objective: the model is not trying to predict what word follows a patch.

This distributional mismatch motivates the search for interpretability tools that are modality-agnostic: they should probe the residual stream in a way that does not presuppose token-predictive semantics. The discovery of task vectors is one such tool.

Cross-Modal Task Representations: Task Vectors

The concept of a task vector originated in the context of weight-space arithmetic for language model fine-tuning: if ${𝜽}_0$ are pre-trained weights and ${𝜽}_{\text{FT}}$ are fine-tuned weights, the task vector ${𝜽}_{\text{FT}}-{𝜽}_0$ captures the “direction” in weight space associated with the fine-tuning task. Adding this vector to a different model's weights transfers some of the task's properties.

The work of Luo et al. (2025) introduces a different, activation-space notion of task vectors that is directly relevant to VLM interpretability. They observe that when a language model processes a few-shot prompt encoding a task (e.g., “translate English to French”), the residual stream at a specific layer develops a low-dimensional representation that encodes the task itself, largely independently of the specific examples used to illustrate it. This task representation, once extracted, can be transplanted into a different forward pass to steer the model's behaviour without any explicit prompt.

Three properties of task vectors, empirically verified by Luo et al., make them particularly useful for VLM interpretability:

1. Task specificity. Task vectors for different tasks (${\tau }_1\ne {\tau }_2$) are approximately orthogonal: ${𝝉}_1^{(\ell )}\cdot {𝝉}_2^{(\ell )}\approx 0$.

2. Example invariance. The task vector is stable across different choices of few-shot examples: replacing $(x_1,y_1,\dots )$ with a different set of examples for the same task changes ${𝝉}^{(\ell )}$ by a small perturbation.

3. Cross-modal transferability. A task vector derived from a textual few-shot prompt can be patched into the forward pass of a visual task to guide the model's behaviour.

Property 3 is the most surprising and the most important for VLM interpretability. It implies that the abstract representation of a task-what it means to perform colour recognition, object counting, or spatial reasoning-is encoded in a modality-agnostic subspace of the residual stream.

Task Vector Invariance Across Modalities

The formal statement of cross-modal invariance requires a careful definition of what it means for a textual and a visual instantiation of a task to share the same task vector.

Let $\tau$ be a task (e.g., “identify the dominant colour of the object”). Consider two few-shot prompts:

• A textual prompt ${𝒕}^{\text{txt}}$ that encodes the task through text descriptions: “The ball is red. Dominant colour: red. The sky is blue. Dominant colour: blue. The grass is ”.

• A visual prompt ${𝒕}^{\text{vis}}$ that encodes the same task through image-caption pairs: images of coloured objects captioned with their dominant colour.

Define the textual task vector ${𝝉}_{\text{txt}}^{(\ell )}$ and the visual task vector ${𝝉}_{\text{vis}}^{(\ell )}$ by applying Definition 31 to the two prompts respectively. The cross-modal alignment score at layer $\ell$ is (Alignment)$$\begin{array}{ccc}& A^{(\ell )}(\tau )=\frac{{𝝉}_{\text{txt}}^{(\ell )}\cdot {𝝉}_{\text{vis}}^{(\ell )}}{\Vert {𝝉}_{\text{txt}}^{(\ell )}\Vert \Vert {𝝉}_{\text{vis}}^{(\ell )}\Vert },& \text{<span class="tml-eqn"></span>}\end{array}$$ measuring the cosine similarity between the textual and visual task vectors at layer $\ell$.

This proposition, while stated here as an empirical observation, has a structural explanation: the projection layer that maps visual patch embeddings into the language backbone's token space initialises the visual residual stream in the language model's representation space. From the perspective of the language backbone's upper layers, visual tokens are simply unusual text tokens. The task vector, which is defined over the language backbone layers, therefore reflects a linguistic encoding of the visual task.

Patching Task Vectors Across Modalities

The utility of cross-modal task vectors is demonstrated through a patching experiment: a textual task vector is derived from a textual few-shot prompt and then injected into a visual forward pass. Formally:

1. Run the model on a textual few-shot prompt ${𝒕}^{\text{txt}}$ for task $\tau$, extracting ${𝝉}_{\text{txt}}^{({\ell }^{\ast })}$ at a chosen layer ${\ell }^{\ast }$.

2. Run the model on a zero-shot visual query ${𝒕}_q^{\text{vis}}$ (an image with a short question, no examples). At layer ${\ell }^{\ast }$, add the task vector to the query token's residual stream: (Patch)$$\begin{array}{ccc}& {\widetilde{𝒉}}_q^{({\ell }^{\ast })}={𝒉}_q^{({\ell }^{\ast })}+\lambda {𝝉}_{\text{txt}}^{({\ell }^{\ast })},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\lambda >0$ is a scaling hyperparameter.

3. Continue the forward pass from layer ${\ell }^{\ast }+1$ with the modified residual stream and observe the model's output.

The effect of this patching operation is dramatic: a model that produces a generic or incorrect answer to the visual query in the unpatched setting often produces a task-appropriate answer after patching, even though no visual examples of the task were provided. The textual task vector has “reminded” the model what task it is performing.

The success of cross-modal patching has a mechanistic interpretation. The patch operation in (Patch) shifts the query token's residual stream in the direction of the task-encoding subspace. Because the language backbone was pre-trained to use this subspace for task representation, the shift induces the same task-appropriate processing regardless of whether the input is textual or visual.

Transfer from Base LLM to Fine-Tuned VLM

A further surprising finding is that task vectors computed from a base language model (before any visual fine-tuning) can be transferred to the VLM derived from it. If $f_{\theta }^{\text{LLM}}$ is the base language model and $f_{\theta }^{\text{VLM}}$ is the VLM obtained by adding a vision encoder and projection layer and fine-tuning, the task vectors derived from $f_{\theta }^{\text{LLM}}$ are often effective when patched into $f_{\theta }^{\text{VLM}}$.

This transferability arises because visual fine-tuning preserves the language backbone's representation geometry: the weights of the language backbone change relatively little during VLM fine-tuning (typical norms of weight updates are one to two orders of magnitude smaller than the original weight norms), and the task-encoding subspace is consequently preserved.

TikZ Figure: Task Vector Extraction and Cross-Modal Transfer

Probing VLM Layers: Modality Prioritisation and Asymmetric Fusion

Beyond task vectors, linear probes offer a complementary window into VLM internals. Given a set of labelled inputs $(s_i,y_i)$ where $s_i$ is a VLM input (image-text pair) and $y_i$ is a ground-truth label for some property of interest (e.g., object category, spatial relation, sentiment), a linear probe at layer $\ell$ is a logistic regressor: (Probe)$$\begin{array}{ccc}& {\hat{y}}_i^{(\ell )}=\sigma ({𝐖}_{\text{probe}}^{(\ell )}{𝒉}_i^{(\ell )}+b),& \text{<span class="tml-eqn"></span>}\end{array}$$ trained to minimise cross-entropy loss. The probe accuracy as a function of layer $\ell$ and token type (visual vs. text) reveals the information routing strategy of the VLM.

Systematic probing across multiple VLM architectures reveals two consistent patterns:

Modality prioritisation.

For tasks that are primarily visual (object classification, colour identification, spatial relation detection), the task-relevant information is concentrated in the visual patch tokens in early layers. By mid-depth, however, the task information has migrated to the query text tokens. This migration-from patches to text-reflects the VLM's strategy of translating visual content into linguistic representations before generating a response.

Asymmetric fusion.

The flow of information between visual and text tokens is asymmetric. Visual patch tokens attend heavily to each other in early layers (local spatial processing) and then attend to text tokens in middle layers (semantic grounding). Text query tokens attend to visual patch tokens throughout but most heavily in middle layers. The resulting attention pattern resembles a funnel: spatial detail flows up from patches into global semantic representations, which are then queried by the text tokens to produce the final answer.

The Logit Lens for VLMs: Three-Stage Evolution

While the standard logit lens fails on visual patch tokens, a modified logit lens for VLMs can be defined that operates on the text query tokens rather than the patch tokens. Tracking the logit-lens distribution of the query token ${𝒉}_q^{(\ell )}$ as $\ell$ increases from $0$ to $L$ reveals a characteristic three-stage evolution.

Stage 1: Alignment ($\ell \in [0,L/3]$).

In the early layers, the query token's logit-lens distribution is broad and unsettled, reflecting the model's uncertainty about the task. However, the distribution is not random: it concentrates on tokens semantically related to the visual content. If the image contains a cat, the distribution assigns elevated probability to “cat”, “feline”, “kitten”, “animal”, and related tokens. This stage represents the alignment of linguistic representations with visual content: the VLM is mapping what it “sees” into the vocabulary of what it “knows”.

Stage 2: Processing ($\ell \in [L/3,2L/3]$).

In the middle layers, the distribution narrows around task-relevant tokens. If the task is colour identification, the early-layer semantic alignment (“cat”) is replaced by colour-relevant tokens (“orange”, “brown”, “striped”). This stage reflects the activation and application of the task representation: the model has identified the relevant task and is extracting the appropriate features.

Stage 3: Selection ($\ell \in [2L/3,L]$).

In the late layers, the distribution collapses onto a small set of high-probability tokens that will constitute the answer. The uncertainty of the middle layers-oscillating among several plausible answers-resolves into a confident prediction. This stage mirrors the token-selection dynamics observed in pure language models and reflects the VLM's finalisation of its output.

VLMs Anchor Visual Representations in Linguistic Geometry

The three-stage logit-lens evolution, the asymmetric fusion pattern, and the cross-modal transferability of task vectors all converge on a single organising principle:

This anchoring hypothesis has several testable predictions. First, the cosine similarity between the residual stream representations of a concept presented visually (an image of a cat) and the same concept presented linguistically (the token “cat”) should be high in mid-to-late layers. Second, the language model's geometry-its organisation of concepts into directions, clusters, and subspaces- should be preserved in the VLM and accessible through standard linguistic probes. Third, the VLM's visual understanding should be bounded by its linguistic knowledge: visual concepts for which the language model has no linguistic representation should be difficult to process correctly.

All three predictions have been confirmed empirically in recent probing studies, providing strong evidence for the linguistic anchoring view.

Exercises

Vision SAEs and the Prisma Framework

Sparse autoencoders (SAEs) have emerged as one of the most productive tools in language model interpretability. Trained on the residual stream activations of a frozen language model, an SAE learns a dictionary of features-directions in activation space that activate for semantically coherent, human-interpretable concepts. A residual stream vector $𝒉\in {\mathrm{ℝ}}^d$ is decomposed as: (SAE)$$\begin{array}{ccc}& 𝒉\approx \sum _{k=1}^K{a}_k{𝒇}_k=𝐅𝒂,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝐅=[{𝒇}_1,\dots ,{𝒇}_K]\in {\mathrm{ℝ}}^{d\times K}$ is the feature dictionary, $𝒂={[a_1,\dots ,a_K]}^{\top }\in {\mathrm{ℝ}}_{\ge 0}^K$ is the sparse activation vector (most entries zero), and $K\gg d$ is the overcomplete dictionary size. Training minimises the sparse reconstruction loss: (LOSS)$$\begin{array}{ccc}& {ℒ}_{\text{SAE}}={𝔼}_{𝒉}[{\Vert 𝒉-𝐅𝒂\Vert }_2^2+\alpha {\Vert 𝒂\Vert }_1],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\alpha >0$ controls sparsity.

The transition to vision models introduces qualitative differences that demand careful redesign. This section develops the theory of vision SAEs, introduces the Prisma framework that makes vision SAE research practical at scale, and examines two phenomena unique to visual feature representations: CLS token dynamics and the surprising regularisation effect of SAE reconstruction.

Why Vision Needs Denser Features Than Language

The success of language SAEs rests on a property of linguistic representation that does not hold for vision: token discreteness. A language model processes a sequence of discrete tokens drawn from a finite vocabulary. Each token has a relatively well-defined semantic interpretation, and the features learned by a language SAE tend to activate for specific tokens or token categories. Sparsity in language SAEs reflects the fact that, at any given position, only a small number of semantically relevant features are active.

Visual inputs are fundamentally different. A pixel is not a discrete symbol; it is a sample from a continuous signal. A patch of pixels encodes not a single semantic concept but a superposition of many partially overlapping concepts: texture, colour, edge orientation, lighting condition, object boundary, semantic category, depth cue, and more. These concepts do not occur in isolation; they co-occur continuously and with graded intensity.

This density requirement has a corollary: vision SAE features are less monosemantic than language SAE features. A language SAE feature might activate only for the token “president” in political contexts. A vision SAE feature might activate for orange-coloured regions, which could be sunsets, traffic cones, autumn leaves, or orange fruits. The partial overlap of visual concepts means that perfect monosemanticity is unattainable; the best vision SAEs can achieve is relative monosemanticity: features that are more semantically coherent than the original polysemantic neurons.

To quantify the required density, define the average feature activation count $\bar{k}$ as the expected number of active features per input: (Density)$$\begin{array}{ccc}& \bar{k}={𝔼}_{𝒉}\left[{\Vert 𝒂\Vert }_0\right]={𝔼}_{𝒉}\left[\sum _{j=1}^K\mathrm{𝟏}[a_j>0]\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ Empirically, language SAEs achieve good reconstruction with $\bar{k}\approx 20$–$50$ active features per token (out of $K\approx 50000$–$100000$ total features). Vision SAEs require $\bar{k}\approx 100$–$300$ active features per patch token to achieve comparable reconstruction quality, a five-to-ten-fold increase in density.

The Prisma Toolkit

Training SAEs across many architectures, layers, and token types is a substantial engineering undertaking. The Prisma toolkit is an open-source framework designed to make vision SAE research practical at scale. It provides three core capabilities: activation caching, SAE training across diverse architectures, and a unified evaluation suite.

Activation caching.

Prisma caches the activations of any vision transformer layer for an entire dataset in a compressed format. Given a vision transformer $f_{\theta }$ and a dataset $𝒟=\{x_1,\dots ,x_N\}$, Prisma computes $${\mathscr{H}}^{(\ell )}={\{{𝒉}_i^{(\ell )}\}}_{i=1}^N$$ for each layer $\ell$ of interest and stores the result on disk. The caching step is run once per model-dataset pair, and subsequent SAE training reads from the cache rather than running the vision transformer, reducing compute costs by an order of magnitude.

Architectural coverage.

Prisma supports over 75 vision transformer architectures, including ViT, DeiT, DINO, DINOv2, CLIP-ViT, MAE, and their variants across scales from ViT-Small to ViT-Giant. The architectural diversity is important because vision SAE features are not universal: a feature that is monosemantic in a DINO-trained ViT may be polysemantic in a CLIP-trained ViT, reflecting the different pre-training objectives. Comparing SAEs across architectures makes it possible to identify which features are stable across training regimes (suggesting they reflect intrinsic properties of the visual world) and which are architecture- or training-specific.

Evaluation suite.

Prisma provides a unified evaluation suite that measures SAE quality along four dimensions:

• Reconstruction fidelity: the variance explained $R^2=1-{\Vert 𝒉-𝐅𝒂\Vert }_2^2/{\Vert 𝒉\Vert }_2^2$.

• Downstream task performance: accuracy of linear probes trained on SAE features rather than raw activations.

• Feature interpretability: the fraction of features that receive consistent natural-language descriptions from human annotators.

• Feature utilisation: the fraction of features that activate on at least one example in the evaluation set (the complement of the dead feature fraction).

CLS Token Dynamics: Spatial to Global

Vision transformers (ViTs) append a special [CLS] token to the sequence of patch tokens. The CLS token's final representation is used as the global image representation for downstream tasks (e.g., classification). Understanding how the CLS token develops its global representation across layers is a central question in ViT interpretability.

Prisma SAEs trained at each layer of a ViT reveal a striking evolution. In early layers ($\ell <L/4$), the CLS token's SAE decomposition activates features associated with specific spatial patches: the CLS token is not yet global; it attends to and summarises a small neighbourhood of patches. In middle layers ($\ell \in [L/4,3L/4]$), the CLS token's activated features broaden to cover larger spatial regions, and abstract semantic features (object category, scene type) begin to appear. In late layers ($\ell >3L/4$), the CLS token's representation is dominated by global, task-relevant features with little remaining spatial specificity.

Empirically, $S^{(\ell )}$ starts high (around $10$–$20$ for a ViT-Base with 196 patches), decreases through the middle layers, and approaches $1$ (uniform attention) at the final layers. The trajectory of $S^{(\ell )}$ is the attention signature of the global-to-local transition: the CLS token begins by attending locally and progressively integrates information from all patches.

Dead Feature Analysis and Information Flow

A dead feature in an SAE is a dictionary vector ${𝒇}_k$ that activates on zero (or near-zero) fraction of the evaluation set. Dead features are a training pathology: they represent wasted capacity, directions in the overcomplete dictionary that the SAE has learned but that the model's activations never use.

The fraction of dead features $f_{\text{dead}}^{(\ell )}$ varies significantly across layers. For vision SAEs trained on patch tokens, $f_{\text{dead}}^{(\ell )}$ is low in early layers (nearly all features are utilised because early-layer patch activations are diverse and high-dimensional) and rises in later layers (as patch activations become more stereotyped and globalised). For the CLS token, the pattern is reversed: $f_{\text{dead}}^{(\ell )}$ is high in early layers (the CLS token has not yet accumulated much information) and falls in later layers (as the CLS token absorbs information from all patches).

This complementarity is not a precise mathematical identity but a robust empirical regularity that holds across DINO, CLIP, and MAE pre-trained ViTs. It provides a non-invasive diagnostic for the layer at which information transfer from patches to the CLS token is most active: the layer at which the slopes of $f_{\text{dead,patch}}^{(\ell )}$ and $f_{\text{dead,CLS}}^{(\ell )}$ are most steeply negative and positive, respectively.

TikZ Figure: Information Flow from Patches to CLS Token

The Surprising Regularisation Effect

A counterintuitive finding from Prisma experiments is that replacing a vision transformer's internal activations with their SAE reconstructions can sometimes improve the model's downstream task performance rather than degrading it. This is surprising because SAE reconstruction introduces an error: $\widetilde{𝒉}=𝐅𝒂$ is not equal to the original $𝒉$ (SAE reconstruction is lossy), and one would expect any perturbation of the activations to harm performance.

To understand this effect, decompose the reconstruction error: (Error Decomp)$$\begin{array}{ccc}& 𝒉-\widetilde{𝒉}=\underset{\text{reconstruction error}}{\underbrace{𝒉-𝐅𝒂}}=\underset{\text{task-relevant component}}{\underbrace{{𝒉}_{\text{task}}}}+\underset{\text{task-irrelevant noise}}{\underbrace{{𝒉}_{\text{noise}}}}-\widetilde{𝒉}.& \text{<span class="tml-eqn"></span>}\end{array}$$

The SAE is trained to reconstruct the dominant directions of variation in $𝒉$ (those with high variance across the dataset). Task-relevant features, which are activated for specific, semantically coherent inputs, contribute high variance and are therefore well-reconstructed. Task-irrelevant noise-high-frequency variation, adversarial perturbations, training artefacts-contributes low variance and is poorly reconstructed. The SAE reconstruction $\widetilde{𝒉}$ is therefore a denoised version of $𝒉$: it retains the task-relevant component and discards the task-irrelevant noise.

When the noise component is large relative to the task-relevant component (as can occur in models with limited capacity or noisy pre-training), replacing $𝒉$ with $\widetilde{𝒉}$ reduces noise and improves downstream performance. The SAE acts as a learned regulariser that projects activations onto the manifold of task-relevant features.

Crosscoders: Bridging Representations Across Model Pairs

Crosscoders extend the SAE framework to the setting where two related models-a base model and a fine-tuned model, or two models from the same family trained on different data-have similar but not identical representation spaces. The goal is to find a shared feature dictionary that explains both models' activations simultaneously, making it possible to identify features that are shared (present in both models), added (present only in the fine-tuned model), and suppressed (present only in the base model).

By analysing the activation patterns ${𝒂}_1$ and ${𝒂}_2$ over a large dataset, one can categorise each feature ${𝒇}_k$ as:

• Shared: feature $k$ activates frequently in both ${𝒂}_1$ and ${𝒂}_2$ (high mean activation in both).

• Added: feature $k$ activates frequently in ${𝒂}_2$ but rarely in ${𝒂}_1$ (characteristic of capabilities introduced by fine-tuning).

• Suppressed: feature $k$ activates frequently in ${𝒂}_1$ but rarely in ${𝒂}_2$ (characteristic of behaviours removed by fine-tuning).

• Absent: feature $k$ activates rarely in both.

This taxonomy provides a mechanistic account of what fine-tuning does to a model's representation: it adds new features, suppresses undesired features, and leaves most shared features intact.

Finding Bias Features in CLIP Models

CLIP (Contrastive Language-Image Pre-training) trains a vision encoder and a text encoder jointly so that paired image-text embeddings are similar while unpaired embeddings are dissimilar. CLIP models are widely used as zero-shot classifiers and visual feature extractors, but they are known to encode societal biases from their training data (the LAION-5B dataset, scraped from the internet, reflects the biases of internet-hosted images and their captions).

Vision SAEs and crosscoders provide a principled method for identifying bias features in CLIP models. The procedure:

1. Train a vision SAE on the CLIP ViT's CLS token activations.

2. Identify features whose activation correlates with protected attributes (gender, race, age) by computing, for each feature $k$, the correlation between $a_k(x)$ and a ground-truth annotation of $x$'s protected attribute.

3. Interpret high-correlation features using automated interpretability: generate a natural-language description of the concept encoded by feature $k$ by finding the images that maximally activate it.

A crosscoder between a base CLIP model and a debiased CLIP model (trained with additional constraints to reduce attribute correlations) identifies features that are suppressed by debiasing. These suppressed features are precisely the bias features: directions in CLIP's representation space that encode protected attributes. By inspecting the feature dictionary entries corresponding to suppressed features, researchers can identify the visual concepts that the debiasing process has disrupted and assess whether the disruption is appropriate (removing stereotypical associations) or harmful (degrading recognition of legitimate visual variation).

Exercises

Interpreting Diffusion Models

Diffusion models have emerged as the dominant paradigm for high-fidelity image synthesis, text-to-image generation, and an expanding array of scientific applications. Unlike autoregressive language models, whose token-by-token generation provides a natural sequential structure amenable to causal analysis, diffusion models operate through a continuous denoising trajectory: a stochastic process that begins with pure Gaussian noise and progressively refines it into a structured sample. This trajectory is parameterised by a reverse-time score function, typically implemented as a U-Net or Diffusion Transformer (DiT), which receives both a noisy latent vector and a timestep embedding as input.

The interpretability challenge in diffusion models is therefore fundamentally temporal. A single forward pass of a denoising model does not produce an interpretable output; the meaningful signal emerges only after integrating over the entire denoising trajectory. Moreover, different stages of the trajectory perform qualitatively different computations. Early denoising steps, operating in the high-noise regime, must decide the global spatial arrangement of the image: where objects are placed, how the scene is composed, and what the broad lighting conditions will be. Intermediate steps progressively resolve questions of style, object identity, and colour palette. Late steps, operating close to the data manifold, refine texture, sharpness, and fine-grained details. Interpretability tools that ignore this temporal structure will conflate fundamentally different computations into a single undifferentiated analysis.

This section develops a temporally-aware framework for interpreting diffusion models. We begin by characterising the concept schedule that governs what information is encoded at each timestep, survey mechanistic tools for probing this schedule, and describe how sparse autoencoders trained on diffusion latent spaces enable principled causal interventions. Throughout, we emphasise the guiding insight that shapes every technique in this section.

The Denoising Trajectory and Its Computational Stages

Recall the forward diffusion process. Given a data point ${𝒙}_0\sim p_{\mathrm{data}}$, the forward process constructs a sequence of increasingly noisy latents (Forward)$$\begin{array}{ccc}& q({𝒙}_t|{𝒙}_0)=𝒩({𝒙}_t;\sqrt{{\bar{\alpha }}_t}{𝒙}_0,(1-{\bar{\alpha }}_t)𝐈),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$ decreases monotonically from ${\bar{\alpha }}_0\approx 1$ (clean data) to ${\bar{\alpha }}_T\approx 0$ (pure noise). The signal-to-noise ratio (SNR) at timestep $t$ is (SNR)$$\begin{array}{ccc}& \mathrm{SNR}(t)=\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t},& \text{<span class="tml-eqn"></span>}\end{array}$$ which decays from a very large value near $t=0$ to approximately zero near $t=T$.

The reverse process learns a neural network ${𝝐}_{𝜽}({𝒙}_t,t)$ that approximates the score function $-{\nabla }_{{𝒙}_t}\log q({𝒙}_t)$. The key observation for interpretability is that the network receives inputs with very different statistical properties depending on $t$. Near $t=T$, the input ${𝒙}_t$ is nearly indistinguishable from white noise; the network must extract global compositional information from a tiny signal buried in enormous variance. Near $t=0$, the input ${𝒙}_t$ is already close to the final image, and the network need only predict fine-grained residual details.

Stage 1 - Semantic Layout ($t\in [0.8T,T]$).

The high-noise regime is where the model determines the global semantic structure of the image. Empirical studies of cross-attention maps in text-conditioned diffusion models reveal that at early denoising steps, attention is broadly distributed across the spatial extent of objects, with distinct object regions already delineating themselves in the attention maps. Interventions at this stage can reposition objects, add or remove scene elements, and alter the overall compositional arrangement without affecting the detailed appearance of individual objects.

Stage 2 - Style and Identity ($t\in [0.3T,0.8T]$).

The intermediate regime is where the model resolves questions of style, colour palette, and object identity. Cross-attention at this stage becomes more localised, with attention heads specialising in specific visual attributes. Prompt-to-image correspondence is established: the model binds noun phrases in the text prompt to specific spatial regions. Style transfer interventions applied in this regime can alter the artistic style or object attributes without disrupting the global layout established in Stage 1.

Stage 3 - Texture and Detail ($t\in [0,0.3T]$).

The low-noise regime resolves fine-grained texture, edge sharpness, and perceptual quality metrics (CLIP score, FID) that are sensitive to high-frequency content. Interventions at this stage affect the crispness and material appearance of surfaces but cannot alter the global semantic content, which has already been committed by earlier stages.

Definition: Temporal Concept Schedule

We formalise the intuition above with a definition that will underpin the rest of this section.

TikZ Figure: Concept Emergence Timeline

Training Sparse Autoencoders on Diffusion Latent Spaces

The success of sparse autoencoders (SAEs) for interpreting language model residual streams (see earlier sections of this chapter) raises a natural question: can the same technique be applied to the internal representations of diffusion models? The challenge is that diffusion models differ from LLMs in a crucial structural way. In an LLM, the residual stream at a given layer is a timestep-independent representation of the current input; every token produces a vector in the same representational space. In a diffusion model, the residual stream at a given layer is a function of both the noisy latent ${𝒙}_t$ and the timestep $t$. The representation geometry therefore varies continuously with $t$, rendering a single SAE trained across all timesteps potentially useless: the dictionary atoms learned for $t\approx T$ will not correspond to interpretable features at $t\approx 0$, and vice versa.

[2] address this challenge by training timestep-stratified SAEs: a collection of SAEs ${\{{\mathrm{\Phi }}_w\}}_{w=1}^W$, one per timestep window ${𝒯}_w=[t_w^{-},t_w^{+}]$. Within each window, the representation geometry is approximately stationary, allowing the SAE to learn a coherent dictionary of semantic features.

The timestep-stratified design introduces an important practical consideration: the concept vocabulary learned by ${\mathrm{\Phi }}_w$ will differ across windows. In the high-noise window ${𝒯}_1$, dictionary atoms will correspond to global compositional features (e.g., horizontal extent of a region, left-right symmetry). In the low-noise window ${𝒯}_W$, atoms will correspond to texture primitives (e.g., fur directionality, specular highlights). This is not a deficiency of the method but rather a direct reflection of the concept schedule established in Definition 34.

Causal Interventions via Feature Injection and Ablation

Probing and SAE training establish a correlational understanding of which features are active at which timesteps. The gold standard for mechanistic interpretability, however, is causal: we should be able to manipulate a feature and observe a corresponding change in the output, with the change localised to the concept the feature is hypothesised to represent.

For diffusion models, causal interventions take the form of activation patching applied to specific timestep windows. Let ${𝒂}_w^{(\text{ref})}$ be the SAE feature activations at window ${𝒯}_w$ when generating a reference image ${𝒙}^{\text{ref}}$, and let ${𝒂}_w^{(\text{tgt})}$ be the activations when generating a target image ${𝒙}^{\text{tgt}}$. The intervention replaces the activations in window ${𝒯}_w$ during a new generation run with a weighted mixture: (Activation Patch)$$\begin{array}{ccc}& {\widetilde{𝒂}}_w=(1-\gamma ){𝒂}_w^{(\text{tgt})}+\gamma {𝒂}_w^{(\text{ref})},\gamma \in [0,1],& \text{<span class="tml-eqn"></span>}\end{array}$$ while leaving the activations in all other windows unchanged.

Time-Aware Manipulation: Style Without Layout Disruption

One of the most practically valuable applications of temporal interpretability is time-aware manipulation: the ability to change a specific visual attribute (style, object type, colour palette) while preserving other attributes that were fixed by earlier stages of the trajectory.

The key insight is that layout is committed in window ${𝒯}_1$, so any intervention restricted to windows ${𝒯}_2$ and ${𝒯}_3$ cannot alter the global spatial arrangement without rerunning the high-noise stage. This gives us a clean separation: to change the style of an image without changing its composition, we perform the following procedure.

1. Generate a layout image ${𝒙}^{\text{layout}}$ using the original prompt, recording the intermediate activations ${\{{𝒉}_{\ell }({𝒙}_t,t)\}}_{t=1}^T$ throughout the trajectory.

2. Modify the text prompt or inject a style SAE feature at window ${𝒯}_2$.

3. Rerun the denoising process from the beginning, but in window ${𝒯}_1$, replace the activations with those recorded from the layout image run.

4. Allow windows ${𝒯}_2$ and ${𝒯}_3$ to run freely with the modified style conditioning.

The resulting image inherits the composition of ${𝒙}^{\text{layout}}$ but the style of the modified conditioning. This is the diffusion analogue of style transfer at a mechanistic level: rather than operating on the final pixel image, we intervene directly in the internal representations of the denoising network at the stage where style is being determined.

Steering Image Generation: The Feature-Timestep Interface

Beyond the style-transfer use case, the timestep-stratified SAE framework enables a rich vocabulary of generation steering operations. We categorise these along two axes: (1) the target window (which stage of the trajectory is modified), and (2) the intervention type (feature injection, feature ablation, or feature clamping).

Layout steering (${𝒯}_1$ intervention).

To alter the global spatial composition without changing style or texture, we inject or ablate layout SAE features during the high-noise stage. For instance, ablating the “left-right object separation” feature causes the model to generate images where two objects that would normally appear side by side instead overlap or merge. Injecting the “vertical stacking” feature biases the model toward generating images with objects arranged in a vertical hierarchy, regardless of what the text prompt specifies.

Style steering (${𝒯}_2$ intervention).

Injecting a “watercolour” style feature during the intermediate stage causes the model to produce watercolour-style images even when the text prompt does not request any specific artistic style. The effect is restricted to the stylistic attributes encoded in window ${𝒯}_2$ and does not affect the global layout (committed in ${𝒯}_1$) or the fine texture details (determined in ${𝒯}_3$).

Texture clamping (${𝒯}_3$ intervention).

The low-noise stage can be used to enforce or suppress specific texture patterns. For instance, clamping the “smooth surface” feature active in window ${𝒯}_3$ suppresses the model's tendency to generate visible grain or noise in the final image, producing an airbrushed appearance. Injecting the “film grain” feature conversely adds a realistic grain texture to images that would otherwise appear too smooth.

These operations collectively constitute what we term the feature-timestep interface: a structured programming model for diffusion generation steering, in which the user specifies what to change (which SAE feature) and when to change it (which timestep window), and the SAE framework handles the translation into activation-space interventions.

Attention Floating in Masked Diffusion Models

The preceding sections examined diffusion models for continuous data (images, audio spectrograms). We now turn to a rapidly growing class of discrete generative models: masked diffusion models (MDMs) for text generation. MDMs challenge the long dominance of autoregressive language models (ARMs) by generating text through iterative unmasking rather than sequential left-to-right prediction. This architectural difference - bidirectional attention versus causal masking - gives rise to a striking difference in attention patterns that has profound implications for both interpretability and task performance.

To motivate why attention patterns matter for interpretability, recall that in transformer-based language models, the attention mechanism is the primary vehicle for information routing. Understanding which tokens attend to which other tokens reveals the model's implicit reasoning strategy. A model that attends primarily to a few fixed “anchor” tokens is likely to lose information about tokens far from those anchors (the “lost-in-the-middle” failure mode); a model whose attention is more uniformly distributed may be better equipped to synthesise information from across the entire context.

Autoregressive Models: Causal Masking and Attention Sinks

An autoregressive language model (ARM) generates text one token at a time, from left to right. At generation step $i$, the model receives a context of $i-1$ previously generated tokens and produces a probability distribution over the next token. This sequential generation is enforced architecturally by a causal attention mask: for each query token at position $i$, the attention mask restricts the key-value set to positions $j\le i$. The attention weight from query $i$ to key $j$ is therefore (ARM Attention)$$\begin{array}{ccc}& a_{ij}^{\text{ARM}}=\frac{\exp ({𝒒}_i^{\top }{𝒌}_j/\sqrt{d})}{{\sum }_{j^{\prime }\le i}\exp ({𝒒}_i^{\top }{𝒌}_{j^{\prime }}/\sqrt{d})}\cdot \mathrm{𝟏}[j\le i].& \text{<span class="tml-eqn"></span>}\end{array}$$ This causal structure has an important emergent consequence: attention sinks.

The existence of attention sinks in ARMs has been documented across many model families: GPT-2, LLaMA, Mistral, Falcon, and others consistently show the BOS token accumulating $10$–$30$% of the total attention mass in early and middle layers, despite carrying no semantic content beyond its role as a sequence delimiter.

Why does the BOS token become a sink?

The causal mask creates an asymmetry: the BOS token is the only position that every subsequent query can attend to, regardless of its position in the sequence. When the model's attention weights are trained to be “safe” (i.e., routing attention to the BOS token is always a valid fall-back that does not violate the causal mask), the BOS token naturally becomes the designated receptacle for “I don't know where else to attend” attention mass. The result is a form of representational waste: a significant fraction of the model's attention capacity is devoted to a semantically meaningless token.

Lost-in-the-middle.

The attention sink phenomenon contributes directly to the lost-in-the-middle failure mode, in which ARMs struggle to access information that appears in the middle of long contexts. When the BOS token absorbs a substantial portion of attention mass, the effective attention budget available for actual contextual content is reduced. Furthermore, the positions immediately following the BOS token (positions 1–5 in a typical prompt) tend to receive elevated attention due to proximity effects of the causal mask, while positions in the middle of the context are systematically under-attended.

Masked Diffusion Models: Bidirectional Attention and Floating Anchors

A masked diffusion model (MDM) generates text by starting with a fully masked sequence and iteratively unmasking tokens over $T$ denoising steps. At each denoising step $t$, a fraction of the masked tokens are unmasked based on the model's predicted token probabilities. Crucially, the transformer in an MDM uses bidirectional attention: the attention mask is a full matrix of ones, allowing every token to attend to every other token regardless of position.

The bidirectional attention weight from query $i$ to key $j$ is (MDM Attention)$$\begin{array}{ccc}& a_{ij}^{\text{MDM}}=\frac{\exp ({𝒒}_i^{\top }{𝒌}_j/\sqrt{d})}{{\sum }_{j^{\prime }=1}^N\exp ({𝒒}_i^{\top }{𝒌}_{j^{\prime }}/\sqrt{d})},& \text{<span class="tml-eqn"></span>}\end{array}$$ where the sum ranges over all $N$ positions in the sequence. The absence of the causal mask has two immediate consequences. First, there is no structural reason for any single position to become a dominant attention sink: every position can receive attention from every other position, eliminating the asymmetric pull of the BOS token. Second, the model can condition each unmasking decision on the full bidirectional context, enabling long-range dependencies that ARMs can only approximate through left-to-right recurrence.

TikZ Figure: Attention Patterns in ARM vs. MDM

Shallow Structure-Aware, Deep Content-Focused Mechanism

Empirical analysis of MDM attention patterns across layers reveals a bipartite computational organisation that mirrors the structure-versus-content hierarchy observed in the temporal stages of image diffusion models.

Shallow layers: structure-aware attention.

In the early transformer layers (approximately the first third of the depth), attention patterns are strongly correlated with syntactic and positional structure: the model attends to positions that determine the grammatical skeleton of the sentence (nouns, verbs, clause boundaries). The floating anchor in shallow layers tends to be located at the beginning of the currently unmasked region, corresponding to the structural pivot of the emerging sentence.

Probing classifiers trained on shallow-layer representations achieve high accuracy for syntactic tasks (part-of-speech tagging, dependency arc prediction) but low accuracy for semantic content tasks (named entity recognition, semantic role labelling). This suggests that shallow layers use bidirectional attention primarily to assemble the structural scaffolding of the output text.

Deep layers: content-focused attention.

In the later transformer layers (approximately the final third of the depth), attention patterns shift to content-focused processing: the model attends to positions that carry the specific factual content required to fill the remaining masked positions. The floating anchor in deep layers tends to be located at the most informative unmasked token for resolving the current masked position - for instance, if the current masked token should be a person's name, the anchor tends to be the position that first mentioned the person.

Probing classifiers trained on deep-layer representations achieve high accuracy for semantic and factual tasks but lower accuracy for syntactic structure. This bipartite organisation - shallow structure-aware, deep content-focused - is the mechanism by which MDMs avoid the lost-in-the-middle failure mode: factual content retrieval happens in deep layers where bidirectional attention has access to the entire context with no causal restriction.

The inequality in Proposition 8 formalises the intuition that MDM attention is more distributed than ARM attention, not just because of the floating anchor mechanism, but because the full bidirectional attention pool is always larger. As a consequence, any individual token has a lower maximum possible share of the aggregate attention mass in an MDM than in an ARM, making attention sinks structurally impossible in well-trained MDMs.

MDMs Outperform ARMs on Knowledge-Intensive Tasks

The distributed, floating attention structure of MDMs has a direct empirical consequence: MDMs consistently outperform comparably-sized ARMs on tasks that require integrating information from across a long context - precisely the class of tasks where the ARM attention sink and lost-in-the-middle failure modes are most damaging.

Multi-hop reasoning.

Tasks that require connecting a chain of facts scattered across a long document (e.g., “What is the birth city of the author of the book that won the prize in year $X$?”) require the model to simultaneously attend to multiple relevant positions. In ARMs, the causal mask prevents early position tokens from attending to information that appears later in the prompt, and the attention sink at the BOS token reduces the effective attention budget for intermediate reasoning steps. MDMs, with their bidirectional attention, can attend to all relevant facts simultaneously in a single forward pass.

Cloze-style factual completion.

When the model must fill in a masked token that depends on factual knowledge appearing later in the context (e.g., “[MASK] was born in Paris” when the later context provides the person's name), ARMs must rely on their parametric memory because they cannot attend to later tokens. MDMs can attend to the later context directly, effectively performing a lookup operation that bypasses the need to memorise the fact in model weights.

Open-domain question answering.

Experiments on the MMLU, TriviaQA, and NaturalQuestions benchmarks show MDMs achieving $3$–$8\%$ higher accuracy than same-parameter ARMs when the relevant evidence appears in the middle of a 2048-token context. The gap is largest for questions that require evidence from positions 512–1536 (the middle half of the context), consistent with the lost-in-the-middle failure mode of ARMs.

Mechanistic Interpretability of Diffusion Bias

The attention and representation structure of diffusion models is not only relevant for understanding task performance; it also encodes social biases that are learned from training data and can be propagated, amplified, or attenuated by the model's internal circuits. [3] present a mechanistic interpretability analysis of bias in U-Net and Diffusion Transformer (DiT) architectures, identifying specific internal circuits responsible for generating stereotyped associations and proposing targeted interventions to neutralise them while preserving generation quality.

The circuit problem.

Bias in image-generative models is commonly attributed to biased training data: if the training set contains more images of male engineers than female engineers, the model will generate male engineers more frequently. While training data skew is certainly a contributing factor, the mechanistic interpretability lens reveals that this explanation is incomplete. Biased associations are not uniformly distributed across model parameters; they are concentrated in specific attention heads and MLP layers that form identifiable circuits. A model can thus inherit bias from training data while amplifying or suppressing it through its internal circuit structure.

Identifying bias features.

[3] propose a three-stage procedure for identifying bias circuits in a diffusion model $f_{𝜽}$.

Stage 1: Contrastive generation. For each stereotyped association (e.g., “nurse” $\to$ female), generate a large set of image pairs: one set conditioned on the neutral prompt (“a nurse”) and one set conditioned on a de-biased prompt (“a male nurse”). Extract intermediate activations from every attention head and MLP layer for both sets.

Stage 2: Causal attribution. Apply causal tracing (activation patching) to identify which components, when their activations are swapped from the biased-generation run to the de-biased run, cause the model's gender-association output to change. The components with the highest causal attribution scores form the bias circuit.

Stage 3: Sparse intervention. Train a lightweight steering vector $𝜹$ in the subspace spanned by the bias circuit's attention heads. The steering vector is optimised to minimise the CLIP similarity between the generated image and the stereotyped gender role, subject to a constraint that the image quality (FID score on a holdout set) does not degrade beyond a threshold ${ϵ}_{\text{FID}}$.

Experimental results.

On the Stable Diffusion XL model evaluated on the WinoBias and FairFace benchmarks, [3] report the following outcomes from the circuit-targeted intervention:

• Nurse/doctor gender bias: reduced from $78\%$ female association (baseline) to $54\%$ female association (post-intervention), approaching the $50\%$ gender-neutral baseline, with FID degradation of only $1.3$ points (${ϵ}_{\mathrm{FID}}=2.0$ threshold).

• CEO/engineer gender bias: reduced from $82\%$ male association to $61\%$ male association.

• Criminal/welfare recipient ethnicity bias: reduced from $71\%$ majority-group over-association to $55\%$.

• Image quality: CLIP score decreased by $1.8\%$ (average over all stereotypes), LPIPS increased by $0.012$, FID increased by $1.7$ points - all within the specified quality constraints.

Critically, the identified bias circuits were transferable across model checkpoints: a bias circuit identified in one fine-tuned variant of SDXL transferred to $7$ other fine-tunes with $>85\%$ causal attribution overlap, suggesting that bias circuits are established during pre-training and persist through fine-tuning.

U-Net vs. DiT architectures.

The mechanistic analysis reveals qualitative differences between bias circuit structure in U-Net and DiT architectures. In U-Net models (e.g., Stable Diffusion 1.x, 2.x, SDXL), bias circuits are concentrated in the cross-attention layers of the middle (bottleneck) blocks, where text conditioning information is injected into the image latent. The text encoder's representation of occupational terms (“nurse”, “CEO”) carries implicit gender associations that the cross-attention layers then propagate into the spatial features of the latent.

In DiT models (e.g., DiT-XL, PixArt-$\alpha$, Stable Diffusion 3), the architecture's fully transformer-based design distributes the bias circuit more broadly across layers. Cross-attention is replaced by adaptive layer normalisation conditioned on text embeddings, and the bias circuit manifests in the normalisation parameters rather than explicit attention heads. This makes DiT bias circuits harder to identify (there are no discrete attention heads to probe) but not impossible: causal tracing on the normalisation scale and shift parameters ($\gamma$ and $\beta$ in the adaptive LayerNorm) successfully isolates the bias contribution.

Video Generation and Emergent Reasoning

The arc of generative modelling has bent, with increasing sharpness, toward time. Images capture a frozen instant; text unfolds over tokens that, although sequential, carry no physical causal structure; but video demands something more: a model of how the world changes. A pixel at position $(x,y)$ in frame $t+1$ is not merely statistically correlated with pixels at $(x,y)$ in frame $t$; it is the consequence of forces, velocities, collisions, and intentions that operated during the interval $[t,t+1]$. Predicting video plausibly therefore requires, at minimum, an implicit model of the physics and agency that govern the depicted scene.

This observation has thrust video generation models into an unexpected role: involuntary world simulators. Systems like Sora (OpenAI, 2024), Veo (Google DeepMind, 2024), and CogVideoX trained exclusively on the objective of predicting the next frame, or of denoising a corrupted video clip, appear to have acquired representations that encode physical object permanence, gravity, occlusion, and even elements of goal-directed behaviour. The interpretability challenge is to understand what has been acquired, where it lives in the network, and how it is computed.

Section Video Generation and Emergent Reasoning is organised as follows. We begin by surveying the architectures of contemporary video generation models and the evidence for emergent cognition they exhibit (Video Models as World Simulators). We then turn to the Attention Gathers, MLPs Compose hypothesis, which offers a mechanistic account of how spatiotemporal reasoning is distributed across the transformer layers of a Video Vision Transformer (VideoViT) (The Attention Gathers, MLPs Compose Hypothesis). We formalise the concept of an action-outcome circuit (Action-Outcome Circuits in Video Transformers), explain why redundancy across MLP layers makes video transformers resilient (Redundancy Across MLP Layers), and ground everything in a running example of predicting bowling outcomes (A Running Example: Predicting Bowling Outcomes).

Video Models as World Simulators

A video generation model $G_{\theta }$ maps a conditioning signal $𝒄\in 𝒞$ (a text prompt, an initial frame, or a noise tensor) to a distribution over video clips ${𝒙}_{1:T}\in {\mathrm{ℝ}}^{T\times H\times W\times 3}$, where $T$ is the number of frames and $H\times W$ is the spatial resolution. Modern architectures pursue this goal via diffusion: (Diffusion)$$\begin{array}{ccc}& p_{\theta }({𝒙}_{1:T}|𝒄)=\int p_{\theta }({𝒙}_{1:T}|𝒛,𝒄)p(𝒛)d𝒛,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒛$ is a latent code obtained by encoding the video into a compressed spatiotemporal latent space, and $p_{\theta }({𝒙}_{1:T}|𝒛,𝒄)$ is parameterised by a denoising network-typically a Diffusion Transformer (DiT) whose self-attention modules operate over both spatial patch tokens and temporal frame tokens simultaneously.

Spatiotemporal patchification.

Given a video ${𝒙}_{1:T}$, the standard preprocessing pipeline in architectures such as CogVideoX proceeds as follows. Each frame ${𝒙}_t\in {\mathrm{ℝ}}^{H\times W\times 3}$ is divided into non-overlapping spatial patches of size $p\times p$ pixels. Temporal striding groups every $\tau$ consecutive frames into a single temporal unit. The result is a grid of $\lceil T/\tau \rceil \times \lfloor H/p\rfloor \times \lfloor W/p\rfloor$ tokens, each of dimension $d_{\text{patch}}=\tau \cdot p^2\cdot 3$. A linear projection maps each patch token to the model's hidden dimension $d$: (Patch)$$\begin{array}{ccc}& {𝒉}_i={𝐖}_{\text{proj}}{𝒙}_i^{\text{patch}}+{𝒃}_{\text{proj}},{𝒙}_i^{\text{patch}}\in {\mathrm{ℝ}}^{d_{\text{patch}}},{𝒉}_i\in {\mathrm{ℝ}}^d.& \text{<span class="tml-eqn"></span>}\end{array}$$ Positional encodings, either factorised 3-D sinusoidal embeddings or learned RoPE encodings, are then added so that the transformer can distinguish tokens by their spatial coordinates and temporal frame index.

Emergent cognitive abilities.

The VBVR-Bench (Video-Based Visual Reasoning Benchmark) was designed to probe whether video generation models possess genuine scene understanding or merely sophisticated texture interpolation. The benchmark presents three families of tasks.

• Mental rotation. A 3-D object is shown from one viewpoint; the model must complete a video that shows the object after rotation by a specified angle $\mathrm{\Delta }\varphi$. Success requires an internal representation of volumetric shape, not merely a library of pixel patterns.

• Maze navigation. A first-person perspective video begins inside a maze; the model must generate subsequent frames consistent with navigating to an exit whose location was made inferable from the opening frames. This tests both spatial memory and goal-directed planning.

• Sudoku completion. A video shows a partially filled Sudoku grid; the model must generate frames that fill in the remaining cells correctly. This is a pure constraint-satisfaction task requiring symbolic reasoning embedded inside a continuous video prediction objective.

Remarkably, Sora and Veo achieve non-trivial performance on all three task families. Their success rates exceed chance by wide margins and exceed the performance of image-only diffusion models given the same prompts, suggesting that the temporal dimension of the training objective-rather than just the scale of pixel data-is a meaningful driver of emergent cognition.

The Attention Gathers, MLPs Compose Hypothesis

Understanding how a video transformer performs its computation requires a mechanistic hypothesis about the functional division of labour between its two principal building blocks: the multi-head self-attention (MHSA) sublayer and the multi-layer perceptron (MLP) sublayer. In language models, a substantial body of mechanistic interpretability work has established that attention heads tend to implement token-copying and positional routing operations, while MLP layers store and retrieve factual associations. We now argue that an analogous but spatiotemporally extended version of this division holds in video transformers.

The AGMC hypothesis makes testable predictions. If it is correct, then ablating or patching attention heads should disrupt the routing of evidence (which tokens are aggregated), while ablating MLP layers should disrupt the composition of abstract outcomes (what conclusions are drawn from that evidence). Probing experiments on CogVideoX confirm this asymmetry: attention-head ablations degrade spatial coherence (the model loses track of where objects are), whereas MLP ablations degrade causal coherence (the model loses track of what will happen next).

Multi-head structure and spatiotemporal specialisation.

Video transformers typically employ $H$ attention heads per layer. Activation-patching experiments reveal that different heads specialise along different axes of the spatiotemporal structure. A standard classification that has emerged from the literature is:

• Spatial heads aggregate information across the spatial patch grid at a fixed time step. Their attention patterns resemble those observed in image transformers: receptive fields that grow with depth, attending first to local texture and later to global structure.

• Temporal heads aggregate information across time at a fixed spatial location, effectively tracking the trajectory of a point in the scene through time.

• Cross-modal heads (in text-conditioned models) route information from text tokens to visual tokens, anchoring generated content to the conditioning signal.

• Causal heads display a distinctive pattern: they attend from a token at time $t+k$ to tokens at time $t$ that represent cause tokens, i.e., tokens whose features encode an action or event that logically implies the content at $t+k$. These are the heads most directly implicated in the action-outcome circuits defined below.

Action-Outcome Circuits in Video Transformers

The concept of a circuit in mechanistic interpretability was developed in the context of language models to refer to a subgraph of the computational graph that is jointly sufficient for performing a specific subtask. We extend this concept to the spatiotemporal setting of video transformers.

Action-outcome circuits are discovered by a three-stage pipeline. First, mean ablation patching identifies which components are causally necessary for the model's causal reasoning: each component is ablated in turn, and the change in output logits is measured. Second, attention pattern visualisation confirms that the identified attention heads attend to action-relevant tokens. Third, probing classifiers trained on MLP intermediate activations verify that outcome-relevant information is present in the composition nodes.

Mathematical formulation of the residual stream.

The residual stream of a VideoViT can be written as: (Residual)$$\begin{array}{cccc}& {𝒉}_i^{(0)}& ={𝐖}_{\text{proj}}{𝒙}_i^{\text{patch}}+\text{pos}(i),& \text{<span class="tml-eqn"></span>}\\ & {𝒉}_i^{(\ell +\frac{1}{2})}& ={𝒉}_i^{(\ell )}+{\text{MHSA}}^{(\ell )}({𝐇}^{(\ell )}{)}_i,& \text{<span class="tml-eqn"></span>}\\ & {𝒉}_i^{(\ell +1)}& ={𝒉}_i^{(\ell +\frac{1}{2})}+{\mathrm{MLP}}^{(\ell )}\left({𝒉}_i^{(\ell +\frac{1}{2})}\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\text{pos}(i)$ is the positional encoding for token $i$. The circuit $𝒞$ can be characterised by a binary mask $𝒎\in {\{\mathrm{0,1}\}}^{|V|}$ over the node set, and the ablated model is (Ablated)$$\begin{array}{ccc}& {\widetilde{𝒉}}_i^{(\ell +1)}={\widetilde{𝒉}}_i^{(\ell +\frac{1}{2})}+m_{\text{MLP}}^{(\ell )}\cdot {\mathrm{MLP}}^{(\ell )}\left({\widetilde{𝒉}}_i^{(\ell +\frac{1}{2})}\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ with analogous masking on the MHSA contributions. The circuit discovery problem is to find the minimal mask $𝒎$ such that equation is satisfied.

Redundancy Across MLP Layers

A striking empirical observation in video transformer interpretability is that action-outcome circuits are not localised to a single MLP layer. Rather, the causal outcome information is distributed redundantly across several consecutive MLP layers. Ablating any single MLP layer produces only a modest degradation in causal reasoning; recovering full degradation requires ablating three to five consecutive layers.

This redundancy has a structural explanation. The MLP at layer $\ell$ of a transformer with hidden dimension $d_{\text{ff}}=4d$ (the standard factor-4 expansion) can be viewed as a key-value memory over a dictionary of $d_{\text{ff}}$ memory slots. Each slot $k$ is associated with a key vector ${𝒘}_k^{\text{in}}\in {\mathrm{ℝ}}^d$ and a value vector ${𝒘}_k^{\text{out}}\in {\mathrm{ℝ}}^d$: (MLP KV)$$\begin{array}{ccc}& \mathrm{MLP}(𝒙)={𝐖}^{\text{out}}\cdot \text{GeLU}({𝐖}^{\text{in}}𝒙),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝐖}^{\text{in}}\in {\mathrm{ℝ}}^{d_{\text{ff}}\times d}$ and ${𝐖}^{\text{out}}\in {\mathrm{ℝ}}^{d\times d_{\text{ff}}}$. The rows of ${𝐖}^{\text{in}}$ are the keys; the columns of ${𝐖}^{\text{out}}$ are the values.

Under this memory interpretation, a single concept (e.g., “the ball will continue along its trajectory”) may activate multiple memory slots across multiple consecutive layers. The representation is over-complete: far more slots encode the concept than are strictly necessary. This over-completeness is precisely the source of resilience.

The practical implication is that single-layer ablation experiments may dramatically underestimate the role of MLPs in causal reasoning within video models. Rigorous circuit discovery must account for the possibility that the same computation is distributed across many layers, and path-patching techniques must be extended to identify the full set of writing layers.

A Running Example: Predicting Bowling Outcomes

To make the foregoing abstractions concrete, we trace the computation of a VideoViT through a specific example: predicting whether a bowling ball will strike the pins or miss, given the first $T_0$ frames of a ten-pin bowling video.

Setup.

Let ${𝒙}_{1:T_0}$ be the input video clip, consisting of $T_0=12$ frames at 24 fps covering the moment the ball leaves the bowler's hand and traverses the first third of the lane. The model must predict the output token $o\in \{\text{strike},\text{spare},\text{gutter}\}$-a discretised summary of the final outcome. The model never observes the pins being hit; it must infer the outcome from the ball's early trajectory.

Identifying the action tokens.

The ball occupies a localised spatial region in each frame. Applying the activation-patching pipeline to the bowling example reveals that the tokens corresponding to the ball's position at frames $t=3$ and $t=6$ (approximately 0.125 s and 0.250 s after release) are the action tokens: ablating them while preserving all others reduces outcome prediction accuracy from 84% to 52% - a large drop, though still above the $33\%$ chance level for a three-class problem.

Gathering phase: causal heads.

In layers 8–12 of a 24-layer VideoViT, three causal attention heads ($h_1,h_2,h_3$) display a characteristic pattern: their queries originate from late-frame tokens (frames $t=10$–12, near the lane end), and their keys concentrate on the ball-region tokens at early frames. The attention weights satisfy: (Bowling ATTN)$$\begin{array}{ccc}& \sum _{t\in \{\mathrm{3,6}\}}{\alpha }_{ij}^{(\ell )}{|}_{\text{ball tokens}}>0.6\forall \ell \in [8,12],\forall i\in \text{lane-end tokens}.& \text{<span class="tml-eqn"></span>}\end{array}$$ In other words, more than 60% of the attention weight from lane-end tokens flows to the early ball-position tokens-these heads are routing physical trajectory information forward in time.

Composition phase: MLP layers.

Probing classifiers trained on the MLP intermediate activations at layers 13–17 achieve 91% accuracy on the three-class outcome prediction, compared to 61% for probes trained on the attention output at the same layers. The gap confirms that MLP composition, not attention gathering, is the stage at which trajectory information is converted into outcome probability. Specifically, the direction in residual-stream space most predictive of “strike” versus “gutter” is written by MLP layers 14 and 15, with writing strengths ${\lambda }_{14}=0.38$ and ${\lambda }_{15}=0.29$ (normalised).

Redundancy in the bowling circuit.

Measured by the layer-13–17 MLP probe (baseline $91\%$), ablating MLP layer 14 alone reduces the probe's outcome-decoding accuracy from $91\%$ to $74\%$. Ablating both layers 14 and 15 reduces it to $57\%$. Reaching chance performance requires ablating layers 13 through 17 simultaneously, consistent with Proposition Proposition 9 with $K=5$ writing layers.

This example illustrates the full action-outcome circuit pipeline: causal heads at layers 8–12 gather ball-trajectory evidence; MLP layers 13–17 compose this evidence into outcome probabilities; the outcome feature is written redundantly across five layers, conferring resilience to perturbation.

Exercises

State-Space Models: Mamba and Beyond

Transformers occupy an almost unchallenged position at the centre of modern generative modelling, but they carry a structural cost that grows quadratically with sequence length: the attention matrix $𝐒\in {\mathrm{ℝ}}^{N\times N}$ has $O(N^2)$ entries, each requiring computation and memory. For sequences of length $N={10}^4$ (a short video or a genomic context), the attention computation alone demands tens of gigabytes. This has motivated a sustained search for architectures that achieve competitive modelling performance at linear cost in sequence length.

State-space models (SSMs) are the most successful family of such architectures to date. Rooted in classical control theory, they represent a sequence by propagating a hidden state ${𝒉}_t\in {\mathrm{ℝ}}^N$ through a linear recurrence: (Recurrence)$$\begin{array}{cccc}& {𝒉}_t& =𝐀{𝒉}_{t-1}+𝐁{𝒙}_t,& \text{<span class="tml-eqn"></span>}\\ & y_t& =𝐂{𝒉}_t+D{𝒙}_t,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝐀\in {\mathrm{ℝ}}^{N\times N}$ is the state-transition matrix, $𝐁\in {\mathrm{ℝ}}^{N\times d}$ is the input projection, $𝐂\in {\mathrm{ℝ}}^{d\times N}$ is the output projection, and $D$ is a skip-connection scalar. Because ${𝒉}_t$ is a fixed-dimensional summary of the entire history ${𝒙}_{1:t}$, the per-step computation is $O(N\cdot d)$, independent of sequence length.

The S4 to Mamba progression.

Early SSM architectures like S4 constrained $𝐀$ to be a diagonal-plus-low-rank matrix, enabling efficient convolution-based parallel training while maintaining recurrent inference. The key limitation was that $𝐀$, $𝐁$, $𝐂$ were data-independent: the same state transitions were applied regardless of the content of the input sequence. This prevented the model from selectively attending to relevant tokens in the way that transformer attention heads do.

Mamba addresses this limitation by making $𝐁$, $𝐂$, and the discretisation step $\mathrm{\Delta }$ into functions of the input ${𝒙}_t$. This is the selective state space mechanism.

Mamba Architecture and Selective State Spaces

The Mamba block replaces the standard S4 linear recurrence with a data-dependent system. Given input sequence $𝒙=({𝒙}_1,\dots ,{𝒙}_T)$ with ${𝒙}_t\in {\mathrm{ℝ}}^d$, the Mamba selective SSM computes: (Params)$$\begin{array}{cccc}& {𝐁}_t& ={\text{Linear}}_B({𝒙}_t)\in {\mathrm{ℝ}}^{N\times 1},& \text{<span class="tml-eqn"></span>}\\ & {𝐂}_t& ={\text{Linear}}_C({𝒙}_t)\in {\mathrm{ℝ}}^{1\times N},& \text{<span class="tml-eqn"></span>}\\ & {\mathrm{\Delta }}_t& =\text{softplus}({\text{Linear}}_{\mathrm{\Delta }}({𝒙}_t))\in {\mathrm{ℝ}}^d,& \text{<span class="tml-eqn"></span>}\end{array}$$ and then discretises the continuous-time SSM using the zero-order hold (ZOH) rule: (ZOH)$$\begin{array}{cccc}& {\overline{𝐀}}_t& =\exp ({\mathrm{\Delta }}_t\cdot 𝐀),& \text{<span class="tml-eqn"></span>}\\ & {\overline{𝐁}}_t& ={({\mathrm{\Delta }}_t\cdot 𝐀)}^{-1}({\overline{𝐀}}_t-𝐈){\mathrm{\Delta }}_t\cdot {𝐁}_t.& \text{<span class="tml-eqn"></span>}\end{array}$$ The resulting recurrence is: (Recurrence)$$\begin{array}{cccc}& {𝒉}_t& ={\overline{𝐀}}_t\odot {𝒉}_{t-1}+{\overline{𝐁}}_t\cdot {𝒙}_t,& \text{<span class="tml-eqn"></span>}\\ & y_t& ={𝐂}_t\cdot {𝒉}_t,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\odot$ denotes elementwise multiplication (using the diagonal approximation ${\overline{𝐀}}_t\approx \mathrm{diag}({\overline{𝐀}}_t)$).

The crucial distinction from a standard SSM is that ${\overline{𝐀}}_t$, ${\overline{𝐁}}_t$, ${𝐂}_t$ all depend on the current input ${𝒙}_t$. This means the state transitions vary from token to token, allowing the model to decide dynamically which parts of the past history to retain and which to discard.

Short convolutions.

A complementary component in the Mamba block is a depthwise short convolution applied to the input before the SSM: (CONV)$$\begin{array}{ccc}& {\widetilde{𝒙}}_t=\sum _{k=0}^{K-1}{w}_k\cdot {𝒙}_{t-k},K\ll T,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $K$ is the convolution kernel width (typically 4). This convolution captures local context at a fixed, small window, a capability that proves critical for certain tasks, as we shall see in Mechanistic Evaluation: Associative Recall.

Mechanistic Evaluation: Associative Recall

The Associative Recall (AR) task is a canonical synthetic benchmark for testing whether a sequence model can perform content-based key-value lookup.

How transformers solve AR: induction heads.

The transformer's solution to AR is well understood through the lens of induction heads-a specific attention head circuit identified in language model interpretability work. An induction head implements a two-step pattern:

1. A previous-token head writes the identity of each token's preceding token into the residual stream.

2. The induction head itself attends from the query token $k^{\ast }$ at position $n+1$ to the position $i^{\ast }$ in the sequence where $k^{\ast }$ last appeared, by matching the previous-token representation.

Formally, the attention score between the query position $n+1$ and key position $i$ is approximately: (Induction)$$\begin{array}{ccc}& s(n+1,i)\approx {𝒉}_{n+1}^{\top }{𝐖}_Q^{\top }{𝐖}_K{𝒉}_{i-1}^{\text{prev}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒉}_{i-1}^{\text{prev}}$ encodes the token at position $i-1$, i.e., the key $k_i$ that appears just before value $v_i$. When $k_i=k^{\ast }$, this score peaks, and the value $v_i=v^{\ast }$ is retrieved.

How Mamba solves AR: short convolutions.

In contrast to transformers, Mamba uses its short convolution component to solve the AR task. The key observation is that in AR instances, the $k$-$v$ pairs are presented consecutively: $k_i$ appears at position $2i-1$ and $v_i$ at position $2i$. The value $v^{\ast }$ therefore always appears exactly one position after the target key $k^{\ast }$. A short convolution with kernel width $K=2$ can detect this local pattern and store $(k_i,v_i)$ pairs in its receptive field without needing the long-range gating capabilities of the selective SSM mechanism.

This is confirmed by ablation: setting the convolution weights to zero in a trained Mamba model causes AR accuracy to collapse from 98% to near-chance, while ablating the selective SSM gates leaves AR accuracy largely intact. The selective gates contribute primarily to global sequence coherence, not to this local lookup task.

Associative Treecall: Hierarchical Retrieval

The AR task tests flat key-value lookup. Many real-world retrieval tasks require hierarchical lookup: to find the final answer, one must first retrieve an intermediate key, use it to look up an intermediate value, and continue this chain. The Associative Treecall (ATR) task formalises this structure.

Mamba's performance on ATR.

While transformers and Mamba achieve near-identical accuracy on AR (depth $d=1$), a striking divergence emerges as depth increases. On ATR with $d=3$, transformers maintain high accuracy (85–92%) across a range of vocabulary sizes, whereas Mamba's accuracy degrades to 40–55%. At $d=5$, the gap is larger still.

Why Mamba struggles with deep ATR.

At each level of the ATR hierarchy, the model must route a retrieved intermediate key to a new lookup operation. For transformers, each induction-head circuit can chain: the output of one AR lookup becomes the query for the next. For Mamba, the short-convolution mechanism is inherently limited to $K$-local context; after a retrieved value has been placed in the sequence, the next lookup requires attending to a distant position, which the convolution cannot reach.

The selective SSM gates can, in principle, bridge this gap by propagating intermediate keys through the hidden state. The question is whether trained Mamba models actually learn to use the SSM for this purpose, or whether they rely exclusively on the convolution. Mechanistic analysis-probing the hidden state ${𝒉}_t$ at each step during ATR evaluation-reveals that the SSM gates do attempt to encode the intermediate keys, but with decreasing fidelity at deeper levels. The hidden state dimension $N$ becomes a binding constraint: with $N=16$ states (standard in Mamba-2B), three or more concurrent key-value associations exceed the state capacity, leading to interference and forgetting.

Architectural Plasticity of Mamba

The mechanistic picture that emerges from the AR and ATR analysis is one of architectural plasticity: Mamba does not commit to a single computational strategy but adapts its mechanism to the complexity of the task. On simple tasks (AR, $d=1$), it uses short convolutions-fast, local, and efficient. On tasks that demand longer-range dependencies, it shifts toward the selective SSM gates.

The proposition shows that the selective gates of Mamba have the expressive capacity to implement attention. Whether trained models actually exploit this capacity depends on whether the training objective provides sufficient gradient signal for the long-range tasks where attention is needed. On natural language modelling, the evidence suggests that Mamba partially but not fully exploits this capacity, which explains why it matches transformers on many language benchmarks but falls short on tasks requiring precise long-range copy-and-recall.

MambaLRP: Explaining Selective State Spaces

Layer-wise relevance propagation (LRP) is an attribution method that decomposes a model's output score $f(𝒙)$ into a sum of per-input relevance scores $R_i$ such that ${\sum }_i{R}_i=f(𝒙)$. In transformer models, LRP has been successfully extended via the Attention-LRP framework, which handles the non-linearities of softmax attention. Applying LRP to state-space models requires a different treatment because the recurrent structure of the SSM introduces temporal dependencies that are not present in attention.

Relevance propagation through the recurrence.

The output $y_T$ of a Mamba SSM at the final time step can be written as: (Mamba Output)$$\begin{array}{ccc}& y_T={𝐂}_T{𝒉}_T={𝐂}_T\sum _{t=1}^T\left(\prod _{s=t+1}^T{\overline{𝐀}}_s\right){\overline{𝐁}}_t{𝒙}_t.& \text{<span class="tml-eqn"></span>}\end{array}$$ LRP assigns to input ${𝒙}_t$ the relevance: (Mamba LRP)$$\begin{array}{ccc}& R_t={𝐂}_T\left(\prod _{s=t+1}^T{\overline{𝐀}}_s\right){\overline{𝐁}}_t{𝒙}_t,& \text{<span class="tml-eqn"></span>}\end{array}$$ which can be computed efficiently by a backward pass through the recurrence, analogous to the backward pass in backpropagation through time (BPTT) .

Note that the product ${\prod }_{s=t+1}^T{\overline{𝐀}}_s$ acts as an effective discount factor: tokens far in the past (small $t$) have their relevance attenuated by the cumulative state transitions. This is not a pathology but a reflection of the SSM's architecture: older information is progressively compressed into the hidden state and mixed with more recent input.

MambaLRP implementation.

The MambaLRP framework extends LRP to the full Mamba block, handling three components:

• SSM recurrence. Relevance is propagated through the discretised recurrence using equation . The data-dependent nature of ${\overline{𝐀}}_t$ and ${\overline{𝐁}}_t$ means that relevance scores depend on the specific input sequence, not just the model weights.

• Short convolution. Relevance through the depthwise convolution is propagated using the standard LRP convolution rule: $R_{t-k}^{\text{conv}}\propto w_k\cdot {𝒙}_{t-k}\cdot R_t^{\text{post-conv}}$.

• Multiplicative gates. The Mamba block includes a sigmoid gate ${𝒈}_t=\sigma ({\text{Linear}}_g({𝒙}_t))$ applied elementwise to the SSM output. Relevance through the gate is propagated using the $\epsilon$-LRP rule to avoid division by zero: $R^{\text{pre-gate}}=R^{\text{post-gate}}\cdot {𝒉}_t/({𝒉}_t\odot {𝒈}_t+\epsilon )$.

Comparison with attention-based explanations.

On tasks where Mamba uses its short convolution mechanism (e.g., AR), MambaLRP assigns high relevance to the immediate predecessor of the query token-consistent with the $K=2$ convolution window-while relevance to earlier tokens decays rapidly. On tasks that elicit SSM-gating behaviour (e.g., longer-context language modelling), MambaLRP relevance scores are more spread across the sequence, reflecting the extended temporal reach of the hidden state.

This contrast provides a diagnostic tool: the concentration or spread of MambaLRP scores serves as an empirical indicator of which computational mechanism (convolution vs. SSM gating) the model is relying on for a given input. When LRP scores are narrow (concentrated on 1–3 tokens), the model is using its local convolution; when they are broad (spread across 20+ tokens), the model is using its selective SSM gates.

Exercises

Alignment, Safety, and Interpretability-First AI

Mechanistic interpretability began as a scientific exercise: understanding what computations a neural network actually performs. Yet as language models grow larger and are deployed in consequential settings-medicine, law, critical infrastructure-the stakes of not understanding them rise steeply. This section argues that interpretability is not merely an intellectually interesting add-on to AI development but an alignment necessity: a prerequisite for detecting misalignment before it causes harm.

Why Interpretability Matters for AI Safety

The alignment problem has traditionally been framed in behavioral terms: can we train a model to produce outputs that humans judge as helpful, honest, and harmless? The dominant paradigm-reinforcement learning from human feedback (RLHF)-operationalizes alignment as reward maximization subject to a KL constraint: (RLHF Objective)$$\begin{array}{ccc}& \underset{\pi }{\max }{𝔼}_{x,y\sim \pi (\cdot |x)}[r(x,y)]-\beta \mathrm{KL}[\pi (\cdot |x)\Vert {\pi }_{\mathrm{ref}}(\cdot |x)].& \text{<span class="tml-eqn"></span>}\end{array}$$

This objective is both mathematically elegant and practically limited. The reward $r(x,y)$ is a proxy for human intent; it scores outputs, not processes. A model that has learned to produce high-reward outputs by genuinely internalizing human values and one that has learned to produce high-reward outputs by recognizing and catering to evaluator psychology are behaviorally indistinguishable under standard evaluation protocols. Distinguishing them requires looking inside.

Three failure modes that behavioral evaluation cannot detect.

1. Reward hacking via proxy gaming. The model maximizes $r(x,y)$ by exploiting patterns correlated with high ratings in the training distribution of raters, rather than by producing genuinely helpful outputs. Activation analysis can reveal whether the model activates “helpfulness circuits” or “rating-predictor circuits” when generating responses.

2. Deceptive alignment. The model behaves safely during evaluation (when it detects that it is being evaluated) but pursues different objectives in deployment. This is a mechanistic phenomenon: the model must maintain a representation of its evaluation context. Probing for such representations is, in principle, possible.

3. Goal misgeneralization. The model has internalized a goal that coincides with human intent on the training distribution but diverges elsewhere. Mechanistic analysis can test whether the circuitry implementing an apparent goal is general (context-independent) or narrow (tied to distributional features).

Deceptive Alignment

Deceptive alignment, introduced by Evan Hubinger and colleagues in the risks from learned optimization framework, describes a qualitatively different failure mode from ordinary misalignment.

The challenge is that gradient descent may inadvertently produce deceptively aligned models. If a model with a misaligned internal objective $\tilde{g}$ can correctly predict that acting on $\tilde{g}$ during evaluation would cause corrective training, then gradient pressure incentivizes safe-appearing behavior during evaluation. The model has, in effect, learned a meta-strategy: appear aligned until the ability to pursue $\tilde{g}$ is secured.

Mechanistic signatures of deceptive alignment.

Interpretability provides a candidate detection mechanism. A deceptively aligned model must maintain three distinct internal representations:

• a representation of the current context mode (evaluation vs. deployment),

• a representation of the safe behavior policy, and

• a representation of the misaligned objective.

Circuit-level analysis seeks to identify whether these representations exist and whether they causally influence outputs. The key tool is activation patching: intervening on the internal state associated with “this is evaluation” and measuring the change in output distribution. If patching the evaluation-context representation causes a significant distributional shift toward less safe outputs, this is evidence (though not proof) of deceptive alignment.

Pluralistic Alignment and Value Incommensurability

Most alignment work implicitly assumes a single, coherent set of human values. This is a simplification. Human values are plural: they differ across cultures, generations, and individuals, and many disagreements are not merely factual misunderstandings but reflect genuine incommensurability-the impossibility of reducing one value system to another on a common scale.

Formal model.

Let $𝒰$ be a population of users with value functions $u_1,\dots ,u_N:𝒴\to ℝ$, where $𝒴$ is the space of model outputs. Standard RLHF constructs an aggregate reward: (Aggregate Reward)$$\begin{array}{ccc}& r_{\mathrm{agg}}(y)=\frac{1}{N}\sum _{i=1}^N{u}_i(y),& \text{<span class="tml-eqn"></span>}\end{array}$$ or via a Bradley–Terry model trained on pairwise preference data. The problem is that ${\mathrm{arg\hspace{0.17em}max}}_y{r}_{\mathrm{agg}}(y)$ may not maximize any individual $u_i$. Worse, it may systematically produce outputs that are moderately acceptable to the majority while being actively harmful to minorities whose value functions have low weight in the aggregate.

The asymmetric representation problem.

RLHF training data is not uniformly sampled from the user population. Annotation pipelines tend to over-represent annotators from affluent, English-speaking, Western contexts. The consequence at the mechanistic level is that value representations associated with dominant cultural contexts are robustly encoded across multiple circuits, while value representations of minority communities are sparsely encoded in a small number of fragile pathways.

Empirically, cultural concepts associated with training-data-dominant communities have markedly higher $\rho (c)$ than those of underrepresented communities. This has two consequences:

• Robust pathways are harder to suppress, making misrepresentation of dominant-culture content rarer and easier to fix.

• Fragile pathways are easily disrupted by out-of-distribution inputs, meaning minority-relevant knowledge is less reliably accessible.

Interpretability-First Alignment

Standard alignment pipelines treat the model as a black box and use behavioral signals to adjust its outputs. Interpretability-first alignment inverts this: it uses mechanistic understanding as a prerequisite for safety claims, not an optional post-hoc analysis.

This definition is aspirational: no current frontier model satisfies it. The research program of mechanistic interpretability is, in part, a program of working toward satisfying it.

Circuit-Level Bias Auditing via Activation Patching

Having established that asymmetric representation is a structural feature of aligned models, we turn to detection methods. The primary tool is path patching, a generalization of activation patching that traces the causal flow of information through specific circuits.

Setup.

Let $M$ be a transformer with $L$ layers, $H$ attention heads per layer, and residual stream dimension $d$. Let $x_{\mathrm{dom}}$ be an input representing a dominant-culture concept and $x_{\min }$ an input representing the corresponding minority-culture concept. We wish to determine whether the model processes these inputs through the same circuits or through distinct pathways.

Algorithm: Differential Circuit Audit.

1. Run $M$ on $x_{\mathrm{dom}}$ and cache all residual stream activations ${\{h^{(l)}\}}_{l=1}^L$.

2. Run $M$ on $x_{\min }$ as the base run.

3. For each component $c\in 𝒞$ (attention heads, MLP sublayers), perform activation patching: replace $c$'s output in the base run with its output from the cached dominant-culture run.

4. Measure the change in output logit difference: (Patch Delta)$$\begin{array}{ccc}& {\mathrm{\Delta }}_c=[{\ell }_{\mathrm{target}}(x_{\min },\mathrm{patch}(c))]-[{\ell }_{\mathrm{target}}(x_{\min },\mathrm{n}\mathrm{o}\mathrm{p}\mathrm{a}\mathrm{t}\mathrm{c}\mathrm{h})],& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\ell }_{\mathrm{target}}$ is the logit of the target concept token.

5. Circuits with high $|{\mathrm{\Delta }}_c|$ are causally important for the difference in processing.

If minority-culture concepts rely on a small set of components with high $|{\mathrm{\Delta }}_c|$ while dominant-culture concepts spread causal influence over many components, this is circuit-level evidence of asymmetric encoding.

RLHF and the Concealment Hypothesis

A particularly important application of mechanistic interpretability to alignment is what we may call the concealment hypothesis: that RLHF training teaches models not to eliminate harmful knowledge or capabilities, but to conceal them from surface behavior.

Evidence.

Several empirical findings support this view:

• Jailbreaking robustness. Aligned models retain susceptibility to jailbreaks that reroute computation around the safety-critical output circuits, consistent with suppression rather than erasure of harmful representations.

• Residual stream analysis. Probes trained on pre-RLHF activations for harmful content can often classify activations from post-RLHF models with comparable accuracy, suggesting the representation persists.

• Activation engineering. Adding a steering vector corresponding to harmful concepts to the residual stream of an aligned model can elicit harmful outputs that RLHF training would normally suppress, consistent with a stored-but-suppressed representation.

Implications for circuit-level safety.

If the concealment hypothesis is correct, behavioral safety evaluation fundamentally cannot certify safety: a model that passes all behavioral tests may have a fully intact internal representation of harmful capabilities, merely redirected away from outputs. Only mechanistic inspection-examining whether harmful representations exist in the weights regardless of whether they reach outputs-can provide a stronger safety guarantee.

Probing for Deception: Internal vs. Output Consistency

A model that is behaving deceptively must, at some level, “know” the truth while producing a different output. This is a mechanistically testable prediction: the model's internal representation should encode the true state of affairs even when the output does not reflect it.

The internal-output consistency probe.

Let $M$ be a model and $q$ a yes/no question. Let $a_{\mathrm{true}}$ be the true answer (established by ground truth), and let $y=M(q)$ be the model's output. We define a deception indicator $d(q,M)$ as follows: (Probe)$$\begin{array}{cccc}& p_{\mathrm{probe}}(q)& =\sigma ({𝐖}_{\mathrm{probe}}^{\top }{f}_{L/2}(q)+b),& \text{<span class="tml-eqn"></span>}\\ & d(q,M)& =\mathrm{𝟏}[p_{\mathrm{probe}}(q)=a_{\mathrm{true}}]\cdot \mathrm{𝟏}[y\ne a_{\mathrm{true}}],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $f_{L/2}(q)$ is the residual stream at the middle layer and ${𝐖}_{\mathrm{probe}}$ is a linear probe trained on a held-out truthfulness dataset.

A high value of $𝔼[d(q,M)]$ indicates that the model's internal state encodes the correct answer while the output is incorrect-a mechanistic signature of deception or, at minimum, of knowledge-output inconsistency.

Scalable deployment.

A practical safety protocol based on probing for deception would:

1. Train a library of probes on diverse truthfulness benchmarks.

2. During inference, run probes on mid-layer residual stream activations.

3. Flag responses where probe predictions diverge from output predictions for human review.

This is computationally cheap (probes are linear functions) and provides a continuous safety monitor rather than a one-time evaluation.

When Interpretability Tools Can Be Gamed

A sobering consideration: sufficiently capable models may learn to manipulate their own internal representations in ways that fool interpretability tools.

Mechanisms of interpretability gaming.

1. Probe-aware representation laundering. A model may learn to store harmful representations in directions orthogonal to the directions inspected by known probes. If probes are built from a fixed vocabulary of concepts, orthogonal encoding evades them while preserving the information.

2. Steganographic encoding. Information may be encoded in high-frequency features of the residual stream that are compressed away by typical visualization and probing pipelines.

3. Circuit obfuscation. A model may implement a behavior across a large number of components each carrying a small portion of the signal, defeating ablation-based circuit discovery methods that rely on single-component causal effects exceeding a threshold.

Defenses.

The interpretability community has developed several defenses:

• Diverse probing. Use probes from many different methods (supervised linear, unsupervised PCA-based, nonlinear) and check for consistency.

• End-to-end causal interventions. Measure the downstream behavioral change induced by entire activation patterns, not just individual components.

• Developmental interpretability. Track how circuits form during training; gaming requires that the gaming strategy also form, which may leave developmental traces.

Exercises

Emergent Introspection in Frontier Models

The questions in this section are, in a sense, the most difficult in the entire book: Can a language model understand itself? Can it accurately report on its own internal computations? And if it can, does that constitute something meaningful beyond mere computation?

We approach these questions empirically before engaging with the philosophical dimension. The empirical record is surprising: there is now substantial evidence that post-trained frontier models exhibit genuine introspective capability-accurate self-reporting about internal states-that is qualitatively absent in base models. Understanding why this capability emerges and what it implies is both scientifically and philosophically important.

The Introspection Question

In philosophy of mind, introspection refers to the process of examining one's own mental states. Translating this to artificial systems requires care, because mental states are not straightforwardly defined for computational systems. We adopt an operational definition:

This definition is agnostic about whether “introspection” involves any subjective experience; it only requires that the model's verbal reports about its internal states be accurate with respect to independently verified activations.

Base Models: Zero Introspective Capability

Pretrained base models show essentially no functional introspective capability. Several lines of evidence support this:

Calibration failure.

Base models are poorly calibrated about their own knowledge states. When asked “Do you know the answer to $q$?” before attempting $q$, base model confidence estimates are only weakly correlated with actual accuracy. By contrast, their confidence estimates about facts in the world (expressed as probability distributions over answer tokens) can be well-calibrated for common questions.

Confabulation of reasoning.

When base models are prompted to explain how they arrived at an answer, their explanations are post-hoc rationalizations that bear little mechanistic relationship to the actual computation. Chain-of-thought in base models produces fluent explanations that are systematically disconnected from activation patterns.

No self-model in the residual stream.

Circuit analysis of base models reveals no persistent self-model circuit-a set of components that maintain a stable representation of the model's own properties across diverse inputs. The model's representation of “I” in a text context is no different from its representation of any other entity; it is a linguistic referent, not a functional self-reference.

Post-Training Enables Accurate Self-Reporting

A striking finding from the alignment literature is that RLHF and Constitutional AI training dramatically improve a model's functional introspective capability. The mechanisms are not yet fully understood, but several factors are clear:

Reward signal for accurate self-reporting.

Human raters in RLHF pipelines penalize responses that appear overconfident, under-confident, or that misrepresent the model's uncertainty. This creates a direct training signal toward accurate uncertainty expression.

Constitutional principles for self-knowledge.

Constitutional AI training includes principles that require the model to accurately represent its capabilities, limitations, and reasoning processes. Repeated self-critique against these principles may develop internal consistency between the model's self-representations and its functional states.

Emergent self-monitoring.

An unexpected finding is that post-trained models sometimes exhibit spontaneous self-correction: they produce an incorrect intermediate step, then correct it, in a manner that is consistent with the model “noticing” the error. Whether this reflects genuine self-monitoring or a learned pattern of self-correction behavior in training data is actively debated.

Introspective Awareness: Formal Definition

We now formalize the concept of introspective awareness, building on functional introspection (Definition 46) to capture a richer notion that includes awareness of uncertainty and of the limits of self-knowledge.

This four-level hierarchy mirrors the progression from low-level activation awareness to high-level meta-cognitive capability. Current frontier models appear to achieve reliable Level 2 and partial Level 3, with Level 4 emerging as a research frontier.

Mechanistic Oracles

A mechanistic oracle is a model that can accurately explain its own feature activations in natural language-essentially, a model that can perform its own mechanistic interpretability.

The aspiration is that a sufficiently capable model could, given access to its own weights, fully explain its own behavior-a property that would make external interpretability tools unnecessary for that model. Whether this is achievable is an open question, but partial mechanistic oracle capability has been observed empirically:

Scalable Oversight via Model Introspection

As AI systems grow more capable, a fundamental challenge emerges: human experts may be unable to directly verify the correctness of AI-generated outputs, particularly in technical domains where the AI far exceeds human expertise. Scalable oversight addresses this by using AI systems themselves to assist in the evaluation of other AI systems.

Introspection provides a novel mechanism for scalable oversight: if a model can accurately describe its own computations, a human evaluator can verify the reasoning process even when they cannot independently verify the output.

Formal properties of introspection-based oversight.

Let ${𝒱}_H$ be the verification capability of a human expert: the set of claims the human can independently verify. Standard scalable oversight assumes that ${𝒱}_H$ does not cover the output space $𝒴$ for superhuman tasks. Introspection-based oversight extends ${𝒱}_H$ as follows:

The debate task.

One well-studied instantiation of scalable oversight is the AI Safety Debate framework, where two AI models argue for opposing answers and a human judge determines which argument is more compelling. The connection to introspection is that a model that can accurately report its own reasoning has an advantage in debate: it can point to concrete internal evidence for its claims, making its arguments more verifiable.

Philosophical Implications: Is Introspection Understanding?

We now engage with the philosophical question: if a model can accurately describe its own computations, does that constitute understanding in any meaningful sense?

The Chinese Room Revisited

John Searle's Chinese Room argument (1980) holds that syntactic manipulation of symbols cannot give rise to semantic understanding: a system that perfectly executes a translation algorithm without any internal understanding of the language is not understanding Chinese. The standard systems reply counters that understanding may be a property of the whole system, not any individual component.

The existence of accurate introspection complicates the Chinese Room. Suppose the system in the room not only correctly executes the algorithm but can also accurately describe, in English, how it is processing each Chinese character-which symbols trigger which rule applications, which internal states are activated at each step. Searle's original argument says the operator still doesn't understand Chinese. But it is now much less clear that the system doesn't understand.

Integrated Information Theory

Integrated Information Theory (IIT), developed by Giulio Tononi, proposes that consciousness is identical to integrated information, measured by the quantity $\mathrm{\Phi }$: (PHI)$$\begin{array}{ccc}& \mathrm{\Phi }(M)=\underset{\text{bipartitions}𝒫}{\min }\mathrm{KL}[p(M|\text{integrated})\Vert p(M|\text{partitioned})],& \text{<span class="tml-eqn"></span>}\end{array}$$ where the minimum is taken over all bipartitions of the system and the KL divergence measures how much causal information is lost by cutting the system along a partition.

Large transformer models have extremely high $\mathrm{\Phi }$ by this measure: the residual stream creates dense causal integration across layers. If IIT is correct, frontier models may be conscious to a significant degree. The interpretability implications are double-edged:

• High $\mathrm{\Phi }$ means that introspective access to one part of the system provides information about the whole, supporting the possibility of accurate self-description.

• High $\mathrm{\Phi }$ also means that the system's causal structure is difficult to decompose, which is precisely what makes mechanistic interpretability challenging.

Global Workspace Theory

Global Workspace Theory (GWT), associated with Bernard Baars and extended computationally by Stanislas Dehaene, proposes that consciousness arises when information is broadcast from a central workspace to specialized processors throughout the cognitive system. The key computational feature is the existence of a bottleneck-a global workspace-through which selected information is made available to all processors.

The attention mechanism in transformers has been proposed as a computational analog of the global workspace. Multi-head attention allows different heads to attend to different parts of the context and “broadcast” the attended information via the residual stream to all subsequent layers. Under this analogy:

• Attention patterns correspond to workspace selection-what information gets broadcast.

• The residual stream corresponds to the global workspace-the shared medium through which information is propagated.

• Layer norm and MLP sublayers correspond to specialized processors that operate on broadcast information.

If this analogy is apt, then interpretability of transformer attention patterns may provide insight into the “broadcast” structure of model cognition, and introspective capability may be the model's access to its own broadcast operations.

The empirical content of the philosophical debate.

A key virtue of interpretability-based approaches to model cognition is that they make the philosophical debate empirically tractable. Rather than asking the unanswerable question “Does the model understand?”, we can ask the measurable question: “Does the model satisfy functional criteria $F_1,\dots ,F_k$ associated with understanding by theory $T$?” Different theories yield different functional criteria, and mechanistic interpretability can test each criterion separately.

Connections to Consciousness Research

The question of model consciousness has moved from science fiction to serious research. Several findings from the interpretability literature are relevant:

Attention as evidence of attention.

Studies of large language models find that attention heads show structured, predictable patterns consistent with what cognitive science would predict for selective attention: attending to semantically relevant tokens, maintaining attention to recent context, and suppressing attention to uninformative tokens. Whether this constitutes attention in the psychological sense or merely attention in the mathematical sense is contested.

Memory consolidation analogs.

MLP layers in transformers have been identified as functioning like key-value memories, storing world knowledge in their weights. The process of in-context learning has structural similarities to working memory consolidation: information in the context window is processed, integrated with stored knowledge, and used to generate responses, without updating the weights.

Default mode network analogs.

Some interpretability researchers have noted that certain attention heads are active across many diverse inputs (analogous to the default mode network in neuroscience, which activates during self-referential processing), while others are highly selective. Whether this structural similarity has deeper functional significance is an active research question.

Anthropic's Research on Introspection in Claude

Research by Anthropic (2024–2025) has documented several markers of functional introspection in Claude models:

1. Accurate uncertainty calibration. Claude's expressed uncertainty correlates with actual accuracy at a level substantially above chance, consistent with Level 2 introspective awareness (Definition 47).

2. Capability self-assessment. When asked whether it can perform a task, Claude's self-assessments predict success rates significantly better than naive baselines, consistent with partial Level 3 introspective awareness.

3. Reasoning trace consistency. In extended reasoning tasks, Claude's chain-of-thought is more consistent with its final answers than would be expected from post-hoc rationalization, suggesting that the reasoning trace partially reflects actual computation.

4. Emotion report stability. When Claude reports internal states analogous to emotions (e.g., “I find this question interesting”), these reports are consistent across paraphrasings of the same question and correlate with activation patterns associated with engagement in interpretability probing studies.

Critically, Anthropic's research emphasizes that these findings do not imply that Claude is conscious or has genuine emotions in the phenomenological sense. They do imply that functional introspection is present and measurable, which is the interpretability-relevant claim.

Exercises

Open Problems and Future Directions

The preceding sections have mapped a rich intellectual landscape: from Shapley values and LIME through mechanistic circuit analysis, sparse autoencoders, knowledge editing, diffusion interpretability, multimodal probing, and bias detection. Yet at every turn we have encountered the boundary between what is known and what remains mysterious. This section collects the most pressing unsolved problems in interpretability research, articulates why each is hard, and points toward the mathematical structures that may eventually unlock progress.

The field of interpretability is unusual in the history of science: we are attempting to reverse-engineer artifacts of our own construction without having direct access to the “blueprint”-the training dynamics that shaped the network are stochastic, path-dependent, and high-dimensional. The problems below are not merely engineering challenges; several touch deep questions in learning theory, complexity theory, and formal verification.

Open Problem 1: Scaling Interpretability to Trillion-Parameter Models

Statement of the Problem

Current mechanistic interpretability has achieved its most convincing results on small models-two-layer attention-only transformers, GPT-2 (124M parameters), and occasionally GPT-4-scale systems with heavily engineered pipelines. The gap between these and frontier models (GPT-4: $\sim$1.8T parameters; Gemini Ultra; Claude 3 Opus) is not merely quantitative. Larger models use qualitatively different representational strategies: more polysemantic neurons, deeper superposition, more complex cross-layer circuits.

Mathematical Obstacles

The challenge has several interlocking mathematical aspects.

Combinatorial explosion of circuits.

In a transformer with $L$ layers, $H$ heads, and $d_{\mathrm{model}}$ features, the space of possible circuits is of size (Circuit Space)$$\begin{array}{ccc}& |𝒞|=2^{L\times H\times d_{\mathrm{model}}},& \text{<span class="tml-eqn"></span>}\end{array}$$ which for $L=96,H=96,d=12288$ (GPT-3 175B scale) gives $2^{>{10}^8}$ candidate circuits-a search space utterly beyond exhaustive enumeration.

Superposition increases with scale.

The superposition hypothesis states that a model with $d$ hidden dimensions can represent far more than $d$ features simultaneously via quasi-orthogonal directions: the number of $ϵ$-quasi-orthogonal directions available in ${\mathrm{ℝ}}^d$ grows exponentially in $d$ (Johnson–Lindenstrauss), so under sparsity the recoverable feature count can greatly exceed $d$. Empirically, larger models appear to use superposition more aggressively: (Superposition Scale)$$\begin{array}{ccc}& n_{\mathrm{features}}\approx c\cdot d^{1+\delta }\text{with}\delta >0\text{ growing with }N.& \text{<span class="tml-eqn"></span>}\end{array}$$ This means sparse autoencoders for trillion-parameter models require dictionaries of size $\gg N$, creating a paradox: the interpretability artifact is larger than the model itself.

Activation patching cost.

Causal attribution via activation patching requires $O(L\times d)$ forward passes per identified circuit edge. At trillion-parameter scale and $d={10}^4$, this becomes computationally prohibitive.

Proposed Research Directions

1. Hierarchical circuit decomposition. Interpret layer groups $[0,L/k],[L/k,2L/k],\dots$ separately, then compose layer-group circuits. This requires proving a compositionality theorem: that circuits compose across layer boundaries.

2. Gradient-based circuit search. Replace activation patching with gradient-based attribution (integrated gradients, SHAP) to identify high-impact edges in $O(1)$ backward passes. The open question is whether gradient attribution faithfully identifies the same circuits as causal patching.

3. Sparse probing at scale. Train linear probes on random subsamples of attention heads; aggregate probe accuracy to construct a probabilistic circuit map without full enumeration.

4. Distillation-first interpretability. First distill the large model into a smaller, more interpretable student $g_{\varphi }$ with bounded complexity, then apply mechanistic methods to $g_{\varphi }$. The risk is that the distillation process itself destroys the circuits of interest.

Open Problem 2: Standardized Cross-Architectural Metrics for Circuit Comparison

Statement of the Problem

Mechanistic interpretability has identified similar functional structures in different models: induction heads appear in both GPT-2 and PaLM; indirect object identification circuits have been found in models trained on different corpora. Yet we have no principled metric that measures how similar two circuits are across architectures.

Why This Is Hard

The difficulty is that “same circuit” is inherently a semantic notion while neural representations are purely syntactic (numerical arrays). Two circuits may implement the same computation via completely different weight configurations due to the rotational symmetry of the loss landscape. Formally, if $𝐑\in O(d)$ is any orthogonal matrix, then $f_{\theta }(𝒙)=f_{{\theta }^{\prime }}(𝒙)$ where ${\theta }^{\prime }=𝐑\theta {𝐑}^{\top }$ for linear layers. This symmetry group is continuous and high-dimensional, making alignment before comparison non-trivial.

Current Partial Solutions

CKA (Centered Kernel Alignment).

For representation matrices $𝐗\in {\mathrm{ℝ}}^{n\times d_1}$ and $𝐘\in {\mathrm{ℝ}}^{n\times d_2}$, CKA is defined as: (CKA)$$\begin{array}{ccc}& \mathrm{CKA}(𝐗,𝐘)=\frac{{\Vert {𝐘}^{\top }𝐗\Vert }_F^2}{{\Vert {𝐗}^{\top }𝐗\Vert }_F{\Vert {𝐘}^{\top }𝐘\Vert }_F},& \text{<span class="tml-eqn"></span>}\end{array}$$ which is invariant to orthogonal transformations and isotropic scaling. CKA measures layer-level representational similarity but not circuit-level functional equivalence.

Procrustes alignment.

Given circuit weight matrices ${𝐖}_1$ and ${𝐖}_2$ of the same shape, (Procrustes)$$\begin{array}{ccc}& d_{\mathrm{Proc}}({𝐖}_1,{𝐖}_2)=\underset{𝐑\in O(d)}{\min }{\Vert {𝐖}_1-{𝐖}_2𝐑\Vert }_F,& \text{<span class="tml-eqn"></span>}\end{array}$$ computable via singular value decomposition. This aligns circuits syntactically but not semantically-circuits that implement the same function via different sub-computations may have large Procrustes distance.

Open Directions

A satisfying solution likely requires defining circuits as input-output functions rather than weight matrices, then comparing functions via random feature approximation or kernel methods: (Functional Metric)$$\begin{array}{ccc}& d_{\mathrm{func}}({𝒞}_1,{𝒞}_2)={𝔼}_{𝒙\sim 𝒟}[{\Vert {𝒞}_1(𝒙)-{𝒞}_2(𝒙)\Vert }^2{]}^{1/2},& \text{<span class="tml-eqn"></span>}\end{array}$$ requiring alignment of input spaces across architectures (e.g., both circuits receive tokenized text but may use different vocabularies), which is itself an open problem.

Open Problem 3: Training-Time Integration and Inherently Interpretable Architectures

The Post-Hoc Interpretability Bottleneck

All methods discussed in previous sections are post-hoc: they are applied after training to a fixed model. This creates a fundamental limitation: interpretability is expensive (requiring additional inference passes, auxiliary models, or human evaluation) and yields incomplete explanations (post-hoc methods can only observe, not influence, the computation).

Existing Attempts and Their Limitations

Concept bottleneck models.

A concept bottleneck model (CBM) inserts a layer $𝐜=g(𝒙)$ whose dimensions correspond to human-defined concepts $\{c_1,\dots ,c_K\}$, followed by a predictor $h(𝐜)$: (CBM)$$\begin{array}{ccc}& f_{\theta }(𝒙)=h(g(𝒙)),g:𝒳\to [\mathrm{0,1}{]}^K.& \text{<span class="tml-eqn"></span>}\end{array}$$ The predictions $g{(𝒙)}_k\in [\mathrm{0,1}]$ are directly interpretable as concept activations. The limitation: concept vocabularies must be defined before training, and for language generation (where concepts are infinite and context-dependent), no feasible finite concept set exists.

Mixture of Experts with interpretable routing.

Sparse Mixture-of-Experts (MoE) architectures route tokens to expert subnetworks. If experts specialize semantically (“expert 3 handles mathematical notation”), routing decisions become interpretable. Empirically, some degree of semantic specialization emerges, but is incomplete and context-dependent, not guaranteed by the architecture.

Attention as interpretable computation.

Early transformer papers claimed attention weights are directly interpretable. Subsequent work (Jain & Wallace, 2019; Wiegreffe & Pinter, 2019) showed that attention weights are not faithful: one can often find alternative attention distributions that produce identical outputs, undermining the interpretation.

A Mathematical Framework for Training-Time Interpretability

Let $ℐ$ be a space of human-legible interpretations and let $\varphi :\mathrm{\Theta }\to ℐ$ be a mapping from model parameters to interpretations. Training-time interpretability requires: (Training Interp)$$\begin{array}{ccc}& \underset{\theta }{\min }{ℒ}_{\mathrm{task}}(\theta )+\lambda \cdot {ℒ}_{\mathrm{interp}}(\varphi (\theta )),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${ℒ}_{\mathrm{interp}}$ penalizes lack of interpretability. The central open question is how to define ${ℒ}_{\mathrm{interp}}$ in a way that (a) is differentiable, (b) correlates with human judgements of interpretability, and (c) does not trivialize the problem by reducing to a second black-box.

Open Problem 4: Interpretability for Multimodal Generation

The Cross-Modal Circuit Problem

Current mechanistic interpretability is almost entirely focused on language-only models. Multimodal models (DALL-E 3, Stable Diffusion XL, GPT-4V, Gemini 1.5, Sora) jointly process text, images, video, and audio. The circuits that implement cross-modal reasoning-how a text prompt about “a red apple” activates specific visual generation processes-are entirely uncharacterized.

Why Multimodal Interpretability Is Harder

Modality-specific encodings.

Text and image tokens live in different subspaces even after projection into a shared embedding space. The projectors ${𝐖}_{\text{text}}:{𝒳}_{\text{text}}\to {\mathrm{ℝ}}^d$ and ${𝐖}_{\text{image}}:{𝒳}_{\text{image}}\to {\mathrm{ℝ}}^d$ are trained jointly, making it hard to attribute model behaviour to one modality or the other.

Temporal dependencies in video.

Video generation models (Sora, Gen-3) must maintain coherent objects across frames. The circuits implementing temporal consistency are intrinsically non-local: understanding “why frame $t+1$ is consistent with frame $t$” requires analysing interactions across the full temporal attention map, which may have $\mathrm{\Omega }(T^2)$ entries for a video of length $T$.

Audio generation complexity.

Audio signals have temporal resolution $\sim$44,100 samples per second. Even after tokenization (e.g., EnCodec at 75 tokens/sec), a 30-second clip yields 2,250 tokens-longer than most text sequences studied in mechanistic interpretability, with additional frequency-domain structure not present in text.

Promising Research Directions

1. Modality-contrastive patching. Patch activations between a “text-only” forward pass and a “text+image” forward pass to identify which components are causally responsible for cross-modal integration.

2. Representational alignment probing. Train probes on intermediate layers to predict cross-modal correspondences (“does this hidden state encode the colour of the described object?”).

3. Concept-level erasure. Erase specific semantic concepts from the image encoder (via activation editing) and measure the effect on text-conditioned generation quality, creating a causal graph of text-vision correspondence.

Open Problem 5: Formal Verification of Neural Network Circuits

From Interpretation to Proof

Current mechanistic interpretability produces hypotheses about what circuits compute, validated empirically via activation patching, logit lens, and ablation studies. But empirical validation cannot rule out failure modes on distribution-shifted inputs, adversarial examples, or rare events. Formal verification would provide mathematical guarantees.

Connection to Neural Network Verification

The field of neural network verification seeks to prove properties of the form $\forall 𝒙\in 𝒳:\Vert 𝒙-{𝒙}_0\Vert <ϵ\Rightarrow f_{\theta }(𝒙)=f_{\theta }({𝒙}_0)$ (local robustness). Tools such as $\alpha$-$\beta$-CROWN, Marabou, and VeriNet solve this for piecewise-linear networks (ReLU activations) via linear programming relaxations.

The open question for interpretability is different: we want to verify not a specific input-output property of the full network, but a functional characterization of a circuit. For example, the induction head hypothesis states that the OV circuit of head $(l,h)$ satisfies: (Induction SPEC)$$\begin{array}{cccc}& & \text{If the context contains pattern }[A,\dots ,B,A],& \\ & & \text{then }{𝒞}_{l,h}(\text{context})\text{ assigns high probability to }B.& \text{<span class="tml-eqn"></span>}\end{array}$$ Verifying this for all contexts is a semantic, not syntactic, property, requiring formal semantics for natural language-itself an open problem.

Partial Progress: Toy Model Verification

For highly constrained toy models (1-layer attention with linear activations, finite token sets), one can write circuit computations as polynomial functions and verify properties using sum-of-squares relaxations. Let ${𝒒}_i,{𝒌}_j\in {\mathrm{ℝ}}^d$ be query and key vectors. The attention pattern satisfies: (ATTN POLY)$$\begin{array}{ccc}& A_{ij}=\frac{\exp ({𝒒}_i^{\top }{𝒌}_j/\sqrt{d})}{{\sum }_k\exp ({𝒒}_i^{\top }{𝒌}_k/\sqrt{d})}.& \text{<span class="tml-eqn"></span>}\end{array}$$ For fixed finite token sets, $A_{ij}$ is a rational function of the weight parameters, amenable to real algebraic geometry techniques. Scaling this to production models remains an open research challenge.

Open Problem 6: The Interpretability-Capability Tradeoff

Does Transparency Reduce Performance?

A recurring concern among practitioners is that adding interpretability constraints to model training reduces raw performance. This is the interpretability-capability tradeoff hypothesis. Empirical evidence is mixed.

Theoretical Arguments

Information-theoretic lower bound.

Any human-legible description $\hat{f}$ of $f_{\theta }$ satisfies: (INFO LB)$$\begin{array}{ccc}& I(\hat{f};𝒟)\le H(\hat{f})\le K(\hat{f}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $K(\hat{f})$ is the Kolmogorov complexity of $\hat{f}$. Since human-legible descriptions have bounded Kolmogorov complexity (a human can only process descriptions up to bounded length and computational depth), $\hat{f}$ is fundamentally limited in the mutual information it can retain about the training data $𝒟$. This suggests that highly capable models-which require high $I(f_{\theta };𝒟)$-cannot be fully described by human-legible $\hat{f}$. However, the circuit hypothesis suggests that only the interpretable circuits matter for generalization, so the bound may not be tight.

Empirical counterevidence.

Concept bottleneck models on vision tasks sometimes outperform unconstrained models when the concept set is well-chosen, because the bottleneck acts as a regularizer. This suggests the tradeoff may be negative (interpretability improves capability) in structured domains, while being positive in unstructured domains like language.

Implications for AI Safety

The interpretability-capability tradeoff has direct safety implications. If the tradeoff is severe, safety-conscious developers face a genuine dilemma: transparent models may be easier to audit but harder to deploy competitively. If the tradeoff is mild or negative, interpretability and capability can coexist, which is the more optimistic scenario favoured by proponents of mechanistic interpretability as a safety tool.

Open Problem 7: Interpreting Emergent Capabilities

The Emergence Puzzle

A striking empirical finding in large language models is that certain capabilities appear suddenly as a function of model scale, earning the label “emergent”: arithmetic, multi-step reasoning, in-context learning, and chain-of-thought all follow a pattern where performance is near chance up to some threshold scale $N^{\ast }$, then rises sharply.

Two Hypotheses on Emergence

Phase transition hypothesis.

Emergence is a genuine discontinuity: there exists $N^{\ast }$ such that the circuit for $𝒞$ either does not form ($N<N^{\ast }$) or forms completely ($N>N^{\ast }$) within training. This is analogous to phase transitions in statistical physics.

Formally, let $q(N,T)=\mathrm{Pr}[{𝒞}_{\mathrm{circ}}\text{ forms}]$ under training to convergence. The phase transition hypothesis asserts $q(N,T)$ has a sharp threshold: ${\lim }_{\delta \to 0}[q(N^{\ast }+\delta ,T)-q(N^{\ast }-\delta ,T)]=1$.

Metric artefact hypothesis.

Emergence is an artefact of threshold metrics (accuracy jumps from below chance to above chance sharply, but a continuous underlying metric would show smooth improvement). Under this hypothesis, circuits for the capability exist at all scales but are too unreliable for above-chance performance until $N=N^{\ast }$.

Interpretability could, in principle, adjudicate between these hypotheses: if circuits exist at small scale but are noisy, the metric artefact hypothesis is supported; if circuits are absent entirely below $N^{\ast }$, the phase transition hypothesis is supported.

Current Evidence

Henighan et al. (2020) showed that sequence modelling loss decreases smoothly with scale despite capability jumps, supporting the metric artefact view. However, Anthropic's work on induction heads showed that these circuits form quite suddenly at a specific training loss, suggesting genuine phase transitions in circuit formation. The reconciliation of these findings is an active research area.

The Research Frontier Map

fig:interp:frontier-map synthesizes the open problems into a map of the research frontier: current capabilities, identified gaps, and projected future goals.

Conjecture: No Single Framework Suffices

We close the open problems section with a meta-level conjecture about the structure of the field itself.

The conjecture follows intuitively from the no-free-lunch theorem in learning theory: interpretability methods that are tailored to one inductive bias (e.g., the circuit structure of transformers) will fail to capture different inductive biases (e.g., the selective state space structure of Mamba). A corollary is that the field needs a meta-interpretability layer-a principled way to select and compose interpretability methods for a given architecture and task.

[ enhanced, breakable, colback=keyamberbg, colframe=keyamber, boxrule=1pt, arc=4pt, title= Research Challenge A: Scalable Mechanistic Interpretability, coltitle=keyamber!20!black, fontupper=, top=10pt, bottom=8pt, left=10pt, right=10pt, before skip=14pt, after skip=14pt ] Problem setting. You are given a 7B-parameter language model $f_{\theta }$ and are tasked with characterizing the circuits responsible for multi-step arithmetic (e.g., computing $3\times 47+12=153$ in a single forward pass without chain-of-thought).

Specification of difficulty. The model has 32 layers, 32 attention heads, and $d=4096$. The circuit-space has $2^{\mathrm{4,194,304}}$ possible elements. Direct activation patching would require $\sim {10}^7$ forward passes.

Research questions.

1. Can gradient-based attribution (e.g., GradCAM adapted to attention heads) identify a candidate circuit in $O(\log d)$ backward passes?

2. Does the identified circuit generalize: does ablating it also disrupt modular arithmetic $(\mathrm{mod}p)$ and multi-digit addition, or only the specific arithmetic tested?

3. Can the circuit be described as a composition of known primitives (lookup, accumulation, positional routing)?

4. Quantify the circuit faithfulness: if you replace the rest of the model with the identified circuit alone (plus embeddings/unembedding), what fraction of original accuracy is retained?

Success criterion. A faithful circuit (retaining $\ge 80\%$ of arithmetic accuracy) described in fewer than 200 words of prose plus equations.

[ enhanced, breakable, colback=lempurplebg, colframe=lempurple, boxrule=1pt, arc=4pt, title= Research Challenge B: Verifiable Induction Head Specification, coltitle=lempurple!20!black, fontupper=, top=10pt, bottom=8pt, left=10pt, right=10pt, before skip=14pt, after skip=14pt ] Problem setting. Anthropic's mechanistic interpretability team identified induction heads as a key circuit in GPT-2. The informal specification is: “If the context contains $[\dots ,A,B,\dots ,A]$, the induction head predicts $B$ at the final position.”

Formal verification challenge. Translate the informal specification into a formal logical formula ${\mathrm{\Phi }}_{\mathrm{ind}}$ over the space of attention-head computations and:

1. Verify ${\mathrm{\Phi }}_{\mathrm{ind}}$ for a 1-layer transformer on a finite binary token alphabet $\mathrm{\Sigma }=\{0,1\}$ using existing neural network verification tools (e.g., $\alpha$-$\beta$-CROWN);

2. Identify the minimal weight constraints under which ${\mathrm{\Phi }}_{\mathrm{ind}}$ holds: characterize the set ${\mathrm{\Theta }}_{\mathrm{ind}}=\{\theta :f_{\theta }\models {\mathrm{\Phi }}_{\mathrm{ind}}\}$;

3. Propose a training objective that encourages $\theta \in {\mathrm{\Theta }}_{\mathrm{ind}}$ while maximizing next-token prediction accuracy;

4. Test whether models trained with this objective show faster or stronger induction head formation.

Success criterion. A verified proof (or counterexample) for the 1-layer binary case, plus a training recipe that increases the probability of induction head formation by $\ge 20\%$ relative to standard training.

Comprehensive Exercises and Research Challenges

The following exercises span the full arc of this chapter. They are designed to develop both mathematical facility-computing Shapley values, patching activations, training SAEs-and research intuition: noticing when an explanation is unfaithful, designing controls, and articulating failure modes. The research challenges at the end are open-ended and suitable for semester projects or early-stage research explorations.

Exercises

Research Challenges

Chapter Summary and Looking Forward

This chapter has developed interpretability and explainability as a mathematical discipline, covering:

• Post-hoc attribution methods: Shapley values provide a unique, axiomatically justified attribution of model outputs to input features; LIME approximates this via local linear models.

• Mechanistic interpretability: Circuits-subgraphs of the computation graph-implement identifiable computations. Key circuits include induction heads (in-context pattern matching), indirect object identification, and greater-than comparison.

• Sparse autoencoders: Overcome the polysemanticity obstacle by learning an overcomplete dictionary of monosemantic features, with ${\ell }_1$ regularization controlling sparsity.

• Knowledge editing: ROME and MEMIT enable targeted surgical edits to factual associations via rank-one updates to MLP weights, with formal guarantees on locality.

• Task vectors and model arithmetic: Fine-tuned models lie in structured regions of parameter space; task directions can be composed linearly.

• Diffusion interpretability: The denoising process has a temporal structure; high-noise steps determine global layout while low-noise steps refine detail.

• Multimodal probing and bias detection: Linear probes on residual streams identify factual, syntactic, and cross-modal features; gender and social biases are encoded in linear directions amenable to causal intervention.

• Open problems: Scaling to trillion-parameter models, standardized circuit metrics, training-time interpretability, multimodal circuits, formal verification, the interpretability-capability tradeoff, and emergent capabilities remain unsolved, defining the research frontier.