This post will probably not be very informative to readers who don't know what anSAE (Sparse Autoencoder)is.
An interpretability researcher once told me, "SAEs can never fail, they can only be failed." But recently, SAEs have been suffering failure after failure, and as a result, the field has grown increasingly pessimistic about SAE work. Korznikov et al.'s Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? purports to be an example of such failures.
However, I have some disagreements with its methodology. The paper presents a synthetic setup with known ground truth features, and claims that SAEs fail to recover features in the "simplest toy setting with fully aligned hyperparameters." In this post, I argue that their setting is not forgiving to SAEs, as the authors imply, but is instead actively hostile in ways that real settings are not.
Can SAEs actually recover real features? Hypothetically, if we knew what the "true" features of the data were, would an SAE be able to identify them?
Do trained SAEs actually do better than completely random SAEs?
This post will be focused on their answer to the first question.
The authors train an SAE on a synthetic toy model, with known ground-truth features, and according to their abstract, SAEs only recover 9% of true features. On the face of it, this seems pretty damning!
Now, earlier toy model experiments did show that SAEs could recover ground truth features, but the authors argue this is because they used a small expansion factor, where there were only twice as many "true" features as there were model dimensions, and that this is unrealistic. Their setup uses an expansion factor of 32, which matches real-world use cases much more closely.
So, how does their setup work?
They don't actually construct a synthetic neural network; instead, they create synthetic activations directly. First, they create their feature vectors, by sampling 3200 vectors from the unit sphere in 100-dimensional space. To create an activation, they just select about 20 feature vectors at random, scale each active feature by a magnitude drawn from LogNormal(0, 0.25),[1] and sum the scaled vectors.
They test two different regimes, one where each feature is equally likely to be active, and another where feature activation frequency has a long tail, so some features fire more often than others. (We'll focus on the constant regime for the most part)
They match their SAE to the setup. There are about 20 expected active features, so they target L0 = 20. There are 3200 true features, so they create 3200 SAE latents. Can the SAE, when given every possible advantage, recover the true features?
These results are, again, pretty damning! They claim that the SAE achieves high reconstruction accuracy while only recovering 3/3200 true features—personally, I disagree that 0.67 explained variance is "high reconstruction accuracy".[2] However, the SAE clearly only does slightly better than random guesswork at recovering the true features.
Now, one obvious objection is that L0/d_model here is 0.2, which is unusually dense. I'd expect there to be a lot of interference here, and in the real world L0/d_model is probably going to be closer to 0.05. Also, 100 dimensions is just not a lot, and maybe you can't really get blessing of dimensionality from it like you would in real models.
A less obvious objection is that if you scroll all the way down to Appendix B, they describe some other SAE architectures they used for this task, like TopK and Matryoshka-on-BatchTopK.
In contrast to BatchTopK, TopK recovers 99.9% of features.[3]
So, I did what anyone would do, and started replicating the paper. They didn't actually include all the hyperparameters they used,[4] so I made reasonable assumptions.
I was able to replicate this result! TopK recovered 99.6% of features,[5] and the other architectures failed miserably.[6]
Why does BatchTopK fail?
This is kind of strange, if you think about it. TopK and BatchTopK are the same architecture, with the only difference being the selection rule.[7] Such a dramatic gap is really weird.
Peeking a bit closer at my training run, in every batch, most latents never fired. That means a small group of monopolist latents were dominating every time. By step 5,000, the same ~6% of latents took over, and they remained in charge for the remaining 295,000 steps. In comparison, with TopK, after the first 20k steps, >99% of latents activated on every batch.
This is the standard dead latents problem: There's no gradient on latents that aren't selected, so once they fall behind in performance it's hard for them to improve. A small group of generalist latents dominate before specialists have the chance to form. The standard fix, of course, is an auxiliary loss, AuxK. But even with AuxK, the dead latent fraction barely budges. I tried sweeping different hyperparameters, but that didn't help either. I even tried resampling, and that did nothing.
Looking closer at the difference between TopK and BatchTopK gives us a hint. In theory, TopK should also be suffering from this problem. The fact that it doesn't suggests that the problem lies in sample-level variance. Certain samples in BatchTopK could receive a significantly larger amount of latents, while other samples are starved. An obvious thought is that input vectors with greater magnitude get more latents assigned to them than those with less magnitude. Checking the data, this does appear to be the case. So what if we normalized the input?
Bussmann's original BatchTopK implementation actually normalized per-sample, folding back in the original magnitude and mean during reconstruction.[8] I copied this method, and lo and behold! BatchTopK achieved 99.3% recovery, comparable to (my) TopK's 99.6% recovery. This also fixes JumpReLU with 97.2% feature recovery.
But while this was the original design, per-sample normalization isn't common anymore! GemmaScope and SAELens don't use it (although they do exclude outliers like BOS tokens), but they still work, without such drastic dead latent rates.
Despite the paper's claims, this synthetic model doesn't demonstrate the failure of SAEs to capture real features. Real SAEs don't have such drastic dead latent rates. Real SAEs are much better at reconstructing in-distribution data than out-of-distribution data. The SAE trained on synthetic data seems to be a fully general vector encoder, reconstructing random vectors with matching norms just as well as it does in-distribution data.
So all this model demonstrates is that this particular synthetic toy model is hostile to many SAE architectures (but not TopK). This is a weakness of these SAE architectures, in that they are not very robust to these types of environments, but this particular failure mode has an obvious signature (the dead latent rates).
Why is this synthetic environment hostile to SAEs?
I hope that I have demonstrated sufficient evidence of the failure of Sanity Checks' synthetic environment. The rest of this post will be dedicated to my personal explorations trying to understand precisely what it is about their environment that causes this divergence from real SAEs.[9]
d_model
Part of it is the density/dimensionality. If we raise d_model up to 800 while keeping the expansion factor of 32 (which means we're now using 25,600 features and latents), L0/d_model shrinks to 0.025, which is much more realistic. Under these conditions, TopK with auxk gets 99.9% recovered, and BatchTopK with auxk gets 53.5% with 46.3% dead latents. JumpReLU recovers 28.3% with 71.6% dead. So, it explains some of the difference, but not all of it.
While resampling didn't work at d_model = 100, when we try it at d_model = 800, it succeeds, raising BatchTopK's recovery to 85.6%. Anthropic resamples using 0.2x average active encoder weight norm, but when we try it with 1x instead, it actually raises recovery to 97.8%, close to TopK's.
So does this mean we've solved it? No! Real SAEs don't need resampling to avoid killing off half their latents.
I could continue to raise d_model, and I expect with More Dakka I could narrow the gap further. But GPT-2 only has 768 dimensions, and SAEstrainedonit work just fine (in that they don't have abysmal dead latent rates). So, we should look for other divergences from real models that our synthetic setup has, while keeping d_model at 800.
Magnitude Variance
An obvious guess is that our environment has unrealistic variance in per-sample magnitude, causing architectures like BatchTopK to perform poorly, and that real models have lower variance. While real models do have outlier tokens like BOS, SAEs typically exclude these, so I sampled GPT-2 activations on OpenWebText at various layers, excluding BOS. I subtracted out the dataset mean, since the decoder bias will usually be pretty close to it and is often initialized to the mean (or the geometric median).
Layer
Coefficient of Variation
GPT-2 L1
.224
GPT-2 L4
.181
GPT-2 L7
.140
GPT-2 L10
.129
Toy model (sigma 0.25)
.131
Toy model (sigma 0.5)
.185
As you can see, our toy model falls within the range of real GPT-2 activations, and in fact is on the low end. This means magnitude variance cannot explain the gap.
Feature Loudness
Why doesn't everybody normalize their SAE inputs? When you normalize, you're tossing away the magnitude, which could perhaps serve as an important signal. For example, maybe some features naturally activate with more magnitude than others, and total magnitude serves as a signal for that. So, maybe giving different features different natural "loudness" would improve performance.[10]
To simulate this, I draw a "loudness" variable for every feature from LogNormal(0, 0.5 sigma). Then, when creating activations, I scale each active feature not only by their activation strength (drawn from LogNormal(0, 0.25 sigma)), but also their loudness.
Running this with BatchTopK, we get 19.6% of features recovered, with 78.3% dead latents. Instead of improving performance, it actively degrades it. So feature loudness alone cannot explain the gap. (Also, my motivation for this was really dumb, see footnote 10 for why)
More Features!
Real LLMs might have many more features than their SAEs have latents. Perhaps, this gives latents more opportunities to grab onto "free" features. Alternatively, this might be distracting, making each individual feature appear less often, or trapping latents in an unfortunate middle ground where they can never lock onto a feature closely enough to converge on it.
So, I ran BatchTopK SAEs on setups with 4x and 2x more features than latents. We're still at 800 dimensions, 25,600 latents, and on average 20 features are active at a time.
On 4x, we recover 0.1% of features (This, of course, means 0.4% of latents captured features, since there are 4x more features than latents), with 58% dead latents. Even per-sample normalization[11] performs similarly poorly, recovering only 0.2% of features. They both reconstruct data only slightly better than noise.
On 2x, we recover 16.3% of features (which means 32.6% of latents capture features), with 59.6% dead latents. Per-sample normalization actually works this time, though, recovering 49.1% of features (so, 98.2% useful latents), and only 1.7% dead latents.
So, having more features turns out to be distracting, in this case.[12]
More Active Features
Real LLMs might have higher average active feature counts than the L0s of their SAEs. What if we tried doubling the active feature count to 40, while keeping target L0 the same?
Unfortunately, this also did not improve performance. BatchTopK only recovered 25.2% of features, compared to 53.5% previously.
Other Things I Tried
I experimented with many other things over the course of my research. Here is a non-exhaustive list of them:
Messing around with the variable setting. The paper's variable setting uses to select feature frequencies, but I also experimented with various Zipfian distributions. In general, I found that whether or not a feature was captured was very dependent on how often it appeared.
Correlating feature activations with each other.
Changing the variance in how many active features there are in each activation.
Magnitude Variance Revisited
At this point, I was feeling pretty frustrated. I knew that the toy was exhibiting failure modes that real SAEs don't typically have, but I wasn't sure entirely why.
As I mentioned in footnote 1, when I was first running these experiments, I incorrectly assumed that the active feature magnitude was drawn with a sigma of 0.5, or a variance of 0.25. I later discovered my mistake, and reran my experiments to match.
But this made me even more confused. Originally, when I was running these experiments, I found that increasing active feature count to 40 actually improved performance, but with sigma 0.25, as shown above, it was harmful to it! This, and other such oddities, made me revisit the question of magnitude variance.
It turns out that there are two types of magnitude variance to worry about: Sample-level magnitude variance is the variance in the magnitude of activations. Feature-level magnitude variance is the variance in the magnitude of a feature, over every activation of that feature.
Under our current setup, these two are correlated. But it seems very possible that sample-level magnitude variance is damaging to architectures like BatchTopK SAE, while feature-level magnitude variance is generally protective, making it easier for latents to capture features.[13]
We can fairly easily measure an objective sample-level magnitude variance for GPT-2, and we already did so. But, in order to find out what the feature-level magnitude variance is, we would need to know what the features are. Thankfully, there are existing GPT-2 SAEs, so I could just base the variance off of Bloom's SAEs.[14]
This is, of course, creating an environment for an SAE by using metrics pulled from an SAE, which is cheating. Thus, I wouldn't consider this to serve as proper evidence that SAEs recover true features, and this is overall not a very principled approach.
GPT-2 Layer 7 has a sample-level CV of 0.140, and a feature-level CV of 1.16.[15] We can target a specific sample-level CV by normalizing and rescaling to a magnitude from a LogNormal distribution with the correct mean and CV. We can target a specific feature-level CV by changing the sigma. However, rescaling does also lower feature-level CV.
I targeted .15 sample-level CV and sigma 1, which, after rescaling, ends up at 1.05 feature-level CV. Training a BatchTopK SAE, we get 96.1% feature recovery with 2.7% of latents dead, without having to do anything unconventional in the training process like per-sample normalization or resampling.
Limitations
When I tried targeting CV .18 and sigma 1, 31.2% of my latents died. GPT-2 Layer 4 has a CV of .18, and Bloom's L4 SAE has 0.3% dead latents. So there's something else that's protective about reality, and we only have one piece of the full story.
The main takeaway here is that it is very challenging to construct a realistic synthetic environment for LLMs, and we should be skeptical of evidence from toy models. My modified synthetic environment doesn't match reality. Sanity Checks' certainly doesn't. SynthSAEBench is more realistic than either of ours, and interestingly enough it also reports BatchTopK SAEs struggling with dead latents at low L0, although I'm not sure this is the same phenomenon.
Takeaways for real SAEs(?)
SAEs can never fail, they can only be failed. So how can we avoid failing them?
If your dead latent rates are really bad, that probably means your SAE is really bad. After all, SAEs can only capture features using live latents, so if most of your latents are dead, you're not capturing a lot of features. I hope this was already obvious, but in case it wasn't, now you know.
Consider using normalization for your SAEs, folding the magnitude back in during reconstruction, after encoding. I don't know if this actually improves performance in real SAEs, but it might. I'll consider testing this empirically, but only if I can get someone to fund it (SAE training is expensive, but SAE benchmarking is even more expensive).
A good health check for SAEs is how much better they are at reconstructing real data than they are at reconstructing random vectors with in-distribution norms. If they're only a little better, that's bad!
I originally assumed the 0.25 here meant variance, but then I asked the authors and actually it meant sigma. Unfortunately, this meant I had a bunch of experiments run under 0.5 sigma, which I then had to repeat with 0.25 sigma. This largely got similar results, although as a happy accident, the parts where it differed were very enlightening. All the results in this post are from 0.25 sigma unless stated otherwise.
In the paper, the authors define a feature being "recovered" as having an SAE latent with >=0.80 cosine similarity. I use this convention throughout the post.
To my chagrin, they don't mention this important finding outside the appendix. However, the ICML listing for their paper displays an updated abstract, which does mention it. For the life of me, though, I can't find any version of their paper which contains this updated abstract, so it's specifically only in the ICML listing for the paper.
The paper itself says, "We will publicly release all code, hyperparameters, and trained SAEs to ensure full reproducibility", but they only say this in Case Study #2, which isn't what I'm studying. The Case Study #2 SAEs with their training configs are on HuggingFace, but https://github.com/AntonKorznikov/random_sae, which is presumably where the code would be released, is empty.
Why only 99.6% instead of 99.9%? Probably down to slight differences in implementation and hyperparameters (perhaps they trained for longer). It's a very small difference and doesn't matter.
In TopK, you select the top k latent activations in each sample, so targeting L0 is ~guaranteed by the architecture. BatchTopK has you select the top batch size ⋅ k latent activations over the entire training batch, which allows for variance in the active latent count (since L0 targets the average latents active per sample, some samples may have more active latents and some samples may have fewer).
Bussmann uses LayerNorm-style normalization, not L2, and L2 is honestly more justifiable here given my reasoning for doing normalization. So I tried it, and L2 normalization gets very similar performance, recovering 99.2% of features on BatchTopK and 98.9% on JumpReLU.
Throughout this process I frequently worried that I had found something which failed to recover features but looked like what a real healthy SAE should look like. However, most failing configurations had extremely high dead latent rates, typically close to the fraction of features not recovered, and the ones that didn't had other signs of failure, like low EV, or a failure to reconstruct real data better than noise.
This is a stupid theory and should serve as an example of how not to do reasoning. In reality, SAEs do not use the total magnitude to decide what latents to activate. I don't know what I was thinking at the time, but nevertheless I am honestly reporting what my motivation was.
Bloom reads SAE activations from the start of each layer (hook_resid_pre) rather than the end. This is no longer conventional, so to avoid confusion I relabel them by the end of the previous layer, which is the same set of activations. For example, when I mention his L4 SAE, I'm referring to what he called his L5 SAE.
Specifically, I measured the median latent's activation magnitude CV. This might be unusually high because Bloom's SAEs are vanilla SAEs, but I don't know.
This post will probably not be very informative to readers who don't know what an SAE (Sparse Autoencoder) is.
An interpretability researcher once told me, "SAEs can never fail, they can only be failed." But recently, SAEs have been suffering failure after failure, and as a result, the field has grown increasingly pessimistic about SAE work. Korznikov et al.'s Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? purports to be an example of such failures.
However, I have some disagreements with its methodology. The paper presents a synthetic setup with known ground truth features, and claims that SAEs fail to recover features in the "simplest toy setting with fully aligned hyperparameters." In this post, I argue that their setting is not forgiving to SAEs, as the authors imply, but is instead actively hostile in ways that real settings are not.
Introduction
Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? asks two important questions:
This post will be focused on their answer to the first question.
The authors train an SAE on a synthetic toy model, with known ground-truth features, and according to their abstract, SAEs only recover 9% of true features. On the face of it, this seems pretty damning!
Now, earlier toy model experiments did show that SAEs could recover ground truth features, but the authors argue this is because they used a small expansion factor, where there were only twice as many "true" features as there were model dimensions, and that this is unrealistic. Their setup uses an expansion factor of 32, which matches real-world use cases much more closely.
So, how does their setup work?
They don't actually construct a synthetic neural network; instead, they create synthetic activations directly. First, they create their feature vectors, by sampling 3200 vectors from the unit sphere in 100-dimensional space. To create an activation, they just select about 20 feature vectors at random, scale each active feature by a magnitude drawn from LogNormal(0, 0.25),[1] and sum the scaled vectors.
They test two different regimes, one where each feature is equally likely to be active, and another where feature activation frequency has a long tail, so some features fire more often than others. (We'll focus on the constant regime for the most part)
They match their SAE to the setup. There are about 20 expected active features, so they target L0 = 20. There are 3200 true features, so they create 3200 SAE latents. Can the SAE, when given every possible advantage, recover the true features?
These results are, again, pretty damning! They claim that the SAE achieves high reconstruction accuracy while only recovering 3/3200 true features—personally, I disagree that 0.67 explained variance is "high reconstruction accuracy".[2] However, the SAE clearly only does slightly better than random guesswork at recovering the true features.
Now, one obvious objection is that L0/d_model here is 0.2, which is unusually dense. I'd expect there to be a lot of interference here, and in the real world L0/d_model is probably going to be closer to 0.05. Also, 100 dimensions is just not a lot, and maybe you can't really get blessing of dimensionality from it like you would in real models.
A less obvious objection is that if you scroll all the way down to Appendix B, they describe some other SAE architectures they used for this task, like TopK and Matryoshka-on-BatchTopK.
In contrast to BatchTopK, TopK recovers 99.9% of features.[3]
So, I did what anyone would do, and started replicating the paper. They didn't actually include all the hyperparameters they used,[4] so I made reasonable assumptions.
I was able to replicate this result! TopK recovered 99.6% of features,[5] and the other architectures failed miserably.[6]
Why does BatchTopK fail?
This is kind of strange, if you think about it. TopK and BatchTopK are the same architecture, with the only difference being the selection rule.[7] Such a dramatic gap is really weird.
Peeking a bit closer at my training run, in every batch, most latents never fired. That means a small group of monopolist latents were dominating every time. By step 5,000, the same ~6% of latents took over, and they remained in charge for the remaining 295,000 steps. In comparison, with TopK, after the first 20k steps, >99% of latents activated on every batch.
This is the standard dead latents problem: There's no gradient on latents that aren't selected, so once they fall behind in performance it's hard for them to improve. A small group of generalist latents dominate before specialists have the chance to form. The standard fix, of course, is an auxiliary loss, AuxK. But even with AuxK, the dead latent fraction barely budges. I tried sweeping different hyperparameters, but that didn't help either. I even tried resampling, and that did nothing.
Looking closer at the difference between TopK and BatchTopK gives us a hint. In theory, TopK should also be suffering from this problem. The fact that it doesn't suggests that the problem lies in sample-level variance. Certain samples in BatchTopK could receive a significantly larger amount of latents, while other samples are starved. An obvious thought is that input vectors with greater magnitude get more latents assigned to them than those with less magnitude. Checking the data, this does appear to be the case. So what if we normalized the input?
Bussmann's original BatchTopK implementation actually normalized per-sample, folding back in the original magnitude and mean during reconstruction.[8] I copied this method, and lo and behold! BatchTopK achieved 99.3% recovery, comparable to (my) TopK's 99.6% recovery. This also fixes JumpReLU with 97.2% feature recovery.
But while this was the original design, per-sample normalization isn't common anymore! GemmaScope and SAELens don't use it (although they do exclude outliers like BOS tokens), but they still work, without such drastic dead latent rates.
Despite the paper's claims, this synthetic model doesn't demonstrate the failure of SAEs to capture real features. Real SAEs don't have such drastic dead latent rates. Real SAEs are much better at reconstructing in-distribution data than out-of-distribution data. The SAE trained on synthetic data seems to be a fully general vector encoder, reconstructing random vectors with matching norms just as well as it does in-distribution data.
So all this model demonstrates is that this particular synthetic toy model is hostile to many SAE architectures (but not TopK). This is a weakness of these SAE architectures, in that they are not very robust to these types of environments, but this particular failure mode has an obvious signature (the dead latent rates).
Why is this synthetic environment hostile to SAEs?
I hope that I have demonstrated sufficient evidence of the failure of Sanity Checks' synthetic environment. The rest of this post will be dedicated to my personal explorations trying to understand precisely what it is about their environment that causes this divergence from real SAEs.[9]
d_model
Part of it is the density/dimensionality. If we raise d_model up to 800 while keeping the expansion factor of 32 (which means we're now using 25,600 features and latents), L0/d_model shrinks to 0.025, which is much more realistic. Under these conditions, TopK with auxk gets 99.9% recovered, and BatchTopK with auxk gets 53.5% with 46.3% dead latents. JumpReLU recovers 28.3% with 71.6% dead. So, it explains some of the difference, but not all of it.
While resampling didn't work at d_model = 100, when we try it at d_model = 800, it succeeds, raising BatchTopK's recovery to 85.6%. Anthropic resamples using 0.2x average active encoder weight norm, but when we try it with 1x instead, it actually raises recovery to 97.8%, close to TopK's.
So does this mean we've solved it? No! Real SAEs don't need resampling to avoid killing off half their latents.
I could continue to raise d_model, and I expect with More Dakka I could narrow the gap further. But GPT-2 only has 768 dimensions, and SAEs trained on it work just fine (in that they don't have abysmal dead latent rates). So, we should look for other divergences from real models that our synthetic setup has, while keeping d_model at 800.
Magnitude Variance
An obvious guess is that our environment has unrealistic variance in per-sample magnitude, causing architectures like BatchTopK to perform poorly, and that real models have lower variance. While real models do have outlier tokens like BOS, SAEs typically exclude these, so I sampled GPT-2 activations on OpenWebText at various layers, excluding BOS. I subtracted out the dataset mean, since the decoder bias will usually be pretty close to it and is often initialized to the mean (or the geometric median).
Layer
Coefficient of Variation
GPT-2 L1
.224
GPT-2 L4
.181
GPT-2 L7
.140
GPT-2 L10
.129
Toy model (sigma 0.25)
.131
Toy model (sigma 0.5)
.185
As you can see, our toy model falls within the range of real GPT-2 activations, and in fact is on the low end. This means magnitude variance cannot explain the gap.
Feature Loudness
Why doesn't everybody normalize their SAE inputs? When you normalize, you're tossing away the magnitude, which could perhaps serve as an important signal. For example, maybe some features naturally activate with more magnitude than others, and total magnitude serves as a signal for that. So, maybe giving different features different natural "loudness" would improve performance.[10]
To simulate this, I draw a "loudness" variable for every feature from LogNormal(0, 0.5 sigma). Then, when creating activations, I scale each active feature not only by their activation strength (drawn from LogNormal(0, 0.25 sigma)), but also their loudness.
Running this with BatchTopK, we get 19.6% of features recovered, with 78.3% dead latents. Instead of improving performance, it actively degrades it. So feature loudness alone cannot explain the gap. (Also, my motivation for this was really dumb, see footnote 10 for why)
More Features!
Real LLMs might have many more features than their SAEs have latents. Perhaps, this gives latents more opportunities to grab onto "free" features. Alternatively, this might be distracting, making each individual feature appear less often, or trapping latents in an unfortunate middle ground where they can never lock onto a feature closely enough to converge on it.
So, I ran BatchTopK SAEs on setups with 4x and 2x more features than latents. We're still at 800 dimensions, 25,600 latents, and on average 20 features are active at a time.
On 4x, we recover 0.1% of features (This, of course, means 0.4% of latents captured features, since there are 4x more features than latents), with 58% dead latents. Even per-sample normalization[11] performs similarly poorly, recovering only 0.2% of features. They both reconstruct data only slightly better than noise.
On 2x, we recover 16.3% of features (which means 32.6% of latents capture features), with 59.6% dead latents. Per-sample normalization actually works this time, though, recovering 49.1% of features (so, 98.2% useful latents), and only 1.7% dead latents.
So, having more features turns out to be distracting, in this case.[12]
More Active Features
Real LLMs might have higher average active feature counts than the L0s of their SAEs. What if we tried doubling the active feature count to 40, while keeping target L0 the same?
Unfortunately, this also did not improve performance. BatchTopK only recovered 25.2% of features, compared to 53.5% previously.
Other Things I Tried
I experimented with many other things over the course of my research. Here is a non-exhaustive list of them:
Magnitude Variance Revisited
At this point, I was feeling pretty frustrated. I knew that the toy was exhibiting failure modes that real SAEs don't typically have, but I wasn't sure entirely why.
As I mentioned in footnote 1, when I was first running these experiments, I incorrectly assumed that the active feature magnitude was drawn with a sigma of 0.5, or a variance of 0.25. I later discovered my mistake, and reran my experiments to match.
But this made me even more confused. Originally, when I was running these experiments, I found that increasing active feature count to 40 actually improved performance, but with sigma 0.25, as shown above, it was harmful to it! This, and other such oddities, made me revisit the question of magnitude variance.
It turns out that there are two types of magnitude variance to worry about:
Sample-level magnitude variance is the variance in the magnitude of activations.
Feature-level magnitude variance is the variance in the magnitude of a feature, over every activation of that feature.
Under our current setup, these two are correlated. But it seems very possible that sample-level magnitude variance is damaging to architectures like BatchTopK SAE, while feature-level magnitude variance is generally protective, making it easier for latents to capture features.[13]
We can fairly easily measure an objective sample-level magnitude variance for GPT-2, and we already did so. But, in order to find out what the feature-level magnitude variance is, we would need to know what the features are. Thankfully, there are existing GPT-2 SAEs, so I could just base the variance off of Bloom's SAEs.[14]
This is, of course, creating an environment for an SAE by using metrics pulled from an SAE, which is cheating. Thus, I wouldn't consider this to serve as proper evidence that SAEs recover true features, and this is overall not a very principled approach.
GPT-2 Layer 7 has a sample-level CV of 0.140, and a feature-level CV of 1.16.[15] We can target a specific sample-level CV by normalizing and rescaling to a magnitude from a LogNormal distribution with the correct mean and CV. We can target a specific feature-level CV by changing the sigma. However, rescaling does also lower feature-level CV.
I targeted .15 sample-level CV and sigma 1, which, after rescaling, ends up at 1.05 feature-level CV. Training a BatchTopK SAE, we get 96.1% feature recovery with 2.7% of latents dead, without having to do anything unconventional in the training process like per-sample normalization or resampling.
Limitations
When I tried targeting CV .18 and sigma 1, 31.2% of my latents died. GPT-2 Layer 4 has a CV of .18, and Bloom's L4 SAE has 0.3% dead latents. So there's something else that's protective about reality, and we only have one piece of the full story.
The main takeaway here is that it is very challenging to construct a realistic synthetic environment for LLMs, and we should be skeptical of evidence from toy models. My modified synthetic environment doesn't match reality. Sanity Checks' certainly doesn't. SynthSAEBench is more realistic than either of ours, and interestingly enough it also reports BatchTopK SAEs struggling with dead latents at low L0, although I'm not sure this is the same phenomenon.
Takeaways for real SAEs(?)
SAEs can never fail, they can only be failed. So how can we avoid failing them?
I originally assumed the 0.25 here meant variance, but then I asked the authors and actually it meant sigma. Unfortunately, this meant I had a bunch of experiments run under 0.5 sigma, which I then had to repeat with 0.25 sigma. This largely got similar results, although as a happy accident, the parts where it differed were very enlightening. All the results in this post are from 0.25 sigma unless stated otherwise.
In the paper, the authors define a feature being "recovered" as having an SAE latent with >=0.80 cosine similarity. I use this convention throughout the post.
To my chagrin, they don't mention this important finding outside the appendix. However, the ICML listing for their paper displays an updated abstract, which does mention it. For the life of me, though, I can't find any version of their paper which contains this updated abstract, so it's specifically only in the ICML listing for the paper.
The paper itself says, "We will publicly release all code, hyperparameters, and trained SAEs to ensure full reproducibility", but they only say this in Case Study #2, which isn't what I'm studying. The Case Study #2 SAEs with their training configs are on HuggingFace, but https://github.com/AntonKorznikov/random_sae, which is presumably where the code would be released, is empty.
Why only 99.6% instead of 99.9%? Probably down to slight differences in implementation and hyperparameters (perhaps they trained for longer). It's a very small difference and doesn't matter.
I didn't bother with Matryoshka here because it's built on top of BatchTopK and fails similarly.
In TopK, you select the top k latent activations in each sample, so targeting L0 is ~guaranteed by the architecture. BatchTopK has you select the top batch size ⋅ k latent activations over the entire training batch, which allows for variance in the active latent count (since L0 targets the average latents active per sample, some samples may have more active latents and some samples may have fewer).
Bussmann uses LayerNorm-style normalization, not L2, and L2 is honestly more justifiable here given my reasoning for doing normalization. So I tried it, and L2 normalization gets very similar performance, recovering 99.2% of features on BatchTopK and 98.9% on JumpReLU.
Throughout this process I frequently worried that I had found something which failed to recover features but looked like what a real healthy SAE should look like. However, most failing configurations had extremely high dead latent rates, typically close to the fraction of features not recovered, and the ones that didn't had other signs of failure, like low EV, or a failure to reconstruct real data better than noise.
This is a stupid theory and should serve as an example of how not to do reasoning. In reality, SAEs do not use the total magnitude to decide what latents to activate. I don't know what I was thinking at the time, but nevertheless I am honestly reporting what my motivation was.
I used L2 normalization for this.
For fun, I also tried twice as many latents as features (51,200 latents, 25,600 features), which got 58.4% feature recovery and 70.5% dead latents.
Based on my experiments, this does seem to be approximately the case.
Bloom reads SAE activations from the start of each layer (hook_resid_pre) rather than the end. This is no longer conventional, so to avoid confusion I relabel them by the end of the previous layer, which is the same set of activations. For example, when I mention his L4 SAE, I'm referring to what he called his L5 SAE.
Specifically, I measured the median latent's activation magnitude CV. This might be unusually high because Bloom's SAEs are vanilla SAEs, but I don't know.