Sparse autoencoders, hands on

Individual neurons respond to several unrelated things at once, which is what makes them hard to read. A sparse autoencoder is the standard attempt to unpick that, and this section trains a small one on real activations: collecting the data, the two terms in the loss that pull against each other, the features that stop firing early in training and never recover, and how to write down what a feature means. That last part is the work, since a label taken from the twenty strongest examples is a hypothesis rather than a finding.

The idea again, briefly#

Neurons are hard to read because the model packs many features into overlapping directions. A sparse autoencoder learns a larger set of directions such that each activation is rebuilt from only a few of them. The hope is that each direction is one understandable feature.

A minimal version#

import torch, torch.nn as nn

class SAE(nn.Module):
    def __init__(self, d_model, n_features):
        super().__init__()
        self.enc = nn.Linear(d_model, n_features)
        self.dec = nn.Linear(n_features, d_model)

    def forward(self, x):
        f = torch.relu(self.enc(x))          # feature activations, mostly zero
        return self.dec(f), f

sae = SAE(768, 768 * 8)
opt = torch.optim.Adam(sae.parameters(), lr=1e-4)
for x in activation_batches:                  # e.g. layer 6 residual stream, many tokens
    x_hat, f = sae(x)
    loss = ((x_hat - x) ** 2).mean() + 1e-3 * f.abs().sum(-1).mean()
    opt.zero_grad(); loss.backward(); opt.step()

Two terms fight: reconstruct the activation well, and use as few features as possible. The penalty weight sets the balance. Track both the reconstruction error and the average number of active features per token.

Training one where the answer is known#

The difficulty with SAEs on a real model is that nobody knows what the right answer is, so every measurement of success is indirect. The way out for learning purposes is to build the activations yourself. The toy model below packs 64 sparse features into 16 dimensions, exactly the superposition situation an SAE is meant to undo, and because you built it you know every true feature direction. Train an SAE on its hidden activations and you can ask the question you can never ask on GPT-2: did it find them.

import torch

torch.manual_seed(0)
N_FEATURES, D_HIDDEN, SPARSITY = 64, 16, 0.95

class ToyModel(torch.nn.Module):
    """64 sparse features squeezed through 16 dimensions and reconstructed."""
    def __init__(self):
        super().__init__()
        self.W = torch.nn.Parameter(torch.randn(N_FEATURES, D_HIDDEN) * 0.1)
        self.b = torch.nn.Parameter(torch.zeros(N_FEATURES))

    def forward(self, x):
        return torch.relu((x @ self.W) @ self.W.T + self.b)

def batch(n):
    return torch.rand(n, N_FEATURES) * (torch.rand(n, N_FEATURES) < 1 - SPARSITY)

toy = ToyModel()
opt = torch.optim.Adam(toy.parameters(), lr=1e-2)
importance = 0.97 ** torch.arange(N_FEATURES)
for _ in range(3000):
    x = batch(1024)
    loss = (importance * (toy(x) - x) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()

W = toy.W.detach()
live = W.norm(dim=1) > 0.5
true = W[live] / W[live].norm(dim=1, keepdim=True)   # the ground-truth directions
print(f"{live.sum().item()} of {N_FEATURES} features represented in {D_HIDDEN} dimensions")
print(f"mean active features per sample: {(batch(4096) > 0).float().sum(-1).mean():.2f}")
# 57 of 64 features represented in 16 dimensions
# mean active features per sample: 3.25

Now train the autoencoder at four penalty weights and measure four things each time: how much of the activation survives the round trip, how many features fire per sample, how many never fire at all, and how close the learned directions come to the true ones. The last of those needs a control, which the next heading explains.

class SAE(torch.nn.Module):
    def __init__(self, d_in, n_feat):
        super().__init__()
        self.enc = torch.nn.Linear(d_in, n_feat)
        self.dec = torch.nn.Linear(n_feat, d_in, bias=False)

    def forward(self, h):
        f = torch.relu(self.enc(h))      # feature activations, mostly zero
        return self.dec(f), f

def train_sae(penalty, steps=3000):
    torch.manual_seed(1)
    sae = SAE(D_HIDDEN, 256)
    opt = torch.optim.Adam(sae.parameters(), lr=3e-3)
    for _ in range(steps):
        with torch.no_grad():
            h = batch(2048) @ W                        # the activations to explain
        h_hat, f = sae(h)
        recon = ((h_hat - h) ** 2).sum(-1).mean()
        loss = recon + penalty * f.abs().sum(-1).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():                          # unit-norm decoder columns
            sae.dec.weight.data /= sae.dec.weight.data.norm(dim=0, keepdim=True) + 1e-8
    return sae

def report(sae):
    with torch.no_grad():
        h = batch(16384) @ W
        h_hat, f = sae(h)
    fired = (f > 0).any(0)
    directions = sae.dec.weight.data.T[fired]
    best = (true @ directions.T).abs().max(-1).values      # recovery of each true feature
    torch.manual_seed(2)
    R = torch.randn(int(fired.sum()), D_HIDDEN); R = R / R.norm(dim=1, keepdim=True)
    control = (true @ R.T).abs().max(-1).values            # the same count, at random
    return (((h_hat - h) ** 2).sum(-1).mean().item(), (f > 0).float().sum(-1).mean().item(),
            int((~fired).sum()), best.mean().item(), control.mean().item())

print("penalty  reconstruction   L0   dead/256   recovery   random control")
for penalty in (0.01, 0.03, 0.10, 0.30):
    r, l0, dead, rec, ctrl = report(train_sae(penalty))
    print(f"  {penalty:4.2f}       {r:8.4f}  {l0:5.1f}      {dead:3d}       {rec:.2f}           {ctrl:.2f}")

penalty  reconstruction   L0   dead/256   recovery   random control
  0.01         0.0009   53.3       82       0.73           0.67
  0.03         0.0048   36.1      100       0.76           0.66
  0.10         0.0485   20.4      131       0.93           0.65
  0.30         0.1698    5.0      165       0.98           0.63

Read across the rows. Reconstruction error rises by a factor of 190 from the top row to the bottom, L0 falls from 53 to 5, the dead share doubles from 82 features to 165, and recovery of the true directions goes from 0.73 to 0.98. There is no setting that is good at everything, which is the actual content of the phrase 'the two terms fight'. The bottom row is the one that recovers the features, and it is also the one that throws away the most of the activation and wastes the most of the dictionary. Anyone showing you an SAE with one number has chosen which of these to tell you about.

The L0 column has a reference point here that a real model does not give you: the data genuinely has about 3.25 active features per sample, so the bottom row at 5.0 is roughly right and the top row at 53.3 is describing nothing. On a real layer you have no such anchor, which is why L0 is usually chosen by what stays interpretable rather than by what is correct.

The control that decides whether recovery means anything#

Recovery above is the best cosine similarity between each true direction and any of the SAE's directions. That statistic has a serious problem: with enough candidate directions in few enough dimensions, a random dictionary scores well. The control column is exactly that, the same number of directions drawn at random, and at a penalty of 0.01 the SAE's 0.73 against a random 0.67 is close to nothing. Only the bottom two rows clear their control by a margin worth reporting.

How large the random baseline is depends almost entirely on the width of the space, and it is worth knowing the shape of that curve before quoting any cosine similarity.

torch.manual_seed(2)
print("\nbest |cos| to the nearest of 100 random directions, by dimension")
for d in (5, 16, 64, 256):
    A = torch.randn(57, d); A = A / A.norm(dim=1, keepdim=True)
    B = torch.randn(100, d); B = B / B.norm(dim=1, keepdim=True)
    print(f"  {d:4d} dimensions: {(A @ B.T).abs().max(-1).values.mean():.2f}")
# best |cos| to the nearest of 100 random directions, by dimension
#      5 dimensions: 0.92
#     16 dimensions: 0.61
#     64 dimensions: 0.34
#    256 dimensions: 0.17

Dead features#

Many features stop activating on anything early in training and never recover. Check what share never fired over a large sample. Published SAE work uses various tricks, such as resampling dead features, to reduce this.

The dead column above is worth sitting with. At the setting that recovers the true features best, 165 of 256 dictionary entries never fire once in 16,384 samples. That is 64% of the capacity you trained, paid for and will store, doing nothing. It also quietly changes what the other numbers mean, because a dictionary of 256 that behaves as a dictionary of 91 is a different object from the one described in the configuration. Report the count of live features alongside the nominal size, always.

Interpreting a feature#

1. 1.Collect the text snippets where it activates most strongly and read them.
2. 2.Also read snippets where it activates weakly. A feature that looks like 'French' at the top may be 'any non-English text' lower down.
3. 3.Check causally: add the feature's direction to the stream and see whether output shifts the way the label predicts.

What to carry, and what to skip#

• Carry: reconstruction error, L0 and the live feature count are one measurement in three parts, and any of the three can be made to look good by sacrificing the others.
• Carry: what an SAE is for. It is an attempt to recover a sparse basis for a layer, not a decoder ring and not an explanation of a behaviour. A feature label is a hypothesis to be tested by intervention.
• Carry: any similarity number needs the random baseline for that dimensionality beside it.
• Skip: chasing the newest encoder variant. The measurement discipline above transfers to all of them; the architectures do not stay current.
• Skip: training an SAE on a real model as a first project. Train one where you know the answer, as above, then move up.