Prompt Categories in the SAE Feature Basis of GPT-2

A language model that answers a math question and a model that answers a creative writing prompt are running the same weights. Somewhere in the forward pass, the difference gets represented. We wanted to know where, and whether that representation is doing any work or is just something an observer can read off the side.

GPT-2 Small has 768 residual dimensions and far more than 768 concepts to represent, so it stores them as overlapping directions. Individual neurons come out polysemantic: one neuron might fire for Python syntax, the word "function", and formal register all at once. Reading meaning off single neurons does not work.

A sparse autoencoder gives you a different basis. It learns an overcomplete dictionary over the activations, 24,576 entries for 768 dimensions, with an L1 penalty that keeps almost all of them at zero for any given input. The expansion plus the sparsity pushes each dictionary entry toward one concept rather than a blend.

We did not train the dictionaries. Joseph Bloom's gpt2-small-res-jb release provides one per layer, loaded through sae_lens, which is what made this fit on a laptop GPU.

The prompts had to be real

An earlier version of this ran on 120 prompts we wrote ourselves. Our instructor pointed out the obvious problem: a hand-written math prompt encodes our idea of what a math question looks like, not what anyone actually types.

So we replaced all of them. 300 prompts across six categories, drawn from Search Arena 24k and WildChat, both real user queries with existing intent labels. WildChat is 27 GB, so we streamed it through the HuggingFace API rather than downloading it. Everything passed five filters before sampling: length bounds, an alphabetic ratio to reject URL soup, a type-token ratio to reject repetition spam, a token ceiling so nothing truncated mid-extraction, and exact-match dedup across both sources.

The categories kept their real-world messiness. Median length runs from 8 words for factual to 40 for emotional, and the corpus is multilingual because the sources are.

Category is encoded early, and sparsely

For each prompt we read the residual stream at the last token position, which under causal attention has seen the whole prompt, then encoded it through the SAE at layers 2, 6 and 10.

At layer 2 the number of active features per category spans more than an order of magnitude: 26 for factual, 34 for math, 344 for code. That ordering is intuitive. Answering "who won the 1998 World Cup" is one retrieval. Answering "why does my React component re-render twice" needs syntax, library semantics, control flow and intent at once.

By layer 10 the spread has collapsed to 157–205 across every category, which is what you would expect if the late layers have stopped representing the prompt and started shaping a next-token distribution.

Every category has features that fire for it and almost nothing else. Two turned out to be shared rather than exclusive: feature #22431 is the top feature for both reasoning and emotional, which suggests it tracks something closer to "this needs deliberation" than to either label.

Then we tried to break it

Selectivity is a correlational claim. A feature that reliably lights up on math prompts has not been shown to do anything. Activation patching is the test: run a clean prompt and a corrupt one, overwrite part of the corrupt run with the clean run's activations, and measure how far the output moves back toward clean.

Causal influence accumulates monotonically with depth, with no reversals anywhere in the twelve-layer sweep. Which sets up the result that actually mattered.

At layer 6, against a no-patch baseline of 3.61 KL: patching the entire 768-dimensional residual stream gets you to 0.64. Patching only the ten most math-selective feature directions gets you to 3.48.

That is 4.5% of the available effect, from the ten features most confident about math.

What that means

The features are real. They are selective, they are reproducible, and they are not enough. Math-relevant computation at layer 6 lives in hundreds of directions, not ten. This is superposition, measured at the level of task category rather than individual tokens.

There is a tension worth sitting with. Layer 2 separates categories best and moves the output least. Layer 11 does the opposite. Where a property is most legible is not where it is most load-bearing, and an interpretability method that only looks for legibility will keep finding the wrong layer.

We also reported the result that did not flatter us. Fisher's discriminant ratio across all 24,576 dimensions comes out around 0.03, and a UMAP projection at layer 6 splits the prompts into two clusters that cut straight across all six categories rather than separating them. Category is recoverable from these activations, but it is not the largest source of variance in them.

Limits

GPT-2 Small is 124M parameters, and concept storage in larger models may be more localized, in which case the distributed result here does not transfer. We read one token position per prompt, so any category signal living mid-prompt is invisible to us. The dictionary is fixed, so features gpt2-small-res-jb never learned cannot appear no matter how real they are. Fifty prompts per category is thin for separating categories that genuinely overlap, like reasoning and factual.

The next thing worth doing is running benign and adversarial prompts against each other to see whether jailbreaks leave an early-layer signature. If they do, it is a cheap runtime filter. The layer-2 result suggests it is worth checking.

Credits

Built with two others for a machine learning course at the University of San Francisco. transformer_lens for hooked forward passes, sae_lens for the pre-trained dictionaries, seed 42 everywhere.