sae-lens 6.29.1__tar.gz → 6.31.0__tar.gz

{sae_lens-6.29.1 → sae_lens-6.31.0}/PKG-INFO

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

{sae_lens-6.29.1 → sae_lens-6.31.0}/pyproject.toml

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

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/__init__.py

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

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/config.py

@@ -82,6 +82,7 @@ class LoggingConfig:
     log_to_wandb: bool = True
     log_activations_store_to_wandb: bool = False
     log_optimizer_state_to_wandb: bool = False
+    log_weights_to_wandb: bool = True
     wandb_project: str = "sae_lens_training"
     wandb_id: str | None = None
     run_name: str | None = None
@@ -107,7 +108,8 @@ class LoggingConfig:
             type="model",
             metadata=dict(trainer.cfg.__dict__),
         )
-        model_artifact.add_file(str(weights_path))
+        if self.log_weights_to_wandb:
+            model_artifact.add_file(str(weights_path))
         model_artifact.add_file(str(cfg_path))
         wandb.log_artifact(model_artifact, aliases=wandb_aliases)
 
@@ -557,6 +559,12 @@ class CacheActivationsRunnerConfig:
                 context_size=self.context_size,
             )
 
+        if self.context_size > self.training_tokens:
+            raise ValueError(
+                f"context_size ({self.context_size}) is greater than training_tokens "
+                f"({self.training_tokens}). Please reduce context_size or increase training_tokens."
+            )
+
         if self.new_cached_activations_path is None:
             self.new_cached_activations_path = _default_cached_activations_path(  # type: ignore
                 self.dataset_path, self.model_name, self.hook_name, None

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/evals.py

@@ -335,7 +335,7 @@ def get_downstream_reconstruction_metrics(
 
     batch_iter = range(n_batches)
     if verbose:
-        batch_iter = tqdm(batch_iter, desc="Reconstruction Batches")
+        batch_iter = tqdm(batch_iter, desc="Reconstruction Batches", leave=False)
 
     for _ in batch_iter:
         batch_tokens = activation_store.get_batch_tokens(eval_batch_size_prompts)
@@ -430,7 +430,7 @@ def get_sparsity_and_variance_metrics(
 
     batch_iter = range(n_batches)
     if verbose:
-        batch_iter = tqdm(batch_iter, desc="Sparsity and Variance Batches")
+        batch_iter = tqdm(batch_iter, desc="Sparsity and Variance Batches", leave=False)
 
     for _ in batch_iter:
         batch_tokens = activation_store.get_batch_tokens(eval_batch_size_prompts)

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/loading/pretrained_sae_loaders.py

@@ -575,6 +575,8 @@ def _infer_gemma_3_raw_cfg_dict(repo_id: str, folder_name: str) -> dict[str, Any
         "model_name": model_name,
         "hf_hook_point_in": hf_hook_point_in,
     }
+    if "transcoder" in folder_name or "clt" in folder_name:
+        cfg["affine_connection"] = "affine" in folder_name
     if hf_hook_point_out is not None:
         cfg["hf_hook_point_out"] = hf_hook_point_out
 
@@ -614,11 +616,11 @@ def get_gemma_3_config_from_hf(
     if "resid_post" in folder_name:
         hook_name = f"blocks.{layer}.hook_resid_post"
     elif "attn_out" in folder_name:
-        hook_name = f"blocks.{layer}.hook_attn_out"
+        hook_name = f"blocks.{layer}.attn.hook_z"
     elif "mlp_out" in folder_name:
         hook_name = f"blocks.{layer}.hook_mlp_out"
     elif "transcoder" in folder_name or "clt" in folder_name:
-        hook_name = f"blocks.{layer}.ln2.hook_normalized"
+        hook_name = f"blocks.{layer}.hook_mlp_in"
         hook_name_out = f"blocks.{layer}.hook_mlp_out"
     else:
         raise ValueError("Hook name not found in folder_name.")
@@ -643,7 +645,11 @@ def get_gemma_3_config_from_hf(
 
     architecture = "jumprelu"
     if "transcoder" in folder_name or "clt" in folder_name:
-        architecture = "jumprelu_skip_transcoder"
+        architecture = (
+            "jumprelu_skip_transcoder"
+            if raw_cfg_dict.get("affine_connection", False)
+            else "jumprelu_transcoder"
+        )
         d_out = shapes_dict["w_dec"][-1]
 
     cfg = {

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/pretrained_saes.yaml

@@ -4148,6 +4148,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_17_width_16k_l0_medium
     path: resid_post/layer_17_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/17-gemmascope-2-res-16k
   - id: layer_17_width_16k_l0_small
     path: resid_post/layer_17_width_16k_l0_small
     l0: 20
@@ -4166,6 +4167,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_17_width_262k_l0_medium
     path: resid_post/layer_17_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/17-gemmascope-2-res-262k
   - id: layer_17_width_262k_l0_medium_seed_1
     path: resid_post/layer_17_width_262k_l0_medium_seed_1
     l0: 60
@@ -4178,6 +4180,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_17_width_65k_l0_medium
     path: resid_post/layer_17_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/17-gemmascope-2-res-65k
   - id: layer_17_width_65k_l0_small
     path: resid_post/layer_17_width_65k_l0_small
     l0: 20
@@ -4187,6 +4190,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_22_width_16k_l0_medium
     path: resid_post/layer_22_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/22-gemmascope-2-res-16k
   - id: layer_22_width_16k_l0_small
     path: resid_post/layer_22_width_16k_l0_small
     l0: 20
@@ -4205,6 +4209,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_22_width_262k_l0_medium
     path: resid_post/layer_22_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/22-gemmascope-2-res-262k
   - id: layer_22_width_262k_l0_medium_seed_1
     path: resid_post/layer_22_width_262k_l0_medium_seed_1
     l0: 60
@@ -4217,6 +4222,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_22_width_65k_l0_medium
     path: resid_post/layer_22_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/22-gemmascope-2-res-65k
   - id: layer_22_width_65k_l0_small
     path: resid_post/layer_22_width_65k_l0_small
     l0: 20
@@ -4226,6 +4232,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_29_width_16k_l0_medium
     path: resid_post/layer_29_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/29-gemmascope-2-res-16k
   - id: layer_29_width_16k_l0_small
     path: resid_post/layer_29_width_16k_l0_small
     l0: 20
@@ -4244,6 +4251,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_29_width_262k_l0_medium
     path: resid_post/layer_29_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/29-gemmascope-2-res-262k
   - id: layer_29_width_262k_l0_medium_seed_1
     path: resid_post/layer_29_width_262k_l0_medium_seed_1
     l0: 60
@@ -4256,6 +4264,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_29_width_65k_l0_medium
     path: resid_post/layer_29_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-4b-it/29-gemmascope-2-res-65k
   - id: layer_29_width_65k_l0_small
     path: resid_post/layer_29_width_65k_l0_small
     l0: 20
@@ -4265,6 +4274,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_9_width_16k_l0_medium
     path: resid_post/layer_9_width_16k_l0_medium
     l0: 53
+    neuronpedia: gemma-3-4b-it/9-gemmascope-2-res-16k
   - id: layer_9_width_16k_l0_small
     path: resid_post/layer_9_width_16k_l0_small
     l0: 17
@@ -4283,6 +4293,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_9_width_262k_l0_medium
     path: resid_post/layer_9_width_262k_l0_medium
     l0: 53
+    neuronpedia: gemma-3-4b-it/9-gemmascope-2-res-262k
   - id: layer_9_width_262k_l0_medium_seed_1
     path: resid_post/layer_9_width_262k_l0_medium_seed_1
     l0: 53
@@ -4295,6 +4306,7 @@ gemma-scope-2-4b-it-res:
   - id: layer_9_width_65k_l0_medium
     path: resid_post/layer_9_width_65k_l0_medium
     l0: 53
+    neuronpedia: gemma-3-4b-it/9-gemmascope-2-res-65k
   - id: layer_9_width_65k_l0_small
     path: resid_post/layer_9_width_65k_l0_small
     l0: 17
@@ -14491,6 +14503,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_12_width_16k_l0_medium
     path: resid_post/layer_12_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/12-gemmascope-2-res-16k
   - id: layer_12_width_16k_l0_small
     path: resid_post/layer_12_width_16k_l0_small
     l0: 20
@@ -14509,6 +14522,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_12_width_262k_l0_medium
     path: resid_post/layer_12_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/12-gemmascope-2-res-262k
   - id: layer_12_width_262k_l0_medium_seed_1
     path: resid_post/layer_12_width_262k_l0_medium_seed_1
     l0: 60
@@ -14521,6 +14535,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_12_width_65k_l0_medium
     path: resid_post/layer_12_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/12-gemmascope-2-res-65k
   - id: layer_12_width_65k_l0_small
     path: resid_post/layer_12_width_65k_l0_small
     l0: 20
@@ -14530,6 +14545,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_15_width_16k_l0_medium
     path: resid_post/layer_15_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/15-gemmascope-2-res-16k
   - id: layer_15_width_16k_l0_small
     path: resid_post/layer_15_width_16k_l0_small
     l0: 20
@@ -14548,6 +14564,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_15_width_262k_l0_medium
     path: resid_post/layer_15_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/15-gemmascope-2-res-262k
   - id: layer_15_width_262k_l0_medium_seed_1
     path: resid_post/layer_15_width_262k_l0_medium_seed_1
     l0: 60
@@ -14560,6 +14577,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_15_width_65k_l0_medium
     path: resid_post/layer_15_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/15-gemmascope-2-res-65k
   - id: layer_15_width_65k_l0_small
     path: resid_post/layer_15_width_65k_l0_small
     l0: 20
@@ -14569,6 +14587,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_5_width_16k_l0_medium
     path: resid_post/layer_5_width_16k_l0_medium
     l0: 55
+    neuronpedia: gemma-3-270m-it/5-gemmascope-2-res-16k
   - id: layer_5_width_16k_l0_small
     path: resid_post/layer_5_width_16k_l0_small
     l0: 18
@@ -14587,6 +14606,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_5_width_262k_l0_medium
     path: resid_post/layer_5_width_262k_l0_medium
     l0: 55
+    neuronpedia: gemma-3-270m-it/5-gemmascope-2-res-262k
   - id: layer_5_width_262k_l0_medium_seed_1
     path: resid_post/layer_5_width_262k_l0_medium_seed_1
     l0: 55
@@ -14599,6 +14619,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_5_width_65k_l0_medium
     path: resid_post/layer_5_width_65k_l0_medium
     l0: 55
+    neuronpedia: gemma-3-270m-it/5-gemmascope-2-res-65k
   - id: layer_5_width_65k_l0_small
     path: resid_post/layer_5_width_65k_l0_small
     l0: 18
@@ -14608,6 +14629,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_9_width_16k_l0_medium
     path: resid_post/layer_9_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/9-gemmascope-2-res-16k
   - id: layer_9_width_16k_l0_small
     path: resid_post/layer_9_width_16k_l0_small
     l0: 20
@@ -14626,6 +14648,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_9_width_262k_l0_medium
     path: resid_post/layer_9_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/9-gemmascope-2-res-262k
   - id: layer_9_width_262k_l0_medium_seed_1
     path: resid_post/layer_9_width_262k_l0_medium_seed_1
     l0: 60
@@ -14638,6 +14661,7 @@ gemma-scope-2-270m-it-res:
   - id: layer_9_width_65k_l0_medium
     path: resid_post/layer_9_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-270m-it/9-gemmascope-2-res-65k
   - id: layer_9_width_65k_l0_small
     path: resid_post/layer_9_width_65k_l0_small
     l0: 20
@@ -18727,6 +18751,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_13_width_16k_l0_medium
     path: resid_post/layer_13_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/13-gemmascope-2-res-16k
   - id: layer_13_width_16k_l0_small
     path: resid_post/layer_13_width_16k_l0_small
     l0: 20
@@ -18745,6 +18770,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_13_width_262k_l0_medium
     path: resid_post/layer_13_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/13-gemmascope-2-res-262k
   - id: layer_13_width_262k_l0_medium_seed_1
     path: resid_post/layer_13_width_262k_l0_medium_seed_1
     l0: 60
@@ -18757,6 +18783,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_13_width_65k_l0_medium
     path: resid_post/layer_13_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/13-gemmascope-2-res-65k
   - id: layer_13_width_65k_l0_small
     path: resid_post/layer_13_width_65k_l0_small
     l0: 20
@@ -18766,6 +18793,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_17_width_16k_l0_medium
     path: resid_post/layer_17_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/17-gemmascope-2-res-16k
   - id: layer_17_width_16k_l0_small
     path: resid_post/layer_17_width_16k_l0_small
     l0: 20
@@ -18784,6 +18812,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_17_width_262k_l0_medium
     path: resid_post/layer_17_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/17-gemmascope-2-res-262k
   - id: layer_17_width_262k_l0_medium_seed_1
     path: resid_post/layer_17_width_262k_l0_medium_seed_1
     l0: 60
@@ -18796,6 +18825,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_17_width_65k_l0_medium
     path: resid_post/layer_17_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/17-gemmascope-2-res-65k
   - id: layer_17_width_65k_l0_small
     path: resid_post/layer_17_width_65k_l0_small
     l0: 20
@@ -18805,6 +18835,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_22_width_16k_l0_medium
     path: resid_post/layer_22_width_16k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/22-gemmascope-2-res-16k
   - id: layer_22_width_16k_l0_small
     path: resid_post/layer_22_width_16k_l0_small
     l0: 20
@@ -18823,6 +18854,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_22_width_262k_l0_medium
     path: resid_post/layer_22_width_262k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/22-gemmascope-2-res-262k
   - id: layer_22_width_262k_l0_medium_seed_1
     path: resid_post/layer_22_width_262k_l0_medium_seed_1
     l0: 60
@@ -18835,6 +18867,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_22_width_65k_l0_medium
     path: resid_post/layer_22_width_65k_l0_medium
     l0: 60
+    neuronpedia: gemma-3-1b-it/22-gemmascope-2-res-65k
   - id: layer_22_width_65k_l0_small
     path: resid_post/layer_22_width_65k_l0_small
     l0: 20
@@ -18844,6 +18877,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_7_width_16k_l0_medium
     path: resid_post/layer_7_width_16k_l0_medium
     l0: 54
+    neuronpedia: gemma-3-1b-it/7-gemmascope-2-res-16k
   - id: layer_7_width_16k_l0_small
     path: resid_post/layer_7_width_16k_l0_small
     l0: 18
@@ -18862,6 +18896,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_7_width_262k_l0_medium
     path: resid_post/layer_7_width_262k_l0_medium
     l0: 54
+    neuronpedia: gemma-3-1b-it/7-gemmascope-2-res-262k
   - id: layer_7_width_262k_l0_medium_seed_1
     path: resid_post/layer_7_width_262k_l0_medium_seed_1
     l0: 54
@@ -18874,6 +18909,7 @@ gemma-scope-2-1b-it-res:
   - id: layer_7_width_65k_l0_medium
     path: resid_post/layer_7_width_65k_l0_medium
     l0: 54
+    neuronpedia: gemma-3-1b-it/7-gemmascope-2-res-65k
   - id: layer_7_width_65k_l0_small
     path: resid_post/layer_7_width_65k_l0_small
     l0: 18

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/saes/temporal_sae.py

@@ -4,7 +4,7 @@ TemporalSAE decomposes activations into:
 1. Predicted codes (from attention over context)
 2. Novel codes (sparse features of the residual)
 
-See: https://arxiv.org/abs/2410.04185
+See: https://arxiv.org/pdf/2511.01836
 """
 
 import math

{sae_lens-6.29.1 → sae_lens-6.31.0}/sae_lens/training/activation_scaler.py

@@ -28,7 +28,9 @@ class ActivationScaler:
     ) -> float:
         norms_per_batch: list[float] = []
         for _ in tqdm(
-            range(n_batches_for_norm_estimate), desc="Estimating norm scaling factor"
+            range(n_batches_for_norm_estimate),
+            desc="Estimating norm scaling factor",
+            leave=False,
         ):
             acts = next(data_provider)
             norms_per_batch.append(acts.norm(dim=-1).mean().item())