Pre-Training Protein Encoder via Siamese Sequence-Structure Diffusion Trajectory Prediction

Summary

The authors propose to use the diffusion on both the protein sequence and the protein structure for pretraining. Additionally, they also take the correlation between different conformers of the same protein into consideration and maximize the mutual information between their trajectories via mutual denoising. Specifically, different conformers of the same protein are generated by perturbations on the side-chain torsional angles. Experiments on Enzyme Commission classification and four Atom3D tasks demonstrate the benefit of diffusion-based pretraining and mutual denoising.

Strengths

1. The idea of maximizing the mutual information between diffusion trajectories of different conformers is interesting and inspiring. The authors also show the objective can be transformed into a loss function for mutual denoising between two relevant trajectories. 2. The evaluation is comprehensive. The authors compare with different pretraining strategy and conduct many ablation studies. They also explore the effect of the size of the pretraining data and different diffusion models.

Weaknesses

1. Random perturbations on the side-chain torsional angles may produce conformations with high energy, which does not reflect the actual "physics underlying the conformational change" as claimed in line 51-52. This may also downgrade the effect of the mutual denosing scheme. As shown in the ablation studies, discarding the mutual information maximization scheme only has very minor effect on the performance. 2. The organization of the paper can be improved. The main insight lies in how to maximize the mutual information between representations of Siamese Trajectories. However, the major illustration of this contribution is put into the Appendix while the authors spend a great deal of time introducing the preliminaries on diffusion.

Questions

1. Do you think the mutual denoising scheme can also work on homologs as you have mentioned the problem in line 50? There might not be a 1-to-1 mapping on the node representations when it comes to homologs since they may have different number of residues. 2. What do you think is the reason that diffusion-based pretraining produces better results than denoising score matching pretraining in your experiments? In Section 4.4, you have mentioned that perturbing distance matrices may produce negative values, however, it is also natural to do denoising score matching pretraining directly on the absolute coordinates [1][2]. Could the denosing score mathching baseline be underestimated because it is implemented on the distance matrix instead of the absolute coordinates? 3. Can you theoretically show how much deviation will be introduced by the approximation in line 743 (Appendix C.1)? For example, can you derive the bound of the gap between the lowerbound of mutual information and the loss you derived from the RHS of the approximation? 4. What is the value of $t$ in the diffusion model in the fine-tuning phase? [1] Pre-training via Denoising for Molecular Property Prediction [2] Energy-Motivated Equivariant Pretraining for 3D Molecular Graphs

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Random perturbations on the side-chain torsional angles may not conform to the actual distribution of the conformational change.

Summary

The authors propose a novel pre-training method for proteins by jointly modeling sequences and structures with a diffusion model (DiffPreT). The encoder, which is the noise prediction network, learns representations for the sequence and the structure, respectively. This representation can then be used for downstream tasks. Additionally, the authors sugest an extension to their method to capture the correlation between conformer structures, based on the maximization of mutual information between diffusion trajectories for different simulated conformers, obtained by perturbing side chain rotations.

Strengths

The authors bring forward some very interesting ideas and the results are promising, possibly signifying a step forwards within the field of protein representations. Moreover, it is clear that an extensive amount of work has gone into running experiments, including ablation studies and additional results in the appendix.

Weaknesses

The paper is confusing and unclear in some places, both in terms of text and in terms of mathematical notation. There are some essential parts of the method that are not properly explained and/or motivated, and not all results seem to indicate a significant improvement. Finally, I think the results could be presented in a more diverse way, rather than only showing big tables. For comments and suggestions on how to imrpove the manuscript, I refer to the "Questions" part.

Questions

- Main questions / comments: 1. In the DiffPreT setting (and to some extent also in the SiamDiff setting), it is unclear to me what noise level is used to get the representation for downstream tasks. It is mentioned that mutiple levels of noise are used, but it remains vague wether this means concatenating respresentations at multiple noise levels, all noise levels, or picking a noise level at random when getting the representation for the downstream task. 2. Maybe even more generally, could you give some intuition on why representations on noisy structures could serve better for downstream tasks? 3. Could you comment on the size of the latent representations that you get? I might have overlooked it, but I missed what the size of $d$ is. This seems like important information, especially when the SiamDiff setting is used where all these representations seem to be concatenated. 4. For the siamDiff setting, it is not intuitive to me what kind of distribution is really being modelled. Randomly rotating the side chains without perturbing the backbone, and simply throwing away structures where there are clashes seems a bit messy. Is the resulting something "biologically relevant"? 5. When considering the standard deviation of some of the results, the improvements are not always statistically significant. Perhaps it would be good to acknowledge this. 6. The structuring of the paper is a bit off in some places, e.g. 3.4 is showing results even before all methods are explained, and the discussion in 4.4 is more a "related works" section, which has mostly been discussed in the introduction already, and which somewhat disrupts the flow of the paper. 7. From the main text, it took me a long time to discover that the encoder $\phi$ is the noise predicting network, and even longer to see that it is GearNet-Edge (and GVP in the appendix). This could be more clear earlier on in the manuscript. 8. In the appendix it is stated that GearnNet-Edge is a structural encoder. However, from the main paper it seems like you get a representation for both the sequence and the structure. How does this work? Are these separate networks? 9. All results are presented as big tables, it would help to include some intuitive graphical results. Perhaps visualize the protein structure / sequence reconstructions of the diffusion model, or somehow visualize the representations (for example see if clusters appear when PCA is done on the representations or something similar). - Other questions / comments: 10. The term "encoder" can be confusing, as the forward process of a diffusion model can also be seen as a series of unparameterized encoders. It would help if this distinction is clear from the beginning. 11. Perhaps a missing reference for protein structure-sequence co-design: Lisanza et al. [2023] 12. The task explanation in the main paper is very minimal. I appreciate that there is a more extensive description in the appendix, but it could still be much more elaborate. 13. Sometimes abbreviations are introduced before their meaning (e.g. EC, ESM). 14. The "Mean Rank" seems like a strange metric to me that does not add much to all other results. 15. The results using GVP are interesting and would maybe be worth including in the main paper, including a discussion on why some scores are much higher than for GearNet-Edge (except for PSR). 16. Some discussion on computational cost is missing, either in the main paper or in the appendix.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors discuss limitations in the appendix. It would be nice if some of this discussion is transferred to the main paper.

Summary

In this work, the authors perform a thorough investigation of different pre-training strategies on joint sequence-structure diffusion models for representation learning, rather than generative modeling. They evaluate on an EC prediction task and four tasks from the Atom3D benchmark. They find that joint sequence-structure pre-training with the SiamDiff method for maximizing mutual information between diffusion trajectories of protein conformers consistently outperforms other methods on downstream task evaluation.

Strengths

Great observation that prior pre-training strategies excel for particular tasks and have shortcomings in others; and while SiamDiff does not yield huge improvements over other strategies, it does consistently outperform them across tasks. This work takes a thorough and critical look at the benefits and tradeoffs of considering sequence- or structure-only methods, and systematically evaluates different pre-training approaches.

Weaknesses

Lack of evaluation against other downstream predictors and backbone models for downstream prediction.

Questions

The sizable gap in sequence denoising accuracy for sequence vs joint diffusion is very interesting, especially considering that in the ablation, the model without structure diffusion is competitive with the others. Can the authors comment on the additional complexity of joint diffusion compared to sequence-only diffusion, and why this accuracy gap is so pronounced? The authors mention model capacity, but clearly capacity is sufficient for downstream tasks.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

The paper proposes to use joint protein sequence and structure diffusion as a pretraining task, which they call DiffPreT. In order to account for the fact that proteins can exist as ensembles of conformers, the paper further proposes to generate pairs of conformers, use the diffusion forward process to corrupt them, and then train the reverse process across diffusion trajectories. They call this pretrained model SiamDiff. The paper then shows that DiffPreT and SiamDiff are effective on a panel of protein function prediction tasks.

Strengths

The paper addresses the important problem of protein function prediction using novel and effective methods. Using diffusion as a pretraining task is novel, interesting, and apparently quite effective. The SiamDiff objective is creative and also empirically effective. The experiments are well-done and convincing. The writing is clear, and the paper is easy to follow. I appreciate the uncertainties on metrics and the attempt to benchmark across a wide variety of tasks. Using two-stage noise-scheduling is also an intuitive way to adapt a model framework originally meant for generation to pretraining.

Weaknesses

### Major In general, this is a very good paper that clearly describes a promising method for protein pretraining. However, it only evaluates on Atom3D and EC number. The paper would be much more significant if they evaluated GearNet on zero-shot fitness tasks such as those in ProteinGym, protein engineering tasks such as those in FLIP, and more general structure prediction tasks such as those used in CAFA. In addition, the results would be more general if the paper considered different diffusion schemes. For structure generation, diffusion on [frames](https://arxiv.org/abs/2301.12485) or [angles](https://arxiv.org/abs/2209.15611) seems to outperform diffusion directly on atoms. In sequence diffusion, D3PM with an absorbing state isn't as computationally efficient as [autoregressive diffusion](https://arxiv.org/abs/2110.02037) while not allowing iterative refinement like D3PM with a uniform prior over amino acids. The paper would be slightly stronger if it discussed these choices, and much stronger if it evaluated the effect of these choices. However, I realize that this is probably not feasible during the discussion period! In general, I think the exposition of the diffusion model could be shortened, as that is not the main contribution of the paper, in order to make more space for contributions that would increase the paper's significance. ### Minor - The paper should include ablations on EC - It's not clear to me why ESM-2-650M-GearNet doesn't require structure as input.

Questions

- How is EC a residue-level task? Don't the labels apply to the entire protein? - Is there a way to use this model on proteins without a structure? - How does the choice of diffusion process effect performance? - How well does this work on function prediction, zero-shot fitness prediction, or protein engineering tasks?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

- SiamDiff requires structures at inference time. - The paper only considers one diffusion process out of many possible choices.

Thank you for the rebuttal

I am impressed with the additional experiments the authors managed to do during the rebuttal period and think that they increase the contribution of the paper. However, on GB1 2-vs-rest, there are some stronger baseline results in Table 3 here: https://www.biorxiv.org/content/10.1101/2022.05.19.492714v4.full.pdf In general, I think this is a strong and interesting paper that deserves acceptance.

Thanks for all the reviewers' suggestions and feedback!

We'd like to extend our gratitude to all the reviewers for their insightful questions, prompt responses, and constructive feedback. We truly value our discussions with you. Your suggestions are helpful for us to add more interesting experiments and discussions to the paper. We will continue to follow your suggestions on revising the manuscript.