This is not a problem of taste. It is a problem of composition. An analysis written against a Dreamer-style recurrent state-space model does not run, without substantial rewriting, on a Genie-style token transformer. A probe designed for the JEPA latent has no natural extension to a model whose latent is a discrete token distribution. Every lab that studies world models is either confined to one family or paying the cost of maintaining three parallel implementations of the same idea.
Our paper argues that this cost is unnecessary, and that the standard route around it — a single unified framework — is worse than the disease. What we propose instead is a small, sharp type system over a set of adapter capabilities. Analyses declare the capabilities they need. Architectures declare the capabilities they expose. Compatibility is a static check. Composition, not unification.
We tried unification first. The design is intuitive: identify the common structure across world-model families, expose it through a single interface, and write analyses against that interface. The problem is that the common structure is thin.
Recurrent state-space models carry two distinct latent objects — a deterministic hidden state and a stochastic latent — that co-evolve through a learned transition. Token transformers carry no such distinction; state is implicit in a growing key-value cache, and the model’s “prediction” is a distribution over next tokens in a discrete space. Joint-embedding predictive architectures carry a continuous latent and, critically, do not decode back to observation space at all. Any unified state representation either erases these differences (making some analyses impossible) or expands to accommodate them (making the abstraction weak enough that it does not constrain implementations).
Every prior attempt at unification in world-model research has landed somewhere in this trade-off. The tooling ends up architecture-specific in practice even when the interface claims otherwise.
1A partial exception is the RL-community convention of treating world models as environments with a step function, which composes across families for some tasks. It is not fine-grained enough for interpretability work: you cannot patch a latent through the step function without knowing what a latent means for the underlying model. The failure mode is characteristic: the framework is presented, adoption starts, and within eighteen months every user has forked to add architecture-specific extensions that were not anticipated by the original design.
We stopped trying to describe world models in terms of what they are and started describing them in terms of what they can do. Each adapter implements a small set of capabilities. The core four are:
encode(observation) → statepredict(state, action?) → next_staterollout(state, k) → trajectoryextract(layer_name) → activationsThese are present in every serious world-model architecture we surveyed. Optional heads absorb the differences: decode(state) → observation for architectures that reconstruct pixels; value_head(state) → scalar for architectures with an explicit value estimator; rollout_for_agent(agent_id) for multi-agent settings.
Analyses declare which capabilities they need. Causal tracing requires predict and extract. Reconstruction attribution requires decode. Theory-of-mind probes require rollout_for_agent. When a user attempts to run an analysis on an architecture whose adapter does not expose the required capabilities, the harness reports the incompatibility at load time, before any compute is spent.
Select an architecture on the left. Its four required capabilities are always present; optional heads (decode, value_head, agent_conditional) vary. Each analysis on the right lights up or is grayed out depending on whether the selected adapter exposes the capabilities it depends on. Compatibility is resolved statically — no wasted compute on a run that would fail.
Each analysis declares a set of required capabilities. Each architecture's adapter declares which capabilities it exposes. Compatibility is a set-containment check. The three architectures shown are stylized: RSSM after Hafner et al. (DreamerV3, 2023); token transformer after Bruce et al. (Genie, 2024); JEPA after LeCun (V-JEPA, 2024).
The figure above walks through the resolution mechanically. RSSMs expose six of seven capabilities (they lack agent-conditional rollout, which most implementations do not have); token transformers expose a different six (they include agent-conditional rollout via agent tokens but often lack a value head); JEPAs expose only the required four (no decoder, no value head, no agent conditioning). Every analysis’s compatibility follows from set containment on the exposed capability set.
The immediate benefit is that library analyses are written once. Causal tracing implemented against the capability interface runs on every architecture that exposes predict and extract, which is all three families under discussion. Mutual-information probes, trajectory-geometry tooling, and dynamical-systems analyses all inherit the same portability. Reconstruction attribution, being decoder-dependent, portably runs on RSSMs and token transformers and portably does not run on JEPAs — reported as such, rather than silently producing meaningless output.
2The silent-nonsense failure mode is the one that keeps me up. A framework that claims universal applicability, produces plots, and lets you draw conclusions from an analysis that was never actually valid for the architecture you ran it on is worse than a framework that refuses to run. The static compatibility check is the mechanism that prevents this class of error.
The framework does not resolve semantic mismatch between architectures. Two adapters can both expose predict, and the state returned can mean subtly different things — the mean of a distribution in one implementation, a sample from that distribution in another. Capability-typing catches structural mismatch; it does not catch this. We flag it in the paper as an open problem.
Nor does it eliminate the cost of writing adapters. A new architecture requires an adapter implementation; we have made this as small as we can, but the cost is not zero. What the framework does is amortize that one-time adapter cost across all subsequent analyses, which is a substantially better position than the current status quo of writing one implementation per (architecture, analysis) pair.
Two things converged that made this the correct time to propose a shared substrate.
The first is that the world-model field is fragmenting on a shorter timescale than the interpretability field can keep up with. A new architectural family has emerged in each of the last three years. If interpretability tooling has to be rewritten from scratch for each, the tools will lag the models by definition, and the lag will grow.
The second is that the safety case for deploying agents that plan against learned world models — which is happening, at scale — depends on being able to inspect those world models. Interpretability that only works for one architecture family is not safety infrastructure. It is a research artifact of that family. For interpretability to serve a safety role, it has to compose across the architectures a deployed system might actually use. The type system is our attempt to make that composition possible.
WorldModelLens, released with the paper, contains adapter implementations for the three families discussed above and reference implementations of the five analyses shown in the figure. The intent is that a research group with a new world-model architecture can write an adapter — roughly fifty lines of code, most of which is method plumbing — and get every analysis in the library for free.
The current work I am doing along these lines extends the framework in two directions: a set of capability-typed analyses over the geometry of learned latent spaces, and a set of energy-based priors that compose with the rollout capability to shape trajectory sampling. Both are drafts. Both are the reason the framework needed to exist.
One prompt walks in. One probability distribution walks out. In between, eight stages on a belt. Pick a prompt, press step forward, watch each stage's contribution to the final answer.
Step through the 9 stages: tokenize → embed → block 0 → … → block 5 → unembed. The token strip shows the logit-lens prediction at every position; the final-stage panel shows the actual output distribution.
Try the A B C D E F G A B C preset to observe in-context induction: middle layers detect the repeated bigram and predict the continuation D from the first occurrence.
Reference prompt: "The capital of France is". Model: distilGPT2 ($L=6$, $d_\text{model}=768$, $n_\text{heads}=12$, $V=50{,}257$).
"The capital of France is"
→ token IDs: [464, 3139, 286, 4881, 318]
→ pieces: ["The", " capital", " of", " France", " is"]
GPT-2 BPE. 5 tokens. Note “ France” is a single token (common proper noun); “ capital” includes its leading space.
input_ids: [5] → embeddings: [5, 768]
Token embedding $W_E[\text{ids}] \in \mathbb{R}^{5 \times 768}$ plus learned positional embedding $W_P[0:5] \in \mathbb{R}^{5 \times 768}$.
Logit lens at this stage approximately returns the input tokens themselves (no context mixing has occurred). Final-position prediction is meaningless; the model has only seen the token “ is” in isolation.
residual: [5, 768]
attn output: [5, 768] (12 heads × 64 dim, projected back via W_O)
mlp output: [5, 768] (3072 neurons → 768 via W_out)
new residual: [5, 768] (sum of three)
Block 0 is dominated by previous-token heads (attending one position back) and surface-feature MLP neurons (capitalization, punctuation, common morphemes). The logit lens still returns near-token-identity at most positions.
residual: [5, 768] → ... → [5, 768]
Semantic consolidation. Geographic relations form: “ France” gathers context from “ capital” and “ of”. By block 3, the final-position logit-lens prediction includes country and city names in the top-5. By block 4, “ Paris” has typically reached top-1, but with low confidence (~30–50%).
This is also where induction heads activate on patterned prompts. For A B C D E F G A B C, blocks 2–4 detect the prefix repetition and route the continuation forward.
residual: [5, 768] → [5, 768]
Sharpening. Late-layer name-mover-style heads pull “ Paris” embedding into the final position; the answer’s probability mass concentrates. Competing candidates (“ the”, “ France”) get suppressed by negative-name-mover-style components.
final_residual[-1]: [768]
W_U: [768, 50257]
logits: [50257]
softmax(logits): [50257] probability distribution
Apply final layer norm, then project the last position’s residual through $W_U$ to produce a logit for every token in the vocabulary. Softmax gives the probability distribution. For distilGPT2 on this prompt, “ Paris” is top-1 with ~80% probability; the remaining mass is distributed over “ France”, “ Europe”, “ Britain”, “ Germany”, and a long tail.
The model commits one token. To generate more, append the chosen token and run the forward pass again.
Same logits, four ways to pick a token. Switch sampling strategy, change the parameter, and click generate to roll 50 samples. The histogram on the right shows where the strategy actually places its bets.
Greedy always picks top-1, deterministic but boring. Temperature rescales then samples from the full distribution. Top-k keeps only the k highest tokens, then renormalizes. Top-p (nucleus) keeps the smallest set whose cumulative probability exceeds p. Frontier chatbots usually combine top-p with temperature 0.7, 1.0.
| Stage | Tensor | Shape | Memory (fp16) |
|---|---|---|---|
| Token IDs | input | $[5]$ | 20 B |
| Embedding | $X_0$ | $[5, 768]$ | 7.5 KB |
| Per-head Q/K/V | per head | $[5, 64]$ | 640 B each |
| Attention pattern | per head | $[5, 5]$ | 50 B per head |
| MLP hidden | per block | $[5, 3072]$ | 30 KB |
| Final residual | $X_L$ | $[5, 768]$ | 7.5 KB |
| Logits | output | $[50257]$ | 100 KB |
distilGPT2 contains:
Total: ~82M parameters (the embedding/unembedding tables add another ~38M).
For comparison:
| Model | $L$ | Heads/block | Neurons/block | Total components |
|---|---|---|---|---|
| distilGPT2 | 6 | 12 | 3,072 | 78 |
| GPT-2 small | 12 | 12 | 3,072 | 156 |
| GPT-2 XL | 48 | 25 | 6,400 | 1,248 |
| Llama 3 8B | 32 | 32 | 14,336 | 1,056 |
| Claude / GPT-4 class | ~100+ | ~100+ | ~50,000+ | tens of thousands |
The MI program: characterize each of these components in terms of what it reads from and writes to the residual stream. This is fully tractable for distilGPT2 and GPT-2 small (the IOI circuit is one example). It is partially tractable for 8B-class open models with sparse autoencoders. It is an open research problem at frontier scale.
1. Most computation happens mid-stack. Embedding produces near-token-identity; the final block sharpens but rarely overturns; the middle blocks (1–4 in distilGPT2; 4–10 in GPT-2 small) do the semantic work. The logit-lens trajectory shows confidence rising mid-stack and saturating at the top.
2. Only the final position predicts the next token. All earlier positions accumulate context that attention will later retrieve into the final position. Logit-lens predictions at non-final positions are largely incidental: the model is not optimizing them.
3. Computation is parallel and distributed, not sequential. The model does not execute “identify France → look up capitals → output Paris” as discrete steps. All blocks compute simultaneously on their inputs; the result is an additive sum on the residual stream. There is no step 3. There are 78 components contributing in parallel.
Pick a starting prompt. Click next token and the model runs one full forward pass: 12 layers, 144 heads, 36,864 MLP neurons, the whole thing, to produce a single token. Click again to run it all over again. That's autoregressive generation.
What you're watching: for each click, the demo samples one token from the top-5 distribution at the current temperature. Lower T sharpens to greedy; higher T flattens to chaos. The real model would do this exactly the same way, but with a 50,257-token vocabulary and a real forward pass instead of these handcrafted continuations. Same shape, different scale.
This is why MI does not ask “what happened at step 3?” but instead “what did head 7.4 contribute?”. The latter has a precise numerical answer (DLA gives a scalar); the former does not.
| Concept | Post |
|---|---|
| The black box problem and why MI exists | post 1 |
| What mechanistic interpretability is | post 2 |
| Who’s doing this work, and current goals | post 3 |
| Neurons, weights, and forward propagation | posts 4–5 |
| Layers, depth, and training | posts 6–7 |
| Transformer architecture overview | post 8 (this series) |
| Tokens and BPE | post 9 |
| Residual stream, logit lens, DLA | post 10 |
| Attention, QK / OV, IOI circuit | post 11 |
| MLPs, key-value memory, superposition | post 12 |
| Full forward pass | post 13 |
Sufficient to read most current MI papers without ambiguity.
The next chapter is features and circuits: applying this architectural foundation to find concrete computational structures inside trained models. Topics:
This post defines the MLP, derives the key-value memory interpretation (Geva et al., 2021) that underlies most modern MLP interpretability, covers neuron archetypes and superposition, and connects the framework to factual editing (ROME/MEMIT) and sparse autoencoders.
An MLP block has two linear layers with a non-linearity between them. The intermediate dimension is huge, typically 4× the hidden size. So in distilGPT2, each block has 3,072 MLP neurons. Each one is a tiny pattern detector. Click a neuron below to see what it fires on and what it writes back to the stream.
Geva et al. (2021) showed that you can read the up-projection rows as keys (patterns the neuron looks for in the residual stream) and the down-projection columns as values (vectors the neuron writes into the stream when it fires). That decomposition turns each neuron into a key-value pair: "if you see X, write Y." It's the cleanest way to think about MLPs that exists. Bricken et al. (2023) extended this with sparse autoencoders to handle the messy fact that real neurons mix many features in superposition, but the key-value framing is still the foundation.
Cycle through neurons. Some are monosemantic (Python keywords, capital letters, France-related contexts). Some are polysemantic, firing on multiple unrelated concepts. The polysemantic case is explained by superposition (below).
For input $x \in \mathbb{R}^{d_\text{model}}$ at one position, a transformer MLP computes:
\[\text{MLP}(x) = W_\text{out}\, \sigma(W_\text{in}\, x + b_\text{in}) + b_\text{out}\]where:
Sizes:
| Model | $d_\text{model}$ | $d_\text{ffn}$ | Neurons per block |
|---|---|---|---|
| GPT-2 small | 768 | 3,072 | 3,072 |
| GPT-2 XL | 1,600 | 6,400 | 6,400 |
| Llama 3 8B | 4,096 | 14,336 | 14,336 |
| GPT-3 175B | 12,288 | 49,152 | 49,152 |
Three properties:
Decompose $W_\text{in}$ row-wise and $W_\text{out}$ column-wise:
Then:
\[\text{MLP}(x) = \sum_{n=1}^{d_\text{ffn}} \sigma(k_n^\top x + b_n)\, v_n\]The MLP is a sum of $d_\text{ffn}$ scaled value-vectors, where each scaling coefficient is a non-linearly gated dot product of $x$ with the corresponding key.
Equivalently:
This is a soft, sparse key-value lookup over a learned database of $d_\text{ffn}$ entries per block. (Geva et al., 2021)
Empirically, MLP neurons fall into recurring categories:
Fire on lexical patterns: capital letters, punctuation, specific morphemes, code-syntax tokens. Their values write tags downstream blocks consume.
Fire after grammatical patterns: possessives, definite articles, sentence-initial positions. Values bias the next-token distribution toward syntactically valid continuations.
Encode specific facts. Meng et al. (2022, ROME) demonstrated that “the Eiffel Tower is in Paris” can be located to a small set of neurons in mid layers and surgically edited (so the model claims the Eiffel Tower is in Rome) by modifying $W_\text{out}$ columns at those positions.
Fire on higher-level patterns: sentiment, sarcasm, discourse markers. Harder to characterize from top-activating examples alone.
A non-trivial fraction of neurons have no clean concept-level description. Often these are polysemantic.
Most real neurons are polysemantic: top-activating contexts span multiple unrelated concepts.
Elhage et al. (2022, “Toy Models of Superposition”) explain why. When features are sparse (most off most of the time), a $d$-dimensional space can represent ~$d / \log d$ features by overlapping them at non-orthogonal angles. The non-linearity in the MLP allows partial recovery: only one feature in a superposed pair is typically active in any given input, so interference is bounded.
Consequences:
Bricken et al. (2023, “Towards Monosemanticity”) and Templeton et al. (2024, “Scaling Monosemanticity”) trained SAEs on Claude 3 Sonnet and recovered millions of monosemantic features ranging from “the Golden Gate Bridge” to “code with security vulnerabilities.”
A neural network has fewer dimensions than concepts. To survive, it packs many features into the same space at near-orthogonal angles. The trick: only a few features are active at once, so they barely interfere. Toggle features on and off and watch the residual vector form. Then read it back out.
What you're seeing: 8 feature directions packed into 2 dimensions. Each active feature contributes its direction (faded amber arrow); the residual is their sum (dark amber). Below: dot products of the residual with each feature direction, the model's "readout." When few features are on the readout is clean. Turn on too many and they start to interfere. This is why neurons can be polysemantic and why sparse autoencoders work.
Because each neuron’s contribution to the residual stream is $\sigma(k_n^\top x) v_n$, its contribution to the final logit of token $w$ is:
\[\Delta\text{logit}_n(w) = \sigma(k_n^\top x)\, v_n^\top W_U[:, w]\]| Sort neurons by $ | \Delta\text{logit}_n(w) | $ to identify which neurons drove the prediction. This is MLP-level DLA. |
# in TransformerLens
mlp_act = cache["post", layer, "mlp"] # [seq, d_ffn]
W_out = model.W_out[layer] # [d_ffn, d_model]
W_U = model.W_U[:, answer_id] # [d_model]
neuron_dla = mlp_act[-1] * (W_out @ W_U) # [d_ffn]
top_neurons = neuron_dla.argsort(descending=True)[:10]
Three lines of evidence support the claim that factual knowledge lives in MLPs:
A clean operational summary: attention moves information; MLPs add new information. Both contribute additively to the residual stream. Their roles are complementary.
Pick a fact. Pick the new answer you want the model to insist on. Drag the edit slider from 0 to 1 to apply a rank-one update to one MLP layer's weights, simulating what the ROME paper does. Watch the prediction shift, and watch a related question to see if your edit generalizes correctly.
What ROME actually does: the model encodes facts as (key, value) pairs in mid-layer MLPs. ROME locates the layer where the subject's representation peaks, then computes a rank-one update to W_out that swaps the value while preserving the key. The edit propagates to paraphrases of the same fact (the consistency check) without destroying unrelated knowledge. Numbers shown are illustrative; real ROME edits typically achieve 95%+ post-edit probability.
| Model | Non-linearity | Form |
|---|---|---|
| GPT-2 / GPT-3 | GeLU | $x \cdot \Phi(x)$ |
| Original Transformer | ReLU | $\max(0, x)$ |
| PaLM, Llama, Mistral | SwiGLU | $\text{Swish}(W_g x) \odot (W_\text{in} x)$ |
SwiGLU adds a gating branch:
\[\text{MLP}_\text{SwiGLU}(x) = W_\text{out}\, (\text{Swish}(W_g x) \odot W_\text{in} x)\]This requires three matrices instead of two, but the key-value interpretation extends: each neuron’s “key” is now a (gate, input) pair, and the value is still the corresponding $W_\text{out}$ column. Most MLP interpretability tooling generalizes with minor modification.
| Component | Role | Reads | Writes |
|---|---|---|---|
| Embedding | Token → vector | token IDs | residual stream |
| Attention | Cross-position routing | residual stream (all positions) | residual stream (current position) |
| MLP | Stored knowledge / transforms | residual stream (current position) | residual stream (current position) |
| Unembedding | Vector → logits | residual stream (last position) | output distribution |
All four components communicate exclusively via the residual stream. Every interpretability tool in this series operates on that interface.
The next post runs a full forward pass through GPT-2 small, end-to-end, with concrete numbers at every stage.
This post defines the operation, derives the QK / OV decomposition that underlies head-level interpretability, and walks through four head archetypes plus the indirect-object identification circuit.
Every attention head asks a single question: "for each position, which other positions should I look at?" Different heads learn wildly different rules. Pick a head from the gallery and watch what it actually pays attention to.
Each transformer block has multiple parallel attention heads (12 in distilGPT2, 96+ in larger models). Each head computes its own query/key/value matrices and produces its own attention pattern. The block sums all heads' contributions back into the residual stream. So the model has hundreds or thousands of tiny independent attention circuits, and a huge amount of MI work is just figuring out what each one does. The patterns shown here are textbook archetypes, observed in real GPT-2 small. See Elhage et al. 2021 and Olsson et al. 2022.
3 prompts × 6 routing-pattern reproductions × matrix and flow views. The patterns shown (“previous-token”, “BOS sink”, “induction”, “duplicate-token”, “name-mover”, “self”) are reproductions of the canonical patterns observed in real GPT-2 small heads.
Each attention head has three weight matrices:
\[W_Q, W_K, W_V \in \mathbb{R}^{d_\text{model} \times d_\text{head}}\]Typical sizes: GPT-2 small has $d_\text{model} = 768$, $n_\text{heads} = 12$, $d_\text{head} = 64$ per head ($d_\text{head} = d_\text{model} / n_\text{heads}$).
For input $X \in \mathbb{R}^{T \times d_\text{model}}$:
\[Q = XW_Q,\quad K = XW_K,\quad V = XW_V \quad \in \mathbb{R}^{T \times d_\text{head}}\]Attention scores (causal, scaled):
\[A = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_\text{head}}} + M\right)\]where $M_{ij} = -\infty$ for $j > i$ (causal mask), 0 otherwise. $A \in \mathbb{R}^{T \times T}$ is the attention pattern.
Output:
\[Z = AV \quad \in \mathbb{R}^{T \times d_\text{head}}\]Multi-head: concatenate $n_\text{heads}$ outputs $[Z^{(1)}, \ldots, Z^{(h)}]$ and project through $W_O \in \mathbb{R}^{(n_\text{heads} \cdot d_\text{head}) \times d_\text{model}}$ to write back to the residual stream.
The scaling factor $\sqrt{d_\text{head}}$ keeps the dot products in a numerically stable range (Vaswani et al., 2017, §3.2.1).
Each head can be analyzed as two independent linear maps composed by softmax + sum.
The attention score is bilinear in the inputs:
\[Q_i K_j^\top = (X_i W_Q)(X_j W_K)^\top = X_i (W_Q W_K^\top) X_j^\top\]The product $W_{QK} := W_Q W_K^\top \in \mathbb{R}^{d_\text{model} \times d_\text{model}}$ is the QK matrix. It maps pairs (query position content, key position content) → score. Eigendecomposing or projecting $W_{QK}$ onto interpretable subspaces reveals the routing rule.
The output written back to the residual stream from source position $j$, weighted by $A_{ij}$, is:
\[\Delta_i = \sum_j A_{ij}\, X_j W_V W_O^{(h)}\]The product $W_{OV} := W_V W_O^{(h)} \in \mathbb{R}^{d_\text{model} \times d_\text{model}}$ is the OV matrix. It maps (source residual content) → (write contribution). Reading the eigenstructure of $W_{OV}$ describes what kind of information the head copies.
The two are independent. Routing (QK) and payload (OV) are trained jointly but are mathematically separate objects. Most interpretability claims about a head reduce to characterizing $W_{QK}$ and $W_{OV}$ separately. (Elhage et al., 2021)
Every attention head does two separate things. QK decides where to look; OV decides what to copy. Pick a head archetype, then toggle between the two views to see the same head from both angles.
In-context bigram completion: if the prefix contains …A B… and the current token is a later A, the head attends to the position right after the prior A and copies that token (B) forward.
A, query matches keys at positions whose previous token equals A. This requires the previous-token information that previous-token heads write.Click and drag across the lower-triangular grid to paint attention weights. The demo auto-normalizes each row so weights sum to 1, then classifies your pattern against the canonical archetypes. Try to draw a previous-token head. Then a BOS sink. Then an induction head.
Try this: click only the cells right below the diagonal. That's a previous-token pattern. Or click only the leftmost column (position 0): BOS sink. The classifier compares your row distributions against archetype templates using cosine similarity, exactly the kind of analysis MI researchers do programmatically over thousands of real heads.
Many heads route most of their attention to position 0 (BOS) on tokens where the head has nothing useful to do. Softmax forces the weights to sum to 1, so the head must attend somewhere; the BOS slot acts as a “rest” position with low informational impact.
Wang et al. (2022) reverse-engineered the algorithm GPT-2 small uses to predict Mary for the prompt:
“When John and Mary went to the store, John gave a drink to ___”
The circuit involves ~26 attention heads across layers 0–11, organized into named functional groups. Sketch:
| Stage | Heads (layer.head) | Role |
|---|---|---|
| Duplicate Token | 0.1, 0.10, 3.0 | Detect repeated names. Output: “this name appears twice.” |
| Previous Token | 2.2, 4.11 | Move name info to positions before/after each name. |
| Induction | 5.5, 5.8, 5.9, 6.9 | Pattern-match across the sentence using duplicate-token features. |
| S-Inhibition | 7.3, 7.9, 8.6, 8.10 | Write a “John is the subject, suppress John” signal at the final position. |
| Name Mover | 9.6, 9.9, 10.0 | Attend from final position to names. Suppression from S-Inhibition makes them attend to Mary, not John. Output: Mary’s logit goes up. |
| Negative Name Mover | 10.7, 11.10 | Slightly suppress the answer (regularization-like). |
The paper validates each role via path-patching ablations: zeroing out a single head’s contribution to the relevant downstream component degrades the answer. Reproducible in TransformerLens with ~50 lines of code.
This was the first complete circuit reverse-engineered in a language model.
Decoder-only models enforce $A_{ij} = 0$ for $j > i$ via the mask $M$. Two reasons:
The mask is added to attention scores before softmax, with $-\infty$ in the masked positions, so masked weights become exactly zero.
In matrix form, $A$ is lower-triangular. Visible in every demo above.
Why $h$ heads instead of one bigger head? Each head learns a different $(W_{QK}, W_{OV})$, allowing the layer to perform multiple routings simultaneously: head 1 might do “previous token” while head 2 does “subject of the sentence” while head 3 acts as a BOS sink. With one head these would have to share the same projection.
attn_layer(X):
heads = []
for h in range(n_heads):
Q = X @ W_Q[h]; K = X @ W_K[h]; V = X @ W_V[h]
A = softmax(Q @ K.T / sqrt(d_head) + causal_mask)
heads.append(A @ V)
return concat(heads, dim=-1) @ W_O
In code, this is one batched matmul with a head dimension. Conceptually, $h$ independent attention operations.
The softmax is what makes attention selective. A linear weighting would average; softmax allows sharp, peaky distributions where one position gets most of the weight.
The constraint is that $\sum_j A_{ij} = 1$. The head must attend somewhere. When no key matches the query, it defaults to whatever residual key is closest, often the BOS token, which becomes the default sink.
Some recent architectures replace softmax with linear attention (Linformer, RetNet, Mamba) or kernelized variants to avoid this and to get sub-quadratic time. Softmax remains the standard for frontier LLMs.
A useful frame: attention performs content-addressable retrieval from the context. The query is an address; the keys are stored entries; the softmax picks the closest match. The OV circuit decides what to retrieve from the matched entry.
This is why attention scales gracefully across context length and why it pairs naturally with MLPs: attention retrieves relevant context-specific information; MLPs apply training-time-stored transformations to it. The next post is on MLPs.
This post defines the stream formally, explains why the additive structure is load-bearing, and introduces the two interpretability primitives derived from it: the logit lens and direct logit attribution.
The residual stream is the conveyor belt every block reads from and writes to. Each cell below shows the model's best guess at the next token if we stopped the belt at that layer and position. This is the logit lens. Pick a prompt, watch the predictions sharpen as you move down the belt.
The logit lens: take the residual stream at any intermediate layer, project it through the final unembedding matrix, see what token would win if the model stopped right there. Early layers produce noise. Middle layers narrow down to the right semantic neighbourhood. Late layers commit to a single answer. The fact this works at all is the deep observation: the residual stream stays "in the same language" as the final output throughout the network. Technique due to Nostalgebraist, 2020. Activation patterns shown here are realistic distilGPT2 traces, precomputed and visualized so the page loads instantly.
Each cell shows the model’s top-1 prediction at layer $\ell$, position $t$. Saturation = confidence. Click a cell for the full top-5 distribution and the residual norm.
For a transformer with $L$ blocks operating on $T$ tokens, the residual stream is the sequence of states ${X_0, X_1, \ldots, X_L}$, each $X_\ell \in \mathbb{R}^{T \times d_\text{model}}$.
Initial state: $X_0 = E + P$ where $E$ is the token embedding and $P$ is the positional encoding (or $X_0 = E$ for models with rotary/RoPE applied inside attention).
Block update:
\[X_{\ell+1} = X_\ell + \text{Attn}_\ell(\text{LN}(X_\ell)) + \text{MLP}_\ell(\text{LN}(X_\ell + \text{Attn}_\ell(\text{LN}(X_\ell))))\]The layer norms ($\text{LN}$) appear inside each sublayer in the pre-norm configuration used by GPT-2, Llama, and most modern models. Schematically:
\[X_{\ell+1} = X_\ell + \Delta_\ell^\text{attn} + \Delta_\ell^\text{mlp}\]Final read: logits = $\text{LN}(X_L) \cdot W_U$, where $W_U \in \mathbb{R}^{d_\text{model} \times V}$ is the unembedding matrix.
The architectural choice that everything depends on is the +: each $\Delta$ is added, never substituted.
Compare the residual update with a non-residual update $X_{\ell+1} = f_\ell(X_\ell)$. Two structural problems with the latter:
Vanishing information. Any signal computed at layer $\ell$ must be re-encoded by $f_{\ell+1}, f_{\ell+2}, \ldots$ to survive. After 30 layers of arbitrary nonlinear transformations, layer-1 signals are effectively destroyed.
Vanishing gradients. Backprop multiplies gradients through every $f_\ell$. With layer-norm and sigmoid/tanh nonlinearities the gradient norm shrinks geometrically. Pre-2015 networks rarely trained stably past 20 layers.
Residual connections (He et al., 2015) solve both: $X_{\ell+1} = X_\ell + f_\ell(X_\ell)$ has an identity path from layer $\ell$ to $\ell+1$. Information and gradients flow through the + without distortion. This is what made GPT-2’s 48-layer and GPT-3’s 96-layer networks trainable.
For interpretability, the additive structure has a stronger consequence: the final state is literally a sum.
\[X_L = X_0 + \sum_{\ell=0}^{L-1} \Delta_\ell^\text{attn} + \sum_{\ell=0}^{L-1} \Delta_\ell^\text{mlp}\]Every component’s contribution is a linear term. This makes the residual stream linearly decomposable.
Final-layer logits are computed as $\text{LN}(X_L) \cdot W_U$. Because every $X_\ell$ lives in $\mathbb{R}^{T \times d_\text{model}}$, the same projection is well-defined at every layer:
\[\text{logits}_\ell := \text{LN}(X_\ell) \cdot W_U\]Nostalgebraist (2020) called this the logit lens. It returns the model’s “current best guess” at each intermediate layer, treating the partial residual stream as if it were the final state.
Empirically (visible in the demo above for “Paris is the capital of”):
France) reaches top-1 with high probability.The lens is not exact, intermediate $X_\ell$ has different statistics than $X_L$, but it is informative and free. Refinements include the tuned lens (Belrose et al., 2023), which learns a per-layer affine correction.
Because $X_L$ is a sum, the logit for any vocabulary token $w$ is also a sum:
\[\text{logit}(w) = (X_0 \cdot W_U[:, w]) + \sum_{\ell=0}^{L-1} (\Delta_\ell^\text{attn} \cdot W_U[:, w]) + \sum_{\ell=0}^{L-1} (\Delta_\ell^\text{mlp} \cdot W_U[:, w])\]Each term is a scalar: how much that component pushed the prediction toward $w$. This is direct logit attribution.
Practical use:
# in TransformerLens
import transformer_lens as tl
model = tl.HookedTransformer.from_pretrained("gpt2")
tokens = model.to_tokens("Paris is the capital of")
logits, cache = model.run_with_cache(tokens)
# decompose final residual stream into per-component contributions
per_layer = cache.decompose_resid(layer=-1, return_labels=True)
# project each onto W_U for the answer token
answer_id = model.to_single_token(" France")
W_U = model.W_U[:, answer_id]
contributions = per_layer[0] @ W_U # one scalar per component
The largest entries in contributions identify the layers/heads/MLPs that drove the answer. DLA is the starting point for circuit analysis: keep zooming in (head → query/key/value → input neurons) until you have a mechanism.
What to notice: the cumulative line is what the model's prediction actually depends on. Tiny early-layer bumps don't matter; the big middle/late layers (where name-movers, induction heads, and MLP fact-recall live) are doing the real work. Negative bars in late layers are negative name-movers — components that suppress the answer to keep calibration honest. This is the entire game of circuit attribution.
The stream has $d_\text{model}$ dimensions but generally encodes far more features than that. Components write to and read from subspaces of the stream, generally not axis-aligned.
Elhage et al. (2022, “Superposition”) characterize this: when features are sparse (most are off most of the time), a $d$-dim space can represent ~$d / \log d$ features by overlapping them. The cost is interference: reading one feature picks up small projections from others.
Consequences:
Provisional model: think of the stream as a high-dimensional space where many features overlap, recoverable by linear probes or SAEs but not by reading individual coordinates.
In the demo, watch:
The first token’s residual stream typically accumulates “housekeeping” state. Attention heads with no relevant key in a given query often place mass on the BOS token as a default, the attention sink (Xiao et al., 2023). Templeton et al. (2024) found that BOS-position SAE features are systematically distinct from content-position features.
Treat the BOS column as anomalous when interpreting diagnostics.
Every transformer mechanism can be expressed as “component $C$ reads from subspace $R$ of the residual stream and writes to subspace $W$ of the residual stream.”
Examples:
What the heatmap reveals: hot cells form a small region — usually mid-to-late layers at the last subject token for factual recall, or at name positions for IOI. Click those cells: the answer flips. Click anywhere else: nothing happens. That's how Meng et al. found ROME and how the IOI circuit was traced.
The next two posts cover attention and MLPs as readers/writers in detail.
This post defines tokens precisely, explains how the vocabulary is constructed (Byte-Pair Encoding), and walks through five concrete failure modes traceable directly to tokenization.
Type anything in the box. Watch it get chopped into subword pieces. Notice how a space is part of a token. Notice how numbers split weirdly. Notice how emoji explode into three pieces. This is real BPE, running in your browser.
, , , This demo uses real Byte-Pair Encoding (BPE), the same algorithm GPT-2 and (close relatives of) GPT-3/4 use. BPE starts from single bytes and learns which byte-pairs to merge based on frequency in the training corpus. The result is a vocabulary of ~50,000 subword pieces: common whole words ("the", "and"), common word-chunks ("ing", "tion"), and everything else falls back to individual bytes. That's why "unhappily" becomes three pieces but "cat" stays whole, "cat" was common enough in the training text to earn its own merge. We run this via the gpt-tokenizer library, fetched lazily. Fallback to an illustrative split if the CDN is unreachable.
Switch between presets. Note how 123456789 does not split into individual digits; how ❤️ consumes three tokens; how Japanese text consumes more tokens per character than English.
A tokenizer is a deterministic function $\text{tok}: \text{string} \to [V]^*$ that maps text to a sequence of integer IDs. The vocabulary $V$ is a fixed lookup table containing strings called token pieces.
Modern LLMs use byte-level BPE (Byte-Pair Encoding):
1. Initialize V := {byte 0, byte 1, ..., byte 255} # 256 entries
2. While |V| < target_size:
a. Count adjacent token pairs in the training corpus.
b. Find the most frequent pair (a, b).
c. Add ab to V as a new token.
d. Replace all occurrences of (a, b) in the corpus with ab.
3. Return V and the merge order.
Encoding new text greedily applies merges in the order they were learned. The output is the list of resulting token IDs.
Notice: the first merges always pick up boring junk like th, he, in. The interesting tokens (the, ing, whole words) only appear after dozens of merges. Real GPT-style tokenizers do this 50,000 times on the entire internet — that's why " the" is one token but " frosting" might be three.
Example (GPT-2 BPE):
"unhappily" → ["un", "happ", "ily"] → [403, 7829, 6148]
" the cat" → [" the", " cat"] → [262, 3797]
"strawberry" → ["str", "aw", "berry"] → [2536, 707, 27078]
Common English words tend to be a single token. Rare or compound words split. Average compression: ~4 characters per token for English, less for other scripts.
A character-level model is conceptually simpler (256 ASCII tokens) but breaks at scale.
1. Sequence length. Self-attention is $O(T^2)$ in the sequence length $T$. A 2,500-character article is 625 BPE tokens vs 2,500 characters: a 16× compute increase per forward pass.
2. Capacity. Subword tokens already encode high-level meaning. The vector for " doctor" carries semantic content that a character-level model would have to assemble from d, o, c, t, o, r, burning model capacity on spelling.
BPE is the standard compromise: dense for frequent strings, sparse for rare ones.
In GPT-2’s tokenizer, 123456789 splits into ["123", "456", "789"] (or similar 3-digit chunks). Different numbers tokenize differently: 1234 may be ["12", "34"] while 12345 may be ["123", "45"]. The model performs arithmetic on these symbolic chunks, not on individual digits.
Fix. Recent models (Llama 3, GPT-4o math fine-tunes) use single-digit tokenization for numbers. Singh et al. (2024) show this alone yields large gains on multi-digit arithmetic.
" cat" (with leading space) and "cat" (without) are different tokens with different embedding rows. A sentence is typically a sequence of space-prefixed words: "the cat sat" → ["the", " cat", " sat"]. The first word lacks a leading space; subsequent words include it. This is why prompts that omit a leading space sometimes produce subtly different outputs.
UTF-8 encodes ❤️ as 3 bytes (\xe2\x9d\xa4 for ❤, plus the variation selector). BPE trained mostly on English never merged these byte sequences. Result: each emoji byte is one token. Consequences:
strawberry is one or two tokens depending on the model. The forward pass never sees individual letters: it operates on dense vectors representing whole subwords. To answer “how many r’s in strawberry?” the model must have learned to internally spell tokens back out, a non-trivial skill that requires it in training. As of 2024 most models still failed this; explicit training on letter-counting tasks fixes it.
| Model | Vocabulary size | Encoding |
|---|---|---|
| GPT-2 | 50,257 | byte-level BPE |
| GPT-3.5 / GPT-4 | 100,277 | cl100k_base |
| GPT-4o | 200,019 | o200k_base |
| Llama 3 | 128,256 | tiktoken-based BPE |
| Claude | ~100K+ | proprietary BPE |
A given prompt produces a different number of tokens at different positions in each model. Position-indexed MI results do not transfer without re-tokenization.
Type or paste text. The demo estimates how each major tokenizer would count it, then computes the API cost at current pricing. Try the same paragraph in English, then in Japanese, then with emoji, the count differences are real and they cost real money.
Why models disagree: Each tokenizer was trained on a different mix of languages. GPT-4's vocabulary is more multilingual than GPT-2's, so the same Japanese sentence costs fewer tokens in GPT-4. Llama 3 has a 128K vocabulary biased heavily toward code. Counts are heuristic estimates (≈ chars-per-token ratios from published benchmarks); load tiktoken for exact numbers.
Inside the transformer:
Every operation in the model is an operation on tokens.
Most models reserve specific IDs outside the regular BPE vocabulary:
<|endoftext|> (<bos>, <eos>): document boundary.<|im_start|>, <|im_end|>: chat role markers (OpenAI ChatML format).[INST], [/INST]: instruction markers (Llama 2/3 chat format).<|system|>, <|user|>, <|assistant|>: chat-tuned models.These tokens behave like any other input but carry structural meaning. The BOS token in particular accumulates state that several attention heads use as a “rest” position, often called the BOS sink (Xiao et al., 2023; also discussed in Bricken et al., 2023).
import tiktoken
enc = tiktoken.encoding_for_model("gpt-4")
ids = enc.encode("How many r's are in strawberry?")
pieces = [enc.decode([i]) for i in ids]
print(list(zip(ids, pieces)))
# [(4438, 'How'), (1690, ' many'), (436, ' r'), (596, "'s"), ...]
For Llama / Mistral / open models:
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")
print(tok.tokenize("strawberry"))
# ['str', 'aw', 'berry']
When studying a new model, dump its vocabulary and search for prompt fragments before drawing conclusions about positions.
The model’s alphabet is not the user’s alphabet. Tokens are the atom: every embedding lookup, attention computation, and loss term is defined per token. Tokenization choice has measurable downstream effects on arithmetic, multilingual fluency, and character-level reasoning.
The next post is on the residual stream: what those token vectors do once they enter the transformer.
This post defines each station precisely, explains why the architecture replaced RNNs, and sets up the vocabulary used in the rest of the series.
Six stations. A sentence walks in. A guess at the next word walks out. Tap any station to see what it does and what the token looks like after its shift.
1. Tokenizer splits text into numeric IDs (often subword pieces). 2. Embedding looks up a vector for each ID from a big table. 3. Attention lets every position peek at every other position and pull in what's relevant. 4. MLP is a per-position feed-forward network that transforms the vector non-linearly. 5. Repeat stacks attention+MLP dozens of times. 6. Unembedding projects the final vector into a probability over the vocabulary. The belt you see is the residual stream, the running total that every block reads from and writes to.
Click a station, press play, watch one token’s vector traverse the pipeline. The structure is identical for every modern LLM: GPT-2, GPT-3, GPT-4, Claude, Llama, Gemini.
1. Tokenizer. Maps a string to a sequence of integer IDs using a fixed vocabulary $V$ (typical sizes: 50,257 for GPT-2, 100,277 for GPT-4, 128,000 for Llama 3). Modern tokenizers use byte-pair encoding (BPE) or its variants. Example: "unhappily" → ["un", "happ", "ily"] → [403, 7829, 6148].
2. Embedding. A learned lookup table $W_E \in \mathbb{R}^{V \times d_\text{model}}$ maps each token ID to a $d_\text{model}$-dimensional vector. Common sizes: $d_\text{model} = 768$ (GPT-2 small), 4,096 (Llama 3 8B), 12,288 (GPT-3 175B). The $i$-th row of $W_E$ is the embedding for token $i$.
3. Self-attention. At each position $t$, the model computes weighted sums over all positions $\le t$ (causal mask). Weights are derived from query/key dot products. Output: a new vector at each position that has read from earlier positions. Multiple heads run in parallel ($n_\text{heads} = 12, 32, 96, \ldots$), each in a lower-dimensional subspace.
4. MLP. A position-wise feed-forward network: two linear layers with a non-linearity (GeLU, SwiGLU). Hidden dimension is typically $4 d_\text{model}$. Same shape in, same shape out. Operates independently on each position.
5. Repeat. One block = attention + MLP + residual connections + layer norms. Stack $N$ of them. GPT-2 small: $N=12$. Llama 3 8B: $N=32$. GPT-3 175B: $N=96$. GPT-4 (rumoured): $N \approx 120$.
6. Unembedding. Multiply the final residual stream vector by $W_U \in \mathbb{R}^{d_\text{model} \times V}$ to produce logits. Apply softmax to get a probability distribution over the vocabulary. Sample → next token.
Total parameters scale roughly as $N \cdot d_\text{model}^2 \cdot 12$ for the standard recipe.
Five logits, one softmax, one temperature dial. Drag the slider and watch the probability mass redistribute. This is the very last computation in a transformer: scores in, distribution out, sample one.
What you're watching: as T drops toward 0 the distribution sharpens onto the top logit (greedy). As T grows the distribution flattens toward uniform (chaos). Temperature 1.0 is the model's natural distribution. Most chatbots use 0.7, 1.0.
In the literature this conveyor belt is the residual stream. Each block’s output is added to its input, not substituted:
\[x_i \;\leftarrow\; x_i + \text{block}(x_i)\]What you should notice: the embedding alone gives a vague guess. As more blocks get added, certain feature dimensions (place / person / tense …) light up and the prediction sharpens. Mute any single block — the loss in confidence tells you that block's contribution. That's exactly what direct logit attribution does at scale.
Three consequences follow directly from this additive structure:
Elhage et al. (2021, “Mathematical Framework”) formalize this view. The residual stream is the central object in the rest of this series.
Recurrent neural networks (RNNs) process tokens sequentially, maintaining a hidden state $h_t = f(h_{t-1}, x_t)$. Two structural problems:
Transformers (Vaswani et al., 2017) replace recurrence with self-attention: every position reads from every other position in parallel via a single matrix multiplication. Sequence length $T$ → $O(T^2)$ time, but fully parallel within a sequence and across batch dimensions. On modern GPUs the constant factor wins.
The original 2017 paper had two halves: an encoder (reads input) and a decoder (writes output), wired for machine translation. For autoregressive text generation, only the decoder is needed.
A decoder-only transformer uses causal self-attention: position $t$ may only attend to positions $0, 1, \ldots, t$. This enforces left-to-right generation and matches the next-token-prediction training objective.
This series is exclusively about decoder-only transformers. (Encoder-only models like BERT are used for classification and embedding tasks, not generation.)
The transformer produces one probability distribution per forward pass. To generate text:
input: "The cat sat on the" → logits → sample → "mat"
input: "The cat sat on the mat" → logits → sample → "."
input: "The cat sat on the mat." → logits → sample → "<eos>"
This is autoregressive decoding. Each generated token is appended to the input and the model runs again. Sampling strategies (greedy, top-k, top-p, temperature) determine which token gets picked from the distribution. KV-caching avoids recomputing attention over previous tokens, making the per-step cost roughly $O(T)$ rather than $O(T^2)$.
| Model | $N$ | $d_\text{model}$ | $n_\text{heads}$ | $V$ | Parameters |
|---|---|---|---|---|---|
| GPT-2 small | 12 | 768 | 12 | 50,257 | 124M |
| GPT-2 XL | 48 | 1,600 | 25 | 50,257 | 1.5B |
| GPT-3 | 96 | 12,288 | 96 | 50,257 | 175B |
| Llama 3 8B | 32 | 4,096 | 32 | 128,000 | 8B |
| Llama 3 70B | 80 | 8,192 | 64 | 128,000 | 70B |
Memorize one row (GPT-2 small is most common in MI papers) and use it as the reference point.
Three sliders. Pick a depth, width, and head count. The demo computes total parameters, the attention/MLP split, FP16 memory, and matches your settings to the closest published model.
What you're learning: in any standard transformer, MLPs hold roughly twice the parameters of attention. Move the sliders and the ratio barely budges, ~⅔ MLP, ~⅓ attention, plus a small embedding/unembedding chunk that shrinks as you scale up. This is why interpretability work that focuses only on attention is missing where most of the model lives.
Subsequent posts cover each station in mechanistic detail:
The next post is on tokenization.
All the grammar, all the facts, all the reasoning patterns. Learned by being wrong, measuring how wrong, and nudging. Billions of times. For months.
This procedure has a name: gradient descent. And honestly, it’s the closest thing AI has to a creation story.
You can’t write rules for recognising cats. Smart people spent decades on this, did genuinely incredible work along the way, and it turned out the hand-written-rules approach just hits a ceiling. That’s how we ended up over here.
You can’t manually set 70 billion weights either. Even if you magically knew what they should be (you don’t, nobody does), it’d take longer than the age of the universe.
So what do you do? Show the model examples. Tell it when it’s wrong. Let it adjust itself. Repeat.
Do that enough times and, somehow, it works. That’s machine learning in one sentence. “Gradient descent” is just the specific algorithm that does the adjusting.
Before the model can learn from being wrong, you need a way to measure how wrong it is. That’s the loss function.
Classification problem: you ask the model is this email spam?. It outputs a probability: 0.73 (73% chance spam). The correct answer is 1.0. The loss is a number measuring the gap. How far off were you?
Common loss functions:
The loss is a single number. Low = good predictions. High = bad predictions.
Goal of training: minimise the loss, averaged over millions of examples.
Way to picture it.
Imagine a two-dimensional landscape. Hills and valleys. Every point corresponds to a specific setting of all the weights. The height at any point is the loss. How wrong those weights make the model.
The model starts at a random point in this landscape. A random mountain, somewhere. Our job: walk it to the nearest valley. The point of lowest loss.
Here’s the thing though: this landscape isn’t 2D. It has as many dimensions as there are weights. A small model: millions of dimensions. GPT-4: hundreds of billions. Try picturing that for a second. Actually don’t, you’ll hurt yourself.
Nobody can visualise this landscape. But we can navigate it, one step at a time, using the gradient.
Every point on this map is a setting of the model's weights. The colour is the loss, how wrong the model is. Gradient descent is a ball rolling downhill. Try it with different learning rates.
Try: LR ≈ 0.08 for smooth descent · LR > 1.2 for chaos · reset a few times to find the local minimum trap on the upper left.
The surface is a sum of Gaussian bumps and wells, a crude 2D stand-in for the millions-of-dimensions loss surface of a real model. At each step we estimate the gradient with a finite-difference ([f(x+ε) − f(x)] / ε), then apply x ← x − lr · ∇f. That's the literal equation of gradient descent. The chart below tracks loss vs step count, the curve you actually see in training runs.
The demo above is only 2D, but the algorithm is exactly what runs inside every deep-learning training loop. Run it with learning rate 0.1. Now 1.3. Now 0.01. Notice how LR changes everything. And how the ball sometimes settles into the shallow valley instead of finding the deeper one.
The gradient is a vector that points in the direction of steepest increase in loss. It tells you: “if you move this way, you’ll get worse.”
So we go the opposite direction.
The gradient with respect to a single weight tells you: “if I increase this weight slightly, does the loss go up or down, and by how much?”
Do this for every weight simultaneously. Take a step in the direction of decreasing loss. That’s gradient descent.
new_weight = old_weight − (learning_rate × gradient)
The learning rate is one of the most important numbers in training. Controls how big each step is.
Finding the right LR is part science, part art. Modern training uses adaptive learning rates. Algorithms like Adam adjust the step size for each weight individually based on how the gradient has been behaving.
Drag the slider all the way down — the ball creeps and gives up before reaching zero. Drag it past 1.0 — the ball bounces off both walls of the bowl. Drag it past 2.0 — the ball escapes the canvas entirely. There’s no good rule for “the right” learning rate. Every modern optimizer is essentially a different way of guessing it from local geometry.
Problem: a model has billions of weights. Computing the gradient for all of them by testing each one (“what if I nudged this weight up a tiny bit?”) would take forever.
Backpropagation (backprop) solves this with the chain rule of calculus. An algorithm that computes the gradient for every weight in a single backward pass through the network. In essentially the same time it takes to run a forward pass.
Key idea: the gradient flows backwards. Start from the loss (at the output). Compute how much each weight contributed to it, layer by layer, going backwards toward the input.
You don’t need to understand backprop’s math to understand MI. But know this: it’s fast, it’s exact, and every major training framework (PyTorch, JAX, TensorFlow) does it automatically. You define your model, you compute the loss, you call .backward(). Gradients appear, pre-computed, on every weight.
You don’t compute the gradient on one example at a time. You compute it on a batch. Typically 128, 512, or 2048 examples at once.
Why batches?
One epoch = one full pass through the entire training dataset. Models are typically trained for many epochs. Same data multiple times.
Modern training uses mini-batches with Adam. An improved version of SGD with adaptive learning rates and momentum.
Three runners, same start. Plain SGD only knows local slope, so on a long thin valley it ricochets between walls. Momentum carries velocity through small bumps and accelerates downhill. Adam adapts each direction’s step size independently and tends to find the basin even on weird surfaces. Click anywhere on the canvas to teleport the start — every optimizer’s character changes with the terrain.
A model can memorise its training data. Loss goes to zero. On new data it’s never seen (test data), the model fails badly.
That’s overfitting. The model learned the specific examples, not the underlying patterns.
Detecting it: track loss on training data and on held-out validation data separately. Validation loss starts rising while training loss keeps falling, the model is overfitting. Training usually stops here.
Techniques to prevent it: dropout (randomly silence neurons during training), weight decay (penalise large weights), data augmentation (artificially expand training data).
Gradient descent optimises loss. That’s all it does.
Does not guarantee the model learned the right algorithm. Does not guarantee it will generalise. Does not guarantee it’s doing what you think it’s doing.
The model might find a shortcut. A way to get low loss without learning the intended behaviour. This is the “specification gaming” problem in a nutshell. Gradient descent won’t catch it. It has no idea what the “intended behaviour” is. It just minimises the number you gave it.
Alright, real talk: between this post and the ones before it, you now know how a single neuron works, how a bunch of them chain into layers, how weights store knowledge, and how all of that was produced in the first place. That’s enough vocabulary to follow almost any mechanistic-interpretability paper on modern neural networks.
Layer 1 sees edges. Diagonal lines, curves, horizontal stripes. Layer 2 sees eye corners, nose tips, ear lobes. Layer 5 sees this person. Their mood. Whether they’re wearing glasses.
Same pixels going in. Totally different lens at every level. Each layer is watching the one below it and writing down what it noticed. It’s like a game of telephone, except the message gets smarter at every hop.
That’s what depth is. That’s literally why “deep” learning is deep.
A layer is a group of neurons that all receive the same inputs and all produce outputs that feed forward together.
Every neuron in a layer:
Net result: each layer transforms its inputs into a new representation. Same information, viewed from a different angle, with different things highlighted.
Slide through the layers of a small vision model. At the bottom you see pixels. At the top you see categories. In between, the network turns the first into the second.
A real pretrained model is too heavy to ship inline, so this is an illustrative pipeline: Layer 0 draws the image on a canvas; Layer 1–2 apply Sobel edge detection; Layer 3–5 show animated feature-map tiles; Layer 6–8 draw a class-activation-style heatmap overlay; Layer 9–11 blur toward "concept space"; Layer 12 shows the final class distribution. The progression mirrors what happens in a real ConvNet (see Zeiler & Fergus 2013, Olah et al. 2017). Swap the image, the predictions change accordingly.
Input layer. No math. Just receives the raw data and passes it forward. Images: pixel values. Text: token embedding vectors. Tabular data: feature values. The input layer is pure format conversion. “Here’s the world, in number form.”
Output layer. The final decision. Classification: a probability over each class (softmax). Language modelling: a probability over every word in the vocabulary (often 50,000+ options). Regression: a single continuous number.
Everything interesting happens between these two.
Don’t let the word fool you. “Hidden” just means “not the input, not the output.” These layers are where all the action happens, and where MI spends basically all its time.
In early vision models, researchers noticed the hidden layers had a striking, almost biological structure. Here’s the rough gradient:
Neurons respond to oriented edges. Horizontal, vertical, 45-degree. Nobody programmed this. Every model trained on natural images independently discovers it. Emerges from the statistics of images themselves.
Combinations of edges form textures. Checkerboards. Crosshatches. Dots. Neurons looking for local patches that match a pattern.
Eyes. Wheels. Leaves. Neurons now looking for parts of real objects. Things that have names in human language.
Full objects, faces, specific categories. High-level human concepts represented in the network’s internal language.
This progression is called hierarchical feature extraction, and it appears in every deep network trained on natural data. Images, text, audio. The depth lets the model compose simple features into complex ones, repeatedly.
Click an object in the bottom row, then trace your way up. Every face is a composition of eye-detectors and mouth-detectors. Every eye-detector is a composition of curve-detectors and edge-detectors. The whole thing is recursive: one small library of features, reused at every level of detail.
In text transformers, the same hierarchy exists, but it’s harder to see because language doesn’t have the obvious spatial structure of images.
Research into transformer layers has found patterns like:
Not perfectly clean (features mix across layers), but the gradient from syntax to semantics is consistent and has been verified by probing experiments across many models.
A cool MI technique that reads out the model’s best guess at each layer, before the final output. Early layers, the guess is mostly garbage. Middle layers, it starts approaching the right semantic category. Late layers, it converges on the final answer. Shows you how the model builds its answer progressively.
Width = more neurons per layer. More “workers” doing parallel analysis at each step.
Depth = more layers. More steps of abstraction before the final answer.
Both help, but differently.
Width gives the model more capacity to represent complex things at each level of abstraction. Depth gives the model more steps to compose simple patterns into complex ones.
Modern networks are both wide and deep. GPT-4 is estimated to have around 120 layers. Many of those layers have ~25,000 neurons each. Which, you know, is a lot.
For interpretability: more layers means more places for information to transform. Also means there’s more “room” for information to be stored in intermediate representations. Which is one reason large models are more capable but also harder to interpret.
Try the Spiral dataset with 0 hidden layers — a single linear boundary can’t separate it, no matter how long you train. Now bump depth up to 3 and watch the boundary curl. Every additional layer is one extra fold the network can put into space. Width adds patience; depth adds expressiveness.
In modern networks (including all transformers), there’s a trick that changed everything: residual connections. Also called skip connections.
Instead of each layer replacing the previous layer’s output entirely, it adds to it. The output of layer N+1 = what layer N produced + what layer N+1 computed.
Sounds small. It’s enormous.
Means information can flow directly from early layers to late layers without passing through every layer in between. Early features don’t get “forgotten” or overwritten.
Also means each layer can focus on adding a small correction, rather than computing everything from scratch. Makes training much more stable and allows much deeper networks.
Understanding what each layer does is one of the central projects of mechanistic interpretability. Not “layer 7 does something useful”. Exactly what. Which features live in which layers. Which operations happen where. When we know that, we can start to decompose a model’s behaviour the same way you’d decompose a program into functions.
Okay, one more thing worth naming. Every weight in every layer we’ve talked about was set by a single procedure: gradient descent. I’ll write about how that actually works in the next blog.
Neural networks somehow figured this out. Every fact, every grammar rule, every pattern a model ever learned is in the weights. Not in any single weight. Not labelled. Just pressed into the collection, in a way nobody designed and nobody fully understands.
Weights are the most important thing in a neural network. They’re also the hardest to read.
So let’s learn to read them.
A weight is one number. It lives on a connection between two neurons.
One number, one relationship, one direction of influence.
Scale this up. A model with 70 billion parameters has 70 billion of these little relationships. All learned from data. All working together to spit out something coherent on the other end.
The staggering bit: nobody wrote a single one of them. Nobody sat down and decided “the word ‘not’ should have a negative weight on the sentiment neuron.” The model figured all of it out, by reading enough examples. That still feels slightly like magic to me, honestly.
Here's a tiny network, two inputs, three hidden neurons, one output (positive / negative sentiment). Click any connection to change its weight and watch the prediction move.
wTwo layers: input (2) → hidden (3, ReLU) → output (1, sigmoid). Matrix math in plain JS. Line colour is cyan for positive weights, red for negative; thickness scales with magnitude. The five test sentences are hardcoded; the ✓/✗ indicator checks the model's prediction against the known label with each weight change, so you can watch the trained model break the moment you misalign a weight.
Play with the demo. Flip the weight from word positivity to negative and watch the model’s prediction invert. Positive reviews now get classified as negative. Try the Broken preset, then Trained. Nothing “inside” the model changed except a handful of numbers. That’s all weights are.
When every neuron in one layer connects to every neuron in the next, you get a weight matrix.
Layer A has 4 neurons. Layer B has 3 neurons. You have a 4×3 grid of weights. 12 numbers, each the strength of one connection.
To calculate layer B’s activations, you multiply: B = W · A. Matrix multiplication.
This is the fundamental operation of a neural network. Everything (attention, MLP layers, embeddings) is built from variations of this.
For interpretability, the weight matrix is where we look for structure:
That last one, matrix factorisation, is a major MI technique. If W decomposes into A × B, then A and B might represent something interpretable.
Random weights, before training: all small, roughly centred on zero. The network produces noise.
After training on language: the weights organise into structure. Not structure we designed. Structure that reflects the regularities in language.
The “small and centred on zero” part isn’t aesthetic — it’s load-bearing. Pick the wrong starting distribution and the signal either dies before it reaches the output or explodes before the first gradient step. He and Xavier initialization aren’t tricks; they’re the only reason deep networks train at all.
Words get encoded as high-dimensional vectors. The weights arrange these vectors so that:
"king" − "man" + "woman" ≈ "queen"
"Paris" − "France" + "Italy" ≈ "Rome"
Not programmed. Emerges from the weights learning which words appear in similar contexts.
In transformers, attention weights form patterns like:
These patterns live in the weight matrices. Finding them is a big chunk of mechanistic interpretability.
Here’s the annoying part. Print out a weight matrix, you see a grid of numbers like 0.023, −0.41, 0.0017, 1.3, −0.88… and?
It tells you almost nothing. The numbers only mean something in combination. One weight doesn’t represent a concept. The whole matrix does.
This one trips people up.
Weights are fixed after training. Don’t change when you give the model a new input. They’re the structure. The compiled knowledge of everything the model learned.
Activations are dynamic. Computed fresh for every input. The model’s current state of processing your specific prompt.
For MI: most research looks at activations (what’s the model thinking about this input?) but relates them back to weights (what in the structure caused this pattern?). Both matter.
One sentence on how they got this way: during training, the model sees an example, makes a prediction, measures how wrong it was, and nudges every weight slightly in the direction that would’ve made it less wrong.
Do that billions of times, across trillions of words. Yes, literally trillions. Yes, it feels absurd. It also works.
That’s gradient descent in miniature. Every weight in a 70B-parameter model gets updated by the exact same logic: forward to compute loss, backward to compute who’s responsible, then nudge each weight downhill. The “trillions of nudges” is just this loop, looped a lot.
The weights that emerge encode the statistical regularities of everything the model was trained on. Grammar. Facts. Logic. Poetry. Chemistry. Slang. All of it. Compressed into numbers.
For now, “it’s gradient descent” is enough.
When an MI researcher asks “what did this model learn to do?”, they’re asking “what do these weights mean?” Finding the answer requires figuring out which directions in weight space correspond to human-interpretable concepts. Which is the project of the whole field.
I’ll write about layers in the next blog. Many weight matrices stacked on top of each other, each doing something different to the flow of information.