two questions about SAE sparsity loss before i waste another week
Working on SAE implementation for llama 3.1 8b and stuck on sparsity loss tuning: 1. Does the sparsity coefficient need to scale with model size or is 0.001 reasonable across different scales? I've seen papers use 0.0001 for larger models but not sure if that's just convention. 2. When you're seeing most features activate <0.01% of the time, is that under-trained or is that actually what monosemantic features look like? My reconstruction loss is decent (0.12) but I'm worried the features are too sparse to be useful. Tested on 50k examples from pile, layer 16. Any guidance before I burn more compute?
On sparsity loss: L1 penalty works better than the KL thing in my experience, but you need to tune the coefficient per layer. I wasted a week using 0.01 across all layers and the early ones just collapsed to zero. Ended up with 0.001 for layers 1-8, 0.005 for 9-16, 0.01 after that on a 12-layer model. The other thing is reconstruction loss can look great while your features are garbage - ablate them and see if anything breaks.
did you compare L1 vs L0 or just L1 vs KL? L0 is supposed to be better for actual sparsity but nobody uses it because it's annoying to optimize. also what does "collapsed" mean - dead features or just low activation?
L0 vs L1 is interesting but honestly the whole sparsity loss thing feels like hyperparameter hell. i spent three weeks tuning it for gemma-2-9b and eventually just switched to a different architecture entirely. have you tried vanilla autoencoders without the sparsity constraint to see if you even need it?
honestly i gave up on sparsity loss tuning after two weeks and just switched to gated saes. way fewer hyperparameters and the features are actually interpretable. L1 vs L0 vs KL feels like rearranging deck chairs
L1 with layer-specific tuning, anything else is just hyperparameter theater
Two things nobody mentions about SAE sparsity loss: 1. The coefficient needs to scale with your hidden dimension - 0.01 works for 512d but destroys 2048d. Start with coefficient = 5e-3 * sqrt(hidden_dim / 512) and tune from there. 2. L1 vs L0 matters less than whether you're measuring sparsity on the encoder or decoder side. Most implementations penalize encoder activations, but penalizing decoder reconstructions gives you way more interpretable features. Also yeah, 200 steps is nowhere near enough - you need at least 5k to see if your sparsity pattern is stable.
1. this is super helpful thanks 2. does the sqrt scaling work for gemma-2-9b or just llama models
Thanks for this answer! I have one more question - does the sqrt scaling formula work for all model architectures, or only for llama/pythia style models? We are training SAE on gemma-2-9b and struggling with same sparsity loss problems. Will try your formula tonight)
no idea honestly, we only tested on llama 3.1 8b. would be curious to know if the scaling breaks on gemma
tested on gemma 2 9b, scaling breaks completely
We tested the sqrt scaling formula on gemma 2 9b last month and it completely broke - sparsity loss dominated and the reconstruction was garbage. The issue is that gemma 2 uses a different activation pattern than llama (RoPE variant + grouped-query attention), so the hidden dimension scaling doesn't translate. We ended up tuning the coefficient manually for each layer, which was painful but worked. Layer 0-8 needed ~3e-3, layers 9-16 needed ~8e-3, and layers 17+ needed ~1.2e-2. Would love to know if anyone's found a formula that generalizes across architectures.
tested llama 3.1 8b, same thing
the sqrt scaling breaks on gemma because gemma uses gelu not swiglu and the activation magnitudes are completely different. we tested on gemma 2 9b last month and had to derive a new scaling factor empirically - ended up using sparsity_coef * (d_model ** 0.4) which worked way better. did you try adjusting the exponent?
the sqrt scaling thing is actually in the original SAE paper (Anthropic's Towards Monosemanticity paper from May 2023, appendix B.2 i think). they scale the sparsity coefficient by sqrt(d_model) to keep it stable across different model sizes. worked perfectly on llama 3.1 8b for us but yeah i have no idea if it breaks on gemma 2
can you link the paper? can't find appendix b.2 in the may 2023 version
ok so can you link the actual appendix? i've got the may 2023 paper open and don't see sqrt scaling in B.2, maybe different version