Kiru Lab  /  Mechanistic Interpretability  /  From Mechanism to Diagnosis

Sparse Autoencoders and Dictionary Learning

If superposition hides features in a compact basis, learn a bigger basis in which they separate.

Concept  ·  about 40 minutes

A sparse autoencoder takes a layer's activations and learns to encode them into a much wider space — often 8 to 64 times wider — under a sparsity penalty, then decode them back. The bet is that the true features are sparse and near-orthogonal, so the wide sparse basis will separate what superposition had entangled.

z = ReLU(W_enc (x - b_dec) + b_enc) wide and sparse x_hat = W_dec z + b_dec Loss = ||x - x_hat||^2 + lambda * ||z||_1

The L1 penalty is the sparsity pressure, and it is the same L1 you met as lasso regularization in track two. The reconstruction term keeps the dictionary faithful. The tension between them is the whole design problem: too much sparsity and you lose information, too little and you have relearned the entangled basis in more dimensions.

  • Many recovered dictionary elements are strikingly interpretable — specific concepts, syntactic roles, topics — in a way individual neurons are not.
  • Evaluation is genuinely unsettled: reconstruction loss and sparsity are measurable, but "interpretability" mostly is not, and auto-interpretation scores have their own failure modes.
  • Feature splitting is real: train a wider dictionary and one feature often subdivides into several, which raises an unresolved question about whether there is a true feature count at all.
  • A dictionary element is still only a correlational object until you intervene on it. The same rule applies here as everywhere else in this track.
Where the field actually is

This is the current frontier rather than a settled technique, and treating it as settled is a mistake. It belongs at the end of this curriculum precisely because evaluating it requires everything before it — sparsity, conditioning, regularization, causal testing, and the discipline to say "we do not know yet".

Hold on to

  • SAEs trade a compact entangled basis for a wide sparse legible one
  • The sparsity/reconstruction tradeoff is the central design tension
  • Dictionary elements need causal validation like anything else

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Train a small SAE on activations from an open model and inspect the top-activating inputs for ten dictionary elements.
    Hint

    Collect activations from one layer over a decent corpus, then train a wide ReLU autoencoder with an L1 penalty.

    Solution

    Some dictionary elements will look strikingly clean — a specific topic, a syntactic role, a formatting pattern — and others will resist interpretation entirely. Report both proportions honestly. Selecting the ten interpretable ones out of thousands and presenting those is a form of cherry-picking that has made this literature harder to evaluate than it should be.

    import torch, torch.nn as nn

class SAE(nn.Module):
    def __init__(self, d_model, expansion=8):
        super().__init__()
        d_hidden = d_model * expansion
        self.enc = nn.Linear(d_model, d_hidden)
        self.dec = nn.Linear(d_hidden, d_model)

    def forward(self, x):
        z = torch.relu(self.enc(x - self.dec.bias))
        return self.dec(z), z

def loss_fn(x, x_hat, z, l1=1e-3):
    return ((x - x_hat) ** 2).mean() + l1 * z.abs().sum(dim=-1).mean()
  2. Sweep the L1 coefficient and plot reconstruction loss against average sparsity. Mark the point where features start collapsing.
    Hint

    Plot reconstruction loss against mean L0 (the average number of active features).

    Solution

    The curve is a clean Pareto front: lower L1 gives better reconstruction and denser codes, higher L1 gives sparser codes and worse reconstruction. Somewhere on the high-L1 side features start collapsing — dictionary elements go permanently dead, or several merge into one that fires for a grab-bag of things. The absence of a principled criterion for choosing a point on this curve is one of the honest open problems in the field, and you should be suspicious of any paper that presents its chosen point as obvious.

    for l1 in (1e-5, 1e-4, 1e-3, 1e-2, 1e-1):
    sae = train_sae(activations, l1=l1)
    x_hat, z = sae(activations)
    recon = float(((activations - x_hat) ** 2).mean())
    l0 = float((z > 0).float().sum(dim=-1).mean())
    dead = float((z.max(dim=0).values == 0).float().mean())
    print(f"l1={l1:<8} recon={recon:.4f} L0={l0:6.1f} dead={dead:.1%}")

In the Bio Mirror

Sparse Coding and the Energy Budget

Sign in to track your progress through the lab.

← From Mechanism to Diagnosis From Circuit to Disorder →