Info
Last Execution: 2026-07-27
| Package | Version |
|---|---|
| nnsight | 0.8 |
| dictionary_learning | 0.1 |
| circuitsvis | 1.43 |
Dictionary Learning¶
📗 You can find an interactive Colab version of this tutorial here.
Polysemanticity¶
Mechanistic interpretability tries to understand a network in terms of its individual components. That gets hard when a single neuron responds to several unrelated inputs — a phenomenon called polysemanticity. One neuron might fire for images of car tires and for pictures of rubber ducks.
Packing many features into a few neurons helps the network use its parameters efficiently, but it makes those neurons hard for a human to read. In this tutorial we use a sparse autoencoder (SAE) to pull the tangled features back apart into directions that each mean one thing. For the underlying ideas, see Anthropic's Towards Monosemanticity and Scaling Monosemanticity.
Sparse Autoencoders¶
A sparse autoencoder learns to reconstruct a model's activations while forcing that reconstruction to route through a much larger, mostly-inactive hidden layer. Because only a handful of hidden units may be active at once, each one is pushed toward a single, interpretable feature. This is a form of dictionary learning: the SAE's hidden units form a dictionary of directions, and any activation is explained as a sparse combination of them. Those directions tend to be far more monosemantic than the model's raw neurons.
We won't train an SAE here — we load a pretrained one and use nnsight to read the activations it explains. To edit a model along SAE directions instead, see Setting Activations, and to wire an SAE permanently into a model see Model Editing.
📚 This tutorial is adapted from work by Samuel Marks and Aaron Mueller (see their GitHub repository and Alignment Forum post), which follows Anthropic's approach detailed here.
Setup¶
If you are on Colab, install the libraries this tutorial uses — nnsight, the
dictionary_learning SAE library, and circuitsvis for the visualizations:
from IPython.display import clear_output
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight dictionary_learning circuitsvis
clear_output()
We pull a pretrained autoencoder straight from the Hugging Face Hub. The
saprmarks/pythia-70m-deduped-saes repository ships the trained dictionaries as a zip;
we download it, unpack it, and load the SAE for the layer-0 MLP output of
pythia-70m-deduped. Each SAE is tied to one specific model, layer, and submodule —
this one only explains the activations at that exact spot.
import os
import zipfile
import torch
from huggingface_hub import hf_hub_download
from dictionary_learning.dictionary import AutoEncoder
# Download and unpack the pretrained dictionaries
zip_path = hf_hub_download(
"saprmarks/pythia-70m-deduped-saes",
"dictionaries_pythia-70m-deduped_10.zip",
)
extract_dir = os.path.join(os.path.dirname(zip_path), "extracted")
with zipfile.ZipFile(zip_path) as z:
z.extractall(extract_dir)
# The SAE trained on the layer-0 MLP output
ae_path = os.path.join(
extract_dir,
"dictionaries/pythia-70m-deduped/mlp_out_layer0/10_32768/ae.pt",
)
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
The SAE maps a 512-dimensional MLP activation into a 32768-dimensional feature space
(a 64× expansion) and back. from_pretrained reads those dimensions off the saved
weights for us:
ae = AutoEncoder.from_pretrained(ae_path)
print("activation dim:", ae.activation_dim)
print("dictionary size:", ae.dict_size)
activation dim: 512 dictionary size: 32768
Applying the SAE¶
We load pythia-70m-deduped with TransformersModel. device_map="auto" lets
HuggingFace place the weights on the GPU when one is available, and dispatch=True
loads them right away rather than on the first trace. We move the SAE onto the same
device so its matmuls line up with the activations we capture.
from nnsight import TransformersModel
model = TransformersModel("EleutherAI/pythia-70m-deduped", device_map="auto", dispatch=True)
tokenizer = model.tokenizer
# Put the SAE where the model's activations live
device = next(model._module.parameters()).device
ae.to(device)
AutoEncoder( (encoder): Linear(in_features=512, out_features=32768, bias=True) (decoder): Linear(in_features=32768, out_features=512, bias=False) )
Now we run a prompt through the model and grab the layer-0 MLP output with nnsight.
Inside model.trace(...), reading .output on a module hands us the exact tensor it
produced during the forward pass; .save() keeps it alive after the block exits so we
can work with it. Here that activation is a (batch, sequence, 512) tensor.
prompt = """
Call me Ishmael. Some years ago--never mind how long precisely--having little or no money in my purse, and nothing particular to interest me on shore, I thought I would sail about a little and see the watery part of the world.
"""
# Extract the layer-0 MLP output from the base model
with model.trace(prompt):
mlp_0 = model.gpt_neox.layers[0].mlp.output.save()
# Encode the activation into the SAE's sparse feature space
features = ae.encode(mlp_0)
print("MLP activation:", tuple(mlp_0.shape))
print("SAE features: ", tuple(features.shape))
MLP activation: (1, 54, 512) SAE features: (1, 54, 32768)
Reading activations with .save()
.output is only meaningful inside a trace — it's a handle to a value produced
during the forward pass. Calling .save() on it keeps the concrete tensor available
after the with block ends; without it, the value is freed when the trace completes.
See Setting Activations for the full mechanics of
reading and writing module outputs.
Each of the 32768 features is a candidate monosemantic direction. To find the ones this prompt actually excites, we sum each feature's absolute activation across all token positions and take the top 20.
# Rank features by total activation across the prompt
summed_activations = features.abs().sum(dim=1)
top_activations_indices = summed_activations.topk(20).indices
# Gather the per-token activation of each top feature
compounded = torch.stack(
[features[:, :, i.item()].cpu()[0] for i in top_activations_indices[0]],
dim=0,
)
Visualization¶
With the autoencoder¶
Let's look at what each of the top 20 features responds to across the prompt. In the visualization below, each column is one feature and each token is shaded by how strongly that feature fires on it. Notice how sharply each feature keys off a single token — that specificity is exactly what makes SAE features interpretable.
from circuitsvis.tokens import colored_tokens_multi
tokens = tokenizer.encode(prompt)
str_tokens = [tokenizer.decode(t) for t in tokens]
# Visualize the top 20 features
colored_tokens_multi(str_tokens, compounded.T)
Without the autoencoder (comparison)¶
For contrast, we run the same ranking on the model's raw MLP neurons instead of the SAE features. These top neurons light up (or go negative) across many unrelated tokens — the polysemanticity we set out to untangle — which makes any single neuron much harder to assign a clean meaning to.
# Rank the raw MLP neurons the same way
summed_activations_or = mlp_0.abs().sum(dim=1)
top_activations_indices_or = summed_activations_or.topk(20).indices
compounded_orig = torch.stack(
[mlp_0[:, :, i.item()].cpu()[0] for i in top_activations_indices_or[0]],
dim=0,
)
# Visualize the top 20 raw neurons
colored_tokens_multi(str_tokens, compounded_orig.T)
Where to go next¶
- Setting Activations — edit activations, including steering a model along an SAE feature direction.
- Model Editing — attach an SAE (or other custom module) so it runs as part of the model on every forward pass.
- Attribution Patching and the other tutorials for more ways to use these activations.