Paper: Backdiff: a diffusion model for generalized transferable protein backmapping (ICLR-2024)

Reviewer comment (Theory):
The Equivariance handling approach in the method constructs a reference coordinate system using the first three amino acids of each protein sequence. This implies that if the positions of the first three amino acids vary, the reference system will also differ, which may not be an ideal approach.

Author response:
Thank you for insightful questions regarding the SE(3) equivariance. Our model in the first draft indeed does not use an equivariant neural network to preserve the SE(3) equivariance. Instead, we built a coordinate framework internally for each CG configuration and we used the framework in the diffusion model. However, as you pointed out, this approach might limit the model's expressivity. Thus, we have re-implemented an SE(3) equivariant BackDiff and have updated all results using the new equivariant model. We also prove that with an equivariant score function, we can achieve the desired SE(3)-equivariance.

Reviewer comment (Evaluation):
In Table 3, the Mean Absolute Error (MAE) of bond length for BackDiff (cons) can reach 0, which appears too good to be true.

Author response:
Thank you for highlighting this point about the Mean Absolute Error (MAE) of bond length in Table 3.  The bond length MAE does not actually reach 0, but should be expressed more rigorously as "<0.001", which I have fixed in the revised draft. This extremely low MAE for bond length is expected.  Applying manifold constraint on score-based diffusion model is similar to applying a restraint potential to certain degrees of freedom in molecular dynamics, restraining the variation on these DOF. The effectiveness of manifold constraint is better for constraining bond length compared to bond angle. This observation is consistent with recent studies [Improving Diffusion Models for Inverse Problems using Manifold Constraints] and [DIFFUSION POSTERIOR SAMPLING FOR GENERAL NOISY INVERSE PROBLEMS] which implied that, empirically, a linear condition mapping (like bond lengths) tends to be more easily matched than a nonlinear condition mapping (like bond angles). Additionally, as shown in Table 6, the choice of manifold constraint weight plays a significant role. A small weight can lead to error in condition, while a large weight might over-deviate the sampling. The weight we choose is optimized by balancing the condition error and the sampling path.

Reviewer comment (Evaluation):
In Table 3, the numerical values of the standard deviation (std) in the second and third rows are almost in the same range as the mean values, which is strange.

Author response:
Thank you for raising questions regarding the bond angle MAE. Like in Table 1 and 2, we evaluate the bond angle MAE of generated all-atom configurations for each testing frame, and compute the mean and std across all testing frames. Thus, a large std should be expected as they are results from backmapping different CG configurations.

Reviewer comment (Theory):
Is the model used in the method SE3 equivariant?

Author response:
Thank you for insightful questions regarding the SE(3) equivariance. Our model in the first draft indeed does not use an equivariant neural network to preserve the SE(3) equivariance. Instead, we built a coordinate framework internally for each CG configuration and we used the framework in the diffusion model. However, as you pointed out, this approach might limit the model's expressivity. Thus, we have re-implemented an SE(3) equivariant BackDiff and have updated all results using the new equivariant model. We also prove that with an equivariant score function, we can achieve the desired SE(3)-equivariance.

Reviewer comment (Theory):
Given that the model learns the displacement of omitted atoms from alpha carbons, why not directly learn displacement in the local coordinate system of each amino acid? This approach could ensure that the representation is SE3 equivariant.

Author response:
Thank you for your excellent advice regarding the use of local coordinate systems for each amino acid to ensure SE(3)-equivariance. This approach presents a valuable alternative to our current method. We appreciate this advice and will explore its feasibility in future research.

Reviewer comment (Experiments):
I am not familiar with the PED dataset. Why were only 92 proteins selected out of 227 for training and testing data?

Author response:
I apologize for any confusion regarding the selection of proteins. We select proteins in the PED dataset following the guidelines discussed in the GenZProt paper. Metal ion-binding complexes, nucleotide-binding complexes, cofactor-binding complexes, PTM-including proteins except phosphorylation and oxidation, D-amino acid protein, and proteins simulated or experimentally measured in unnatural conditions are excluded. The 92 proteins selected are those computed from reliable MD simulations and sampling methods.

Reviewer comment (Experiments):
What do you mean by single- and multi-protein experiments? What is the primary difference? When frames are used for data partitioning, is it possible for different frames of the same protein to appear in both the training and testing sets, potentially leading to data leakage?

Author response:
We apologize for the unclear presentation in the manuscript. To clarify: single-protein experiments involve using data from a single protein. We partition ensembles of frames from this single protein into training and testing datasets. This tests the model's ability to predict structures within a single protein's conformational space. In contrast, multi-protein experiments involve training and testing the model on multiple different proteins. This tests the model's ability to generalize across various proteins. Due to the large conformation space of long-chain proteins, and the comparatively small numbers of ensembles in our dataset, we believe data leakage is not an issue. Still, to make our results more convincing, in the new training of the single-protein experiments, we have reduced the number of training data.

Reviewer comment (Experiments):
Have you re-implemented the baseline models (especially CGVAE) to condition them on other CG variables such as N and sidechain COM for the UNRES benchmark? If BackDiff is conditioned on C alpha, N, and side chain COM, while CGVAE is conditioned only on C alpha as in its original paper, it would be hard to tell if the performance difference is coming from the method or the difference in information given to the models. The same applies for MARTINI and Rosetta benchmarks. Alternatively, you could report the performance of BackDiff conditioned on C alpha only, with no CG auxiliary variable constraints on side chain COM.

Author response:
Thank you for highlighting the importance of a fair comparison in our experiments. We did reimplement the baseline models. I maintained all the original GenZProt architecture, except that the residue-residue message-passing and residue-atom message-passing will contain additional particles corresponding to additional CG beads other than alpha carbons. I think this is a fair comparison, as our method BackDiff, and the baseline GenZProt and modified Torsional Diffusion all utilize the same model architecture (e3nn package with E(3) equivariant convolution as message-passing layers).

Reviewer comment (Experiments):
How did you select the PED entries for testing? The three test proteins all look pretty linear and disordered. How does the model perform on a globular protein?

Author response:
Thank you for your insightful observation regarding our selection of PED entries for testing.  Including only disordered proteins may limit the scope of our evaluation. To address this, we have revised our test set to incorporate a broader range of protein structures. Specifically, we have replaced one of the previously used test proteins (PED00025, a disordered protein) with PED00055, a well-characterized globular protein. This choice was guided by the aim of assessing the model's performance across a more diverse range of protein structures, including those with more defined and compact conformations.

Reviewer comment (Evaluation):
Could you report the diversity of the generated structures conditioned on the same CG structure? For example, in the referred baseline [1], the authors reported quantitative metrics for diversity, such as the Earth Mover’s Distance for side chain torsion angles.

Author response:
Diversity is an important metric to evaluate, and we have included such a measurement (DIV) in the new draft. The result shows the BackDiff trained with transferable CG protocol generates samples with the highest diversity overall.

Reviewer comment (Evaluation):
Could you provide a speed analysis of your method, for example how much time required to backmap a frame, similar to what was done in [1]?

Author response:
Certainly! We have added a brief speed analysis in Appendix I: Sampling Efficiency. In general, the diffusion model is slow in sampling, especially with the manifold constraint technique. BackDiff roughly requires around 5 minutes to generate 100 samples for a single frame.

Reviewer comment (Experiments):
It could be interesting to see how the model performance changes as we increase the CG resolution (the number of CG atoms and CG auxiliary variables), especially in terms of the diversity of generated all atom conformations. This could provide insights into the CG system’s entropy. This is not a requirement, but just a curiosity.

Author response:
This is indeed an interesting question to investigate. Although we do not put such a comparison in the text, we can observe Table 2 (multi-protein experiment for UNRES, a CG of lower resolutions) and Table 8 (multi-protein experiment for MARTINI, a CG of higher resolutions).  Backmapping from CG of lower resolutions will generate samples with higher diversity (lower DIV score). We suspect that for BackDiff, since MARTINI utilizes multiple sidechain beads, and all this information is used for manifold constraint, the generated samples all get closer to the original all-atom configuration.