progen2_large_L24_SAE_k350

progen2_large_L24_SAE_k350

A BatchTopK sparse autoencoder (SAE) trained on layer 24 residual-stream activations of hugohrban/progen2-large, together with the precomputed dashboard cache needed to browse its features interactively.

This repository accompanies a bachelor thesis on interpretability of protein language models and is published so that reviewers can reproduce the reported numbers and inspect the features directly. Training, analysis and the dashboard build on the InterPLM framework, extended here with a ProGen2 embedder.

SAE configuration

Reported metrics

Reconstruction fidelity on a held-out 10k-sequence test set (final_evaluation.yaml):

Concept-F1 against UniProtKB annotations (734 concepts). Full per-concept numbers are in results_*_counts/, summaries in metrics_summary.json / valid_metrics_summary.json:

Best-covered concept types on test: Signal peptide (2/2), Coiled coil (1/1), Disulfide bond (1/1), Domain (90/294), Zinc finger (4/12), Transit peptide (1/3).

Repository contents

├── ae.pt                              # trained SAE, raw (activation_rescale_factor = 1)
├── ae_normalized.pt                   # same weights + per-feature activation rescaling
├── config.yaml                        # full training config (required to load either .pt)
├── final_evaluation.yaml              # reconstruction fidelity
├── max_activations_per_feature.pt     # per-feature max activation
├── Per_feature_statistics.yaml        # activation frequency statistics
├── Per_feature_max_examples.yaml      # top activating proteins per feature
├── Per_feature_quantile_examples.yaml # lower-quantile examples per feature
├── feature_stats/max.npy
├── results_valid_counts/
│   ├── concept_f1_scores.csv.gz       # full concept x feature x threshold sweep (12.1M rows)
│   └── valid_metrics_summary.json
├── results_test_counts/
│   ├── concept_f1_scores.csv.gz       # full concept x feature x threshold sweep (12.1M rows)
│   ├── heldout_top_pairings.csv
│   ├── heldout_all_top_pairings.csv
│   └── metrics_summary.json
└── dashboard_cache/
    ├── cache_level_metadata.yaml
    └── layer_24/
        ├── layer_info.yaml
        ├── SAE_features.yaml
        ├── Per_feature_statistics.yaml
        ├── Per_feature_max_examples.yaml
        └── Per_feature_quantile_examples.yaml

ae.pt vs ae_normalized.pt

Both files contain identical encoder, decoder, bias and threshold tensors. They differ only in the 10240-element activation_rescale_factor buffer:

• ae.pt — all ones. The raw training output.
• ae_normalized.pt — per-feature scaling (min 2.1353, max 94.0764, mean 7.7289), produced by interplm/sae/normalize.py. Feature activations are divided by this factor so that they live on a comparable scale across features.

ae_normalized.pt is the file used for every analysis and steering result reported in the thesis, and is what most InterPLM scripts expect. Use ae.pt only if you specifically want unnormalized activations.

Loading the SAE

load_sae reads config.yaml from the same directory to reconstruct the architecture, so keep the files together.

from huggingface_hub import snapshot_download
from interplm.sae.inference import load_sae

sae_dir = snapshot_download(repo_id="hugohrban/progen2_large_L24_SAE_k350")
sae = load_sae(model_dir=sae_dir, model_name="ae_normalized.pt", device="cuda:0")

# acts: (n_residues, 2560) layer-24 activations from progen2-large
features = sae.encode(acts)   # -> (n_residues, 10240), at most k=350 nonzero per residue

Running the dashboard

The dashboard needs one extra step: DashboardCache loads the SAE from dashboard_cache/layer_24/SAE.pt by exact filename. In the original cache that file was a separate copy of the normalized SAE, but it is the same model as ae_normalized.pt — identical encoder, decoder, bias and threshold tensors, with activation_rescale_factor agreeing to a maximum relative difference of 3.3e-06 (float32 round-trip noise). It is therefore not shipped twice; copy it into place instead:

git clone https://huggingface.co/hugohrban/progen2_large_L24_SAE_k350
cd progen2_large_L24_SAE_k350
cp ae_normalized.pt dashboard_cache/layer_24/SAE.pt

streamlit run interplm/dashboard/app.py -- --cache_dir /path/to/progen2_large_L24_SAE_k350/dashboard_cache

Pass the dashboard_cache directory itself, not dashboard_cache/layer_24 — the loader reads cache_level_metadata.yaml at that level and treats each subdirectory as one layer.

Feature browsing, per-feature statistics and max-activating examples work from the cache alone.

The first load takes about 10 minutes and appears to hang — this is expected. Nearly all of it is yaml.unsafe_load parsing the 251 MB SAE_features.yaml; the cache was verified to load in 572 s on a warm filesystem.

Paths the cache expects

The cache YAMLs embed relative paths, resolved from your working directory:

Protein names and sequences in the dashboard come from the SwissProt TSV; without it that metadata is unavailable. The analysis embeddings are only needed for pages that recompute activations live.

Concept Explorer

The Concept Explorer page is not fully usable from this repository alone. It reads sae_dir/results_{split}_counts/concept_f1_scores.csv (uncompressed — run gunzip -k results_test_counts/concept_f1_scores.csv.gz first), but it also requires cached rank_eval/ results, which are not published here, and otherwise falls back to recomputing rank evaluation from the analysis embeddings and UniProtKB annotations on a GPU.

The F1 CSVs are included as a results artifact — so the reported concept numbers can be checked per concept — rather than as a working dependency. They are gzipped because each is ~1.8 GB uncompressed (12,069,138 rows over 748 concepts x 10,240 features x several thresholds; most rows are near-zero F1 and exist only to make the sweep complete).

Security note

The dashboard cache is loaded with yaml.unsafe_load and the YAML files embed Python object references (interplm.sae.dictionary.BatchTopKSAE, interplm.dashboard.protein_metadata.UniProtMetadata, pathlib.PosixPath). Loading them therefore requires the interplm package to be importable, and carries the same trust assumptions as unpickling: only load this cache if you trust this repository.

License

MIT, following the InterPLM framework this work builds on.