feat: Topk SAE training

[chanind]

Description

This PR implements topk SAE training by adding the topk auxiliary loss. This PR makes several design choices:

• topk is specified as an architecture for training, so it's similar to gated and jumprelu SAEs. This seems fit the idea of an SAE architecture since it has its own custom training routine and losses, and it seems strange to call "jumprelu" and "gated" architectures, but not topk.
• This PR is based heavily on /~https://github.com/EleutherAI/sae/blob/main/sae/sae.py#L191 rather than the original topk SAEs paper. The Eleuther version requires no additional config options for the auxiliary loss, which means less changes to our existing config, while allowing us to add those config options later if we want. We may also want to refactor the way our config / SAE classes work now that we support so many different architectures anyway.
• This PR makes the following changes to Eleuther's topk aux loss:
  • Removes the variance normalization. Eleuther divides all their losses by variance, but we don't do that anywhere, so I removed this normalization to be consistent
  • Take the mean loss over batches. Eleuther takes the sum, but in our other losses we take a mean over batches rather than sum over batches, so I changed this to be consistent with that.

Our implementation of topk training is likely less efficient than Eleuther's as they use a custom sparse kernel for the SAE decoder (see /~https://github.com/EleutherAI/sae/blob/main/sae/kernels.py). We can try to support something like this in the future, but it will likely require a bit of refactoring before we can support a special decoder kernel just for topk.

I'm currently running some test training runs to make sure things look decent, and will upgrade this PR from draft when those are complete.

Fixes #202

training test run dashboard: https://api.wandb.ai/links/chanind/zju8dl70

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• 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)

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

[codecov]

Codecov Report

Attention: Patch coverage is 84.09091% with 7 lines in your changes missing coverage. Please review.

Project coverage is 72.88%. Comparing base (17506ac) to head (b49fece).
Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
sae_lens/training/training_sae.py	85.18%	2 Missing and 2 partials ⚠️
sae_lens/sae_training_runner.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #370      +/-   ##
==========================================
+ Coverage   72.74%   72.88%   +0.13%     
==========================================
  Files          22       22              
  Lines        3266     3297      +31     
  Branches      431      438       +7     
==========================================
+ Hits         2376     2403      +27     
- Misses        762      764       +2     
- Partials      128      130       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[curt-tigges]

Looks good other than the two comments I made.

[curt-tigges]

This should probably use whatever device is set through the config, rather than being hardcoded to CUDA.

[curt-tigges]

If this function is intended to be imported by other files/classes, probably should not be an underscore/internal function.

[chanind]

It's only imported by its own unit test, it should not be imported externally which is why it has an underscore