Internal-Coordinate Density Modelling of Protein Structure: Covariance Matters

This is our continued response to reviewer U1Uk.

• "Full" covariance matrix

  The full covariance matrix refers to the covariance over all dihedrals and bond angles (named $\boldsymbol{\kappa}$ in the paper), as opposed to a more standard setting of only considering the diagonal (i.e. the variance). For $M$ atoms, this corresponds to a matrix with $(2 \times M - 5)^2$ entries. The reason we emphasize the "fullness" of the covariance matrix is that we can infer it from predicting only M Lagrange multipliers.

• Prior

  We agree that the term "prior" can be a bit confusing, since there is also a prior for the latent space of the VAE. We now clearly separate this "$z$-prior" from the "$\kappa$-prior", where the latter is our prior belief over the distribution of $\Delta \kappa$, based on the variance over these internal coordinates in the training set. The reciprocal of this variance, makes up the diagonal of our prior precision matrix. This is shown in Figure 3 and explained in Section 3.1 (Setup) and Section 3.5 ("VAE model architecture").

• TICA explanation

  We have now elaborated our explanation for the TICA plots, please see general response.

• Density modelling: an open problem?

  Density modelling is an established and very active field of research, e.g. [1, 2, 3, 4] (as also mentioned in response to reviewer jsg1), and is not considered a solved problem.

• “We obtain samples that are guaranteed to fulfil physical constraints of the protein topology (e.g. bond lengths, and bond angles)”

  What we are saying in this specific sentence is that, by construction, the sampled structures will be guaranteed to adhere to a correct local protein topology, purely by choosing to model the protein in internal coordinate space. This means e.g. valid bond lengths and no local clashes. Chemical integrity is one of the main benefits of working with internal coordinate representations, as was also outlined in the second paragraph of Section 2.1. The approximations made in the paper only relate to global constraints. We have added the word ``local’’ to this sentence to make the distinction more clear.

References

1) Noé, Frank, et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning." Science 365.6457 (2019): eaaw1147.

2) Zhong, Ellen D., et al. "CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks." Nature methods 18.2 (2021): 176-185.

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

4) 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.

We thank reviewer U1Uk for their interest in our paper and their in-depth review. We have attempted to pool some of the concerns and questions together below and we hope to have addressed all points, but please let us know if something is missing. We look forward to further discussions.

• Central goal of the method and role of the VAE

  The central problem in this paper is density modelling, i.e. modelling protein fluctuations. We do density modelling in terms of internal coordinates, while imposing constraints in 3D space. These 3D constraints are global constraints which determine the overall shape of the protein, as opposed to local constraints which ensure e.g. valid bond lengths and no local clashes. We show how we can infer a covariance matrix over internal coordinates from a single mean structure using the Lagrange formalism. Subsequently, we train a reconstruction VAE on a collection of protein structures, also called a protein ensemble. A protein ensemble can be a collection of structures that exhibit small fluctuations (unimodal setting) or larger conformational changes (multimodal setting). The VAE decoder outputs a mean over internal coordinates, which is used to infer the covariance matrix corresponding to that mean, and together they parameterize a multivariate Gaussian distribution, from which we can draw samples. Through encoding and decoding angles (reconstruction), the decoder learns to map from $z$ to valid internal coordinate means, such that in generative mode, we can simply draw samples from the prior over $z$ and get valid output distributions. Through the latent variable, the VAE can capture multiple conformational modes by predicting different means and corresponding covariance matrices.

• Baselines

  We agree that the results are more meaningful with additional baselines, please see general response. We have also elaborated the explanation of the baselines to be more understandable.

• Predicting Lagrange multipliers

  First of all, we would like to emphasize that in our setting the constraints, $C$, are not user-defined, but are to be inferred by the model based on the input data. Having said that, we absolutely concur that in an ideal setting, one would solve for the Lagrange multipliers given the constraints, whether they are user-defined or predicted by a model. However, for the constraints we want to impose we cannot solve for $\lambda$ in closed form. Therefore we opted for predicting the Lagrange multipliers directly, which is essentially an indirect way of predicting $C$ (in other words, the lambdas are a proxy for $C$). To illustrate this, we have included a new appendix C.1 (see also general response), which shows the correspondence between constraints $C$ (calculated using Eq. 9) and atom fluctuations across VAE samples.

  Moreover, one of the newly added baselines, "$\kappa$-prior (learned)", is a VAE trained to predict the variance directly, without including our constraints. This baseline does much worse compared to the full setup with constraints, which indicates that the constraints we incorporate are meaningful.

• Mathematical derivations

  We apologize for any lack of clarity and rigor regarding the mathematical derivations. As described in the general response, we have now elaborated the relevant sections (3.1-3.4) for improved understanding. We have incorporated all suggestions to the best of our ability, and would be happy to make further clarifications if necessary.

• Ramachandran distributions

  Ramachandran plots are a common way to visualize dihedral distributions by plotting the torsional angles around the $N-C_\alpha$ bonds ($\phi$) and $C_\alpha-C$ bonds ($\psi$) against each other (See section 4.2, "Unimodal"). The resulting distribution shows clusters for different secondary structural elements, e.g. beta sheets, right-handed helices, and left-handed helices. In our application, we can check qualitatively if the distribution is valid by comparing to the reference distribution (either simulation-based: MD, or experimental: NMR). We have updated Figure 4 to have more clear labels concerning the Ramachandran distributions.

We sincerely thank reviewer jsg1 for their interesting questions and helpful suggestions. We will answer in detail below.

• Density modelling (ML models trained per system)

  In this paper, we are training one model per protein using simulation data or experimental data to model the density over internal coordinates. Density modelling is an established and active field of research, e.g. [1, 2, 3, 4], with direct and indirect applications. Capturing a distribution from which we can sample cheaply (which is enhanced by the fact we model the backbone only) is useful for data augmentation, especially in a low data setting. Moreover, density modelling per protein can be seen as a step towards modelling complexes, doing cheap simulations, analyzing function and, ultimately, generalize across proteins.

• TIC plots and quantitative metrics

  We thank the reviewer for the great suggestions, and refer to the general response for our answer.

• How could this method help us understand protein fluctuations?

  This is a very interesting point, and we would argue that our method works in the opposite direction. We incorporate our knowledge of protein fluctuations as an inductive bias in the proposed method, thereby creating a more meaningful way to incorporate constraints.

• Regularizer

  We agree that our auxiliary loss corresponds to multiplying with a Laplace prior (L1 loss). We chose this auxiliary loss because our approximation only holds for smaller fluctuations. Without the regularizer,the global constraints can become too weak, and the covariance matrix can become invalid. The hyperparameters $a$ and $w_{\mathrm{aux}}$ can be used for tuning (see appendix D).

• Figure 4: capturing the actual variance

  We agree that our method does not always capture the true variance of the data. We believe this to be partly due to hyperparameter choice, the first-order approximation for fluctuations (Eq. 7) and, indeed, the Gaussian assumption. However, we stress that modelling a complex protein like 1pga in a limited data regime is highly challenging, and it might be that our mean encoder and decoder are not expressive enough in this setting. One possible solution could be introducing a hierarchical VAE to model larger local fluctuations, but this is beyond the scope of the current paper.

• Proteins with large degrees of freedom

  We believe our multimodal test cases to exhibit large degrees of freedom, as folding and unfolding involves dramatic conformational changes. However, we would like to emphasize that our general-purpose method is designed for improved modelling of smaller fluctuations, and that large degrees of freedom can be realized through model architecture choices. For example, our proof-of-concept VAE realizes large conformational changes through different predicted means (and corresponding covariance matrices).

References

1) Noé, Frank, et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning." Science 365.6457 (2019): eaaw1147.

2) Zhong, Ellen D., et al. "CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks." Nature methods 18.2 (2021): 176-185.

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

4) 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.

First of all, we would like to thank reviewer QUDF for their insightful comments, which we will address point-by-point below.

• Baselines:

  We agree that adding more baselines is very valuable. Please refer to our general response for the newly included baselines.

• Auxiliary loss on $\lambda$

  It is indeed important to have a regularizer on $\lambda$ values, which we employ in the form of an MAE auxiliary loss (see Section 3.5).

• Balancing the representation power of the model

  In the current setup, balancing the representation power of the model is mostly done through hyperparameter tuning (see also appendix D). Through the latent variable model (VAE), we gain some flexibility by predicting different means (and corresponding covariance matrices). However, in a low data regime and especially for more complex proteins like 1pga, it is still non-trivial to do perfect density modelling regardless of the chosen architecture.

• Unimodal scenario vs multimodal scenario

  We stress that our main contribution is to model small fluctuations correctly, which is also where we see the most distinct difference between our method and the baselines. We included the multimodal scenario to show how our general-purpose method can potentially be used to model larger conformational changes, which is mostly realized through the latent variable mapping to different means.

• Generalization across proteins

  Generalization to new systems is currently beyond the scope of the paper, as this would require a large and diverse dataset, which we don’t have at our disposal. However, we suggest that our method for incorporating constraints in internal-coordinate modelling could facilitate the next steps towards generalization when incorporated in a highly expressive model that is trained on sufficient amounts of data.

To all reviewers, ACs, and readers,

We thank the reviewers for their comments, which we have meticulously addressed in the current version of the manuscript, as outlined in the general response above. The main new contributions feature: 1) more baselines, including a flow-based model that is state-of-the-art in internal-coordinate density modelling, 2) a new test case (protein 1fsd, NMR dataset), 3) quantitative results, and 4) extended explanations for mathematical derivations and TIC analysis.

We highly value your feedback, and would be happy to engage in further discussions and provide clarifications where necessary.