Repository: EleutherAI/delphi
Pull request #100: Log dead feature counts and firing rates
Author: login: alexandersumer
Created: 2025-03-08T07:39:43Z | Merged: 2025-03-09T12:34:42Z



Diff:
diff --git a/delphi/__main__.py b/delphi/__main__.py
index 18bbd3aa..4186b284 100644
--- a/delphi/__main__.py
+++ b/delphi/__main__.py
@@ -240,6 +240,11 @@ def populate_cache(
     Populates an on-disk cache in `latents_path` with SAE latent activations.
     """
     latents_path.mkdir(parents=True, exist_ok=True)
+
+    # Create a log path within the run directory
+    log_path = latents_path.parent / "log"
+    log_path.mkdir(parents=True, exist_ok=True)
+
     cache_cfg = run_cfg.cache_cfg
     tokens = load_tokenized_data(
         cache_cfg.cache_ctx_len,
@@ -268,9 +273,13 @@ def populate_cache(
         hookpoint_to_sparse_encode,
         batch_size=cache_cfg.batch_size,
         transcode=transcode,
+        log_path=log_path,
     )
     cache.run(cache_cfg.n_tokens, tokens)
 
+    # Save firing counts to the run-specific log directory
+    cache.save_firing_counts()
+
     cache.save_splits(
         # Split the activation and location indices into different files to make
         # loading faster
diff --git a/delphi/latents/cache.py b/delphi/latents/cache.py
index 1530d983..f9bf52bf 100644
--- a/delphi/latents/cache.py
+++ b/delphi/latents/cache.py
@@ -182,6 +182,7 @@ def __init__(
         batch_size: int,
         transcode: bool = False,
         filters: dict[str, Float[Tensor, "indices"]] | None = None,
+        log_path: Path | None = None,
     ):
         """
         Initialize the LatentCache.
@@ -190,7 +191,9 @@ def __init__(
             model: The model to cache latents for.
             hookpoint_to_sparse_encode: Dictionary of sparse encoding functions.
             batch_size: Size of batches for processing.
+            transcode: Whether to transcode the model outputs.
             filters: Filters for selecting specific latents.
+            log_path: Path to save logging output.
         """
         self.model = model
         self.hookpoint_to_sparse_encode = hookpoint_to_sparse_encode
@@ -198,6 +201,8 @@ def __init__(
         self.batch_size = batch_size
         self.width = None
         self.cache = Cache(filters, batch_size=batch_size)
+        self.hookpoint_firing_counts: dict[str, Tensor] = {}
+        self.log_path = log_path
         if filters is not None:
             self.filter_submodules(filters)
 
@@ -234,10 +239,12 @@ def filter_submodules(self, filters: dict[str, Float[Tensor, "indices"]]):
             filters: Filters for selecting specific latents.
         """
         filtered_submodules = {}
-        for hookpoint in self.hookpoint_to_sae.keys():
+        for hookpoint in self.hookpoint_to_sparse_encode.keys():
             if hookpoint in filters:
-                filtered_submodules[hookpoint] = self.hookpoint_to_sae[hookpoint]
-        self.hookpoint_to_sae = filtered_submodules
+                filtered_submodules[hookpoint] = self.hookpoint_to_sparse_encode[
+                    hookpoint
+                ]
+        self.hookpoint_to_sparse_encode = filtered_submodules
 
     def run(self, n_tokens: int, tokens: token_tensor_shape):
         """
@@ -270,8 +277,16 @@ def run(self, n_tokens: int, tokens: token_tensor_shape):
                                 latents
                             )
                             self.cache.add(sae_latents, batch, batch_number, hookpoint)
+                            firing_counts = (sae_latents > 0).sum((0, 1))
                             if self.width is None:
+                                self.hookpoint_firing_counts[hookpoint] = (
+                                    firing_counts.cpu()
+                                )
                                 self.width = sae_latents.shape[2]
+                            else:
+                                self.hookpoint_firing_counts[
+                                    hookpoint
+                                ] += firing_counts.cpu()
 
                 # Update the progress bar
                 pbar.update(1)
@@ -386,3 +401,16 @@ def save_config(self, save_dir: Path, cfg: CacheConfig, model_name: str):
                 config_dict = cfg.to_dict()
                 config_dict["model_name"] = model_name
                 json.dump(config_dict, f, indent=4)
+        self.save_firing_counts()
+
+    def save_firing_counts(self):
+        """
+        Save the firing counts for the cached latents.
+        """
+        if self.log_path is None:
+            log_path = Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt"
+        else:
+            log_path = self.log_path / "hookpoint_firing_counts.pt"
+
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(self.hookpoint_firing_counts, log_path)
diff --git a/delphi/log/result_analysis.py b/delphi/log/result_analysis.py
index ce52ce18..bd0ca176 100644
--- a/delphi/log/result_analysis.py
+++ b/delphi/log/result_analysis.py
@@ -3,6 +3,7 @@
 import numpy as np
 import orjson
 import pandas as pd
+import torch
 from torch import Tensor
 
 
@@ -176,7 +177,12 @@ def parse_score_file(file_path):
     return df
 
 
-def build_scores_df(path: Path, target_modules: list[str], range: Tensor | None = None):
+def build_scores_df(
+    path: Path,
+    target_modules: list[str],
+    hookpoint_firing_counts: dict,
+    range: Tensor | None = None,
+):
     metrics_cols = [
         "accuracy",
         "probability",
@@ -200,7 +206,14 @@ def build_scores_df(path: Path, target_modules: list[str], range: Tensor | None
     ]
     df_data = {
         col: []
-        for col in ["file_name", "score_type", "latent_idx", "module"] + metrics_cols
+        for col in [
+            "file_name",
+            "score_type",
+            "latent_idx",
+            "firing_counts",
+            "module",
+        ]
+        + metrics_cols
     }
 
     # Get subdirectories in the scores path
@@ -213,10 +226,7 @@ def build_scores_df(path: Path, target_modules: list[str], range: Tensor | None
             for score_file in list(score_type_path.glob(f"*{module}*")) + list(
                 score_type_path.glob(f".*{module}*")
             ):
-                if "latent" in score_file.stem:
-                    latent_idx = int(score_file.stem.split("latent")[-1])
-                else:
-                    latent_idx = int(score_file.stem.split("feature")[-1])
+                latent_idx = int(score_file.stem.split("latent")[-1])
                 if range is not None and latent_idx not in range:
                     continue
 
@@ -226,6 +236,9 @@ def build_scores_df(path: Path, target_modules: list[str], range: Tensor | None
                 df_data["file_name"].append(score_file.stem)
                 df_data["score_type"].append(score_type)
                 df_data["latent_idx"].append(latent_idx)
+                df_data["firing_counts"].append(
+                    hookpoint_firing_counts[module][latent_idx].item()
+                )
                 df_data["module"].append(module)
                 for col in metrics_cols:
                     df_data[col].append(df.loc[0, col])
@@ -254,7 +267,20 @@ def plot_line(df: pd.DataFrame, visualize_path: Path):
 
 
 def log_results(scores_path: Path, visualize_path: Path, target_modules: list[str]):
-    df = build_scores_df(scores_path, target_modules)
+    hookpoint_firing_counts: dict[str, Tensor] = torch.load(
+        Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt", weights_only=True
+    )
+    df = build_scores_df(scores_path, target_modules, hookpoint_firing_counts)
+
+    # Calculate the number of dead features for each module which will not be in the df
+    num_dead_features = sum(
+        [
+            (hookpoint_firing_counts[module] == 0).sum().item()
+            for module in target_modules
+        ]
+    )
+    print(f"Number of dead features: {num_dead_features}")
+
     plot_line(df, visualize_path)
 
     for score_type in df["score_type"].unique():
diff --git a/delphi/tests/conftest.py b/delphi/tests/conftest.py
index 5c614416..06014158 100644
--- a/delphi/tests/conftest.py
+++ b/delphi/tests/conftest.py
@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import cast
 
 import pytest
@@ -95,9 +96,65 @@ def cache_setup(tmp_path_factory, mock_dataset: torch.Tensor, model: PreTrainedM
     # Save the cache config
 
     cache.save_config(temp_dir, cache_cfg, "EleutherAI/pythia-70m")
+    hookpoint_firing_counts = torch.load(
+        Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt", weights_only=True
+    )
     return {
         "cache": cache,
         "tokens": tokens,
         "cache_cfg": cache_cfg,
         "temp_dir": temp_dir,
+        "firing_counts": hookpoint_firing_counts,
     }
+
+
+def test_hookpoint_firing_counts_initialization(cache_setup):
+    """
+    Ensure that hookpoint_firing_counts is initialized as an empty dictionary.
+    """
+    cache = cache_setup["cache"]
+    assert isinstance(cache.hookpoint_firing_counts, dict)
+    assert len(cache.hookpoint_firing_counts) == 0  # Should be empty before run()
+
+
+def test_hookpoint_firing_counts_updates(cache_setup):
+    """
+    Ensure that hookpoint_firing_counts is properly updated after running the cache.
+    """
+    cache = cache_setup["cache"]
+    tokens = cache_setup["tokens"]
+    cache.run(cache_setup["cache_cfg"].n_tokens, tokens)
+
+    assert (
+        len(cache.hookpoint_firing_counts) > 0
+    ), "hookpoint_firing_counts should not be empty after run()"
+    for hookpoint, counts in cache.hookpoint_firing_counts.items():
+        assert isinstance(
+            counts, torch.Tensor
+        ), f"Counts for {hookpoint} should be a torch.Tensor"
+        assert counts.ndim == 1, f"Counts for {hookpoint} should be a 1D tensor"
+        assert (counts >= 0).all(), f"Counts for {hookpoint} should be non-negative"
+
+
+def test_hookpoint_firing_counts_persistence(cache_setup):
+    """
+    Ensure that hookpoint_firing_counts are correctly saved and loaded.
+    """
+    cache = cache_setup["cache"]
+    cache.save_firing_counts()
+
+    firing_counts_path = Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt"
+    assert firing_counts_path.exists(), "Firing counts file should exist after saving"
+
+    loaded_counts = torch.load(firing_counts_path, weights_only=True)
+    assert isinstance(
+        loaded_counts, dict
+    ), "Loaded firing counts should be a dictionary"
+    assert (
+        loaded_counts.keys() == cache.hookpoint_firing_counts.keys()
+    ), "Loaded firing counts keys should match saved keys"
+
+    for hookpoint, counts in loaded_counts.items():
+        assert torch.equal(
+            counts, cache.hookpoint_firing_counts[hookpoint]
+        ), f"Mismatch in firing counts for {hookpoint}"
diff --git a/delphi/tests/e2e.py b/delphi/tests/e2e.py
index 073ef3d0..6150ab0d 100644
--- a/delphi/tests/e2e.py
+++ b/delphi/tests/e2e.py
@@ -58,7 +58,10 @@ async def test():
 
     # Performs better than random guessing
     scores_path = Path("results") / run_cfg.name / "scores"
-    df = build_scores_df(scores_path, run_cfg.hookpoints)
+    hookpoint_firing_counts = torch.load(
+        Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt", weights_only=True
+    )
+    df = build_scores_df(scores_path, run_cfg.hookpoints, hookpoint_firing_counts)
     for score_type in df["score_type"].unique():
         score_df = df.query(f"score_type == '{score_type}'")
 


Reviews:
user: login: luciaquirke
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-08T14:53:22Z
user: login: luciaquirke
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-08T14:56:36Z
user: login: luciaquirke
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-08T14:56:47Z
user: login: luciaquirke
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-08T14:57:43Z
user: login: luciaquirke
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-09T10:18:37Z

Review comments:
diff_hunk: @@ -95,9 +96,67 @@ def cache_setup(tmp_path_factory, mock_dataset: torch.Tensor, model: PreTrainedM
     # Save the cache config
 
     cache.save_config(temp_dir, cache_cfg, "EleutherAI/pythia-70m")
+    cache.save_firing_counts()
path: delphi/tests/conftest.py
user: login: luciaquirke
body: Let's call this at the end of the caching function instead of here to simplify usage
created_at: 2025-03-08T14:53:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -386,3 +395,11 @@ def save_config(self, save_dir: Path, cfg: CacheConfig, model_name: str):
                 config_dict = cfg.to_dict()
                 config_dict["model_name"] = model_name
                 json.dump(config_dict, f, indent=4)
+
+    def save_firing_counts(self):
+        """
+        Save the firing counts for the cached latents.
+        """
+        log_path = Path.cwd() / "results" / "log" / "hookpoint_firing_counts.pt"
path: delphi/latents/cache.py
user: login: luciaquirke
body: medium confidence: maybe we should pass the logging output path for the run into the latents cache init, and then use it here, and then in the CLI in \_\_main\_\_ we can include the name of the run in this path so the firing counts get saved to the particular run folder
created_at: 2025-03-08T14:56:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -52,6 +53,7 @@ def latent_balanced_score_metrics(
         print(f"F1 Score: {metrics['f1_score']:.3f}")
         print(f"Precision: {metrics['precision']:.3f}")
         print(f"Recall: {metrics['recall']:.3f}")
+        # print(f"Dead Feature Count: {(df['hookpoint_firing_counts'] == 0).sum()}")
path: delphi/log/result_analysis.py
user: login: luciaquirke
body: Remove print statement
created_at: 2025-03-08T14:56:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -226,6 +237,9 @@ def build_scores_df(path: Path, target_modules: list[str], range: Tensor | None
                 df_data["file_name"].append(score_file.stem)
                 df_data["score_type"].append(score_type)
                 df_data["latent_idx"].append(latent_idx)
+                df_data["hookpoint_firing_counts"].append(
path: delphi/log/result_analysis.py
user: login: luciaquirke
body: optional nit: call this "firing_count"
created_at: 2025-03-08T14:57:43Z
author_association: CONTRIBUTOR

PR comments:
user: login: CLAassistant
created_at: 2025-03-08T07:39:48Z
body: [![CLA assistant check](https://cla-assistant.io/pull/badge/signed)](https://cla-assistant.io/EleutherAI/delphi?pullRequest=100) <br/>All committers have signed the CLA.
author_association: NONE
user: login: SrGonao
created_at: 2025-03-08T13:50:00Z
body: Is this to adress #96 ?
author_association: COLLABORATOR
user: login: luciaquirke
created_at: 2025-03-09T05:33:59Z
body: LGTM with requested changes
author_association: CONTRIBUTOR
user: login: luciaquirke
created_at: 2025-03-09T12:34:47Z
body: Thanks for the contribution!
author_association: CONTRIBUTOR