diff --git a/sae_lens/evals.py b/sae_lens/evals.py index d91e513cf..b04d059cd 100644 --- a/sae_lens/evals.py +++ b/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, diff --git a/tests/test_evals.py b/tests/test_evals.py index f3e260c33..94249c1b9 100644 --- a/tests/test_evals.py +++ b/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,