Paper: Internal-Coordinate Density Modelling of Protein Structure: Covariance Matters (ICLR-2024)

Reviewer comment (Theory):
The math in this paper is sloppy and incomplete, and I’m not convinced that the method is correctly derived. I suspect that the Lagrangian is incorrectly derived, but I would be happy to see full derivations of the method to verify if it is. The results have some issues, which implies that something is perhaps wrong.

Author response:
* **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.

Reviewer comment (Writing):
I could not understand the generative model with VAEs and Unets, or the experiments. I’m not sure what the paper or the models try to demonstrate or achieve, despite the clear problem statement of finding safe distributions of angles wrt coordinates. I think the paper deviates from the original problem quite far into somewhere else. I can’t really interpret if the results are good or not.

Author response:
* **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.

Reviewer comment (Experiments):
The ML significance of the method is ultimately a bit light. There are only a few experiments, and no comparisons to other ML models. It seems that this paper is about finding variations in proteins. These are conformers, and there is already ML works on doing generative models for those. Comparisons to those are needed, and the paper should better positions itself to wider literature. Currently the paper takes too narrow and isolated perspective to the problem domain.

Author response:
* **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.

Reviewer comment (Writing):
How does an i’th angle relate to m’th atom i’th coordinate? How come you have the same i'th index for two different things (phi/phi and x/y/z)?

Author response:
Thank you for pointing out the inconsistency for the index variable ("$m$" to "$n$" in eq 5), we agree that this can be confusing.  They are the same entity. We apologize and have now updated the manuscript correspondingly (the index is now $m$ everywhere).

Reviewer comment (Theory):
I’ve lost track where eq 9 comes from. What is this? Aren't Cm user defined constants: why are we now solving for their values? Also aren’t you solving for the lambdas? Why are there none here? It seems that again quite a lot of intermediate derivations are missing, and I can’t verify that this is correct.

Author response:
* **Role of $C$**

  Your intuition is correct - you can think of it as choosing an expected fluctuation level $C_m$ for each atom in the chain. Rather than saying that these values are "user defined", we just phrase it in terms of "fitting" these fluctuations such that they best fit the the target density. In other words, the model fits the level of fluctuation for each individual atom such that it corresponds best to the distribution of the training data.

  One additional complication is that since the Lagrange formalism above is employed as part of the decoder in a variational autoencoder, all "predictions" we discussed above are conditioned on a latent variable $z$. Different choices of $z$ will correspond to different modes of the distribution we are trying to fit, and some will require large fluctuations than others. We thus have $C_m$ values for each atom $m$, but these are dependent on $z$.

Reviewer comment (Theory):
What does it mean that you predict lambdas with a network? Err… Lagrangian multipliers can’t be just chosen at will or predicted: the lagrangian formalism dictates under which coefficient values you can say something about your solution optimality. If you don’t follow the lagrangian approach rigorously, you have no guarantees of what you end up with, and it might be of little use, or not do what you think it does. Are you perhaps not even doing a lagrange method but instead perhaps doing a penalty method? Can you clarify the setting and the goals and the overall workflow here?

Author response:
* **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.

Reviewer comment (Writing):
I’m lost on what are you trying to achieve with the VAE, and what’s the point. So you take angles, encode them into latents, and reconstruct the angles. Why is this useful? What are we trying to achieve? If the reason you have a VAE is to evaluate derivatives, then something is wrong. Surely you should be able to just compute your derivatives directly?

Author response:
* **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.

Reviewer comment (Theory):
So the VAE reconstruction as fed to a Unet that outputs lagrange multipliers… This feels adhoc. Why is it useful to predict lagrange multipliers? Intuitively this sounds wrong: the lagrange multipliers need to be at equilibrium point under lagrange conditions (etc), and not just predicted by some neural network. Can you elaborate what are you trying to do, and what role the lagrange has here? How much rigor are you even aiming at?

Author response:
* **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.

Reviewer comment (Theory):
It seems that in the end you only produce a kind of soft regularisation by something that is inspired or slightly flavored with Lagrangians, but you are not actually doing Lagrangian optimisation. Given that the section 3 is named “constraints” I was under the impression that you are actually doing Lagrangian constraints, but I don’t think this is the case. Can you clarify what role the Lagrangian has?

Author response:
* **Lagrangian**

  We apologize for not making this clear. It is indeed just the method of Lagrange multipliers we are using. Actually, it is very similar to the approach under the header "Example 4 (Entropy)" on the wikipedia page you mention (https://en.wikipedia.org/wiki/Lagrange_multiplier). In our case, it is a density rather than a discrete distribution, and its a relative entropy (i.e. KL divergence) rather the entropy, but otherwise the idea is the same. We write up the lagrangian, differentiate it and set to zero to find the stationary point. This "MaxEnt" technique for finding maximally non-informative distributions under a given set of constraints was pioneered by ET Jaynes in a paper in 1957 [1] for the discrete case, and later generalized to the continuous case in the 1960ies. A brief introduction is given on this wikipedia page: https://en.wikipedia.org/wiki/Principle_of_maximum_entropy. 

  As you can see on the page above (see "Continuous case"), solutions turn out to take an exponential form. Since we recognize the quadratic form in our exponent as a Gaussian, we don't have to solve for the normalization constant. All that is left to do is determine the value of the lagrange multipler lambda in terms of the constraint $C_m$. Normally, these $C_m$ constraints would be "user-defined" as you say, and it would be appropriate to find the corresponding lambda to realize the constraint on $C_m$. However, in our case, our network predicts the allowed level of fluctuations per atom (see comment below). Since we have a direct relationship between $C_m$ and the lagrange multiplier (eq 9 in our paper) we can therefore directly predict lambdas instead of the $C_m$ values. Doing so implicitly sets the fluctuation level for each atom.

Reviewer comment (Writing):
What are local and global constraints? What is “valid” ramachandran distribution? What does “improved” 3D variance mean? What do you improve? What is “full covariance matrix”? Can you give a definition. Why is it full? How big is it? Covariance of what..? “that approximates distribution better…” Which distribution? Can you give a definition? What it the prior? Is it the eq 1? What is sigma_k,data? “underestimates the fluctuations” What fluctuations? Of whose? What is the “standard estimator”? What are you trying to achieve in the experiments? Can you clarify these points?

Author response:
* **"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").

Reviewer comment (Theory):
How does the VAE impose global constraints? What are global constraints?

Author response:
* **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.

Reviewer comment (Presentation):
Fig4: What is ramachandran reference? Where does the reference come from? What does it represent? What are the Rama samples? Where do they come from? Do these have something to do with data, or do they come from the model? What is MD reference? What does it represent or do?

Author response:
* **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.

Reviewer comment (Presentation):
Fig5: What does this figure show? Where do these energy landscapes come from? What did you learn here? Did you just learn to match the MD reference panel with the VAE? What are the axes of this panel? I’m pretty lost what you did in the experiments, or what are you trying to demonstrate. The experiments do not discuss constraints or varying the angles in coordinate-safe ways at all, so I’m wondering if this is the goal anymore. Can you help me understand?

Author response:
The reported results on fluctuations along the atom chain (Figure 4) and TICA landscapes (Figure 5) are the result of sampling by our VAE (see also appendix A for experimental details and appendix C.2 and C.3 for sample visualizations), which generates conformers based on the learned density. In the current version of the manuscript, we also compare the conformer-generating performance against several baselines, including a flow-based model [1] which is, to the best of our knowledge, state-of-the-art for internal-coordinate density modelling (see general response).

Reviewer comment (Theory):
“We obtain samples that are guaranteed to fulfil physical constraints..”. With the amount of \approx in the paper, this is very hard to believe. Do you really guarantee this? In what part of the paper is this guarantee made?

Author response:
* **“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.

Reviewer comment (Experiments):
The experiments do not include enough baselines to validate the practical usefulness of this method. Many related DL methods are mentioned in related work, but they are not compared in the experiments. Although the authors mentioned the main focus of this work is internal-coordinate modeling, I am still interested in a more thorough comparison.

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

Reviewer comment (Reproducibility):
Section 3.5: U-net is used to estimate the values for the Lagrange multipliers \lambda for each constraint. Is there any supervision or auxiliary loss on \lambda?

Author response:
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).

Reviewer comment (Experiments):
Restricting the fluctuation will restrict the representation power of the model. How to balance that except for hand-tuning the hyper-parameters?

Author response:
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.

Reviewer comment (Experiments):
Any analysis to compare the models learned on the unimodal scenario and multimodal scenario?

Author response:
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.

Reviewer comment (Experiments):
Can the model generalize to generate conformers for new proteins?

Author response:
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.