Towards Monosemanticity — Visual Experiments¶

Goal: Train sparse autoencoders on superposed activations and reproduce the paper's findings: the sparsity trade-off, feature splitting, monosemanticity, and dead features.

What we'll explore¶

#	Experiment	Question answered
1	Sparsity trade-off	How does the L1 coefficient λ trade reconstruction against sparsity?
2	Feature splitting	Does a bigger dictionary split coarse features into finer ones?
3	Monosemanticity	Are SAE features more single-concept than neurons?
4	Dead features	Do some features never fire, and does resampling revive them?

---

Linked theory: research/02-towards-monosemanticity.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import warnings
warnings.filterwarnings('ignore')

torch.manual_seed(0)
np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11,
    'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35,
})

d, n_true = 16, 40                       # 40 true features superposed in 16 dims
F = torch.randn(d, n_true); F /= F.norm(dim=0, keepdim=True)


def gen_acts(batch, p=0.06):
    x = (torch.rand(batch, n_true) < p).float() * torch.rand(batch, n_true)   # sparse feature values
    return x @ F.T, x                                                         # activations, true features


class SAE(nn.Module):
    def __init__(self, m):
        super().__init__()
        self.b_dec = nn.Parameter(torch.zeros(d))
        self.enc = nn.Linear(d, m)
        self.D = nn.Parameter(torch.randn(m, d) * 0.1)

    def forward(self, x):
        f = torch.relu(self.enc(x - self.b_dec))
        Dn = self.D / (self.D.norm(dim=1, keepdim=True) + 1e-8)   # unit-norm dictionary
        return f, self.b_dec + f @ Dn


def train_sae(m, lam, steps=2500):
    sae = SAE(m); opt = torch.optim.Adam(sae.parameters(), lr=2e-3)
    for _ in range(steps):
        x, _ = gen_acts(512)
        f, xr = sae(x)
        loss = ((x - xr) ** 2).sum(1).mean() + lam * f.abs().sum(1).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return sae


print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt import torch import torch.nn as nn import warnings warnings.filterwarnings('ignore') torch.manual_seed(0) np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) d, n_true = 16, 40 # 40 true features superposed in 16 dims F = torch.randn(d, n_true); F /= F.norm(dim=0, keepdim=True) def gen_acts(batch, p=0.06): x = (torch.rand(batch, n_true) < p).float() * torch.rand(batch, n_true) # sparse feature values return x @ F.T, x # activations, true features class SAE(nn.Module): def __init__(self, m): super().__init__() self.b_dec = nn.Parameter(torch.zeros(d)) self.enc = nn.Linear(d, m) self.D = nn.Parameter(torch.randn(m, d) * 0.1) def forward(self, x): f = torch.relu(self.enc(x - self.b_dec)) Dn = self.D / (self.D.norm(dim=1, keepdim=True) + 1e-8) # unit-norm dictionary return f, self.b_dec + f @ Dn def train_sae(m, lam, steps=2500): sae = SAE(m); opt = torch.optim.Adam(sae.parameters(), lr=2e-3) for _ in range(steps): x, _ = gen_acts(512) f, xr = sae(x) loss = ((x - xr) ** 2).sum(1).mean() + lam * f.abs().sum(1).mean() opt.zero_grad(); loss.backward(); opt.step() return sae print('Setup complete.')

Setup complete.

---

Experiment 1: Sparsity trade-off¶

Question: The SAE loss is reconstruction + λ·L1. How does λ trade reconstruction quality against how many features fire at once (sparsity)?

We train at several λ (fixed dictionary size) and plot reconstruction error against the average number of active features (L0).

In [2]:

lams = [0.02, 0.05, 0.1, 0.2, 0.4, 0.8]
mse, l0 = [], []
x_test, _ = gen_acts(2000)
for lam in lams:
    sae = train_sae(m=128, lam=lam)
    with torch.no_grad():
        f, xr = sae(x_test)
    mse.append(((x_test - xr) ** 2).sum(1).mean().item())
    l0.append((f > 1e-3).float().sum(1).mean().item())
print('lambda :', lams)
print('MSE    :', np.round(mse, 3))
print('L0     :', np.round(l0, 2))

lams = [0.02, 0.05, 0.1, 0.2, 0.4, 0.8] mse, l0 = [], [] x_test, _ = gen_acts(2000) for lam in lams: sae = train_sae(m=128, lam=lam) with torch.no_grad(): f, xr = sae(x_test) mse.append(((x_test - xr) ** 2).sum(1).mean().item()) l0.append((f > 1e-3).float().sum(1).mean().item()) print('lambda :', lams) print('MSE :', np.round(mse, 3)) print('L0 :', np.round(l0, 2))

lambda : [0.02, 0.05, 0.1, 0.2, 0.4, 0.8]
MSE    : [0.003 0.014 0.037 0.081 0.154 0.318]
L0     : [28.65 21.11 13.02  6.8   3.07  1.92]

In [3]:

fig, ax = plt.subplots(figsize=(8.5, 5.2))
fig.suptitle('The sparsity–reconstruction trade-off (set by λ)', fontsize=13, fontweight='bold')

ax.plot(l0, mse, '-o', color='seagreen', lw=2)
for lam, a, b in zip(lams, l0, mse):
    ax.annotate(f'λ={lam}', (a, b), textcoords='offset points', xytext=(6, 5), fontsize=8)
ax.set_xlabel('avg active features per input (L0, sparsity)')
ax.set_ylabel('reconstruction MSE')
ax.set_title('Lower-left is unattainable; λ slides along the frontier')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8.5, 5.2)) fig.suptitle('The sparsity–reconstruction trade-off (set by λ)', fontsize=13, fontweight='bold') ax.plot(l0, mse, '-o', color='seagreen', lw=2) for lam, a, b in zip(lams, l0, mse): ax.annotate(f'λ={lam}', (a, b), textcoords='offset points', xytext=(6, 5), fontsize=8) ax.set_xlabel('avg active features per input (L0, sparsity)') ax.set_ylabel('reconstruction MSE') ax.set_title('Lower-left is unattainable; λ slides along the frontier') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: A clear Pareto frontier — small λ gives low reconstruction error but many active features (dense, less interpretable); large λ gives few active features (sparse, interpretable) but higher reconstruction error.
2. Why it looks this way: λ directly weights the L1 penalty against reconstruction, so raising it forces fewer features to explain each input at the cost of fidelity.
3. Key takeaway: There is no free lunch — choosing λ is choosing a point on the sparsity/fidelity frontier, the central tuning decision when training an SAE.

---

Experiment 2: Feature splitting¶

Question: The paper finds that a larger dictionary splits one coarse feature into several finer ones. Does dictionary size act like a resolution knob?

We make true features come in close pairs (fine variants of a concept). A small dictionary should merge each pair into one atom; a large dictionary should split them into two distinct atoms.

In [4]:

torch.manual_seed(2)
n_pairs = 8
base = torch.randn(d, n_pairs)
# each concept has 2 close variants (cos ~ 0.9)
Fp = []
for j in range(n_pairs):
    for _ in range(2):
        v = base[:, j] + 0.3 * torch.randn(d)
        Fp.append(v / v.norm())
Fp = torch.stack(Fp, 1)   # d x 16


def gen_pairs(batch, p=0.08):
    x = (torch.rand(batch, Fp.shape[1]) < p).float() * torch.rand(batch, Fp.shape[1])
    return x @ Fp.T


def distinct_atoms_per_pair(m, steps=2500):
    sae = SAE(m); opt = torch.optim.Adam(sae.parameters(), lr=2e-3)
    for _ in range(steps):
        x = gen_pairs(512); f, xr = sae(x)
        (((x - xr) ** 2).sum(1).mean() + 0.1 * f.abs().sum(1).mean()).backward()
        opt.step(); opt.zero_grad()
    D = sae.D.detach(); D = D / D.norm(dim=1, keepdim=True)
    Fn = (Fp / Fp.norm(dim=0, keepdim=True)).T          # 16 x d
    counts = []
    for j in range(n_pairs):
        a0 = int((Fn[2*j] @ D.T).argmax()); a1 = int((Fn[2*j+1] @ D.T).argmax())
        counts.append(2 if a0 != a1 else 1)             # distinct atoms for the 2 variants?
    return np.mean(counts)


split_small = distinct_atoms_per_pair(m=10)
split_large = distinct_atoms_per_pair(m=64)
print(f'avg distinct atoms per concept-pair  small dict (m=10): {split_small:.2f}   large dict (m=64): {split_large:.2f}')

torch.manual_seed(2) n_pairs = 8 base = torch.randn(d, n_pairs) # each concept has 2 close variants (cos ~ 0.9) Fp = [] for j in range(n_pairs): for _ in range(2): v = base[:, j] + 0.3 * torch.randn(d) Fp.append(v / v.norm()) Fp = torch.stack(Fp, 1) # d x 16 def gen_pairs(batch, p=0.08): x = (torch.rand(batch, Fp.shape[1]) < p).float() * torch.rand(batch, Fp.shape[1]) return x @ Fp.T def distinct_atoms_per_pair(m, steps=2500): sae = SAE(m); opt = torch.optim.Adam(sae.parameters(), lr=2e-3) for _ in range(steps): x = gen_pairs(512); f, xr = sae(x) (((x - xr) ** 2).sum(1).mean() + 0.1 * f.abs().sum(1).mean()).backward() opt.step(); opt.zero_grad() D = sae.D.detach(); D = D / D.norm(dim=1, keepdim=True) Fn = (Fp / Fp.norm(dim=0, keepdim=True)).T # 16 x d counts = [] for j in range(n_pairs): a0 = int((Fn[2*j] @ D.T).argmax()); a1 = int((Fn[2*j+1] @ D.T).argmax()) counts.append(2 if a0 != a1 else 1) # distinct atoms for the 2 variants? return np.mean(counts) split_small = distinct_atoms_per_pair(m=10) split_large = distinct_atoms_per_pair(m=64) print(f'avg distinct atoms per concept-pair small dict (m=10): {split_small:.2f} large dict (m=64): {split_large:.2f}')

avg distinct atoms per concept-pair  small dict (m=10): 1.00   large dict (m=64): 1.75

In [5]:

fig, ax = plt.subplots(figsize=(7.5, 5))
fig.suptitle('A larger dictionary splits coarse features into fine ones', fontsize=13, fontweight='bold')

ax.bar(['small dict\n(m=10)', 'large dict\n(m=64)'], [split_small, split_large],
       color=['crimson', 'seagreen'])
ax.axhline(1, color='gray', ls='--', lw=1, label='1 = variants merged')
ax.axhline(2, color='black', ls=':', lw=1, label='2 = variants split')
ax.set_ylabel('avg distinct atoms per concept-pair'); ax.set_ylim(0, 2.3)
ax.set_title('Resolution rises with dictionary size')
ax.legend()

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(7.5, 5)) fig.suptitle('A larger dictionary splits coarse features into fine ones', fontsize=13, fontweight='bold') ax.bar(['small dict\n(m=10)', 'large dict\n(m=64)'], [split_small, split_large], color=['crimson', 'seagreen']) ax.axhline(1, color='gray', ls='--', lw=1, label='1 = variants merged') ax.axhline(2, color='black', ls=':', lw=1, label='2 = variants split') ax.set_ylabel('avg distinct atoms per concept-pair'); ax.set_ylim(0, 2.3) ax.set_title('Resolution rises with dictionary size') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: With a small dictionary, the two variants of each concept tend to share one atom (≈1 distinct atom per pair); with a large dictionary, they get separate atoms (≈2 per pair).
2. Why it looks this way: A small dictionary lacks capacity, so it represents each concept coarsely with a single shared direction. Extra capacity lets the SAE allocate separate atoms to fine variants — feature splitting.
3. Key takeaway: The number of features is a resolution knob, not an absolute count — bigger dictionaries reveal finer-grained concepts, exactly as the paper reports.

---

Experiment 3: Monosemanticity¶

Question: Are SAE features more single-concept (monosemantic) than the raw neurons? We measure, for each unit, how concentrated its responses are on a single true feature.

Specificity = the share of a unit's activation explained by its single best-matching true feature (1.0 = perfectly monosemantic).

In [6]:

torch.manual_seed(3)
sae = train_sae(m=128, lam=0.15)
x, xt = gen_acts(4000)
with torch.no_grad():
    feats, _ = sae(x)
feats = feats.numpy(); xt = xt.numpy(); acts_neuron = x.numpy()


def specificity(unit_acts):
    # correlation of a unit's activation with each true feature; share of the top one
    unit_acts = unit_acts - unit_acts.mean()
    corr = np.abs((xt - xt.mean(0)).T @ unit_acts)
    corr = corr / (corr.sum() + 1e-9)
    return corr.max()


alive = np.where((feats > 1e-3).mean(0) > 1e-3)[0]
sae_spec = [specificity(feats[:, i]) for i in alive]
neuron_spec = [specificity(acts_neuron[:, j]) for j in range(d)]
print(f'median specificity  SAE features: {np.median(sae_spec):.2f}   neurons: {np.median(neuron_spec):.2f}')

torch.manual_seed(3) sae = train_sae(m=128, lam=0.15) x, xt = gen_acts(4000) with torch.no_grad(): feats, _ = sae(x) feats = feats.numpy(); xt = xt.numpy(); acts_neuron = x.numpy() def specificity(unit_acts): # correlation of a unit's activation with each true feature; share of the top one unit_acts = unit_acts - unit_acts.mean() corr = np.abs((xt - xt.mean(0)).T @ unit_acts) corr = corr / (corr.sum() + 1e-9) return corr.max() alive = np.where((feats > 1e-3).mean(0) > 1e-3)[0] sae_spec = [specificity(feats[:, i]) for i in alive] neuron_spec = [specificity(acts_neuron[:, j]) for j in range(d)] print(f'median specificity SAE features: {np.median(sae_spec):.2f} neurons: {np.median(neuron_spec):.2f}')

median specificity  SAE features: 0.36   neurons: 0.07

In [7]:

fig, ax = plt.subplots(figsize=(8.7, 5))
fig.suptitle('SAE features are far more monosemantic than neurons', fontsize=13, fontweight='bold')

bins = np.linspace(0, 1, 25)
ax.hist(neuron_spec, bins=bins, alpha=0.6, color='crimson', density=True, label=f'neurons (median {np.median(neuron_spec):.2f})')
ax.hist(sae_spec, bins=bins, alpha=0.6, color='seagreen', density=True, label=f'SAE features (median {np.median(sae_spec):.2f})')
ax.set_xlabel('specificity (share on single best true feature)'); ax.set_ylabel('density')
ax.set_title('Higher = more single-concept')
ax.legend()

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8.7, 5)) fig.suptitle('SAE features are far more monosemantic than neurons', fontsize=13, fontweight='bold') bins = np.linspace(0, 1, 25) ax.hist(neuron_spec, bins=bins, alpha=0.6, color='crimson', density=True, label=f'neurons (median {np.median(neuron_spec):.2f})') ax.hist(sae_spec, bins=bins, alpha=0.6, color='seagreen', density=True, label=f'SAE features (median {np.median(sae_spec):.2f})') ax.set_xlabel('specificity (share on single best true feature)'); ax.set_ylabel('density') ax.set_title('Higher = more single-concept') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: SAE-feature specificity (green) clusters near 1 — each feature responds mostly to a single true feature — while neuron specificity (red) is much lower and spread out.
2. Why it looks this way: Neurons carry superposed mixtures, so each correlates partially with many true features. The sparse, overcomplete SAE disentangles these into directions that each track one concept.
3. Key takeaway: The SAE turns polysemantic neurons into monosemantic features — the paper's headline result, quantified here as a specificity gap.

---

Experiment 4: Dead features¶

Question: Some SAE features never activate (dead), wasting capacity — the same pathology as dead codes in VQ. Does resampling them to high-error inputs revive them?

We count dead features after training, then reinitialize them toward poorly-reconstructed inputs and retrain briefly.

In [8]:

torch.manual_seed(4)
m = 200
sae = train_sae(m=m, lam=0.2, steps=2000)
x, _ = gen_acts(4000)
with torch.no_grad():
    f, xr = sae(x)
dead = (f > 1e-3).float().mean(0) < 1e-4
n_dead_before = int(dead.sum())

# resample: point dead encoder rows / dict atoms at the worst-reconstructed inputs
with torch.no_grad():
    err = ((x - xr) ** 2).sum(1)
    worst = x[torch.topk(err, n_dead_before).indices]
    idx = torch.where(dead)[0]
    sae.D.data[idx] = (worst - sae.b_dec) / ((worst - sae.b_dec).norm(dim=1, keepdim=True) + 1e-8)
    sae.enc.weight.data[idx] = (worst - sae.b_dec)
    sae.enc.bias.data[idx] = 0.0
opt = torch.optim.Adam(sae.parameters(), lr=2e-3)
for _ in range(800):
    xb, _ = gen_acts(512); fb, xrb = sae(xb)
    (((xb - xrb) ** 2).sum(1).mean() + 0.2 * fb.abs().sum(1).mean()).backward()
    opt.step(); opt.zero_grad()
with torch.no_grad():
    f2, _ = sae(gen_acts(4000)[0])
n_dead_after = int(((f2 > 1e-3).float().mean(0) < 1e-4).sum())
print(f'dead features  before resampling: {n_dead_before}/{m}   after: {n_dead_after}/{m}')

torch.manual_seed(4) m = 200 sae = train_sae(m=m, lam=0.2, steps=2000) x, _ = gen_acts(4000) with torch.no_grad(): f, xr = sae(x) dead = (f > 1e-3).float().mean(0) < 1e-4 n_dead_before = int(dead.sum()) # resample: point dead encoder rows / dict atoms at the worst-reconstructed inputs with torch.no_grad(): err = ((x - xr) ** 2).sum(1) worst = x[torch.topk(err, n_dead_before).indices] idx = torch.where(dead)[0] sae.D.data[idx] = (worst - sae.b_dec) / ((worst - sae.b_dec).norm(dim=1, keepdim=True) + 1e-8) sae.enc.weight.data[idx] = (worst - sae.b_dec) sae.enc.bias.data[idx] = 0.0 opt = torch.optim.Adam(sae.parameters(), lr=2e-3) for _ in range(800): xb, _ = gen_acts(512); fb, xrb = sae(xb) (((xb - xrb) ** 2).sum(1).mean() + 0.2 * fb.abs().sum(1).mean()).backward() opt.step(); opt.zero_grad() with torch.no_grad(): f2, _ = sae(gen_acts(4000)[0]) n_dead_after = int(((f2 > 1e-3).float().mean(0) < 1e-4).sum()) print(f'dead features before resampling: {n_dead_before}/{m} after: {n_dead_after}/{m}')

dead features  before resampling: 63/200   after: 22/200

In [9]:

fig, ax = plt.subplots(figsize=(7.5, 5))
fig.suptitle('Resampling revives dead SAE features', fontsize=13, fontweight='bold')

ax.bar(['before\nresampling', 'after\nresampling'], [n_dead_before, n_dead_after], color=['crimson', 'seagreen'])
ax.set_ylabel(f'dead features (of {m})')
ax.set_title('Dead-feature count drops after resampling')
for i, v in enumerate([n_dead_before, n_dead_after]):
    ax.text(i, v, str(v), ha='center', va='bottom', fontsize=11)

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(7.5, 5)) fig.suptitle('Resampling revives dead SAE features', fontsize=13, fontweight='bold') ax.bar(['before\nresampling', 'after\nresampling'], [n_dead_before, n_dead_after], color=['crimson', 'seagreen']) ax.set_ylabel(f'dead features (of {m})') ax.set_title('Dead-feature count drops after resampling') for i, v in enumerate([n_dead_before, n_dead_after]): ax.text(i, v, str(v), ha='center', va='bottom', fontsize=11) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: A sizeable fraction of features are dead after training (red); after resampling them toward high-error inputs and a short retrain, the dead count drops sharply (green).
2. Why it looks this way: Features whose direction never wins activation receive no gradient and stay dead. Reseeding them onto inputs the SAE reconstructs poorly puts them where they can capture unexplained variance and start firing.
3. Key takeaway: Dead features are the SAE analogue of VQ codebook collapse, and the same fix — resampling — recovers wasted capacity.

---

Summary¶

Experiment	Key insight
1. Sparsity trade-off	λ slides the SAE along a reconstruction/sparsity Pareto frontier.
2. Feature splitting	Bigger dictionaries split coarse features into finer ones — size is a resolution knob.
3. Monosemanticity	SAE features track single concepts; neurons stay polysemantic.
4. Dead features	Some features never fire; resampling toward high-error inputs revives them.

Connection to the broader theory¶

These reproduce the method and findings of Towards Monosemanticity: a sparse, overcomplete dictionary extracts monosemantic features from superposition, with feature splitting and dead-feature pathologies. The next topic, Scaling Monosemanticity, takes this method to Claude 3 Sonnet and focuses on how to evaluate feature quality at scale.