Prioritizing the Rare Hits: A Top-Focused Ranking Framework for Extrapolating Protein Fitness

Protein fitness prediction requires extrapolating from low-order mutants to unseen higher-order combinations, where epistatic interactions and the scarcity of high-fitness variants make generalization challenging. Conventional regression and ranking objectives optimize global performance across the full fitness distribution, thereby underemphasizing the rare, top-ranked variants that are most valuable for protein engineering. To address this limitation, we introduce TREP, a Top-focused Ranking framework for Extrapolating Protein fitness. TREP maps protein representations into an activity-oriented latent space spanned by two orthogonal anchors, while percentile-aware pairwise weighting and stratified sampling focus learning on the high-fitness tail. We further augment training with synthetic inactive variants generated by recombining mutations associated with the bottom tail of the measured fitness distribution. Experiments across several two-to-many protein fitness benchmarks demonstrate that TREP consistently improves both overall ranking performance and the retrieval of rare high-fitness variants.