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

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

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

Extract interpretable features from protein sequences:

For detailed training and analysis examples, see the GitHub README.

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.