Sparse Autoencoders › Code Implementation
Language AI HandbookAll chapters Sparse Autoencoders › Code ImplementationOne section · Sections in this chapter8 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a minimal SAE and train it on synthetic activation data that simulates the superposition of independent features. This implementation follows the design used in Anthropic's open-source SAE research, simplified for clarity.
Setup and Data Generation
import numpy as np
import torch
# Simulation parameters
d_model = 16 # Dimensionality of model activations (small for illustration)
n_features_true = (
64 # True number of underlying features (much larger than d_model)
)
n_samples = 50000 # Number of activation vectors to generate
feature_prob = 0.05 # Probability each feature is active on a given token
# Generate random "ground truth" feature directions in d_model space
# These simulate the directions that the language model uses to represent features
feature_directions = torch.randn(d_model, n_features_true)
feature_directions = feature_directions / feature_directions.norm(
dim=0, keepdim=True
)
# Generate sparse feature activations
# Each row is a token; each column is a feature (mostly zeros)
feature_acts = torch.zeros(n_samples, n_features_true)
active_mask = torch.rand(n_samples, n_features_true) < feature_prob
feature_acts[active_mask] = (
torch.rand(active_mask.sum()) * 2 + 0.5
) # Positive magnitudes
# Construct "model activations" as sum of active feature directions + noise
activations = feature_acts @ feature_directions.T # (n_samples, d_model)
activations = activations + 0.05 * torch.randn_like(activations)Activation tensor shape: torch.Size([50000, 16]) True features: 64, Model dimension: 16 Average features active per token: 3.20 Sparsity: 95% of features zero per token
We generate 50,000 synthetic activation vectors where each is a superposition of a small number of the 64 true features, compressed into a 16-dimensional space. This simulates the superposition hypothesis at small scale: more features than dimensions, with each token using only a few of those features at once. Each true feature has a probability of 0.05 of being active on any given token, so on average about 3.2 features are active simultaneously (roughly 5% of 64). This is deliberately a challenging setting: the 64 features are packed into 16 dimensions, and the interference between simultaneously active features creates overlapping activation patterns.
SAE Model Definition
import torch
import torch.nn as nn
class SparseAutoencoder(nn.Module):
def __init__(self, d_model: int, d_sae: int):
"""
d_model: dimension of model activations (input/output)
d_sae: dictionary size (> d_model for overcomplete dictionary)
"""
super().__init__()
self.d_model = d_model
self.d_sae = d_sae
# Pre-encoder bias to center activations
self.b_pre = nn.Parameter(torch.zeros(d_model))
# Encoder: maps from model space to sparse feature space
self.W_enc = nn.Parameter(torch.randn(d_model, d_sae) * 0.01)
self.b_enc = nn.Parameter(torch.zeros(d_sae))
# Decoder: maps from sparse feature space back to model space
self.W_dec = nn.Parameter(torch.randn(d_sae, d_model) * 0.01)
def encode(self, x: torch.Tensor) -> torch.Tensor:
"""Compute sparse feature activations via ReLU encoder."""
x_centered = x - self.b_pre
pre_activation = x_centered @ self.W_enc + self.b_enc
return torch.relu(pre_activation)
def decode(self, z: torch.Tensor) -> torch.Tensor:
"""Reconstruct model activations from sparse features."""
return z @ self.W_dec + self.b_pre
def forward(self, x: torch.Tensor):
z = self.encode(x)
x_hat = self.decode(z)
return x_hat, z
def normalize_decoder(self):
"""Enforce unit norm on decoder columns (feature dictionary vectors)."""
with torch.no_grad():
norms = self.W_dec.norm(dim=1, keepdim=True).clamp(min=1.0)
self.W_dec.data /= normsSAE architecture: 16d -> 128d -> 16d Expansion factor: 8x Total parameters: 4,240 W_enc: torch.Size([16, 128]) = 2,048 params W_dec: torch.Size([128, 16]) = 2,048 params b_enc: 128 params, b_pre: 16 params
The expansion factor of 8x gives us 128 dictionary features to explain patterns in 16-dimensional activation space. Notice that the encoder and decoder are separate matrices (not tied weights), which allows the encoder to learn a different geometry for feature detection than the decoder uses for reconstruction. The encoder rows are optimized to respond strongly to their target feature direction and weakly to others, while the decoder columns are optimized to reconstruct activations accurately when linearly combined. In real language models, expansion factors of 4x to 64x are common; larger factors trade computational cost for the ability to represent more fine-grained features.
Training Loop
import torch.optim as optim
def train_sae(
sae: SparseAutoencoder,
activations: torch.Tensor,
lambda_l1: float = 0.001,
n_epochs: int = 20,
batch_size: int = 512,
lr: float = 2e-4,
):
optimizer = optim.Adam(sae.parameters(), lr=lr)
n_batches = len(activations) // batch_size
history = {"total_loss": [], "recon_loss": [], "l1_loss": [], "mean_l0": []}
for epoch in range(n_epochs):
# Shuffle data each epoch
perm = torch.randperm(len(activations))
epoch_total, epoch_recon, epoch_l1, epoch_l0 = 0.0, 0.0, 0.0, 0.0
for i in range(n_batches):
batch = activations[perm[i * batch_size : (i + 1) * batch_size]]
optimizer.zero_grad()
x_hat, z = sae(batch)
# Reconstruction loss: mean squared error per sample
recon_loss = ((batch - x_hat) ** 2).sum(dim=-1).mean()
# L1 sparsity penalty on feature activations
l1_loss = z.abs().sum(dim=-1).mean()
loss = recon_loss + lambda_l1 * l1_loss
loss.backward()
optimizer.step()
# Normalize decoder after each step to maintain unit-norm columns
sae.normalize_decoder()
# Track effective L0: count of non-zero features per token
l0 = (z > 0).float().sum(dim=-1).mean().item()
epoch_total += loss.item()
epoch_recon += recon_loss.item()
epoch_l1 += l1_loss.item()
epoch_l0 += l0
history["total_loss"].append(epoch_total / n_batches)
history["recon_loss"].append(epoch_recon / n_batches)
history["l1_loss"].append(epoch_l1 / n_batches)
history["mean_l0"].append(epoch_l0 / n_batches)
return history
history = train_sae(
sae, activations, lambda_l1=0.01, n_epochs=30, batch_size=512
)Final reconstruction loss: 0.0421 Final mean L0 (active features): 69.00 Sparsity: 46.1% of features zero per token Variance explained (R^2): 0.9948
The SAE learns to reconstruct activation vectors with high fidelity, but this minimal setup does not yet achieve a sparse code. The score above 0.95 indicates that the learned dictionary captures most of the activation-space structure, while the mean remains well above the desirable range of 5-15 active features out of 128. This combination is a useful warning: excellent reconstruction alone does not guarantee interpretable sparsity. The 128-feature dictionary is larger than the 64 true features in the synthetic data, giving the SAE extra capacity that the penalty must constrain.
The following figure shows training progress across epochs, tracking reconstruction loss and mean effective simultaneously. In a well-trained SAE, both curves should decrease and then plateau, indicating convergence.
Evaluating Feature Sparsity and Dictionary Utilization
After training, it is important to check that the learned features are used. A common failure mode is feature collapse: some features become very popular (active for most tokens) while others are never activated at all. The distribution of feature activation frequencies tells you how well the dictionary is utilized.
# Compute activation frequency for each feature across the dataset
with torch.no_grad():
all_z = []
for i in range(0, len(activations), 512):
batch = activations[i : i + 512]
_, z_batch = sae(batch)
all_z.append(z_batch)
all_z = torch.cat(all_z, dim=0) # (n_samples, d_sae)
# Frequency: fraction of tokens where each feature activates at all
feature_freq = (all_z > 0).float().mean(dim=0) # (d_sae,)
dead_features = (feature_freq == 0).sum().item()
active_features = (feature_freq > 0).sum().item()Total dictionary features: 128 Active features (fire on at least one token): 128 Dead features (never activate): 0 Median activation frequency: 0.5051 Max activation frequency: 0.8601 Features active on >10% of tokens: 128
Dead features represent wasted capacity: they contribute nothing to reconstruction and have not found a useful region of activation space. In practice, dead features are addressed with techniques like "neuron resampling" (reinitializing dead feature vectors toward high-loss data points) or simply using a larger dictionary. Features active on more than 10% of tokens are "ultra-common" features and often correspond to very general properties (e.g., "token is part of English text" rather than anything semantically specific). The ideal feature frequency distribution is a heavy-tailed distribution where most features fire rarely (specific concepts) and a few fire moderately often (general properties).
The histogram below shows the distribution of feature activation frequencies across the dictionary. A healthy SAE has most features firing on a small fraction of tokens, with a long tail of very specific features.
