Linear Fitness Subspace in Protein Language Models Enables Sample-Efficient Directed Evolution

SiYuan Ma; Canran Xiao; Zikai Xiao; Albert Gao; Liang He; Xuan-Yu Wang; Shuying Cao; Xiaojun Jia

q-bio.QMOct 6, 2026

Linear Fitness Subspace in Protein Language Models Enables Sample-Efficient Directed Evolution

Authors: SiYuan Ma, Canran Xiao, Zikai Xiao, Albert Gao, Liang He, Xuan-Yu Wang, Shuying Cao, Xiaojun Jia

Organizations: College of Computing and Data Science, Nanyang Technological University, Singapore · Sun Yat-sen University, China · Zhejiang University, China · Carnegie Mellon University, USA · Shanghai Institute of Optics and Fine Mechanics, China · Zhongnan Hospital, Wuhan University, China · University of Southern California, USA

Abstract

Model-guided directed evolution seeks to identify high-fitness protein variants under limited oracle budgets. Protein language models (PLMs) provide rich representations for this task, but task-agnostic zero-shot scores can be misaligned with a target assay, while supervised search in high-dimensional embedding spaces can make surrogate modeling and uncertainty estimation sample-inefficient. We propose the Linear Fitness Subspace (LFS) hypothesis: within mutation-induced residue-level representation changes, a compact, assay-specific set of directions makes fitness variation linearly accessible from few labeled variants. This is a local, supervision-recoverable statement rather than a claim that protein fitness landscapes or global PLM geometry are universally linear. Building on this observation, we introduce Subspace-Guided Evolutionary Search (SGES), which estimates an LFS from a small initial sample and performs surrogate modeling, uncertainty estimation, and acquisition in the learned subspace. Across 10 core ProteinGym assays, 87 extended static-validation assays, and an 18-assay budgeted-search evaluation, SGES improves fitness prediction and search efficiency over zero-shot PLMs and recent ML-guided protein optimization baselines. Controlled comparisons with PCA, random projections, label-shuffled PLS, classical mutation features, and acquisition ablations further isolate the benefit of a fitness-aligned site-delta coordinate.

Figures & tables

Figure 1: Overall pipeline of Subspace-Guided Evolutionary Search (SGES). SGES starts from the wild-type sequence and a small initial oracle budget, constructs mutation candidates, extracts mutation-aware PLM features, and performs iterative search through surrogate prediction, uncertainty-aware acquisition, oracle evaluation, and population update. This overview connects the representation-level discovery in Section 3 with the search algorithm introduced later: fitness prediction is not treated as generic high-dimensional embedding regression, but is guided by mutation-induced representation changes that expose a compact fitness-relevant structure.

Assay	Ridge $\delta$	Mean-pool	MLP	Ratio
SPG1_STRSG	0.945	0.620	0.973	0.971
BLAT_Firnberg	0.865	0.730	0.903	0.958
PABP_YEAST	0.840	0.580	0.916	0.917
P53_HUMAN	0.810	0.650	0.874	0.927
RASH_HUMAN	0.755	0.590	0.825	0.915
AMIE_PSEAE	0.758	0.550	0.823	0.921

Table 1: Linear accessibility of mutational fitness. Ridge regression on site-delta features recovers most of the nonlinear MLP performance, while mean-pooled features are substantially weaker.

Figure 2: LFS validation. Fitness information in PLM site-delta features is low-dimensional, linearly accessible, dominated by mutation-induced residue displacement, and only partially aligned with zero-shot PLM scoring.

Figure 3: Search performance. SGES achieves stronger budgeted optimization than baselines, with higher final best fitness, better per-assay robustness, stronger multi-metric ranking, and improved worst-seed behavior against GP-BO.

Method	Avg. Best Fitness	Avg. Rank
SGES (Ours)	1.471	1.38
GP-BO ( Soldát and Kléma, 2024 )	1.328	2.38
AlphaDE ( Yang et al., 2025 )	1.267	2.63
LatentDE ( Tran et al., 2025 )	1.223	3.50
EVOLVEpro ( Jiang et al., 2024 )	1.183	5.00
BOES ( Soldát and Kléma, 2024 )	1.173	4.75

Table 2: Compact comparison with recent protein optimization methods. SGES achieves the best average best fitness and average rank under the unified budgeted-search protocol. Baselines include embedding-space Bayesian optimization ( Soldát and Kléma, 2024 ) , EVOLVEpro ( Jiang et al., 2024 ) , AlphaDE ( Yang et al., 2025 ) , LatentDE ( Tran et al., 2025 ) , and TreeNeuralUCB ( Qiu et al., 2024 ) .

Figure 4: Ablation and subspace analysis. SGES benefits from coupling subspace-aware proposal with uncertainty-aware acquisition, learns biologically related cross-assay directions, and improves with PLM scale. Panel D summarizes Expected Calibration Error (ECE) by calibration-quality bins (Good/Medium/Poor; lower ECE is better); the density curves denote quality groups rather than method identities.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Assay	Ours 50%	Ours 5%	ESM2-650M	ESM2-15B	Tranc.-L	TrEVE-L	EVE	MSA-T	GEMME	ProSST
SPG1_STRSG	0.973	0.915	0.301	0.337	0.289	0.290	0.272	0.171	0.283	0.743
BLAT_Firnberg	0.903	0.785	0.737	0.424	0.637	0.736	0.729	0.741	0.683	0.766
PABP_YEAST	0.916	0.800	0.716	0.695	0.688	0.684	0.648	0.663	0.674	0.716
P53_HUMAN	0.874	0.713	0.432	0.497	0.400	0.413	0.427	0.291	0.423	0.425
RASH_HUMAN	0.825	0.606	0.498	0.313	0.450	0.487	0.480	0.446	0.434	0.619
AMIE_PSEAE	0.823	0.625	0.557	0.613	0.438	0.481	0.464	0.605	0.558	0.486

Appendix

Table S1: Zero-shot and few-shot comparison on 10 assays. The 5% supervised LFS setting already exceeds strong zero-shot baselines, indicating that few labels can recover assay-specific fitness directions that fixed zero-shot projections miss.

Variant / Feature	Panel	$k$	Surrogate	Acquisition	Avg. $\rho$	Avg. Best
Random projection	A	8	Deep ensemble	UCB	0.512	1.08
PCA projection	A	8	Deep ensemble	UCB	0.624	1.22
Label-shuffled PLS	A	8	Deep ensemble	UCB	0.485	1.02
PLS-LFS projection	A	8	Deep ensemble	UCB	0.804	1.45
Full site-delta feature	A	1280	Deep ensemble	UCB	0.763	1.32
PLS-LFS + ridge top- $k$	B	8	Ridge	Greedy mean	0.686	1.15

Appendix

Table S2: Controls separating dimensionality reduction, representation choice, and acquisition effects. Panel A holds the ensemble and UCB fixed while changing the projection; Panel B adds SGES components progressively; Panel C compares mutation representations without search acquisition.

Method	Budget	Initial set	Candidate / proposal setting	Representation	Surrogate / acquisition	Seeds
Random Search	500	Same $N_{\mathrm{init}}=32$	Same mutation search space	No learned PLM surrogate	Random	5
AdaLead	500	Same $N_{\mathrm{init}}=32$	Local evolutionary proposal	No explicit LFS	Fitness-guided local search	5
GP-BO	500	Same $N_{\mathrm{init}}=32$	Same candidate space	Full PLM / high-dimensional	GP Bayesian optimization	5
EVOLVEpro	500	Same $N_{\mathrm{init}}=32$	PLM-guided proposal	PLM supervised signal	Method-specific search	5
AlphaDE	500	Same $N_{\mathrm{init}}=32$	Directed-evolution proposal	Method-specific	Method-specific acquisition	5
BOES	500	Same $N_{\mathrm{init}}=32$	Budgeted search	Embedding-space	BO / evolutionary search	5

Appendix

Table S3: Controlled budgeted-search protocol. Oracle budget, initial labeled-set size, and random-seed count are matched; proposal mechanisms remain method-specific where required by the original algorithm.

Assay set	Count	Primary role
Core assays	10	Detailed LFS analysis / ablations
Extended assays	87	Static large-scale prediction validation
Search-statistics assays	18	Multi-seed search metrics / Wilcoxon tests
Multi-mutant benchmarks	5	Combinatorial transfer evaluation

Appendix

Table S4: Relationship among assay sets. The sets have different evaluation roles and should not be interpreted as inconsistent counts of the same experiment.

Assay	SGES $\rho$	EVOLVEpro $\rho$	SGES Best	EVOLVEpro Best
SPG1_STRSG	0.973	0.865	–	–
BLAT_Firnberg	0.903	0.812	1.691	1.410
PABP_YEAST	0.916	0.808	1.170	1.055
P53_HUMAN	0.874	0.755	2.690	2.150
RASH_HUMAN	0.825	0.720	1.525	1.180
AMIE_PSEAE	0.823	0.705	0.785	0.610

Appendix

Table S5: Full EVOLVEpro comparison. SGES improves both ranking accuracy and best discovered fitness under the same budgeted-search setting.

Search Metric	SGES vs GP-BO	SGES vs AdaLead	SGES vs Random	SGES vs EVOLVEpro
AUC (Best-so-far)	$p=0.007$ $(W^{+}=144)$	$p<0.001$ $(W^{+}=155)$	$p<0.001$ $(W^{+}=171)$	$p<0.001$ $(W^{+}=168)$
Best Fitness Found	$p=0.016$ $(W^{+}=138)$	$p<0.001$ $(W^{+}=152)$	$p<0.001$ $(W^{+}=171)$	$p<0.001$ $(W^{+}=165)$
Budget-to-Top-5%	$p=0.016$ $(W^{+}=138)$	$p<0.001$ $(W^{+}=149)$	$p<0.001$ $(W^{+}=171)$	$p<0.001$ $(W^{+}=162)$

Appendix

Table S6: Extended Wilcoxon signed-rank tests on 18 assays. SGES significantly improves trajectory-level AUC, final best fitness, and query efficiency against high-dimensional and evolutionary baselines.

Assay	SGES	AlphaDE	BOES	LatentDE	TreeNeuralUCB	GP-BO	EVOLVEpro
BLAT_Firnberg	$\mathbf{1.691\pm 0.024}$	$1.485\pm 0.112$	$1.390\pm 0.105$	$1.420\pm 0.145$	$1.285\pm 0.210$	$1.550\pm 0.160$	$1.410\pm 0.120$
P53_HUMAN	$2.690\pm 0.154$	$2.450\pm 0.210$	$2.210\pm 0.315$	$2.380\pm 0.285$	$2.105\pm 0.420$	$\mathbf{2.850\pm 0.280}$	$2.150\pm 0.200$
PABP_YEAST	$1.170\pm 0.091$	$1.085\pm 0.085$	$1.015\pm 0.090$	$1.050\pm 0.110$	$0.985\pm 0.145$	$\mathbf{1.250\pm 0.150}$	$1.055\pm 0.100$
RASH_HUMAN	$\mathbf{1.525\pm 1.014}$	$1.210\pm 0.850$	$0.985\pm 0.620$	$1.140\pm 0.745$	$0.850\pm 0.915$	$0.880\pm 0.450$	$1.180\pm 0.300$
AMIE_PSEAE	$\mathbf{0.785\pm 0.018}$	$0.685\pm 0.045$	$0.640\pm 0.055$	$0.665\pm 0.060$	$0.590\pm 0.095$	$0.720\pm 0.150$	$0.610\pm 0.080$
DYR_ECOLI	$\mathbf{1.065\pm 0.221}$	$0.840\pm 0.145$	$0.765\pm 0.110$	$0.810\pm 0.135$	$0.710\pm 0.185$	$0.850\pm 0.170$	$0.795\pm 0.120$

Appendix

Table S7: Full comparison with recent ML-guided protein optimization methods. All methods use the same budgeted-search protocol. SGES achieves the strongest average performance and rank.

Method	Runtime / Seed	GPU Mem.	Avg. Best
	(min)	(GB)	Fitness
Random Search	2	1.2	1.01
AdaLead ( Sinai et al., 2020 )	15	2.5	1.25
EVOLVEpro ( Jiang et al., 2024 )	42	4.8	1.18
SGES (Ours)	28	3.5	1.68
GP-BO ( Soldát and Kléma, 2024 )	245	18.5	1.33

Appendix

Table S8: Computational cost. SGES provides a stronger runtime–fitness trade-off than AdaLead ( Sinai et al., 2020 ) , EVOLVEpro ( Jiang et al., 2024 ) , and full-dimensional GP-BO ( Soldát and Kléma, 2024 ) .

Method	Mean $\rho$	Median $\rho$	$\rho>0.7$ Ratio	vs. ZS
Ridge (site-delta)	0.584	0.612	35.6%	+0.146
Ridge (mean-pooled)	0.385	0.392	8.0%	-0.053
Zero-shot ESM2	0.438	0.451	14.9%	Baseline

Appendix

Table S9: Summary statistics on 87 ProteinGym assays. Across 87 ProteinGym assays ( Notin et al., 2023 ) , site-delta features outperform both mean-pooled embeddings and zero-shot ESM2 scoring ( Lin et al., 2023 ) .

Figure S1: Pareto frontier of runtime and fitness. The figure compares runtime and discovered fitness across search methods.

Figure S2: Linear probe scatter plots for additional assays. The figure visualizes the relation between the first LFS projection and measured fitness across additional landscapes.

Figure S3: Cross-assay cosine similarity heatmap. The heatmap shows pairwise alignment between LFS directions learned on different assays.

Figure S4: Dense budget scaling curves. The figure shows best-so-far fitness as a function of oracle budget.

Figure S5: Acquisition-function ablation. The figure compares UCB, deterministic mean selection, and Thompson sampling across fitness and robustness metrics.

Figure S6: Biological interpretation of LFS. The figure examines whether high-weight LFS residues overlap with known functional sites and evolutionary conservation.

Group	Mean cos	Max cos	Min cos
Same Protein (BLAT F vs BLAT J)	0.68	0.68	0.68
Same Function (Enzyme vs Enzyme)	0.22	0.35	0.12
Different Function	0.05	0.11	0.01
All Pairs	0.14	0.68	0.01

Appendix

Table S10: Cosine similarity statistics for LFS directions. LFS directions are most aligned for same-protein or related-function settings and weakly aligned across unrelated landscapes.

Train/Test	BLAT-F	BLAT-J	DYR	AMIE	RASH	P53	PTEN	SPG1	PABP	HIS7
BLAT-F	0.865	0.624	0.045	0.012	-0.015	-0.032	0.041	-0.018	0.022	0.015
BLAT-J	0.735	0.650	-0.012	0.025	0.010	0.014	0.025	-0.005	-0.011	0.030
DYR	-0.022	0.015	0.675	0.085	0.042	0.025	0.055	-0.022	0.014	0.045
AMIE	0.010	0.033	0.051	0.758	0.022	-0.015	0.018	0.010	-0.005	0.020
RASH	-0.015	0.011	0.022	-0.018	0.755	0.085	0.042	0.015	0.035	0.012
P53	-0.005	-0.020	0.018	0.010	0.045	0.810	0.065	0.022	-0.012	0.008

Appendix

Table S11: Cross-assay subspace transfer matrix. Each entry reports Spearman

\rho

when an LFS direction learned on the training assay is evaluated on the test assay.

Assay	#Muts	SGES MAE $\downarrow$	ZS MAE $\downarrow$	SGES $r{\uparrow}$	ZS $r{\uparrow}$
PABP_YEAST	2–3	0.450	1.240	0.520	0.180
SPG1_STRSG	2–3	0.280	0.850	0.650	0.220
GFP	2–8	0.685	1.585	0.420	0.115
GB1	2–4	0.515	1.120	0.485	0.140
AAV	2–7	0.590	1.345	0.465	0.125
Average	–	0.504	1.228	0.508	0.156

Appendix

Table S12: Multi-mutant generalization. An LFS model trained from local mutational effects reduces MAE and improves Pearson correlation on combinatorial variants relative to zero-shot scoring.

Assay	SGES $\rho$	PCNN $\rho$	SGES Best	PCNN Best
BLAT_Firnberg	0.903	0.780	1.691	1.350
P53_HUMAN	0.874	0.715	2.690	2.100
RASH_HUMAN	0.825	0.680	1.525	1.150
PTEN_HUMAN	0.593	0.520	1.939	1.550
Average	0.825	0.685	1.471	1.215

Appendix

Table S13: Comparison with PCNN under budget constraints. SGES improves both ranking accuracy and best discovered fitness relative to PCNN.

Assay	Top Mutations	Predicted Mechanism	Score
BLAT_Firnberg	M182T, E104K, G238S	Stabilizing / extended spectrum	9.2
P53_HUMAN	V143A, Y220C	Core packing restoration	8.5
AMIE_PSEAE	S293T, F154Y	Active-site volume optimization	8.8

Appendix

Table S14: Functional annotation of representative top mutations discovered by SGES. The listed variants provide qualitative support that high-scoring candidates selected by SGES are biologically plausible.

Explore similar work

ProtLingo: Efficient Protein Language Modeling via Conditional Memory and Expert Routing
Sep 7, 2026Mingrui Li, Sixian Shen, Minzhang Li +4Language Modeling
ShanghaiTech University
Towards A Generative Protein Evolution Machine with DPLM-Evo
Apr 30, 2026Xinyou Wang, Liang Hong, Jiasheng Ye +4Alphaevolve Evolutionary Optimization Methods
ByteDance Seed · Nanjing University · CUHK +1
PFArena: Benchmarking Language Models for Protein Modification
Sep 24, 2026Yawen Ouyang, Xinbo Zhang, Ziyuan Ma +15Protein Quality-Driven Selective Mutation
Shanghai Artificial Intelligence Laboratory · Generative Symbolic Intelligence Lab (GenSI), Tsinghua University · Institute for AI Industry Research (AIR), Tsinghua University +1

Sep 7, 2026, cs.AIJ/K move · Enter open · S save

ProtLingo: Efficient Protein Language Modeling via Conditional Memory and Expert Routing

Mingrui Li, Sixian Shen, Minzhang Li, Ruiyi Zhang+3

ShanghaiTech University

Language Modeling

Open paper AskPDF arXiv

Proteins perform diverse cellular functions, and even single amino-acid substitutions can alter stability, activity, or molecular interactions. Protein language models (PLMs) provide a scalable approach for modeling such sequence--function relationships from unlabeled sequences, but increasing the size of dense Transformer backbones often brings substantial computational cost without consistently improving mutation-sensitive prediction. We introduce ProtLingo, an efficient PLM framework that augments a pretrained single-sequence backbone with conditional local memory and sparse expert routing. ProtLingo maps contextual residue representations into route-specific discrete codes, composes centered local windows into latent

N

-gram addresses, and retrieves reusable residual signals associated with recurring local sequence contexts. In parallel, selected feed-forward blocks are upcycled into sparse Mixture-of-Experts layers with shared and routed experts, enabling residue-dependent computation while activating only a subset of parameters. Experiments on protein fitness prediction, FLIP benchmarks, and supervised contact prediction show that ProtLingo achieves competitive performance with a 150M-scale backbone, including strong parameter efficiency on mutation-effect prediction and preserved long-range structural representations.

Abstract

Explore similar work

ProtLingo: Efficient Protein Language Modeling via Conditional Memory and Expert Routing

Towards A Generative Protein Evolution Machine with DPLM-Evo

PFArena: Benchmarking Language Models for Protein Modification