Add logic to train JumpReLU SAEs

[anthonyduong9]

Description

Adds logic to train JumpReLU SAEs. JumpReLU is a state-of-the-art architecture, so users should be able to train JumpReLU SAEs. Currently, they can load pre-trained JumpReLU SAEs and perform inference with them.

Fixes #330

Type of change

Please delete options that are not relevant.

• [x] Bug fix (non-breaking change which fixes an issue)
• [x] New feature (non-breaking change which adds functionality)

Checklist:

• [x] I have commented my code, particularly in hard-to-understand areas
• [ ] I have made corresponding changes to the documentation
• [x] My changes generate no new warnings
• [x] I have added tests that prove my fix is effective or that my feature works
• [x] New and existing unit tests pass locally with my changes
• [x] I have not rewritten tests relating to key interfaces which would affect backward compatibility

You have tested formatting, typing and unit tests (acceptance tests not currently in use)

• [x] I have run make check-ci to check format and linting. (you can run make format to format code if needed.)

Performance Check.

If you have implemented a training change, please indicate precisely how performance changes with respect to the following metrics:

• [ ] L0
• [ ] CE Loss
• [ ] MSE Loss
• [ ] Feature Dashboard Interpretability

Please links to wandb dashboards with a control and test group.

[anthonyduong9]

Some minor comments but core logic looks great! I'll try training a JumpReLU SAE to make sure this works.

It feels odd that the l1_coefficient is used for the JumpReLU sparsity loss coefficient, and that it gets logged as l1_loss. Maybe in a follow-up PR we should rename l1_coefficient to sparsity_coefficient instead? I added a PR #357 which lets the forward pass log any losses it wants, so if that gets merged, it would probably make sense to call the loss sparsity_loss since there's not technically any L1 loss in jumprelu.

Thanks! Let me know if there's anything else you need from me.

And I'm guessing unless you think we should keep the l1_coefficient field for the other architecture(s), and add a sparsity_coefficient that's just used for JumpReLU, your suggestion makes sense.

[chanind]

Trying a test-run now, but one more thing is that the typing for architecture for LanguageModelSAERunnerConfig needs to be updated as well - it's currently set to "standard" | "gated" so throws a typing error when setting "jumprelu".

[chanind]

It seems like the general paradigm of most of the recent SAE architectures is having some sort of sparsity-inducing loss, either l1 in normal SAEs or l0 in jumprelu (or nothing for topk), and then an optional auxiliary loss. I'd be for just calling it sparsity_coefficient to handle the architecture differences, but should be discussed in another follow-up issue probably.

[chanind]

I tried training for 1B tokens with sparsity coeff of 1e-3 (which I think is reasonable based on the gemma scope paper). Looks like L0 is still coming down by the end of training (Gemma Scope trains on 4B tokens), so I'm assuming this would plateau at a reasonable loss if I kept going.

[anthonyduong9]

LGTM! Great work with this!

Thanks!