Repository: jbloomAus/SAELens
Pull request #432: fix: base k_aux on d_in instead of d_sae in topk aux loss
Author: login: chanind
Created: 2025-02-22T04:15:17Z | Merged: 2025-03-17T19:34:00Z

# Description

This PR fixes a minor issue with our topk aux loss implementation, where we calculate `k_aux` using `d_sae` instead of the correct `d_in`. This likely doesn't make a huge difference in practice, but can't hurt to fix. This PR also calls `detach()` on the residual error before calculating topk aux loss, similar to [dictionary_learning's implementation](https://github.com/saprmarks/dictionary_learning/blob/main/dictionary_learning/trainers/top_k.py#L297). This should help ensure that the aux loss only pulls dead latents towards the SAE error, and doesn't accidentally pull live latents towards dead latents.

This PR also adds a test asserting that our topk aux loss matches the sparsity implementation aside from a normalization factor.

As a side note, we should probably add an optional `aux_coefficient` to the config to further customize this, but this is something that probably makes sense to hold off on until the refactor is done. 

## Type of change

Please delete options that are not relevant.

- [x] 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:

- [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 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 tests

- [x] I have run `make check-ci` to check format and linting. (you can run `make format` to format code if needed.)


Diff:
diff --git a/pyproject.toml b/pyproject.toml
index 3be7adcad..f8de674a1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,6 +58,7 @@ mkdocstrings = "^0.25.2"
 mkdocstrings-python = "^1.10.9"
 tabulate = "^0.9.0"
 ruff = "^0.7.4"
+sparsify = {git = "https://github.com/EleutherAI/sparsify"}
 
 [tool.poetry.extras]
 mamba = ["mamba-lens"]
diff --git a/sae_lens/training/training_sae.py b/sae_lens/training/training_sae.py
index ad2bda018..8bb44b974 100644
--- a/sae_lens/training/training_sae.py
+++ b/sae_lens/training/training_sae.py
@@ -469,31 +469,28 @@ def calculate_topk_aux_loss(
     ) -> torch.Tensor:
         # Mostly taken from https://github.com/EleutherAI/sae/blob/main/sae/sae.py, except without variance normalization
         # NOTE: checking the number of dead neurons will force a GPU sync, so performance can likely be improved here
-        if (
-            dead_neuron_mask is not None
-            and (num_dead := int(dead_neuron_mask.sum())) > 0
-        ):
-            residual = sae_in - sae_out
+        if dead_neuron_mask is None or (num_dead := int(dead_neuron_mask.sum())) == 0:
+            return sae_out.new_tensor(0.0)
+        residual = (sae_in - sae_out).detach()
 
-            # Heuristic from Appendix B.1 in the paper
-            k_aux = hidden_pre.shape[-1] // 2
+        # Heuristic from Appendix B.1 in the paper
+        k_aux = sae_in.shape[-1] // 2
 
-            # Reduce the scale of the loss if there are a small number of dead latents
-            scale = min(num_dead / k_aux, 1.0)
-            k_aux = min(k_aux, num_dead)
+        # Reduce the scale of the loss if there are a small number of dead latents
+        scale = min(num_dead / k_aux, 1.0)
+        k_aux = min(k_aux, num_dead)
 
-            auxk_acts = _calculate_topk_aux_acts(
-                k_aux=k_aux,
-                hidden_pre=hidden_pre,
-                dead_neuron_mask=dead_neuron_mask,
-            )
+        auxk_acts = _calculate_topk_aux_acts(
+            k_aux=k_aux,
+            hidden_pre=hidden_pre,
+            dead_neuron_mask=dead_neuron_mask,
+        )
 
-            # Encourage the top ~50% of dead latents to predict the residual of the
-            # top k living latents
-            recons = self.decode(auxk_acts)
-            auxk_loss = (recons - residual).pow(2).sum(dim=-1).mean()
-            return scale * auxk_loss
-        return sae_out.new_tensor(0.0)
+        # Encourage the top ~50% of dead latents to predict the residual of the
+        # top k living latents
+        recons = self.decode(auxk_acts)
+        auxk_loss = (recons - residual).pow(2).sum(dim=-1).mean()
+        return scale * auxk_loss
 
     def calculate_ghost_grad_loss(
         self,
diff --git a/tests/conftest.py b/tests/conftest.py
index a48d30e71..6bd563ba0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,3 +1,4 @@
+import os
 import random
 import shutil
 from pathlib import Path
@@ -11,6 +12,9 @@
 
 torch.set_grad_enabled(True)
 
+# sparsify's triton implementation breaks in CI, so just disable it
+os.environ["SPARSIFY_DISABLE_TRITON"] = "1"
+
 
 @pytest.fixture(autouse=True)
 def reproducibility():
diff --git a/tests/training/test_training_sae.py b/tests/training/test_training_sae.py
index 82e0b9ac8..34a20880a 100644
--- a/tests/training/test_training_sae.py
+++ b/tests/training/test_training_sae.py
@@ -3,6 +3,7 @@
 
 import pytest
 import torch
+from sparsify import SparseCoder, SparseCoderConfig
 
 from sae_lens.sae import SAE
 from sae_lens.training.training_sae import (
@@ -101,28 +102,74 @@ def test_calculate_topk_aux_acts_k_less_than_dead():
     assert torch.allclose(result, expected)
 
 
-def test_TrainingSAE_calculate_topk_aux_loss():
-    # Create a small test SAE with d_sae=4, d_in=3
+def test_TrainingSAE_topk_aux_loss_matches_unnormalized_sparsify_implementation():
+    d_in = 128
+    d_sae = 192
     cfg = build_sae_cfg(
-        d_in=3,
-        d_sae=4,
+        d_in=d_in,
+        d_sae=d_sae,
         architecture="topk",
+        activation_fn_kwargs={"k": 26},
         normalize_sae_decoder=False,
     )
 
+    sae = TrainingSAE(TrainingSAEConfig.from_sae_runner_config(cfg))
+    sparse_coder_sae = SparseCoder(
+        d_in=d_in, cfg=SparseCoderConfig(num_latents=d_sae, k=26)
+    )
+
+    with torch.no_grad():
+        # increase b_enc so all features are likely above 0
+        # sparsify includes a relu() in their pre_acts, but
+        # this is not something we need to try to replicate.
+        sae.b_enc.data = sae.b_enc + 100.0
+        # make sure all params are the same
+        sparse_coder_sae.encoder.weight.data = sae.W_enc.T
+        sparse_coder_sae.encoder.bias.data = sae.b_enc
+        sparse_coder_sae.b_dec.data = sae.b_dec
+        sparse_coder_sae.W_dec.data = sae.W_dec  # type: ignore
+
+    dead_neuron_mask = torch.randn(d_sae) > 0.1
+    input_acts = torch.randn(200, d_in)
+    input_var = (input_acts - input_acts.mean(0)).pow(2).sum()
+
+    sae_out = sae.training_forward_pass(
+        sae_in=input_acts,
+        current_l1_coefficient=0.0,
+        dead_neuron_mask=dead_neuron_mask,
+    )
+    comparison_sae_out = sparse_coder_sae.forward(
+        input_acts, dead_mask=dead_neuron_mask
+    )
+    comparison_aux_loss = comparison_sae_out.auxk_loss.detach().item()
+
+    normalization = input_var / input_acts.shape[0]
+    raw_aux_loss = sae_out.losses["auxiliary_reconstruction_loss"].item()  # type: ignore
+    norm_aux_loss = raw_aux_loss / normalization
+    assert norm_aux_loss == pytest.approx(comparison_aux_loss, abs=1e-4)
+
+
+def test_TrainingSAE_calculate_topk_aux_loss():
+    # Create a small test SAE with d_sae=3, d_in=4
+    cfg = build_sae_cfg(
+        d_in=4,
+        d_sae=3,
+        architecture="topk",
+        normalize_sae_decoder=False,
+    )
     sae = TrainingSAE(TrainingSAEConfig.from_sae_runner_config(cfg))
 
     # Set up test inputs
     hidden_pre = torch.tensor(
-        [[1.0, -2.0, 3.0, -4.0], [1.0, 0.0, -3.0, -4.0]]  # batch size 2
+        [[1.0, -2.0, 3.0], [1.0, 0.0, -3.0]]  # batch size 2
     )
-    sae.W_dec.data = torch.tensor(2 * torch.ones((4, 3)))
-    sae.b_dec.data = torch.tensor(torch.zeros(3))
+    sae.W_dec.data = torch.tensor(2 * torch.ones((3, 4)))
+    sae.b_dec.data = torch.tensor(torch.zeros(4))
 
-    sae_out = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    sae_in = torch.tensor([[2.0, 1.0, 3.0], [5.0, 4.0, 6.0]])
-    # Mark neurons 1 and 3 as dead
-    dead_neuron_mask = torch.tensor([False, True, False, True])
+    sae_out = torch.tensor([[1.0, 2.0, 3.0, 4.0], [4.0, 5.0, 6.0, 7.0]])
+    sae_in = torch.tensor([[2.0, 1.0, 3.0, 4.0], [5.0, 4.0, 6.0, 7.0]])
+    # Mark neurons 1 and 2 as dead
+    dead_neuron_mask = torch.tensor([False, True, True])
 
     # Calculate loss
     loss = sae.calculate_topk_aux_loss(
@@ -133,15 +180,17 @@ def test_TrainingSAE_calculate_topk_aux_loss():
     )
 
     # The loss should:
-    # 1. Select top k_aux=2 (half of d_sae) dead neurons
-    # 2. Decode their activations (should be 2x the sum of the activations of the dead neurons)
-    # thus, (-12, -12, -12), (-8, -8, -8)
-    # and the residual is (1, -1, 0), (1, -1, 0)
-    # Thus, squared errors are (169, 121, 144), (81, 49, 64)
-    # and the sums are (434, 194)
-    # and the mean of these is 314
-
-    assert loss == 314
+    # 1. Select top k_aux=2 (half of d_in=4) dead neurons
+    # 2. Decode their activations (should be 2x the activations of the dead neurons)
+    # For batch 1: dead neurons are [-2.0, 3.0] -> activations [-4.0, 6.0] -> sum 2.0 for each output dim
+    # For batch 2: dead neurons are [0.0, -3.0] -> activations [0.0, -6.0] -> sum -6.0 for each output dim
+    # Residuals are: [1.0, -1.0, 0.0, 0.0], [1.0, -1.0, 0.0, 0.0]
+    # errors are: [1.0, 3.0, 2.0, 2.0], [-7., -5., -6., -6.]
+    # Squared errors are: [1.0, 9.0, 4.0, 4.0], [49.0, 25.0, 36.0, 36.0]
+    # Sum over features: 18.0, 146.0
+    # Mean over batch: 82.0
+
+    assert loss == 82
 
 
 def test_TrainingSAE_forward_includes_topk_loss_with_topk_architecture():


Reviews:
user: login: anthonyduong9
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-01T08:15:50Z

Review comments:
diff_hunk: @@ -473,10 +473,10 @@ def calculate_topk_aux_loss(
             dead_neuron_mask is not None
             and (num_dead := int(dead_neuron_mask.sum())) > 0
         ):
-            residual = sae_in - sae_out
+            residual = (sae_in - sae_out).detach()
 
             # Heuristic from Appendix B.1 in the paper
-            k_aux = hidden_pre.shape[-1] // 2
+            k_aux = sae_in.shape[-1] // 2
path: sae_lens/training/training_sae.py
user: login: anthonyduong9
body: Since we're changing `calculate_topk_aux_loss()`, want to also refactor it to use an early return?

It'd be

```
    def calculate_topk_aux_loss(
        self,
        sae_in: torch.Tensor,
        sae_out: torch.Tensor,
        hidden_pre: torch.Tensor,
        dead_neuron_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        # Mostly taken from https://github.com/EleutherAI/sae/blob/main/sae/sae.py, except without variance normalization
        # NOTE: checking the number of dead neurons will force a GPU sync, so performance can likely be improved here
        if dead_neuron_mask is None or int(dead_neuron_mask.sum()) == 0:
            return sae_out.new_tensor(0.0)

        num_dead = int(dead_neuron_mask.sum())
        residual = (sae_in - sae_out).detach()

        # Heuristic from Appendix B.1 in the paper
        k_aux = sae_in.shape[-1] // 2

        # Reduce the scale of the loss if there are a small number of dead latents
        scale = min(num_dead / k_aux, 1.0)
        k_aux = min(k_aux, num_dead)

        auxk_acts = _calculate_topk_aux_acts(
            k_aux=k_aux,
            hidden_pre=hidden_pre,
            dead_neuron_mask=dead_neuron_mask,
        )

        # Encourage the top ~50% of dead latents to predict the residual of the
        # top k living latents
        recons = self.decode(auxk_acts)
        auxk_loss = (recons - residual).pow(2).sum(dim=-1).mean()
        return scale * auxk_loss
```

, which has a lot less indentation/is a lot easier to read.
created_at: 2025-03-01T08:07:53Z
author_association: COLLABORATOR
diff_hunk: @@ -101,28 +102,68 @@ def test_calculate_topk_aux_acts_k_less_than_dead():
     assert torch.allclose(result, expected)
 
 
-def test_TrainingSAE_calculate_topk_aux_loss():
-    # Create a small test SAE with d_sae=4, d_in=3
+def test_TrainingSAE_topk_aux_loss_matches_unnormalized_sparsify_implementation():
     cfg = build_sae_cfg(
-        d_in=3,
-        d_sae=4,
+        d_in=128,
+        d_sae=192,
path: tests/training/test_training_sae.py
user: login: anthonyduong9
body: Want to extract all the `128`s and `192`s to variables `d_in` and `d_sae`, respectively?
created_at: 2025-03-01T08:10:11Z
author_association: COLLABORATOR
diff_hunk: @@ -101,28 +102,68 @@ def test_calculate_topk_aux_acts_k_less_than_dead():
     assert torch.allclose(result, expected)
 
 
-def test_TrainingSAE_calculate_topk_aux_loss():
-    # Create a small test SAE with d_sae=4, d_in=3
+def test_TrainingSAE_topk_aux_loss_matches_unnormalized_sparsify_implementation():
     cfg = build_sae_cfg(
-        d_in=3,
-        d_sae=4,
+        d_in=128,
+        d_sae=192,
         architecture="topk",
+        activation_fn_kwargs={"k": 26},
         normalize_sae_decoder=False,
     )
 
+    sae = TrainingSAE(TrainingSAEConfig.from_sae_runner_config(cfg))
+    comparison_sae = SparseCoder(d_in=128, cfg=SparseCoderConfig(num_latents=192, k=26))
path: tests/training/test_training_sae.py
user: login: anthonyduong9
body: I think `sparse_coder` might be a better name than `comparison_sae` as someone reading the test could quickly tell it's an instance of `SparseCoder`.
created_at: 2025-03-01T08:15:44Z
author_association: COLLABORATOR

PR comments:
user: login: curt-tigges
created_at: 2025-03-11T19:35:37Z
body: @chanind could you apply these changes to the new TopKTrainingSAE class in the refactor PR (once merged) as well?
author_association: CONTRIBUTOR
user: login: chanind
created_at: 2025-03-17T19:05:37Z
body: Will do! Once I manage to fix this strange triton error in tests...
author_association: COLLABORATOR
user: login: codecov[bot]
created_at: 2025-03-17T19:09:13Z
body: ## [Codecov](https://app.codecov.io/gh/jbloomAus/SAELens/pull/432?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joseph+Bloom) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 74.10%. Comparing base [(`4f4a7b7`)](https://app.codecov.io/gh/jbloomAus/SAELens/commit/4f4a7b7265bae33a5e28adb9263d113b8b6c20dd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joseph+Bloom) to head [(`7ad80bf`)](https://app.codecov.io/gh/jbloomAus/SAELens/commit/7ad80bfc522ce2dc8f616b7b283e38dcdb0d3886?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joseph+Bloom).
> Report is 11 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #432      +/-   ##
==========================================
+ Coverage   74.06%   74.10%   +0.04%     
==========================================
  Files          19       19              
  Lines        3165     3174       +9     
  Branches      457      460       +3     
==========================================
+ Hits         2344     2352       +8     
- Misses        665      667       +2     
+ Partials      156      155       -1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/jbloomAus/SAELens/pull/432?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joseph+Bloom).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Joseph+Bloom).

<details><summary>🚀 New features to boost your workflow: </summary>

- ❄ [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE