Sparse Autoencoders for Low-$N$ Protein Function Prediction and Design
This work addresses the problem of data-scarce protein function prediction and design for researchers in computational biology, representing an incremental improvement by applying SAEs to an existing bottleneck.
The study tackled the challenge of predicting protein function with limited data by evaluating sparse autoencoders (SAEs) on fine-tuned ESM2 embeddings, finding that SAEs outperform or match baselines in fitness prediction with as few as 24 sequences and achieve top-fitness variants in 83% of cases for protein design.
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.