PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.07607v1q-bio.QMcs.AIcs.LG

TL;DR

Zero-shot PLM scores may be misaligned with assay-specific fitness, while supervised search in high-dimensional embeddings can be sample-inefficient under limited oracle budgets. This paper learns a compact assay-specific subspace from mutation-induced PLM changes and reports improved prediction and search across benchmark settings.

  • Problem

    Under limited oracle budgets, zero-shot PLM scores may misalign with assay fitness, while high-dimensional supervised search can make surrogate modeling and uncertainty estimation sample-inefficient.

    Source: §1 · p. 1

  • Method

    SGES estimates a compact, assay-specific Linear Fitness Subspace from mutation-local PLM site-delta features, then performs surrogate modeling, uncertainty estimation, and acquisition in that subspace.

    Source: §1 · p. 2

  • Results

    Across 10 core ProteinGym assays and 87 extended assays, SGES improves over zero-shot PLMs and multiple ML-guided search baselines under matched oracle budgets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support mutation-aware, low-dimensional representations as useful coordinates for sample-efficient protein optimization across the evaluated benchmarks.

    Source: §6 · p. 8

    LFS discovery relies mainly on single-mutant observations and assumes the frozen PLM encodes the target property; strongly epistatic or nonsequence-dependent settings may need interaction or multimodal extensions.

    Source: §7 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Model-guided directed evolution must find high-fitness variants with limited oracle budgets, but zero-shot PLM scores can be misaligned with assay fitness and high-dimensional supervised search can be sample-inefficient. This paper proposes an assay-specific Linear Fitness Subspace (LFS) in mutation-local PLM representations and uses it to guide search.

Motivation

  • Limited oracle budgets make accurate fitness prediction and efficient exploration of protein sequence space central challenges for model-guided directed evolution.Each candidate requires experimental or computational oracle evaluation.Source: §1 · p. 1

Linear Fitness Subspace

  • In mutation-local site-delta features, a compact, assay-specific set of supervised directions makes a substantial component of fitness linearly accessible.The LFS claim is local and does not assert that global PLM embeddings or complete fitness landscapes are universally linear.Source: §1 · p. 2

Geometric view of zero-shot scoring

  • Zero-shot PLM scoring is a fixed task-agnostic direction, while few-shot supervision can align the projection with assay-specific fitness directions.The paper relates poor zero-shot performance to misalignment with the target assay direction.Source: §1 · p. 2

Subspace-guided search

  • SGES estimates the LFS from a small initial query budget, then performs surrogate modeling, uncertainty estimation, and acquisition in the learned subspace.The framework also uses diversity filtering.Source: §1 · p. 2

Controlled attribution and broad validation

  • Across 10 core ProteinGym assays and 87 extended assays, SGES improves over zero-shot PLMs and multiple search baselines under matched oracle budgets.Matched-dimensionality controls, weaker mutation-feature baselines, component ablations, and acquisition ablations help separate fitness alignment from dimensionality reduction and acquisition effects.Source: §1 · p. 2

2 Related Work

Prior work uses PLMs for zero-shot mutation scoring and supervised prediction, while directed-evolution methods combine surrogates, acquisition, and sequence-generation strategies. This paper focuses more narrowly on locating assay-specific fitness directions in mutation-local PLM representations and searching within them.

PLM fitness prediction

  • PLMs support zero-shot mutation scoring and supervised prediction, but global or mean-pooled embeddings can obscure mutation-local residue perturbations.Classical evolutionary and structural signals also remain relevant comparison baselines.Source: §2 · pp. 2–3

ML-guided directed evolution

  • Directed-evolution methods combine surrogate models, acquisition functions, and proposal mechanisms, but often treat PLM embeddings as high-dimensional black-box features.When many dimensions are weakly related to fitness, surrogate fitting and uncertainty estimation can be sample-inefficient.Source: §2 · p. 3

Uncertainty estimation

  • SGES estimates uncertainty after projecting site-delta features into the learned LFS, concentrating exploration on directions that empirically covary with fitness.This design addresses concern that uncertainty in high-dimensional PLM spaces may reflect irrelevant embedding variation.Source: §2 · pp. 3–4

Representation geometry

  • LFS makes a narrower claim than general representation-linearity work: assay-specific mutation effects occupy supervised directions of residue-level site-delta displacements.Those directions also serve as coordinates for budgeted search.Source: §2 · p. 4

Comparison with PPDE

  • PPDE uses a sequence-space sampler, whereas SGES estimates an assay-specific LFS for prediction and uncertainty modeling.Figure 1 shows SGES iterating through surrogate prediction, uncertainty-aware acquisition, oracle evaluation, and population update.Source: Figure 1 · p. 3, §2 · p. 4

3 The Linear Fitness Subspace Hypothesis

The section defines mutation-induced site-delta features and proposes that assay fitness is locally linearly accessible in a compact subspace of those features. Empirical results support linear accessibility and low dimensionality, motivating search in that learned coordinate system.

3.1 Problem Setup and Notation

  • For a mutation at position p, the feature δ(s) is the mutant’s residue-level PLM representation minus the wild type’s representation at that site.This local displacement isolates mutation-induced change and removes most sequence-level background shared by mutant and wild type.Source: §3.1 · p. 4

3.2 Hypothesis Statement

  • The LFS hypothesis proposes compact, supervised directions in which mutation-induced fitness variation is approximately linearly accessible, without claiming globally linear fitness landscapes.Mutants are projected as z(s) = W^Tδ(s); zero-shot scores are analyzed as a fixed one-dimensional direction and compared with a supervised assay direction.Source: §3.2 · p. 4

3.3 Empirical Validation

  • 0.763 average Spearman correlation for ridge regression on site-delta features, versus 0.825 for nonlinear MLP; the average linearity ratio is 0.921, while mean-pooled features score 0.574.These comparisons support mutation-aware site-delta features rather than mean-pooled features as the coordinate for linear fitness prediction.Source: Table 1 · p. 5, §3.3 · p. 5
  • Performance saturates at about 5–10 LFS components, and the first direction orders variants by fitness in representative assays.Feature ablation identifies site-delta as the dominant predictor, while zero-shot–LFS alignment helps explain when zero-shot scoring fails.Source: §3.3 · p. 5

4 Method: Subspace-Guided Evolutionary Search

SGES learns an assay-specific subspace from mutation-induced site-delta features, then performs search in that lower-dimensional coordinate system. Its surrogate and uncertainty estimates operate on fitness-aligned directions rather than generic high-dimensional PLM embeddings.

4 Method: Subspace-Guided Evolutionary Search

  • SGES first discovers an assay-specific fitness subspace from a small labeled set, then trains its search surrogate on projected site-delta features.This distinguishes the method from standard PLM-guided search using generic high-dimensional embeddings.Source: §4 · p. 5

4.1 Subspace Discovery

  • PLS learns the projection from labeled site-delta features and fitness, selecting directions that covary with fitness; SGES uses k = 8 in its main experiments.The chosen dimension follows saturation analysis: performance typically saturates around k = 5–10, with k@95% = 8.Source: §4.1 · p. 5

4.2 Subspace-Guided Search

  • In each search iteration, SGES projects candidates into the LFS and scores them with a deep-ensemble surrogate and upper-confidence-bound acquisition.For multi-mutant assays it applies Hamming-distance diversity filtering and re-estimates the subspace every 50 oracle evaluations.Source: §4.2 · pp. 5–6

4.3 Why Subspace Search Helps

  • Projecting into the LFS reduces the effective dimension from d to k, focusing prediction and uncertainty on coordinates with empirical covariance with fitness.This is intended to reduce data demands for surrogate fitting and focus exploration on coordinates more likely to change fitness.Source: §4.3 · p. 6

5 Experiments

Across the budgeted-search evaluation, SGES performs strongly against recent optimization baselines, and controlled analyses favor fitness-aligned LFS coordinates over generic dimensionality reduction. Additional analyses examine biological interpretability, multi-mutant transfer, and performance against PCNN.

5 Experiments

  • The evaluation separates 10 core assays for detailed analyses, 87 extended assays for static prediction, and 18 assays with complete multi-seed budgeted-search traces.Budgeted search uses Bmax = 500, Ninit = 32, and five random seeds for each supervised method.Source: §5 · p. 7

5.1 Main Search Results

  • SGES achieves the best average best fitness and average rank; against EVOLVEpro, average Spearman correlation rises from 0.725 to 0.825 and average best fitness from 1.183 to 1.471.On 18 search-evaluation assays, SGES significantly improves AUC, best fitness, and Budget-to-Top-5% against GP-BO, AdaLead, Random Search, and EVOLVEpro.Source: Table 2 · p. 7, §5.1 · p. 7

5.2 Ablation and Subspace Analysis

  • At k = 8 with the same deep-ensemble UCB pipeline, PLS-LFS reaches average Spearman 0.804 and average best fitness 1.45, exceeding PCA, random projection, and label-shuffled PLS.The reported comparisons are 0.624/1.22 for PCA, 0.512/1.08 for random projection, and 0.485/1.02 for label-shuffled PLS; component ablations show progressive gains from LFS, ensemble uncertainty, UCB, and diversity filtering.Source: §5.2 · p. 7

5.3 Additional Analyses

  • High-weight LFS residues concentrate near known functional sites in BLAT, P53, and PTEN and are enriched at annotated active or binding residues.The authors interpret this as evidence that the subspace is aligned with biologically meaningful constraints.Source: §5.3 · p. 7
  • Across five multi-mutant benchmarks, the LFS model reduces average MAE from 1.228 to 0.504 and raises average Pearson correlation from 0.156 to 0.508 versus zero-shot scoring.The authors caution that this captures a strong additive component, not a complete model of higher-order epistasis.Source: §5.3 · pp. 7–8
  • On four evaluated assays, SGES improves average Spearman correlation from 0.685 to 0.825 and average best fitness from 1.215 to 1.471 versus PCNN.Source: §5.3 · p. 8

6 Conclusion

The conclusion presents LFS as an assay-specific, supervision-recoverable coordinate in mutation-induced PLM site-delta space. It positions SGES as a low-dimensional search framework and reports improved prediction and budgeted search across ProteinGym benchmarks.

6 Conclusion

  • LFS is an assay-specific, supervision-recoverable coordinate in mutation-induced PLM site-delta space where fitness is locally linearly accessible.Source: §6 · p. 8
  • SGES combines low-dimensional surrogate modeling with uncertainty-aware acquisition in the learned subspace.Source: §6 · p. 8
  • Across ProteinGym benchmarks, SGES improves prediction and budgeted search over zero-shot PLMs, high-dimensional GP-BO, EVOLVEpro, and recent optimization baselines.The authors describe these results as supporting low-dimensional, mutation-aware representations for sample-efficient protein optimization.Source: §6 · p. 8

7 Limitations

The limitations bound SGES’s coverage of epistasis, calibration, and properties not encoded by its frozen PLM. Supplementary experiments compare zero-shot and few-shot prediction and use controlled comparisons to separate fitness alignment from generic compression and acquisition effects.

7 Limitations

  • Single-mutant subspace discovery captures additive directions but may miss higher-order epistasis, and strongly epistatic settings may require interaction modeling.Multi-mutant results nevertheless suggest that LFS remains useful.Source: §7 · p. 8
  • Assay-dependent calibration makes SGES primarily a ranking and acquisition method rather than a fully calibrated probabilistic model.Source: §7 · p. 8
  • SGES assumes the frozen PLM encodes the target property; assay conditions, molecular partners, or conformational states may require structural or multimodal extensions.Source: §7 · p. 8

A.1 Zero-shot and Few-shot Fitness Prediction

  • ProSST’s strongest zero-shot average is ρ = 0.597, versus ρ = 0.686 with 5% supervised LFS and ρ = 0.825 with 50% supervised LFS.The comparison tests whether few labels recover assay-specific directions that fixed zero-shot projections miss.Source: §A.1 · p. 12

A.2 Controlled Attribution and Protocol Clarifications

  • Matched-dimensionality controls show that PLS-LFS outperforms PCA, random projection, and label-shuffled PLS, while site-delta PLM features exceed the tested mutation-feature alternatives.The controls also evaluate progressively adding ensemble uncertainty, UCB, and diversity filtering; search rows match oracle budget, initial set, applicable space, and seeds.Source: §A.2 · p. 12
  • The zero-shot alignment analysis compares a direction regressed from zero-shot scores with supervised PLS directions, using cos(w_zs, w_lfs) as the reported alignment.The main search uses k = 8, selected from a saturation analysis with k@95% = 8.Source: §A.2 · p. 12

A.3 Search Baseline Comparisons

The comparisons test whether SGES’s search advantage comes from its learned coordinate rather than supervision alone or a standard search component. Across the reported comparisons, SGES performs strongly against supervised PLM-guided and other optimization baselines, while not winning every assay.

A.3 Search Baseline Comparisons

  • Table S1 compares zero-shot and few-shot prediction on 10 assays; its caption reports that 5% supervised LFS exceeds strong zero-shot baselines.Source: Table S1 · p. 13
  • Table S2 separates projection choice, progressive SGES components, and mutation representations, while the search comparisons distinguish Random Search, AdaLead, EVOLVEpro, and full-dimensional GP-BO.Source: §A.3 · p. 12, Table S2 · p. 13
  • The controlled budgeted-search protocol matches oracle budget, initial labeled-set size, and seed count, while allowing method-specific proposal mechanisms where required.The assay sets have different evaluation roles rather than representing inconsistent counts of one experiment.Source: Table S3 · p. 13, Table S4 · p. 13
  • SGES raises average ranking correlation from 0.725 to 0.825 and average best fitness from 1.183 to 1.471 versus EVOLVEpro.The comparison supports a learned low-dimensional site-delta coordinate as a more effective search representation than generic PLM feature use.Source: §A.3 · pp. 13–14
  • SGES wins 6 out of 8 assays and has the best average rank, while GP-BO wins on P53 and PABP.Source: §A.3 · p. 14

A.4 Computational Cost, Implementation Details, and Large-scale Validation

Extended experiments assess generalization beyond the core assays, computational cost, implementation, and supplementary diagnostics. They report broad site-delta prediction gains and favorable search outcomes alongside evidence about runtime, linear structure, assay-specific geometry, and search trajectories.

A.4 Computational Cost, Implementation Details, and Large-scale Validation

  • SGES offers a favorable runtime–fitness trade-off among nontrivial optimizers, costing substantially less than GP-BO while achieving higher average best fitness.GP-BO requires nearly an order of magnitude more runtime and much higher memory.Source: §A.4 · p. 14, Table S8 · p. 15
  • SGES extracts frozen-PLM site-delta features, learns an LFS projection by partial least squares, and fits a deep-ensemble surrogate in that subspace for UCB candidate ranking.It updates the active set after oracle rounds and periodically re-estimates the projection as search moves from the wild-type neighborhood.Source: §A.4 · pp. 14–15
  • Across 18 assays, SGES significantly improves trajectory-level AUC, final best fitness, and query efficiency against high-dimensional and evolutionary baselines, and achieves the strongest average performance and rank in the recent-method comparison.Source: Table S6 · p. 15, Table S7 · p. 15
  • Across 87 ProteinGym assays, site-delta Ridge improves mean correlation by 0.146 over zero-shot ESM2 and raises the fraction above ρ > 0.7 from 14.9% to 35.6%.Source: Table S9 · p. 15, §A.4 · p. 15

B.1 Runtime–Fitness Trade-off

  • The runtime–fitness frontier shows SGES’s favorable compromise relative to high-dimensional optimization methods, avoiding GP-BO’s computational burden while retaining strong optimization quality.Source: §B.1 · p. 16

B.2 Additional Linear Probe Scatter Plots

  • Additional assay scatter plots show repeated ordering of variants along the first LFS direction, although the relationship is not equally tight across assays.The passages attribute variation in tightness to possible epistasis or measurement noise.Source: §B.2 · p. 16

B.3 Cross-Assay Subspace Geometry

  • Related assays have stronger alignment between learned LFS directions, while unrelated functional classes remain weakly aligned.The pattern supports structured assay-specific geometry rather than a universal direction or arbitrary noise.Source: §B.3 · p. 16

B.4 Dense Budget Scaling

  • SGES improves early and maintains stronger best-so-far fitness across the query budget, indicating gains along the search trajectory rather than only at the endpoint.The dense budget curves test early exploration as well as final fitness.Source: §B.4 · p. 16

B.5 Acquisition-function Ablation

The acquisition ablation favors UCB when uncertainty is computed in the LFS. The biological interpretation analysis links high-weight residues to functional sites and conservation, while stopping short of a complete mechanistic map.

B.5 Acquisition-function Ablation

  • UCB provides the most reliable stability–fitness trade-off; deterministic mean selection can exploit early predictions too aggressively, while Thompson sampling adds stochasticity.In the LFS, uncertainty is concentrated along directions with empirical covariance with fitness.Source: §B.5 · p. 16

B.6 Biological Interpretation

  • High-weight LFS residues cluster near known functional sites and associate with conservation, supporting biologically meaningful constraints without implying a complete mechanistic map.The figure examines overlap with functional sites and evolutionary conservation.Source: §B.6 · pp. 16–17, Figure S6 · p. 19

C Cosine Similarity and Cross-Assay Transfer

LFS directions are assay-specific rather than globally universal, yet show stronger alignment for same-protein or functionally related assays. Cross-assay transfer is selective, with the strongest off-diagonal transfer between the two BLAT assays.

C Cosine Similarity and Cross-Assay Transfer

  • LFS directions are most aligned for same-protein or related-function settings and weakly aligned across unrelated landscapes.This pattern is consistent with directions adapting to each assay while reflecting shared biological structure.Source: C Cosine Similarity and Cross-Assay… · p. 20, Table S10 · p. 20
  • Different-function assay pairs have average cosine similarity 0.05, while the same-protein BLAT pair reaches 0.68.These statistics support an assay-specific but biologically structured view of LFS directions.Source: C Cosine Similarity and Cross-Assay… · p. 20
  • Cross-assay transfer is selective: the strongest off-diagonal transfer is between the two BLAT assays, while most unrelated train–test pairs are near zero or negative.The transfer matrix reports Spearman ρ for an LFS learned on a training assay and evaluated on a test assay.Source: C Cosine Similarity and Cross-Assay… · p. 20, Table S11 · p. 21

D Multi-mutant and Spectral-method Results

The multi-mutant experiments find that LFS learned from local single-mutation effects remains informative on combinatorial variants, and SGES improves over PCNN under budget constraints. These results support additive transfer, not a claim that LFS resolves all epistasis.

D Multi-mutant and Spectral-method Results

  • The multi-mutant findings support additive transfer but do not establish invariance to large epistasis residuals or that LFS fully solves epistasis.The reported multi-mutant model uses the additive representation, and remaining error is expected to reflect higher-order mutation interactions requiring explicit modeling.Source: D Multi-mutant and Spectral-method… · p. 21
  • Average MAE decreases from 1.228 to 0.504, and average Pearson correlation increases from 0.156 to 0.508 on combinatorial variants.The aggregate result is evidence that an LFS model trained from local mutational effects transfers beyond single mutants relative to zero-shot scoring.Source: D Multi-mutant and Spectral-method… · p. 21, Table S12 · p. 21
  • SGES improves average Spearman correlation from 0.685 to 0.825 and average best fitness from 1.215 to 1.471 relative to PCNN under budget constraints.PCNN models fitness-landscape topology, whereas SGES models mutation-induced PLM representation geometry.Source: D Multi-mutant and Spectral-method… · p. 21, Table S13 · p. 21
  • Annotations of representative SGES-selected mutations are consistent with stability improvement, altered catalytic or binding geometry, and core-packing effects.The listed top variants provide qualitative support for biological plausibility and align with enrichment patterns in Figure S6.Source: D Multi-mutant and Spectral-method… · pp. 21–22, Table S14 · p. 22
Loading 2610.07607v1…