goodfire-ai · danbraunai-goodfire · Dec 5, 2025 · Dec 5, 2025 · Dec 5, 2025 · Dec 8, 2025
diff --git a/.gitignore b/.gitignore
@@ -2,6 +2,9 @@ spd/scripts/sweep_params.yaml
 docs/coverage/**
 notebooks/**
 
+# Script outputs (generated files, often large)
+scripts/outputs/
+
 **/out/
 neuronpedia_outputs/
 .env

diff --git a/spd/configs.py b/spd/configs.py
@@ -13,7 +13,34 @@
 
 from spd.base_config import BaseConfig
 from spd.log import logger
-from spd.spd_types import CiFnType, ModelPath, Probability
+from spd.spd_types import GlobalCiFnType, LayerwiseCiFnType, ModelPath, Probability
+
+
+class LayerwiseCiConfig(BaseConfig):
+    """Configuration for layerwise CI functions (one per layer)."""
+
+    mode: Literal["layerwise"] = "layerwise"
+    fn_type: LayerwiseCiFnType = Field(
+        ..., description="Type of layerwise CI function: mlp, vector_mlp, or shared_mlp"
+    )
+    hidden_dims: list[NonNegativeInt] = Field(
+        ..., description="Hidden dimensions for the CI function MLP"
+    )
+
+
+class GlobalCiConfig(BaseConfig):
+    """Configuration for global CI function (single function for all layers)."""
+
+    mode: Literal["global"] = "global"
+    fn_type: GlobalCiFnType = Field(
+        ..., description="Type of global CI function: global_shared_mlp"
+    )
+    hidden_dims: list[NonNegativeInt] = Field(
+        ..., description="Hidden dimensions for the global CI function MLP"
+    )
+
+
+CiConfig = LayerwiseCiConfig | GlobalCiConfig
 
 
 class ScheduleConfig(BaseConfig):
@@ -397,13 +424,10 @@ class Config(BaseConfig):
         ...,
         description="Number of stochastic masks to sample when using stochastic recon losses",
     )
-    ci_fn_type: CiFnType = Field(
-        default="vector_mlp",
-        description="Type of causal importance function used to calculate the causal importance.",
-    )
-    ci_fn_hidden_dims: list[NonNegativeInt] = Field(
-        default=[8],
-        description="Hidden dimensions for the causal importance function used to calculate the causal importance",
+    ci_config: CiConfig = Field(
+        ...,
+        description="Configuration for the causal importance function. "
+        "Use LayerwiseCiConfig for per-layer CI functions or GlobalCiConfig for a single global CI function.",
     )
     sampling: SamplingType = Field(
         default="continuous",
@@ -614,8 +638,6 @@ def microbatch_size(self) -> PositiveInt:
         "pretrained_model_name_hf": "pretrained_model_name",
         "recon_coeff": "ci_recon_coeff",
         "recon_layerwise_coeff": "ci_recon_layerwise_coeff",
-        "gate_type": "ci_fn_type",
-        "gate_hidden_dims": "ci_fn_hidden_dims",
     }
 
     @model_validator(mode="before")

diff --git a/spd/experiments/ih/ih_config.yaml b/spd/experiments/ih/ih_config.yaml
@@ -34,8 +34,10 @@ stochastic_recon_layerwise_coeff: 1
 importance_minimality_coeff: 1e-2
 pnorm: 0.1
 output_loss_type: kl
-ci_fn_type: "vector_mlp"
-ci_fn_hidden_dims: [128]
+ci_config:
+  mode: layerwise
+  fn_type: vector_mlp
+  hidden_dims: [128]
 
 n_examples_until_dead: 10000
 

diff --git a/spd/experiments/lm/gpt2_config.yaml b/spd/experiments/lm/gpt2_config.yaml
@@ -7,8 +7,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "vector_mlp"
-ci_fn_hidden_dims: [12]
+ci_config:
+  mode: layerwise
+  fn_type: vector_mlp
+  hidden_dims: [12]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "transformer.h.1.attn.c_attn"

diff --git a/spd/experiments/lm/ss_gpt2_config.yaml b/spd/experiments/lm/ss_gpt2_config.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "vector_mlp"
-ci_fn_hidden_dims: [12]
+ci_config:
+  mode: layerwise
+  fn_type: vector_mlp
+  hidden_dims: [12]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "transformer.h.1.mlp.c_fc"

diff --git a/spd/experiments/lm/ss_gpt2_simple-1L.yaml b/spd/experiments/lm/ss_gpt2_simple-1L.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 1
 n_mask_samples: 1
-ci_fn_type: "shared_mlp"
-ci_fn_hidden_dims: [550]
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "h.*.mlp.c_fc"

diff --git a/spd/experiments/lm/ss_gpt2_simple-2L.yaml b/spd/experiments/lm/ss_gpt2_simple-2L.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 1
 n_mask_samples: 1
-ci_fn_type: "shared_mlp"
-ci_fn_hidden_dims: [550]
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "h.*.mlp.c_fc"

diff --git a/spd/experiments/lm/ss_gpt2_simple_config.yaml b/spd/experiments/lm/ss_gpt2_simple_config.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "shared_mlp"
-ci_fn_hidden_dims: [1000]
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [1000]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "h.*.mlp.c_fc"

diff --git a/spd/experiments/lm/ss_gpt2_simple_noln_config.yaml b/spd/experiments/lm/ss_gpt2_simple_noln_config.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "vector_mlp"
-ci_fn_hidden_dims: [12]
+ci_config:
+  mode: layerwise
+  fn_type: vector_mlp
+  hidden_dims: [12]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "h.*.mlp.c_fc"

diff --git a/spd/experiments/lm/ss_llama_simple-1L.yaml b/spd/experiments/lm/ss_llama_simple-1L.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 550
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ss_llama_simple-2L.yaml b/spd/experiments/lm/ss_llama_simple-2L.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 550
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ss_llama_simple_config.yaml b/spd/experiments/lm/ss_llama_simple_config.yaml
@@ -6,8 +6,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "shared_mlp"
-ci_fn_hidden_dims: [1000]
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [1000]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "h.*.mlp.gate_proj"

diff --git a/spd/experiments/lm/ss_llama_simple_mlp-1L.yaml b/spd/experiments/lm/ss_llama_simple_mlp-1L.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 550
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ss_llama_simple_mlp-2L-wide.yaml b/spd/experiments/lm/ss_llama_simple_mlp-2L-wide.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 1250
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [1250]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ss_llama_simple_mlp-2L.yaml b/spd/experiments/lm/ss_llama_simple_mlp-2L.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 550
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [550]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ss_llama_simple_mlp.yaml b/spd/experiments/lm/ss_llama_simple_mlp.yaml
@@ -3,9 +3,10 @@ wandb_run_name: null
 wandb_run_name_prefix: ''
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-- 800
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [800]
 sampling: continuous
 sigmoid_type: leaky_hard
 module_info:

diff --git a/spd/experiments/lm/ts_config.yaml b/spd/experiments/lm/ts_config.yaml
@@ -9,8 +9,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "vector_mlp"
-ci_fn_hidden_dims: [8]
+ci_config:
+  mode: layerwise
+  fn_type: vector_mlp
+  hidden_dims: [8]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "transformer.h.3.mlp.c_fc"

diff --git a/spd/experiments/resid_mlp/resid_mlp1_config.yaml b/spd/experiments/resid_mlp/resid_mlp1_config.yaml
@@ -7,8 +7,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "mlp"
-ci_fn_hidden_dims: [16]
+ci_config:
+  mode: layerwise
+  fn_type: mlp
+  hidden_dims: [16]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "layers.*.mlp_in"

diff --git a/spd/experiments/resid_mlp/resid_mlp1_global_config.yaml b/spd/experiments/resid_mlp/resid_mlp1_global_config.yaml
@@ -0,0 +1,85 @@
+# ResidualMLP 1 layer - Global CI
+# --- WandB ---
+wandb_project: spd
+wandb_run_name: null
+wandb_run_name_prefix: ""
+
+# --- General ---
+seed: 0
+n_mask_samples: 1
+ci_config:
+  mode: global
+  fn_type: global_shared_mlp
+  hidden_dims: [400, 300]
+sigmoid_type: "leaky_hard"
+module_info:
+  - module_pattern: "layers.*.mlp_in"
+    C: 100
+  - module_pattern: "layers.*.mlp_out"
+    C: 100
+identity_module_info: null
+use_delta_component: true
+
+# --- Loss config ---
+loss_metric_configs:
+  - classname: "ImportanceMinimalityLoss"
+    coeff: 1e-5
+    pnorm: 2.0
+    beta: 0
+  - classname: "StochasticReconLayerwiseLoss"
+    coeff: 1.0
+  - classname: "StochasticReconLoss"
+    coeff: 1.0
+output_loss_type: mse
+
+# --- Training ---
+batch_size: 2048
+eval_batch_size: 2048
+steps: 20_000
+lr_schedule:
+  start_val: 2e-3
+  fn_type: constant
+  warmup_pct: 0.0
+
+# --- Faithfulness Warmup ---
+faithfulness_warmup_steps: 200
+faithfulness_warmup_lr: 0.01
+faithfulness_warmup_weight_decay: 0.1
+
+# --- Logging & Saving ---
+train_log_freq: 100
+eval_freq: 500
+n_eval_steps: 100
+slow_eval_freq: 5_000
+slow_eval_on_first_step: true
+save_freq: null
+ci_alive_threshold: 0.1
+n_examples_until_dead: 204_800
+eval_metric_configs:
+  - classname: "CIHistograms"
+    n_batches_accum: 5
+  - classname: "ComponentActivationDensity"
+  - classname: "PermutedCIPlots"
+    identity_patterns: ["layers.*.mlp_in"]
+    dense_patterns: ["layers.*.mlp_out"]
+  - classname: "IdentityCIError"
+    identity_ci:
+      - layer_pattern: "layers.*.mlp_in"
+        n_features: 100
+    dense_ci:
+      - layer_pattern: "layers.*.mlp_out"
+        k: 50
+  - classname: "CI_L0"
+    groups: null
+  - classname: "CIMeanPerComponent"
+  - classname: "StochasticHiddenActsReconLoss"
+
+# --- Pretrained model info ---
+pretrained_model_class: "spd.experiments.resid_mlp.models.ResidMLP"
+pretrained_model_path: "wandb:goodfire/spd-pre-Sep-2025/runs/pziyck78"
+
+# --- Task Specific ---
+task_config:
+  task_name: resid_mlp
+  feature_probability: 0.01
+  data_generation_type: "at_least_zero_active"
diff --git a/spd/experiments/resid_mlp/resid_mlp2_config.yaml b/spd/experiments/resid_mlp/resid_mlp2_config.yaml
@@ -7,9 +7,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: shared_mlp
-ci_fn_hidden_dims:
-  - 256
+ci_config:
+  mode: layerwise
+  fn_type: shared_mlp
+  hidden_dims: [256]
 sigmoid_type: leaky_hard
 module_info:
   - module_pattern: "layers.*.mlp_in"

diff --git a/spd/experiments/resid_mlp/resid_mlp3_config.yaml b/spd/experiments/resid_mlp/resid_mlp3_config.yaml
@@ -7,8 +7,10 @@ wandb_run_name_prefix: ""
 # --- General ---
 seed: 0
 n_mask_samples: 1
-ci_fn_type: "mlp"
-ci_fn_hidden_dims: [128]
+ci_config:
+  mode: layerwise
+  fn_type: mlp
+  hidden_dims: [128]
 sigmoid_type: "leaky_hard"
 module_info:
   - module_pattern: "layers.*.mlp_in"