Sparse Autoencoders for Low-$N$ Protein Function Prediction and Design

TL;DR

This paper demonstrates that sparse autoencoders trained on limited protein sequence data can effectively predict protein function and guide design, outperforming baseline models and capturing meaningful biological features.

Contribution

It systematically evaluates the effectiveness of sparse autoencoders on low-$N$ protein function prediction using fine-tuned ESM2 embeddings, showing their superiority in data-scarce regimes.

Findings

SAEs outperform or match ESM2 baselines with as few as 24 sequences.

SAEs encode compact, biologically meaningful representations.

Steering latent spaces yields high-fitness protein variants in 83% of cases.

Abstract

Predicting protein function from amino acid sequence remains a central challenge in data-scarce (low-$N$) regimes, limiting machine learning-guided protein design when only small amounts of assay-labeled sequence-function data are available. Protein language models (pLMs) have advanced the field by providing evolutionary-informed embeddings and sparse autoencoders (SAEs) have enabled decomposition of these embeddings into interpretable latent variables that capture structural and functional features. However, the effectiveness of SAEs for low-$N$ function prediction and protein design has not been systematically studied. Herein, we evaluate SAEs trained on fine-tuned ESM2 embeddings across diverse fitness extrapolation and protein engineering tasks. We show that SAEs, with as few as 24 sequences, consistently outperform or compete with their ESM2 baselines in fitness prediction, indicating that their sparse latent space encodes compact and biologically meaningful representations that generalize more effectively from limited data. Moreover, steering predictive latents exploits biological motifs in pLM representations, yielding top-fitness variants in 83% of cases compared to designing with ESM2 alone.