Elana/InterPLM-esm2-650m

Sparse Autoencoders for ESM-2 (650M)

Interpret protein language model representations using sparse autoencoders trained on ESM-2-650m layers. These models decompose complex neural representations into interpretable features, enabling deeper understanding of how protein language models process sequence information.

• —📊 Model details in the InterPLM pre-print
• —👩‍💻 Training and analysis code in the GitHub repo
• —🧬 Explore features at InterPLM.ai

Model Details

• —Base Model: ESM-2 650M (33 layers)
• —Architecture: Sparse Autoencoder
• —Input Dimension: 1,280
• —Feature Dimension: 10,240

Available Models

We provide SAE models trained on different layers of ESM-2-650M:

All models share the same architecture and dictionary size (10,240). You can find SAEs trained on ESM-2 8M here. The 650M SAEs capture more known biological concepts than the 8M but require additional compute for both ESM embedding and SAE feature extraction.

Usage

from interplm.sae.inference import load_sae_from_hf
from interplm.esm.embed import embed_single_sequence

# Get ESM embeddings for protein sequence
embeddings = embed_single_sequence(
   sequence="MRWQEMGYIFYPRKLR",
   model_name="esm2_t33_650M_UR50D",
   layer=18  # Choose ESM layer (1,9,18,24,30,33)
)

# Load SAE model and extract features 
sae = load_sae_from_hf(plm_model="esm2-650m", plm_layer=18)
features = sae.encode(embeddings)

For details on training and analyzing SAEs on PLMs, see the GitHub README.

Model Normalization

The SAEs we've trained have arbitrary scales between features since encoder/decoder weights could be linearly scaled without changing reconstructions. To make features comparable, we normalize them to activate between 0-1 based on max activation values from Swiss-Prot (since this is our primary analysis dataset). By default, use our pre-normalized SAEs (ae_normalized.pt). As this might not perfectly scale features not present in Swiss-Prot proteins, for custom normalization use ae_unnormalized.pt with this code.