DePLM: Denoising Protein Language Models for Property Optimization

Authors: zeyuan  wang
keyan  ding
ming  qin
xiaotong  li
xiang  zhuang
yu  zhao
jianhua  yao
qiang  zhang
huajun  chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 14:45:34

Abstract: Protein optimization is a fundamental biological task aimed at enhancing theperformance of proteins by modifying their sequences. Computational methodsprimarily rely on evolutionary information (EI) encoded by protein languagemodels (PLMs) to predict fitness landscape for optimization. However, thesemethods suffer from a few limitations. (1) Evolutionary processes involve thesimultaneous consideration of multiple functional properties, often overshadowingthe specific property of interest. (2) Measurements of these properties tend to betailored to experimental conditions, leading to reduced generalizability of trainedmodels to novel proteins. To address these limitations, we introduce DenoisingProtein Language Models (DePLM), a novel approach that refines the evolutionaryinformation embodied in PLMs for improved protein optimization. Specifically, weconceptualize EI as comprising both property-relevant and irrelevant information,with the latter acting as “noise” for the optimization task at hand. Our approachinvolves denoising this EI in PLMs through a diffusion process conducted in therank space of property values, thereby enhancing model generalization and ensuringdataset-agnostic learning. Extensive experimental results have demonstrated thatDePLM not only surpasses the state-of-the-art in mutation effect prediction butalso exhibits strong generalization capabilities for novel proteins.

Main paper: 
 introduction: 
 Proteins play vital roles in numerous biological processes, shaping their structure and function over billions of years of evolution. This evolutionary diversity presents significant opportunities for advancing fields such as drug discovery and materials science [24, 60] . However, the inherent properties of existing proteins, such as thermostability, often fall short of practical requirements in various scenarios. Consequently, researchers have dedicated themselves to optimizing proteins to enhance their properties of interest. Protein optimization is the task that involves modifying protein sequences and efficiently identifying well-performing proteins.Traditional deep mutational scans (DMS) and directed evolution (DE) rely on expensive wet-lab experiments [15, 6, 58] . Recently, computational approaches that accurately model the relationship between proteins and their property fitness, often termed a "fitness landscape" [51] , have become crucial for efficient protein optimization. These approaches, powered by machine learning, enable rapid evaluation of mutation effects and are pivotal for efficient protein optimization.One widely explored avenue involves leveraging the "Evolutionary Information" (EI), which can be instantiated by the likelihood of an amino acid appearing at a certain position of protein sequences, to infer the mutational effects [41, 49, 16] . This stems from the observation that as organisms evolve through natural selection, mutations that improve functional properties become more prevalent. Therefore, the likelihood of a mutation occurring is directly linked to its impact on biological function [20] . To compute the informative likelihoods of mutating one amino acid to another, the predominant methods involve protein language models (PLMs) trained on millions of protein sequences, which capture the EI in a self-supervised manner [50, 43, 39] . Thanks to their strong generalization capabilities, PLMs have been utilized to guide the artificial selection of beneficial mutations with notable efficacy [24] .With the advent of high-throughput experimentation, large-scale and diverse annotated datasets for DMS are becoming increasingly accessible [44] . Consequently, researchers have extended the use of self-supervised EI into supervised prediction settings [18, 17, 23, 21, 11, 65] . Specifically, they usually fine-tune PLMs on experimentally annotated datasets, with the objective of minimizing the disparities between predicted and experimental fitness values [8, 2, 46] . However, two critical aspects are often overlooked. Firstly, it fails to account for the removal of irrelevant EI. Evolution optimizes multiple properties simultaneously to meet survival needs, often overshadowing the optimization target of interest [35] . Therefore, conventional fine-tuning methods using the whole evolutionary information are suboptimal. Secondly, the prevalent learning objective incorporates dataset-specific information that is often overfitted to the training data at hand, hindering the model's ability to generalize toward new proteins. This limitation is significant since DMS experimental techniques often encounter constraints regarding their applicability across a wide range of proteins [58] .In this work, we introduce a novel Denoising Protein Language Model (DePLM) tailored for protein fitness prediction. The central concept revolves around perceiving the EI captured by PLMs as a blend of property-relevant and irrelevant information, with the latter akin to "noise" for the targeted property, necessitating its elimination. To achieve this, drawing inspiration from denoising diffusion models that refine noisy inputs to generate desired outputs [55, 10, 29] , we devise a rank-based forward process to extend the diffusion model for denoising EI, as illustrated in Figure 1 . Specifically, we refine the likelihood of mutations provided by PLMs. To parameterize this framework, we initially extract protein representations considering both primary and tertiary structures. Subsequently, we utilize this representation to guide the denoising process. In pursuit of dataset-agnostic learning and robust model generalization, we conduct the diffusion process in the rank space of property values and replace the conventional objective of minimizing numerical errors with maximizing rank correlation. Extensive experiments have demonstrated that the introduced rank-based denoising process significantly improves the protein fitness prediction performance, and simultaneously maintains strong generalization ability for novel proteins. Our contributions can be summarized as follows:• We introduce DePLM, a novel approach for refining evolutionary information captured by protein language models to predict mutation effects, effectively filtering out irrelevant information to improve mutation effect predictions.• We design a rank-based forward process within the denoising diffusion framework, extending the diffusion process to the rank space of mutation likelihoods. • We shift the learning objective from minimizing numerical errors to maximizing rank correlation, fostering dataset-agnostic learning and ensuring robust generalization. • DePLM significantly outperforms state-of-the-art models for mutational effect prediction and demonstrates strong potential for optimizing unseen proteins.2 Background and Related Works 
task formulation and evaluation: 
 Protein optimization seeks to elucidate the impacts of sequence mutations on protein properties. Formally, given a widetype protein sequencex wt = [x wt 1 , • • • , x wt n , • • • , x wt N ] with N amino acids, a mutation µ = {µ n : x wt n → x mt n , n ∈ [[1, N ]]} refers to the substitution of the amino acid x wt n at certain positions n with another amino acid x mt n . Note that a mutation can affect multiple positions simultaneously. The task of protein optimization is mathematically formulated as learning a function F θ , parameterized by θ: F θ (x wt , µ) = y, where y denotes the impact of the mutation µ on the wildtype sequence x wt , i.e., the property value of the mutated protein sequence. With numerous (µ, y) pairs, one can approximate the fitness landscape of the protein x wt for the target property.To evaluate the consistency between predicted and ground-truth fitness landscape, one often uses Spearman's rank correlation coefficient, which prioritizes relative rankings over absolute values [44] . Specifically, given the set of predictions Y = {y} and ground-truth Y * = {y * }, both sets are first converted to their respective ranks R(Y ) and R(Y * ), then the coefficient ρ is calculated asEQUATIONwhere cov(•, •) is the covariance of the ranked variables, σ is the standard deviations of these ranks.Self-supervised modeling of sequence mutation effects. Evolutionary information provides valuable insights into how sequence mutations affect biological functions [14] . A straightforward approach to capture EI is through Multiple Sequence Alignments (MSAs). For example, the SIFT model [41] predicted the mutation effect by performing position-specific statistical analysis of aligned sequences. Riesselman et al. [49] and Frazer et al. [16] employed variational autoencoder trained on protein-specific MSAs to detect patterns of interaction among positions, achieving higher prediction performance. However, the MSA approach is limited by its protein-specific nature, rendering it less effective for proteins, such as orphan proteins, that lack sufficient homologous sequences.Instead, many studies [2, 40, 50, 39, 4, 37] explored PLMs trained on evolutionary-scale data to capture EI, which can generalize beyond specific proteins. Further, several studies aim to merge PLMs with MSAs to harness the benefits of both approaches. Rao et al. [48] proposed MSA Transformer to apply language modeling to aligned sequences, and PoET [59] modeled the distribution over protein families rather than sequences. Tan et al. [56] suggested that mutation effects are related to their structural context, and Notin et al. [44] showed that the likelihood derived from structure-based design models [28, 9] can be complementary to those generated by PLMs. Despite the utility of EI, the information captured by PLMs is often entangled with multiple protein properties, leading to suboptimal optimization of the target property when used directly.Supervised modeling of fitness landscape. The utility of EI extends into supervised prediction scenarios [18, 17, 23, 21] . Dallago et al. [8] demonstrated that straightforward fine-tuning of the representations from PLMs holds the potential to predict fitness. Hie et al. [25] used PLMs to predict the evolutionary trajectory of protein families. Yang et al. [65] advocated for using machine learning tools for small datasets and neural networks for larger datasets. In response to sparse labels, Elnaggar et al. [12] proposed a lightweight parametric model called ConvBERT to avoid overfitting. ProteinNPT [46] leveraged a non-parametric Transformer that combines masked language modeling and fitness prediction tasks, demonstrating excellence in low-resource scenarios. However, these methods often lose the generalization capabilities of PLMs after dataset-specific fine-tuning. In this study, we show that combining EI with experimental data in a way that minimizes reliance on dataset-specific information is crucial for enhancing performance while maintaining generalization.Denoising diffusion models. Several diffusion models have been applied to protein research. Diffusion processes in discrete token spaces have been promising for designing protein sequences [3, 1, 22] . Jing et al. [32] focused on torsion for generating conformations, while Corso et al. [7] studied the product space R × SO(3) × SO(2) m for protein docking. Additionally, backbone generation requires defining SE(3) diffusion process [66, 61, 30] . In contrast to these works, we operate the diffusion process in the rank space of mutation likelihoods to predict the fitness landscape.The ELBO can be calculated as follows:E[log p θ (r 0 |x wt )] = E log E q(r 1:T -1 |R(Π ⋆ ),R( Π)) p θ (r 0:T |x wt , Π) × p( Π) q(r 1:T -1 |R(Π ⋆ ), R( Π)) ≥ E q log p θ (r 0:T |x wt , Π T ) × p(Π T ) q(r 1:T -1 |R(Π 0 ), R(Π T )) = E q log p(Π T ) - T t=1 log p θ (R(Π t-1 )|x wt , Π t ) q(r t |R(Π t-1 ), R(Π T )) = E q log p(Π T ) -log p θ (R(Π 0 )|x wt , Π 1 ) q(r 1 |R(Π 0 ), R(Π T )) - T t=2 log p θ (R(Π t-1 )|x wt , Π t ) q(r t-1 |R(Π 0 ), R(Π t ) + log q(r t-1 |R(Π 0 ), R(Π T )) q(r t |R(Π 0 ), R(Π T )) = E q log p(R(Π T )) q(r T |R(Π 0 ), R(Π T )) -log p θ (R(Π 0 )|x wt , Π 1 ) - T t=2 log p θ (R(Π t-1 )|x wt , Π t ) q(r t-1 |R(Π 0 ), R(Π t )) = -E q [(1 -ρ(q(r T |R(Π 0 ), R(Π T ))||p(R(Π T ))) -log p θ (R(Π 0 )|x wt , Π 1 )) + T t=2 1 -ρ(q(r t-1 |R(Π 0 ), R(Π t ))||p θ (R(Π t-1 )|x wt , Π t ))The first term, 1 -ρ(q(r T |R(Π 0 ), R(Π T ))||p(R(Π T )), is a constant and can therefore be omitted from the objective function. Following Xu et al. [64] , we combine the second term, log p θ (R(Π 0 )|x wt , Π 1 )), with the final term. This leads us to the expression for ELBO L ELBO = T t=1 (1 -ρ(q(r t-1 |R(Π 0 ), R(Π t ))||p θ (R(Π t-1 )|x wt , Π t ))) as in Eq. 4.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In Section 5, we discuss the limitations of this method and directions for further improvements.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 3.1 Overall Framework DePLM, as depicted in Figure 2 , is designed to filter out irrelevant information from the noisy evolutionary likelihoods produced by PLMs. Given a wildtype protein x wt , the evolutionary likelihood produced by a PLM can be denoted asΠ = [ π1 , • • • , πn , • • • , πN ],where πn ∈ R 20 denotes the probability of 20 amino acids occurring at the position n. This likelihood can be decomposed into the target property likelihood Π ⋆ and additive noise Π ϵ introduced by irrelevant properties, such that Π = Π ⋆ + Π ϵ . The DePLM takes the noisy likelihood Π as input and refines it to isolate the desired likelihood Π ⋆ via a rank-based denoising diffusion process. A comprehensive explanation of the symbols and operations utilized is available in Appendix A.Denoising diffusion models consist of two main processes: a forward corruption process and a learned reverse denoising process. In the forward corruption process, small amounts of noise are progressively added to the ground truth. The reverse denoising process then learns to recover the ground truth by gradually eliminating the accumulated noise. When applying these models to denoise the mutation likelihood Π in protein optimization, however, there are two significant challenges. First, the relationship between actual property values and experimental measurements frequently exhibits nonlinearity, stemming from the diversity of experimental approaches [44] . Consequently, reliance on minimizing discrepancies between predicted and observed values for denoising purposes risks overfitting to the specific dataset utilized, thereby diminishing the model's generalization capabilities. Second, unlike those conventional denoising diffusion models, our final noisy state Π is deterministic, requiring the accumulated noise to converge [36] . rank-based controlled forward corruption and learned denoising backward processes. In the corruption step, we use sorting algorithms to generate trajectories, shifting from the rank of property-specific likelihood to that of the evolutionary likelihood. DePLM is trained to model the backward process. Right: We illustrate the alteration of the Spearman coefficient during the transformation from evolutionary likelihood to property-specific likelihood via the sorting algorithm.To navigate these challenges, we propose a rank-based denoising diffusion process that focuses on maximizing rank correlation (see Figure 3 ). Let r 0 = R(Π ⋆ ) ∈ N 20×N + be the rank of ground-truth likelihood of the target property, r T = R( Π) ∈ N 20×N + be the rank of noising evolutionary likelihood. Intermediate rank variables, r t ∈ N+ 20×N for t = 1, . . . , T -1, are generated along this sequence.Forward process. The forward process operates in the rank space, gradually transitioning from the initial rank r 0 through increasingly chaotic states r t , culminating in the final state r T after T steps. Unlike traditional models where corruption can be random, the corruption process here must carefully manage the progression because the initial state r 0 and the final state r T are pre-defined.To manage this transition, we leverage the QuickSort algorithm [26] to create a feasible space for sampling intermediate rank variables r t . At each time step t, we apply the sort algorithm to generate sorting trajectories from the variable r t-1 towards the end variable r T . The rank variable r t chosen along these trajectories ensures a progressive decrease in rank correlation from the initial to the intermediate states, and converse increase towards the final state. This approach ensures a controlled and meaningful progression through the rank space, as detailed in Appendix B.1. The forward process can be conceptualized as a Markov chain:EQUATIONBackward process. Given that the rank variables are non-differential, the backward process operates in the likelihood space rather than the rank space. This process constructs the likelihood Π 0 from the noisy evolutionary likelihood Π T = Π, ensuring that R(Π 0 ) equals R(Π ⋆ ). Importantly, while Π 0 and Π ⋆ are aligned in rank, they do not need to be identical in value. The key focus is on the rank variables, as the relative rankings are more critical for optimizing protein performance than the exact numerical values of the likelihoods. To facilitate this process, we employ the protein x wt as a guiding signal, and model it as a conditional Markov chain with learned transitions:EQUATIONwhere p θ (Π t-1 |x wt , Π t ) is the learnable transition kernel, θ denotes its parameters, and Π t is the intermediate likelihood variables during the backward process, with r t = R(Π t ).The forward process introduces noise in rank space, while the reverse process denoises in likelihood space. To effectively link these two processes, we leverage Spearman's rank correlation ρ defined in Eq. ( 1). Building upon this, we modify the variational Evidence Lower BOund (ELBO) as our learning objective function:EQUATIONDetailed derivations are provided in the Appendix B.2. Overall, given a protein x wt , its property likelihood Π ⋆ is generated by first drawing evolutionary likelihood Π from PLMs, and then iteratively refined through p θ (Π t-1 |x wt , Π t ). We predict y according to Eq. 10 in Appendix C.1.To parameterize the transition kernel p θ (Π t-1 |x wt , Π t ), as shown in the Denoising Module of Figure 2 , DePLM first learns a protein representation (Feature Encoder of Figure 2 ), and then utilizes it to guide the process of discerning and eliminating noise (Denoising Block of Figure 2 ).Feature encoder. We encode features from both sequences and structures because they complementarily describe the impact of mutations. Sequence information is derived from representations h e and attention weights M ∈ R N ×N ×m generated by PLMs, where m is the head number. Meanwhile, structural information is obtained through trained structure encoders like ESM-IF [28] , which process the protein backbone to produce structural representations h s . We merge the two sets of representation using Multi-Layer Perceptrons (MLPs), yielding a unified representation h in R N ×d , where d represents the hidden dimension.The feature encoder leverages multiple stacked layers to update h and M . Let's denote the output feature vectors of representations and attention weights in layer l as h l and M l , respectively. Initially, h 0 equals h and M 0 equals M . We update the representations using the standard attention mechanism, incorporating a bias derived from the attention weights:EQUATIONwhere q l , k l , and v l are the linear projection of h l and B l M is the linear projection of M l . Then, the attention weights are updated by communicating with the sequence representation through both an outer product ⊗ and an outer difference ⊖:EQUATIONEQUATIONwhere || means the concatenation operation. The final layer output h L is used in the denoising block.Denosing block. Given the intermediate likelihood variable at the t step Π t , we employ a denoising block to implement Eq. ( 3). The central premise is that h L should encapsulate only property-specific protein information. By subtracting h L from the hidden representation of the noisy Π t , we isolate the hidden representation of the noise. This hidden representation is then transformed into the likelihood space using a PLM head (PLMHead),EQUATIONHere MLP(Π t ) denotes a trainable MLP that maps Π t to the hidden representation space, while PLMHead maps hidden presentations back to the likelihood space. It is important to use a frozen PLMHead to ensure consistency between noise and noisy likelihoods in a unified space. This approach effectively denoises Π t to obtain Π t-1 .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: There is no societal impact of the work performed. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 In this section, we extensively evaluate DePLM across various datasets and demonstrate its superior performance and robust generalization capabilities. Specifically, we aim to address the following key questions. Performance comparison (Q1): Can DePLM beat SOTA on protein fitness prediction tasks? Generalization ability (Q2): Does DePLM maintain its generalization ability post-training? Ablation study (Q3): What is the extent of improvement achievable for each component? Analysis (Q4): Does the assumption that EI contains noise hold? We begin by outlining the general experimental setups used in our evaluations. We use ESM-IF [28] as the structure encoder and the structures are predicted by AlphaFold2 [33] . DePLM comprises 42.2 million trainable parameters and involves 3 diffusion steps. We set the learning rate at 0.0001, with a weight decay of 0.005, utilizing AdamW as the optimizer. All models are trained on four Nvidia V100 32G GPUs for up to 100 epochs by default.Datasets and Baselines. We conducted a thorough study across four benchmarks, including Prote-inGym [44] , β-Lactamase (Abbr., β-lact.) and Fluorescence (Abbr., Fluo.) from PEER [63] , and GB1 (utilizing a 2-vs-rest split) from FLIP [8] , where the latter two involve multiple mutants. We compare DePLM with nine baselines, including 1) four protein sequence encoders trained from scratch (CNN [53] , ResNet [47] , LSTM [47] , and Transformer [47] ) as naive baselines, 2) five extended baselines that incorporate self-supervised models (OHE [27] , fine-tuned versions of ESM-1v [39] , ESM-MSA [48] , and Tranception [43] , as well as ProteinNPT [46] ). More details about datasets and baselines can be found in Appendix C.1 and C.2.Results. We present the evaluation results in Table 9 . DePLM achieves better performance compared to the baselines, affirming the advantage of integrating evolutionary information with experimental data for protein engineering tasks. It is worth noting that ESM-MSA and Tranception exhibit enhanced EI compared to ESM-1v due to the introduction of MSAs. By comparing their results, we demonstrate that higher-quality EI significantly improves outcomes after fine-tuning. However, even with these improvements, their performance still falls short of that achieved by DePLM. We attribute this difference to the architecture employed by our model, which enables more efficient utilization of experimental data. We also notice that DePLM yields better performance than ProteinNPT, underscoring the efficacy of the proposed denoising training process.Datasets and Baselines. Computational techniques that can generalize across different proteins are essential, given the limitations of DMS experimental methods in handling various proteins [58] . In our study, we leverage ProteinGym to test this generalization capability. Specifically, ProteinGym categorizes DMS datasets into five coarse categories based on the protein properties they measure: stability, fitness, expression, binding, and activity. Given a testing dataset, we randomly select an additional 40 datasets from the same category for training. Importantly, we ensure that the sequence similarity between the datasets used for training and testing is kept below 50% to prevent potential data leakage. In this experiment, we compare DePLM with four self-supervised baselines (ESM-1v [39] , ESM-2 [34] , and TranceptEVE [45] ), two structure-based baselines (ESM-IF [28] and ProteinMPNN [9] ), and three supervised baselines (CNN [8] , fine-tuned version of ESM-1v and ESM-2). The choice of baselines is elaborated in Appendix C.1.Results. As shown in Table 2 , one can observe that DePLM consistently outperforms all baseline models. This finding underscores the inadequacy of baselines that rely solely on unfiltered evolutionary information, which often dilutes target properties due to concurrent optimization of multiple objectives. By eliminating the influence of irrelevant factors, DePLM enhances its performance significantly. In addition, the performance of baselines trained to minimize the disparity between predicted and experimental scores (ESM1v (FT) and ESM2 (FT)) falls significantly short of that achieved by our DePLM. This observation highlights that optimizing the model in a rank space introduces less bias from specific datasets and yields superior generalization. Furthermore, we observe that protein structure information contributes to the stability and binding properties, whereas evolutionary information enhances fitness and activity attributes.The ablation study aims to validate the efficacy of the modules devised for DePLM. These modules include the use of structural information, the feature encoder, the denoising module, and the rank objective. The results in Table 8 show the absence of any of these modules leads to a decline in performance, indicating their collective significance. Additionally, we observe that the rank objective has the largest impact on performance, highlighting the importance of reducing dataset-specific information. Further ablation studies are detailed in Appendix C.4.Necessity of Filtering Property-Irrelevant Information in EI. To ascertain the importance of filtering out property-irrelevant information, we analyzed the impact of training with datasets targeting various optimization objectives. As illustrated in Figure 4 (left), we observe that using training datasets with characteristics divergent from those of the test dataset leads to diminished performance across stability, expression, and activity properties. This reduction in performance underscores the detrimental interference among different properties, emphasizing the need to mitigate these adverse effects. Interestingly, we observed that training on datasets targeting alternative properties can improve performance in the binding property. This improvement is likely due to the scarce availability of datasets specifically focusing on the binding property within ProteinGym. This observation suggests a beneficial cross-utilization of data, where leveraging information from unrelated attributes can enhance inference for properties with limited data. In Figure 4 (right), we illustrate the influence of dataset size on performance. We found that even with a minimal number of datasets (K=2), DePLM significantly boosts performance, indicating its proficiency in filtering out irrelevant information. Furthermore, as the number of training datasets increases, there is a corresponding improvement in performance, showcasing the model's capability to continuously enhance its filtering efficacy with more data.Differentiating the Noisy Evolutionary Likelihood from the Property-specific Likelihood. Denoising the evolutionary likelihood is helpful in identifying protein sequences that manifest specific properties. In Figure 5 (left), we utilize entropy to gauge the degree of conservation at each position and illustrate the likelihood attributed to hydrophobic amino acids within the structure. Residues with outward-facing side chains on alpha helices are associated with higher entropy, whereas inwardfacing positions exhibit lower entropy. Figure 5 (right) compares the evolutionary likelihood and the property-specific likelihood and visualizes the differences. We notice that binding and active share similar offsets. This suggests that the indole-3-glycerolphosphate synthase functions by attaching to other molecules, which makes sense considering its role as an enzyme [54] . We also observe that the property-specific likelihoods exhibit a more uniform entropy when compared to the noisy evolutionary likelihood. We assume this may arise due to a bias towards inward-facing positions across all properties, potentially overshadowing the evolutionary significance of outward-facing side chains. This imbalance is addressed through the denoising process described in Section 3.2. 
conclusion and limitations: 
 In this paper, we propose DePLM, a simple yet effective fine-tuning approach that leverages a feature encoder to obtain expressive protein representations and then uses them to extract property-specific likelihood from the noisy evolutionary likelihood for mutational effect prediction. Our experiments demonstrate that DePLM not only surpasses state-of-the-art baselines but also shows exceptional generalization capabilities. Additionally, our analysis confirms that utilizing sufficiently large datasets or incorporating data from other relevant properties can significantly enhance performance.Due to limited resources, our experiments are conducted using wild-type marginal probability. This approach predicts the impacts of all mutations in a single forward pass, summing the effects of individual mutations to estimate the consequences of multiple mutations simultaneously. However, this method is not ideal, as it overlooks the complex interactions between mutations. Our method can potentially achieve better performance in predicting the effects of multiple mutations, leveraging more effective prediction techniques such as masked marginal probability. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper does not pose such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: In Appendix C.3, we have explicitly mentioned and properly respected all the creators or original owners of assets used in this paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: To ease reading and facilitate understanding of our DePLM, in Table 4 , we summarize the symbols and notations employed throughout the paper. We provide the pseudo codes as follows, to help readers easily understand how to construct the space of rank variables in the forward process.Data: The ranks of likelihoods at time steps t 1 and t 2 (where t 1 < t 2 ), represented by r t1 and r t2 ; the number of sampling trajectories η. Result: The feasible space of rank variables S r t1:t2 between r t1 and r t2 . S r t1:t2 ← ∅, ξ ← ∅, i ← 0. // Variable initialization Compute sorting index I t1 so that r t1 [I t1 ] is monotonically increasing and r t1[I t1 ][I -1 t1 ] = r t1 . r t1 ← r t1 [I t1 ], r t2 ← r t2 [I t1 ]. ξ ← ξ ∪ {[0, len(r t1 ) -1]} // Set left index ϕ to 0 and right index ψ to len(r t1 ) -1 while i < η do while Stack ̸ = ∅ do τ ← ∅. for [ϕ, ψ] ∈ ξ do r t2 , φ = Sort(r t2 , ϕ, ψ) // No element in [ϕ, φ] is greater than any element in [φ, ψ]. τ ← τ ∪ {[ϕ, φ -1]} if ϕ < φ -1. τ ← τ ∪ {[φ + 1, ψ]} if ψ > φ + 1. end for ξ ← τ S r t1:t2 ← S r t1:t2 ∪ r t2 [I -1 t1 ] end while i ← i + 1 end while B.C.1 Baseline details C.1.1 Self-supervised methodsSelf-supervised methods can decode the link between protein sequences and their function by analyzing evolutionary-scale data. This is because protein properties influence the choice of sequences throughout evolution. Leveraging advancements from Natural Language Processing (NLP), researchers have leveraged the Transformer architectures combined with the Masked Language Modeling (MLM) learning objective to effectively capture the complexity of protein sequences.During training, each input sequence x is altered by replacing some amino acids with a special mask token. The network learns to predict these missing tokens from the modified sequence:EQUATIONwhere M represents a set of masked positions and x /M denotes the masked protein sequence. Specifically, MLM replaces 80% of these positions with the special "[MASK]" token, 10% with a randomly chosen alternative amino acid token, and 10% remain as the original input tokens.With the MLM learning objective, PLMs can output the likelihood of each amino acid occurring at masked positions given the context of unmasked amino acids. This can be directly applied to the prediction of mutational effects by comparing the likelihood of the mutant sequence to the wildtype sequence [39] . Specifically, one uses the amino acid in the wildtype protein as a reference state:EQUATIONThe probability p in Eq. 10 is actually implemented as Π * with our DePLM framework. Note that although the property y (e.g., thermostability) has its own physical concepts and units, which are usually different from the likelihood on the right-hand side of Eq. 10, one usually cares the relative ranks of property values. Hence the above equivalence holds.Supervised methods use experimental data to fit neural networks for predicting protein fitness landscapes. Most existing methods treat fitness landscape prediction as a regression problem. The process begins by transforming the wildtype sequence x wt and its mutations (µ 1 , . . . , µ n ) into mutant sequences (x mt 1 , . . . , x mt n ). These sequences are then input into a protein encoder to extract features. A predictor is subsequently employed to estimate the property value y, optimized using mean squared error (MSE) loss.Performance Comparison. We adopt nine supervised baselines to compare with DePLM. Following Xu et al. [63] , we use four well-known protein sequences encoders: shallow CNN and ResNet, which focus on short-range interactions, and LSTM and Transformer, which focus on long-range interactions. Subsequently, we utilize ridge regression to forecast fitness using features extracted by the aforementioned models. We summarize each model in Table 5 Table 5 We also compare DePLM with five extended baselines that incorporate self-supervised models. Specifically, we investigate OHE [27] , which utilizes ridge regression supplemented with protein likelihood from DeepSequence [49] . We also assess fine-tuned versions of ESM-1v [39] , ESM-MSA [48] , and Tranception [43] , which predicts the fitness score by inputting their representations and likelihoods as features into an trainable MLPs. Furthermore, our approach is compared against ProteinNPT [46] , a state-of-the-art baseline equipped with a non-parametric transformer specifically designed for label-scarce settings. To ensure fair comparisons, we keep the parameters of these unsupervised models frozen.Generalization ability. We conduct a comprehensive comparison of model generalization capabilities, including three types of models. 1) Self-supervised models trained on sequence variations. Among these, ESM-1v [39] and ESM-2 [34] are leading PLMs based on MLM objective, which form the foundation of DePLM. ProtSSN [56] integrates both evolutionary and structural features, enhancing the model's ability to capture complex protein characteristics. We also assess TranceptEVE [45] , the current state-of-the-art in unsupervised ProteinGym benchmarking. 2) Inverse Folding models. Since the function on proteins is dependent on its shape, amino acid likelihoods derived based on backbone can be used to infer mutation effects. Therefore, we employ ESM-IF [28] and ProteinMPNN [9] as baseline models.3) Supervised models. We also emply three supervised baselines, including CNN, which has previously been acknowledged as the best naive baseline. Due to computational By leveraging wildtype marginal probability, our method can predict the fitness scores of all possible single mutants in a single forward pass. In contrast, the state-of-the-art model ProteinNPT requires (D/B) forward passes to predict the fitness landscape of an assay, where D is the number of data points and B is the batch size. For predicting the A4GRB6_PSEAI_Chen_2020 assay's fitness landscape, ProteinNPT requires 58.5M GMACs with 219M parameters, while our DePLM only requires 347. 16 GMACs with 834M parameters (792M non-trainable and 42.2M trainable). The detailed calculation process is described in Table 7 . Diffusion steps We investigate how the performance changes with the number of diffusion steps T . As illustrated in Figure 6 , when T = 0, substituting the proposed denoising module with MLPs consistently led to lower performance, indicating the advantage of the denoising module. Interestingly, we observe that although our model can effectively learn how to denoise through the rank-based diffusion framework, the performance does not always improve with increasing T . The few diffusion steps in DePLM can be attributed to the following reasons: Wellinformative Initialization: Standard diffusion models transform uninformative Gaussian noise into a complex target distribution, requiring numerous steps to capture the transformation accurately. In contrast, DePLM starts with an initial distribution that represents an informative protein evolutionary likelihood. This initial distribution needs only minor adjustments to reach a property-specific likelihood. Thus, DePLM requires fewer diffusion steps compared to those in standard diffusion models. Efficient Noise Sampling: In standard diffusion models, Gaussian noise is injected independently into each data. However, in DePLM, noise sampling considers the overall difference between the current and target distribution. A quick sorting algorithm is employed to generate a sampling pool from which we draw noises. This approach allows each step to transform the distribution more efficiently, thereby reducing the number of steps needed.Increasing the number of diffusion steps leads to a deterioration in model performance. This decline occurs because a higher number of diffusion steps enhances the model's fitting capability, which increases the risk of overfitting to the training data as reported in [42, 5] . To further elucidate this point, we present the performance metrics of the model on both the training and test datasets in the Table 7 . Justification: In Section 1, we point out this paper aims to solve property optimization problems, and our contribution lies in proposing a novel approach that refines EI.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The authors propose a new fine-tuning technique for protein language models, using a rank-based diffusion model. The reviewers highlighted that the paper is well written, that it contains interesting, novel ideas, and that the importance of the methodological advances is backed up by empirical results. Some concerns were raised regarding the computational efficiency of the method, the ablation studies, and the number of baselines, but the reviewers all stated that their concerns were (mostly) addressed by the rebuttal.  In conclusion, the paper presents several novel ideas, and obtains results superior to the current state-of-the art.

Reviews: review: paper_summary: paper_summary: In this work, the authors propose Denoising Protein Language Models (DePLM) to enhance protein optimization by refining evolutionary information (EI) in protein language models (PLMs). Since traditional methods struggle with considering multiple functional properties simultaneously and lack generalizability to novel proteins due to experimental condition-specific measurements. DePLM addresses these issues by denoising EI to remove irrelevant information, improving model generalization and ensuring dataset-agnostic learning. Experimental results demonstrate DePLM's superior performance in mutation effect prediction and generalization to new proteins.
strength_weakness: strength_weakness: 1. The computational cost of the proposed method has been well discussed.
questions: questions: 1. The framework relies on a sorting algorithm to define the forward diffusion process. What is the computational cost? Will it be an obstacle for applying the proposed method to problems at scale?
2. Following the previous question, how does the computational cost (flops and # parameters) compare with baseline models?
3. DePLM only uses 3 diffusion steps, which is much less than standard diffusion model. I wonder what would happen if more steps are used? Also, what is the reason that DePLM doesn't need multiple sampling steps?
4. In Table 1, DePLM (ESM2) achieves superior performance. Does the authors by any chance have the results with ESM2 alone?
5. In Table 3, it is shown that structural information doesn't make big difference in performance. What could be the cause of this?
In Figure 3c, it seems the correlation increases with forward diffusion process. Should it be in the opposite way? Maybe I didn't fully understand it, buts when noise is added, the correlation between the optimal $\Pi^*$ and current $\Pi$ should decrease.
limitations: limitations: The authors sufficiently addressed the limitations in the submission.
paper_summary: paper_summary: In this work authors tackle the problem of optimizing protein sequences towards a given property. They outline limitations of existing methods using Protein Language Models within an optimization loop to optimize property as those pLMs are not tailored towards a given property. They introduce a rank-based diffusion model to fine-tune pLMs on mutation effect prediction. This scheme can be used to optimize a wild type sequence towards a given property. They evaluate their method on four datasets and reach state-of-the-art performance. They also show that their method is able to generalize across different datasets, allowing to overcome the common problem of data scarcity in protein optimization.
strength_weakness: strength_weakness: **Choice of Architecture**
- The ablation study shows very marginal improvement of using the structural information. Since this information is typically obtained through modules more expensive than sequence modules, the benefit of this part of the model is unclear.
- The detailed architecture of DePLM is not provided and some hyperparameter choices look arbitrary. For instance, the choice of 3 diffusion steps is surprising and only justified by an ablation studies on 3 small datasets.
questions: questions: - Could you clarify and quantify the cost of training DePLM without the different modules in the ablation studies, to put in perspective with the associated performance gain ?
- Could you provide detailed choice of parameters for DePLM, notably for the different components of the denoising module ?
- The generalization experiment still shows a significant gap with training and testing on the same data source. I understand gathering new labeled data induces wet-lab costs. Would it make sense from a real-world application point of view, for a given test dataset, to combine data from other datasets and the corresponding training dataset to fine-tune DePLM ?   

Some typos:
- l.80 "widetype"
- l.200 denosing
limitations: limitations: I believe authors adequately address limitations of their work.
paper_summary: paper_summary: This paper proposes a new method, DePLM, for supervised fine-tuning of protein language models (PLMs) for fitness prediction tasks. DePLM uses a denoising framework with a rank correlation objective to iteratively denoise PLM likelihoods and only retain the component of the likelihood that corresponds to the fitness property of interest. The paper reports that DePLM improves on previous methods for fitness prediction on ProteinGym benchmarks and a few other single-protein benchmarks.
strength_weakness: strength_weakness: **Major:**
1. **More baselines and ablation studies are needed to show whether DePLM is state of the art.** DePLM incorporates structural information, but in the experiments claiming that DePLM outperforms baseline methods for protein fitness prediction (Table 1), none of the other methods are given structural information. Methods that leverage structural information, such as SaProt, should be included as other relevant baselines. Similarly, in Table 2 we see that DePLM is the only method that includes evolutionary, structural, and experimental labels, but other methods that do include evolutionary and structural information could be naturally finetuned to include experimental labels. The claim that “the architecture employed by our model … enables more efficient utilization of experimental data” does not seem sufficiently supported by the shown results.
2. **The datasets used in the experiments in the paper should be described more clearly.** In particular, the experimental results for baseline methods on ProteinGym are significantly different between this paper (Table 1) and previous papers, as well as the ProteinGym leaderboard available at https://proteingym.org/. For example, ProteinNPT reports an average Spearman correlation of 0.65 on ProteinGym when performing cross-validation on random splits of the data, but in Table 1 of this paper, ProteinNPT has Spearman correlation higher than 0.65 for every subset of ProteinGym. The paper does not explain why there is this discrepancy.
3. **Describing the method under the diffusion framework is confusing and obscures what methods should be relevant comparisons to DePLM.** The theoretical basis of diffusion models relies on specifying a forward noise process that, as the time step t grows, turns complex data distributions into a known distribution that we can sample from. Learning the reverse process allows us to sample from the complex data distribution. In this paper, there is not a stochastic forward noise process, there is no noisy distribution as t -> infinity (there is a deterministic rank ordering), and there is no data distribution at t=0 that the method is attempting to sample from (again there is a deterministic rank ordering). This method seems accurately described as a denoising model, but it is confusing to me to use the diffusion framework.
4. **There are no error bars on experimental results.** It is well known that some ProteinGym datasets are harder than others, so performance must vary significantly based on the random cross-validation split, and it is unclear how many of the results are statistically significant.
5. Overall, I think that this paper presents an interesting novel method, but needs more thorough experiments with stronger baselines to show that there is a real gain in performance from their complex architecture and modeling pipeline, compared to simpler ways of integrating sequence, structure, and experimental information together.

**Minor:**
* Line 133 “addictive” -> additive
* Line 299 “Limitation” -> Limitations
questions: questions: 1. Do you have comparisons between DePLM and structure-aware models like SaProt and ProtSSN on protein fitness prediction? (Either with finetuning, or at least a simple baseline where they are used as part of a One Hot Encoding supervised method?)
2. Why are the baseline results on ProteinGym different between previously published results and what you show in Table 1?
3. Can you add error bars or other statistical analysis of the results?
4. Why does adding more denoising steps beyond 3 produce worse results? Can you determine the right number of denoising steps without ground truth labels to do hyperparameter tuning?
5. Do you apply the QuickSort algorithm with stochasticity?
limitations: limitations: The authors discuss some limitations of their work. Here are some additional points of consideration:
1. The computational complexity of their method is not directly compared with other state of the art methods, but based on the sentence, “models are trained on four Nvidia V100 32GB GPUs for 100 epochs”, it seems possible that their method is computationally much more expensive than other baselines.
2. The paper often uses “protein optimization” and “protein fitness prediction” interchangeably, but these are importantly different tasks, and the paper uses metrics traditionally associated with fitness prediction, not optimization. Spearman correlation metrics measure fitness prediction accuracy across both low and high fitness sequences, while NDCG and other similar metrics emphasize finding the best/optimal sequences out of a set. It is not clear from the experiments in the paper that you will find more optimal protein sequences using their method; generally, such experiments involve an actual optimization procedure to generate candidate sequences, and validation that these candidates are better than the candidates produced by baseline methods.

Comments: title: Thanks!
Answer to authors' rebuttal
Thanks!
Thanks!
Response to rebuttal
Thank you
comment: Thank you for your valuable feedback! We are glad that our responses have successfully resolved your concerns.
I appreciate the authors' clarification on the ablation studies and believe these additional results will further strengthen their work. Authors also adequately addressed my other comments.
This confirms my initial assessment and I believe the paper should be accepted.
Thank you immensely for your feedback! We are gratified to know that we have successfully addressed your concerns.
Thank you very much for your feedback! We are pleased that we could address your concerns.
I thank the authors for the detailed response which have addressed most of my concerns. I still positive about this submission.
Thank you to the authors for their additional results and careful response.

I still feel that the use of the diffusion terminology could be better motivated and explained. As I stated in the original review, I think the denoising framework is appropriate. Denoising does not require diffusion though, and it felt more confusing than illustrative to describe the addition of noise as a diffusion process, when it has a fixed end point.

With that being said, based on the additional baseline results and the other reviewer's comments, I am willing to increase my score.