🤖 machine learning

Exemplar Partitioning for Mechanistic Interpretability

This paper introduces Exemplar Partitioning (EP), an unsupervised method that constructs interpretable feature dictionaries from language model activations using observed exemplars rather than learned weights, achieving comparable interpretability performance to sparse autoencoders with approximately 1,000 times fewer training tokens while enabling direct cross-model comparisons and providing built-in out-of-distribution detection.

Original authors: Jessica Rumbelow

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jessica Rumbelow

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, chaotic library where every book represents a single thought or sentence a computer model has ever processed. Inside this library, the "ideas" aren't written in words; they are hidden as complex patterns of light and electricity.

For a long time, scientists trying to understand these models have used a method called Sparse Autoencoders (SAEs). Think of this like hiring a team of expensive, highly trained librarians to read every single book, sort them into thousands of specific categories, and write a new index. It takes a massive amount of time and money (computing power) to train these librarians, and the categories they create are based on what they learned to look for.

Exemplar Partitioning (EP), the method introduced in this paper, is a completely different, much cheaper approach. Instead of training librarians, it uses a simple, automated sorting machine.

Here is how it works, using everyday analogies:

1. The "First-Come, First-Served" Sorting Machine

Imagine you are walking through a crowded room (the stream of data). You want to group people based on how similar they look.

The Old Way (SAE): You hire a designer to create a perfect, pre-defined set of 10,000 "bins." You then spend weeks training a system to figure out exactly which person goes into which bin to minimize errors.
The New Way (EP): You just walk through the room. When you see the first person, you say, "Okay, you are the Leader of Group A." Everyone who looks very similar to this Leader joins Group A.
- If you see someone who looks different from the current Leaders, you say, "You are the Leader of a brand new Group B."
- You don't decide in advance how many groups there will be. The groups form naturally based on how the people in the room actually look.

2. The "Anchor" Concept

In this new method, every group is anchored by a real person you actually saw (an "Exemplar").

Why this matters: In the old method, a group might be defined by an "average" of thousands of people, which is a ghostly, imaginary person that doesn't actually exist. In the new method, every group is tied to a real, specific example. If you want to know what "Group A" is about, you just look at the first person who started it. You can even point to that specific person and say, "If we remove this person's influence, the group disappears."

3. The "Free Map" of the Room

Because the groups are formed by simple distance (how similar people look), the method creates a perfect map of the room.

The "Out-of-Place" Detector: If you walk into the room and see someone who looks nothing like any of the Leaders, the system immediately knows, "This person doesn't belong to any group we've seen before." This is a free way to spot strange or unexpected inputs without any extra training.

4. What the Paper Actually Found

The authors tested this on a specific AI model (Gemma-2-2B) and found some surprising things:

It's incredibly fast and cheap: They built these dictionaries using about 1,000 times less computing power than the traditional method. It took minutes on a single computer chip instead of weeks of training.
It works just as well for finding concepts: When they asked the system to find specific ideas (like "math," "code," or "refusal to answer"), it performed almost as well as the expensive, trained method.
It found the "Refusal" switch: They discovered a specific group of inputs where the AI decides to say "No" (refuse a request). By looking at the "Leader" of that group, they could prove that if you remove that specific pattern, the AI stops refusing. This shows the groups are real, causal parts of how the AI thinks.
It sees the world differently: The new method groups things by density (clumps of similar things), while the old method groups things by lines (mathematical separation). They agree on the most obvious, clear ideas (the "core"), but they split up the messy, complex ideas differently.

5. The "Training" vs. "Watching" Difference

The most important takeaway is that this method doesn't learn anything. It doesn't try to guess the best way to sort things. It just watches the data stream, picks the first few unique examples it sees, and sorts everything else based on how close it is to those examples.

Because it doesn't rely on a specific training run, you can compare the "map" of a model before it was trained and after it was trained, and you can see exactly which groups stayed the same and which ones changed. It's like comparing a map of a city before and after a new highway is built, using the same landmarks.

In summary: This paper introduces a way to understand AI models by simply organizing their thoughts into natural clusters based on real examples, rather than forcing them into pre-made boxes. It's faster, cheaper, and just as good at finding the "meaning" inside the machine.

Technical Summary: Exemplar Partitioning for Mechanistic Interpretability

Problem Statement

Mechanistic interpretability seeks to reverse-engineer neural networks into human-legible algorithms. Current approaches face distinct limitations:

Circuit Analysis: Isolates task-specific subgraphs but fails to produce a model-wide feature inventory.
Supervised Probing: Limited to testing concepts researchers already know to look for.
Sparse Autoencoders (SAEs): The dominant unsupervised approach, SAEs learn a fixed-size dictionary by optimizing a reconstruction objective that bundles sparsity, reconstruction, and dictionary size. This training process is computationally expensive (requiring billions of tokens) and introduces commitments (e.g., fixed dictionary size, linear basis learning) that may obscure intrinsic properties of the model's activation geometry.

The central question posed is: How much of the feature structure SAEs uncover is due to the training objective itself, versus how much is already present in the model's activation geometry and accessible via a more direct, non-learned construction?

Methodology: Exemplar Partitioning (EP)

Exemplar Partitioning is an unsupervised method for constructing interpretable feature dictionaries from Large Language Model (LLM) activations without training. It treats the dictionary as a Voronoi partition of the activation space, built via leader clustering on a single streaming pass.

Core Construction

Streaming Leader Clustering: Instead of averaging activations to form centroids (which can create incoherent directions and cause clustering instability), EP uses fixed exemplars. An exemplar is a real, observed activation from the data stream.
- When a new activation arrives, it is centered and normalized.
- If its distance to the nearest existing exemplar is below a threshold, it joins that region.
- If it exceeds the threshold, it becomes a new exemplar, seeding a new region.
Calibration and Thresholding: The distance threshold $\theta_p$ is not an absolute value but the $p$ -th percentile of pairwise distances computed on a small calibration batch. This allows the resolution of the dictionary to be tuned (smaller $p$ = finer granularity; larger $p$ = coarser) while remaining comparable across layers and models.
Saturation Criterion: The dictionary grows until a "saturation window" passes with no new exemplars. This means the dictionary size is not prespecified but determined by the activation geometry at the chosen resolution.
Inference: New activations are assigned to their nearest exemplar. The dictionary is natively 1-sparse (one-hot) but can be extended to $n$ -sparse or dense assignments based on distances.

Key Properties

Observability: Since regions are anchored by observed exemplars rather than learned weights, dictionaries built from the same data stream are directly comparable across layers, models, and training checkpoints via exemplar matching.
Stability Metric: The paper introduces a statistic $D_i = \log_{10}(N_i c_i^2)$ (where $N_i$ is member count and $c_i$ is member coherence) to predict the cross-seed stability of a region. High $D_i$ regions are more likely to replicate across different streaming orders.

Key Contributions and Results

1. Efficiency and Scalability

EP constructs dictionaries with $\sim 10^3\times$ fewer tokens than comparable SAEs.

Build Cost: A p1 dictionary for Gemma-2-2B builds in minutes on a single GPU using $\sim 3.6 \times 10^6$ tokens, whereas comparable SAEs require training on billions of tokens.
No Backpropagation: The method requires zero backward passes.

2. Interpretability and Causal Intervention

Causal Ablation: In Gemma-2-2B-it, refusal behavior concentrates in a specific EP region. Ablating the exemplar direction (the first-arrival activation) of this region collapses held-out refusal by up to $\Delta = -0.96$ . This demonstrates that the "refusal direction" is a precise axis captured by a single exemplar, which is more effective for intervention than the region's mean direction.
Cross-Checkpoint Matching: Matching base and instruction-tuned dictionaries reveals that instruction tuning largely re-anchors the directional structure of the model, preserving only a small universal subset. It also shows that the refusal direction in the tuned model is promoted from a within-prompt harmful-content direction present in the base model.

3. Benchmark Performance

AxBench (Concept Detection): On Gemma-2-2B-it L20, an EP dictionary at resolution $p=1$ $p = 1$ achieves a mean AUROC of 0.881 across 500 latent concepts.
- This is +0.126 over the canonical GemmaScope SAE entry (0.755).
- It is within 0.030 of the strongest baseline, SAE-A (0.911), despite using a smaller candidate pool and significantly less compute.
SAEBench (Information Retention): A p10 EP dictionary ( $\ell_0=1$ ) retains $\sim 97\%$ of raw-activation probe accuracy (0.930 vs. 0.959 for IdentitySAE). This suggests that the linearly decodable identity of concepts is largely preserved by the density structure alone, even when discarding almost all raw coordinates.

4. Geometric Insights

Density vs. Linear Separability: EP and SAEs make different geometric commitments. EP commits to density (regions are real cones in activation space), while SAEs commit to linear separability (sparse reconstruction).
Shared Core: Approximately 20% of EP regions match an SAE feature at $F1 > 0.5$ (at $p=10$ ). They agree where directions are both density-concentrated and linearly separable.
Out-of-Distribution (OOD) Signal: The distance to the nearest exemplar provides a free OOD signal. Random tokens are consistently further from their nearest exemplar than in-distribution tokens, with the gap widening at finer resolutions.

Significance and Claims

The paper claims that Exemplar Partitioning offers a new class of interpretability objects that are:

Efficient: Accessible to researchers without frontier-lab compute budgets.
Direct: They reveal properties of activation geometry (density, stability, causal directions) that may be obscured by the reconstruction objectives of SAEs.
Commensurable: The use of observed exemplars allows for direct tracking of feature evolution across training checkpoints and model variants without complex alignment procedures.

The authors emphasize that while EP captures a "shared core" with SAEs, it also exposes broad, content-anchored regions that SAEs splinter, and conversely, SAEs capture sparse linear directions that EP does not resolve at a given resolution. The paper concludes that density structure alone carries much of the signal measured by probing, and that the method's efficiency and open dictionary interface make it a practical addition to the interpretability toolkit.

Limitations Acknowledged:

EP is seed-dependent; region identity shifts with streaming order, though stability can be predicted via $D_i$ .
The current implementation assumes centred cosine geometry, which may not capture features dependent on activation magnitude.
The experiments are exploratory; the paper does not claim EP is a complete replacement for SAEs but rather a complementary, highly efficient alternative for specific interpretability tasks.