DePLM: Denoising Protein Language Models for Property Optimization

Many thanks for the confirmation of our methodological novelty, and the constructive and valuable comments on the experiments. We have conducted addtional experiments to make the results and conclusion stronger. Here, we provide details on the comments below.

Experiments

1. More baselines and ablation studies are needed to show whether DePLM is state of the art.
2. There are no error bars on experimental results.

Thanks for your suggestions, we have included fine-tuned versions of SaProt and ProtSSN as additional baselines in both the fitness prediction task and the generalization ability evaluation. Our revised analysis now features comprehensive statistical assessments, including mean Spearman correlation coefficients and standard deviations for each dataset, to highlight result variability. These results, detailed in the updated PDF (Tables 1 and 2), show that DePLM outperforms SaProt and ProtSSN, which incorporate evolutionary, structural, and experimental labels. Notably, DePLM achieves the smallest standard deviation on 4 out of 5 datasets, underscoring its robustness. We believe these additional results support our claim that DePLM is state-of-the-art and will incorporate them into the final draft of our manuscript.

The datasets used in the experiments in the paper should be described more clearly. Why are the baseline results on ProteinGym different between previously published results and what you show in Table 1?

The discrepancy arises from differences in dataset composition. ProteinNPT in paper was evaluated on 100 DMS assays, while the leaderboard includes 216 assays, resulting in an average Spearman correlation coefficient of 0.73. Our paper focuses on a subset of 201 DMS assays due to PLM context length limitations (Appendix C.2), explaining the higher Spearman correlation observed in Table 1.

Diffusion Process

1. Describing the method under the diffusion framework is confusing and obscures what methods should be relevant comparisons to DePLM.
2. Do you apply the QuickSort algorithm with stochasticity?

We address the concerns regarding our use of the diffusion framework with the following points:

1. Evolution optimizes multiple properties simultaneously, often obscuring the specific optimization objective of interest. Modeling the removal of irrelevant properties as a denoising process is reasonable. By framing it this way, we can leverage denoising models to effectively address the problem, aligning with the overall methodology and objectives of our research.
2. Extending the diffusion model to handle the likelihood order, deterministic at t=0 and t=\infty, represents a significant challenge and innovation in our work. We introduce a novel approach by using a sorting algorithm to identify the noise sampling space. The randomness inherent in the pivot index selection during applying the quicksort algorithm ensures that the forward process integrates the necessary stochasticity, aligning with the principles of the diffusion framework. Thus, the answer to the question about applying QuickSort with stochasticity is affirmative.

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?

This decline occurs because a higher number of diffusion steps enhances the model's fitting capability, which also increases the risk of overfitting to the training data, as previously reported in [1, 2].

Determining the optimal number of diffusion steps without ground truth labels is challenging. However, empirical evidence suggests that DePLM requires fewer denoising steps than standard diffusion models. This efficiency stems from two key advantages:

1. Well-informative Initialization: Unlike standard diffusion models that transform uninformative Gaussian noise into a complex target distribution, DePLM starts with an informative evolutionary likelihood. This requires only minor adjustments to reach a property-specific likelihood, resulting in requiring fewer diffusion steps compared to those in standard diffusion models.
2. 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 allows each step to transform the distribution more efficiently, thereby reducing the number of steps needed.

Clarifying the computational complexity of DePLM

DePLM offers significant computational efficiency. It predicts the fitness scores of all possible single mutants in one forward pass, while ProteinNPT requires (D/B) forward passes (D for the number of mutants and B for the batch size). For the A4GRB6_PSEAI_Chen_2020 assay, DePLM requires only 347.16 GMac, compared to ProteinNPT's 58M GMac: 11,724.83 GMac per mutant * 5001 mutants. These calculations, performed using the ptflops package, highlight DePLM's efficiency. The detailed calculation process is described in Table 4 of the uploaded PDF.

Clarification of "protein optimization" and "protein fitness prediction"

Thank you for highlighting the distinction between protein optimization and protein fitness prediction. We acknowledge the important differences between these two tasks, though there are strong correlations between them. In the final version of the paper, we will revise the terminology to more accurately reflect our focus, ensuring clarity and precision in our language. Additionally, we will include a discussion on the potential for future work to explore optimization procedures that generate candidate sequences, providing a pathway to validate improvements over baseline methods.

[1] Improved Denoising Diffusion Probabilistic Models. Nichol et al. ICML 2021.

[2] Extracting Training Data from Diffusion Models. Carlini et al. 2023.

We appreciate Reviewer ZGbQ's constructive comments, which have significantly improved our paper. Below, we address each comment in detail.

Marginal improvement from structural information

The marginal improvement observed using the structural information can be attributed to the dataset selection. For GB1 and Fluorescence, the extensive number of training mutants leads to only a slight enhancement from the inclusion of structural data. Furthermore, in our additional evaluations using the data-sparse ProteinGym assays, we observed a consistent improvement in the Spearman correlation coefficient by approximately 2.7%. The updated results are presented in Table 3 of the uploaded PDF file.

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.

The choice of 3 diffusion steps in DePLM is contributed to the following reasons:

1. Well-informative 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.
2. 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.

It is important to note that while increasing the number of diffusion steps can enhance the model's fitting ability, it also introduces a risk of overfitting as reported in [1, 2]. This results in an initial improvement in model performance, followed by a decline as the number of diffusion steps continues to rise. Therefore, we empirically determined that setting the number of steps to 3 provided the optimal trade-off. To further elucidate this point, we have included a table below that demonstrates the effects of varying the number of diffusion steps:

Clarify and quantify the cost of training DePLM, to put in perspective with the associated performance gain.

In our experimental setup, training is conducted solely with the Feature Encoder and the Denoising Block. Each training process involves one forward pass through the Protein Language Model (PLM), incurring a cost of 180.56 GMac, and one forward pass through the Structure Encoder, incurring a cost of 77.55 GMac. The computational overhead of the PLM is essential as it provides the evolutionary distribution as the initial state for denoising. From the analysis of the additional results, we observe that the structural encoder contributes a performance improvement of approximately 0.02.

Over 100 epochs, the computational cost for the Feature Encoder amounts to 4101 GMac (= 100 epochs * 41.01 GMACs per epoch). For the Denoising Block, the cost totals 4803 GMac (= 100 epochs * 3 steps * 16.01 GMACs per forward). These components contribute performance improvements of 0.08 and 0.34, respectively, on the ProteinGym dataset.

Overall, these designs are beneficial in enhancing performance. The additional computational cost of the different modules is minimal, making it a worthwhile trade-off for the observed performance improvement.

Could you provide detailed choice of parameters for DePLM, notably for the different components of the denoising module?

For the feature encoder in DePLM, we set the sequence state dimension to 1280 and the attention head size to 32. The pairwise-residue state dimension is 32, with a matching attention head size. We apply a dropout rate of 0.2. In the denoising block, the MLP hidden dimension for converting likelihood to representation is set to 1280, using GELU activation. These parameter choices were systematically determined through extensive experimentation and cross-validation to optimize the performance of DePLM while ensuring computational efficiency. For a comprehensive description of the implementation details, please refer to the Supplementary Material (Code: archive>src>models>DePLM_module.py).

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?

To explore this, we performed experiments using two different mixing strategies to construct the training dataset: (1) combining data from other datasets with 1/4 of the corresponding training data, and (2) using only 1/4 of the corresponding training data. The results of these experiments are summarized in the table below:

The results indicate that when the corresponding training data is insufficient, incorporating data from other datasets with similar properties significantly improves model performance. We will include these findings in the final draft to emphasize the practical benefits of data integration.

[1] Improved Denoising Diffusion Probabilistic Models. Nichol et al. ICML 2021.

[2] Extracting Training Data from Diffusion Models. Carlini et al. 2023.

We sincerely thank Reviewer 5NpU for your insightful feedback. We have addressed your concerns below and hope our responses provide clarity:

Computational Cost

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?

We utilize the QuickSort algorithm, which has a time complexity of $O(n\log n)$. Given the sparsity of labels, sorting the assay with the most labels in ProteinGym (~500k mutants) takes only 1.45 seconds on 2.8GHz Quad-Core Intel Core i7, while sorting the assay with the median number of labels (~5k mutants) takes just 0.0056 seconds.

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 4 of the uploaded PDF.

Overall, DePLM proves to be an effective and efficient method for predicting protein fitness landscapes and the inclusion of the ranking algorithm does not hinder scalability compared to baseline models.

Diffusion Step

1. 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?
2. 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^{\star}$ and current $\Pi$ should decrease.

The few diffusion steps in DePLM can be attributed to the following reasons:

1. Well-informative 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.
2. 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, as illustrated in Figure 8. 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 [1, 2]. To further elucidate this point, we present the performance metrics of the model on both the training and test datasets in the table below:

Figure 3c illustrates the relationship between the Spearman coefficient of evolution likelihood and intermediate rank variables as a function of the number of forward steps. The caption might be a bit confusing but we will correct it the final version.

Model Architecture

1. In Table 1, DePLM (ESM2) achieves superior performance. Does the authors by any chance have the results with ESM2 alone?
2. In Table 3, it is shown that structural information doesn't make big difference in performance. What could be the cause of this?

The PLM-based results reported in Table 1 are sourced from the ProteinGym Leaderboard https://proteingym.org/, which does not include ESM2-based results. Here, we report ESM2 results obtained by our own experiments and will include these in the final draft.

Regarding the second question, the minimal performance difference observed with structural information can be attributed to the dataset selection. When considering label-richGB1 and Fluorescence datasets, DePLM only shows a slight improvement by incorporating the structural data. To further investigate the role of structural information, we conducted additional evaluations using the label-sparse ProteinGym assays. The updated results are presented in Table 3 of the uploaded PDF file, which demonstrates a consistent enhancement in the Spearman correlation coefficient of approximately 2.7% when incorporating structures.

[1] Improved Denoising Diffusion Probabilistic Models. Nichol et al. ICML 2021.

[2] Extracting Training Data from Diffusion Models. Carlini et al. 2023.