Generalizing Protein Fitness Prediction Across Distribution Shifts via Structural and Co-evolutionary Priors

Abstract

Protein fitness prediction is essential for protein engineering. Recent fitness-landscape benchmarks such as FLIP2 have shown that methods performing well under conventional train-test splits can fail when the test set is shifted out of distribution (OOD), including mutations at unseen positions and variants from unseen protein backbones. In particular, current protein language models (pLMs) often suffer substantial performance degradation under such distribution shifts, limiting their practical use in protein engineering. We therefore develop a framework that augments sequence representations with structural and co-evolutionary information for fitness prediction. Specifically, we pair structure-aware pLMs with features pooled from an AlphaFold-style trunk applied to each sequence and its multiple sequence alignment (MSA), and use a lightweight prediction probe to map these frozen representations to fitness scores. Our framework consistently outperforms the best previously reported FLIP2 scores across several types of distribution shift. We then systematically study the contribution of each component in our framework to generalization. Structural priors yield the largest performance gains on held-out protein backbones, whereas trunk features improve generalization to held-out mutation positions. Rather than proposing a single best method, we characterize which modalities are effective for different protein families and distribution shifts, providing practical guidance for choosing representations according to the generalization regime. These findings support the development of more generalizable fitness predictors that can serve as effective reward functions or oracles for protein design toward desired functions.