L EARNING PROTEIN STRUCTURE WITH A DIFFERENTIABLE SIMULATOR
Anonymous authors
Paper under double-blind review
A BSTRACT
The Boltzmann distribution is a natural model for many systems, from brains to materials and biomolecules, but is often of limited utility for ﬁtting data because
Monte Carlo algorithms are unable simulate it in available time.
This gap between the expressive capabilities and sampling practicalities of energy-based models is exempliﬁed by the protein folding problem, since energy landscapes underlie contemporary knowledge of protein biophysics but computer simulations are often unable to fold all but the smallest proteins from ﬁrst-principles.
In this work we aim to bridge the gap between the expressive capacity of energy functions and the practical capabilities of their simulators by using an unrolled Monte Carlo simulation as a model for data.
We compose a neural energy function with a novel and efﬁcient simulator based on Langevin dynamics to build an end-toend-differentiable model of atomic protein structure given amino acid sequence information.
We introduce techniques for stabilizing backpropagation under long roll-outs and demonstrate the model’s capacity to make multimodal predictions and to (sometimes) generalize to unobserved protein fold types when trained on a large corpus of protein structures.
I NTRODUCTION
Many natural systems, such as ants in a colony, cells in a tissue, or atoms in a protein, organize into complex structures from simple underlying interactions.
Explaining and predicting how macroscopic structures such as these arise from simple interactions is a major goal of science and, increasingly, machine learning.
The Boltzmann distribution is a foundational model for relating local interactions to global system behaviors, but can be challenging to ﬁt to data.
Given an energy function U ✓ [ x ] , the probability of a system conﬁguration x scales exponentially with energy as
[Equation]
Importantly, simple energy functions U ✓ [ x ] consisting of weak interactions can collectively encode complex system behavior, such as the structures of materials, macromolecules, tissues, or when endowed with latent variables, images, sound, and text (Ackley et al., 1985; Salakhutdinov &
Larochelle, 2010).
Unfortunately, generating samples x ⇠ p ✓ ( x ) and learning model parameters ˆ ✓ with the Boltzmann distribution can require Monte Carlo simulations that frequently are unable to produce a useful result in available time.
These difﬁculties have driven a shift towards generative models that are easier to learn and sample from, such as directed latent variable models (VAEs) and autoregressive models (Goodfellow et al., 2016).
The protein folding problem provides a prime example of both the power of energy-based models at describing complex relationships in data as well as the challenge of generating samples from them.
Decades of research in biochemistry and biophysics support an energy landscape theory of protein folding (Dill et al., 2017), in which the folds that natural protein sequences adopt are the minima of a free energy surface.
Without the availability of external information such as coevolutionary information or homologous structures to constrain the energy function, however, contemporary simulations are challenged to generate globally favorable low-energy structures in available time.
[Figure]
How can we get the representational beneﬁts of energy-based models with the sampling efﬁciency of directed models?
Here we explore a potential solution of directly training an unrolled simulator of an energy function as a model for data.
By directly training the sampling process, we eschew the question ‘when has the simulator converged’ and instead demand that it produce a useful answer in a ﬁxed amount of time.
Leveraging this idea, we construct an end-to-end differentiable model of protein structure that is trained by backpropagtion through folding (Figure 1).
NEMO (Neural energy modeling and optimization) can learn at scale to generate 3D protein structures consisting of hundreds of points directly from sequence information.
Our main contributions are:
• Neural energy simulator model for protein structure that composes a deep energy function, unrolled Langevin dynamics, and an atomic imputation network for an end-to-end differentiable model of protein structure given sequence information •
Efﬁcient sampling algorithm that is based on a transform integrator for efﬁcient sampling in transformed coordinate systems • Stabilization techniques for long roll-outs of simulators that can exhibit chaotic dynamics and, in turn, exploding gradients during backpropagation • Systematic analysis of combinatorial generalization with a new dataset of protein sequence and structure
1. R ELATED W ORK
Protein modeling Our model builds on a long history of coarse-grained modeling of protein structure (Kolinski et al., 1998; Kmiecik et al., 2016).
Recently, multiple groups have demonstrated how to learn full force ﬁelds using likelihood-based approaches (Jumper et al., 2017; Krupa et al., 2017), similar to our maximum likelihood loss (but without backpropagtion through folding for fast sampling).
While this work was in progress, two groups reported neural models of protein structure (AlQuraishi, 2018;
Anand & Huang, 2018), where the former focused on predicting structure via angles and the latter via distances.
We show how an energy function provides a natural framework to integrate both kinds of constraints, which in turn is important to sample-efﬁcient structural generalization.
Learning to infer or sample Structured prediction includes a long history of casting predictions in terms of energy minimization (LeCun et al., 2006).
Recently, (Wang et al., 2016; Amos &
Kolter, 2017; Belanger & McCallum, 2016) have built hybrid neural networks that use differentiable optimization as a building block in neural architectures.
Structured Prediction Energy Networks (SPENs) with unrolled optimization (Belanger et al., 2017) are a highly similar approach to ours, differing in terms of the use of optimization rather than sampling.
Additional methodologically related work includes approaches to learn energy functions and samplers simultaneously (Wang &
Liu, 2016; Chen et al., 2018a; Song et al., 2017), to learn efﬁcient MCMC operators (Song et al., 2017; Levy et al., 2017), to build expressive approximating distributions with unrolled Monte Carlo simulations (Salimans et al., 2015; Titsias, 2017), and to learn the parameters of simulators with
[Figure]
implicitly deﬁned likelihoods 1 (Mohamed & Lakshminarayanan, 2016; Dai et al., 2017; Tran et al., 2017).
M ODEL
Overview NEMO is an end-to-end differentiable model of protein structure X conditioned on sequence information s consisting of three components (Figure 1): (i) a neural energy function
U ✓ [ x ; s ] for coarse grained structure x given sequence, (ii) an unrolled simulator that generates approximate samples from U via internal coordinate Langevin dynamics ( § 2.3), and (iii) an imputation network that generates an atomic model X from the ﬁnal coarse-grained sample x ( T ) ( § 2.4).
All components are trained simultaneously via backpropagation through the unrolled process.
2. R EPRESENTATION
Proteins Proteins are linear polymers (sequences) of amino acids that fold into deﬁned 3D structures.
The 20 natural amino acids have a common monomer structure [-(N-H)-(C-R)-(C=
O)-] with variable side-chain R groups that can differ in properties such as hydrophobicity, charge, and ability to form hydrogen bonds.
When placed in solvent (such as water or a lipid membrane), interactions between the side-chains, backbone, and solvent drive proteins into particular 3D conﬁgurations (‘folds’), which are the basis for understanding protein properties such as biochemical activity, ligand binding, and interactions with drugs.
Coordinate representations We predict protein structure X in terms of 5 positions per amino acid: the four heavy atoms of the backbone ( N , C ↵ , and carbonyl C
=O ) and the center of mass of the side chain R group.
While it is well-established that the locations of C ↵ carbons are sufﬁcient to reconstruct a full atomic structure (Kmiecik et al., 2016), we include these additional positions for evaluating backbone hydrogen bonding (secondary structure) and coarse side-chain placement.
The differentiable simulator generates an initial coarse-grained structure (1-position-per-amino-acid) with the loss function targeted to the midpoint of the C ↵ carbon and the side chain center of mass.
Sequence conditioning We consider two modes for conditioning our model on sequence information: (1) 1-seq , in which s is an L ⇥ 20 matrix containing a one-hot encoding of the amino acid sequence, and (2) Proﬁle , in which s is an L ⇥
40 matrix encoding both the amino acid sequence and a proﬁle of evolutionarily related sequences ( § B.6).
Internal coordinates In contrast to Cartesian coordinates x , which parameterize structure in terms of absolute positions of points x i 2 R 3 , internal coordinates z parameterize structure in terms of relative distances and angles between points.
We adopt a standard convention for internal coordinates of chains (Parsons et al., 2005) where each point x i is placed in a spherical coordinate system deﬁned by the three preceding points x i  1 , x i  2 , x i  3 in terms of a radius (bond length 2 ) b i 2 (0 , 1 ) , a polar angle (bond angle) a i 2 [0 , ⇡ ) , and an azimuthal angle (dihedral angle)
d
i 2 [0 , 2 ⇡ ) (Figure 2B).
We deﬁne z i = { ˜ b i , ˜
a
i , d i } , where ˜ b i , ˜ a i are unconstrained parameterizations ( § A.1).
The transformation x = F ( z ) from internal coordinates to Cartesian is then deﬁned by the recurrence
[Equation]
where ˆ u i =
x i  x i
|| x i  x i  1 || is a unit vector from x i  1 to x i and ˆ n i
= ˆ u i  1 ⇥ ˆ u
i ||
ˆ u
i  1 ⇥ ˆ u i || is a unit vector normal to each bond plane.
The inverse transformation z = F  1 ( x ) is simpler to compute, as it only involves local (and fully parallelizable) calculations of distances and angles ( § A.1).
2. N EURAL ENERGY FUNCTION
Deep Markov Random Field We model the distribution of a structure x conditioned on a sequence s with the Boltzmann distribution, p ✓ ( x | s ) = 1 Z exp
(  U ✓ [ x ; s ]) , where U ✓ [ x ; s ] is a sequenceconditioned energy function parameterized by a neural network.
Our approach is compatible with any differentiable energy function U [ x ; s ] , though we focus on a decomposition
U ✓ [ x ; s ] =
X i l i ( s ; ✓ ) f i ( x ; ✓ ) , (2) which is a Markov Random Field with coefﬁcients { l i ( s ; ✓ ) }
Mi =1 computed by a sequence network and structural features { f i ( x ; ✓ ) }
Mi =1 computed by a structure network (Figure 2A).
This decomposition facilitates (i) increased interpretability, as the (learned) structural features are independent of sequence, and (ii) increased computational efﬁciency, as the sequence-based coefﬁcients can be computed once.
Sequence network The sequence network takes as input one-dimensional sequence information s and outputs: (1) Energetic coefﬁcients , a set of 1- and 2-dimensional sequence features { l i ( s ; ✓ ) }
Mi =1 , (2) Simulator initial state z (0) ,
(3) Simulator hyperparameters preconditioning matrix C , and (4)
Predicted secondary structure (Figure 6).
It is parameterized by a combination of 1D, 2D, and graph convolutions (Gilmer et al., 2017) ( § A).
Structure network The structure network takes as input a coarse-grained structure x and outputs a set of 1D and 2D structural features { f i ( x ; ✓ ) } Mi =
(Figure 6).
We design the energy function to be invariant to rigid body motions (rotations and translations in SE(3) ) by leveraging a set of invariant base features (Figure 2C) which are:
Internal coordinates z All internal coordinates except 6 are invariant to rotation and translation
and we mask these in the energy loss.
Distances D ij =
k x i  x j k between all pairs of points.
We further process these by 4 radial basis functions with (learned) Gaussian kernels.
Orientation vectors ˆ v ij , which are unit vectors encoding the relative position of point x j in a local coordinate system of x i with base vectors
ˆ u i
ˆ u i
+ k
ˆ u i
{+ˆ u i +
+}
k , ˆ n i +1 , and the cross product thereof
Input : State z (0) , energy U ( x ) , step ✏ , time T , scale C
Output: Trajectory x (0) , . . .
, x ( T )
Initialize x (0)
F ( z (0) ) ; while t < T do
Compute forces f z =  @ x @ z T r x U ;
Sample  z ⇠ N  12 ✏ C f z , ✏ C  ; ⌅ z ( t + ✏ )
z ( t ) +  z ; ⌅ x ( t + ✏ )
F ( z ( t + ✏ ) ) ; t t + ✏ ; end
Input : State z (0) , energy U ( x ) , step ✏ , time T , scale C
Output: Trajectory x (0) , . . .
, x ( T )
Initialize x (0)
F ( z (0) ) ; while t < T do
Compute forces f z =  @ x @ z T r x U ;
Sample  z ⇠ N  12 ✏ C f z , ✏ C  ; ⌅ ˜x x ( t ) + @ x @ z ( t )  z ( t ) ;
⌅ x ( t + ✏ ) x
( t ) +⇣ @ x @ z ( t ) + @
˜x @
z ⌘  z ( t ) ; t t + ✏ ; end
2. E FFICIENT SIMULATOR
Langevin dynamics The Langevin dynamics is an SDE that aymptotically samples from the
Boltzmann distribution (Equation 1).
It is typically simulated by a ﬁrst-order discretization as x ( t + ✏ ) x
( t )  ✏
[Equation]
Internal coordinate dynamics
The efﬁciency with which Langevin dynamics explore conformational space is highly dependent on the geometry (and thus parameterization) of the energy landscape
U ( x ) .
While Cartesian dynamics are efﬁcient at local structural rearrangement, internal coordinate dynamics much more efﬁciently sample global, coherent changes to the topology of the fold (Figure 2B) .
We interleave the Cartesian Langevin dynamics with preconditioned Internal Coordinate dynamics,
[Equation]
where C is a preconditioning matrix that sets the relative scaling of changes different degrees of freedom.
For all simulations we unroll T = 250 time steps, each of which is comprised of one
Cartesian step followed by one internal coordinate step (Equation 9, § A.3).
Transform integrator Simulating internal coordinate dynamics is often computationally intensive as it requires rebuilding Cartesian geometry x from internal coordinates z with F ( z ) (Parsons et al., 2005) because it is intrinsically sequential.
Here we bypass the need for recomputing coordinate transformations at every step by instead computing on-the-ﬂy transformation integration (Figure 3).
The idea is to directly apply coordinate updates in one coordinate system to another by numerically integrating the Jacobian.
This can be favorable when the Jacobian has a simple structure, such as in our case where it requires only distributed cross products.
2. A TOMIC IMPUTATION
Local reference frame reconstruction The imputation network builds an atomic model X from the ﬁnal coarse coordinates x ( T ) .
Each atomic coordinate X i,j of atom type j at position i is placed in a local reference frame as
[Equation]
where e i,j ( z ; ✓ ) and r i,j ( z ; ✓ ) are computed by a 1D convolutional neural network (Figure 6).
T RAINING
We train and evaluate the model on a hierarchically and temporally split set of ⇠ 60,000 protein structures.
The model is trained by gradient descent using a composite loss that combines terms from likelihood-based and empirical-risk minimization-based training.
3. D ATA
Structural stratiﬁcation There are several scales of generalization in protein structure prediction, which range from predicting the structure of a sequence that differs from the training set by a few letters to predicting a 3D structure with a fold topology with no representative in the training set.
To test these various levels of generalization systematically across many different protein families, we built a training and test set on top of the CATH hierarchical classiﬁcation of protein folds (Orengo et al., 1997).
CATH hierarchically organizes proteins from the Protein Data Bank (Berman et al., 2000) into domains (individual folds) that are classiﬁed at the levels of C lass, A rchitecture, T opology, and H omologous superfamily.
We collected protein domains from CATH up to length 200 and hierarchically and temporally split this set into training ( ⇠ 35k folds), validation ( ⇠ 15k folds), and test sets ( ⇠ 10k folds).
Test subsets The ﬁnal test set is subdivided into four subsets: C , A , T , and H , based on the level of maximal similarity between a given test domain and domains in the training set.
For example, domains in the C or A sets may share class and potentially architecture classiﬁcations with train but will not share topology (i.e. fold).
3. L OSS
Likelihood The gradient of the data-averaged log likelihood of the Boltzmann distribution is
[Equation]
which amounts to minimizing the average energy of samples from the data relative to samples from the model.
In an automatic differentiation setting, we implement a Monte Carlo estimator for this gradient by adding the energy gap ,
[Equation]
to the loss, where ? is an identity operator that sets the gradient to zero 4 .
Empirical Risk
In addition to the likelihood loss, which backpropagates through the energy function but not the whole simulation, we developed an empirical risk loss composing several measures of protein model quality.
It takes the form
[Equation]
schematized in Figure 6.
Our combined loss sums all of the terms
L =
L ER + L ML without weighting.
3. S TABILIZING BACKPROPAGATION THROUGH TIME
We found that the long roll-outs of our simulator were prone to chaotic dynamics and exploding gradients, as seen in previous work on meta-learning (Maclaurin et al., 2015).
Unfortunately, when chaotic dynamics do occur, it is typical for all gradients to explode (across learning steps) and standard techniques such as gradient clipping (Pascanu et al., 2013) are unable to rescue learning ( § B.4).
To stabilize training, we developed two complimentary techniques that regularize against chaotic simulator dynamics while still facilitating learning when they arise.
They are
• Lyapunov regularization We regularize the simulator time-step function (rather than the energy function) to be approximately 1-Lipschitz.
(if exactly satisﬁed, this eliminates the possibility of chaotic dynamics)
[Figure]
lass
[Figure]
Generalizationlevels
[Table]
• Damped backpropagation through time We exponentially decay gradient accumulation on the backwards pass of automatic differentiation by multiplying each backwards iteration by a damping factor  .
We adaptively tune  to cancel the scale of the exploding gradients.
This can be thought of a continuous relaxation of and quantitatively tunable alternative to truncated backpropagation through time.
R ESULTS 4. G ENERALIZATION ACROSS CATH
For each of the 10,381 protein structures in our test set, we sampled 100 models from NEMO, clustered them by structural similarity, and selected a representative structure by a consensus
algorithm (Ginalski et al., 2003).
For evaluation of performance we focus on the TM-Score (Zhang & Skolnick, 2005), a measure of structural similarity between 0 and 1 for which TM > 0.5 is typically considered an approximate reconstruction of a fold.
Calibrated uncertainty
We ﬁnd that, when the model is conﬁdent (i.e. the number of distinct structural clusters is low ⇠ 1-3), it is also accurate with many predictions having TM > 0.5 (Figure 4, left).
Unsurprisingly, the conﬁdence of the model tends to go with the difﬁculty of generalization, with the most conﬁdent predictions from the H test set and the least conﬁdent from C .
Structural generalization
However, even when sequence identity is low and generalization difﬁculty is high (Figure 4, center), the model is still able to make some accurate predictions of 3D structure.
Figure 5 illustrates some of these successful predictions at C and A levels, speciﬁcally
[Figure]
and sheet the predictive distribution is multimodal with non-trivial differences between the clusters representing alternate packing of the chain.
In some of the models there is uneven distribition of uncertainty along the chain, which sometimes corresponded to loosely packed regions of the protein.
Comparison to an end-to-end baseline We constructed a baseline model that is a non-iterative replica of NEMO which replaces the coarse-grained simulator module (and energy function) with a two-layer bidirectional LSTM that directly predicts coarse internal coordinates z (0) (followed by transformation to Cartesian coordinates with F ).
We trained this baseline across a range of hyperparameter values and found that for difﬁcult C , A , and T tasks, NEMO generalized more effectively than the RNNs (Table 1).
For the best performing 2x300 architecture, we trained two additional replicates and report the averaged perfomance in Figure 4 (right).
Additionally, we report the results of a sequence-only NEMO model in Table 1.
Paralleling secondary structure prediction (Rost & Sander, 1993; McGufﬁn et al., 2000), we ﬁnd that the availability of evolutionary information has signiﬁcant impact
4. A DVANTAGES AND DISADVANTAGES
This work presents a novel approach for protein structure prediction that combines the inductive bias of simulators with the speed of directed models.
A major advantage of the approach is that model sampling (inference) times can be considerably faster than conventional approaches to protein structure prediction (Table 4).
There are two major disadvantages, both of which have to do with training.
First, the computational cost of training is higher than that of angle-predicting RNNs such as our baseline or AlQuraishi (2018), and consequently those methods have been feasibly scaled to larger datasets than ours (in protein length and diversity) that are more relevant to protein structure prediction tasks.
Second, the instability of backpropagating through long simulations is unavoidable and only partially remedied by our approaches of Lipschitz regularization and gradient damping which can lead to slower learning and less expressive energy functions.
Methods for efﬁcient (i.e. subquadratic) N -body simulations and for more principled stabilization of deep networks may be relevant to addressing both of these challeges in the future.
C ONCLUSION
We described a model for protein structure given sequence information that combines a coarse-grained neural energy function and an unrolled simulation into an end-to-end differentiable model.
To realize this idea at the scale of real proteins, we introduced an efﬁcient simulator for Langevin dynamics in transformed coordinate systems and stabilization techniques for backpropagating through long simulator roll-outs.
We ﬁnd that that model is able to predict the structures of protein molecules with hundreds of atoms while capturing structural uncertainty, and that the model can structurally generalize to distant fold classiﬁcations more effectively than a strong baseline
A PPENDICES A M ODEL A. C OORDINATE SYSTEMS
Inverse transformation
The inverse transformation z = F  1 ( x ) involves fully local computations of bong lengths and angles.
[Equation]
Jacobian The Jacobian @ x @ z deﬁnes the inﬁnitesimal response of the Cartesian coordinates x to perturbations of the internal coordinates z .
It will be important for both converting Cartesian forces into angular torques and bond forces as well as the development of our transform integrator. It is [Table]
Jacobian matrix
[Table]
deﬁned element-wise as
[Equation]
The Jacobian has a simple form that can be understood by imagining the protein backbone as a robot arm that is planted at x 0 (Figure 2B).
Increasing or decreasing the bond length
b i extends or retracts all downstream coordinates along the bonds axis, moving a bond angle a i drives circular motion of all downstream coordinates around the bond normal vector
ˆ n i centered at x i  1 , and moving a dihedral angle d i drives circular motion of downstream coordinate x j around bond vector
ˆ u i centered at x i  1 .
Unconstrained representations Bond lengths and angles are subject to the constraints b i > 0 and 0 < a i < ⇡ .
We enforce these constraints by representing these degrees of freedom in terms of fully unconstrained variables ˜
b i and ˜ a i via the transformations b i = log ⇣ 1 + e ˜ b
i ⌘ and a i = ⇡
1+ e  ˜ ai .
All references to the internal coordinates z and Jacobians @ x @ z will refer to the use of fully unconstrained representations (Table 2).
A. E NERGY FUNCTION
Figure 6 provides an overall schematic of the model, including the components of the energy function.
CNN primitives All convolutional neural network primitives in the model schematic (Figure 6) follow a common structure consisting of stacks of residual blocks.
Each residual block includes consists of a layer of channel mixing (1x1 convolution), a variable-sized convolution layer, and a second layer of channel mixing.
We use dropout with p = 0 .
9 and Batch Renormalization (Ioffe, 2017) on all convolutional layers.
Batch Renormalization rather than Normalization was necessary rather owing to the large variation in sizes of the structures of the proteins and resulting large variation in mini-batch statistics.
A. I NTERNAL COORDINATE DYNAMICS WITH A TRANSFORM INTEGRATOR
Why sampling vs. optimization Deterministic methods for optimizing the energy U ( x ; s ) such as gradient descent or quasi-Newton methods can effectively seek local minima of the energy surface, but are challenged to optimize globally and completely ignore the contribution of the widths of energy minima (entropy) to their probability.
We prefer sampling to optimization for three reasons: (i) noise in sampling algorithms can facilitate faster global conformational exploration by overcoming local minima and saddle points, (ii) sampling generates populations of states that respect the width (entropy) of wells in U and can be used for uncertainty quantiﬁcation, and (iii) sampling allows training with an approximate Maximum Likelihood objective (Equation 5).
Langevin Dynamics
The Langevin dynamics are a stochastic dynamics that sample from the canonical ensemble.
They are deﬁned as a continuous-time stochastic differential equation, and are simulated in discrete time with the ﬁrst order discretization
[Equation]
Each time step of ✏ involves a descent step down the energy gradient plus a perturbation of Gaussian noise.
Importantly, as time tends toward to inﬁnity, the time-distribution of the Langevin dynamics converges to the canonical ensemble.
Our goal is to design a dynamics that converge to an approximate sample in a very short period of time.
Coordinate systems and preconditioning The efﬁciency with which Langevin dynamics explore conformational space is highly dependent on the geometry of the energy landscape U
( x ) , which in turn depends on how the system is parameterized.
Molecular energy functions in Cartesian coordinates tend to exhibit strong correlations between variables that result from the requirement that underlying molecular geometries satisfy highly stereotyped bond lengths and angles.
As a result, simulations of naive Cartesian Langevin dynamics require a small time step to satisfy these constraints and tend to be dominated by high-frequency, localized vibrations of the chain.
The large, global motions that are essential to protein folding can require thousands to millions of times steps to manifest.
A well-known solution to the complex dependencies of Cartesian coordinates is to carry out optimization and simulation in internal coordinates, which directly parameterize molecular geometries in terms of the bond lengths and angles (Parsons et al., 2005).
Internal coordinate parameterizations possess the advantages that (i) bond length and angle constraints are easy to satisfy and (ii) small changes to a single angle can drive large, coherent rearrangements of the chain (Figure 2B).
For example, simply replacing x ’s with z ’s in Equation 8 yields the dynamics
[Equation]
The advantages and disadvantages of the two coordinate systems are complementary: Cartesian dynamics efﬁciently sample local structural rearrangements and inefﬁciently sample global chain motions, while internal coordinate dynamics efﬁciently sample global, correlated motions of the chain but are challenged to make precise local rearrangements.
The time dynamics of these alternative parameterizations need not be kinetically realistic to converge to the correct distribution over conformational space.
Different coordinate systems warp the local geometry of the energy landscape and will in turn rescale and redirect which global vibrational and local vibrations dominate the dynamics.
This relative rescaling can be further optimized by applying a global linear transformation to the energy landscape with a preconditioning ‘inverse mass’ matrix
C , giving the update
[Equation]
[Table]
Transform integrator
The need to rebuild Cartesian geometry x from internal coordinates z with
F ( z ) at every time step is one of the major costs of conformational sampling codes based on Internal coordinates (Parsons et al., 2005) because it is intrinsically sequential.
Here we show how it is possible to bypass the need for geometry reconstruction at every step by instead computing on-the-ﬂy geometry modiﬁcation .
Imagine following a change to the internal coordinates  z ( t ) along a straight path from z ( t ) to z ( t + ✏ ) and tracking the corresponding nonlinear path of the Cartesian coordinates from x ( t ) to x ( t + ✏ ) .
If this path is indexed by u 2 ( t, t + ✏ ) , then the dynamics of x with respect to u are given by @ x @ u = @ x @ z @
z @ u = @ x @ z 1 ✏  z .
Integrating the dynamics of x gives
[Equation]
This illustrates that it is possible to convert coordinate changes in one coordinate system (e.g. Internal
Coordinates) to coordinate changes in another (e.g. Cartesian) by integrating an autonomous system of ODEs with dynamics governed by the Jacobian.
Since ✏ is small, we integrate this system with a single step of Heun’s method (improved Euler), where we ﬁrst substitute an Euler approximation to predict x ( t + ✏ ) as
[Equation]
and then substitute the Jacobian evaluated at the predicted state ˜x ( t + ✏ ) to form trapezoidal approximation
[Equation]
The comparison of this algorithm with naive integration is given in Figure 7.
The corrector step is important for eliminating the large second-order errors that arise in curvilinear motions caused by angle changes (Figure 2B and Figure 7).
In principle higher-order numerical integration methods or more time steps could increase accuracy at the cost of more evaluations of the Jacobian, but we found that second-order effects seemed to be the most relevant on our timescales.
Mixed integrator Cartesian dynamics favor local structural rearrangements, such as the transitioning from a helical to an extended conformation, while internal coordinate dynamics favor global motions such as the change of the overall fold topology.
Since both kinds of structural rearrangements are important to the folding process, we form a hybrid integrator (Algorithm 3) by taking one step with each integrator per force evaluation.
Translational and rotational detrending Both Cartesian and Internal coordinates are overparameterized with 3 L degrees of freedom, since only 3 L  6 degrees of freedom are necessary to encode a centered and un-oriented structure 5 .
As a consequence, a signiﬁcant fraction of the per time-step changes  x can be explained by rigid translational and rotational motions of the entire structure.
We isolate and remove these components of motion by treating the system { x 1 , . . .
, x L } as a set of particles with unit mass, and computing effective structural translational and rotational velocities by summing point-wise momenta.
The translational component of motion is simply the average displacement across positions  x Trans
i = h
x i i .
For rotational motion around the center of mass, it is convenient to deﬁne the non-translational motion as  ¯ x i =
x i   x
Trans i and the centered Cartesian coordinates as ¯ x i
= x i  h
x
i i .
The point-wise angular momentum is then l i =
¯ x i ⇥  ¯ x i and we deﬁne a total angular velocity of the structure !
by summing these and dividing by the moment of inertia as !
=
P i l i ) / P i || ¯
x i || .
We convert the angular velocity !
into Cartesian displacements with an unrolled Heun integration as  x Rot i
[Equation] i ) , which leaves the isolated structural motions as  x Struct
i [Equation]
[Figure] Input : Initial state z (0) , energy U ( x ) , time steps ✏ x , ✏ z , total time T , preconditioners C x , C z , Output : Trajectory x (0) , . . . , x ( T )
Initialize x (0) F ( z (0) ) ; while t < T do f x r x U ;  x ( Cart ) CartesianStep ( x ( t ) , f x , ✏ x , C x ) ;  x ( Int ) ClippedInternalStep ( x +  x ( Cart ) , f x , ✏ z , C z ) ; x x + Detrend
(  x ( Cart )
+  x ( Int ) ) ; t t + ✏ ; end
Speed clipping We found it helpful to stabilize the model by enforcing a speed limit on overall structural motions for the internal coordinate steps.
This prevents small changes to the energy function during learning from causing extreme dynamics that in turn produce a non-informative learning signal.
To accomplish this, we translationally and rotationally detrend the update of the predictor step  x and compute a hypothetical time step ˆ ✏ z that would limit the fastest motion to 2 Angstroms per iteration.
We then compute modiﬁed predictor and corrector steps subject to this new, potentially slower, time step.
While this breaks the asymptotics of Langevin dynamics, (i) it is unlikely on our timescales that we achieve stationarity and (ii) it can be avoided by regularizing the dynamics away from situations where clipping is necessary.
In the future, considering non-Gaussian perturbations with kinetic energies similar to Relativistic Monte Carlo (Lu et al., 2016) might accomplish a similar goal in a more principled manner.
The ﬁnal integrator combining these ideas is presented in Figure 3.
B A PPENDIX B: T RAINING B. D ATA
For a training and validation set, we downloaded all protein domains of length L  200 from C lasses ↵ ,  , and ↵ /  in CATH release 4.1 (2015), and then hierarchically purged a randomly selected set of A , T , and H categories.
This created three validation sets of increasing levels of difﬁculty: H , which contains domains with superfamilies that are excluded from train (but fold topologies may be present), T , which contains fold topologies that were excluded from train (fold generalization), and A which contains secondary structure architectures that were excluded from train.
For a test set, we downloaded all folds that were new to CATH release 4.2 (2017), which (due to a propensity of structural biology to make new structures of previously solved folds), provided 10, test domains at varying levels of difﬁculty.
B. L OSS
We optimize the model using a composite loss containing several terms, which are detailed as follows.
Distance loss We score distances in the model with a contact-focused distance loss
[Equation]
where the contact-focusing weights are
[Equation]
⌘ and  ( u ) = 1 1+exp(  u ) is the sigmoid function.
Angle loss We use the loss
[Equation]
where H ( z ) =
[cos( a i ) sin( a i ) cos( d i ) sin( a i ) sin( d i )] T are unit length feature vectors that map the angles { a i , d i } to the unit sphere.
Other angular losses, such as the negative log probability of a Von-Mises Fisher distribution, are based on the inner product of the feature vectors H ( z a ) · H ( z b ) rather than the Euclidean distance ||H ( z a )  H ( z b ) || between them.
It is worth noting that these two quantities are directly related by ||H ( z a )  H ( z b ) ||
= p 2(1  H ( z a ) · H ( z b )) .
Taking z a as ﬁxed and z b as the argument, the Euclidean loss has a cusp at z a whereas the Von-Mises Fisher loss is smooth around z a .
This is analogous to the difference between L 1 and L 2 losses, where the cusped L 1 loss favors median behavior while the smooth L 2 loss favors average behavior.
A third loss similar to the Euclidean loss would be the Great Circle distance, arccos ( H ( z a ) · H ( z b )) .
Trajectory loss In a further analogy to reinforcement learning, damped backpropation through time necessitates an intermediate loss function that can criticize transient states of the simulator.
We compute this by featurizing the per time step coordinates as the product D ij ˆ v ij (Figure 2C) and doing the same contact-weighted averaging as the distance loss.
Template Modelling (TM) Score The TM [Equation] score (Zhang & Skolnick, 2005),
[Equation]
is a measure of superposition quality between two protein structures on [0 , 1] that was presented as an approximately length-independent alternative to RMSD.
The TM-score is the best attainable value of the preceding quantity for all possible superpositions of two structures, where D i =
|| x ( Model )  x ( Data ) || .
This requires iterative optimization, which we implemented with a sign gradient descent with 100 iterations to optimally superimpose the model and target structure.
We backpropagate through this unrolled optimization process as well as that of the simulator.
Hydrogen bond loss We determine intra-backbone hydrogen bonds using the electrostatic model of DSSP (Kabsch & Sander, 1983).
First, we place virtual hydrogens at 1 Angstroms along the negative angle bisector of the C i  1  N i  C ↵ i bond angle.
Second, we compute a putative energy U h-bond ij (in kcal / mol ) for each potential hydrogen bond from an amide donor at i to a carbonyl acceptor at j as
[Equation]
where D ab = ||
X i,a  X j,b || is the Euclidean distance between atom a of residue i and atom b of residue j .
We then make hard assignments of hydrogen bonds for the data with
[Equation]
We ‘predict’ the probabilities of hydrogen bonds of the data given the model via logisitic regression of soft model assignments as
[Equation]
where a, b, c are learned parameters with the softplus parameterizations enforcing a, b > 0 and  ( u )
= 1 / (1 + exp
(  u ) is the sigmoid function.
The ﬁnal hydrogen bond loss is the cross-entropy between these predictions and the data,
[Equation]
Secondary Structure Prediction We output standard 8-class predictions of secondary structure and score them with a cross-entropy loss.
B. S TABILIZING BACKPROPAGATION THROUGH TIME
The combination of energy function, simulator, and reﬁnement network can build an atomic level model of protein structure from sequence, and our goal is to optimize (meta-learn) this entire procedure by gradient descent.
Before going into speciﬁcs of the loss function, however, we will discuss a challenges and solutions for computing gradients of unrolled simulations in the face of chaos.
B. C HAOS AND EXPLODING GRADIENTS
Gradient-based learning of iterative computational procedures such as Recurrent Neural Networks (RNNs) is well known to be subject to the problems of exploding and vanishing gradients (Pascanu et al., 2013).
Informally, these occur when the sensitivities of model outputs to inputs become either extremely large or extremely small and the gradient is no longer an informative signal for optimization.
We ﬁnd that backpropagation through unrolled simulations such as those presented is no exception to this rule.
Often we observed that a model would productively learn for tens of thousands of iterations, only to suddenly and catastrophically exhibit diverging gradients from which the optimizer could not recover - even when the observed simulation dynamics exhibited no obvious qualitative changes to behavior and the standard solutions of gradient clipping (Pascanu et al., 2013) were in effect.
Similar phenomena have been observed previously in the context of meta-learning (Maclaurin et al., 2015).
In Figure 8, we furnish a minimal example that illustrates how chaos can lead to irrevocable loss of learning.
We see that for even a simple particle-in-a-well, some choices of system parameters
[Figure] Step size b Chaotic Periodic (such as too large a time step) can lead to chaotic dynamics which are synonymous with explosive gradients.
This example is hardly contrived, and is in fact a simple model of the distance potentials between coordinates in our simulations.
Moreover, it is important to note that chaos may not be easy to diagnose: for learning rates ↵ 2 [1 . 7 , 1 . 8] the position of the particle x remains more or less conﬁned in the well while the sensitivities diverge to 10 200 .
It seems unlikely that meta-learning would be able to recover after descending into chaos.
The view per time step Exploding gradients and chaotic dynamics involve the same mechanism: a multiplicative accumulation of sensitivities.
In dynamical systems this is frequently phrased as ‘exponentially diverging sensitivity to initial conditions’.
Intuitively, this can be understood by examining how the Jacobian of an entire trajectory decomposes into a product of Jacobians as
[Equation]
When the norms of the per time-step Jacobians @ x ( t ) @
x ( t  1) are typically larger than 1, the sensitivity || @
x ( T ) @
x (0) || will grow exponentially with T .
Ideally, we would keep these norms well-behaved which is the rationale recent work on stabilization of RNNs (Henaff et al., 2016; Chen et al., 2018b).
Next we will offer a general-purpose regularizer to approximately enforce this goal for any differentiable computational iteration with continuous state.
Approximate Lipschitz conditions One condition that guarantees that a deterministic map F :
R N !
R N , x t = F ( x t  1 , ✓ ) cannot exhibit exponential sensitivity to initial conditions is the condition of being non-expansive (also known as 1-Lipschitz or Metric).
That is, for any two input points x a , x b 2 R N , iterating the map cannot increase the distance between them as | F ( x a , ✓ )  F ( x b , ✓ ) |  | x a  x b | .
Repplying the map to the bound immediately implies
[Equation]
Input : Initial state x (0) , time-stepping function F ( x , s , ✓ ) , external inputs s 1 , . . .
, s T parameters ✓ , Loss function L ( x 1 , . . .
, x T ) , Damping factor 0 <<  <
Output : Exponentially damped gradient r ✓ L
Initialize x (0)
F ( z (0) ) ; for t 2 , . .
, T do
Compute time step ˜x t
F ( x t  1 , s t , ✓ ) ;
Decay the gradient † x t (1   ) ?
( ˜x t )
+  ˜x t ; end
Compute loss L ( x 1 , . . .
, x T ) ;
Compute gradient r ✓ L
AutoDiff ( L , ✓ ) ; † where ?
( · ) is the stop gradient function.
for any number of iterations t .
Thus, two initially close trajectories iterated through a non-expansive mapping must remain at least that close for arbitrary time. We approximately enforce non-expansivity by performing an online sensitivity analysis within simulations. At randomly selected time-steps, the current time step x ( t ) is rolled back to the preceding state and re-executed with small Gaussian perturbations to the state  ⇠ N (0 , [Equation]
4 I ) 6 .
We regularize the sensitivity by adding
[Equation]
to the loss.
Interestingly, the stochastic nature of this approximate regularizer is likely a good thing - a truly non-expansive map is quite limited in what it can model.
However, being ‘almost’ non-expansive seems to be incredibly helpful for learning.
Damped Backpropagation through Time
The approximate Lipschitz conditions (or Lyapunov regularization) encourage but do not guarantee stable backpropagation.
When chaotic phasetransitions or otherwise occur we need a fall-back plan to be able to continue learning.
At the same time, we would like gradient descent to proceed in the usual manner when simulator dynamics are stable.
To this end we introduce a damping factor to backpropagation that can adaptively combat exponentially diverging gradients with exponential discounting (Algorithm 4).
Damped backpropagation can be seen as a continuous alternative to the standard approach of
Truncated Backpropagation through Time.
Rather than setting the gradient to 0 after some ﬁxed intervals of time-steps, we decay it on the backwards pass of reverse-mode differentiation by a factor of  .
This is mildly evocative of the notion of discounted future rewards in reinforcement learning.
During backpropagation this causes a biased estimate of Jacobians that favors short term sensitivities (or rewards) as
[Equation]
B. M ULTIPLE S EQUENCE A LIGNMENT G ENERATION
We use multiple sequence alignments of evolutionarily related sequences for both proﬁle construction ( § B.6) and (ii) data augmentation ( § B.7).
For every domain in the dataset, we extracted the sequence from the PDB and then used jackhmmer (Eddy, 2011), to iteratively search the Uniprot90 database (Suzek et al., 2014) (release 4/2016) with 5 iterations and a length-normalized bitscore threshold of 0.3.
We then removed sequences with over 50% gaps relative to the query sequence and then redundancy-reduced the alignment with hhﬁlter (Remmert et al., 2012) such that all sequences are at least a normalized Hamming distance of 0.8 away from one another. [Table]
B. P ROFILE G ENERATION
We brieﬂy describe how we construct evolutionary proﬁles, or position-speciﬁc scoring matrices (PSSMs), for each protein domain.
Let S
=
{ S (1) , . . .
, S ( L ) } be the set of L columns of a multiple sequence alignment over M sequences where each column S ( i ) is an M ⇥ q matrix that one-hot encodes the sequence data at position
i (for an alphabet of size q ).
The regularized empirical frequency of letter a at site i is then
[Equation]
where ↵ is a pseudocount that we set to 10.
We compute our PSSM features for letter a at site i as
[Equation]
where  ( u ) = 1 1+exp(  u ) is the logistic sigmoid and B a is the average frequency of amino acid a in
UniProt (Apweiler et al., 2004).
B. E VOLUTIONARY D ATA A UGMENTATION
To reduce our reliance on alignments and the generation of proﬁles for inference of new sequences while still leveraging evolutionary sequence data, we augmented our training set by dynamically spiking in diverse, related sequence into the model during training.
Given a set of M sequences in the alignment we sample a sequence t based on its normalized Hamming distance d t with probability
[Equation]
where  EDA is a scaling parameter that we set to 5.
When the alternate sequence contains gaps, we construct a chimeric sequence that substitutes those sites with the query.
This strategy increased the number of available sequence-structure pairs by several orders of magnitude, and we used it for both proﬁle and 1-seq based training.
C A PPENDIX C: R ESULTS C. S TRUCTURE GENERATION AND PROCESSING
For each sequence from the CATH release 4.2 dataset, 100 structures were generated from both the proﬁle and sequence-only models, while a single structure was generated from the RNN baseline models.
The reported TM-scores were calculated using Maxcluster (Siew et al., 2000).
A single representative structure was chosen from the ensemble of 100 structures using 3D-Jury (Ginalski et al., 2003).
A pairwise distance matrix of TM-scores was calculated for all of the 100 structures in the ensemble.
Clusters were determined by agglomerative hierarchical clustering with complete linkage using a TM-score threshold of 0.5 to determine cluster membership. [Figure] sc [Figure]
M
[Figure]
sc o r e
[Figure]
T M
[Figure]
sc o r e
( s equen c eon l y )
[Figure]
[Figure]