Internal-Coordinate Density Modelling of Protein Structure: Covariance Matters
Authors: Marloes Arts
Jes Frellsen
Wouter Boomsma
Venue: Accepted by TMLR ()

Abstract: After the recent ground-breaking advances in protein structure prediction, one of the remaining challenges in protein machine learning is to reliably predict distributions of structural states. Parametric models of fluctuations are difficult to fit due to complex covariance structures between degrees of freedom in the protein chain, often causing models to either violate local or global structural constraints. In this paper, we present a new strategy for modelling protein densities in internal coordinates, which uses constraints in 3D space to induce covariance structure between the internal degrees of freedom. We illustrate the potential of the procedure by constructing a variational autoencoder with full covariance output induced by the constraints implied by the conditional mean in 3D, and demonstrate that our approach makes it possible to scale density models of internal coordinates to full protein backbones in two settings: 1) a unimodal setting for proteins exhibiting small fluctuations and limited amounts of available data, and 2) a multimodal setting for larger conformational changes in a high data regime.

Published in Transactions on Machine Learning Research (05/2024)
Internal-Coordinate Density Modelling of Protein Structure: Covariance Matters
Marloes Arts
ma ku, dk
Deportment of Computer Science University of Copenhagen
Jes Frellsen
jefr dk
Department of Applied Mathematics and Computer Science
Technical University of Denmark
Wouter Boomsma
wb.Bii dk
Department of Computer Science University of Copenhagen
Reviewed on OpenReview: https: openreview. net/ forum? 90099X22Z22 9XR2tZRmEB
Abstract
After the recent ground-breaking advances in protein structure prediction, one of the remain- ing challenges in protein machine learning is to reliably predict distributions of structural states. Parametric models of fluctuations are difficult to fit due to complex covariance structures between degrecs of freedom in the protein chain, often causing models to either violate local or global structural constraints. In this paper, we present a new stralegy for modelling protein densities in internal coordinates, which uses constraints in 3D space to induce covariance structure between the internal degrees of freedom. We illustrate the poten- tial of the procedure by constructing variational autoencoder with full covariance output induced by the constraints implied by the conditional mean in 3D, and demonstrate that our approach makes it possible to scale demsity models of internal coordinates to full protein backbones in two settings: 1) a unimodal sctting for proteins exhibiting small fluctuations and limited amounts of available data, and 2) a multimodal setting for larger conformational changes in a high data regime.
1 Introduction
Protcins macro-molecules that are involved in nearly all cellular processes. Most proteins adopt compact 3D structure, also referred to the native state. This structure is rich source of knowledge about the protein, since it provides information about how the protein can engage biochemically with other proteins to conduct its function. The machine learning community has made spectacular progress in recent years on the prediction the native state from the amino acid sequence a protein (Jumper et al., 2021; Senior et al., 2020; Wu et al., 2022b; Back et al., 2021; Wu et al., 2022a). However, the static picture the structure of a protein is mislcading: in reality a protein is continuously moving, experiencing both thermal fluctuations and larger conformational changes, both of which affect. its function. One of the remaining challenges in machine learning for structural biology is to reliably predict these distributions of states, rather than just the most probable state. We discuss the state of the density modelling field in Section 5 (Related work).
Modelling the probability density of protein structure is non-trivial, due to the strong constraints imposed by the molecular topology. The specific challenges depend on the chosen structural representation: if a structure is represented by the 3D coordinates of all its atoms, these atom positions cannot be sampled independently without violating the physical constraints of, e.g., the bond lengths separating the atoms. In addition, an arbitrary decision must be made about how the structure placed in global coordinate system, which

---

Published in Transactions on Machine Learning Research (05/2024)
implics that operations done this representation should preferably be invariant or equivariant to this choice. An alternative is to parameterize the structure using internal coordinates, i.e. in terms of bond lengths, bond angles and dihedrals (rotations around the bonds). The advantage of this representation is that internal degrees of freedom can be sampled independently without violating the local bond constraints of the molecule. It also makes possible to reduce the number of degrees of freedom to be sampled for instance fixing the bond lengths ideal values, since they fluctuate much less than the torsion angles and bond angles.
For the reasons given above, an internal coordinate representation would appear to be attractive choice for density modelling. However, one important problem reduces the appeal: small fluctuations in internal coordinates will propagate down the chain, leading to large fluctuations remotely downstream in the protein. As consequence, internal-coordinntee density modelling necessitates careful modelling of the covariance structure between the degrees of freedom in order to ensure that small fluctuations in internal coordinates result in small perturbations of the 3D coordinates of the protein. Such covariance structures are typically highly complex, making direct cstimation difficult.
In this paper, investigate whether density modelling of full-size protein backbones in internal coordinates is feasible. We empirically demonstrate the difficulty in estimating the covariance structure of internal coordinates from data, and instead propose technique for inducing the covariance structure by imposing constraints on downstream atom movement using the Lagrange formalism. Rather than estimating ing the covariance structure from scratch, we can instead modulate the covariance structure by choosing appropriate values for allowed fluctuations downstream atoms (Fig. 1). We demonstrate the procedure in the context. of variational autoencoder.
K1 K2 K3
Constraints
(foed)
X
x2
N (x Ex)
K
Model
Figure A protein structure ensemble is modelled in internal coordinate space (dihedrals in green, bond angles in orange, and fixed bond lengths in purple), while imposing constraints atom fluctuations in Euclidean space. The resulting full covariance structure can be used to construct multivariate Gaussian distribution, from which samples can be drawn that reflect the local and global fluctuations of the protein.
Given a prior on the internal coordinate fluctuations and a predicted mean, we impose constraints on the atom fluctuations in 3D space to obtain a full covariance structure over the internal coordinates. We show that this allows us to generate valid structures in terms of both internal and Cartesian coordinates. Our regimes: low data regime for proteins that exhibit small, unimodal fluct uations,
method validated in and high data regime for proteins that exhibit multimodal behavior. We anticipate that this method could serve as building block applicable more generally for internalccordinatee density estimation, for instance intornal-coordinate denoising diffusion models.
Our main contributions are:
We formulate a procedure for inducing full protein backbone covariance structure in internal coordi- nates, based constrainis alom fluctuations in 3D space.
Rather than predicting a full covariance matrix over internal coordinates, our proposed method only requires to prodict one Lagrange multiplier for each atom, from which the full covariance matrix can be constructed. For M atoms with fixed bond lengths, this corresponds to reduction from (2 M - 5)2 to simply M predicted values.
2

---

Published in Transactions on Machine Learning Research (05/2024)
We design variational autoencoder which models fluctuations for full-length protein backbones in internal coordinates. Even though constraints are formulated in Fuclidean space, the model is not. dependent on a global reference frame (i.e. it is rotationally invariant).
We demonstrate that our model provides meaningful density estimates on ensemble data for proteins obtained from experiment and simulation, even when data is limited. For more complex densities and large amount of data, show that the expressiveness of the chosen base model plays a key role.
Scope. The focus of this paper is modelling distributions of protein structure states in internal coordinates, with the emphasis on the low data regime, We are thus concerned with thermodynamic ensembles, rather than the detailed dynamics that a molecule undergoes or generalization across different proteins. Dynamics could potentially be modelled on top of our approach, for instance by fitting a discrete Markov model to describe transitions bet ween states, and using our approach to model the thermal fluctuations within a state, but this is beyond the scope of the current work. The main contribution of this paper is technique for modclling internal coordinate covariance structure, which has applicability beyond the proof of principle presented here. Another perspective associated with static structure.
approach is that we wish to describe the aleatoric uncertainty
2 Background
2.1
Cartesian vs internal coordinates
As stated above, Cartesian coordinates and internal coordinates each have advantages and disadvantages. Assume we have a 3D protein structure in Euclidean space with atom positions æ. Throughout this paper, we only consider backbone atoms N, Ca and C, which gives a total number of atoms M = 3 x L, where L is the number of amino acids. The Euclidean setting thus results in 3 x M coordinates. Even though in this setting each of the atoms can fluctuate without affecting other atoms the backbone chain, there no guarantee for chemical integrity, i.e. conservation of bond lengths and respecting van der Waals forces. This can lead to local backbone crossings and generally unphysical protein structures.
One way to ensure chemical integrity is to paramcterize protein structure in internal coordinate space using dihedrals K1, bond angles K2 and bond lengths K3. Here, dihedrals are torsional angles that twist the protein around the bond between two consecutive atoms, bond angles are angles within the plane that is formed by two consecutive bonds, and bond lengths are the distances between two consecutive backbone atoms (left panel Fig. 1). Since bond length distributions have very little variance (Creighton, 1993), we choose to fix them, thereby reducing the number of variables over which we need to estimate the covariance. We will refer to the remaining two internal coordinates together as to avoid notation clutter. As dihedrals are defined by four points (the dihedral is the angle between the plane defined by the first three points and the plane defined by the last three points) and bond angles are defined by three points, the resulting protein structure representation will have (2 x M) - 5 coordinates. Not only does this result in less coordinates to determine a full covariance structure over, the coordinates are also automatically rotation and translation invariant, as opposed to Cartesian coordinates.
The remaining problem is that small changes in one internal coordinate have large consequences for the global structure of the protein, since all atoms downstream of the internal coordinate will move together, acting as rigid body. It is therefore challenging to preserve global structure while altering internal coordinates.
2.2 Standard precision estimators do not capture global fluctuations
Because of the limitations of internal coordinates mentioned in Scction 2.1, it is a highly non-trivial task to capture full covariance structure over k which also conforms to constraints in Euclidean space that are inherent to the protein, especially when the amount of data is limited. As an example, we use a standard estimator to get a precision matrix (i.c. the inverse of the covariance matrix) over & for a short molecular dynamics (MD) simulation of the B1 domain of protein G, based on PDB entry Ipga (Fig. 2). Details about the simulation can be found in Appendix A. We see that when take samples from a multivariate Gaussian
3

---

Published in Transactions on Machine Learning Research (05/2024)
with estimated mean and precision matrix, the samples exhibit atom fluctuations that are much higher than the original simulation, and with very different, pattern along the chain.
Much of the covariance structure a direct consequence of the internal coordinate representation. For instance, preserving a simple distance between two atoms in a chain requires correlated changes between all the dihedral angles separating them. However, given specific structure of the chain, these correlations are predictable, and we can therefore attempt to express them collectively in terms of much smaller number of parameters. Specifically, will consider the correlations in internal coordinates that are induced by uncorrelated fluctuations in the atomic coordinates in Euclidean space.
Varfance (mean over 3306]
10000
3.0
Standard estimator
K2
9100
2.5
2.0-
1.5
1.0
standard estimator Xx
K1
25
50
75
100
125
150
175
Atom position
M
Figure 2: When standard estimator used to gel the precision structure over internal coordinates, resulting atom fluctuations significantly deviate from MD simulations. Blue arrows and red helices represent secondary structural elements. The variance calculated as the mean of the variances over the x, y and z axis, in A
3
Internal-coordinate density modelling with constraints
3.1 Setup
We paramctorize 3D protein structure in terms of internal coordinates (i.c. dihodrals and bond angles, while bond lengths are kept fixed due their low variance), which logether will be referred to as k. Our aim is obtain a multivariate Gaussian distribution p(Ax) over the deviations Ax from the mean, centered zero, with full precision structure. This target distribution is subject constraints over atom fluctuations, enforcing the preservation of global structure. Our point of departure is to specify Gaussian prior q (ARk), which we will call the k-prior, over the internal coordinate distribution, where the mean is zero and the -2 precision is a diagonal matrix with the diagonal filled by the inverse variance over all k values o w,data" This prior sets the overall scale of fluctuations expected for the different degrees freedom. For convenience, we will set this equal to the empirical variance of the degrees of freedom observed in the data 1 The strength of the s-prior can be tuned using hyperparameter a. The k-prior is thus defined
(Ax) = exp (- 1 TA AR (-2
k.prior Ak
(1)
Zg
Σ-1 where 2, = xxdaaa 2TT is the normalization constant for the k-prior distribution and N..prior a diag(o dataa) is the precision matrix. Our approach will be to construct new distribution p which is as close possible to q, but which fulfills constraint that prohibits the downstream 3D coordinates from fluctuating excessively. We thus wish to minimize the Kullback-Leibles divergence between the objective distribution and x-prior:
p (A)
(KL(pla) = p(Ak)) In q( (Ax)
dAk.
subject constraints on the expected value over squared alom displacements of each ream atom:
E EAkxp(AR) A = | - = Cm
(3)
When used like this the x-prior should be thought base distribution, rather than Bayeslan prior, similar how the prior concept used variational autoencoder.
4

---

Published in Transactions on Machine Learning Research (05/2024)
where Arm is the squared displacement of atom m and Cn is tolerance value the expected value of this fluctuation (i.e. the variance of in Cartesian coordinates, assuming isotropic Gaussian). We will see later that by tuning the On values (blue circles in Fig. 1), we can modulate the covariance structure in internal coordinates.
3.2 Lagrange formalism to incorporate constraints
Employing Jaynes' maximum entropy principle (Jaynes, 1957), the Lagrange formalism to incorporate constraint for each of our M atoms, under the conditions that our probability density p(Ak) has zero mean and sums to one". This leads to the Lagrangian
D(plo) = |((Ak) In p(Ak) dAk + Ao (((A))AAK-1) (AA
M
(fseholAn)ddn Cm
Σ Am m==
Next, we take the functional derivative and set to zero: result (Jaynes, 1957; Kesavan & Kapur, 1989)8:
aD(p.g) 0p(Ax)
0, leading to the following well-cstablished
0 In p(Ak) +1 Ao + Σ Am Ar 7 p(Dn) Zp" (Avr)
1 ((Ax) exp Σ Am A dn
exp
(4)
8 D(p.g)
where Zp = exp(-1- - Xo) is the normalization constant of the target distribution. Note that since p(Ax] is positive, we know our solution will indeed be minimum.
3p(Ax))
3.3
First order approximation for atom fluctuations
In order use Eq. (4) to satisfy the imposed constraints, we need to express Ar in terms of As. Using first order Taylor expansion (small angle approximation), can express the 3D displacement vector A of each atom a linear transformation of the angular perturbation:
dzm Ak;
Amm ~
Σ
5
ōk;
where Im the position vector of the mhth atom, assuming that the location of atom m is downstream of the 4th internal coordinate in the direction in which the changes propagated through the chain. The partial derivatives @rm are known in closed form (Bottaro et al., 2012). From Eq. (5) it follows that the squared dx, displaccment distance can be approximated by
dæm
dæm AN; Ora
dam AN2) = Ar ART GmAKK
Az t
Σ
Orss
AKs
=
Σ
dzg
(6)
where G
0am dw.
Dem De,
is a symmetric and positive-definite matrix.
Substituting Eq. (6) and our s-prior expression from Eq. (1) into target distribution from Eq. (4) gives new Gaussian distribution:
1 (-AAn (Ex,prior Ex.conste) DK) =V(O0. Eix)
p(AR)=
Z exp
where E-1 k.cousstr 2 M Am Gm and is the new normalization constant.
Fvven though throughout this derivation we have included the normalization constant for completeness, in practice we work with unnormalized densities and normalize post hoc, since we recognize the final result to Gaussian.
Note that By (yIn s/q Aoly - 1) EM Am (Ar Cm)) In 2 +
5 M
m=1 Am Az?
7
5

---

Published in Transactions on Machine Learning Research (05/2024)
3.4 Satisfying the constraints
The final step in the constrained optimization is to establish the values for the Lagrange multipliers. closed form solution for this is not readily available, but using the findings from Section 3.3, we can now rewrite the constraints from Eq. (3) as
Cm = AA [A [A AmA Gm)
(8)
where the last simplification step comes from standard expectation calculus on a quadratic form (AR G Ar), where Åx has zero mean (Eq. 378 in Petersen et al., 2008). Although is non-trivial to express the Lagrange multipliers À in terms of atom fluctuations C, we thus see that it possible to evaluate C' given a set of Lagrange multipliers A. In the following, we will therefore construct our models such that our networks predict A, directly. The predicted set of Lagrange multipliers À will serve as a proxy for the atom fluctuations Cm (as demonstrated in Appendix C.1). We thereby predict the scale for the fluctuation levels allowed for all atoms in Cartesian coordinates, which in turn induces a full covariance structure over internal coordinates.
3.5
VAE pipeline
VAE model architecture. To demonstrate how our method works within a modelling context, we use & variational autoencoder (VAE), as illustrated in Fig. 3. We chose this model because of its simplicity, inspectable latent space, and stable training behavior, making modelling in low data regime feasible. The VAE has a simple linear encoder that takes internal coordinates (dihedrals and bond angles, bond lengths are kept fixed) as input and maps to latent space z, where have standard Gaussian a prior the latent space, which we call the z-prior to avoid confusion with the s-prior. The decoder outputs the mean over K, which is converted into Cartesian coordinates using PNeRF (AIQuraishi, 2018). This mean structure in 3D coordinates is used for two purposes. First, using the structure we can evaluate the partial dorivatives of atom positions with respect to the individual к as in Eq. (6). Second, the predicted mean structure is used to obtain a pairwise distance matrix d. This rotationally and translationally invariant representation reflects the steric constraints between non-covalently bonded atoms, and was therefore chosen as the input to U-Net (Ronneberger al., 2015), from which we estimate values for the Lagrange multipliers for each constraint. This allows the variational autoencoder, conditioned on the latent state 2, to modulate the allowed fluctuations. Implementation wise, the U-net is concluded with an average pooling operation that for each row-column combination computes one Lagrange multiplier X. Togother with our fixod-variance k-prior over & and hyperparameter a determining the strength of this s-prior, a new precision matrix is formed according to Eq. (7). We train a new VAE for each protein, where each trained model can generate new structures through simple ancestral sampling: first generating z from the standard normal z-prior, and subsequently sampling from a multivariate Gaussian distribution with the decoded mean and the constructed precision matrix. We refer to Appendix A for more details on data handling, model architecture, and other model settings.
Loss. Following standard practice we optimize the evidence lower bound (ELBO) using the Gaussian likelihood on the internal degrees of freedom as constructed above. While this loss encourages fit to the target к density, does not ensure that the predicted Lagrange multipliers are within the range within which our first order approximation of the fluctuations is valid. To ensure this, add the inverse of the L1 norm of the Lagrange multipliers as an auxiliary regularizing loss, which prevents the k-prior from dominating, By tuning the weight Waux on the auxiliary loss, we can influence the strength of the constraints. Our total loss C is thus defined as
+ DKL where DKL is the Kullback-Leibler divergence, and and & are the parameters of the encoder and decoder, respectively.
C= -Ego
|| p(z)) + Waux |X|-1,
6

---

Published in Transactions on Machine Learning Research (05/2024)
N(ODD
z
E-1 K.prlor
MLP
a O k,data
2
Uk
14x
E-1 k,constr
Σ
pNeRF
K1 K2
Pdd
U-Net
atoms
atoms
N Ux xx)
Figure 3: Model overview. The encoder (left) embeds internal coordinates into the latent space. The decoder (right) predicts mean, from which constraints extracted to obtain precision matrix. Together with the k-prior over the precision matrix based the input data, new precision matrix is formed which can be used to sample from multivariate Gaussian.
4 Experiments
4.1
Test cases
Unimodal setting in low data regime. We consider three test proteins for small fluctuations in a low data regime: lunc, 1fsd, and lpga. lune corresponds the solution structure of the human villin C-terminal headpiece subdomain. This protein contains 36 residues with 108 backbone (N, C& and C) atoms. This solution nuclear magnetic resonance (NMR) dataset is freely available from the Protein Data Bank and contains 25 conformers. 1fsd, beta beta alpha (BBA) motif, is also freely available NMR dataset. containing 41 structures. This system has 28 residues with 84 backbone atoms. 1pga, corresponding to B1 mumunoglobulin--indiin domain protein G, is a 56 amino acid long protein with 168 backbone atoms. We have a short in-house molecular dynamics (MD) simulation, which is 20ns long and structures were saved at a 50ps interval, resulting 400 structures for this protcin. Soe Appendix A for more details about the simulation.
Multimodal setting in high data regime. We also include two test cases for larger fluctuations following multimodal distribution in a high data regime. Both are known as "fast-folders" and the MD datasets are available upon request from Lindorff-Larsen et al. (2011). We refer the reader to this work for detailed descriptions the simulations. Chignolin (chn025) is a peptide with a hairpin structure, containing 10 residues and thus 30 backbone atoms. The simulation 106 μs long, saved a 200 ps interval, resulting 534.743 data points. The second test case, 2f4k, the chicken villin headpiece, with residues and 105 backbone atoms. The simulated trajectory is 125 us and also saved every 200 ps, yielding 629.907 structures.
4.2 Metrics
For the unimodal setting, we choose two simple measures of local and global structure, respectively. To evaluate local structure fluctuations, we show Ramachandran plots, well-known visualization tool in the context of protein structures, where & and dihedrals, which are the torsional angles around the N - Ca and C. C bonds, are plotted against each other. As a global measure to evaluate the overall 3D structural fluctuations, we report the variance over atom positions, averaged over three dimensions, across superposed (i.e. structurally aligned) samples.
For the multimodal sctting, we report free energy landscapes, parameterized by the first two components of time-lagged independent component analysis (TICA) (Molgedey & Schuster, 1994). Similar to e.g. PCA, TICA fits linear model to map high-dimensional input to lower-dimensional output, but TICA also
7

---

Published in Transactions on Machine Learning Research (05/2024)
incorporates the time axis. In this case, the high-dimensional representation consists of all dihedrals and pairwise distances within the protein backbone. The resulting components are ranked according to their capacity to explain the slowest, low-frequency, modes of motion. Taking the first two components corresponds to selecting reaction coordinates that underlic the slowest protein conformational changes, which is highly correlated with (un)folding behavior. This projection to a lower dimensional space can be considered a free energy landscape that reflects which conformational changes are favorable over others. We fit the TICA model on the time-ordered MD data, and pass samples from the VAE and baselines through the fitted model to create free energy landscapes.
Baselines. Apart from comparing the generated samples from our model to reference distributions from MD or NMR, we include four baselines. The first two baselines are VAE-based, using an encoder-decoder architecture identical to our own VAE, but without incorporating the meaningful 3D constraints. The first bascline, named "k-prior (fixed)" xM is a VAE trained to prodict fx given a fixed covariance matrix that is equal to p.ppior In other words, this is the same as our full VAE setup, but omitting the imposed 3D constraints. The second baseline is "k "k-prior (learned)", which corresponds to more standard VAE-setting where the decoder directly outputs a mean and a variance (i.c. diagonal covariance matrix). The third baseline does not involve a VAE, but samples structures from a multivariate Gaussian with mean based on the dataset and a precision matrix computed by a standard estimator. This is an emperical estimator for MD datasets, and an Oracle Approximating Shrinkage (OAS) estimator (Chen et al., 2010) for NMR datascts, since emperical estimators do not converge for such low amounts of samples. Finally, we include the flow-based model from Köhler el al. (2023) as fourth baseline, which, to the best of our knowledge, is the current state of the art in density modelling for internal coordinate representations.
4.3 Internal-coordinate density modelling results
Unimodal setting in low data regime. The unimodal, low data regime test cases exhibit small fluctuations around the native protein structure, where the largest fluctuations correspond to the loops connecting different secondary structure elements. Fig. 4 demonstrates that lunc, 1fsd and Ipga structures sampled from the VAF conform to local and global constraints, with valid Ramachandran distributions compared to the reference as well as improved atom position variance along the chain with respect to the baselines (see quantitative results in Table A2). Even in the extremely low data regime of 25 and 41 data points for lune and Ifsd, respectively (top two rows in Fig. 4), the VAE is able to estimate a full covariance matrix that approximates the distribution better than the baselines, especially in loop regions. Consequently, the baselines show higher tendency to sample unphysical structures with backbone crossings, even though the local sucture is preserved. These effects can also be observed in the 3D visualization of the sampled ensembles in Appendix C.2.
The third test case, 1pga (bottom row in Fig. 4), has a more complex structure with two ß-strands at the N-terminus forming a sheet together with two ß-strands from the C-terminus. These global constraints are not captured well by the baselines in this low data regime (400 data points), resulting in very high fluctuations in loop regions which violate the native structure (additional visualizations can be found in Appendix C.2). For our VAF, we see the benefits of imposing global constraints, resulting in much better density estimation compared to the basclines. Morcover, we can control the interplay between local and global constraints by adjusting the hyperparameters of our model, as exemplified in Appendix D.1. However, the complexity of this protein prevents perfect density estimation in a low data regime. Interestingly, we show in Appendix C.1 that the variance of the atom positions highly correlates to the imposed constraints calculated from a set. of predicted Lagrange multipliers using Eq. (7).
Multimodal setting in high data regime. Here, we explore the use of our approach for modelling more complex behavior in a high data regime. Fig. 5 shows the free energy landscape in terms of the two first TICA components for cln025 and 214k. When comparing the VAE and the baselines to the MD reference (see also quantitative results in Table A3), it is clear that the learned prior and standard estimator do not capture all modes in the free energy landscape. The blow- model performs best, suggesting that in this multimodal setting with plenty of available data, our proof-of-concept VAE is not as expressive as this state-of the-rr-rtte-tttt model. Moreover, the benefit of imposing 3D constraints on top of the fixed s-prior (see baseline) seems
8

---

Published in Transactions on Machine Learning Research (05/2024)
plats
Standard
VAE
VAS e-prise
Lunc backbone data points (KNR)
1.50
20
0.75
0.00
VMM
Vallance aloug the stom diain (A7
1fsd backbone data points (NNR)
130
8.00
0.00
0.00
2.00
0.00
www.m0
WWW.mnX
Tpp 168 backhore atoms
240
1:0 [
140
120
125
130
1.00
0.55
50
V00
VU0L
25
1.500
17.
Figure 4: Modelling fattoos in the unimodal setting for Ipga, 1sdd, and lunc. Left: structure visualization, with a-helices in red and ß-sheets as blue arrows. Middle: Ramachandran plots for the MD reference and VAE samples. Right: variance along the atom chain for VAE samples, MD reference, and baselines. Secondary structure elements are indicated along the x-axis.
Beference (MD)
k-peior (Iied)
prior (leurned)
Standurd estimator
Flow
CLN025
30 badibone atomn 081.269 data
R
105 backhone atoms 565.117 data points
Figure 5: Modelling (un) folding behavior in the multimodal setting for cln025 and 2f4k. Left: structure visualization. Right: TICA free energy landscapes for MD reference, VAE, and baselines.
beneficial for cln025, but the effect is as strong for the a-helical 2f4k, where local constraints might dominate (a similar effect can be seen when comparing the VAE to the fixed s-priir baseline for lunc in the unimodal setting). However, simple VAE setup is evidently able to model large conformational changes through its latent space, demonstrating how our general-purpose method for modeling fluctuations using 3D constraints can be incorporated into more expressive models to model complex behavior.
Similar to the unimodal case, there is a tradeoff between local and global constraints which we can modulate using hyperparamcters, as demonstrated Appendix D.2. In addition, we show in Appendix C.3 how distinct regions of the VAE latent space map different clusters in the TICA free energy landscape, and visualize the corresponding structures.
9

---

Published in Transactions on Machine Learning Research (05/2024)
5
Related work
There is a large body of work on models for analyzing trajectories of molecular dynamics simulations, either through Markov state models (Chodera & Noé, 2014; Singhal & Pande, 2005; Sarich et al., 2013; Schütte et al., 1999; Prinz et al., 2011), or more complex modelling stralegies (Mardt et al., 2018; Hernández el al., 2018; Sultan et al., 2018; Mardt et al., 2020; Xie et al., 2019). Typically, these focused on dimensionality reduced representations of the molccular structures, and are therefore not density models from which samples can be drawn. Similarly, interesting work has been conducted on generative models for contact maps, which are not directly realizable as 3D structures (Anand & Huang, 2018).
To our knowlodge, the first generative density model of full protein coordinates was the Boltzmann generator (Noé et al., 2019), normalizing flow over the Cartesian coordinates of protein ensembles. An extension of this approach was later used to estimate Ca-only coarse-grained force fields for molecular dynamics simulations (Köhler et al., 2023). This method, which uses internal coordinate imputs, demonstrated the ability of augmented flows to sample structural ensembles for proteins up to 35 amino acids A range of works have explored the use of latent variable models. One example is the IG-VAE (Eguchi et al., 2022), which generates structures in 3D coordinates but expresses the loss in terms of distances and internal coordinates to maintain SE(3) invariance. Similar approaches have been used to analyze cryo-EM data, where the task is to generate ensembles of structures given the observed cryo-EM image data. Since cryo-EM data provides information at slightly lower resolution than the full-atomic detail we discuss here, the output of these approaches are often density maps in 3D space (Zhong et al., 2021; Punjani & Fleet, 2021). One example of atomic-level modelling in this space is Rosenbaum et al. (2021), which decodes deterministically into 3D coordinates, but describes the variance in image space. Finally, diffusion models and low-matching have recently provided a promising new approach to density modelling, with impressive examples of density modelling at the scale of full-size proteins (Watson et al., 2022; Ingraham el al., 2022; Anand & Achim, 2022; Arts et al., 2023; Schreiner et al., 2023; Jing et al., 2024).
The primary objective in our paper is to investigate whether cotornal-coordiaate density modelling with a full covariance structure is feasible using a simple, parsimonious setup. Internal-coordinate probabilistic models of proteins have traditionally focused on protein local structure, i.e. correct modelling of angular distributions of the secondary structure elements in proteins. Early work was based on hidden Markov models of small fragments (Camproux et al., 1999; 2004; de Brevern al., 2000; Benros et al., 2006). The discrete nature of the fragments meant that these models did not constitute complete probabilistic model of the protein structure. Later models solved this issue by modelling local structure in internal coordinates, using different sequential models and angular distributions (Edgoose et al., 1998; Bystroff et al., 2000; Hamelryck et al., 2006; Boomsma et al., 2008; 2014; Thygesen et al., 2021). Due to the downstream effects of small internal-coordinate fluctuations, these models are not by themselves capable of modelling the distribution of entire protein structures, but they are useful as proposal distributions in Markov chain Moute Carlo (MCMC) simulations of proteins (Trbäck & Mohanty, 2006; Boomsma et al., 2013). Using deep learning architectures to model the sequential dependencies in the protein chain, recent work has pushed the maximum length of fragments that can be reliably modelled to stretches of 15 residues (Thygesen et al., 2021), where the fragment size is limited due to the challenges in estimating the necessary covariance structure.
Our work was inspired by methods used for constrained Gaussian updates in MCMC simulation, first introduced by Favrin et al. (2001), and later extended by Bottaro et al. (2012). Our method generalizes the approach to global updates of proteins, derives the relationship between the Lagrange multipliers and corresponding fluctuations in Euclidean space, and uses neural networks to govern the level of fluctuations in order to modulate the induced covariance structures.
Recent work has demonstrated that internal-coordinate modelling can also be done using diffusion models (Jing et al., 2022). So far this method has been demonstrated only on small molecules. We believe the method we introduce in this paper might help scale these diffusion approaches to full proteins. In Cartesian space, the Chroma model (Ingraham et al., 2022) demonstrated the benefits of correlated diffusion arising from simple constraints between atoms, e.g. to maintain the structural integrity of the chain. Our method builds on similar considerations in internal coordinate space, where the challenge is to ensure global integrity rather than local integrity.
10

---

Published in Transactions on Machine Learning Research (05/2024)
6 Discussion
Although protein structure prediction is now considered solved problem, fitting the density of structural ensembles remains an active field of research. Many recent activities in the field focus on diffusion models in the Cartesian coordinate representation of a protein. In this paper, we take a different approach and investigate how we can describe small-scale fluctuations in terms distribution over the internal degrees of freedom of protein. The main challenge in this context the complex covariance between different parts of the chain. Failing to model this properly results in models that produce disruptive changes to the global structure, even for fairly minor fluctuations in the internal coordinates. Instead of estimating the covariance matrix from data, we show that it can be induced by imposing constraints on the Cartesian fluctuations. In sense, this represents a natural compromise between internal and Cartesian coordinates: we obtain samples that are guaranteed to fulfill the physical constraints of the local protein topology (e.g. bond lengths, and bond angles), while at the same time producing meaningful fluctuations globally.
We implement the idea in the decoder of a variational autoencoder on two protein systems. This is primarily proof of concept, and this implementation has several limitations. First of all, the standard deviations in the s-prior of the internal degroes of freedom currently set hyperparameter. These could be estimated from data, either directly within the current VAE setup, using a preexisting model of protein local structure. Another limitation of the current model is that the produced fluctuations are generally too small to fully cover the individual modes of the target densitics, mostly duc to the first-order approximation we make for the atom fluctuations. This also comes inlo play in the multimodal setting, where our simple VAE falls short compared the more complex flow-based method when enough data available. Latent variable models, such as VAEs, can capture a complex distribution by decomposing it into continuous mixture of Gaussians. While our proof-of-concept VAE model has demonstrated its ability to capture multimodal distributions, it lacks the expressivity to fully cover all individual modes. This obstacle could potentially be overcome by constructing a hierarchical VAË, where samples are constructed as a multi-step process, similar to the generation process in diffusion models. In fact, we believe that our fundamental approach of induced covariance matrices could be a fruitful way to make diffusion models internal coordinates scale to larger systems, by allowing for larger non-disruptive steps. We leave these extensions for future work.
7 Code and data availability
The code base, NMR datasets and in-house generated MD data are available at this github repository: iance, matters. a
8 Acknowledgements
The work was supported by the Novo Nordisk Foundation through project grant NNF180C0052719 and through the Center for Basic Machine Learning Research in Life Science (MLLS, grant NNF200C0062606), We also acknowledge support from the Pioneer Centre for AI (DRNF grant number P1). Finally, we thank Tone Bengtsen for running the in house MD simulation for Ipga.
11

---

Published in Transactions on Machine Learning Research (05/2024)
References
Mohammed AlQuraishi. pnerf: Parallelized conversion from internal to cartesian coordinates. bioRziv, pp. 385450, 2018.
Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arx:22200. 15019, 2022.
Namrata Anand and Possu Huang. Generative modeling for protein structures. Advances in neural information processing systems, 31, 2018.
Marloes Arts, Victor Garcia Satorras, Chin-Was Wei Huang, Daniel Zuügner, Marco Federici, Ceeilia Clementi, Frank Noć, Robert Pinsler, and Rianne van den Berg. Two for onc: Diffusion models and force fields for coarse-grained molecular dynamics. Journal of Chemical Theory and Compulation, 19(18):6151 6159, 2023.
Minkyung Baek, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee, Jue Wang, Qian Cong, Lisa N Kinch, R Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using three-track neural network. Science, 873(6557) 2021.
C. Benros, A.G. de Brevern, Etchebest, and S. Hazout. Assessing novel approach for predicting local 3D protein structures from sequence. Proteins, 62:865 -8865-880, 2006.
W. Boomsma, K.V. Mardia, C.C. Taylor, J. Ferkinghoff-Borg. A. Krogh, and T. Hamelryck. A generative, probabilistic model of local protein structure. Proc Nall Acad Sci USA, 105(2)):8932 8937, 2008.
Wouter Boomsma, Jos Frellsen, Tim Harder, Sandro Bottaro, Kristoffer E Johansson, Pengfei Tian, Kasper Stovgaard, Christian Andreetta, Simon Olsson, Jan B Valentin, et al. Phaistos: A framework for markov chain monte carlo simulation and inference of protein structure. Journal of computational chemistry, 34 (19):1697-1705, 2013.
Wouter Boomsma, Pengfei Tian, Jes Frellsen, Jesper Ferkinghoff Borg, Thomas Hamelryck, Kresten Lindorff- Larsen, and Michele Vendruscolo. Equilibrium simulations of proteins using molecular fragment replacement and umr chemical shifts. Proceedings of the National Academy of Sciences, 111(3)):33522 13857, 2014.
Sandro Bottaro, Wouter Boomsma, Kristoffer E. Johansson, Christian Andrectta, Thomas Hamelryck, and Jesper Perkinghoff-Borg Subtle monte carlo updates in dense molecular systems. Journal of Chemical Theory and Computation, 8(2):695-705 702, 2012.
C. Bystroff, V. Thorsson, and D. Baker. HMMSTR: a hidden Markov model for local sequence-strucluue correlations in proteins. J Mol Biol, 301 2000.
AC Camproux, P. Tuffery, JP Chevrolat, JF Boisvicux, and S. Hazout. Hiddon Markov model approach for identifying the modular framework of the protein backbone. Protein Eng Des Sel, 12((2):1063-100 1999.
AC Camproux, R. Gautier, and P. Tufféry. A hidden Markov model dcrived structural alphabet for proteins. J Mol Biol, 339(3):5911666 605, 2004.
Yilun Chen, Ami Wiesel, Yonina C Eldar, and Alfred O Hero. Shrinkage algorithms for mmse covariance estimation. IEEE transactions on signal processing, 58(10):5016 5029, 2010.
John D Chodera and Frank Noé. Markov state models of biomolecular conformational dynamics. Current opinion in structural biology, 25:135 144, 2014.
Thomas E Creighton. Proteins: structures and molecular properties. Macmillan, 1993.
AG de Brevern, C. Etchebest, and S. Hazout. Bayesian probabilistic approach for predicting backbone structures terms of protein blocks. Proteins, 41(3):271 287, 2000.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle Beauchamp, Loo-Ping Wang, Andrew Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 13(7):
e1005659, 2017.
12

---

Published in Transactions on Machine Learning Research (05/2024)
T. Edgoose, L. Allison, and DL Dowe. An MML classification of protein structure that knows about angles and sequence. Pac Symp Biocompul, 3:585-500 596, 1998.
Raphael R Eguchi, Christian A Choe, and Po-Ssu Huang. Ig-vae: Generative modeling of protein structure by direct 3d coordinate generation. PLoS computational biology, 18(6))e1010271, 2022.
G. Favrin, A. Irbäck, and F. Sjunnesson. Monte Carlo update for chain molecules: Biased Gaussian steps in torsional space. Chem. Phys., 114:8154 8158, 2001.
T. Hamelryck, J.T. Kent, and A. Krogh. Sampling realistic protein conformations using local structural bias. PLoS Comput Biol, 2(9)x 2006.
Carlos X Hernández, Hannah K Wayment-Steele, Mohammad M Sultan, Brooke E. Husic, and Vijay S Pande. Variational encoding complex dynamics. Physical Review E, 97(6):062412, 2018.
Moritz Hoffmann, Martin Konrad Scherer, Tim Hempel, Andreas Mardt, Brian de Silva, Brooke Flena Husic, Stefan Klus, Hao Wu, J Nathan Kutz, Steven Brunton, and Frank Noé. Deeptime: python library for machine learning dynamical models from time series data. Machine Learning: Science and Technology, 2021.
John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent. Xue, Fritz Obermeyer, Andrew Beam, et al. Illminating protein space with a programmable generative model. bioRaie, pp. 2022 12, 2022.
A. Irbäck and S. Mohanty. Profasi: a monte carlo simulation package for protein folding and aggregation. J. Comput. Chem., 27(13):1548-1555 2006.
Edwin T Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusson for molecular conformer generation. arXiv preprint arXiv:2206.0 01729, 2022.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint 0rXiv:2402.44844 2024.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 589,999, 2021.
Hiremaglur K Kesavan and Jagat Narain Kapur. The generalized maximum entropy principle. IEEE Transactions on systems, Man, and Cybernetics, 19(5):1042 1052, 1989.
Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Compulation, 19 (3):942 952, 2023.
Kresten Lindorff Larsenn, Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science, (31(6655):517--20, 2011.
James A Maier, Carmenza Martines, Koushik Kasavajhala, Lauren Wickstrom, Kevin E. Hauser, and Carlos Simmerling. ff14sb: improving the accuracy of protein side chain and backbone paramcters from f199sb. Journal of chemical theory and compulation, 11(8):3696 3713, 2015.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets for deep learning of molecular kinetics. Nature communications, 9(1):1-11, 2018.
Andrcas Mardt, Luca Pasquali, Frank Noé, and Hao Wu. Deep learning markov and koopman models with physical constraints. In Mathematical and Scientiffc Machine Learning, pp. 451-475. PMLR, 2020.
Lutz Molgedey and Heinz Georg Schuster. Separation of a mixture of independent signals using time delayed correlations. Physical review letters, 72(23):3634, 1994.
13

---

Published in Transactions on Machine Learning Research (05/2024)
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457):eeaw) 2019.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
Jan-Hendrik Prinz, Hao Wu, Marco Sarich, Bettina Keller, Martin Senne, Martin Held, John D Chodera, Christof Schütte, and Frank Noć. Markov models of molccular kinetics: Generation and validation. The Journal of chemical physics, 134(17):174105, 2011.
Ali Punjani and David J Fleet. 3d flexible refinement: structure and motion of flexible proteins from cryo-em. BioRain, 2021.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pp. 234 241. Springer, 2015.
Dan Rosenbaum, Marta Garnelo, Michal Zielinski, Charlic Beattic, Ellen Clancy, Andrea Huber, Pushmeet Kohli, Andrew W Senior, John Jumper, Carl Doersch, et al. Inferring continuous distribution of atom coordinates from cryo-em images using vaes, arXiv preprint arXie:2106. 14108, 2021.
Marco Sarich, Ralf Banisch, Carsten Hartmann, and Christof Schütte. Markov state models for rare events in molecular dynamics. Entropy, 16(1):258 286, 2013.
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time- resolution surrogates for molecular dynamics. arXiv preprint arXiv:2305.18044 2023.
Schrödinger. The PyMOL Molecular Graphics System. version 2.5.2.
Ch Schütle, Alexander Fischer, Wilhelm Huisinga, and Peter Deufthard. A direct approach to conformational dynamics based on hybrid monte carlo. Journal of Computational Physics, 151(1):146-166 1999.
Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek, Alexander WR Nelson, Alex Bridgland, et al. Improved protein structure prediction using potentials from deep learning. Nature, 577(7792):706 710, 2020.
Nina Singhal and Vijay S Pande. Error analysis and efficient sampling in markovian state models for molecular dynamics. The Journal of chemicul physics, 123) 2005.
Mohammad M Sultan, Hannah K Wayment-Steele, and Vijay S Pande. Transferable neural networks for enhanced sampling of protein dynamics. Journal of chemical theory and computation, 14(4):1887-1894, 2018.
Christian B Thygesen, Christian Skjødt Steenmans, Ahmad Salim Al-Sibahi, Lys Sanz Moreta, An- ders Bundgård Sørensen, and Thomas Hamelryck. Efficient generative modelling of protein structure fragments using deep markov model. In International Conference on Machine Learning, pp. 10258-10267. PMLR, 2021.
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Androw J Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRriv, pp- 2022-12, 2022.
Kevin E. Wu, Kevin K Yang, Rianne van den Berg, James Y Zou, Alex X Lu, and Ava P Amini. Protein structure generation via folding diffusion. arXiv preprint ar X2X::2209.15611 2022a.
Ruidong Wu, Fan Ding, Rui Wang. Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, et al. High-resolution de novo structure prediction from primary sequence. BioRxiv, pp. 2022-00 07, 2022b.
14

---

Published in Transactions on Machine Learning Research (05/2024)
Tian Xie, Arthur France-Lanord, Yanming Wang, Yang Shao-Hona and Jeffrey C Grossman. Graph dynamical networks for unsupervised learning of atomic scale dynamics in materials. Noture communicationss, 10(1):
1 9, 2019.
Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis. Cryodrgn: reconstruction of heterogeneous cryo-em structures using neural networks. Nature methods, 18(2):176-185, 2021.
15

---

Published in Transactions on Machine Learning Research (05/2024)
A Experimental details
Model. See Fig. 3 for an overview of the model. The encoder and the decoder of the VAE are simple three-layer MLPs (multilayer perceptrons). Given a protein with M backbone atoms, the encoder takes (2 x .V - 5) x inputs, corresponding to (M - 3) dihedrals and (M-2) 2) bond angles which are fed as pairs of (sin, cos) imputs to avoid periodicity issues. Similarly, the decoder yields (2 x M - 5) x outputs, which can be converted
angles in the [-T, π| interval using the 2-argument arctangent (atan2). The MLP linear layer sizes of the encoder [128, 64, mapping to 16-dimensional latent space, and layer sizes of the decoder [32, 64, 128] (reverse of the encoder).
We use a standard U-Net (Ronneberger et al., 2015) to predict atom fluctuation constraints À from the mean predictions that the decoder outputs. The predicted mean for the internal degrees of freedom fe is translated into a mean structure Pæ using pNcRF (AlQuraishi, 2018), from which we calculate a pairwise distance matrix. This x M matrix with single channel serves as the input to the U-Net, which scales the number of channels up 1024 four steps before scaling back down one channel in four steps.
All datasets were split 90%-10% into training and validation set, with the same split for our VAE and all baselines. The best model is selected based on the validation loss. As we are handling very small datasets and will eventually evaluate per-atom fluctuations, it is not feasible to set aside a reliable test set that reflects the same distribution. Since we are in a generative modelling scenario rather than one where we are need of predictive accuracy, the main focus to fit a parametric model to a given data distribution. We therefore evaluate our generative performance based on the reference data, corresponding to the entire dataset, following related work (Köhler al., 2023; Arts et al., 2023; Schreiner al., 2023).
The weights for the k-prior and auxiliary loss explored with grid search (see Appendix D), values chosen for the models reported in the main paper are shown in Table A1 together with other experimental details. The model training starts with a warm-up phase in two different ways: 1) predicting Ps only, with Σ = I and 2) linearly increasing the weight of the KL-term from 0 to 1. Proteins in the low data regime (unimodal setting) have 100 epoch mean-only warm-up and a 200 epoch KL warm-up, while proteins in the high data regime (multimodal setting) have a 3 epoch mean-only warm-up and an 8 epoch KL warm-up. All models were trained using an Adam optimizer with learning rate of -4, on Nvidía Quadro RTX (48GB) GPU.
Final metrics calculated structures sampled from the model. For the evaluation in the unimodal setting, the number of samples was chosen to be equal to the total number of data points (25, 41 and 400 for lunc, Ipga and 1fsd, respectively) for all models. For the multimodal cases, 400.000 samples were drawn for TIC analysis. TICA was done using the Deeptime library (Hoffmann et al., 2021), using a lagtime of 100 and reducing the dign--dimesiionl input, consisting dihedrals and pairwise distances, to two dimensions. The TICA model is fit on the reference data (ordcred time), from which the resulting linear map stored and applied to sampled structures from the VAE and baselines. All structure visualizations were done using PyMOL (Schrödinger, version 2.5.2).
Table A1: Experimental details for test cases.
# train
# validation
# residues
# epochs
batch size
a
W Waux
lunc
23
2
36
1000
32
50
1fsd
37
4
28
1000
32
1
25
1pga
360
40
56
1000
32
50
25
cln025
481269
53474
10
50
64
25
50
2f4k
565117
62790
35
50
32
50
16

---

Published in Transactions on Machine Learning Research (05/2024)
Molecular dynamics details. The molecular dynamics simulation for lpga done in OpenMM (Fastman et al., 2017), using an Amber forcefield (Maier et al., 2015), water type TIP3P, box geometry "rhombic dodecahedron" and padding of 1 nm on each side of the solvated protein (i.e. 2 nm in total). The simulation is 20ns in total with a 50ps time lag, giving 400 structures. For MD details on cln025 and 2f4k we refer the reader to Linsorf-LLareen al. (2011).
17

---

Published in Transactions on Machine Learning Research (05/2024)
B Quantitative results
B.1 Unimodal setting, low data regime
Table A2: MSE (lower is better) to reference for atom fluctuations, unimodal setting, + standard
deviation over sampling runs.
K-prior (fixed)
s-prior
Standard
VAE
Flow
(learned)
estimator
lunc
0.020 + 0.005
3.268 + 0.791
0.015 + 0.003
4.219 + 0.990
142.297 = 24.191
fisd
0.496 + 0.051
16. + 3.049
10.339 + 1.675
9.583 + 1.954
112.031 + 16.897
1pga
0.039 + 0.002
3.464 + 0.200
1.550 + 0.112
1040.001 + 60.128
3070.602 = 109.194
B.2
Multimodal setting, high data regime
Table A3: Jensen--Shaee Shannon distance (lower is better) between binned Boltzmann distributions, i.e.
exp
free energy ka T
comparing VAE and baselines to the reference, multimodal setting. Mean + standard
deviation over sampling runs.
k-prior
K-prior (learned)
Standard
VAE
Flow
(fixed)
estimator
cln025
0.4555 + 0.0003
0.5384
0.0004
0.6060
0.0001
0.6861
0.0002
0.1931
0.0006
2f4k
0.3740 +: 0.0005
0.2967 + 0.0001
0.3418 + 0.0003
0.5179 + 0.0003
0.1833 0.0003
18

---

Published in Transactions on Machine Learning Research (05/2024)
C VAE sampling
C.1 Comparing constraints to atom fluctuations across samples
As derived Eq. (7), we can evaluate the constraint value Cn for each atom m given set of Lagrange multipliers. These constraints were placed the squared atom displacements, which is equivalent to the variance along the atom chain. Fig. Al demonstrates that the isotropic fluctuations of 400 Ipga samples drawn from the VAE are indeed quite close to C calculated from separately sampled sets of Lagrange multipliers. Since the constraints are placed on non-superposed (i.e. not structurally aligned) protein structures", this plot shows the variance along the atom chain for non-superposed structurcs.
Variance along the atom chain [
2.00
VAE
1.75
1.50
(non-superposed)
1.25
1.00
0.75
0.50
0.25
0.00
0
25
50
75
100
125
150
175
Atom position
Figure A1: Variance along the atom chain for non-superposed Ipga structures sampled with the VAE (orange) compared to constraints calculatod from predicted X-values (grey dashed). Secondary structure element locations indicated.
C.2
Visualization of sampled structures in the unimodal setting
Fig. A2 shows sampled superposed ensembles for our model and baselines, as well as the MD/NMB reference. This demonstrates that VAE samples, where global constraints were enforced, generally have globally consistent fluctuations compared to the reference data. In contrast, the baselines tend to exhibit fluctuations that are too large, which can lead to unphysical structures containing crossings and, in some cases, lacking secondary structure elements.
C.3 Latent space visualization in the multimodal setting
In this section, we visualize the VAE latent space in the multimodal setting (cln025) in Fig. A3. Moreover, we demonstrate how 100 random samples from latent space map to structure samples in the TICA free energy landscape, and show the 3D structures that correspond these samples. Transitions from the native state more unfolded conformations can be obscrved when going from the cluster the top right of TICA space towards the left. Depending on Σ, fluctuations around the means (which decoded from the latent space samples) vary in size. Therefore, means that are close together in terms of latent space location do not necessarily lead to sampling similar 3D structures. Moreover, used a UMAP to reduce the number of latent dimensions from 16 to 2, and this simplified representation might not capture the full complexity of the latent space. Nonetheless, is apparent that more unfolded struct ures largely originate from the rightmost cluster latent space.
built using (NNeF(AlQuraishi, 2018), which builds the chain step-by-step, thereby corre- sponding our post-rotational constraints.
'Sampled protein
19

---

Published in Transactions on Machine Learning Research (05/2024)
1pga
1fsd
Reference
VAE
x-prior (fixed)
k-prior (learned)
Standard estimator
Flow
Figure A2: Visualization of ensembles for reference data, the VAE model and baselines for 1pga, 1fsd and Junc. Number of samples is equal to the reference ensemble (400, 41 and 25 for lpga, 1fsd and lune, respectively.)
20

---

Published in Transactions on Machine Learning Research (05/2024)
space
Free energy KeT
UMAP1
T1C
Decode: Mx G Infer (m4. E) Sample
Unlabelled
UMAPO
TIC
Figure A3: Top left: UMAP reduction 2D of the originally 16-dimensional VAE latent space, with 100 samples shown in random shapes and colors. The grey scatterplot depicts the aggregated posterior, with the KDE of the aggregated posterior grey lines. Annotated green numbers correspond to boxes in the TICA free energy landscape (all structures corresponding to box are left unlabelled to avoid clutter). Top right: structure samples corresponding to latent space samples visualized in TICA space with the same symbols and colors as the latent space samples. Samples are grouped together green numbered boxes. Bottom row: 3D structures corresponding to the different numbered boxes in the TICA plot.
21

---

Published in Transactions on Machine Learning Research (05/2024)
D Ablation for hyperparameters
The two main hyperparameters that need to be chosen in the VAE setting are the strength of the k-prior and the weight of the regularizing loss Waux- These two weights can be set to prioritize local or global constraints in different ways. We demonstrate the effect on a unimodal case (protein G, 1pga) and a multimodal case (chignolin, cln025). In both cases, results are shown for a gridsearch over a = [1, 25, 50) and Waux = [1, 25, 50].
D.1 Unimodal
Fig. A4 shows results for the ablation on a and Waux in the unimodal setting. Increasing the strongth of the k-prior through a while keeping Waux constant corresponds to narrower distributions in the Ramachandran plot and bond angle distributions. higher weight Waux for constant a leads stronger global constraints, as demonstrated by the fluctuations along the atom chain.
D.2 Multimodal
To understand the impact of hyperparameters in the multimodal setting, we first consider the impact on samples drawn from the VAE that was trained with fixed k-prior, which depends on hyperparameter a. Fig. A5 illustrates how the distribution in the TIC free energy landscape changes when strengthening the prior. For = I, there is preference towards the metastable cluster the top left, while increasing the value of leads to stronger preference for the lowest energy cluster the top right.
When sampling from the VAF, where constraints are imposed on top of the k-prior, there is interplay between a and Waux, as shown in Fig. A6. Even though the exact trend is less clear here, the relative values of the hyperparameters have an observable influence on e.g. the width of the "bridge" between the lopmost two clusters, the size of the engryr- downward extrusion of the top left cluster, and the spread towards the less populated cluster the bottom right.
22

---

Published in Transactions on Machine Learning Research (05/2024)
Apgu, badibone Jt000s
angles
Reference
Standard estimator
Sucundary structure slements
VA
Max
a 1. Mow
Maue 50
2% Year
26 Mau 25
Woss
1.50
240
50 Wext 25
240
50 Masi 50
Figure A4: Ablation of and Wanx for protein (Ipga, structure shown at top left). From left to right: precision matrix example predicted by the VAE, Ramachandran plot, bond angle distributions, fluctuations along the atom chain (sccondary structure elements indicated, VAE æ-prior (fixed) out of scale for 1).
23

---

Published in Transactions on Machine Learning Research (05/2024)
MD reference, CLN025
a
25
a 50
Free Free energy Ka?
TIC1
TIC
Figure A5: Influence of hyperparameter a on samples drawn from the VAE with a fixed s-prior (without
imposing constraints) for chignolin (cIn025), visualized in TIC space.
MD reference, CLN025
1, Waw =1
1, Waux
a 1, Wauy
25, Waux
energy
= 25, 25
= 25, 50
50, Wmx =1
50, Wanz 25
50,
50
TIC 0
Figure A6: Hyperparameter ablation of and wanx VAE samples for chignolin (cIn025), visualized in TIC space.
24


Reviews:
[{"review_id": "YSqBVmske6", "reviewer": "Reviewer_Qxm9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. I appreciate the readability and self-containment of the paper. All terms and the context have been nicely explained without interrupting the flow of the method presentation.\n2.  The authors carefully presented the necessary steps and rationale behind incorporating the constraints on downstream atom movement by leveraging Lagrange formalism \n3. The experiments include relevant and very recent baselines for comparison \n\nWeaknesses: \n\n1.  Quantitative results can be extended and included in the main text\n2.  Seems like the flow baseline is outperforming the proposed method\n3.  In the current form, it is unclear what could be the exact benefits of the proposed models for downstream tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new way of modeling protein structure densities in internal coordinates, namely, bond lengths, bond angles and dihedrals by considering a full protein backbone covariance structure. Such an approach of capturing density over the internal degrees of freedom of a protein, provides guarantees for conforming with physical constraints with regards to protein topology. By computing one  Lagrange multiplier for each atom, the overall complexity is greatly reduced, making the procedure amenable for larger molecules than what has been explored in prior work.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Could you please include standard error or standard deviation in the results? From the tables in appendix B seems like flow is doing better than the proposed VAE approach, and maybe some margin of error could bring more clarification (maybe both methods perform on-par)\n- Is there any argument or setup where the VAE approach can be favorable compared to the flow-base one? For example is smaller data sizes, or maybe comparing the number of parameters and training time for the two models will highlight some advantages of the proposed method. This will make the paper much more convincing.\n- In the introduction, you mention: \"Another perspective on our approach is that we wish to describe the aleatoric uncertainty associated with a static structure.\" I find this quite interesting and I was hoping you could expand more on it if not possible to include some experimental results. This might open more applications to your method and make it a more valuable contribution.\n- The protein sizes are still on the smaller side, 36, 28 and 56 amino acids long for the unimodal setting and 10 amino acids for the multimodal. Since the implementation seems light (mention linear VAE, and the Lagrange trick) could you maybe explore larger molecules? This could also be of advantage when comparing to the flow and other methods from literature that explored only small molecule setups.\n- You mention fixing the bond lengths to ideal values, since they fluctuate much less. Were there any experimental validation of this? I understand running ablation studies take too long for the rebuttal period, but I was wondering if some essential modeling information might be lost, or if the authors have some intuition they can share regarding this choice.", "broader_impact_concerns": "No concerns on broader impact.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "s4mOiOVLai", "reviewer": "Reviewer_wNBB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Uses internal coordinates instead of Cartesian ones for degree of freedom reduction.    \n2. Uses Lagrange formalism to add constraints to assist better covariance matrix prediction.    \n3. Achieves SOTA performance on low data regime test cases.\n    \n\nWeakness:\n1. The model seems to take the MD or NMR data for a specific protein as training data and then predict the conformation of the same protein. Therefore, for a protein without reliable NMR of MD data, it cannot be generalized.\n2. The model has been only tested on proteins up to 56 amino acids, which is much smaller than most functional proteins.\n3. According to Appendix A, the datasets are split into training and validation set in a 90%:10% ratio. The sampled results are compared with the entire dataset, which is used to train and valid the model.  It seems there is no test set, and the results may have a data leakage problem.\n4. The proposed model is worse than a flow-based model in large data regime test cases.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a new protein backbone conformation generation approach. It models internal coordinates instead of Cartesian coordinates to reduce the degree of freedom. Using internal coordinates imposes a challenge of large global structural changes caused by small internal coordinate deviation. To solve this problem, the authors aim to obtain a multivariate Gaussian distribution over the deviations of internal coordinates from the mean and adding constraints from the Cartesian atomic positions using Lagrange formalism. The proposed model uses a VAE framework to predict the mean internal coordinates using target-specific training data. The predicted mean values of internal coordinates are used to construct protein structure in the Cartesian space. Then the pair-wise distance matrix is fed into a U-net architecture to predict the Lagrange multipliers $\\lambda$ of the constraints. The final covariance metric is constructed using internal coordinate covariance from training data and the constraint covariance from predicted $\\lambda$. Sampling from the newly constructed gaussian distribution, the model is able to better fitting the distribution of protein conformation in low data regime comparing to baseline models. But in large data regime, the model is worse than a flow based approach.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the following questions:\n\n1. The test cases only contains relatively small proteins. What is the performance on larger proteins?\n2. What is the formula of the loss?\n3. How are internal coordinates (dihedral and bond angles) embedded when feeding into VAE?\n4. What is the dimension of model input, latent space, and output?\n5. How is the training/validation/test data sets split? Have you try to split them based on conformational diversity to assess if the model can generate unseen conformers in the training data? This is very important, as we want the model to generate poses in the unexplored conformational space of a protein.\n6. How are the results generated by the flow-based model? Is it trained on the same datasets with the same splits?\n7. How is free energy calculated in  large data regime? Is it free energy or potential energy?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0TQQQAPUHK", "reviewer": "Reviewer_A967", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- I enjoyed reading this manuscript as it is well organised and well written. I have given a few suggestions on how to slightly improve the flow in the `Requested Changes` section. \n\n\n- The introduce method is well motivated and rigorous.\n\n- The claims are backed by some relatively convincing experimental evidence.\n\n### Weaknesses\n- Perhaps the lack of closed form formula (or approximate numerical method) to get the Lagrangian multipliers $\\lambda$ given the variance constraints $C_m$ is a limitation. Is there any advantage in predicting $\\lambda$ from $\\mu_d$ (via $\\mu_x$)?\n\n- If I understand correctly the current method is unable to model multimodal distribution over conformation, which albeit not so common, can be limiting as some proteins exhibit such behaviour (and other methods can handle multimodality).\n\n- As the core topic is about enforcing some covariance structure on the generated backbones, I feel that there is perhaps the discussion over the choice covariance is a bit short. I suppose that the choice of a diagonal covariance matrix (defined by the $C_m$) is for tractability reasons. May still be worth discussing this a bit more, e.g. relating to covariances in Chroma?\n- Perhaps the strongest criticism would be regarding the choice of using a VAE and the accompanying experiments. A VAE can be a very relevant deep generative model if the main goal is to learn latent representations, otherwise the field has demonstrated that (continuous) flows based generative models (e.g. diffusion) are better suited for raw capacity of modelling complex distributions such as protein backbone. The experiments are restricted to few and relatively short proteins, so perhaps in this setting a VAE has a stronger prior and thus is better suited for such low data regime?\nTypical metrics to evaluate the model fit via the quality and diversity of the generated proteins are also missing.\n\n### Summary\nIn a summary, the proposed approach is quite likely far from state of the art for generating realistic protein conformations, but the introduced idea is well motivated, rigorous and interesting, thus of interest for the community.\nAs briefly mentioned by the authors, this idea could be used for constructing diffusion/flow models on intrinsic coordinates with a predefined covariance structure, e.g. by building on the following work: 'EigenFold: Generative Protein Structure Prediction with Diffusion Models', Jing et al. 2023 and 'Harmonic Self-Conditioned Flow Matching for Multi-Ligand Docking and Binding Site Design', Stärk et al. 2023.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, the authors tackle the task of building a deep generative model for protein conformations, relying on their intrinsic coordinates representation.\nTo alleviate the stability issue that comes with intrinsic coordinates---small changes in intrinsic coordinates cause large changes in the Cartesian coordinates, especially far away in the chain---they introduce some covariance structure in the intrinsic coordinates via atomic variance constraints.\nThey do so by defining a density which is given by the product of a prior density and a term involving the atomic variance constraints following a Lagrangian formalism.\nThen taking a first order Taylor expansion, they get a new Gaussian distribution over the intrinsic coordinates (displacement to the mean) where the precision matrix is the sum of the prior precision and a term which involve Gram matrix of the Jacobian of the mapping from intrinsic to extrinsic coordinates.\nThen they showcase this method for introducing backbone covariance structure in intrinsic coordinates within a variational auto-encoder framework.\nThey empirically demonstrate the ability of this model to generate samples that not only respect local structures but also the global topology of typical conformations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Minor comments:\n- Sec 2.1: would be useful to have an illustrative figure of a protein backbone with the dihedral/torsion, bond angle and bond length. Figure 2 (left) does partially fulfil this. The dihedral angles would in particular be quite easier to understand from a Figure than from english. Perhaps can simply update Fig 2?\n- Sec 2.2: I am not sure to fully understand this bit. If the Gaussian samples have the same covariance as the 'true' samples covariance (via MD), how come they have different variance? Is this due to the difference between the covariance/precision matrix estimator and the true underlying covariance/precision matrix?\n- Sec 3.1:\n    - It's typically pretty reasonable to assume bond length to be fixed, but perhaps worth motivating this (in a few words), e.g. saying that the variance is very low over some dataset of interest?\n    - Eq 2: Perhaps refer to Fig 1 which has a useful illustration of these $C_m$ variance constraints? Also where does one get these $C_m$ values from?\n- Sec 3.3:\n    - Atoms ordering is arbitrary right (as in could flip the indices)? Does this ordering affect the construction of the covariance? I guess not since $G_m$ and is symmetric and the sum over $m$ is flip invariant (since permutation invariant). \n    - It would be worth saying a word on these partial derivatives $\\frac{\\partial x_m}{\\partial \\kappa_i}$, which I assume are computed automatically by backprop through the extrinsic to intrinsic change of paramterization? Or do we know these in close form? Could there be some numerical instabilities arising from these?\n- Sec 3.5: Worth moving Fig 3 to be on top or on the side of the first paragraph?\n- Figure 4: Would suggest setting the label fontsize according to the TMLR latex template fontsize, currently it's a bit small, especially the dihedral angles in the Ramachandran plots.\n- Sec 4.3:\n    - `\\cref` 'Fig. A2' broken, as in send to Fig 2 when clicking on the link.\n    - Is the same model shared for different proteins? or a new one is trained with a different number of parameters etc? I am not sure that there is a sentence on this.\n    - Any idea why the flow matching model performs so poorly in the low data regime compared to the VAE architectures? Is it modelling the intrinsic or cartesian coordinates? Is the prior/base distribution the same?\n    - Figure 5: I would suggest adding a quantitative metric on this plot that summarize models' performance, which would complement the current Fig.\n- Sec 5:\n    - `IG-VAE` missing reference? Eguchi et al. 2020.\n    - There are some prior work based on 'parametric family densities' e.g. 'A generative, probabilistic model of local protein structure' Boomsma et al. 2008, which is a line of work that might be useful to mention (or not).\n    - There are also earlier 'deep generative models' tackling protein conformation e.g. 'Generative modeling for protein structures', Anand and Huang 2018.\n    - There are been quite a lot of recent work based on diffusion / flow matching, some of which are cited, but being exhaustive would be indeed likely too verbose.\n    - I may have organised the Related Work section differently. Perhaps first a paragraph on models over cartesian representation, then intrinsic and finally on incorporating covariance. The discussion on this last point a bit short, especially given it's the core topic of this manuscript. It would be of interest, here or elsewhere in the paper to discuss the different choices of covariance structure used in Chroma and how they relate to this work for instance.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The general consensus of the reviewers is in the direction of acceptance. Predicting the distribution of states of a protein is challenging because, traditionally, experimental protein structure determination requires crystallization. Therefore, comparison to ground truth experimental evidence is not readily feasible. The comparison to molecular dynamics simulations is a reasonable surrogate for experimental evidence at this stage of the state of experimental capabilities. In summary, paper makes claims that rest upon the evidence provided and there is expected to be an audience for this work within the machine learning community.
Decision: Accept (poster)