Most of what we currently call "feature discovery" in language models is wrapped up in dictionary-learning methods like sparse autoencoders (SAEs) – which work, and which have been scaled to millions of features on frontier-scale models, but which bundle two distinct commitments into a single training objective: a reconstruction loss and a sparsity loss over a fixed size dictionary. Those commitments make sense if your goal is reconstructive decomposition – if you want to take an activation and rebuild it from a sparse code. They make less obvious sense if your aim is to find interpretable structure (directions? features?) in activation space, to retrieve representative examples, identify causal interventions, or measure how representations change across layers and inputs. And it turns out a lot of that doesn't really need the full SAE machinery.

An Exemplar Partitioning dictionary built from Gemma-2-2B activations at p₂ (K = 5,129 regions). Left: eight sample regions, each shown with its member count, its exemplar's logit-lens decode, and an excerpt of a member input with the activating tokens highlighted. Right: a PCA-projected 3D rendering of the Voronoi partition; each cell is one region, with a random selection also labelled with logit-lens decode.

This post is about a method that strips out all of those commitments and just covers the activation manifold with observed exemplars at a calibrated resolution. It has one hyperparameter, makes one streaming pass over the data with no backward passes or gradient descent – and despite that, on the AxBench latent concept-detection benchmark at Gemma-2-2B-it layer 20, EP at p₁ reaches 0.881 mean AUROC across all 500 concepts. That's within 0.03 of SAE-A – AxBench's strongest dictionary-based baseline – with about 1,000× less build compute (EP at p₁ used 3.6 × 10⁶ activation tokens, and does no gradient descent; the canonical GemmaScope 16k SAE on Gemma-2-2B was trained on ~4 × 10⁹ activation tokens with ~10⁶ optimiser steps).

Glossary

Activation: a vector extracted from a forward pass of the model on some input, at a chosen layer and hook point. The thing we're clustering to make the Voronoi partition.

Voronoi partition: a partition of a space into cells (or regions), where each region contains the set of points closest to one designated anchor (in this case, an exemplar activation).

Region: one Voronoi cell of the EP dictionary, anchored by its exemplar (and normally associated with other statistics, like the mean activation across members and the number of members, that you might want to look at later).

Exemplar: a real, observed activation that anchors a region. The first activation to land in a new cell becomes that cell's exemplar – it stays fixed for the lifetime of the dictionary, and serves both as the membership criterion (new activations join if they're close enough) and as a direction you can read off, project onto, or ablate downstream. You can also use the mean activation in the region for these tasks, which is explored further in the paper.

EP Dictionary: the collection of exemplars and their regions, together with the procedure for assigning new activations to regions.

p: the percentile of pairwise activation distances (computed once on a calibration stream) used to set the clustering threshold. Smaller p = tighter threshold = more regions = finer dictionary. Larger p = looser threshold = fewer regions = coarser dictionary.

Exemplar Partitioning

Exemplar Partitioning (EP) constructs a Voronoi partition of activation space. You can read it like a dictionary mapping activations to Voronoi cells (regions), using logit lens readouts (a quick way to interpret intermediate activations by passing them through the model's unembedding), example prompts, important or common tokens, and so on. You build the EP dictionary by leader-clustering (Hartigan, 1975) on activations, which is very simple.

There is one important hyperparameter, a distance threshold which determines when to assign new activations to clusters. Call this p, the distance percentile (see the paper §2 for the calibration procedure). In practice, you probably build multiple EP dictionaries at different p, which you can think of as looking at activation space in higher or lower resolution: typically, a low threshold produces many small regions with more specific features, and a high threshold produces fewer, coarser regions corresponding to more general features.

The EP dictionary is built by streaming a model's activations on any data corpus (just hook a model at some layer, as you would if training an SAE). For each new activation, see if there exists an exemplar within the distance threshold – if so, the new activation joins its region. Otherwise, it becomes the exemplar of a new region. Keep doing that until some large number of activations passes without producing a new exemplar.

For each region, maintain the exemplar, a running average activation, some stats about the region members, and optionally an index of member prompts/activations so you can look at them later. You don't have to keep all activations in memory, so this is computationally tractable, and in fact really fast and token efficient. No backward passes or gradient descent – cheap as chips.

Leader-clustering animation: a streaming sequence of points where each point either joins an existing cluster (within the distance threshold of its anchor) or seeds a new one.

That's the whole algorithm! There are a couple of small subtleties around calibrating the distance threshold, and transformations you might apply to the activations before clustering (e.g. centring, normalising, which I do in the paper to build the partition on the unit sphere) but that's basically it.

Inference

By definition, Voronoi partitions tile the entire space – so at inference (given a new activation, find which region it lives in), every activation maps to exactly one region. You get 1-sparsity for free! But you can also compute the distance from the new activation to every exemplar, so sparsity becomes a readout choice rather than a construction hyperparameter. You can take the top 1, the top 5, or the full distance vector. Or pick the top features via an adaptive threshold on the distance vector (e.g. Otsu's method, which finds the cutoff that maximises between-class variance). It's very flexible.

Properties of the EP dictionary

So, now you've built an EP dictionary. What can you do with this object? It's just a partition, but a lot of the things you'd want to do with an SAE turn out to fall out of the partition directly. I did some very simple experiments to characterise them, but there is way more work to be done (see final section of this post).

The results below are meant to give intuition. Specific choices of layer, model, and data corpus vary between experiments and aren't load-bearing for the picture they collectively paint – full configurations for each are in the paper.

Concept detection (AxBench)

The headline result in the paper is that on AxBench latent concept detection at Gemma-2-2B-it L20, EP reaches mean AUROC 0.881 over 500 concepts – +0.126 over the canonical GemmaScope SAE leaderboard entry (0.755) and within 0.030 of SAE-A's 0.911. EP finds meaningful conceptual structure (according to this benchmark) and does it with about 1,000× less compute. Hurrah! (Full per-method breakdown, including the supervised AxBench baselines and EP across all percentile resolutions, is in the paper §4.)

How EP and SAEs relate

So EP and SAEs both find meaningful structure on AxBench. Are they converging on the same features?

Comparing EP regions against the canonical GemmaScope SAE by which tokens fire each feature, the picture is asymmetric. At a moderate EP resolution, roughly one in five EP regions has a strong SAE counterpart. Sharpen EP further and it just keeps splitting into sub-regions the SAE doesn't carry. There's a small shared core and a lot of disagreement on both sides.

What's the disagreement about? The two methods make different geometric commitments. SAEs commit to linear separability – every feature is a learned decoder direction optimised for sparse reconstruction. EP commits to density – every region is a part of activation space, anchored on an observed activation, summarised by how many activations sit in it and how tightly they cluster. The natural reading is that the shared core sits where both commitments hold simultaneously: directions that are density-concentrated and linearly separable. Outside that intersection, EP captures broad content-anchored regions that the SAE splinters into multiple narrow features, and the SAE captures sparse linear directions for which no single contiguous EP region exists at this resolution.

Find and steer refusal

I built an EP dictionary on a mix of harmful and benign prompts on Gemma-2-2B-it, then scored each prompt's generation for refusal and computed the refusal rate per region. One region absorbs most of the refusing prompts; the rest sit near the build-set base rate. Projecting held-out harmful prompts off that region's exemplar direction collapses refusal almost entirely – from a refusal rate of around 0.98 to around 0.02 on the best seed. This is the same range as dedicated refusal-direction work (Arditi et al., 2024). More detail in the paper §4.

The symmetric intervention – adding the exemplar direction to the activation at every position, scaled by some steering strength – is more interesting and more limited. At low steering strength, generations are indistinguishable from the unsteered baseline. Past a threshold, the model starts to refuse, but the refusals are structurally degenerate: it loops on "I" or opens with apologetic prefixes ("I'm so sorry", "I am assuming") that get counted as refusal but lack discourse-coherent refusal content. Push the strength higher still and generations degrade into single-token loops. So the exemplar direction is causally necessary for refusal – remove it and refusal breaks – but not causally sufficient on its own. This is consistent with refusal being a multi-component output behaviour: the exemplar captures the discriminating axis used in the refusal decision, not the full set of components needed to produce a coherent refusal. (Arditi et al.'s refusal direction does produce coherent refusals when added, so the exemplar here is doing less work on the production side than theirs; it's plausibly more like a "this is the kind of prompt where I should be cautious" axis than the full refusal feature).

A free OOD signal

EP dictionaries are ideally built to saturation (the point at which every new activation joins an existing region, rather than seeding a new one), so every typical in-distribution activation has a nearby exemplar, by construction. It turns out that the distance of an activation to the nearest exemplar in the dictionary looks like quite a good measure of distribution shift.

If you build a Pile-trained EP dictionary on Gemma-2-2B-it, then ask "how far from the nearest exemplar is this activation?", you find that random-token activations (mostly OOD in the Pile?) sit measurably further out than Pile activations, with the gap widening as you sharpen the dictionary's resolution. Bulgarian Wikipedia (under-represented in the Pile but not really OOD) activations sit between the two. The gap also shrinks with depth – late-layer processing seems to pull heterogeneous inputs back towards the typical activation manifold.

So while EP covers every input at inference by assignment, distance-to-nearest-exemplar could give you a free, graded OOD score in the bargain.

Cross-checkpoint drift (base ↔ IT)

Because exemplars are real activations, two dictionaries built under the same protocol can be matched via their exemplars: every region in one dictionary has a nearest counterpart in the other, and the distance between them tells you whether they're anchoring the same region, or different ones. This is essentially impossible with SAEs, where features at the same layer of two different checkpoints have no shared coordinate system – you'd be matching learned decoder columns across two independent training runs.

I built dictionaries on the same Pile data for the base and instruction-tuned (IT) versions of Gemma-2-2B and matched them. Only a handful of regions survive as common to both, mostly general-purpose syntactic patterns - instruction tuning seems to substantially reorganise activation space, at least at the layers I tested.

What is actually happening here? We can see more by building EP dictionaries on the same instruction-formatted prompts (a mix of harmful and benign) for both the base and instruction-tuned models, and looking at how each model's activations (both within the prompt, and at the final token of each prompt where the model is poised to respond) are partitioned.

At the final token, the base model treats almost all of these prompts the same way: their activations don't show a harmful/benign separation that I could find. But the instruction-tuned model splits the same final-token activations into several distinct regions with sharp harmful/benign separation – including the refusal-loaded region from the previous section.

Interestingly, the base model isn't completely blind to the distinction – it just doesn't find it at the final token. The region in the base dictionary whose direction most closely matches the IT refusal region is one populated almost entirely by activations from earlier positions in the prompts. While the base model is reading the harmful content, it does represent it as a distinct direction; it just doesn't carry that representation forward to the final-token activation, where the next output gets decided. Instruction tuning takes that pre-existing within-prompt direction and brings it forward to the final token, where it becomes a refusal decision.

Domain saturation

EP doesn't pick dictionary size in advance. It grows until the stream stops producing new regions, then stops. So the saturated size of a dictionary on a given input stream is itself a measurement of that stream's activation geometry.

Building EP dictionaries (at a fixed p) on three different kinds of input – math, code, and chat – on the same model activations gives three different saturated dictionary sizes, and the differences vary across layers. Chat grows monotonically with depth: the model uses more and more distinct regions to handle conversation as it processes it. Code is essentially flat across the network, and code activations cover a smaller area of activation space than chat does, at every layer. Math is non-monotonic, peaking in the middle layers. Roughly a factor of two separates code from chat at every layer. (Plots in the paper §4.)

Inside the partition

An EP partition is a set of regions of activation space, each anchored on a real activation. What does the neighbourhood between two nearby regions look like?

You can take two close regions in the same dictionary, and ask which other regions sit between them: regions whose exemplars are closer to both anchors than the anchors are to each other. How that neighbourhood populates as you sharpen resolution is itself informative:

Partition neighbourhood between two anchor regions across three resolutions of the Gemma-2-2B L12 dictionary.

At coarse resolution (p₁₆, K=83) the neighbourhood is empty: the anchors are close and no region sits between them. Sharpening to p₁₀ (K=203), two regions appear – one anchored on discourse connectives (to, ahead, in, beforehand) and one on CamelCase code identifiers (SourceChecksum, UnsafeEnabled). At p₈ (K=292) the neighbourhood grows to 55 regions; four illustrative ones are shown, including verb forms (intended, supposed, designed; appear, seems, seem; began, started, begun) – the last of which also drags in a code identifier (ArgsConstructor) inside the same Voronoi cell.

This example is only illustrative – there are many more experiments to be done here to understand how neighbouring regions relate.

Future work

There's a lot to do! In no particular order:

Geometric assumptions

EP uses a particular notion of distance (centred cosine) to decide which activations cluster together. That's a hypothesis about activation geometry – other distance metrics would yield different region structures, and a head-to-head against a non-linear distance might be a nice test of the linear representation hypothesis (the idea that meaningful concepts in LLMs correspond to linear directions in activation space).

Seed sensitivity

Which activation gets to exemplify a region depends on streaming order. EP is sensitive to this. In the paper I propose a measure of region stability, which helps. You could also, during clustering, check new candidates against the running mean of each region, and replace a region's exemplar whenever a more representative one is found. Or, if you maintain a member index for each region, you could select a more representative exemplar post-hoc once the dictionary has saturated. You could also do something like merging unstable regions.

What does EP tell us about neural geometry?

EP imposes very little on the activation space. So you can use it to measure things like, how does region size and density vary (and does this tell us anything useful)? Are similar concepts close by (seems like yes)? Which "features" are rarely used and which are common (look at partition member count)?

Comparative work with SAEs

The shared-core / mostly-different result is the obvious entry point. What exactly makes the shared core special? How robust is it to different seeds? What kinds of features live there? Building bigger, finer-grained EP dictionaries while filtering (or merging? Can we identify regions that should be merged?) unstable regions might close the remaining AxBench gap to the strongest SAE baseline.

Steering and patching

I barely scratched the surface here – how good are EP dictionaries at identifying steering vectors? How dependent is this on p? Are some kinds of regions more steerable than others?

Diffing

Tracking changes in the partition across different layers, or across different models, or across different checkpoints of the same model. What do we gain and lose by finetuning? When are features fragmented or unified at different layers? You can probably trace the same prompt through EP dictionaries built at multiple layers of the same model, and get a discrete path through partition cells.

Refusal and behaviour diagnostics

Build an EP dictionary on a deployed model, identify the regions that fire for behaviours you care about (refusal, jailbreaks, persona switches, deception etc.), and you have ready-made handles for monitoring or ablating them at inference.

OOD monitoring at deployment

Distance-to-nearest-exemplar is a free per-input signal that an activation is unlike anything in the build distribution. It costs nothing extra to compute, and could flag inputs (or completions?) that deserve more attention.

Finetuning audits

Build dictionaries on a model before and after finetuning (RLHF, SFT, DPO, an unlearning pass, whatever) and see which regions get preserved, dropped, or introduced. A direct "what did this finetune actually change?" signal, with no need for a labelled probing dataset.

Also, I only tested on gemma-2-2b and the Pile. So, next up is to validate EP on more models, more datasets, etc etc, which I will do next, if nobody beats me to it.

That's all for now

Initial experiments suggest that EP dictionaries expose interpretable and causal regions of activation space, and – because exemplars are observed activations rather than learned decoder columns – are commensurable across layers, models, and training checkpoints in a way that learned-dictionary methods like SAEs aren't. They are fully unsupervised, a lot cheaper to build than comparable SAEs, seem to perform similarly in some ways, but not in others.

If you would like to build on this work, the code is at github.com/jessicarumbelow/exemplar-partitioning, and the paper is here.

Feel free to get in touch if you have any questions or would like to collaborate.