sae-lens 6.22.2__tar.gz → 6.23.0__tar.gz

{sae_lens-6.22.2 → sae_lens-6.23.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.22.2
+Version: 6.23.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE

{sae_lens-6.22.2 → sae_lens-6.23.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.22.2"
+version = "6.23.0"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.22.2"
+__version__ = "6.23.0"
 
 import logging
 

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/loading/pretrained_sae_loaders.py

@@ -753,10 +753,14 @@ def get_dictionary_learning_config_1_from_hf(
     activation_fn = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
     activation_fn_kwargs = {"k": trainer["k"]} if activation_fn == "topk" else {}
 
+    architecture = "standard"
+    if trainer["dict_class"] == "GatedAutoEncoder":
+        architecture = "gated"
+    elif trainer["dict_class"] == "MatryoshkaBatchTopKSAE":
+        architecture = "jumprelu"
+
     return {
-        "architecture": (
-            "gated" if trainer["dict_class"] == "GatedAutoEncoder" else "standard"
-        ),
+        "architecture": architecture,
         "d_in": trainer["activation_dim"],
         "d_sae": trainer["dict_size"],
         "dtype": "float32",
@@ -905,9 +909,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     )
     encoder = torch.load(encoder_path, map_location="cpu")
 
+    W_enc = encoder["W_enc"] if "W_enc" in encoder else encoder["encoder.weight"].T
+    W_dec = encoder["W_dec"] if "W_dec" in encoder else encoder["decoder.weight"].T
+
     state_dict = {
-        "W_enc": encoder["encoder.weight"].T,
-        "W_dec": encoder["decoder.weight"].T,
+        "W_enc": W_enc,
+        "W_dec": W_dec,
         "b_dec": encoder.get(
             "b_dec", encoder.get("bias", encoder.get("decoder_bias", None))
         ),
@@ -915,6 +922,8 @@ def dictionary_learning_sae_huggingface_loader_1(
 
     if "encoder.bias" in encoder:
         state_dict["b_enc"] = encoder["encoder.bias"]
+    if "b_enc" in encoder:
+        state_dict["b_enc"] = encoder["b_enc"]
 
     if "mag_bias" in encoder:
         state_dict["b_mag"] = encoder["mag_bias"]
@@ -923,6 +932,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     if "r_mag" in encoder:
         state_dict["r_mag"] = encoder["r_mag"]
 
+    if "threshold" in encoder:
+        threshold = encoder["threshold"]
+        if threshold.ndim == 0:
+            threshold = torch.full((W_enc.size(1),), threshold)
+        state_dict["threshold"] = threshold
+
     return cfg_dict, state_dict, None
 
 

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/pretrained_saes.yaml

@@ -14959,3 +14959,81 @@ goodfire-llama-3.1-8b-instruct:
     path: Llama-3.1-8B-Instruct-SAE-l19.pth
     l0: 91
     neuronpedia: llama3.1-8b-it/19-resid-post-gf
+
+saebench_gemma-2-2b_width-2pow12_date-0108:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: adamkarvonen/saebench_gemma-2-2b_width-2pow12_date-0108
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_0_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_1_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_2_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_3_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_4_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_5_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_5
+
+saebench_gemma-2-2b_width-2pow14_date-0107:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: canrager/saebench_gemma-2-2b_width-2pow14_date-0107
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_0_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_1_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_2_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_3_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_4_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_5_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_5
+
+saebench_gemma-2-2b_width-2pow16_date-0107:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: canrager/saebench_gemma-2-2b_width-2pow16_date-0107
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_0_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_1_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_2_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_3_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_4_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_5_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_5

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/saes/gated_sae.py

@@ -89,7 +89,7 @@ class GatedSAE(SAE[GatedSAEConfig]):
     @torch.no_grad()
     def fold_W_dec_norm(self):
         """Override to handle gated-specific parameters."""
-        W_dec_norms = self.W_dec.norm(dim=-1).unsqueeze(1)
+        W_dec_norms = self.W_dec.norm(dim=-1).clamp(min=1e-8).unsqueeze(1)
         self.W_dec.data = self.W_dec.data / W_dec_norms
         self.W_enc.data = self.W_enc.data * W_dec_norms.T
 
@@ -217,7 +217,7 @@ class GatedTrainingSAE(TrainingSAE[GatedTrainingSAEConfig]):
     @torch.no_grad()
     def fold_W_dec_norm(self):
         """Override to handle gated-specific parameters."""
-        W_dec_norms = self.W_dec.norm(dim=-1).unsqueeze(1)
+        W_dec_norms = self.W_dec.norm(dim=-1).clamp(min=1e-8).unsqueeze(1)
         self.W_dec.data = self.W_dec.data / W_dec_norms
         self.W_enc.data = self.W_enc.data * W_dec_norms.T
 

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/saes/jumprelu_sae.py

@@ -167,8 +167,8 @@ class JumpReLUSAE(SAE[JumpReLUSAEConfig]):
         # Save the current threshold before calling parent method
         current_thresh = self.threshold.clone()
 
-        # Get W_dec norms that will be used for scaling
-        W_dec_norms = self.W_dec.norm(dim=-1)
+        # Get W_dec norms that will be used for scaling (clamped to avoid division by zero)
+        W_dec_norms = self.W_dec.norm(dim=-1).clamp(min=1e-8)
 
         # Call parent implementation to handle W_enc, W_dec, and b_enc adjustment
         super().fold_W_dec_norm()
@@ -325,8 +325,8 @@ class JumpReLUTrainingSAE(TrainingSAE[JumpReLUTrainingSAEConfig]):
         # Save the current threshold before we call the parent method
         current_thresh = self.threshold.clone()
 
-        # Get W_dec norms
-        W_dec_norms = self.W_dec.norm(dim=-1).unsqueeze(1)
+        # Get W_dec norms (clamped to avoid division by zero)
+        W_dec_norms = self.W_dec.norm(dim=-1).clamp(min=1e-8).unsqueeze(1)
 
         # Call parent implementation to handle W_enc and W_dec adjustment
         super().fold_W_dec_norm()

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/saes/sae.py

@@ -484,7 +484,7 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
     @torch.no_grad()
     def fold_W_dec_norm(self):
         """Fold decoder norms into encoder."""
-        W_dec_norms = self.W_dec.norm(dim=-1).unsqueeze(1)
+        W_dec_norms = self.W_dec.norm(dim=-1).clamp(min=1e-8).unsqueeze(1)
         self.W_dec.data = self.W_dec.data / W_dec_norms
         self.W_enc.data = self.W_enc.data * W_dec_norms.T
 

{sae_lens-6.22.2 → sae_lens-6.23.0}/sae_lens/saes/topk_sae.py

@@ -531,7 +531,7 @@ def _fold_norm_topk(
     b_enc: torch.Tensor,
     W_dec: torch.Tensor,
 ) -> None:
-    W_dec_norm = W_dec.norm(dim=-1)
+    W_dec_norm = W_dec.norm(dim=-1).clamp(min=1e-8)
     b_enc.data = b_enc.data * W_dec_norm
     W_dec_norms = W_dec_norm.unsqueeze(1)
     W_dec.data = W_dec.data / W_dec_norms