Sparse autoencoders — pulling interpretable features out of a trained network
In February 2026, Goodfire published a write-up describing what it took to do interpretability research on Kimi K2 Thinking, a trillion-parameter mixture-of-experts model. The headline number was striking: in a single overnight run, they harvested 3 billion activations — roughly half the token count of English Wikipedia — and ended up with about 33 TB of compressed tensors on disk. To make that practical they forked the SGLang inference server and added hooks for shipping activations off the GPU without strangling throughput, threading the needle around continuous batching, chunked prefill, and 8-way tensor parallelism.
The obvious question is: 3 billion activations of what, exactly, and why?
The short answer is that they were collecting training data for a sparse autoencoder — the central interpretability tool of the last few years. An SAE is the thing that takes a model’s hidden state, which from the outside looks like an opaque vector of floats, and turns it into a list of recognizable features: “this is a loop”, “this is a vertical stroke”, “this is a politeness signal”, “this is a Python function definition”. Once you have that list, you can read what the model is thinking about at any point, and — more surprisingly — you can edit it.
This article is about the SAE itself. We will not deploy SGLang or harvest billions of tokens. Instead we will use the 3-digit MNIST CNN from the CNN article (digits 0, 1, 2) as a miniature lab. Everything Goodfire built at trillion-parameter scale — the hook, the harvested dataset, the encoder/decoder, the sparsity penalty, the steering intervention — has a direct counterpart at this scale. The point is that the primitives are small. What is hard at frontier scale is the engineering around them, which we will return to at the end.
Four important concepts
Let’s distinguish four terms used in this article:
- Monosemanticity: one neuron responds primarily to one interpretable feature.
- Polysemanticity: one neuron responds to several seemingly unrelated features.
- Redundancy: multiple neurons encode similar information. If one is removed or becomes inactive, others may compensate.
- Superposition: more features are encoded than there are dimensions, using overlapping directions in the representation space.
Monosemanticity and polysemanticity describe an individual neuron’s responses; redundancy and superposition describe how information is represented across neurons or coordinates. These properties can coexist: a network can have both redundant and polysemantic neurons.
The problem SAEs solve
When you train a neural network end-to-end with backpropagation, nothing in the task loss requires each hidden neuron to correspond to a single human-interpretable concept. Training rewards useful predictions, so a neuron can respond to several seemingly unrelated features. Superposition is one mechanism that can produce this polysemanticity, but observing a polysemantic neuron alone does not establish that the representation contains more features than dimensions.
It’s easy to see polysemanticity in our 3-digit MNIST CNN. Recall its rough shape from the CNN article — a couple of conv+pool stages, a flatten, and a small dense layer of, say, 64 units before the 3-class softmax:
model = keras.Sequential([
keras.layers.Conv2D(16, 3, activation="relu", input_shape=(28, 28, 1)),
keras.layers.MaxPool2D(2),
keras.layers.Conv2D(32, 3, activation="relu"),
keras.layers.MaxPool2D(2),
keras.layers.Flatten(),
keras.layers.Dense(64, activation="relu", name="features"),
keras.layers.Dense(3, activation="softmax", name="logits"),
])The features layer produces a 64-dimensional vector — call it — that the classifier then reads. We can ask: what does a single coordinate of — say — respond to?
The standard way to look is to run the whole training set through the model, record the value of for each input, and look at the top-k images that activated this neuron most strongly. With a small network like ours, that pile of top images is rarely clean. You’ll see things like: some 0s with a thick loop, but also a few 2s with a wide curve at the bottom, and maybe one or two 1s drawn with an unusual slant. Neuron 17 isn’t a “loop detector” or a “curve detector” — it’s a mixture detector that responds to several patterns the model has found convenient to bundle together. That bundling is fine for the classifier, which just needs the final softmax to be right, but it’s a disaster for someone trying to read out what the model represents.
This is the central interpretability problem. The model has clearly learned something useful — it gets 99% on the test set — but the basis it uses to store its knowledge is not aligned with concepts we recognize. Concepts live somewhere in , but not on the coordinate axes.
The SAE’s job is to find the basis that is aligned with concepts, and to do so without any labels — purely from the geometry of the activations.
What a sparse autoencoder actually is
An autoencoder, in the most general sense, is a pair of networks: an encoder that maps an input vector to some internal code, and a decoder that maps the code back. The whole thing is trained to reconstruct its input. The classical autoencoder squeezes the input through a bottleneck smaller than the input — that’s the compressive variety used for denoising, dimensionality reduction, and the like.
A sparse autoencoder inverts this. The internal code is wider than the input, often by a factor of 8x or 16x or more — but only a small handful of code entries are allowed to be nonzero for any given input. So instead of compressing, the SAE expands the activation into an over-complete dictionary and forces a sparse selection from that dictionary to reconstruct the original.
Concretely, for our MNIST CNN, the SAE takes the 64-dimensional features activation and maps it to a much wider feature vector — let’s say :
The encoder is one linear layer plus a ReLU. The decoder is one linear layer (no nonlinearity). has 512 columns, one per feature, each a vector in — these columns are the dictionary atoms, the directions the SAE has learned. To reconstruct , the decoder takes a weighted sum of those columns where the weights are the active entries of .
The objective has two terms:
The reconstruction term says: the decoder must put back together accurately. The sparsity term — an L1 penalty on the feature vector — pushes most entries of to be exactly zero. The interplay is the whole point: with no sparsity penalty, the easiest solution is to set (or its left-inverse) and use every feature a little bit, which gains nothing over the original basis. With sparsity, the SAE is forced to find a small set of directions that, when activated, reconstruct — and the cheapest way to do that is to make each direction correspond to a recurring pattern in the activations. A pattern is, by definition, a concept.
There are a few variants worth knowing about:
- L1 SAEs (the classic) use the L1 penalty above. Easy to train, but the L1 term biases feature magnitudes downward, which complicates downstream analysis.
- TopK SAEs replace the L1 penalty with a hard constraint: keep only the largest entries of and zero out the rest before the decoder. No magnitude bias, but must be chosen.
- JumpReLU uses a learnable threshold per feature, getting some of the benefits of both.
For our purposes the L1 version is fine — it’s the simplest to write down, and the conceptual picture is the same across variants.
The expansion factor and the target sparsity (how many features should fire per input) are the two main knobs. For our 64-dimensional MNIST activation, a 512-feature dictionary with ~8 active features per input is a reasonable starting point. For Anthropic’s and Goodfire’s work on frontier models, expansion factors of 16x–64x and active counts in the dozens are typical, which is part of why the training data appetite is so large.
Activation harvesting: the universal primitive
Before we can train the SAE we need a dataset of activations from the target model. This is the step Goodfire spent most of their engineering effort on, because at trillion-parameter scale every part of it is hostile to throughput. At our scale it’s three lines of Keras.
The basic primitive is the same at every scale: attach a forward hook to the target layer, run inputs through the model, and copy the layer’s output into a buffer. In Keras the cleanest way is to build a sub-model that ends at the layer of interest:
# Build a sub-model that returns the `features` layer activation, not the softmax.
feature_extractor = keras.Model(
inputs=model.input,
outputs=model.get_layer("features").output,
)
# Run the training set through it, in batches, and stack the results.
all_activations = feature_extractor.predict(X_train, batch_size=256)
# all_activations.shape == (N, 64)That’s the harvest. The output is an tensor — one row per input image, 64 activation values per row — and that tensor is what we’ll train the SAE on. For our 3-digit subset of MNIST with ~18,000 training images, this is about 4.5 MB of floats. We can hold it in RAM and iterate over it as many times as we want.
In PyTorch the same primitive is spelled differently — you register a forward hook on the target module — but the semantics are identical:
captured = []
def hook(module, inputs, output):
# Detach so we don't drag autograd state along, and move to CPU.
captured.append(output.detach().cpu())
handle = model.features.register_forward_hook(hook)
with torch.no_grad():
for x_batch in loader:
_ = model(x_batch) # forward pass — the hook fires
handle.remove()
all_activations = torch.cat(captured, dim=0) # (N, 64)Goodfire’s pipeline is the same pattern, just under enormous pressure. Kimi K2 Thinking is sharded across 8 GPUs with tensor parallelism, so a naive hook on every rank gives you partial activations that need to be gathered, or — if you hook after the synchronization point — duplicated data across all 8 GPUs. SGLang uses continuous batching, so a batch during prefill isn’t a clean tensor — it’s a flattened tensor with tokens from many unrelated sequences mashed together, and you have to track provenance so you know which activation came from which token of which prompt. Long inputs are split by chunked prefill, so a single document’s activations arrive across multiple captures and have to be stitched back together. None of this changes what an activation is. It just changes how much engineering it takes to copy one off the GPU without killing throughput.
The shape of their solution is worth describing because it’s a useful general pattern. The forward pass itself does the cheapest possible thing: copy the activation into a GPU-side staging buffer preallocated from a pool, enqueue a small metadata record (a tensor handle, provenance, and a CUDA event) into a ring buffer, and continue. A separate sidecar thread waits on the CUDA event, copies the tensor into pinned CPU memory once the GPU work is done, and a writer thread takes pinned-memory tensors and compresses them to disk, reconstructing per-sequence records along the way. Three stages, each running at its own pace, decoupled by buffers. The forward pass never blocks on disk I/O. We don’t need any of this for MNIST, but the architecture is good to know — it’s how you push a hook to its throughput limit on any system where forward-pass speed matters.
Training the SAE
Once we have the activation tensor, training the SAE is a small, self-contained optimization problem. Here’s the loop in Keras, written transparently rather than idiomatically — every step in the loss is visible:
import numpy as np
import tensorflow as tf
from tensorflow import keras
D_INPUT = 64 # width of the layer we harvested
D_FEATURES = 512 # dictionary size — expansion factor of 8x
L1 = 1e-3 # sparsity coefficient
class SAE(keras.Model):
def __init__(self):
super().__init__()
self.encoder = keras.layers.Dense(D_FEATURES, activation="relu")
self.decoder = keras.layers.Dense(D_INPUT, activation=None,
use_bias=True)
def call(self, h):
f = self.encoder(h) # (B, D_FEATURES) — sparse code
h_hat = self.decoder(f) # (B, D_INPUT) — reconstruction
return h_hat, f
sae = SAE()
opt = keras.optimizers.Adam(1e-3)
@tf.function
def train_step(h_batch):
with tf.GradientTape() as tape:
h_hat, f = sae(h_batch)
recon_loss = tf.reduce_mean(tf.reduce_sum((h_batch - h_hat) ** 2, axis=1))
l1_loss = tf.reduce_mean(tf.reduce_sum(tf.abs(f), axis=1))
loss = recon_loss + L1 * l1_loss
grads = tape.gradient(loss, sae.trainable_variables)
opt.apply_gradients(zip(grads, sae.trainable_variables))
return loss, recon_loss, l1_loss
# Iterate over the harvested activations, many epochs.
ds = tf.data.Dataset.from_tensor_slices(all_activations).shuffle(10_000).batch(256)
for epoch in range(50):
for h_batch in ds:
train_step(h_batch)Two things worth noticing. First, the CNN itself is frozen throughout SAE training — we never touch its weights. The SAE is a parasite hanging off the side of the model, fed by the activation tensor we harvested once. Second, the L1 coefficient is the dominant hyperparameter. Too small and the SAE happily uses dozens of features per input and loses interpretability. Too large and the SAE refuses to use enough features to reconstruct , and the reconstruction term blows up. A useful diagnostic during training is the L0 — the average number of features with nonzero activation per input. For our setup we’d target something like L0 ≈ 8 with reconstruction error a small fraction of .
One pitfall worth flagging: SAEs are prone to dead features — dictionary atoms that never activate for any input and just sit there. Production setups (Anthropic’s, Goodfire’s) use various tricks to resuscitate them, including periodic re-initialization of dead features along the direction of the largest reconstruction error. For the toy case it’s enough to know the failure mode exists.
Reading the dictionary
Once the SAE is trained, the dictionary has 512 columns and each column is a direction in the original activation space. We want to know what each column means. The standard procedure has two steps, and neither involves labels.
Step 1: For each feature, find the inputs that activate it most. Run the entire training set through the CNN, then through the SAE encoder, and for each of the 512 features collect the top-k images whose code entry for that feature was largest. This gives you, for each feature, a small gallery of inputs.
# all_activations: (N, 64) — harvested CNN activations
# all_codes: (N, 512) — SAE feature activations
_, all_codes_t = sae(all_activations)
all_codes = all_codes_t.numpy()
# For each feature, the indices of the top 16 most-activating inputs.
top_k = 16
top_indices = np.argsort(-all_codes, axis=0)[:top_k] # (top_k, 512)
# For feature f, the gallery is X_train[top_indices[:, f]].Step 2: Look at the gallery. For our 3-digit MNIST CNN, with a well-trained SAE, you’d typically see features fall into recognizable buckets:
- A handful of features whose galleries are nearly all 0s, all 1s, or all 2s. These are the class-aligned features — directions the model genuinely uses to discriminate, now isolated as single dictionary atoms.
- Features whose galleries are mixed across classes but share a stroke-level pattern — wide-bottom curves (appearing in 2s and the occasional 0), strong vertical strokes (1s and the spine of some 2s), closed loops (mostly 0s).
- A handful of features whose galleries don’t look like much. These are usually dead or near-dead features the model never recruited, or noise features that pick up on idiosyncrasies of the training set.
This is the payoff: the 64-dimensional activation, which on its native axes was polysemantic and unreadable, has been rewritten as a sparse combination of recognizable patterns. We didn’t tell the SAE what to look for. The sparsity penalty forced it to find directions that recur across many inputs, and those recurring directions happen to be the same patterns that humans recognize — because both we and the network are looking at the same dataset and the same task structure.
A few caveats. The interpretability of features depends on the quality of the model and the layer chosen. Early conv layers tend to give edge-and-texture features; the pre-classifier dense layer tends to give class-and-part features; the very last layer before softmax tends to give nearly-class-aligned features that are less interesting. There’s also no guarantee a feature is monosemantic just because its top-16 gallery looks clean — the gallery is a sample, not a proof. Diagnostic tools like feature ablation help separate clean features from coincidental ones, but the field still works largely by visual inspection.
Steering: editing the dictionary
The reason features are interesting is not just that they describe what the model is doing — it’s that they let us change what the model does, surgically, without retraining anything.
The trick is straightforward. We have the SAE encoder and decoder . If we pick a feature index , set to some larger value, and recompute , we have a modified activation vector. If we substitute that vector back into the CNN — replace the features layer output with — and run the classifier on top, we have just steered the model along the direction of feature .
In code, this is most easily done by splitting the CNN at the steering point:
# Two halves of the original CNN, split at the `features` layer.
feature_extractor = keras.Model(model.input, model.get_layer("features").output)
classifier_head = keras.Model(model.get_layer("features").output, model.output)
def steer(x, feature_idx, strength=5.0):
h = feature_extractor(x[None, ...]) # (1, 64)
_, f = sae(h) # (1, 512)
f_steered = f.numpy().copy()
f_steered[0, feature_idx] += strength # nudge the chosen feature
h_steered = sae.decoder(f_steered) # (1, 64) — new activation
return classifier_head(h_steered) # (1, 3) — new probabilitiesIf feature 42 is one of our “vertical-stroke / 1-ish” features, calling steer(image_of_0, feature_idx=42, strength=5.0) will measurably push the classifier’s probability mass toward class 1, even though we never told the classifier anything about 1s — we just added more of a direction that the SAE discovered is correlated with being a 1, and the classifier reacted accordingly. With a large enough strength, the classifier flips. With a moderate strength, the prediction stays correct but with reduced confidence.
This is the smallest-possible version of the result Goodfire described as “real-time, mid-chain-of-thought steering”. In their case the target model is autoregressive, the activation being steered is the residual stream at some transformer layer, and the steering is applied at every decode step as the model is generating — so a feature like “be more cautious about this answer” or “talk less like a coding assistant” can be turned up or down in the middle of a response. The geometry is exactly what we just did: pick a feature, add it to the activation, let the rest of the model continue from the modified state.
Two things make this important rather than a curiosity. First, it shows that the features the SAE finds are causal, not just descriptive — pushing them changes behavior in predictable ways. A feature that only correlated with a behavior, without being part of how the model computed it, would not steer. Second, steering is cheap: no fine-tuning, no RL, no labeled data. The SAE is trained once. After that, every feature in the dictionary is a free behavioral knob.
Why this is hard at frontier scale
Everything we did fits in a few hundred lines of Python on a laptop. The size jump to a trillion-parameter model is what Goodfire’s blog is really about. It’s worth understanding why each part gets harder, because the difficulty isn’t conceptual — it’s operational.
Data volume. Our 64-dim activation tensor for 18,000 inputs is ~4.5 MB. Kimi K2 Thinking’s residual stream is, depending on the layer, several thousand dimensions wide, and SAEs at that scale want billions of tokens of training data to find clean features. 3B tokens × a few thousand floats × 2 bytes (after quantization) lands at ~33 TB compressed. You can’t hold this in RAM. You can’t even comfortably hold it on local SSD.
Throughput. At a hypothetical 100 tokens/sec of harvest throughput, 3B tokens is a year of GPU time. To make it overnight you need tens of thousands of tokens per second. That’s the territory of production inference servers — SGLang, vLLM — which means you have to add your hook to a system that has been pessimized around the assumption that nobody is hooking it.
Inference-server complications. Continuous batching mixes unrelated sequences into one tensor and requires per-token provenance. Chunked prefill splits documents across multiple forward passes. Tensor parallelism shards weights across GPUs and forces you to choose where in the layer’s pipeline to hook so you get complete activations without redundant gathers. CUDA graphs pre-record execution and don’t love conditional hooks. Each of these is solvable; together they push you toward what Goodfire calls a surgical fork — small, targeted patches rather than a rewrite, because the inference server is moving fast and you don’t want to inherit a permanent maintenance burden.
Storage and I/O. Once activations are in pinned CPU memory you have to write them out fast enough to keep up with the forward pass. That means asynchronous, multi-stage pipelines — staging buffers, ring buffers, sidecar threads — so disk I/O latency is decoupled from inference. The mental model is the same as a well-engineered logging pipeline: producers can’t block, consumers run at their own pace, everything is sized to absorb the worst-case mismatch.
None of this changes the SAE itself. The encoder is still one linear layer, the decoder is still one linear layer, the loss is still reconstruction plus sparsity. What changes is the dataset preparation, and that’s where the engineering goes. Goodfire’s write-up is the best public account of what that looks like in detail.
Why this matters
The SAE story is interesting on its own — but the broader point is that interpretability has shifted from “stare at attention weights and hope” to a tractable optimization problem with usable outputs. You can run an SAE on any layer of any model you have weights and activations for, and you can use the resulting dictionary to read, label, and steer behavior. That moves interpretability from “we examined this small model and these are some observations” toward “we have a feature dictionary for the production model and here are the knobs”.
The remaining gap is infrastructure. We harvested 18,000 activations in a single predict call. Goodfire harvested 3 billion in an overnight run by carefully rebuilding the harvest primitive on top of a production inference server. The conceptual machinery is small. The engineering is what makes it land.
If you want to play with this yourself, the workflow is exactly the one above:
- Train (or download) a model. The MNIST CNN from our CNN article works.
- Pick a layer and harvest its activations on a representative dataset. A Keras sub-model or a PyTorch forward hook is enough.
- Train an SAE on those activations with an expansion factor of 8x–16x and an L1 coefficient tuned to give a small target L0.
- Visualize each feature’s top-k activating inputs. Find the recognizable ones.
- Steer by modifying the SAE code vector and running the classifier on the reconstructed activation. Watch the prediction shift.
That’s the whole loop, and it scales — with engineering — all the way up to Kimi K2 Thinking.