Str2Str: A Score-based Framework for Zero-shot Protein Conformation Sampling

Questions:

Shape of the mentioned variables.

• They are dependent on the number of residues: the shape of respective variables for a single protein are: T = (N_residue, ); R = (N_residue, 3, 3); v = (N_residue, 3); x = (N_residue, 37, 3).

Explanation of the sampling diversity.

• Thanks for asking this question. We want to clarify that the proposed Str2str is in a transfer learning setting which does not require distribution data of the test system for training. During sampling, the perturbation (forward) process enforces exploration (say diversity) and the amortized learned score network will anneal (exploit) the perturbed sample to a candidate conformation. Overall, the forward-backward process achieves good balance between exploration and exploitation.

Explanation of the term “dynamics” in Section 3.2.

• Explanation: The semantic meaning of “dynamics” used in the forward-backward dynamics originates from the Langevin dynamics — which the backward process of diffusion/SGMs can belong to. We agree that “dynamics” comes with a time, while the time indicates here the time index of the diffusion process. According to your suggestion, we have altered the writing by using “process” in replace of “dynamics” in the Section 3.2. We hope this could address your concern.

Whether to ensure dissimilarity between train-test sets.

• Thanks for asking this question. In case there may be the least misunderstanding, we want to humbly clarify that such a process is not necessary: (1) amino-acid sequence information is neither used as input nor output; (2) The resulting evaluation is conducted based on the distribution metrics instead of the accuracy for a single structure. In other words, the setting is unsupervised w.r.t. the conformation sampling task; (3) The sampling for each test system corresponds to a new task during inference, and not a single data point belongs to that task is present during training. It is by definition zero-shot learning; (4) From the ML point of view, the usual train-test dissimilarity is enforced to avoid data leakage from supervision, which can make the “generalization” non-sense. This does not happen in our case.
• Therefore, we preprocessed the training data by excluding the PDB entries to contain no a single conformation of all test systems.

We sincerely hope our response can address your concerns. If you have other questions, we are very glad to discuss them during the rebuttal.

Firstly, we sincerely thank the reviewer for providing helpful suggestions to improve our manuscripts. Specifically, we want to address the above concerns as follows:

Weakness:

Seemingly “incremental” to previous methods.

• Thanks for raising this point. We want to humbly argue in the following aspects:
  • Firstly, the focus of this research paper is to present a zero-shot framework for the task of conformation sampling w/o expensive simulation data. As far as we are concerned, Str2str is the first-in-class example that can sample protein-level conformations in the zero-shot manner.
  • One of our main contributions is the successful experimental demonstration that it is able to transfer the generative models pretrained on PDB with score matching objective to the downstream task of conformation sampling, even in a zero-shot manner. This distinguishes our work from those directly applying well-known methods such as diffusion to new data in an old task framework, and thus should not be criticized.
  • Lastly, we are not aiming at proposing new generative methods extending SGM/diffusion [1] or introducing a new structure decoder on top of IPA as in AlphaFold2 [2]. Correspondingly, beyond Diffusion and IPA, our framework can flexibly accommodate other counterparts, such as LDM[3], flow matching[4] or SE(3)-transformers[5], etc, and these can be future work.

More related works for MC/MD methods.

• Thanks for this valuable suggestion. To better contextualize, we have carefully reviewed the related MC/MD literatures and complemented the relevant discussion in the Appendix G.1.

Baselines of MC and MD.

• Firstly, we note that the MD baseline, obtained from D.E. Shaw Research [4], were already shown in our main results. These simulations are among the most authentic and well-recognized studies in the MD domain, and unique in the simulation time scale.
• Secondly, for MC, their validity should be comparable to MD since both of them involved force fields. Since MC is not mainstream in the MD domain (since it is not scaling well), there are very few works studying the protein dynamics using MC in recent decades.
  • Evaluating MC methods on fidelity/diversity requires sampled trajectories, which are not accessible. Running any MC or MD in the same time scale as [5] from scratch is computationally intractable for us. The estimated simulation time for all benchmark systems far exceeded >10,000 GPU days, which is comparable to LLM pre-training such as Llama[7]. The most relevant work [8] conducted a MC-based study using a 100-fold larger timestep (x100 acceleration), it still requires ~100 GPU days and their data is not open-sourced.

References:

[1] Song, Yang, et al. "Score-Based Generative Modeling through Stochastic Differential Equations." International Conference on Learning Representations. 2020.

[2] Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." Nature 596.7873 (2021): 583-589.

[3] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[4] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).

[5] Fuchs, Fabian, et al. "Se (3)-transformers: 3d roto-translation equivariant attention networks." Advances in neural information processing systems 33 (2020): 1970-1981.

[6] Lindorff-Larsen, Kresten, et al. "How fast-folding proteins fold." Science 334.6055 (2011): 517-520.

[7] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models." arXiv preprint arXiv:2302.13971 (2023).

[8] Heilmann, Nana, et al. "Sampling of the conformational landscape of small proteins with Monte Carlo methods." Scientific reports 10.1 (2020): 18211.

Dear Reviewer 6bVL,

Thank you for your helpful suggestions for improving our manuscripts. We sincerely appreciate your feedback and have carefully responded to each question point-by-point.

This is a kind reminder that as the reviewer-author discussion period is ending soon, we look forward to hearing from you about your feedback to our response. (Unlike previous years, there will be no second stage of reviewer-author discussion this year)

In particular, we have complemented related discussions and improved the clarification of writing, according to your advice.

Thank you again for your precious time and effort!

Best regards,

Authors #3973

Questions (continued):

Discussion on the possible sampling of outlier.

• Since there is no criterion defining “valid” conformation, let us denote “valid” as local minima belonging to the underlying distribution. For fast folding proteins, we can infer from the results in [6] that distant (if compared with a folded start) yet valid conformations are transient unfolded states. In Figure S6, large $T_\delta$ tends to destroy the original fold, while keeping some local secondary structures. Especially in the case of WW domain (we complement a Figure S8), the main components become coil (loop) and loosened for large $T_\delta$ .
• On the other hand, using three JS divergences instead of simply comparing the diversity takes the outlier into consideration during evaluation. The “diversity” roughly reflects the radial of the space coverage while the “JS” penalizes the mistakes by sampling outliers in the wrong direction. The current gap between Str2str and MD reference represents one can still improve this imperfection in the future. In our humble view, even conditioning on sequence can neither guarantee the sampling of “valid” conformation, if without enough training data.
• Lastly, as a typical “invalid” or “outlier” example, the idpGAN [7] has a strong tendency to sample coil components, which can belong to the “outlier” for most structural proteins. Therefore, the used fidelity metrics differentiate and penalize them based on the MD reference.

Investigation of reweighting (importance sampling).

• According to your suggestion, we have re-evaluated the ensemble w/ logp(X)=Constant (i.e. only use the FF to reweight). The JS-divergence results are [JS-PwD=0.471; JS-TIC=0.448; JS-Rg=0.487], which is still worse than that w/o reweighting. We dived into the details and have the following preliminary analysis:
  • We noticed the minimized energy for the generated ensembles has much larger standard deviations (~100 kcal/mol) than that in normal MD simulation (<10 kcal/mol) after equilibration. Since the energy is mainly contributed by the side chains, we probe our packing module (FASPR) by re-packing a batch of equilibrated conformation. The minimized energy of the repacked conformations is 50-100 kcal/mol higher than the pre-packed energy. Such results show that, even given perfect backbones, it still fails to generate low-energy side-chain conformation within the minimization steps. Therefore, the importance weight is not stable or accurate, which causes large variance to the estimation. As illustration, energy difference in 5 kcal/mol under 300K will lead to weight $w \approx 4414$. The distribution will be sharply biased toward very few samples, making the estimation much worse by using only 1000 samples. In practice, a more robust packing module can be used to perform effective energy reweighting.

Apo/holo dataset.

• According to your suggestion, we have added the apo/holo benchmark in comparison with EigenFold by strictly following their settings, which is updated in Appendix F.4.

Code / reproducibility.

• We are currently refactoring the codebase into a more user-friendly version. As stated in the Reproducibility Statement paragraph, we promised all source codes, scripts, training data, necessary artifacts and experimental environments (settings) will be released upon the acceptance of this work. The benchmark MD trajectories are obtained from the original authors in D.E. Shaw Research (DESRES, https://www.deshawresearch.com/index.html). According to the DESRES License 3(e), “Licensee shall not…disclose, without DESRES's prior written approval, the Datasets, whether in whole or in part, other than as expressly authorized hereunder…”. It is thus illegal to put them on OpenReview during the double-blind review. For both reproducibility and legal purposes, we will sincerely ask DESRES whether it is proper/how to share the benchmark data upon acceptance.

References:

[6] Lindorff-Larsen, Kresten, et al. "How fast-folding proteins fold." Science 334.6055 (2011): 517-520.

[7] Janson, Giacomo, et al. "Direct generation of protein conformational ensembles via machine learning." Nature Communications 14.1 (2023): 774.

We sincerely hope our response can address your concerns. If you have other questions, we are very glad to discuss them during the rebuttal.

We appreciate the insightful questions and helpful suggestions for improving our manuscript. We have appended the related experiments in the updated PDF. Our response to each question is as follows:

Questions:

Difference between Str2str and Framediff + Forward-backward (FB) dynamics.

• Firstly, according to our experimental results in the paper, the amortized conformation sampling can be competent by a backbone generative model trained on PDB. Intuitively, our framework is ready to accommodate all backbone design methods [1,2,3,4] including FrameDiff [4] in a plug-and-play manner, by simply applying the FB dynamics. Among them, FrameDiff that consists of the Riemannian diffusion on $SE(3)^n$ and the IPA module of AF2[5] was acknowledged by us as a very effective score model to follow.
• Secondly, we argue that Str2str is targeting a totally different task, based on an existing model Framediff: Framediff (and its counterpart [1,2,3]) aims to unconditionally sample novel and designable backbones; Str2str aims to approximate the underlying conformation distribution similar to MD. The model architecture coincidence comes from (as our main contributions) how we designed the zero-shot sampling pipeline and demonstrated the effectiveness of transfer learning and thus should not be criticized.
• Specific implementation differences: (i) the ideal frame geometry (idealized bond lengths and angles) used during training and inference is sequence-specific instead of all alanine; (2) the training data is set to be in length between 10 and 512. (that of framediff is >60, which is larger than most test systems), etc.

Ablation on orientation (rotation matrix) modeling.

• The original purpose of using “frames” is to be able to generate all heavy atom coordinates plus the packing module, where the orientation can be important. According to this suggestion, we newly conducted the training w/o rotation and compared it with the previous checkpoint trained w/ rotation. Given the limited time of rebuttal, both checkpoints are aligned to be 300,000 steps for fair comparison. As shown in Appendix C.4, modeling the orientations can be important for the zero-shot conformation sampling.

Explain the difference between high $T_\delta$ sample and initial sample.

• An increased perturbation scale encourages structural diversity and the samples can be different from the starting conformation. Such phenomenon does not contradict with the resulting small JS-PwD/MAE-TM, because these metrics are compared with long reference MD: JS = JS(p_pred||p_MD); MAE = |diversity(pred) - diversity(MD) |. After sufficiently long simulation, MD trajectory also discovers the conformation states that are quite distant from the starting one. The “small” resulting metrics demonstrate the model can mainly agree with the MD with small discrepancy (JS, MAE).

References:

[1] Watson, Joseph L., et al. "De novo design of protein structure and function with RFdiffusion." Nature 620.7976 (2023): 1089-1100.

[2] Ingraham, John, et al. "Illuminating protein space with a programmable generative model." BioRxiv (2022): 2022-12.

[3] Lin, Yeqing, and Mohammed AlQuraishi. "Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds." arXiv preprint arXiv:2301.12485 (2023).

[4] Yim, Jason, et al. "SE (3) diffusion model with application to protein backbone generation." arXiv preprint arXiv:2302.02277 (2023).

[5] Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." Nature 596.7873 (2021): 583-589.

Dear Reviewer Fe1f,

Thank you for your helpful suggestions and the recognition of the contribution of our paper. We sincerely appreciate your feedback and have carefully responded to each question point-by-point.

This is a kind reminder that as the reviewer-author discussion period is ending soon, we look forward to hearing from you about your feedback to our response. (Unlike previous years, there will be no second stage of reviewer-author discussion this year)

In particular, we have added more experimental results and related discussion per your advice. If you need any clarification, please feel free to contact us and we are very glad to discuss with you.

Thank you again for your precious time and effort !

Best regards,

Authors #3973

Dear authors,

Thank you for your response, which addressed most of my concerns. My evaluation has been updated accordingly. Lastly, would you mind showing the TM$_\mathrm{ens}$ plot (EigenFold Fig 5) in the appendix?

We sincerely thanks for your questions and helpful suggestions for improving our manuscript. We put our response to your concerns accordingly as follows:

Weakness:

The size of the benchmark/test set.

• First we understand this concern. In fact, the authentic protein trajectory data (long enough) is very limited in the MD domain. Most previous MD research works did not make their data disclosed. The 12 fast folding protein trajectories run by the DEShaw Research are one of the most well-recognized. Though N=12 seems far fewer than common ML benchmarking practice, it is still one of the largest scales for test systems. For example, one related work [1] only chose a subset (N=5) out of these 12 to evaluate their performance.

General presentation issue.

• Thanks for your suggestions. Several clarification issues suggested by other reviewers are addressed by us. Meanwhile, we will pay attention to the overall writing of this manuscript.

Questions:

Why is not diversity “the higher the better”?

• This observation comes from the structural characteristics of different proteins. (i) Some proteins can show very deterministic structure, an extreme example is the cyclic peptides. In this case, the motions in the protein dynamics are limited, and thus with small ensemble diversity from MD simulations; (ii) On the other hand, other proteins show more dynamic behaviors like distinguishably different states. They can have relatively larger diversity, which can be loyally reflected in the long MD simulations. For example, the SARS-Cov-2 spike protein was found to have transitions between its open and closed states [2]. So it is reasonable to compare with a MD-reference diversity rather than claiming it is “the higher the better”.

Time profiling is shown in Table 2 for one system only.

• The runtime profiling on one target is enough to demonstrate the efficiency of the proposed Str2str as (1) the runtime scale is significantly smaller; (2) different systems can share a similar maginitude of simulation time. To better illustrate, we profiled for 100us MD the smallest system (Chignolin) which takes ~95.2 GPU days, and the largest system (A3D) which takes ~200.8 GPU days.

The reason to profile 100us MD but not for 100ns.

• In MD domain practice, running for a scale between 100us (0.1ms) and 1ms is considered to be a typical “long” simulation to discover natural dynamic characteristics. For example, in [3,4], all the systems are simulated for at least 100us, up to 1ms. Thus, we see 100us as a doorsill to become “long”. Since longer simulation is considered to uniformly “better” reflect the real dynamics than the shorter one, we thus omit the shorter cases. Also to illustrate, 100ns MD simulation in the GTT case (in explicit solvent) on a single V100 still takes ~3.87 hrs to finish, let alone 100ns is generally considered “not enough” in standard practice.

References:

[1] Arts, Marloes, et al. "Two for one: Diffusion models and force fields for coarse-grained molecular dynamics." Journal of Chemical Theory and Computation 19.18 (2023): 6151-6159.

[2] Gur, Mert, et al. "Conformational transition of SARS-CoV-2 spike glycoprotein between its closed and open states." The Journal of chemical physics 153.7 (2020).

[3] Shaw, David E., et al. "Atomic-level characterization of the structural dynamics of proteins." Science 330.6002 (2010): 341-346.

[4] Lindorff-Larsen, Kresten, et al. "How fast-folding proteins fold." Science 334.6055 (2011): 517-520.

We sincerely hope our responses can address your concerns. If you have other questions, we are very glad to discuss them during the rebuttal.

Dear Reviewer dLpk,

Thank you for your helpful suggestions and the recognition of the soundness and contribution of our paper. We sincerely appreciate your feedback and have carefully responded to each question point-by-point.

Unlike previous years, there will be no second stage of reviewer-author discussion this year. If you have other question or need any clarification during this period, please feel free to contact us and we are very glad to answer.

Thank you again for your precious time and effort !

Best regards,

Authors #3973

Questions

Clarification of the Forward-backward dynamics.

• Firstly, we clarify that the condition is consistently the input across experiments. Here, we only use $p(T|T_0; \theta)$ as a proposal distribution to obtain i.i.d. samples for each test system. A very natural extension echoing this question is that, we can craft a MCMC sampling pipeline by altering the starting point at each MCMC step $i$ and perform sample using $p(x^{[i+1]}|x^{[i]}; \theta)$ where condition $x^{[i]}$ is the sample from the step $i-1$. According to your suggestion, we provide a description of this in the method section (page 4).

Description of side chain packing.

• We are sorry for not making it clear. In section 3.1, the sampling is conceptual and any form of the side-chain modeling can be used. The key idea is to sample atoms step-wise, from backbone to side-chains. In practice, as introduced in Section 3.2, we used the packing module FASPR [3] which performs a rotamer-library search. Other packing methods can also be used, such as Diffpack [4] that predicts the torsion angles and use them to recover the corresponding atoms similar to what is done in AF2 [5].

The effect of $T_\delta$ on ensemble diversity.

• Thanks for bringing up this question. To investigate this behavior, we additionally investigated the ensemble diversity versus different $T_\delta$ in Appendix F.3 / Figure S7. Indeed, the diversity tends to plateau after $T_\delta$=0.8 for PF and SDE, which shows the enhancement of diversity does not apply to all $T_\delta$. We appreciate this empirical finding and have left a remark accordingly in the paper. Moreoever, from a theoretical view, the forward diffusion will not reach the Gaussian even for T=1.0; it only approximates the Gaussian when $t$ goes up and it is not strictly equivalent to sampling from a gaussian noise.

The motivation of using edge layers.

• We here justify the rationale: the DenoisingIPA can be viewed as a variant of the StructureModule (SM) in AF2 that translates frames to frames. The original SM only updates the single representation $s_i$, frames $T_i$ but not pair representation $z_{ij}$. Since the $z_{ij}$ there is the output of a strong encoder model called Evoformer. We can anticipate very informative $z_{ij}$ even without any updating inside the SM. However in our case, the initial $z_{ij}$ is as simple as the input embedding (time, position). So it is better to update the pair representation along inside $s_i, T_i$ in the model.

Typo in algorithm 2.

• Thanks for pointing this out. We have refined the notation in Algorithm 2. Also, we note that a change of variable is applied for the backward time domain $[T_\delta, 2T_\delta]$ so that we can conveniently leverage the prior discretization of time.

Interpretation on why str2str works.

• Thank you for starting us off with this insightful question.
  • Firstly, we summarize that Str2str learns the vector fields, i.e. score, pointing towards a data-like stuff in the whole structure space (protein landscape). Thus, the trained model can (be expected to) push arbitrary perturbed structure to a nearby data-like minima. Str2str demonstrates a typical transfer learning setting by pretraining on PDB dataset, which enables it to do zero-shot conformation sampling for unseen proteins. Unlike EigenFold, str2str does not rely on sequence embedding but only an initial geometric frame during inference.
  • “The reason why the model works” can lie in our key assumption that: the dynamic patterns among conformations can be partially shared with the patterns between different structures in the PDB. That is to say, the common inter-protein structure differences in the training set can be distilled in the score network, which is “unlocked” during sampling to infer the inter-conformation deviation within a specific protein system. This agrees with the fundamental motivation of “generalization” in deep learning.

References:

[3] Huang, Xiaoqiang, Robin Pearce, and Yang Zhang. "FASPR: an open-source tool for fast and accurate protein side-chain packing." Bioinformatics 36.12 (2020): 3758-3765.

[4] Zhan, Yangtian, et al. "DiffPack: A Torsional Diffusion Model for Autoregressive Protein Side-Chain Packing." arXiv preprint arXiv:2306.01794 (2023).

[5] Jumper, John, et al. "Highly accurate protein structure prediction with AlphaFold." Nature 596.7873 (2021): 583-589.

We really hope the above responses and revisions can address your concerns. Please kindly let us know if you have any other questions. We’re always happy to have further discussion and improve the quality of the manuscript.

We firstly appreciate the insightful questions and suggestions for improving our manuscript. We have added corresponding experiments and discussion in the updated PDF. Our responses to the specific concerns are listed below:

Weakness:

Review and contextualize str2str based on prior studies.

• We appreciate this suggestion and have appended a thorough review by contextualizing the prior works, especially the EigenFold in Appendix G.2. To answer the questions:
  • Similarity: both models leverage PDB for training with a denoising score matching objective. Difference: the EigenFold modeled a mapping from the sequence-encoded latent space (OmegaFold embedding ) to the structure space, while the Str2str delineated solely the structure space and sampled conformations in an amortized manner. During inference, EigenFold performs classic conditional generation while Str2str transforms the unconditional distribution into a translation proposal, from which it samples diverse conformation.
  • We note that Str2str does not use ESMFold [1] embedding but only the output structure, which is different from EigenFold. We reckon the performance gap is due to the fact that Eigenfold handles a more challenging conditional generative task, where the embedding condition ($s_i, z_{ij}$) is far more complicated than a scalar label (for example in [2]). Without enough training data for each condition, the resulting diversity is thus limited to a large extent. In contrast, Str2str learns the structure landscape from which the conformations are sampled and can perform better.

Examine bond breaking as validity.

• Thanks for raising this great point! We sincerely adopt this suggestion and investigate the bond-breaking ratio for each model. In specific, we have defined and evaluated a new validity metric (”Val-Bond”) showing the ratio of samples w/o violating a statistical upper bound for adjacent Ca-Ca atoms. Relevant contents are updated in Section 4.1, Appendix E & Table [1, S3-7], and the results show that conformations sampled by Str2str are nearly all valid w.r.t. both clash and bonding.

[1] Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." Science 379.6637 (2023): 1123-1130.

[2] Song, Yang, et al. "Score-Based Generative Modeling through Stochastic Differential Equations." International Conference on Learning Representations. 2020.

Dear Reviewer BVKt,

Thank you for your helpful suggestions and the recognition of the soundness and contribution of our paper. We sincerely appreciate your feedback and have carefully responded to each question point-by-point.

This is a kind reminder that as the reviewer-author discussion period is ending soon, we look forward to hearing from you about your feedback to our response. (Unlike previous years, there will be no second stage of reviewer-author discussion this year)

In particular, we have added more experimental results and discussions per your advice. Please refer to our posted point-to-point response and we are glad to discuss with you.

Thank you again for your precious time and effort !

Best regards,

Authors #3973