aasatorres/esm2-sae-topk-16384-k512

ESM-2 Sparse Autoencoder (TopK, K=16384, k=512)

Sparse Autoencoder (SAE) trained on residue-level embeddings from ESM-2 (650M, layer 33) for interpretability research on protein language models.

Model details

Parameter	Value
Architecture	TopK SAE
Input dim	1280 (ESM-2 layer 33)
Latent dim	16384
TopK k	512
Training proteins	50,000 Swiss-Prot
ESM-2 model	facebook/esm2_t33_650M_UR50D

Files

• best.pt: model weights (PyTorch)
• config.json: architecture config and input normalization statistics (emb_mean, emb_std)
• metrics.csv: training and validation metrics per epoch

Usage

import torch, json
from pathlib import Path

config = json.load(open("config.json"))
checkpoint = torch.load("best.pt", map_location="cpu")
# checkpoint["model_state_dict"] contains the SAE weights

Companion datasets

• Processed dataset: https://doi.org/10.5281/zenodo.20379180
• Alignment scores & logs: https://doi.org/10.5281/zenodo.20384152
• Code: https://github.com/Andre-Satorres/PLM-Interpretability