RankFlow: Property-aware Transport for Protein Optimization
Authors: Lu Yu
Wei Xiang
Kang Han
Gaowen Liu
Ramana Rao Kompella
Venue: ICLR 2026 Poster (2026)
Keywords: protein language models
fitness prediction
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: A key step in protein optimization is modeling the fitness landscape, which maps proteins to functional assay readouts. Existing methods typically either use property-agnostic likelihoods/embeddings from pretrained protein language models (PLMs) for fitness prediction, or assume independent mutational effects, limiting their ability to capture higher-order interactions. In this work, we introduce RankFlow, a conditional flow framework that refines PLM representations to be a property-aligned distribution via a tailored energy function and captures multi-mutation interactions through learnable embeddings. To align optimization with evaluation protocols, we propose the Rank-Consistent Conditional Flow Loss (RC$^2$), a differentiable ranking objective that enforces the correct order of mutants rather than absolute values, which improves out-of-distribution generalization. Finally, we introduce a Property-guided Steering Gate (PSG) that concentrates learning on positions carrying signals for the target property while suppressing unrelated evolutionary biases. Across the ProteinGym, PEER, and FLIP benchmarks, RankFlow obtains state-of-the-art ranking accuracy and superior generalization performance.

Published as a conference paper at ICLR 2026
RANKFLOW: PROPERTY-AWARE TRANSPORT FOR PROTEIN OPTIMIZATION
Lu Yu?.2, Wei Xiang Kang Han', Gaowen Liu", and Ramana Rao Kompella²
1 School of Computing, Engineering and Mathematical Sciences, La Trobe University, Melbourne, Australia
2 Australian Centre for Al and Medical Innovation, La Trobe University, Melbourne, Australia
3 Cisco Research, San Jose, CA
[1. yu, w. xiang, k.han M @latrobe..............o wwwrrobeeeede.. edu. fgaoliu, tkompell/@cisco.com
ABSTRACT
A key step in protein optimization is modeling the fitness landscape, which maps proteins to functional assay readouts. Existing methods typically either use property-agnostic likelihood//emmeeeddnng from pretrained protein language models (PLMs) for fitness prediction, or assume independent mutational effects, limiting their ability capture higher-order interactions. this work, we intro- duce RankFlow, a conditional flow framework that refines PLM representations be property-aligned distribution via a tailored energy function and captures multi-mutation interactions through learnable embeddings. To align optimization with evaluation protocols, we propose the Rank-Consia Consistent Conditional Flow Loss (RC*), a differentiable ranking objective that enforces the correct order of mutants rather than absolute values, which improves -ittt--distributionn generalization. Fi- nally, introduce croperty-guided Steering Gate (PSG) that concentrates learn- ing on positions carrying signals for the target property while suppressing unre- lated evolutionary biases. Across the ProteinGyym PEER, and FLIP benchmarks, RankFlow obtains state-of-the-art ranking accuracy and superior generalization performance.
INTRODUCTION
Proteins enable catalysis, structural support, molecular transport, and cellular signaling. The ability to design protein sequences and structures with tailored properties has a direct impact on sustainable manufacturing and therapeutics. In practice, optimization targets series of properties that can include stability, binding affinity, and enzymatic activity, which are typically measured by high- throughput assays (Biswas et 202T). A central objective is learn property landscapes that map variants to these assay readouts and that explain how single and combined mutations modulate them (Notin et al. 2024). The more accurately we model and navigate such landscapes, the more reliably we can propose edits that deliver targeted behavior in realistic experimental settings.
Protein optimization typically begins by modeling the mapping from sequence and structure to func- tional readouts, often called the filness landscape (Romero & Arnold 2009). A significant obstacle in modeling the fitness landscape the scarcity of experimentally measured labels. This has moti- vated self-supervised representation learning for mutation effect prediction (Hopf et al. 2017 Meier let al. 2021). Early approaches, including Riesselman et al. (2018); (001)):Frazzz Frazer et alla(22222) (2021), ); Laine et al. (2019) have modeled family-specific distributions from multiple sequence alignments (MSAs) and used those priors to score variants. More recently, hybrid designs have combined family-specific
and fanosic information for fitness prediction (Alley al. 2019 Rives al. 2021 Those PLMs have shown that their predicted likelihoods can infer evolutionary trajectories and predict zero-shot mutation effects (Rao al.) 2021 Notin et al. (002). Subsequently, with the growth of tightthroughput assays and broader annotated datasets, many studies have moved to supervised training (Gelman al. [021] Heinzinger et al., 2019). They fine-tune PLMs on experimental read- outs and minimize differences between predicted and measured property values.
1 Corresponning Author

---

Published as a conference paper at ICLR 2026
Regression head
RankFlow
40 training assays
ESM2 SaProl regression
Performance
Ve
Held-out test assay
Rankl'low
Vo
predicted high-fiteess region
regression f/:
RankFlow Y/
Spearman of Reg.
Spearman of RankFlow
Regression RankFlow
(b) Generalization
Figure I: Illustration of RankFlow. (a) In a local region of representation space, a regression head on PLM embeddings overfit when fine-tuned single assay, mapping some truly high-fitness mutants to low scores and vice versa. RankFlow instead reshapes wild-type-conditiones mutant rep- resentations into fitness-aligned distribution, enforcing property-aware landscape. (b) In cross- assay generalization experiment, models are trained on 40 Deep Mutational Scanning (DMS) assays from the same category and evaluated on held-out assay; RankFlow achieves higher Spearman correlation than fine-tuned ESM2/SaProt with regression heads, indicating stronger generalization.
However, two critical aspects remain underaddressed. First, while PLMs capture rich evolutionary regularities, PLM representations property-agnostic (Loell & Nanda 2018): they encode mul- tiple, often competing evolutionary constraints (e.g.. foldability, stability, expression), which can dilute or even oppose the signal of the property of interest. Directly using such embeddings can bias predictions toward wild-type-likk preferences rather than property-aligner improvements. Second, some fitness prediction methods overlook multi-mutation interactions (epistasis) and assume addi- tive effects of individual mutations (Notin et al.) 2023a Wang et al., (024). This leads to wrong predictions, especially for higher-order mutants where non-additive interactions dominate and sig- nificantly influence the target property (Dauparas et al. 2022: Chen et al. 2024).
In this work, we present RankFlow, a property-aware conditional flow framework designed to ad- dress both limitations while preserving the generalization ability of pretrained PLMs. Standard approaches attach a regression head to PLM embeddings and fine-tune on a single DMS assay, which can yield strong in-assay performance but tends to introduce dataset-specifie bias and inferior cross-pross cross-protein generalization. Rather than performing point-to-po to-iint regression, RankFlow keeps the PLM embeddings as a base sequence/structure prior and learns energy-guided conditional flow that reshapes hild-type-conditioned mutant representations into fitness-aligned distribution. This refinement captures multi- mutation interactions through learnable embeddings defined on mutation sets, going beyond additive assumptions. The energy function is constructed from observed fitness scores and steers the learned dynamics toward landscapes where high-fitness mutants consistently outrank fow-fitness ones. Figure | illustrates how RankFlow differs from regression methods. To align training with the evaluation protocols of protein engineering, we propose the Rank-Consisee Conditional Flow Loss (RC*), a differentiable soft-ranking objective that enforces the correct order of mutants rather than absolute values, improving robustness noise and generalization to unseen assays. Finally, we introduce a Property-guiee -uided Steering Gate (PSG), which focuses learning positions carrying signals for the target property and suppresses directions reflecting unrelated evo- lutionary biases, reducing wild-type preference and sharpening the edit signal. We evaluate Rank- Flow on ProteinGym (Notin et al. 2023a), 33- Lactamase and Fluorescence from PEER (Xu et al.) 2022). and GB1 from FLIP (Dallago al. 202T). where it consistently outperforms state-of-the-an methods in ranking accuracy and generalization. Our work makes the following contributions:
We propose RlokFlow, a novel conditional flow framework that refines PLMM-dervee mu- tant representations into property-aligee embeddings, capturing non-additive interactions across mutation sets.
We introduce the Conk-Consistent Conditional Flow Loss (RC2 Loss), a rank-ccosistentt flow objective that aligns training with rank-based evaluation metries and improves out-of- distribution generalization.
We develop the copertty-guides Steering Gate (PSG) mechanism, which focuses learning property-relevant positions and reduces wild-type bias in PLM representations.
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
Protein Representation Learning. Protein representation learning has seen significant advance- ments with the introduction of protein language models (PLMs) that leverage large-scale protein sequence data (Rives et al. (021). Farly works focused on family-specific models trained on MSAs (Riesselman et al. 2018, Hoopf et al. 2017: Laine al., 2019), which capture evolution- ary constraints within protein families. More recent approaches have developed family-agnostic PLMs (Rao al., 202T Lin et al. 2023 Yang al. 2023) that learn from vast and diverse protein sequences, enabling zero-shot prediction of mutational effects and improved generalization across different proteins. Hybrid models that combine Family-specife and -anily-agnostic information have also been proposed (Alley et al.) 2019, Notin et al.) 2022, Hsu et al. 2022] Su et al.) 2024), further enhancing the predictive capabilities of PLMs. These models typically use self-supervised objectives, such as masked language modeling or autoregressive prediction, to learn rich represen- tations that can be fine-tuned for specific downstream tasks.
Protein Fitness Prediction. Protein fitness prediction aims to model the relationship between proteins and their functional properties, such stability and activity (Dallago al. 2021 Gelman et al. 2021). Traditional approaches directly map sequences to measured properties via regression or classification (Gelman al. 2021 Heinzinger
al., 2019). while alignment- baseeddt--bddddd-eee models fo- cus specific protein families and use evolutionary information from MSAs to infer mutational effects (Riesselman al. 2018 Laine et al. 2019 Shin al. 2021). Unsupervised methods predict mutation effects by evaluating changes in likelihood scores (Riesselman et al. 20188] Frazer et al.)
2021 Notin et al., 2022; Nijkamp et al., 2023] Truong Jr & Bepler (023). and more recent work
leverages PLM representations within supervised frameworks for improved fitness prediction (Wang) et al. 2024 Notin et al. 2023b Tan et al. 2025). Despite these advances, existing methods primar- ily optimize for absolute fitness regression, which can limit robustness when extrapolating unseen mutation combinations. Notin et al.] (2023)) introduces a non non-nonnnne con--aaametric Transformer that jointly embeds protein sequences and labels and excels in label-scarce settings. |Beck et al. (2025) propose a meta-learning approach that trains an in-context regressor over more than 100) fitness tasks, but assumes access many related tasks and excludes proteins longer than 750 amino acids due to memory limits. Both methods are also computationally and memory-intensive during training and inference. Kermut (Groth et al., 2024) combines sequence and structure information in a Gaussian process with composite kernels and achieves state-of- sfte-ooftthe-rr performance; however, it inherits the cubic scaling of exact GPs and thus becomes costly on large variant libraries, and requires trun- cation or subsampling for long or dense assays (Groth et al.) 2024). Ronen et al. (2025) further shows that Kermut and related predictors can sensitive to the choice of representation and use prediction-basse based screening plus representation ensembling to stabilize performance and improve -utoo--distriuutinn generalization and uncertainty quantification. In contrast, our method learns lightweight, property-awee aware flow top of PLM representations, and explicitly models multi- mutation interactions to enable strong generalization even on large and diverse variant libraries.
METHOD
This section describes the proposed RankFlow framework. Section 3.1 presents the problem setup and model overview. Section introduces the Property-aware Conditional Flow, which constructs new property-guided flow to generate the target distribution. In Section 3.3 describe the Steering Gate. Section 3.4 details the RankFlow architecture and training and inference procedures.
3.1 PROBLEM SETUP AND MODEL OVERVIEW
Protein optimization seeks to quantify how sequence edits change a property of interest. Let a wild type protein be wt [w ...., wt with N amino acids (AAs). A mutation is a set of substitutions μ == {/m : xx > 2m |m E [I,...,,)} which may act on multiple positions at once. Applying
μ yields mutant sequence me [a:ih: ..... ......N mL] For wild type protein and given assay, typical learning objective is to fit a function Fo Fo(x"" wt = y. where is the measured property value of the mutation ya wt
3

---

Published as a conference paper at ICLR 2026
Multi-medal Fusion Encoder
Property-earee
Property-guidee
Steoring
Widd-tpe sequence
SABVATS-SDDIGG....
|5a AA
Positive
ESM-2 Inter Scace
ESM-IF
E5M-2 Internal Spacee
ESM-2
Steering
pdo
Encoder
Space
|S-1 A
Nogathve
Decoder
Head
ho
Ränk consistent
Flow Matching Model G(0)
PLM Head
Mutti-modal
Energy Function (,(e)
1.05
Rank-Consisseent
Unified Rupresematiin
Score
Flow
Figure 2: Overview of RankFlow for protein fitness prediction. RankFlow takes wild-type- conditioned PLM representations the source distribution Po and learns a property-aware condi- tional flow (Section 3.2 that transports Po to a property-aligned target q. The flow is conditioned on the wild type context (via the Multi-modal Fusion Encoder), the mutation set, and a Property-guided Steering Gate (Section 3.3 that emphasizes positions relevant to the target property. Training uses dual-objective loss that combines property-awer -ware flow matching with the rank-consistemt RC² loss.
In RankFlow, rather than fitting a deterministic predictor Fo that is prone to overfitting in the low- data regime, we train a conditional flow parameterized by 0 that models distribution over protein representations whose density is tilted toward the target property via energy function. This flow takes PLM representations as the source distribution Po and transports it to a property-aware distribution q conditioned the mutated AAs with the PLM's mask token, feed the sequence into the PLM, and extract the final hidden representation ho E R xd prior to the output head. Masking removes self-information about the wild type residues, yielding a mutation-aware yet wild-type-specific state hg ~ which serves the input distribution for RankFlow. The overview of RankFlow is depicted in Figure 2
and y. Specifically, for each mutant a:!tt in the dataset, mask
3.2 PROPERTY-AWARS CONDITIONAL FLOW
Continuous Property-awane Flow. Inspired by Zhang al. (2025). we learn conditional flow that generates samples not directly from the base distribution Pu, but from an energy svergy-guided dis- tribution qh) (o(h) expl- -x))) where (hh) is an energy function that encodes the target property. Following guided flows (Zheng et al. 2023), we model conditional distribution pe(h ho)
time & via Gaussian conditional path:
(h ho) = M( hu, o²
(1)
where pe, oi are differentiable scheduler functions. The property-aware distribution at time t is:
(h) (h) (h))
with & denotes the energy at time L. To realize this distribution, we construct flow via time- dependent velocity field vy(0 (0) describing the velocity of the particle at position that matches the conditional vector field u. ho). For the Gaussian path (h ho), this field admits the closed form with score function log log ho):
u. (h ho) = jr Pe h + (jir 0z σz) ax Pt Vn (h ho),
where in and o, are both the derivative of Pt and 0, with respect to time
Energy function. We construct assay invariant energy function that emphasizes mutants with strong fitness scores while remaining robust to differences across assays. Our design combines two principles: ( mutants with large fitness values should be favored, and (ii) mutants whose fitness

---

Published as a conference paper at ICLR 2026
departs from local substitution patterns should be highlighted (Wylie & Shakhnovich 2011
Let IC
be the training set (K = [.). Mutant ";"" has measured label yi and a mutation set μ. = {(s, a.,k))
where & denotes mutated site and aiRk the amino acid at that site. We measure similarity between two mutants zt and ar by a substitution-aware edit distance:
(4) so that denb j) means zn!t and an differ by single mutational move. For each mutant at"" we define its neighborhood (i) = { E I j + i, (inb S r }. where r is a mutational radius. Within this neighborhood, we compute kernel weights Kij and normalize them:
daub d = | + [ E S, SS, ask # aj,k H/ = (k, (..k) E p
Kij
Kij- = exp ( (i,b((,))))
Ris
(5)
Ejer() Kij
Using the neighborhood weights, we then compute:
U
Σ
Kig
St
Σ JEN(1)
Rig (Us - 0,)
(6)
JEN())
where yi is the normalized fitness score with yi (yi
where jer and our are the mean and deviation. 3/ is the local baseline implied by nearby mutants, and s; is the corresponding variance. Finally, with flow time t [0. 1]. we form the energy function (h) and weight w;(t) by
combining a global magnitude term and local deviation term:
(h)) = (A Tt + (1 - y1
MA
(e(()) x exp(-3E,(h)),
where it denotes the standardized global fitness of variant N. encouraging the flow shift em- beddings toward globally high-fitness regions, (ū-- / quantifies how much mutant i deviates from local substitution patterns. The mixing coefficient balances these two components, and the final weight (t) emphasizes mutants that are simultaneously high in standardized fitness and anomalous relative to local substitution trends.
controls the sharpness of the weighting.
Lank-Consisters Conditional Flow Loss. With energy weights w;(t), we learn the flow by matching the vector field (0) to the target velocity uh ho):
LPFM (0) Ei (t) (w: 0) (h ho) [2]]
(8)
where oi is the vector field. (1) is obtained by applying a softmax to (t). Given infinite data, CPEM(0) is able to achieve the global minimum. However, in task, we observe main limitation with this objective: some assays contain only a few hundred labeled variants, which is insufficient to learn complex transport map that accurately matches property values.
Therefore, we propose a Lank-Consisteet Conditional Flow Loss (RC*) that complements the LPrM(B) (8). Specifically, for each sample x!!". the representation hy from the PLM's internal space is fed into the flow model G()) to produce a predicted target representation h1 We compute target logits with the PLM head, Qter = Ph. and form the predicted score ý by summing logit differences over the mutated sites p4;;
(log Qugt m==!!!
log Qigt (
ye
mEp.
D: can be regarded as the predicted property value for the mutant x:"". This readout is invariant to the scale of logits and focuses on the mutations that drive the property change (Gordon et al.) 2025). For a batch of B mutants, let ÿ = (wx denote the model predictions and y = {us} be the experimental measured labels. Our proposed RC² loss LRFIOw (0) minimizes differentiable surrogate of the Spearman rank correlation between ş and y:
ERHOW (0) = >rank (1 Pson (ū), R(v))),
(10)
R-(-) is a standard differentiable ranking operator used in prior work (Cuturi al. 2019) on differentiable sorting and ranking with temperature > 0, the hard rank operator I?(-) is used for ground-truit labels, and Psoft computes a differentiable correlation between (soft) predicted ranks and the hard ranks of target y. Now the total loss is defined as:
where Arank a
((0) = LrVM (PrM(0) ()) + LRFOOW (8).
(11)
This total loss combines the property-aware flow matching and the rank-wise flow objectives, bal- ancing the need to align with the target property while preserving the correct ranking of mutants.
5

---

Published as a conference paper at ICLR 2026
3.3 REPRESENTATION OPTIMIZATION VIA PROPERTY-GUIDED STEERING GATE
While the training objective Eq. 11 directly optimizes representations for a target property, the evo- lutionary information from PLMs property-agnostic. They encode many objectives once and can steer updates toward directions that are neutral or even adverse for given assay (Notin et al.) 2022). To address this issue, we propose Property-guided Steering Gate (PSG) that concentrates learning on positions carrying signals for the target property and suppresses directions driven by unrelated evolutionary biases.
Let S and S be the top and bottom {-quantiles of the training set by the measured property (we use { 0.3 unless stated), For mutant x,,", we define the wild-type-conditionn conditioond token delta by:
(e) h (x."") h (0) (x=*). Ah?
(12)
where (0) RNxd denotes the token representation from &-layer of PLM. Empirically, we use the final PLM layer (f = L) from the PLM's internal space before its prediction head. This delta ties the signal at each position the effect of mutating that site and reduces wild-type-speciff offsets. We then compute average per-position token representations across sequences in each set and construct a steering matrix V, that points along the mean difference between high- and low-property samples:
Ah!O,
!h
(13)
S+1
S
eES1
(
(2)
(14)
(4) captures the direction that separates high- and low-property mutants in the PLM representation space. Prior to training, we compute V once and cache it. For each mutant a:aat during training, we score position n by a cosine similarity to V():
(Ah (E)
VI(O))
Winn
(15)
IAh
+
where Ahl is the n-th row of Ah, ( tive win indicates the token's 's representation aligns with the high-property direction, while negative Winn indicates opposition. We convert these scores gate vector:
0 is small constant for numerical stability. Large posi-
g: 7 m(...).
(16)
with scale 7 > 0 and sigmoid (-). This gate serves a condition to our flow model concentrate learning on positions that carry signal for the target property and reduces the influence of positions that reflect unrelated evolutionary signals.
3.4 RANKFLOW ARCHITECTURE
The implementation of RankFlow comprises two core components: (1) Multi-modal fusion encoder, which combines structure and sequence representations for a protein sequence; and (2) Conditional flow head: which predicts the velocity field at each time step.
Multi-modal Fusion Encoder. We leverage existing state-of-the-art architectures to encode the structure and sequence modalities separately. For structure encoding, rely on the pretrained structure encoder ESM-IF (Hsu et al. (022).. which can effectively capture the geometric context of the wild-type protein. For sequence encoding, we utilize pretrained PLM ESM-2 (Lin et al. 2023) to extract evolutionary information from wild-type sequences. The outputs from both encoders are projected using two Multi-Layer Perceptrons (MIPs) and then fused via a self-attention block to create unified representation F E RR xd that captures both sequence and structural features.
Conditional Flow Head. The conditional flow head is designed to predict the velocity field v.(.) at each time step t. To refine PLM representations, applicd only at mutated sites with one embedding per position. For a mutation set μ {n, 1M construct c(p) E Lxd, with entries c... (1)) Opes (m) + (pm())))) ], where &pes E R2 Lxd
add learnable embeddings to the flow head,
mm=

---

Published as a conference paper at ICLR 2026
and Paa 20x are trainable modules. (m) will be concatenated ho that the flow learn mutation-specifii adjustments. The flow head takes as imput the current state h2 the mutation set μ4, and the condition C, which includes the fused representation F and the steering gate g. The conditional flow head is now parameterized as o, (, C: 0), where a light stack of U-Net blocks with time emhedding and layer normalization effectively model the dynamics of the flow.
Putting all the components together, RankFlow is trained as a conditional flow-matching model. For each mutant =!= we first obtain its representation h from the frozen PLM, together with the fused sequence-structtue representation F and steering gate obtained from Eq. [6 these are collected into the condition C.. We then sample a time he = =phh" + TrE with fixed scheduler (Mg, ox) and (r, hh) closed form according to Eq. 3 where we omit the index for brevity. The flow head takes (h,, C) as input and predicts velocity (h, 0). The parameters are optimized by minimizing the energy-weighted flow-matche matching loss [PPM Eq. SStogether with the rank-consistent Appendix(A.ssuumaaiize
(4 1), construct noisy state (N, I). and compute the target velocity
loss in Eq. 10 Algorithm in Appendix A. summarizes the overall training procedure. At inference
time, we fix C for a given assay and mutant, integrate the learned vector field from & = I to using a fixed-step Heun solver with N steps (Appendix A.2), and map the final representation to scalar fitness score.
4 EXPERIMENT
4.1 IMPLEMENTATION DETAILS
We first fix the architecture and optimiser, then select method-specific hyperparameters on few representative assays. For the loss compo- sition, we start from LPFM and add LRFIOW only if it improves validation Spearman (which it consistently does), so we recommend using the combined loss. For the energy design of X. perform a coarse sweep 0,0.25,0.5,1 and choose the value the validation optimum (typically 0.5). For scheduling, we compare lin- ear vs. cosine time schedules and pick the bet- ter one (cosine in our case). Once chosen on this validation subset, we reuse the same con- figuration for all ProteinGym assays without per-assay tuning; full training and evaluation configurations are given in Appendix [B.6] We use Spearman's 's rank correlation coefficient be- tween predicted and experimentally measured fitness values as the evaluation metric.
Active site (Residue 58)
Solvent-ex Solvent-exposed areas (Residues 117-120, 121-122)
Predicted top
Experimental
variants
fitness rank
D118N
3
FIISY
7
E122Y
E122W
F115Q
Figure 3: AICDA HUMAN using Mol* Viewer (Sehnal) et al. 2021). The top three predicted mutants by RankFlow are solvent-exy solvent-exposed positions.
Structural visualization of
4.2 PERFORMANCE COMPARISON
Benchmarks. ProteinGym (otin et al. 2023a) serves as a benchmark for mutation effect pre- diction. contains 217 substitution DMS assays and over 2.4 million mutated sequences spanning more than 200 protein families. Additional details are provided in Appendix B. (6-lcc. and Fluo. are widely used benchmarks for protein fitness prediction. Both comprise deep mutational scanning measurements covering thousands of single- and multi-mutants. P-lact. includes 4, 4,158 training, 520 validation, and 520 test samples; Fluo. contains 21,446 training. 5,362 validation, and 27,217 test samples. GBI from FLIP (Dallago et al.) 2021) is derived from deep mutational scanning of the IgG-binding protein G domain B which assayed nearly all single to quadruple mutants. In the 2-vs-rest split, double mutants used for training, while single, triple, and quadruple mutants held out for validation and testing.
Table I reports the performance of RankFlow and other baselines four benchmarks. For the ProteinGym benchmark, we follow DePLM (Wang et al., (024) in excluding datasets whose wild-type proteins exceed 1,024 residues, yielding 201 DMS datasets. Among these methods,
7

---

Published as a conference paper at ICLR 2026
Table 1: Spearman performance of supervised methods on 3-lact., GBI, Fluo., and ProteinGym
under the Random scheme. Results of CNN, ResNet, LSTM, and Transformer are from Wang et al. 20044; OHE, ESM-MSA, Tranception, and ProteinNPT are from Notin al. (2023b). best, second best.
ProteinGym
Model
G-lact.
GB1
Fluo.
Stability
Fitness
Expression
Binding
Activity
CNN
0.788
0.588
0.627
0.599
0.573
0.781
0.502
0.682
ResNet
0.734
0.489
0,521
0.525
0.481
0.152
0.133
0.636
LSTM
0.745
0.413
0.477
0,496
0.408
0.139
-0.002
0.494
Transformer
0,560
0.149
0.156
0.172
0.155
0.261
0.271
0.643
OHE
0.718
0.545
0.573
0.562
0.555
0.823
0.533
0.637
ESM-Iv
0,880
0.566
0.642
0.596
0.572
0.536
0.394
0.438
ESM-2
0,882
0,573
0.645
0.587
0.576
ESM-MSA
0,885
0.568
0.632
0.565
0,600
ProtSSN
0.877
0.692
0.718
0.757
0.678
SaProt
0.882
0.686
0.716
0.749
0.677
Tranception
0.871
0.632
0.704
0.671
0.623
ProteinNPT
0,904
0.668
0.736
0.706
0.680
DePLM (ESMIv)
0.887
0.704
0.738
0.773
0.688
0.900
0.676
0.662
DePLM (ESM2)
0.897
0.707
0.742
0.764
0.693
0.904
0.665
0.662
RankFlow (Ours)
0.911
0.742
0.765
0.781
0.722
0.912
0.689
0.687
ProtSSN Tan et al. 2025), SaProt Su et al. (2024), DePLM (Wang et al. 2024). and our model are sequence+structure approaches that use both amino acid sequences 3D structures, whereas all other methods operate on sequence information only. All models are trained either from scratch or via susk-specific fine-tuning. As shown Table I our RankFlow achieves the saate-of-the-eee per- formance on all benchmarks consistently. Both ES M-MSA (Rao et al. 2021) and Tranception (Notin) et al. 2022) outperform ESM-1v (Meier al., 202T) and ESM-2 (Lin et al., 2023) by leveraging evolutionary information from MSAs. However, RankFlow still outperforms them significantly, demonstrating that reorganizing information in pretrained PLMs can surpass even MSA-based ap- proaches. Compared to DePLM, which also uses generative model, RankFlow directly learns a conditional flow that transports property-agnostic PLM embeddings to property-aligned distri- bution. By modeling mutation sets rather than independent sites, better captures higher-order epistatic interactions that DePLM can only approximate, and consistently surpasses DePLM across benchmarks. We provide a per-assay comparison of Spearman scores in Fig. g We further analyzed which variants are prioritized by our method and found that, as shown in Fig. 3 high-fitness pre- dictions are enriched at solvent-exposed positions and away from the active site, which consistent with established biological observations (Notin et al. 2022).
We further evaluate RankFlow on ProteinGym under the Contiguous and Modulo schemes, with results summarized in Table Kermut (Groth et al.] (024) achieves the best performance on the Contiguous split, where single long region of positions is entirely withheld during training. con- trast, RankFlow achieves the highest Spearman scores both the Random and Modulo splits. The Modulo split suppresses periodic subsets of positions while still preserving distributed positional coverage; RankFlow's conditional-flov formulation appears to leverage this dispersed contextual information more effectively than kernel-based regression. Aggregated over all three evaluation modes, RankFlow yields the highest average performance. This suggests that RankFlow exhibits more stable performance across different forms distribution shift.
4.3 UNCERTAINTY ESTIMATION
We adopt the hybrid uncertainty estimation used in ProteinNPT, combining MC-Dropout and batch- resampling to obtain per-mutation epistemic uncertainty scores. This approach is model-agnostic and does not require any changes to our model architecture. As reported in Table the Spearman correlation between ble to or in some schemes slightly higher than the values reported for Stable CNN, Stable Bayesian Ridge, and Stable Kermut (Ronen et al.) 2025). These results indicate that approach not only achieves strong predictive accuracy but also delivers robust uncertainty quantification on par with stabilized baselines.
uncertainty scores and the absolute prediction errors (Puncertinty) is compara-
8

---

Published as a conference paper at ICLR 2026
Table 2: Spearman performance on the ProteinGym benchmark under Random, Modulo, and Con- tiguous evaluation schemes. Results for Kermut are taken from the original paper (Groth et al. 2024.. while all other baseline results are sourced from the ProteinGym website.
Model OHE
Random
Modulo
Contiguous
Avg.
Std. err.
0.582
0.022
0.059
0.221
0.014
ESM-1v OHE
0.565
0.394
0.396
0.452
0.015
DeepSequence + OHE
0.523
0.393
0,395
0.437
0.017
MSAT OHE
0.540
0,409
0,407
0.452
0.014
Tranception + OHE
0.550
0,413
0,415
0.459
0.012
TranceptEVE OHE
0.552
0,435
0,435
0.474
0.012
ESM-1v Emb.
0.614
0,514
0.479
0.535
0.014
MSAT Emb.
0.670
0.562
0.513
0.581
0.013
Tranception Emb.
0.681
0.525
0,439
0.548
0.008
ProteinNPT
0.741
0.588
0,529
0.619
0.009
Kermut
0.744
0.631
0.591
0.655
0.000
RankFlow (Ours)
0.786
0.635
0.589
0.669
0.018
Table 3: Benchmark results ProteinGym under the Contiguous and Modulo evaluation schemes. Puncertainty the Spearman correlation between the uncertainty scores and the absolute errors of the predictions. All baseline model values taken from Ronen et al. (2025)
Model
Contignous
Spearman
Spearman
Spearman
CNN
0.344=0.013
0.019=00.006
0.344=0.013
0.019=00006
0.365++0.02
0.0133:0:006
Stable CNN
0.402
0.12910..00
0.492100.010
0.129 0.008
0.509100.010
0.133-10.008
Bayesian Ridge
0.422=00004
0.008+0.000
0,422:0.014
0.693++00113
0.005++0.004
Stable Bayesian Ridge
0.597
0.153
0.59710.001
0.008
0.0.51
10.097
Kermut
0.6016 = 0.012
0.110=20.088
0.606:0.002
0.110-20008
0.75810.012
0.10610.0088
Stable Kermut
0.667 +: 0.010
0.164100.10
0.687 = 0.010
0.164 = 0.010
0.785 +: 0.010
0.184 + 0.008
RankFlow (Ours)
0.58910.009
0.135 1 0.009
0.635 0.010
0.132 0.008
0.786 1: 0.012
0.165 1 0.006
4.4 GENERALIZATION
To demonstrate the generalization of RankFlow, we further test it on the ProteinGym bench- mark (Notin et al., 2023a). Following the experimental settings in DePLM (Wang al. 2024). for each test dataset, we randomly select 40 additional datasets from the same category for train- ing, while ensuring that sequence similarity between training and test datasets remains below 50% to avoid data leakage. Table $ summarizes the performance of RankFlow and the baselines on ProteinGym. RankFlow consistently outperforms all methods across categories, indicating strong generalization. Despite using PLMs, it has far fewer trainable parameters than SaProt (37.1M vs. 650M), enabling training within about hour on single A 100 GPU, whereas some large-scale baselines require days. This efficiency and data efficiency make RankFlow well suited for protein engineering applications.
4.5 ABLATION STUDY
We conduct an ablation study analyze the contributions of RankFlow's components, comparing the full model with four variants: (1) energy-guided conditional flow only (RankFlow (Cpem only)); (2) RC² loss only (RankFlow (Crelow only)): (3) RankFlow w/o structure info; and (4) RankFlow w/o steering gate g,, We evaluate all variants on ProteinGym, and report the results by (a) function type and (b) mutation depth. As shown in Fig. 4 our energy-guided objective LPFM yields the largest gains across assays and mutation depths, underscoring the importance of property- awwre flow matching for learning an accurate fitness landscape. The RC² loss is especially helpful at higher mutation depths, where the combinatorial explosion of sequences makes reliable supervision scarce. Removing the sopperty-guided steering gate (RankFlow w/o steering gate gi) improves over the -onl variant, highlighting the value of focusing learning positions most relevant to the target property. We also observe a moderate performance drop when structure information and the multi-modal fusion encoder removed (RankFlow w/o Structure info), particularly on label-sparse

---

Published as a conference paper at ICLR 2026
Table Generalization ability evaluation on ProteinGym. Results for zero-shot and fine-tuned (FT)
baselines are taken from Wang et al. (2024).
Input Modalities
Trainahle parameters
Spearman Expression
Model
Seq.
Struct.
Stability
Fitness
Binding
Activily
Zero- shot methods
ESMIv
0.437
0.395
0,427
0.287
0.415
ESM2
0.523
0.396
0,439
0.356
0.433
ProtSSN
0.560
0.408
0,435
0.362
0.458
TranceptEVEL
0,500
0.477
0,457
0.360
0,487
ESM-IF
0.624
0.346
0,436
0.380
0,412
ProteinMPNN
0.564
0,166
0,209
0.159
0,203
Supervised methods
CNN
2.7M
0.141
0.053
0.043
0.056
0.095
ESMIv (FT)
650M
0.497
0,318
0.301
0,216
0.385
ESM2 (FT)
650M
0.454
0.359
0.338
0.276
0,391
ProtSSN (FT)
148M
0.689
0.478
0.421
0.488
SaProt (FT)
650M
0.703
0.442
0.496
0,391
0.495
DePLM (ESMIv)
42.2M
0.763
0.467
0,506
0,409
0.499
DePLM (ESM2)
42.2M
0.773
0.480
0.510
0.441
0.518
RankFlow (Ours)
37.1M
0.797
0,515
0.534
0.457
0,554
ProteinGym assays (notably in the Fitness and Activity categories), yet the sequence-only variant still surpasses strong sequence-only baselines such as ESM2(FT) and SaProt(FT), indicating that RankFlow remains effective without structures. More details are in Appendix B.6
RankFlow KankFlow Calte only
Rankflow onl) Raidflww laato otly
RankFlow
anke)
RaskFlow (Full)
Placa
Expensiie:
Bicking
Antecty
(6) Function Type
(b) Mutotion Depth
Figure 4: Results of ablation study for analyzing contributions different components. Breakdown performance assays is grouped by (a) function type and (b) mutation depth.
CONCLUSION
We present RankFlow, conditional flow-matching framework for protein fitness prediction that learns a property-aware landscape on top pretrained PLMs. By introducing an energy function tied assay-specift fitness and a rank- objective, it shapes the flow that mutants ordered coherently by their functional properties. In addition, the property-guided steering gate focuses learning relevant positions, improving performance diverse protein engineering tasks. Across Protein and additional protein engineering benchmarks, RankFlow consistently matches or surpasses wate-of-the-ee supervised methods under the same training protocols, using far fewer trainable parameters than full PLM fine-tuning and offering a robust, transferable approach to protein fitness prediction.
Reproducibility Statement. All the datasets used in this work are publicly available and are de- scribed in the main content and Appendix A.1 Detailed descriptions of the model architecture, training, and inference procedures are provided in Section [.4 and Appendix A.T A. and we explain and justify all hyperparameter choices in Appendix A.2 to ensure reproducibility.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi and George M Church. Unified rational protein engineering with sequence-based deep representation learning. Nature methods, 16(12):1155-322, 2019.
Jacob Beck, Shikha Surana, Manus McAuliffe, Oliver Bent, Thomas D Barrett, Juan Jose Garau- Luis, and Paul Duckworth. Metalic: Meta-learning in-context with protein language models. In The Thirteenth International Conference on Learning Representations, 2025.
Surojit Biswas, Grigory Khimulya, Ethan C Alley, Kevin M Esvelt, and George M Church. Low-n protein engineering with data-efficient deep learning. Nature methods, 18(4):389-3)) -968899999,,19 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018.
Yinghui Chen, Yunxin Xu, Di Liu, Yaoguang Xing, and Haipeng Gong. An end-to-end framework for the prediction of protein structure and fitness from single sequence. Nature Communications, 15(1):7400, 2024.
Marco Cuturi, Olivier Teboul, and Jean-Philippe Vert. Differentiable ranking and sorting using optimal transport. In Advances in Neural Information Processing Systems, volume 32, 2019.
Christian Dallago, Jody Mou, Kadina E Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. FLIP: Benchmark tasks in fitness landscape inference for proteins. In Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel, et al. Robust deep learning- based protein sequence design using ProteinMPNN, Science, 378(6615):99-66 2022.
Jonathan Frazer, Pascal Notin, Mafalda Dias, Aidan Gomez, Joseph K Min, Kelly Brock, Yarin Gal, and Debora S Marks. Disease variant prediction with deep generative models of evolutionary data. Nature, 599(7883):91-95 2021.
Sam Gelman, Sarah A Fahlberg, Pete Heinzelman, Philip A Romero, and Anthony Gitter. Neural networks to learn protein sequence- function relationships from deep mutational scanning data. Proceedings of the National Academy of Sciences, 118(49):211448881 18, 2021.
sequence-function
Cade W Gordon, Amy X Lu, and Pieter Abbeel. Protein language model fitness is a matter of preference. In The Thirteenth International Conference Learning Representations, 2025.
Peter Mørch Groth, Mads Kerm, Lars Olsen, Jesper Salomon, and Wouter Boomsma. Kermut: Com-
posite kernel regression for protein variant effects. In Advances in Neural Information Processing Systems, volume 37, pp. 295 19514-29565. 2024.
Michael Heinzinger, Ahmed Elnaggar, Yu Wang, Christian Dallago, Dmitrii Nechaev, Florian Matthes, and Burkhard Rost. Modeling aspects of the language of life through transfer-learning protein sequences. BMC bioinformatics, 20(1):723, 2019.
Thomas A Hopf, John B Ingraham, Frank Poelwijk, Charlotta PI Schärfe, Michael Springer, Chris Sander, and Debora S Marks. Mutation effects predicted from sequence co-variation. Nature biotechnology, 3582):228-135. 2017.
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie. Tom Sercu, Adam Lerer, and Alexan- der Rives. Learning inverse folding from millions of predicted structures. In International Con- ference on Machine Learning, pp.
PMLR, 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. nature, 596(7873):533-58, 2021.
Elodie Laine, Yasaman Karami, and Alessandra Carbone. GEMME: simple and fast global epistatic model predicting mutational effects. Molecular biology and evolution, 36(1 36(11) 2619, 2019.
11

---

Published as a conference paper at ICLR 2026
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637).1123-11300 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Kaiser Loell and Vikas Nanda. Marginal protein stability drives subcellular proteome isoelectric point. Proceedings of the National Academy of Sciences, 115(4)): 2018.
Joshua Meier. Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. In Advances in Neural Information Processing Systems, volume 34, pp. 29287-29303, 2021.
Erik Nijkamp, Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. ProGen2: exploring the boundaries of protein language models. Cell systems, 14(11): (4(11))98-9978 2023.
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena-Hurtadee Aidan N Gomez, Debora Marks, and Yarin Gal. Tranception: protein fitness prediction with autoregressive transform- ers and inference-time retrieval. In International Conference Machine Learning, pp. 16990- 17017, 2022.
Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood Van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, et al. ProteinGym: sarge-scale bench- marks for protein fitness prediction and design. In Advances in Neural Information Processing Systems, volume 36, pp. 64331-64379, 2023a.
Pascal Notin, Ruben Weitzman, Debora Marks, and Yarin Gal. ProteinNPT: erty prediction and design with nonpparametric transformers. In Advances in Neural Information Processing Systems, volume 36, pp.
Improving protein prop-
2023b.
Pascal Notin, Nathan Rollins, Yarin Gal, Chris Sander, and Debora Marks. Machine learning for functional protein design. Nature biotechnology, 42(2):216 2024.
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA transformer. In International Conference on Machine Learning, pp- 8844-8856. PMLR, 2021.
Adam J Riesselman, John B Ingraham, and Debora S Marks. Deep generative models of genetic variation capture the effects mutations. Nature methods, 15(10):816-88) 822, 2018.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15):20012239111 2021.
Philip A Romero and Frances H Arnold. Exploring protein fitesss landscapes by directed evolution. Nature reviews Molecular cell biology, 10(12):866-876, 2009.
Omer Ronen, Alex Y Zhao, Ron Boger, Chengzhong Ye, and Bin Yu. Stabilizing protein fitness predictors via the pes framework. In The Exploration in Al Today Workshop ICML, 2025.
David Schnal, Sebastian Bittrich, Mandar Deshpande, Radka Svobodová, Karel Berka, Václav Bazgier, Sameer Velankar, Stephen K Burley, Jaroslav Koča, and Alexander S Rose. Mol* viewer: modern web app for 3D visualization and analysis of large biomolecular structures. Nucleic acids research, 49 49(W1):W43-w4539 2021.
Jung-Fun Shin, Adam J Riesselman, Aaron W Kollasch, Conor McMahon, Flana Simon, Chris Sander, Aashish Manglik, Andrew Kruse, and Debora Marks. Protein design and variant prediction using autoregressive generative models. Nature communications, 12(1):2403, 2021.
Jin Su, Chenchen Han, Yuyang Zhou, Junjie Shan, Xibin Zhou, and Fajie Yuan. SaProt: Protein language modeling with structure-award vocabulary. In The Tweifth International Conference on Learning Representations, 2024.
12

---

Published as a conference paper at ICLR 2026
Yang Tan, Bingxin Zhou, Lirong Zheng, Guisheng Fan, and Liang Hong. Semantical and geo- metrical protein encoding toward enhanced bioactivity and thermostability. Elife, 13:RP98033, 2025.
Timothy Truong Jr and Tristan Bepler. PoET: A generative model of protein families as sequences- of-sequences. In Advances in Neural Information Processing Systems, volume 36, pp. 77379- 77415, 2023.
Zeyuan Wang. Keyan Ding, Ming Qin, Xiaotong Li, Xiang Zhuang, Yu Zhao, Jianhua Yao, Qiang Zhang, and Huajun Chen. DePLM: Denoising protein language models for property optimization. In Advances Neural Information Prucessing Systems, volume 37, pp. 75302-75328, 2024.
Scott Wylie and Eugene Shakhnovich. A biophysical protein folding model accounts for most mutational fitness effects in viruses. Proceedings of the National Academy of Sciences, 108(24): 9916-9921, 2011.
Minghao Xu, Zuobai Zhang, Jiarui Lu, Zhaocheng Zhu, Yangtian Zhang, Ma Chang, Runcheng Liu, and Jian Tang. PEER: a comprehensive and multi-task benchmark for protein sequence understanding. In Advances in Neural Information Processing Systems, volume 35, pp. 35156-
35173, 2022.
Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. Protein Engineering, Design and Selection, 36:gzad015, 2023.
Shiyuan Zhang, Weitong Zhang, and Quanquan Gu. Energy Energy-weighted flow matching for offline reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025.
Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky TQ Chen. Guided flows for generative modeling and decision making, arXiv preprint arXiv:2311. 133433,, 2023.
13

---

Published as a conference paper at ICLR 2026
Algorithm Training of RankFlow
Input:
Pretrained language model PLM, pretrained structure encoder SE, training data D ((a= xx",, }) batch size B, learning rate y, weight X, max epochs E, rank loss weight Amnkk
time sampler (. (100
Initialize RankFlow parameters 0
h"t )LM(xx*),
s SE(x"t)
for epoch = to E do for each batch { (ax:ttt x"t y)) in do
h = (LM((
P Encoder( wh; (ho 8");;
Sample 4 ~ (0, 1);
Compute embeddings
mutated positions c;(p;:;;
hm con can ch, (at(),
Compute perturbed representation: h, Me, he + dee where ~ N/ 1):
Compute target velocity ue (hs, hmmr)) using Eq. 3
Compute steering gate g using Eq. 16
V FlowHead hi,, ti, [F;; g:]);
Compute property-aware flow matching loss CrrM using Eq. g
Compute wise loss LRrlOW using Eq. I0
CreM + LRNOO
0 0 VoL
end for
end for Return: Trained RankFlow model parameters 0
A METHOD DETAILS
A.1 DETAILS OF TRAINING
Algorithm 1 outlines the training procedure for RankFlow. The model parameters O are initialized, and the training proceeds for specified number of epochs. In each epoch, the training data is pro- cessed in batches. For each assay, the wild-type sequence is encoded using the pretrained language model (PLM) and structure encoder (SE). Because each assay has a single wild-type sequence, its representation is computed once and reused across all batches.
A.2 IMPLEMENTATION DETAILS
For the main results in Table I we use pretrained ESM2-650M as the sequence encoder and ESM-IF as the structure encoder, and structures are predicted by AlphaFold2 (Jumper al. 202T). For Table which includes assays longer than 1,024 residues, we instead use pretrained ESMC-600M- 2024-12 as the sequence encoder, while keeping ESM-IF as the structure encoder. Training uses grouped batches per assay and per protein. We use AdamW with learning rate 10- weight decay 10-2, batch size 256, and early stopping on validation set. Given the variability in the datasets, the number of training epochs is set to maximum of 50. Training typically takes less than 1 hour on single 80 GB A100 GPU. During inference, we use a fixed-step second-order explicit solver (Heun) over / E [1, using the number of steps N 20. From the endpoint, we read scalar score by averaging mutant-vs-wild type logit changes over edited sites as in Eq. g This readout invariant to the scale of logits and focuses on the mutations that drive the property change. For the mutational radius r, we use r=2 by default, and for extremely small assays, neighborhoods are sparse, r will be increased to 3. During experiments, we fix À 0.5 after experimental validation. We use {=0.3 by default to form S and S- If either split has fewer than 20 samples, we switch to top-k selection with PSG by setting all gates to 1. The hlper-parameee size. The temperature т in Eq. 10 is set to 0.1. The weight Ayank in Eq. 10 is set to 0.5. 3 in Eq. is chosen from {1.0, 2.0, 5.0 based validation performance. For each assay, compute Spearman's between predicted scores and ground-truth labels on the test set, and no additional calibration post-processing used.
min (20, [0.3K]] per side. For very small assays (K K80,, we disable
7 is chosen from {0.5, 1.0] based on dataset
14

---

Published as a conference paper at ICLR 2026
Table 5: Comparison of different energy function designs on ProteinGym.
Energy Function
Stability
Fitness
Expression
Binding
Activity
E,(h)
y/
0,902
0.738
0.750
0.768
0.701
((h))
= R(yi)/n
0.901
0.712
0.725
0.742
0.695
(h)) (Ni + (1 A)
R1) yor
(Ours)
0.911
0.742
0.765
0.781
0.722
Hyperparameter selection. We follow the commonly used training setup of flow-matching mod- els and perform coarse sweep on a held-out validation split over the three specific knobs that appear in Table [8] the objective, the energy energy-guidance weight, and the schedule. For the objective, the ablation shows that removing RC or replacing it with another commonly used loss (e.g., kinetic reg.) degrades performance. For the energy-guidant
weight X, the value 0.5 yields the best average Spearman correlation, whereas both smaller and larger values lead to a noticeable drop in perfor- mance, so we fix = 0.5 for all experiments and recommend selecting À on a small validation set from simple grid such as 0.25, 0.5, 1.0 in new applications. For the schedule, cosine scheduling performs slightly better than the linear variant, so we use as the default and recommend it in the text.
The energy-guuee -eridanc weight is robust over a wide range, with À = 0.5 performing best on average. use it in all reported results.
Cosine scheduling consistently outperforms linear scheduling, Overall, RankFlow is robust to reasonable variations of these hyperparameters rather than relying on extensive benchmark-speciffic tuning.
(i) the objective, comparing variants without RC², with RC only, and the full objective with ki- netic regularization; (ii) the energy-guidane weight {0, 0.25, 0.5, 1}; and (iii) the schedule for (He,or)) comparing linear and cosine scheduling. The ablation indicates that the full objective with RC2 and kinetic regularization consistently outperforms the other variants, so we adopt this as the default. For the guidance weight, performance is relatively stable over the range À [0.25, 1), with = 0.5 achieving the best average Spearman correlation across the four ProteinGym properties; we therefore fix À 0.5 for all experiments and recommend À E [0.25, 0.5] as practical range. Cosine scheduling provides a small but consistent improvement over the linear schedule, so we use the cosine schedule in all reported results.
Overall, practitioners can reproduce our configuration by directly adopting these default settings, and the ablations in Table [ show that RankFlow is robust to reasonable variations of these hyper- parameters rather than relying on extensive benchmark-speae sppecifcc tuning.
A.3 DESIGN DETAILS
Design of Energy Function. We analyze the design of the energy function (h) = (A (Agi+
A) - in Eq. 7 by comparing it with two alternatives: (1) a normalized score value (h) = - without neighbor comparison, and (2) rank-based score E (h) = - R(y:)/n, where R(y;) is the rank of y among all n samples. The results are summarized Table The first alternative directly uses the normalized property value as the energy. which may be sensitive to outliers and does not consider the relative ranking among variants. The second alternative uses the rank as the energy, which captures relative ordering but loses fine-grat ine-grained information about the magnitude of differ- ences between variants. Our proposed energy function combines both absolute and relative informa- tion, leading to improved performance across all categories. This demonstrates that incorporating both the magnitude of the predicted property and its deviation from the local average provides a more informative signal for guiding the flow model in learning the fitness landscape.
Vs:
RankFlow Architecture Details. Our flow matching is based on one-dimensional U-Net back- bone (UNetID) tailored for sequential data. The encoder consists four residual blocks with inter- leaved linear attention and downsampling operations, reducing temporal resolution while enriching representations. A bottleneck stage applies two residual blocks with a full attention layer in between, capturing long-range dependencies. The decoder mirrors the encoder with four residual blocks and upsampling layers, combined with skip connections from the encoder. Temporal conditioning is incorporated through sinusoidal time embeddings processed by MLP, while additional MI.P
15

---

Published as a conference paper at ICLR 2026
Table Model architecture details.
Component
Details
Fusion Encoder
2 MLPs with dim= 1280, I-layer Self- Attentionn,on output dim=
U-NetID with 4 ResBlocks (down) 2 ResBlocks (middle)
Flow Head
+ 4 ResBlocks (up), hidden dim=128, output dim=1280
handles auxiliary fitness embeddings. Finally, residual block and 1x ExConvll ID projection layer yield the model's output. In total, UNetiD contains approximately 9.6M trainable parameters. The fusion encoder consists of two MLP layers with ReLU activations, followed by self-attention layer to integrate sequence and structure features. The detailed architecture summarized in Table
B EXPERIMENTAL DETAILS
B. DETAILS OF BENCHMARK
ProteinGym is an extensive set of DMS assays, containing 217 datasets. Each dataset is categorized into one of five groups based on the measured property: Stability, Fitmess, Expression, Binding, and Activity. For fair comparison, we implemented the Random cross- cross-validation method as described in DePLM. Specifically, for each dataset, each mutation the dataset randomly assigned to one of five folds. The model is trained on four folds and evaluated on the remaining fold, and this process is repeated five times to ensure each fold serves as the test set once, The final performance is reported as the average across all five folds. For results reported in Table 2 which include all 217 assays, follow ProteinGym's 's official evaluation protocols for the three splitting schemes: Random, Modulo, and Contiguous. In order to deal with
B.2 EVALUATION METRIC
Following ProteinGym Notin et al. 2023a). we quantify concordance with the experimental land- scape using Spearman's 's p. For each protein assay, let the predicted scores be ÿ (M))== and the experimental measurements be (w:) Spearman's coefficient is the Pearson correlation between the ranked vectors:
cov( (((). R(y))
p
r only)
where R(-)) is the vector of ranks (averaging ties). We compute p per protein assay and report aggregate statistics across assays in the results.
B.3 MULTIPLE MUTANTS PROPERTY PREDICTION
We further evaluate RankFlow 15 protein assays with multiple mutants in ProteinGym. Figure shows the Spearman's rank correlation between model predictions and experimental measurements for these assays. RankFlow consistently outperforms other baselines, demonstrating its effective- ness in capturing the complex interactions among multiple mutations. This highlights RankFlow's potential for guiding protein engineering efforts that involve combinatorial mutagenesis.
B.4 CROSS- VALIDATION RESULTS
The error bars for the cross-validation setting are reported in Table 7 RankFlow surpasses Protein- NPT and DePLM across all categories while using fewer trainable parameters. Relative to DePLM, RankFlow improves average performance by 1.8% (Stability). 3.8% (Fitness), 2.2% (Expression), 1.7% (Binding), and 3.0% (Activity). is also lighter (37. IM trainable parameters RankFlow vs. 42.2M in DePLM) and faster to train (50 training epochs for RankFlow 100 for DePLM). These results underscore RankFlow's effectiveness and parameter-efficien design.
In addition, analyze the correlation of predictions from different methods. As shown in Figure 6. the predictions from RankFlow exhibit less correlation with those from DePLM and ESM-2 on the Fluorescence benchmark, indicating that RankFlow might capture complementary information.
16

---

Published as a conference paper at ICLR 2026
Models
ProteinNPT
RankHoow
QBWV12 Somermeye_
PABP YEAST Melamed 2013
QSWTC7 BRRICC_CNDD__PCCC Somermeyer 2022
D7PM05 CYPGR Somermeyer 2022
A4_ HUMAN LLRiANSSSumaaaaae 2022
GR82 HwL.HHMMAANFFFA Faure 2021
2021
Assay
DLG4 HLL.HUMAN_FAAA Faure, 2021
GFP_AEQVI_Sarre Sarkisaan 2016
#PGI_STRSG Dison_ 2014
YAP1 HUMAN Araya_2012
YEAST_ HIST_YEAST_
HI57
Pokusaeva_ 2019
F7YBW8 MESOW Aakre_2015
GCN4 YEAST_ 22aller22199
SPG1_STRSC
0.5
0.8
1.0
Spearman
Figure 5: Spearman's 's rank correlation between model predictions and experimental measurements, for assays in ProteinGym with multiple mutants.
Table 7:
Model performance ProteinGym. We report standata deviation performance
under the Random scheme.
Model
Stability
Fitness
Expression
Binding
Activity
ProteinNPT DePLM
0.904
0.668+0.035
0.736+40.023
0.706
0.680+10.026
0.89710.013
0.742:10.027 0.707=0.027
0.764=1.004
0.693:10.024
RankFlow (Ours)
0.911
0.742 0.022
0.765110001
0.781:0.033
0.722.1.019
Additionally, the correlation coefficients on GB1 are significantly lower, suggesting that the models may be learning different aspects of the fitness landscape for this particular protein. We also found that ESM-2 and DePLM exhibit a stronger correlation between their predictions, which is consistent with their similar methodological basis. This suggests that RankFlow's 's unique architecture and training objectives enable it to learn distinct aspects of the protein fitness landscape, which may not be fully captured by other models. The moderate correlation also highlights the potential for ensemble approaches that combine predictions from multiple models to further enhance performance protein fitness prediction tasks.
B.5 RESULT ANALYSIS OF 8-LACT.
To understand how RankFlow achieves accurate predictions hyperactive mutants, extracted mutant embeddings of 3-lact. from RankFlow, and visualized the latent space with t-SNE. As shown in Figure 7 mutants with high experimental fitness values (top quantile) are closer in the latent space,
17

---

Published as a conference paper at ICLR 2026
NotMet
0.65
Tes Mit
0.14
04E
0.83
0.81
0.34
0.26
DePu4
0.56
0.63
0.50
0.72
D6PUM
0.25
0.44
0.42
0.56
RaniFoow
0.57
0.48
0.56
0.45
0.48
FarkF
0.17
0.12
0.08
0.25
0.16
ON6
Bastex
Dans
Bonkfiow
( Flaorescence
(b) GBI
Figure 6: The pair-wise Spearman's correlation coefficient of the predicted fitness values from dif- ferent methods.
indicating that RankFlow effectively captures the underlying distribution of the fitness landscape. In contrast, mutants with low fitness values (bottom quantile) are more dispersed, suggesting that Rank- Flow can distinguish between high- and low-fitness variants. The t-SNE visualization also reveals that RankFlow organizes the mutant embeddings which likely contributes to its strong performance in predicting protein fitness. This clustering of high-fitness mutants indicates that RankFlow has learned meaningful representations that correlate with experimental outcomes, enabling it to generalize well to unseen variants.
way that reflects their functional properties,
Fitness score
D-SNE2
20%-0%
40%
60% 80%
80% 90%
Top
Score
D-SNEL
Figure 7: 1-SNE visualization of 3-lact. mutant embeddings from RankFlow. Points are colored by quantile bins of experimentally measured fitness; the bottom panel shows the ground-truth fitness trend.
B.6 SUPPLEMENTARY ABLATION STUDY
Table highlights the contribution each design choice RankFlow. Removing the rank- consistent loss (CRFFww) leads to noticeable degradation across all assay categories, especially in Fitness and Activity, confirming that ranking consistency is crucial under data-scarce conditions. Using RC² alone performs poorly, indicating that flow matching remains necessary anchor the transport map. Adding a simple kinetic regularizer recovers only part of the lost performance, un- derscoring the importance of our rank-aware objective. For the energy design, both extremes (A=1
18

---

Published as a conference paper at ICLR 2026
Table 8: Ablations of RankFlow on ProteinGym.
Model Variant
Stability
Fitness
Expression
Binding 0.781
Activity 0.722
Full RankFlow
0.911
0.742
0.765
Objective Loss
w/o CRFLow
0,885
0.702
0.736
0.744
0.680
LRF only
0,881
0.624
0.676
0.622
0.613
LPEM + kinetic reg.
0.892
0.714
0.745
0.755
0.695
CPFM + pairwise hinge
0.889
0.709
0.740
0.751
0.690
Energy guidance
À
0,902
0.738
0.750
0.768
0.701
À 0
0,896
0.718
0.744
0.756
0.699
==0225
0,900
0.719
0.748
0.759
0.702
00050
0.911
0.742
0.765
0.781
0.722
Scheduling
Linear (Me, or)
0.908
0.736
0.762
0.772
0.718
Cosine (Mi, o+)
0.911
0.742
0.765
0.781
0.722
Table 9: Comparison of computational costs of ProteinNPT and DePLM.
Method
Training (MACs)
Inference (MACs)
Parameter (#)
ProteinNPT
58.630 38.63M
11724.82 x 5001
11.7K
100M 119M
9.16K
180.56 + 77.55
DePLM
347.16 180.56 + 77.55 + 89.05
792M 42.2M
+99.05 100
897.1K
179.25 x 5001
RankFlow (Ours)
19.34
792M +37.1M
+77.55 + 19.34 x 30
using only global magnitude or λ=0 using only local deviation) underperform, while intermediate values achieve stronger results; the best trade-off is obtained at A=0.5, showing that global fitness and local substitution deviation are complementary signals. Finally, cosine scheduling consistently outperforms a linear schedule, particularly in Fitness and Binding, demonstrating that smoother variance growth improves the stability of the learned flow. Overall, the ablations validate that each proposed component: LRFIOW. balanced energy guidance, and cosine scheduling contributes to Rank- Flow's state-of-the-ar performance.
B.7 COMPUTATIONAL COSTS
Following Wang et al. 2024), we use A4GRB6 PSEAI as an example to compare the computa- tional costs of different models. The sequence length is 267, and the dataset size is 5001. We report the number of trainable parameters, peak GPU memory usage, and training time in Table 9 For both DePLM and RankFlow, the sequence encoder requires 180.56G MACs and 179.25G MACs, respectively, as RankFlow only requires internal representation from the sequence encoder the final predicted logits. The structure encoder requires 77. 77.55G MACs for both models. Before training a model on given assay, both ProteinNPT and RankFlow compute and persist to disk the sequence embeddings for all mutated proteins in that assay. During training. we load from disk the embeddings corresponding to each batch. RankFlow trains with higher MACs than DePLM (897. vs. 9.1K as learns a property-aligned velocity field for all mutants. However, for this assay, RankFlow lowers total inference cost by 94.4% compared to DePLM, making it more prac- tical for large-scale screening. RankFlow also uses fewer trainable parameters (37.1M vs. 42.2M) by employing lightweight flow architecture instead of large transformer decoder, reducing ory footprint and speeding up training. Overall, RankFlow achieves a favorable balance of training efficiency and inference speed compared to ProteinNPT and DePLM.
19

---

Published as a conference paper at ICLR 2026
MAX
BAN
Spearman
Spearman
Figure 8: Results under the random cross-validation scheme on ProteinGym. We report the DMS-
level performance of DePLM, ProteinNPT, and our RankFlow.
20

---

Published as a conference paper at ICLR 2026
C BACKGROUND
C.I CONDITIONAL FLOW MATCHING
Continuous Normalizing Flows (CNFs) (Chen et al. 2018 consider the dynamics of probability density via a path p [0, 1] x Rdd - Rgo that transports between data distribution Po and an initial
distribution p1- The flow is governed by a vector field v |0, 1] x Rd + IR through the continuity equation:
Oppr + V (px(h)o(h,1) O, Pt-o = Po,
p1-
(17)
The objective is to learn the time-dependent vector field e describing the velocity of particle at position h. In practice, given the conditional distribution pr ho), which is usually modeled as a Gaussian path, i.e., Pr (h hv) = (wsho, go where P4. is Gaussian scheduler. The vector field v is learned by minimizing the flow matching objective (Lipman et al. 2022):
(0) = (0, hiop ( [ll,p t) u, (h, (h) where ho follows the data distribution po.
(18)
C.2 PROTEIN LANGUAGE MODELS FOR MUTATIONAL. EFFECT PREDICTION
Protein Language Models (PLMs) are deep learning models trained on large corpora of protein sequences using sulf-supervised learning objectives, such as masked language modeling (MLM) or autoregressive modeling. These models learn to capture the statistical patterns and dependencies in protein sequences, enabling them to generate meaningful representations that can be used for various downstream tasks, including mutational effect prediction. Protein language models trained using the masked language modeling objective, such as ESM-2 (Lin et al. 2023), can be used to predict the effects of mutations on protein function. These models trained to predict masked amino acids in sequence based on the surrounding context, allowing them to learn rich representations that capture the relationships between different amino acids and their roles in protein structure and function. Specifically, given a protein sequence a (I,, ................................ .NN.. where N is the length of the sequence and v, represents the amino acid at position i, a PI.M trained with MI.M learns to predict the masked amino acid I' given the context of the other amino acids in the sequence. The model outputs a probability distribution over the 20 standard amino acids for each position in the sequence. For mutant sequence z! and its corresponding wild type sequence ws, can compute the log-odds score for each mutation at position m as follows:
(log P( mh x (m log P( xy (=0 [=))))),, mEj
(19)
where 1 the set of mutated positions, and wit denotes the mutant sequence with the amino acid at position m masked. This log-odds score reflects the model's confidence in the mutant amino acid relative to the wild type amino acid, providing an estimate of the mutation's 's effect on protein func- tion (Meier et al., 2021). By leveraging the learned representations from PLMs, we can effectively predict the impact of mutations protein properties, aiding in tasks such protein engineering and understanding disease-related mutations.
D LIMITATION AND FUTURE WORK
While RankFlow demonstrates strong performance in protein fitness prediction, there are limitations and areas for future work. First, the reliance on pretrained language models means that the quality of the embeddings is contingent on the diversity and representativeness of the training data used for these models. Future work could explore fine-tuning or adapting PLMs specifically for protein en- gineering tasks to enhance their relevance. Second, while RankFlow effectively integrates structural information through ESM-IF, it may not fully capture dynamic aspects of variant structures that are crucial for function. Incorporating more sophisticated structural representations or dynamic simu- lations could further improve performance. Finally, while RankFlow is efficient compared to some baselines, further optimizations in model architecture and training procedures could make it more accessible for large-scale applications in protein engineering. For example, exploring lightweight architectures or distillation techniques for extracting protein representations could reduce computa- tional costs without sacrificing accuracy.
21

---

Published as a conference paper at ICLR 2026
E THE USE OF LARGE LANGUAGE MODELS (LLMS)
Our paper uses LLMs to polish the writing. We did not use LLMs to generate any scientific content, including the main ideas, algorithms, or experimental results. All scientific content was developed and verified by the authors.
22


Reviews:
[{"review_id": "OgRSIOR3Au", "reviewer": "Reviewer_sQ9J", "summary": "This paper presents RankFlow, a property-aware conditional flow model for protein optimization. It refines PLM embeddings into property-aligned representations. It captures complex multi-mutation interactions through an energy-guided flow and a Property-Guided Steering Gate (PSG). The Rank-Consistent Conditional Flow Loss (RC² Loss) enforces correct mutant ranking and ensures out-of-distribution generalization. The model achieves superior performance on various benchmarks.", "questions": "1. The ProteinGym benchmark has a standard evaluation pipeline for unsupervised and supervised fitness prediction tasks. What is the difference between the evaluation procedure this paper follows and the standard ProteinGym evaluation pipeline?\n\n2. Are the unsupervised methods and supervised methods evaluated on the same set of protein sequences in Tables 1 and 2?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Unlike existing methods that simply add the mutational effect of an individual site, RankFlow uses energy-guided conditional transport to capture complex multi-mutation interactions. In addition, optimizing the relative ordering of mutants provides better generalization to unknown protein families.", "weaknesses": "The model achieves better performance than previous models. However, there could be more detailed experimental setup descriptions, especially as the setup is not identical to the standard ProteinGym benchmark. The author could also discuss more on other supervised models, such as ProteinNPT and Metalic.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TuKEYIYDrw", "reviewer": "Reviewer_PgXb", "summary": "The goal of this work is to develop a **fitness predictor** to assess the effects of mutations on a given protein sequence (wild type). This is a central problem in computational biology, with applications in both **genomics** and **computational protein design**.  \n\nThe authors propose a new computational procedure called **RankFlow**, which employs **flow matching** to learn a distribution in the latent space of a protein language model (PLM). This distribution is *tilted* toward mutants that are both high in fitness and anomalous relative to local substitution trends. Samples from this learned distribution are then passed through a PLM head and subsequently used for fitness prediction.  \n\nResults on relevant benchmarks show that RankFlow performs well compared to multiple existing methods.", "questions": "1. Can diffusion be used to provide **uncertainty quantification**? It would be interesting to compare this form of UQ with other approaches [3, 4].  \n2. The authors claim that previous methods do not capture **epistasis**, but PLM embeddings can, in principle, model epistasis. Can the authors provide a comparative analysis of their model on multiple mutants versus other methods?  \n3. The notation is confusing in places. For example:\n   - In **Figure 2**, the unified representation is denoted as *F*, but *F* never appears in the flow-matching equations—can the authors clarify this?\n   - In **Equation (3)**, what are \\(\\dot{y}\\) and \\(\\dot{\\sigma}\\)?  \n4. Can the authors clearly explain how the **datasets are split**? Do they use the random split for ProteinGym?  \n5. Can the authors clearly define how **fitness predictions** are generated according to their method in the main text?  \n6. Is the method sensitive to any **flow-matching hyperparameters**? How does randomness in the flow-matching process affect predictions?\n\n---\n\n[1] Jagota, M., Ye, C., Albors, C., Rastogi, R., Koehl, A., Ioannidis, N., & Song, Y. S. (2023). *Cross-protein transfer learning substantially improves disease variant prediction.* **Genome Biology**, 24(1), 182.  \n\n[2] Groth, P. M., Kerrn, M., Olsen, L., Salomon, J., & Boomsma, W. (2024). *Kermut: Composite kernel regression for protein variant effects.* **Advances in Neural Information Processing Systems**, 37, 29514–29565.  \n\n[3] Ronen, O., Zhao, A. Y., Boger, R., Ye, C., & Yu, B. (2025). *Stabilizing protein fitness predictors via the PCS framework.* In *The Exploration in AI Today Workshop at ICML 2025.*  \n\n[4] Greenman, K. P., Amini, A. P., & Yang, K. K. (2025). *Benchmarking uncertainty quantification for protein engineering.* **PLOS Computational Biology**, 21(1): e1012639.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Good performance on multiple relevant benchmarks.  \n- The method can be pre-trained on multiple DMS assays, thereby benefiting from previously observed cross-protein transfer effects [1]. I believe this is the main strength of the proposed approach.", "weaknesses": "- The evaluation is missing two important baselines: [2] and [3].  \n- The evaluation does not include **out-of-distribution (OOD)** splits from **ProteinGym**.  \n- The authors do not propose any means for **uncertainty quantification (UQ)**, which limits the applicability of this predictor for **Bayesian Optimization**.  \n- Additional **ablations** on architectural choices would strengthen the work. For example:\n  - What happens if the flow-matching module is replaced with a simple MLP layer?  \n  - What happens if the ranking loss is replaced with other commonly used losses?  \n  - How does this method compare with an approach similar to **CPT-1** [1]?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4XSIBIccUd", "reviewer": "Reviewer_3v7D", "summary": "The authors propose a flow-matching-based framework for protein fitness prediction. For that, they design a new energy function that is combined with representations from pLMs. They also add a gate mechanism to leverage mutational information from a dataset to bias mutations in the framework. Results show that the proposed method achieves state-of-the-art performance for the ProteinGym benchmark and three more tasks. Ablation studies show that the proposed components, especially the proposed energy function, are crucial to achieve good performance.", "questions": "Overall, the paper contributes a new flow-based framework for protein fitness prediction that could be a contribution to epistasis modeling. My initial recommendation is borderline acceptance, but I would like to discuss with the authors the following comments.\n\nComments:\n\n1. The proposed method uses structure as input. For the datasets experimented, are all structures experimental or some are predicted? A clarification regarding which benchmark methods use sequence, structure, and sequence-structure seems needed.\n2. Many hyperparameters are included in Section 3, and additional information is given only in the Appendix. As these seem to directly influence the state-of-the-art performance of the proposed method, more discussion about how to choose these hyperparameters is needed.\n3. For the gate, how much assay data is needed for this term to be effective? Does the nature of assay data influence this term?\n4. (lines 315-316) Can the authors elaborate on what it means that the learnable embeddings are applied only at the mutated positions?\n5. Additional clarification regarding the modeling of the energy function in Eq. 7 seems needed. The ablations in the Appendix show that this term is crucial for performance.\n\nMinor Comments (that did not impact the score):\n\n1. The use of \\citet and \\citep is wrong throughout the manuscript. Almost all references are missing parentheses.\n2. The acronym of pLMs is defined twice in text.\n3. Need consistency in the use of hyphens, e.g., “deep-mutational-scanning”; and when using a comma for numbers.\n4. Figure 4 should be improved as it is very hard to read the caption in the printed version.\n5. Typo: “reproducibility”", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors propose a conditional flow framework that learns property-aligned embeddings that improve fitness prediction, especially for cases in which mutants have multiple mutations.\n2. A new energy function (Eq. 7) and loss are proposed for protein fitness prediction.\n3. The proposed RankFlow method achieves state-of-the-art results on important benchmarks such as ProteinGym.", "weaknesses": "1. Hyperparameters that should be tuned to achieve state-of-the-art performance are discussed only in the Appendix.\n2. The implementation of the flow-based method is not clear when following the information in Section 3.2.\n3. Additional analysis of the method for proteins without structure available (without the multimodal fusion encoder) seems needed.\n4. Code is not available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VmZg3gCY5Z", "reviewer": "Reviewer_XW2e", "summary": "The paper presents RankFlow, a method for improving protein fitness prediction using flow matching. The main contribution is reformulating fitness prediction as a generative task rather than a regression problem. While the proposed approach is potentially interesting, the paper requires major revisions, particularly to improve clarity in the overall motivation and experimental setup.", "questions": "1. Is performance difference between RankFlow and DePLM statistically significant?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The idea of applying flow matching to protein fitness prediction appears novel.\n- The proposed components, such as the Rank-Consistent Conditional Flow Loss and the Property-Guided Steering Gate, are original and well-motivated.", "weaknesses": "Major comments\n\n- Although the method’s technical details are described, the overall intuition and practical usage remain unclear. The motivation behind reformulating fitness prediction as a generative problem (rather than regression) is not well explained. A conceptual figure illustrating this idea would be helpful (see, for example, Fig. 2 in [1]). In addition, several typos make the method difficult to follow (see minor comments).\n- The experimental setup is insufficiently described:\n  - Evaluation metrics are not mentioned in the main text.\n  - The values reported in Table 1 and Table 2 do not match those on the official ProteinGym website [2], and no explanation is provided. For example, ProteinNPT achieves a “Stability” score of 0.904 in the paper, while [2] reports 0.776.\n- The evaluation mixes unsupervised models (which do not use any fitness data, e.g., ESM-2, SaProt) with supervised models (e.g., ProteinNPT) without justification or explanation. It is unclear what training data RankFlow uses. The parameter count comparison in the table appears inconsistent. How can RankFlow have fewer parameters than ESM-2 if it uses ESM-2 as its backbone?\n\nMinor comments\n\n- [Line 152] x^{mt} is likely a typo; it should be x^{wt}.\n- [Line 257] R_{\\tau} is mentioned but never explained.\n- [Line 273 and Eq. 13] The notation for S^+ and S^- is inconsistent.\n- [Eq. 16] g_i is introduced here but its role is not clarified in the rest of the technical explanation of the method.\n- [Line 372] The text claims: “As shown in Fig. 3, variants predicted to have high fitness are at solvent-exposed positions and are located away from the active site, which is in line with established biological knowledge (Notin et al., 2022a).” However, Fig. 3 does not demonstrate this.\n\n[1] Corso et al., 2023. DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking.https://arxiv.org/abs/2210.01776\n\n[2] https://proteingym.org/", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)