peal.entrypoints.run_sae_analysis

Fit or load a sparse dictionary on a generator’s latents and evaluate it.

Drives the sparse-dictionary machinery of the diffusion autoencoder generator (load_sparse_dictionary / fit_sparse_dictionary / explain_sae in peal/generators/diffusion_autoencoder.py, dictionary classes in peal/sparse_dictionaries) and the ground-truth evaluation in peal/sparse_dictionaries/sae_evaluation.py.

Invocation:

python run_sae_analysis.py --config <generator.yaml> \
    [--sd_config <sparse_dictionary.yaml>] [--run_name NAME] \
    [--top_k K] [--n_components N] [--lr LR] [--explain_components] \
    [--wandb_project peal_sae_analysis]

Without --explain_components the dictionary is only loaded, or fitted when no weights exist. With it, the per-component counterfactual search over the whole dataset runs as well, which takes hours to days. --top_k, --n_components and --lr override the dictionary config. Unknown flags are ignored, and --is_loaded is parsed with type=bool, so any non-empty string is True.

Reads the generator config, its data config and the checkpoints they name. The dictionary is scored on the validation split, with the test split as a held-out check, and the results go into the dictionary’s base_path: the fitted weights, sae_evaluation_report.txt, sae_evaluation_results.pkl, matching_matrix_diagonalized.png and TensorBoard logs in sae_logs/. The same numbers are logged to W&B when an API key is configured.

Functions

main	Fit/load the dictionary, optionally explain it, then evaluate and log.