progen2_large_L24_SAE_k10

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 is the k=10 (highly sparse) variant. A k=350 SAE trained on the same layer is published separately at hugohrban/progen2_large_L24_SAE_k350. The two sit at opposite ends of a sparsity trade-off: this one reconstructs the layer far worse but its features align with annotated biological concepts substantially better (see below).

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

Architecture	BatchTopKSAE (interplm.sae.dictionary.BatchTopKSAE)
Base model	hugohrban/progen2-large, layer 24
Input dimension	2560
Dictionary size	10240 (expansion factor 4)
Sparsity k	10
Training steps	200,000 (final checkpoint at step 199,999)
Learning rate	5e-5 (10k warmup, decay from 130k)
Batch size	1024
Auxiliary-k alpha	0.015625
Training data	UniRef50 500k-sequence subset (~1.86 epochs)

2 of 10240 features (0.02%) are dead — they never activate, and carry an activation_rescale_factor of 0.

Reported metrics

Reconstruction fidelity — poor, and that is the real result

On a held-out 10k-sequence test set (final_evaluation.yaml):

Metric	Value
CE with SAE patching	2.9875
% cross-entropy loss recovered	0.0

Read the 0.0 carefully. pct_loss_recovered is clipped to [0, 100] by interplm/train/fidelity.py:177, so 0.0 is a floor, not a measurement: the raw value was ≤ 0, meaning patching in this SAE's reconstruction gave a cross-entropy no better than zero-ablating layer 24 entirely. At k=10 the SAE keeps only 10 active features per residue out of 10240, which is far too few to reconstruct the residual stream. This is a genuine property of the model, not a bug or a failed evaluation.

For comparison, the k=350 SAE recovers 95.60% of the loss with a patched CE of 2.5150.

Concept F1 — better than the k=350 model

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

Metric	Validation	Test
Mean F1 per concept	0.3365	0.3552
Median F1 per concept	0.1593	0.1906
90th percentile F1	0.9428	0.9387
Concepts identified (F1 > 0.5)	228 / 734	208 / 734
Concept coverage	31.06%	28.34%
Fraction of features polysemantic	4.36%	3.44%
Rank stability (mean Spearman)	—	0.7286

So despite reconstructing nothing, k=10 nearly doubles the k=350 model's concept coverage (28.34% vs 14.58% on test) and mean F1 (0.3552 vs 0.2083), with fewer polysemantic features (3.44% vs 5.76%). Sparsity buys interpretability here at the cost of fidelity.

Concept types with any coverage on test: Coiled coil (1/1), Disulfide bond (1/1), Signal peptide (2/2), Zinc finger (7/12), Domain (169/294), Transit peptide (1/3), Region (25/100), Compositional bias (1/7), Motif (1/27).

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
│   └── valid_metrics_summary.json
├── results_test_counts/
│   ├── concept_f1_scores.csv.gz       # full concept x feature x threshold sweep (2.0M 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 (max 79.7683, mean 16.8147; 0 for the 2 dead features, smallest live value 3.342), 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_k10")
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=10 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.2e-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_k10
cd progen2_large_L24_SAE_k10
cp ae_normalized.pt dashboard_cache/layer_24/SAE.pt

streamlit run interplm/dashboard/app.py -- --cache_dir /path/to/progen2_large_L24_SAE_k10/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 two minutes (verified at 117 s), nearly all of it yaml.unsafe_load parsing the 48 MB SAE_features.yaml.

Paths the cache expects

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

File	Key	Expected path
cache_level_metadata.yaml	protein_metadata._metadata_path	data/annotations/uniprotkb/proteins-swissprot.tsv.gz
layer_24/layer_info.yaml	aa_embeds_dir	data/analysis_embeddings/progen2_large/layer_24
layer_24/layer_info.yaml	sae_dir	trained_saes/best_progen_large_24_k10

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 ~310 MB uncompressed (2,007,760 rows over 734 concepts x 10,240 features x several thresholds).

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.