DiffPack: A Torsional Diffusion Model for Autoregressive Protein Side-Chain Packing

Summary

The paper proposes DiffPack, a torsional diffusion model that accurately predicts the conformation of protein side-chains given their backbones. DiffPack learns the joint distribution of side-chain torsional angles by diffusing and denoising on the torsional space. To avoid issues arising from simultaneous perturbation of all four torsional angles, the paper proposes autoregressively generating the four torsional angles and training diffusion models for each torsional angle. The method achieves remakrable improvements in angle accuracy on benchmark datasets and enhances side-chain predictions in the AlphaFold2 model.

Strengths

-The paper is well-written and easy to follow, with a clear description of the proposed DiffPack model. While the idea of using torsional diffusion is not entirely novel, the paper builds on the work of Jing et al. at NeurIPS 2022 on small molecule torsion diffusion and applies it to the problem of protein side-chain packing. -The paper's technical contributions lie in the development of a torsional diffusion model that considers the restrictions imposed by covalent bond lengths and angles, and the autoregressive generation of torsional angles. -The paper's results demonstrate the potential of DiffPack in advancing protein structure prediction and design.

Weaknesses

-Is there a metric available for assessing the overall conformational plausibility of the generated structures? This question is relevant because in the context of small molecule conformation generation, deep learning methods have been shown to produce conformations with lower overall RMSD, but there are still significant challenges with conformational plausibility (e.g. generated benzene rings are not necessarily planar). t would be beneficial to have a metric that takes into account both the geometric accuracy and conformational plausibility of the generated structures to ensure that they are reliable for downstream applications. -It would be helpful to include a comparison of model run times between DiffPack and other methods such as Attenpack and force field-based methods. This would enable us to better assess the potential gaps in downstream applications when making proteomic predictions with these methods.

Questions

Please see the above weaknesses.

Rating

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

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors didn't discuss the limitations and potential negative societal impact of their work.

Summary

The paper proposes DiffPack a torsional diffusion model to learn side chain placements. In particular, the authors presents a few modification to vanilla torsional diffusion models that improve the empirical results obtaining strong empirical performance.

Strengths

The authors propose a number of modifications to the vanilla application of torsional diffusion to side chain predictions and these have enables them to obtain state-of-the-art performance on this important task.

Weaknesses

Although the empirical results are very good, not many parts of the method are significantly novel. Moreover, among the few modifications that the authors made the presentation of them with respect to the rest of the field and the justification should be improved: 1. Annealed temperature sampling: the low temperature sampling procedure derived from the assumption of the Gaussian distribution seems to be very similar to the one presented in Ingraham et a. (2022) [20]. The authors should include the reference of this technique in the relevant section and clarify the difference between the methods. Moreover, I believe the statement “ideally the sampling process converges to the global optimum when T→0” is wrong or at least misleading (does not specify what the ideal situation is). 2. Multi-round sampling: this is related to the mixing Langevin steps that multiple papers have proposed before (e.g. Ingraham et a. [20]). However, unlike in those approaches, the authors do not limit themselves to a Langevin step that preserves the expected distribution but instead keep the “ODE-term” making therefore very unclear what the process is “theoretically” doing in this process. The authors should provide further discussion of these challenges. 3. Autoregressive diffusion: some of the comparisons (excluding the inference performance) used to motivate autoregressive diffusion versus vanilla are misleading: - Figure 7: presents the number of steric clashes during the diffusion process not of the generated structures. This is misleading because these steric clashes may not matter or even help the model to detect them to avoid them. As the autoregressive method (that removes following atoms) artificially removes these steric clashes and this may prevent it to reason about them when generating the distributions. The number of steric clashes should be compared based in the resulting generated structures. - Figure 5 compares the loss values of different methods however these are not comparable. E.g. one of the reason for the loss of X1 is lower than X4 might just be fact that the entropy of the distributions of X1 (which is correlated to the lower bound of the score matching loss) is lower than that of X4. Similarly this makes the comparisons of the losses of the different methods not comparable (e.g. the entropy of the conditional distribution obtained by the autoregressive is lower than that of the unconditional obtained by the joint loss).

Questions

The statement “We focus exclusively on protein side-chain prediction under the assumption of a fixed and highly accurate backbone” seems to contradict with the AlphaFold2 results presented in the main text. “We extend DiffPack to accommodate non-native backbones generated from AlphaFold2”. How is this extension done? What is the model trained on? Section B.1 what is the distance threshold?

Rating

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

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

In some of the tables the (e.g. Table 3) the bolding of the best number is wrong and biased to the proposed method. I believe that the discussion on the number of parameters is misleading as this is not a very useful measure of computational feasibility. Instead the authors should compare the methods by runtime (and potentially memory cost).

Summary

This paper focuses on the task of sidechain packing in proteins, where one wishes to predict the positions of the side-chain atoms given the positions of the backbone atoms. To this end, the main contribution of the paper is DiffPack, an extension of recently proposed torsional diffusion models to the side-chain packing task. This is achieved by combining three aspects: i) autoregressive sampling of the four side-chain torsion angles, with a separate diffusion model trained for each, conditioned on the previous torsion angles ii) multi-round and annealed temperature sampling to improve quality of generated samples iii) Modifications to the transition kernel formulation on the torus for residues where the torsion angles have a periodicity of $\pi$. Experimental performance compared to various baselines for side-chain sampling showcases the improved performance offered by DiffPack.

Strengths

1. The paper is very well written, clear and easy to understand. The technical choices made throughout the paper are sound and well motivated, and the application area considered is well suited for the torsional diffusion framework. 2. The authors compare their method to a variety of baselines for side-chain packing and showcase improved experimental performance of their method.

Weaknesses

1. The technical contributions are largely incremental - the formulations regarding torsional diffusion have already been well explored in previous recent papers. 2. Some questions / clarifications regarding the experimental evaluation: * It is unclear to me the benefits offered by using auto-regressive models for each torsion angle as opposed to joint diffusion. From Table 1, 2 and 4, one can see that, for about 4$\textdegree$ variation in MAE of $\chi_1$, the RMSD drops by 0.1. From Table 4, the $\chi_1$ MAE between the autoregressive model and the joint diffusion model is about 2$\textdegree$, giving a much smaller corresponding drop in RMSD. Using auto-regressive models definitely seems to accelerate training / convergence as noted in Fig 5. However, the errors in discretization during sampling, idealization of bond lengths and angles when reconstructing coordinates could end up reducing the improved training performance of the autoregressive models, and eventually offer similar values in RMSD. * Could the authors add a table regarding the run times associated with the evaluations of the different deep learning baselines? How many samples are generated for each protein before selection with the confidence model? 3. An anonymous link for the code submission is not provided, making it harder to verify reproducibility.

Questions

See the above section for questions. Edit: I have read the authors rebuttal, and post the discussion phase, still maintain my assessment about the paper.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

2 fair

Limitations

The authors have addressed limitations associated with their work

Summary

The protein side-chain packing problem consists of predicting the positions of atoms in amino-acid side chains given the backbone structure and residue identities. The paper proposes to do this using a diffusion model that accounts for physical constraints and models side chain structures as joint distributions over torsional angles. Furthermore, they find that generating the torsional angles autoregressively improves the generation quality. DiffPack outperforms existing methods with fewer model parameters and can enhance side chain predictions from AlphaFold2.

Strengths

The paper provides a unique solution to the important and well-studied protein side-chain packing problem. The writing is generally clear, and the work is well-contextualized in the literature. The autoregressive diffusion framework is intuitively effective, and the formulation seems correct. This is backed up by strong empirical results when comparing to previous work and in the ablations. Providing confidence scores, reducing the number of parameters compared to previous models, and being able to refine AlphaFold2 predictions will make the method very useful for downstream practitioners.

Weaknesses

In general, this is a strong submission with few major weaknesses. ### Major - Clarity: it's a bit unclear to me how the model moves between atomic coordinates in GearNet and predicting the scores and confidences on each angle. A few more lines of text or a subfigure could be very helpful here. Likewise, it would be nice to have some more details about model size, training hyperparameters, and training hardware. - Soundness: The paper talks about confidence scores and shows that they improve generations. It would be better to also have a table or figure showing how well the confidence scores are calibrated. - Soundness: are the baselines also given the chance to generate multiple conformations and then to have a confidence model pick the best one? If not, the comparisons are not quite one-to-one. - Significance: While I appreciate the case studies shown in 5.5, the paper would be stronger with more context here. What downstream biological or engineering applications, if any, can DiffPack do that are not accessible with existing methods? ### Minor - There are a few minor points on clarity: The standard in the field seems to be DDPM (denoising diffusion probabilistic models) instead of DPM, as used in line 252. In Figure 4, the legend says blue and yellow, but I see blue and red in the figure. In Figure 5, the colors for the four autoregressive curves are very difficult to tell apart -- it might help to also vary the linestyle.

Questions

- Is there any intuition for choosing the variance exploding SDE instead of variance preserving? - How well-calibrated are the confidence predictions? - What hardware was used to train the model? How long was it trained for? - What downstream biological or engineering applications, if any, can DiffPack do that are not accessible with existing methods? - How exactly does the model move between atomic coordinates and score / confidence predictions? - When generating, how many different conformations are being generated and passed to the confidence model? - What work remains to be done in side-chain packing? Does DiffDock completely solve the problem? - How hard would it be to train an order-agnostic autoregressive diffusion here? Would we expect that to produce better samples than a model with a fixed decoding order, as shown here?

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

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

- The authors should address where DiffDock still does not solve the side chain packing problem, either in general or in specific cases.

Response to rebuttal

Thank you for the thorough rebuttal. I am especially pleased with the analysis of the confidence scores and multiple chances for baseline methods. Pending discussion with the other reviewers, I am inclined to raise my score to 7.

Thanks for the rebuttal

Thank you for the response. It well addressed my concerns.

Response to rebuttal

Thank you for the response, I am glad my comments were helpful and I appreciate the effort in clarifying the relation of the method to previous work. I have raised my score to 6.

Response to Rebuttal

I thank the authors for their efforts with the rebuttal process. Below follows a point-by-point response to the rebuttal: **Q1** Thank you the comments, the increase in scale is indeed a valid point which I missed to include before, but some concerns proposed to deal with the scale still remain (more in Q2). **Q2** Thank you for the clarification! I understand the issues associated with sidechain packing as you outlined and the utility of autoregressive models in this situation, but if they were to indeed mitigate the issues to the extent as expected, the peformance should be considerably better. Could you provide some references as to ranges of RMSD values in sidechain packing that qualify as poor vs satisfactory?. In the protein structure prediction tasks, <1A RMSD is already considered near native for instance. For a task utilizing diffusion models, 0.12A RMSD could just be a byproduct of the noise scale at the final step. **Q3, Q4 and Q5** Thank you, these experiments definitely shed more light on the utility of the DiffPack as a fast & accurate sampler. Time permitting question this, but have the authors carried out any quantitative analysis of the steric effects in the predictions from models with and without autoregressive diffusion?

Thanks

We appreciate the opportunity to address your concerns and are grateful for your valuable input, which has greatly enhanced the quality of our work. Thank you again for your effort!