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

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

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

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

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

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

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

{sae_lens-6.22.3 → 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.3 → 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