fix: GatedSAE needs to prevent the model from cheating L0 over steps

[jasonrichdarmawan]

Description

Summary: This PR fix issue where the model cheat the L0 over steps. There is no dependencies that are required for this change. This PR is committed on top of PR #545

PR #545 fix issue where the loss from pi_via_gate also flow to the W_dec and b_dec

https://github.com/jbloomAus/SAELens/blob/3432f0059bc1d90119fe48f88e17ac07aea3a620/sae_lens/saes/gated_sae.py#L196

This PR fix issue where the model cheat the L0 over steps: the loss from pi_via_gate shrinks the encoder in the first step, and then increase the decoder in the next step to adjust to the changed encoder. Repeat this over many steps, the encoder will be very small and decoder will be very large

Note: The above is an assumption. However, the experiment results below show that GatedSAE with normalized decoder does have higher metrics/explained_variance than without normalization.

Fixes #544

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 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 tests

• 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.

WandB dashboard

Control and test group: Tag Gated No Scale vs Norm Aux

[jasonrichdarmawan]