ESMC-sae-k64-codebook16384

• ESMC SAEs for hidden states (all layers)
• ESMC SAEs for MLP outputs (all layers)
• ESMC SAEs for one layer (different sparsity / codebook size)

The ESMC sparse autoencoders (SAEs) are unsupervised neural networks trained to decompose the learned internal representations from the ESMC model variants into a sparse representation space comprising more biologically interpretable features, revealing what the model "sees" of the user's protein input. Each feature is encouraged to be approximately monosemantic (capturing one interpretable concept) through a large feature space combined with a sparsity constraint, and may represent a specific biologically relevant property of the protein, such as a zinc binding site, beta barrel structure, or transmembrane helix.

Building on top of the ESMC 6B SAEs, the ESM Atlas is a map of 6.8 billion proteins covering the full breadth of life's biodiversity and more than one billion predicted structures. The SAEs enable translation of the model's internal representations into ~16,000 interpretable biological features. Learn more about how to use the ESM Atlas on the Biohub Platform.

Read more about ESMC and SAEs in our paper.

Intended Use

• Decomposing ESMC embeddings into interpretable features
• Visualizing feature activations on sequences and structures

Usage

The SAE model ESMC-6B-sae-layer60-k64-codebook16384 provides users with interpretable, agent-generated feature descriptions. Users can access these feature descriptions through the ESM Atlas or through the Biohub Platform.

While all SAE models can be accessed through Hugging Face, only the following five SAE models are available through the Biohub Platform:

• ESMC-6B-sae-layer60-k64-codebook16384
• ESMC-6B-sae-layer60-k64-codebook65536
• ESMC-600M-sae-layer27-k64-codebook16384
• ESMC-600M-sae-layer27-k64-codebook65536
• ESMC-300M-sae-layer23-k64-codebook65536

Install esm from GitHub (a PyPI release is coming soon):

pip install esm@git+https://github.com/Biohub/esm.git@main

You can access an SAE model through Hugging Face using the code below:

import torch
from transformers import AutoModel, AutoTokenizer

sequence = "MGSNKSKPKDASQRRRSLEPAENVHGAGGGAFPASQTPSKPASADGHRGPSAAFAPAAAEPKLFGGFNSSDTVTSPQRAGPLAGGVTTFVALYDYESRTETDLSFKKGERLQIVNNTEGDWWLAHSLSTGQTGYIPSNYVAPSDSIQAEEWYFGKITRRESERLLLNAENPRGTFLVRESETTKGAYCLSVSDFDNAKGLNVKHYKIRKLDSGGFYITSRTQFNSLQQLVAYYSKHADGLCHRLTTVCPTSKPQTQGLAKDAWEIPRESLRLEVKLGQGCFGEVWMGTWNGTTRVAIKTLKPGTMSPEAFLQEAQVMKKLRHEKLVQLYAVVSEEPIYIVTEYMSKGSLLDFLKGETGKYLRLPQLVDMAAQIASGMAYVERMNYVHRDLRAANILVGENLVCKVADFGLARLIEDNEYTARQGAKFPIKWTAPEAALYGRFTIKSDVWSFGILLTELTTKGRVPYPGMVNREVLDQVERGYRMPCPPECPESLHDLMCQCWRKEPEERPTFEYLQAFLEDYFTSTEPQYQPGENL"

model = AutoModel.from_pretrained("biohub/ESMC-6B", device_map="auto").eval()
tokenizer = AutoTokenizer.from_pretrained("biohub/ESMC-6B")
sae = AutoModel.from_pretrained(
    "biohub/ESMC-6B-sae-k64-codebook16384",
    allow_patterns=["config.json", "layer_30.safetensors", "layer_60.safetensors"],
    device=model.device,
)
sae.initialize_layers([30, 60])
model.add_sae_models([sae.layers["30"], sae.layers["60"]])

inputs = tokenizer(sequence, return_tensors="pt", padding=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}

with torch.inference_mode():
    output = model(**inputs)

# sparse.coo tensor of shape (batch, seq_len, codebook_size)
print(output["sae_outputs"]["layer60"].shape)

Model Details

ESMC SAEs are trained to reconstruct ESMC embeddings at a residue level, meaning for a protein of length L, there are L sparse vectors of SAE features. For the models hosted on the Biohub platform, the embeddings are extracted from the hidden states. To provide different options for protein interpretability, our full set of SAE models contains both models that have been trained on hidden states and models that have been trained directly on the MLP outputs. Training SAE models on MLP outputs generates features specific to that layer's computation, while training the SAE models on hidden states may provide a more global understanding. We use the TopK approach for training SAEs to control sparsity by only allowing the top k features at each position to be active.

There are two critical hyperparameters for the TopK approach:

• k: the number of active features per position
• codebook_size: total number of features the SAE can learn

With smaller codebooks, the SAE may group related concepts together. For example, a single feature might activate for all metal-binding sites. With larger codebooks, the model can split general concepts into more granular features. For example, the model may learn dedicated features for zinc-finger motifs, iron-sulfur clusters, and calcium-binding loops.

Model Naming

There are three different families of models based on the SAE training target.

Hidden states — SAE model for every layer. The first family is trained on hidden states at every layer of the respective ESMC models. The naming convention is:

{esmc-model}-sae-k64-codebook16384

The options for esmc-model are:

• ESMC-300M
• ESMC-600M
• ESMC-6B

MLP outputs — SAE model for every layer. The second family is trained on the per-layer MLP output (before the residual connection) at every layer of the respective ESMC models. The naming convention is:

{esmc-model}-sae-mlp-k64-codebook131072

The options for esmc-model are:

• `ESMC-