chore: Test JumpReLU/Gated SAE and refactor sae forward #328
Conversation
| x = self.run_time_activation_norm_fn_in(x) | ||
|
|
||
| # apply b_dec_to_input if using that method. | ||
| sae_in = self.hook_sae_input(x - (self.b_dec * self.cfg.apply_b_dec_to_input)) |
There was a problem hiding this comment.
we were calling hook_sae_input() here after subtracting b_dec, but in other encode() variants we call hook_sae_input() after reshape_fn_in() instead. I assumed this difference was a bug and not intentional.
| @ self.W_dec | ||
| + self.b_dec, | ||
| d_head=self.d_head, | ||
| ) |
There was a problem hiding this comment.
I assumed the code in these codepaths is trying to exactly recreate all the encode_x() variants followed by decode(), but is duplicating code just to avoid triggering hooks.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #328 +/- ##
==========================================
+ Coverage 62.43% 63.38% +0.95%
==========================================
Files 25 25
Lines 3144 3111 -33
Branches 550 398 -152
==========================================
+ Hits 1963 1972 +9
+ Misses 1077 1038 -39
+ Partials 104 101 -3 ☔ View full report in Codecov by Sentry. |
Description
This PR adds test coverage to the JumpReLU and GatedSAE encode methods.
In doing this, I realized there's a lot of duplication between all the encode variants and cleaned that up as well. I think there were some potential minor bugs in this duplication, for instance in
forward(), when adding error term, we calledrun_time_activation_norm_fn_out()afterreshape_fn_out(), but we do the opposite indecode().It looks like most of the duplication between the
errorsection offorward()and the normalencode() / decode()was there just to avoid triggering hooks, so I added a contextmanager_disable_hooks()which is used to disable hooks in this branch of the code while reusing our existingencode() / decode()methods. This should mean we don't need to worry about these duplicated codepaths being slightly different.If the refactor is out of scope, I can revert the changes to
sae.pyand just leave the test coverage.Fixes #323
Type of change
Please delete options that are not relevant.
Checklist:
You have tested formatting, typing and unit tests (acceptance tests not currently in use)
make check-cito check format and linting. (you can runmake formatto format code if needed.)