fix: honor standalone feature-density evaluation

sae_lens/evals.py

@@ -126,6 +126,29 @@ def run_evals(
         eval_config.batch_size_prompts or activation_store.store_batch_size_prompts
     )
 
+    if (
+        eval_config.compute_kl or eval_config.compute_ce_loss
+    ) and eval_config.n_eval_reconstruction_batches <= 0:
+        raise ValueError(
+            "eval_config.n_eval_reconstruction_batches must be > 0 when "
+            "compute_kl or compute_ce_loss is True."
+        )
+    compute_sparsity_variance_metrics = (
+        eval_config.compute_l2_norms
+        or eval_config.compute_sparsity_metrics
+        or eval_config.compute_variance_metrics
+        or eval_config.compute_featurewise_density_statistics
+    )
+    if (
+        compute_sparsity_variance_metrics
+        and eval_config.n_eval_sparsity_variance_batches <= 0
+    ):
+        raise ValueError(
+            "eval_config.n_eval_sparsity_variance_batches must be > 0 when "
+            "compute_l2_norms, compute_sparsity_metrics, compute_variance_metrics, "
+            "or compute_featurewise_density_statistics is True."
+        )
+
     # TODO: Come up with a cleaner long term strategy here for SAEs that do reshaping.
     # turn off hook_z reshaping mode if it's on, and restore it after evals
     if "hook_z" in hook_name:
@@ -144,11 +167,6 @@ def run_evals(
     }
 
     if eval_config.compute_kl or eval_config.compute_ce_loss:
-        if eval_config.n_eval_reconstruction_batches <= 0:
-            raise ValueError(
-                "eval_config.n_eval_reconstruction_batches must be > 0 when "
-                "compute_kl or compute_ce_loss is True."
-            )
         reconstruction_metrics = get_downstream_reconstruction_metrics(
             sae,
             model,
@@ -189,16 +207,7 @@ def run_evals(
 
         activation_store.reset_input_dataset()
 
-    if (
-        eval_config.compute_l2_norms
-        or eval_config.compute_sparsity_metrics
-        or eval_config.compute_variance_metrics
-    ):
-        if eval_config.n_eval_sparsity_variance_batches <= 0:
-            raise ValueError(
-                "eval_config.n_eval_sparsity_variance_batches must be > 0 when "
-                "compute_l2_norms, compute_sparsity_metrics, or compute_variance_metrics is True."
-            )
+    if compute_sparsity_variance_metrics:
         sparsity_variance_metrics, feature_metrics = get_sparsity_and_variance_metrics(
             sae,
             model,

tests/test_evals.py

@@ -32,8 +32,12 @@
 from sae_lens.saes.batchtopk_sae import (
     BatchTopKTrainingSAE,
 )
-from sae_lens.saes.sae import SAE, TrainingSAE
-from sae_lens.saes.standard_sae import StandardSAE, StandardTrainingSAE
+from sae_lens.saes.sae import SAE, SAEMetadata, TrainingSAE
+from sae_lens.saes.standard_sae import (
+    StandardSAE,
+    StandardSAEConfig,
+    StandardTrainingSAE,
+)
 from sae_lens.saes.topk_sae import TopKTrainingSAE
 from sae_lens.training.activation_scaler import ActivationScaler, ActivationWhitening
 from sae_lens.training.activations_store import ActivationsStore
@@ -162,6 +166,109 @@ def training_sae(cfg: LanguageModelSAERunnerConfig[Any]):  # type: ignore
 ]
 
 
+@pytest.fixture(
+    params=[
+        ("blocks.0.hook_resid_pre", "none"),
+        ("blocks.0.attn.hook_z", "hook_z"),
+    ],
+    ids=["resid_pre", "hook_z"],
+)
+def feature_density_eval_inputs(request: pytest.FixtureRequest):
+    hook_name, reshape_activations = request.param
+    model = HookedTransformer(
+        HookedTransformerConfig(
+            n_layers=1,
+            d_model=8,
+            n_ctx=4,
+            d_head=4,
+            n_heads=2,
+            d_mlp=16,
+            d_vocab=20,
+            act_fn="relu",
+            device="cpu",
+        )
+    )
+    sae = StandardSAE(
+        StandardSAEConfig(
+            d_in=8,
+            d_sae=4,
+            reshape_activations=reshape_activations,
+            metadata=SAEMetadata(
+                hook_name=hook_name,
+                context_size=4,
+                prepend_bos=False,
+            ),
+        )
+    )
+    with torch.no_grad():
+        sae.W_enc.zero_()
+        sae.b_enc.copy_(torch.tensor([1.0, -1.0, 1.0, -1.0]))
+    store = ActivationsStore.from_sae(
+        model,
+        sae,
+        Dataset.from_dict({"tokens": [[0, 2, 3, 4], [0, 5, 6, 7]]}),
+        store_batch_size_prompts=2,
+        total_tokens=2,
+        disable_concat_sequences=True,
+        sequence_separator_token=None,
+    )
+    return model, sae, store
+
+
+@pytest.mark.parametrize("exclude_special_tokens", [False, [0]])
+def test_run_evals_feature_density_only(
+    feature_density_eval_inputs: tuple[
+        HookedTransformer, StandardSAE, ActivationsStore
+    ],
+    exclude_special_tokens: bool | list[int],
+):
+    model, sae, store = feature_density_eval_inputs
+    original_hook_z_reshaping_mode = sae.hook_z_reshaping_mode
+    metrics, feature_metrics = run_evals(
+        sae,
+        store,
+        model,
+        ActivationScaler(),
+        eval_config=EvalConfig(compute_featurewise_density_statistics=True),
+        exclude_special_tokens=exclude_special_tokens,
+    )
+
+    assert sae.hook_z_reshaping_mode is original_hook_z_reshaping_mode
+    assert set(metrics) == {"token_stats"}
+    assert feature_metrics["feature_density"] == [1.0, 0.0, 1.0, 0.0]
+    expected_per_prompt = 3.0 if exclude_special_tokens else 4.0
+    assert feature_metrics["consistent_activation_heuristic"][::2] == [
+        expected_per_prompt,
+        expected_per_prompt,
+    ]
+
+
+@pytest.mark.parametrize("n_batches", [0, -1])
+def test_run_evals_feature_density_requires_positive_batches(
+    feature_density_eval_inputs: tuple[
+        HookedTransformer, StandardSAE, ActivationsStore
+    ],
+    n_batches: int,
+):
+    model, sae, store = feature_density_eval_inputs
+    original_hook_z_reshaping_mode = sae.hook_z_reshaping_mode
+    with pytest.raises(
+        ValueError, match="n_eval_sparsity_variance_batches must be > 0"
+    ):
+        run_evals(
+            sae,
+            store,
+            model,
+            ActivationScaler(),
+            eval_config=EvalConfig(
+                compute_featurewise_density_statistics=True,
+                n_eval_sparsity_variance_batches=n_batches,
+            ),
+            exclude_special_tokens=False,
+        )
+    assert sae.hook_z_reshaping_mode is original_hook_z_reshaping_mode
+
+
 def test_run_evals_base_sae(
     base_sae: SAE[Any],
     activation_store: ActivationsStore,