Representation learning for low-data protein fitness prediction using sequence–structure embeddings and graph neural networks

Abstract

Accurate prediction of mutational fitness effects remains challenging in low-data protein engineering settings, where experimentally characterized variants are limited. Although recent protein language models and structure-aware neural networks have demonstrated strong performance on large-scale datasets, the relative importance of structural information, pre-trained embeddings, and optimization strategies under data-constrained conditions remains insufficiently understood. In this study, we investigate a graph-based framework that integrates AlphaFold2-derived protein structures, Masked Inverse Folding with Sequence Transfer (MIF-ST) embeddings, and graph neural networks for protein fitness prediction using a dataset of 152 TEM \(\beta \) -lactamase variants. Through controlled ablation experiments, we systematically evaluate the effects of structural quality, feature richness, pre-trained embeddings, and automated hyperparameter optimization across both random train–test splits and mutation-specific splits requiring generalization to unseen mutations. Under standard random splits, simpler feature representations perform competitively because training and test samples follow similar distributions. In contrast, under mutation-specific distribution shift, models incorporating structural information and pre-trained sequence–structure embeddings achieve considerably stronger predictive performance and improved generalization. Among the evaluated configurations, the full framework combining AlphaFold2, MIF-ST, and Optuna achieved the highest overall performance across the evaluation datasets. These findings demonstrate that the effectiveness of structure-aware representations depends strongly on the evaluation regime and highlight the importance of mutation-specific generalization when designing data-efficient protein fitness prediction frameworks.