Support clustering

[mivanit]

NOTE: Commit before I (Dan) started putting this in a state for merging to main.

Description

• Implements clustering of subcomponents based on minimum description length (MDL) clustering.
• Refactor slurm scripts for consistency across spd runs, harvest, autointerp, and clustering

Related Issues

#13

Does this PR cause a breaking change?

No

[danbraunai-goodfire]

Old description, for posterity:

NOTE: Commit before I (Dan) started putting this in a state for merging to main.

(migrated from #43)

Description

This PR implements clustering of subcomponents based on minimum description length (MDL) clustering.

The primary interface to this package is via:

python spd/clustering/scripts/main.py --merge-config MERGE_CONFIG [--model-path MODEL_PATH] [--n-batches N_BATCHES] [--batch-size BATCH_SIZE]

This runs a pipeline which:

1. tokenizes and splits data from the dataset used by the specified model
2. runs an ensemble of the clustering algorithm, one run for each batch
3. unifies the histories to account for differences in active components between batches
4. computes and plots distances between the different merge trajectories, giving insight into which merges may be useful for downstream investigations

Related Issues

#13

Depends on:

• fixing state dict mappings #185
• Standalone clustering prereqs #186
• Reduce disk usage in CI #187

TODOs:

Code style

• array type hinting, see clustering #43 (comment)

• get rid of redundant type hints

  • clearer style.md would be useful for this, Claude sometimes omits/removes even the useful type hints

• what to do with the old notebook-style scripts tests/clustering/scripts/cluster_{resid_mlp|ss}.py, see clustering #43 (comment)

  • I find that I am frequently annoyed by these tests failing, which probably means they are good tests. perhaps they need to be refactored in some way though

  • EDIT 2025-10-21: continue to find the script tests to catch issues not caught by any other code. mostly pertaining to CLI and configs

  • EDIT 2025-11-11: keeping these for now, they are useful

• MergeConfig cleanup, see clustering #43 (comment)

  • in particular, the sampling methods implementation

Variable naming

• "components" -> "subcomponents"
• "n_components", "c_components" -> n_subcomponents
• somehow unify naming of "groups", "components", "clusters"
• "current_merge" -> "current_state"?
• consistency between compute_merge_costs and compute_mdl_cost

Implementation

• deciding on shelling out with custom file descriptors vs
  using concurrent.futures for running multiple instances of s2_clustering.py in parallel

  • the shelling out approach may cause issues with wandb because we both have a "summary" run (which is what computes distances between the different merges) and each run of s2_clustering has its own associated wandb run

• there should be a way to call the pipeline without parallel processing for s2

  • this should be the default when running with an ensemble size of 1?

• allow running of specific stages from cli

  I’d like an option to be able to run both:

  1. A specific stage of the pipeline
  2. All stages from a specific point
     It should be able to load the results in previous stages from a cache. The cli could have a use_cache: bool arg and maybe another cli arg or two to handle which stages it should run.

Testing

• more unit tests
• test loading and validating all configs
• partition tests, so clustering tests run only when clustering code is changed

Methodology

Additions

• migrate to slurm for clustering runs ([clustering] Refactor to two-stage process #203)
• alternate distance metric (algorithm by @Laplace418) [clustering] distance computation #213
• multi-batch clustering [clustering] multi-batch merging refactor #206
• adding clustering runs to an ensemble
  interface would look something like:
  • pass existing ensemble id instead of a config to spd-cluster
  • it reads the config and commit hash (errors if no commit hash) from ensemble config
  • start a bunch of new clustering runs based on that info
• proper rank computation?

Checks/Removals

• float vs boolean coactivations: it may be best to stick with just boolean coactivations, but then we should be sure this is reasonable
  • if we keep float activations, there is a bug with activation_mask[:, merge_pair[0]] + activation_mask[:, merge_pair[1]]
  • merge_config.activation_threshold should probably default to 0 and remove the None if we get rid of float acts
• get rid of component popping: we may or may not add this back in the future, but we should try to thin out the features of this PR (completed in [clustering] distance computation #213)
• properly test mcmc sampling