sae-lens 6.6.4__tar.gz → 6.7.0__tar.gz

{sae_lens-6.6.4 → sae_lens-6.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.6.4
+Version: 6.7.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch

{sae_lens-6.6.4 → sae_lens-6.7.0}/pyproject.toml

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

{sae_lens-6.6.4 → sae_lens-6.7.0}/sae_lens/__init__.py

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

{sae_lens-6.6.4 → sae_lens-6.7.0}/sae_lens/evals.py

@@ -718,17 +718,9 @@ def get_recons_loss(
         **model_kwargs,
     )
 
-    def kl(original_logits: torch.Tensor, new_logits: torch.Tensor):
-        original_probs = torch.nn.functional.softmax(original_logits, dim=-1)
-        log_original_probs = torch.log(original_probs)
-        new_probs = torch.nn.functional.softmax(new_logits, dim=-1)
-        log_new_probs = torch.log(new_probs)
-        kl_div = original_probs * (log_original_probs - log_new_probs)
-        return kl_div.sum(dim=-1)
-
     if compute_kl:
-        recons_kl_div = kl(original_logits, recons_logits)
-        zero_abl_kl_div = kl(original_logits, zero_abl_logits)
+        recons_kl_div = _kl(original_logits, recons_logits)
+        zero_abl_kl_div = _kl(original_logits, zero_abl_logits)
         metrics["kl_div_with_sae"] = recons_kl_div
         metrics["kl_div_with_ablation"] = zero_abl_kl_div
 
@@ -740,6 +732,18 @@ def get_recons_loss(
     return metrics
 
 
+def _kl(original_logits: torch.Tensor, new_logits: torch.Tensor):
+    # Computes the log-probabilities of the new logits (approximation).
+    log_probs_new = torch.nn.functional.log_softmax(new_logits, dim=-1)
+    # Computes the probabilities of the original logits (true distribution).
+    probs_orig = torch.nn.functional.softmax(original_logits, dim=-1)
+    # Compute the KL divergence. torch.nn.functional.kl_div expects the first argument to be the log
+    # probabilities of the approximation (new), and the second argument to be the true distribution
+    # (original) as probabilities. This computes KL(original || new).
+    kl = torch.nn.functional.kl_div(log_probs_new, probs_orig, reduction="none")
+    return kl.sum(dim=-1)
+
+
 def all_loadable_saes() -> list[tuple[str, str, float, float]]:
     all_loadable_saes = []
     saes_directory = get_pretrained_saes_directory()

{sae_lens-6.6.4 → sae_lens-6.7.0}/sae_lens/loading/pretrained_sae_loaders.py

@@ -1001,10 +1001,14 @@ def get_sparsify_config_from_disk(
         layer = int(match.group(1))
         hook_name = f"blocks.{layer}.hook_resid_post"
 
+    d_sae = old_cfg_dict.get("num_latents")
+    if d_sae is None:
+        d_sae = old_cfg_dict["d_in"] * old_cfg_dict["expansion_factor"]
+
     cfg_dict: dict[str, Any] = {
         "architecture": "standard",
         "d_in": old_cfg_dict["d_in"],
-        "d_sae": old_cfg_dict["d_in"] * old_cfg_dict["expansion_factor"],
+        "d_sae": d_sae,
         "dtype": "bfloat16",
         "device": device or "cpu",
         "model_name": config_dict.get("model", path.parts[-2]),
@@ -1248,11 +1252,11 @@ def get_mwhanna_transcoder_config_from_hf(
     try:
         # mwhanna transcoders sometimes have a typo in the config file name, so check for both
         wandb_config_path = hf_hub_download(
-            repo_id, "wanb-config.yaml", force_download=force_download
+            repo_id, "wandb-config.yaml", force_download=force_download
         )
     except EntryNotFoundError:
         wandb_config_path = hf_hub_download(
-            repo_id, "wandb-config.yaml", force_download=force_download
+            repo_id, "wanb-config.yaml", force_download=force_download
         )
     try:
         base_config_path = hf_hub_download(
@@ -1326,6 +1330,66 @@ def mwhanna_transcoder_huggingface_loader(
     return cfg_dict, state_dict, None
 
 
+def mntss_clt_layer_huggingface_loader(
+    repo_id: str,
+    folder_name: str,
+    device: str = "cpu",
+    force_download: bool = False,  # noqa: ARG001
+    cfg_overrides: dict[str, Any] | None = None,
+) -> tuple[dict[str, Any], dict[str, torch.Tensor], torch.Tensor | None]:
+    """
+    Load a MNTSS CLT layer as a single layer transcoder.
+    The assumption is that the `folder_name` is the layer to load as an int
+    """
+    base_config_path = hf_hub_download(
+        repo_id, "config.yaml", force_download=force_download
+    )
+    with open(base_config_path) as f:
+        cfg_info: dict[str, Any] = yaml.safe_load(f)
+
+    # We need to actually load the weights, since the config is missing most information
+    encoder_path = hf_hub_download(
+        repo_id,
+        f"W_enc_{folder_name}.safetensors",
+        force_download=force_download,
+    )
+    decoder_path = hf_hub_download(
+        repo_id,
+        f"W_dec_{folder_name}.safetensors",
+        force_download=force_download,
+    )
+
+    encoder_state_dict = load_file(encoder_path, device=device)
+    decoder_state_dict = load_file(decoder_path, device=device)
+
+    with torch.no_grad():
+        state_dict = {
+            "W_enc": encoder_state_dict[f"W_enc_{folder_name}"].T,  # type: ignore
+            "b_enc": encoder_state_dict[f"b_enc_{folder_name}"],  # type: ignore
+            "b_dec": encoder_state_dict[f"b_dec_{folder_name}"],  # type: ignore
+            "W_dec": decoder_state_dict[f"W_dec_{folder_name}"].sum(dim=1),  # type: ignore
+        }
+
+    cfg_dict = {
+        "architecture": "transcoder",
+        "d_in": state_dict["b_dec"].shape[0],
+        "d_out": state_dict["b_dec"].shape[0],
+        "d_sae": state_dict["b_enc"].shape[0],
+        "dtype": "float32",
+        "device": device if device is not None else "cpu",
+        "activation_fn": "relu",
+        "normalize_activations": "none",
+        "model_name": cfg_info["model_name"],
+        "hook_name": f"blocks.{folder_name}.{cfg_info['feature_input_hook']}",
+        "hook_name_out": f"blocks.{folder_name}.{cfg_info['feature_output_hook']}",
+        "apply_b_dec_to_input": False,
+        "model_from_pretrained_kwargs": {"fold_ln": False},
+        **(cfg_overrides or {}),
+    }
+
+    return cfg_dict, state_dict, None
+
+
 NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
     "connor_rob_hook_z": connor_rob_hook_z_huggingface_loader,
@@ -1337,6 +1401,7 @@ NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sparsify": sparsify_huggingface_loader,
     "gemma_2_transcoder": gemma_2_transcoder_huggingface_loader,
     "mwhanna_transcoder": mwhanna_transcoder_huggingface_loader,
+    "mntss_clt_layer_transcoder": mntss_clt_layer_huggingface_loader,
 }