Sixteen million features in GPT-4: OpenAI's TopK sparse autoencoders

Choosing k instead of tuning a penalty

OpenAI posted a sparse autoencoder paper on June 6 with a 16 million latent dictionary trained on GPT-4 residual stream activations over 40 billion tokens. The authors are Leo Gao, Tom Dupré la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Several of them were on the Superalignment team, and this reads as that team's last interpretability paper.

The central method change is small to describe. A standard sparse autoencoder uses a ReLU encoder and adds an L1 penalty on the latents to push most of them to zero. The TopK autoencoder instead keeps the k largest latents for each token and zeroes the rest. Sparsity becomes a number you set rather than a coefficient you tune until the L0 comes out where you wanted.

The reason this is better than a convenience is shrinkage. An L1 penalty pulls every positive activation toward zero, so the reconstruction is biased downward even for features that should fire strongly. The paper shows through refinement experiments that TopK avoids this bias. You get a cleaner reconstruction at the same sparsity because the penalty is no longer fighting the magnitudes.

Dead latents

The obstacle to training very wide dictionaries is that latents die. Once a latent stops firing on any input, nothing in the loss brings it back, and the paper reports up to 90 percent dead latents without mitigation at large scale. Prior work had accepted this as a cost. The comparison the authors give is a 34 million latent dictionary from earlier work in which only 12 million were alive, a 65 percent dead rate.

Two fixes get the number down. The first is tied initialisation, setting each encoder direction parallel to its decoder direction at the start. The second is an auxiliary loss they call AuxK, which asks the top k_aux dead latents to model the reconstruction residual, so that dead latents receive a gradient that encourages them to pick up whatever the live ones missed. AuxK adds about 10 percent compute. With both, the 16 million latent GPT-4 autoencoder ends training with 7 percent of its latents dead.

Scaling laws for a dictionary

The paper treats autoencoders the way a pretraining team treats models. Reconstruction MSE follows a power law in compute budget. The token budget needed to converge scales as roughly n to the 0.6 for GPT-2 small and n to the 0.65 for GPT-4, where n is the number of latents. And there is a joint law relating loss to both the number of latents and the sparsity k, so you can trade width for sparsity along a predictable curve.

One number gives a sense of how much of GPT-4 the dictionary captures. Replacing the residual stream with the autoencoder's reconstruction degrades the model's downstream loss to the level of a GPT-4 trained with about 10 percent of its pretraining compute. That is a large amount of the model's behaviour recovered from a sparse code, and also a reminder that a lot is still missing.

Four ways to score a feature

The evaluation section is the part we expect to be reused most. Reconstruction loss alone tells you nothing about whether latents mean anything, so the paper adds three metrics. Probe loss checks whether the dictionary recovers known features, using 61 classification tasks. Explainability uses an automated method called N2G to generate a pattern for each latent and measures precision and recall of that pattern against the latent's activations. Ablation sparsity asks whether turning off one latent has a sparse effect on the output logits, on the theory that a real feature should affect a few things rather than everything.

All three improve as the dictionary grows. That is the strongest argument in the paper for going wide, because reconstruction improving with width is expected, while interpretability improving with width was not guaranteed.

What the authors say is still broken

The limitations are stated more bluntly than usual. A large fraction of the random activations of features found, especially in GPT-4, are not yet adequately monosemantic. The activations were collected with a context of 64 tokens, which is short. And TopK forces every token to use exactly k latents, which the authors call likely suboptimal, since some tokens plausibly need more features and some fewer.

The release includes training code, trained autoencoders for open models, and a visualiser, so the method is checkable outside OpenAI. What we would try first is the fixed-k assumption. A per-token adaptive k, or a threshold learned per latent, would test whether the rigidity is costing reconstruction or interpretability, and would tell us whether the scaling laws here hold for whichever activation rule wins.