Structure Language Models for Protein Conformation Generation

Dear reviewer X2Sg,

Thank you very much for the helpful suggestion and question. We want to address them correspondingly as follows:

(for metric):

• We have carefully selected a relatively diverse range of tasks and used corresponding metrics that require different capabilities to comprehensively evaluate both the pros and cons of our method, instead of only reporting tasks where our method performs better.
• We also want to highlight that our SLM models are better than other SOTA sequence-based methods across almost all tasks and metrics. This provides enough evidence of the advantages of our SLM methods, as a family of sequence-based methods training using the same PDB data. We also note that MSA-based methods sometimes outperform our method, especially in the case of conformation changing pairs and IDP tasks, while at the cost of being more computationally expensive in practice and strongly dependent on the number of available homologous sequences.
• We believe the “key metrics” are closely related to the exact application (as well as their available format of “ground truth” to compare). There is no universal definition in the current related study. For example, for IDP researchers, the stable reference structure does not exist, so we ground the evaluation on the discrepancy of distribution or statistics; for fold-switchers such as enzyme/kinase-like target, we can curate the ground truth pairs from the database for evaluation, thus pairwise TM-score can play a part. In practice, we state and reason the most relevant metric as follows (those metrics would be highlighted in the next revision if needed):
  • for BPTI (Table 1), we prefer JS-PWD (introduced in Str2Str, the symmetric distributional metric based on the most abundant invariant feature vector of structures);
  • for apo-holo and fold-switch (Table 3), we believe TM-ens (introduced in EigenFold) can truthfully reflect the matching performance for both ground-truth targets;
  • for IDP (Table 4), we use MAE-PwD because the PwD features contain more abundant information (more dimensions) than the other two.
• Last, we understand the concern that it is hard to tell which model is superior without a comprehensive metric (ideally we want some metric like FID or Inception score for image generation). That raises the special need for a comprehensive benchmark covering diverse data of protein types and proper metrics in our community.

(for Atlas data and benchmark):

• We are happy to compare with AlphaFlow, which is a very nice relevant work. For AlphaFlow/ESMFlow, we however think the data setting in their work is not quite aligned with our LM model design. This is because the training/evaluation 100ns MD simulation simply produces the rather limited and local movement of the system. We note that, AlphaFold, as well as the flow matching variant AlphaFlow, show limited diversity but it can be good at capturing the small, local conformation moving with small Euclidean distance (on average 2 angstroms in the trajectory, see below). The geometric space-based models such as AlphaFlow, are easy to learn the distribution over these highly frequent, “local” distributions, since they are closely neighboring in the euclidean space. \
• However, in SLM, since we quantize the structure into distinct tokens, the geometric distance and edit distance of tokens are not smoothly correlated. We have examined the Atlas dataset and unfortunately found a surprisingly high (>50% on average) edit distance between conformation in the same trajectory, even though they show on average 1-2 A RMSD. Based on this limitation, it is promising to point out future work that design better architectures capturing the correlation from geometric distance to properly defined latent token distance.
• As an probing investigation, we firstly reported some trajectory RMSD in mean and max as follows:

Since the targets are general monomeric proteins with hundreds of residues, the average trajectory RMSD ~2 angstrom is so small to reveal significant changes.

(to be continued)

Dear reviewer 2wMe,

Thanks for your supportive comment and time to review our paper! Here is our response to your review.

Response to weakness:

Thanks for your good suggestion. Accordingly, we have added the model size of baselines in the Section 5.4 - Runtime analysis to have a good reference of each baseline and discuss the training data used in detail in the Appendix A.5 by carefully reviewing the mentioned baselines. Unlike the purely language models, we think it is not completely fair to directly compare between different types of model (eg., AF2-based inference requires O (N^3) triangle attention computation). That means smaller model size doesn’t even imply a more “lite” / “efficient” training or inference process.

Response to question:

For this question, we explain in the following sense: The main difference between SLM variants is their architecture since we keep the training/finetuning data fixed. Since in the latent space, the Euclidean protein structure is quantized into discrete tokens and thus the whole structure is represented by a sequence of tokens z, different model variants can learn in a different way and possibly develop different biases towards a distribution of generated {z}. In this sense, for a specific test case of BPTI, we may imagine the kinetic clusters scattering in the different positions of the space. The generated (conditional) distribution p_<model>(z | sequence) can be different depending on <model> due to the imperfect training, which is because the training does not provide the whole information of such distribution over trajectories during training/finetuning. So the inference is in a sense of “transfer learning”: it aims at learning on static PDB data and “guess” (predict) the possible meta-stable states of BPTI during inference.

Let us know if anything is missing or confusing. We are happy to further discuss and keep improving our manuscript!

Dear reviewer 2wMe,

Thank you again for your supportive comments and constructive suggestion! As the author-reviewer discussion deadline is approaching, we want to kindly follow up to see if the response addresses your concerns or if you have any further questions :)

Author #5043

Firstly, we want to kindly thank reviewer yRU8 for the detailed and helpful review. We carefully list our response to your concerns/questions (according to the organized bullet points):

For novelty, our methodology mainly aims at presenting “how and why” it makes sense to use presently popular structural autoencoders (eg., [1,2,3]) + latent LMs with empirical studies and efficiency profiling. For modern AI for Science research, we believe that improving upon well-established foundation models is equally important as developing ad hoc novel models (dominated nowadays). In this sense, we are building up our work and trying to draw more attention to contributions that make the foundation models work well in specific tasks such as conformation generation (providing a recipe). We ponder that a better improvement of our novelty may lie in proposing new invariant architectures in the auto-encoders that are more suitable for latent generative modeling!

For quality, (1) In MD datasets:

Thanks for the suggestion. We first note that the Atlas dataset is a bit special for their data distribution and not fully suitable for the current application of ESM-like model (we will reason below). We tried performing the fine-tuning for ESM3 on the Atlas training dataset (following Jing et al.), it seemed too computationally intensive to fit during rebuttal. We instead compare between the AlphaFlow and ESM3 (zero-shot) in the table below. The evaluation is based in the bb-only mode because the open-source ESM3 cannot model side-chain atoms (thus the exposed residue is also not available).

From the results above, it is evident that the ESM3 struggles with modeling the “local” and “detailed” structural changes in the 100ns simulation Atlas dataset, as many conformations may cluster within the xyz neighborhood. This highlights the advantages of explicitly generative modeling in Euclidean space for this Atlas task (their specific data distribution), where the AlphaFlow model can be more precisely defined here. Building on this observation, we propose that a promising direction for future research lies in developing novel SLM architectures tailored to MD data, since it is not trivial to directly apply to. These architectures should incorporate inductive biases that align with the local exploration characteristics of MD simulations, which could involve designing more effective auto-encoders and integrating advanced structure language model priors.

(to be continued)

Dear reviewer yRU8,

Thanks very much for your follow-up suggestion. We apologize if our response implied any such intention. We did not mean to argue or justify whether Atlas is imperfect. Our prior response "Atlas dataset is a bit special for their data distribution and not fully suitable for the current application of ESM-like model…" was to provide a possible explanation for why language models may fall short of Euclidean-space diffusion models (eg., AlphaFlow), especially in the 100ns MD data setting.

Nevertheless, we believe all these mentioned datasets are valuable and cover different aspects, and thus we discuss them to provide more context in the rebuttal. Surely, for more comprehensive benchmarking purposes, we are happy to include the additional results in the manuscript, as you suggested; we also believe it may provide good context to inspire valuable future works!

To conclude, we have addressed these in the updated pdf:

• (Globally) Rename "minibatch Gibbs sampling" -> Iterative Decoding (ID), according to the naming convention in ESM3 preprint paper (A.1.10 - ESM3 Inference, Hayes et al. [1])
• Added the MD results from rebuttal to Appendix E.4/Table S6 (+will complement the table and enrich the discussion in the final version)
• Incorporated the difference between ESM3 and ESMDiff in Appendix B.4 and also put a reference anchor in the first paragraph of Section 5.

We want to extend our gratitude again for your helpful assessment and supportive comments for improving the quality of our manuscript, as well as pointing out the promising follow-up directions for future research :)

[1] Hayes, Tomas, et al. "Simulating 500 million years of evolution with a language model." bioRxiv (2024): 2024-07.

Dear reviewer FWmx,

We want to extend our heartfelt thanks for your supportive comments on our manuscript and the detailed suggestions. Here is the response to Question 1[major]:

Thank you so much for suggesting the related work. We agree that pure LM with structure coordinates as text is a fairly parallel work to our method given similar data and model architecture. Yet neither of the papers has released their code. The official repos of both LM-CH [a] and BindGPT [b] [https://github.com/danielflamshep/xyztransformer, https://github.com/insilicomedicine/bindgpt ] are yet empty :( We would like to add the comparison once the code is available making our experimental benchmark more comprehensive.

That being said, we add the following discussion regarding the pure LM over coordinates and explain why our approach may be better in the current revision:

• Firstly, formulating the conformation generation problem as pure “text” generation and using atomic coordinates directly as the target requires the LM to handle longer context dependencies due to the tokenization manner, which adds additional challenges to the learning problem[1,2].
• Also, taking [b] for example, the total number of atoms in the “small molecules” is limited only to tens to hundreds, and thus learning on the system may not suffer from long context or scaling. However, without proper coarse-graining, successful models on small molecules may not work as well for large molecules such as proteins of thousands of atoms. How to properly coarse-grain the xyz coordinate file is a non-trivial problem in practice.
• Moreover, without further modifications, cross-entropy loss on the text tokens of xyz digits does not serve as a good measure for the distance between geometric structures and thus the space structure to be learned. After the training, the GPT model may achieve a high likelihood by only capturing high-frequency xyz patterns instead of the low-frequency features that are critical for determining valid structures. For similar reasons, the tasks that use LMs to generate continuous data, such as text-to-image and video translation, often involve the LM predicting quantized features produced by a VQ-VAE [3,4,5].
• Last but not least, the conformation generation requires some configuration information about the protein structure we want to model (in our case, the amino acid types). This is however different from these models which aim to de novo “design” unnatural molecules from scratch. How to place such conditions in these models is still a good open question.

We value this interesting open discussion, and we think it is important to design and encode proper inductive bias into modeling the xyz coordinate in the file data and make these LM models more promising in parallel to our SLM!

[1] Kuratov, Yuri, et al. "BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack." arXiv preprint arXiv:2406.10149 (2024).

[2] Lee, Jinhyuk, et al. "Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More?." arXiv preprint arXiv:2406.13121 (2024).

[3] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017).

[4] Yan, Wilson, et al. "Videogpt: Video generation using vq-vae and transformers." arXiv preprint arXiv:2104.10157 (2021).

[5] Bordes, Florian, et al. "An introduction to vision-language modeling." arXiv preprint arXiv:2405.17247 (2024).

Response to Q2:

Thanks for this point. We believe that joint training for the discrete VAE and LM prior is not very feasible in this context, as it requires backpropagation through latent tokens to LM forward on the fly. We attempt to answer this question from the theoretical perspective (we apologize in advance if we misunderstood the question.):

• Joint training of LM prior and the dVAE may suffer from unstable training. During the early training stage, the dVAE posterior can generate sequences of latent tokens that may be out-of-distribution along with the training. This is because the posterior p(z|x,c) is parameterized by a categorical distribution over latent tokens such that it can more drastically change when the model parameters in p(z|x,c) are updated along with training. This can confuse the LM prior and slow down the training.
• Secondly, we notice that previous practices in the image, video, etcs. [1,2,3] usually choose to place the training for the autoencoder (constructing the latent space) before the prior training. This makes sense because the structure of the latent space is definite and fixed during the training of LM prior, guaranteeing that the training can converge in a stable way. A very long-standing example in joint training of generative models is the GAN[4], which has strong performance yet suffers from unstable training and tricky convergence.

We agree that the end-to-end joint training seems more attractive yet it may practically be difficult for training due to the interplay between different components, in specific the hyperparameter tuning of loss weights, learning schedule, etc. We are with the idea that training separately may result in suboptimal loss minima yet it can converge better in practice, and thus we follow prior works like [1,2,3] to derive the scheme. We kindly hope this can help explain the situation better!

[1] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017). [2] Ramesh, Aditya, et al. "Zero-shot text-to-image generation." International conference on machine learning. Pmlr, 2021. [3] Yan, Wilson, et al. "Videogpt: Video generation using vq-vae and transformers." arXiv preprint arXiv:2104.10157 (2021). [4] Goodfellow, Ian, et al. "Generative adversarial nets." Advances in neural information processing systems 27 (2014).

Response to Q3:

• We sincerely appreciate this biologically relevant suggestion. We agree with the idea that building an atom-level autoencoder including the side chains, and even other biomolecules such as ligands, RNA, etc. will make the generative model much more powerful in downstream applications. We have briefly mentioned this in the Section 6 (Conclusion and limitations) to potentially enlighten valuable future works!