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.

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.

For details on training and analyzing SAEs on PLMs, 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.