Use this skill when working with sparse autoencoders (SAEs), crosscoders, dictionary learning on neural network activations, training SAEs/crosscoders from scratch, loading pretrained dictionaries, caching model activations, or comparing model internals across fine-tuned model pairs using the dictionary_learning / crosscoder_learning library.
Use this skill when working with sparse autoencoders (SAEs), crosscoders, dictionary learning on neural network activations, training SAEs/crosscoders from scratch, loading pretrained dictionaries, caching model activations, or comparing model internals across fine-tuned model pairs using the dictionary_learning / crosscoder_learning library.
This downloads ~2.5 GB of dictionaries (MLP outputs, attention outputs, residual streams for all layers of Pythia-70m-deduped, trained on 2B tokens from The Pile). Directory structure after download:
Demo Scripts
scripts/train_and_evaluate_sae.py
#!/usr/bin/env python3"""
Train and Evaluate a Sparse Autoencoder (SAE) using dictionary_learning.
This script demonstrates the complete workflow for:
1. Loading a language model with nnsight
2. Creating an ActivationBuffer to stream model activations
3. Training a StandardTrainer SAE using trainSAE
4. Evaluating the trained SAE (MSE, L0, L1, variance explained)
5. Saving to disk and loading back
Requirements:
pip install git+https://github.com/jkminder/dictionary_learning
pip install torch nnsight transformers
Usage:
python train_and_evaluate_sae.py
python train_and_evaluate_sae.py --model EleutherAI/pythia-70m-deduped --layer 1 --device cuda:0
"""import argparse
import json
import os
import torch
from typing import Iterator, List# ─── dictionary_learning imports ────────────────────────────────────────────from dictionary_learning import ActivationBuffer, AutoEncoder
from dictionary_learning.trainers import StandardTrainer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate
# ─── nnsight import ─────────────────────────────────────────────────────────try:
from nnsight import LanguageModel
except ImportError:
raise ImportError(
"nnsight is required. Install it with: pip install nnsight"
)
defmake_toy_data_iterator(num_sentences: int = 500) -> Iterator[str]:
"""
Create a simple iterator of text strings for demonstration.
In real training you would replace this with a proper dataset, e.g.:
from datasets import load_dataset
dataset = load_dataset("Eleuther/pile", split="train", streaming=True)
data = (example["text"] for example in dataset)
Args:
num_sentences: Number of repeated sentences in the toy dataset.
Returns:
An iterator that yields strings.
"""
sentences = [
,
,
,
,
,
,
,
,
]
data = [sentences[i % (sentences)] i (num_sentences)]
(data)
():
()
model = LanguageModel(model_name, device_map=device)
model_name.lower():
submodule = model.gpt_neox.layers[layer].mlp
activation_dim =
model_name.lower():
submodule = model.transformer.h[layer].mlp
activation_dim =
:
ValueError(
)
()
()
model, submodule, activation_dim
() -> ActivationBuffer:
data = make_toy_data_iterator(num_sentences=num_sentences)
buffer = ActivationBuffer(
data=data,
model=model,
submodule=submodule,
d_submodule=activation_dim,
n_ctxs=n_ctxs,
ctx_len=ctx_len,
out_batch_size=batch_size,
device=device,
)
(
)
buffer
() -> AutoEncoder:
dictionary_size = expansion_factor * activation_dim
(
)
trainer_cfg = {
: StandardTrainer,
: AutoEncoder,
: activation_dim,
: dictionary_size,
: lr,
: device,
: ,
}
os.makedirs(save_dir, exist_ok=)
ae = trainSAE(
data=buffer,
trainer_configs=[trainer_cfg],
steps=steps,
warmup_steps=warmup_steps,
resample_steps=resample_steps resample_steps > ,
save_dir=save_dir,
log_steps=,
)
()
ae
() -> :
ae.()
ae = ae.to(device)
()
activations = torch.randn(n_samples, activation_dim, device=device)
torch.no_grad():
reconstruction, features = ae(activations, output_features=)
mse = torch.nn.functional.mse_loss(reconstruction, activations).item()
variance = activations.var().item()
variance_explained = (, - mse / variance) *
l1 = features.().mean().item()
l0 = (features > ).().(dim=-).mean().item()
alive_features = (features > ).(dim=).().mean().item()
metrics = {
: mse,
: variance_explained,
: l1,
: l0,
: alive_features,
: features.shape[-],
: activation_dim,
}
()
k, v metrics.
scripts/train_sae_demo.py
#!/usr/bin/env python3"""
Train a Sparse Autoencoder (SAE) on Language Model Activations
This script demonstrates how to use the dictionary_learning library to:
1. Load a language model via nnsight
2. Create an ActivationBuffer to stream activations from MLP layers
3. Train a standard sparse autoencoder (SAE) using StandardTrainer
4. Evaluate the trained SAE on held-out activations
5. Save the trained dictionary to disk and push to Hugging Face Hub
Requirements:
pip install git+https://github.com/jkminder/dictionary_learning
pip install nnsight torch datasets
Usage:
python train_sae_demo.py
"""import torch
from typing import Iterator
from nnsight import LanguageModel
from dictionary_learning import ActivationBuffer, AutoEncoder
from dictionary_learning.trainers import StandardTrainer
from dictionary_learning.training import trainSAE
from dictionary_learning.evaluation import evaluate
# ---------------------------------------------------------------------------# Configuration# ---------------------------------------------------------------------------
MODEL_NAME = "EleutherAI/pythia-70m-deduped"# Any HuggingFace causal LM works here
DEVICE = "cuda:0"if torch.cuda.is_available() else"cpu"
LAYER_INDEX = 1# Which transformer layer's MLP to train the SAE on
ACTIVATION_DIM = 512# Output dimension of Pythia-70m MLP layers
DICT_SIZE = 16 * ACTIVATION_DIM # 8192 — 16x expansion ratio is common
LEARNING_RATE = 1e-3
L1_PENALTY = 8e-4# Sparsity regularization coefficient
N_CTXS = int()
BATCH_SIZE =
SAVE_PATH =
() -> Iterator[]:
sample_texts = [
,
,
,
,
,
,
,
,
,
,
] *
(sample_texts)
() -> ActivationBuffer:
submodule = model.gpt_neox.layers[LAYER_INDEX].mlp
data = get_training_data()
buffer = ActivationBuffer(
data=data,
model=model,
submodule=submodule,
d_submodule=ACTIVATION_DIM,
n_ctxs=N_CTXS,
device=device,
)
(
)
buffer
() -> :
{
: StandardTrainer,
: AutoEncoder,
: ACTIVATION_DIM,
: DICT_SIZE,
: LEARNING_RATE,
: L1_PENALTY,
: device,
}
() -> AutoEncoder:
buffer = build_activation_buffer(model, device)
trainer_cfg = build_trainer_config(device)
()
()
()
()
()
()
()
()
ae = trainSAE(
data=buffer,
trainer_configs=[trainer_cfg],
)
()
ae
() -> :
()
held_out_activations = torch.randn(, ACTIVATION_DIM).to(device)
torch.no_grad():
features = ae.encode(held_out_activations)
reconstruction = ae.decode(features)
reconstruction_v2, features_v2 = ae(
held_out_activations, output_features=
)
mse = torch.nn.functional.mse_loss(reconstruction, held_out_activations)
l1 = features.().mean()
l0 = (features > ).().mean(dim=-).mean()
pct_alive = (features.().(dim=) > ).().mean()
()
()
()
()
torch.allclose(reconstruction, reconstruction_v2, atol=), \
torch.allclose(features, features_v2, atol=), \
()
() -> :
os
os.makedirs(save_path, exist_ok=)
ae.save(save_path)
()
() -> :
()
"The quick brown fox jumps over the lazy dog."
"Sparse autoencoders learn interpretable features from neural network activations."
"Dictionary learning finds a sparse representation of data."
"Language models encode semantic information in their hidden states."
"Feature decomposition helps us understand what a model has learned."
"Mechanistic interpretability studies the internal computations of neural networks."
"The residual stream aggregates information across transformer layers."
"Attention heads route information between token positions."
# Cycle through sentences to get num_sentences total strings
len
for
in
range
return
iter
def
load_model_and_submodule
model_name: str,
layer: int,
device: str,
"""
Load a LanguageModel with nnsight and select an MLP submodule.
Args:
model_name: HuggingFace model identifier, e.g. "EleutherAI/pythia-70m-deduped".
layer: Which transformer layer's MLP to use.
device: PyTorch device string, e.g. "cuda:0" or "cpu".
Returns:
Tuple of (model, submodule, activation_dim)
"""
# Pythia models use gpt_neox.layers[i].mlp (output dim = 512 for 70m)
# GPT-2 models use transformer.h[i].mlp
# Adapt this section for other model families.
if
"pythia"
in
# Pythia-70m MLP output dim; adjust for larger Pythia variants:
# pythia-160m: 1024, pythia-410m: 2048, pythia-1b: 4096, etc.
512
elif
"gpt2"
in
3072
# GPT-2 small MLP output
else
# Generic fallback — adjust activation_dim to match your model
raise
f"Unknown model family for {model_name}. "
"Please extend load_model_and_submodule() for this model."
print
f" Submodule: {type(submodule).__name__} at layer {layer}"
print
f" Activation dim: {activation_dim}"
return
def
build_activation_buffer
model,
submodule,
activation_dim: int,
device: str,
n_ctxs: int = 3000,
ctx_len: int = 128,
batch_size: int = 64,
num_sentences: int = 500,
"""
Construct an ActivationBuffer that streams activations from the given submodule.
The buffer internally processes text through the model, captures the
submodule's output activations, and yields them in batches. When the
buffer is half-depleted it automatically refreshes with new text.
Args:
model: nnsight LanguageModel instance.
submodule: The model submodule whose outputs to capture.
activation_dim: Output dimension of the submodule.
device: Torch device string.
n_ctxs: Number of contexts (sequences) to buffer at once.
Higher = more memory but more diverse batches.
ctx_len: Maximum sequence length per context.
batch_size: Number of activation vectors per yielded batch.
num_sentences: How many toy sentences to generate (demo only).
Returns:
An initialised ActivationBuffer ready to yield activation batches.
"""
buffer: ActivationBuffer,
activation_dim: int,
expansion_factor: int = 8,
lr: float = 1e-3,
device: str = "cpu",
steps: int = 1000,
warmup_steps: int = 100,
resample_steps: int = 500,
save_dir: str = "./trained_sae",
"""
Train a StandardTrainer SAE on activations from the buffer.
Args:
buffer: ActivationBuffer yielding activation batches.
activation_dim: Input/output dimension of the SAE (= submodule output dim).
expansion_factor: Dictionary size = activation_dim * expansion_factor.
lr: Learning rate for ConstrainedAdam optimizer.
device: Torch device string.
steps: Total number of training gradient steps.
warmup_steps: Number of linear LR warmup steps.
resample_steps: Dead neuron resampling interval (0 to disable).
save_dir: Directory to save the trained SAE.
Returns:
The trained AutoEncoder instance.
"""
"""
Evaluate a trained SAE on random activations (demo) or real activations.
Reports:
- MSE loss (average squared reconstruction error)
- L1 sparsity (sum of absolute feature activations)
- L0 sparsity (average number of active features per sample)
- Fraction of alive dictionary features
In real usage you would pass actual model activations as a tensor or
DataLoader, not random tensors.
Args:
ae: Trained AutoEncoder instance.
activation_dim: Dimension of activations.
n_samples: Number of random samples to evaluate on.
device: Torch device string.
Returns:
Dictionary of metric_name -> value.
"""
eval
# ── For demonstration we use random activations ──────────────────────────
# Replace with: activations = next(iter(your_buffer)) for real evaluation
print
f"\nEvaluating SAE on {n_samples} random activation samples ..."
with
True
max
0.0
1.0
100.0
abs
1e-6
float
sum
1
# Fraction of alive features (active on at least one sample)
1e-6
any
0
float
"mse_loss"
"variance_explained_pct"
"mean_l1_sparsity"
"mean_l0_sparsity"
"fraction_alive_features"
"dictionary_size"
1
"activation_dim"
print
"\n=== SAE Evaluation Metrics ==="
for
in
3e4
# ActivationBuffer capacity (contexts held in memory)
"""
Return an iterator of strings for SAE training.
In production, replace this with a large text corpus such as The Pile,
OpenWebText, or a domain-specific dataset. Here we use a small hardcoded
list to keep the demo self-contained and runnable without network access.
Returns:
Iterator[str]: An iterator that yields text strings.
"""
# Replace this list with:
# from datasets import load_dataset
# data = load_dataset("EleutherAI/pile", split="train", streaming=True)
# return (item["text"] for item in data)
"The quick brown fox jumps over the lazy dog."
"Sparse autoencoders decompose neural network activations into interpretable features."
"Dictionary learning finds a sparse representation of data in an overcomplete basis."
"The residual stream in a transformer carries information between layers."
"Mechanistic interpretability aims to reverse-engineer neural networks."
"Features learned by sparse autoencoders often correspond to human-interpretable concepts."
"The encoder maps activations to a sparse feature vector."
"The decoder reconstructs the original activation from the sparse features."
"Dead neurons are dictionary features that never activate on any input."
"Neuron resampling periodically reinitializes dead neurons during training."
500
# Repeat to simulate a larger dataset for demonstration
"""
Construct an ActivationBuffer for MLP output activations at a given layer.
The ActivationBuffer wraps an nnsight LanguageModel, hooks into a specific
submodule (here: the MLP of layer LAYER_INDEX), and maintains a rolling
buffer of activations sampled from the provided text data.
Args:
model: An nnsight LanguageModel instance.
device: Target device string (e.g., "cuda:0" or "cpu").
Returns:
ActivationBuffer: Configured and ready to yield activation batches.
"""
# Access the MLP submodule in Pythia-70m.
# For other model architectures, adjust the attribute path accordingly:
# GPT-2: model.transformer.h[layer].mlp
# LLaMA: model.model.layers[layer].mlp
# Mistral: model.model.layers[layer].mlp
# Output dimension of the MLP
# How many contexts to keep in the buffer
print
f"[ActivationBuffer] Created buffer for layer {LAYER_INDEX} MLP "
"""
Build the trainer configuration dictionary for a StandardTrainer + AutoEncoder.
The config dict is passed as one element of the trainer_configs list to
trainSAE(). Multiple configs can be passed to sweep hyperparameters.
Args:
device: Target device string.
Returns:
dict: Trainer configuration dictionary.
"""
return
"trainer"
"dict_class"
"activation_dim"
"dict_size"
"lr"
"l1_penalty"
"device"
# Optional: warmup_steps=1000 for linear LR warmup
# Optional: resample_steps=25000 to resample dead neurons periodically
def
train_sae
model: LanguageModel, device: str
"""
Train a sparse autoencoder on MLP activations from the language model.
This function:
1. Creates an ActivationBuffer that streams MLP activations.
2. Calls trainSAE() with a StandardTrainer configuration.
3. Returns the trained AutoEncoder.
Args:
model: An nnsight LanguageModel instance.
device: Target device string.
Returns:
AutoEncoder: The trained sparse autoencoder.
"""
"""
Evaluate a trained AutoEncoder on random held-out activations.
Reports MSE loss, L1 sparsity, L0 sparsity, and fraction of live neurons.
Args:
ae: A trained AutoEncoder instance.
device: Target device string.
"""
print
"\n[Evaluation] Evaluating SAE on random activations..."
# In production, use real held-out activations from the model.
# Here we use random tensors as a placeholder to demonstrate the API.
# Verify both forward call styles give the same result
assert
1e-5
"Mismatch between ae.decode(ae.encode(x)) and ae(x)"
assert
1e-5
"Mismatch in features between two forward call styles"
print
" [OK] Both forward call styles are consistent."
def
save_sae
ae: AutoEncoder, save_path: str
None
"""
Save the trained AutoEncoder weights to a local directory.
The saved directory contains:
- ae.pt: The state_dict of the fully trained dictionary.
- config.json: Hyperparameters used to train the dictionary.
Args:
ae: A trained AutoEncoder instance.
save_path: Local directory path to save the dictionary.
"""
import
True
print
f"\n[Save] Dictionary saved to: {save_path}"
def
demo_load_and_use_pretrained
None
"""
Demonstrate loading a pretrained AutoEncoder from a local path or HF Hub
and performing encode / decode operations.
This function shows the recommended API for inference with a pretrained SAE.
"""
print
"\n[Demo] Loading pretrained AutoEncoder from local path..."
# Replace with your actual path to a pretrained dictionary
