Title: LEARNING PROTEIN STRUCTURE WITH A DIFFERENTIABLE SIMULATOR

Abstract: The Boltzmann distribution is a natural model for many systems, from brains to materials and biomolecules, but is often of limited utility for fitting data because Monte Carlo algorithms are unable to simulate it in available time. This gap between the expressive capabilities and sampling practicalities of energy-based models is exemplified by the protein folding problem, since energy landscapes underlie contemporary knowledge of protein biophysics but computer simulations are challenged to fold all but the smallest proteins from first 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 efficient 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, in some cases, generalize to unobserved protein fold types when trained on a large corpus of protein structures.

1. INTRODUCTION:

Many natural systems, such as 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 system behavior, but can be difficult to fit to data. Given an energy function U ✓ [x], the probability of a system configuration x scales exponentially with energy as

EQUATION

where the (typically intractable) constant Z normalizes the distribution. Importantly, simple energy functions U ✓ [x] consisting of weak, local interactions can collectively encode complex system behaviors, such as the structures of materials and molecules or, when endowed with latent variables, the statistics of images, sound, and text (Ackley et al., 1985; Salakhutdinov & Larochelle, 2010) . Unfortunately, learning model parameters ✓ and generating samples x ⇠ p ✓ (x) of the Boltzmann distribution is difficult in practice, as these procedures depend on expensive Monte Carlo simulations that may struggle to mix effectively. These difficulties have driven a shift towards generative models that are easier to learn and sample from, such as directed latent variable models 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 To rapidly sample low-energy configurations, the Langevin dynamics simulator leverages both (i) an internal coordinate parameterization, which is more effective for global rearrangements, and (ii) a Cartesian parameterization, which is more effective for localized structural refinement. (C) The base features of the structure network are rotationally and translationally invariant internal coordinates (not shown), pairwise distances, and pairwise orientations.

2016; Amos & Kolter, 2017; Belanger & McCallum, 2016) . 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 (Kim & Bengio, 2016; Wang & Liu, 2017; Dai et al., 2017; Song et al., 2017; Chen et al., 2018a) , to learn efficient MCMC operators (Song et al., 2017; Levy et al., 2018) , to build expressive approximating distributions with unrolled Monte Carlo simulations (Salimans et al., 2015; Titsias, 2017) , and to learn the parameters of simulators with implicitly defined likelihoods1 (Mohamed & Lakshminarayanan, 2016; Tran et al., 2017) .

2. MODEL:

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 final coarse-grained sample x (T ) ( § 2.4). All components are trained simultaneously via backpropagation through the unrolled process.

2.1. REPRESENTATION:

Proteins Proteins are linear polymers (sequences) of amino acids that fold into defined 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 configurations ('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 sufficient 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. Internally, the differentiable simulator generates an initial coarse-grained structure (1-position-peramino-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) Profile, in which s is an L ⇥ 40 matrix encoding both the amino acid sequence and a profile of evolutionarily related sequences ( § B.7).

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 defined by the three preceding points x i 1 , x i 2 , x i 3 in terms of a radius (bond length2 ) 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 define z i = { bi , ãi , d i }, where bi , ãi are unconstrained parameterizations of b i and a i ( § A.1). The transformation x = F(z) from internal coordinates to Cartesian is then defined by the recurrence

x i = x i 1 + b i [ ûi 1 ni 1 ⇥ ûi 1 ni 1 ] " cos(⇡ a i ) sin(⇡ a i ) cos(d i ) sin(⇡ a i ) sin(d i ) # , where ûi = xi xi 1 ||xi xi 1 || is a unit vector from x i 1 to x i and ni = ûi 1 ⇥ ûi || ûi 1 ⇥

û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.2. NEURAL 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

EQUATION

which is a Markov Random Field with coefficients {l i (s; ✓)} M i=1 computed by a sequence network and structural features {f i (x; ✓)} M i=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 efficiency, as the sequence-based coefficients can be computed once and reused throughout a simulation.

Sequence network:

The sequence network takes as input one-dimensional sequence information s and outputs: (1) Energetic coefficients, a set of 1-and 2-dimensional sequence features {l i (s; ✓)} M i=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; ✓)} M i=1 (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:

1. Internal coordinates z All internal coordinates except 6 are invariant to rotation and translation3 and we mask these in the energy loss.

2. Distances D ij = kx i x j k between all pairs of points. We further process these by 4 radial basis functions with (learned) Gaussian kernels.

3. Orientation vectors vij , which are unit vectors encoding the relative position of point x j in a local coordinate system of x i with base vectors ûi ûi+1 k ûi ûi+1k , ni+1 , and the cross product thereof.

2.3. EFFICIENT SIMULATOR:



Langevin dynamics:

The Langevin dynamics is a stochastic differential equation that aymptotically samples from the Boltzmann distribution (Equation 1). It is typically simulated by a first-order discretization as

EQUATION

Internal coordinate dynamics The efficiency with which Langevin dynamics explores conformational space is highly dependent on the geometry (and thus parameterization) of the energy landscape U (x). While Cartesian dynamics are efficient at local structural rearrangement, internal coordinate dynamics much more efficiently 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 to each degree 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) which is an intrinsically sequential process. Here we bypass the need for recomputing coordinate transformations at every step by instead computing on-the-fly 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.4. ATOMIC IMPUTATION:

Local reference frame reconstruction The imputation network builds an atomic model X from the final 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

X i,j = x i + e i,j (z; ✓) [ ûi ni+1 ni+1 ⇥ ûi ] r i,j (z; ✓),

where e i,j (z; ✓) and r i,j (z; ✓) are computed by a 1D convolutional neural network (Figure 6 ).

3. TRAINING:

We train and evaluate the model on a set of ⇠67,000 protein structures (domains) that are hierarchically and temporally split. The model is trained by gradient descent using a composite loss that combines terms from likelihood-based and empirical-risk minimization-based training.

Algorithm 1: Direct integrator 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 fz = @x @z T rxU ;

Sample z ⇠ N 1 2 ✏Cfz, ✏C ; ⌅ z (t+✏) z (t) + z; ⌅ x (t+✏) F(z (t+✏) ); t t + ✏; end Algorithm 2: Transform integrator 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 fz = @x @z This exchanges the cost of an inner-loop transformation step (e.g. geometry construction F(z)) for an extra Jacobian evaluation, which is fully parallelizable on modern hardware (e.g. GPUs).

T rxU ; Sample z ⇠ N 1 2 ✏Cfz, ✏C ; ⌅ x x (t) + @x @z (t) z (t) ; ⌅ x (t+✏) x (t) + 1 2 ⇣ @x @z (t) + @x @z ⌘ z (t) ; t t + ✏; end

3.1. DATA:

Structural stratification 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 at a few positions to predicting a 3D fold topology that is absent from training set. To test these various levels of generalization systematically across many different protein families, we built a dataset on top of the CATH hierarchical classification 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 classified at the levels of Class, Architecture, Topology, and Homologous superfamily (from general to specific). We collected protein domains from CATH releases 4.1 and 4.2 up to length 200 and hierarchically and temporally split this set ( § B.1) into training (⇠35k folds), validation (⇠21k folds), and test sets (⇠10k folds).

Test subsets:

The final 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 classifications with train but will not share topology (i.e. fold).

3.2. LOSS:

Likelihood The gradient of the data-averaged log likelihood of the Boltzmann distribution is

EQUATION

which, when ascended, will minimize 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 (the negative of) this gradient by adding the energy gap,

EQUATION

to the loss, where ? is an identity operator that sets the gradient to zero4 .

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 schematized in Figure 6 . Our combined loss sums all of the terms L = L ER + L ML without weighting.

EQUATION

3.3. STABILIZING BACKPROPAGATION THROUGH TIME:

We found that the long roll-outs of our simulator were prone to chaotic dynamics and exploding gradients, as seen in other work (Maclaurin et al., 2015; Parmas et al., 2018) . 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.5). 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 satisfied, this eliminates the possibility of chaotic dynamics.)

• 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 as a continuous relaxation of and a quantitatively tunable alternative to truncated backpropagation through time. 

4. RESULTS:



4.1. GENERALIZATION 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 standard 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 find that, when the model is confident (i.e. the number of distinct structural clusters is low ⇠1-3), it is also accurate with some predictions having average TM > 0.5 (Figure 4 , left). Unsurprisingly, the confidence of the model tends to go with the difficulty of generalization, with the most confident predictions from the H test set and the least confident from C.

Structural generalization However, even when sequence identity is low and generalization difficulty is high (Figure 4 , center), the model is still able to make some accurate predictions of 3D structures. Figure 5 illustrates some of these successful predictions at C and A levels, specifically 4ykaC00, 5c3uA02 and beta sheet formation in 2oy8A03. We observe that 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 difficult 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; McGuffin et al., 2000) , we find that the availability of evolutionary information has significant impact on prediction quality.

4.2. ADVANTAGES 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. First, the computational cost of training and sampling is higher than that of angle-predicting RNNs (Figure 10 ) such as our baseline or AlQuraishi (2018). Consequently, those methods have been scaled to larger datasets than ours (in protein length and diversity) which 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. These approaches can also lead to slower learning and less expressive energy functions. Methods for efficient (i.e. subquadratic) N -body simulations and for more principled stabilization of deep networks may be relevant to addressing both of these challenges in the future.

5. CONCLUSION:

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 efficient simulator for Langevin dynamics in transformed coordinate systems and stabilization techniques for backpropagating through long simulator roll-outs. We find 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 classifications more effectively than a strong baseline. The model generates an atomic structure X (top right) from sequence information s (top left) via 3 steps: First, a sequence network takes in the sequence information s, processes it with a combination of 1D, 2D, and graph convolutions (MPNN, bottom left), and outputs energy function weights l as well as simulator hyperparameters (top center). Second, the simulator iteratively modifies the structure via Langevin dynamics based on the gradient of the energy landscape (Forces, bottom center). Third, the imputation network constructs predicted atomic coordinates X from the final simulator time step x (T ) . During training, the true atomic coordinates X (Data) , predicted atomic coordinates X, simulator trajectory x (1) , . . . , x (T ) , and secondary structure predictions SS (Model) feed into a composite loss function (Loss, bottom right), which is then optimized via backpropagation.

APPENDICES A MODEL:

A.1 COORDINATE SYSTEMS

Inverse transformation:

The inverse transformation z = F 1 (x) involves fully local computations of bong lengths and angles.

b i = ||x i x i 1 ||, a i = arccos ( ûi • ûi 1 ) , d i = sign ( ûi 2 • ni ) arccos ( ni 1 • ni ) .

Jacobian The Jacobian @x @z defines the infinitesimal 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 defined element-wise as

@x j @b i = ⇢ ûi i  j 0 i > j , @x j @a i = ⇢ ni ⇥ (x j x i 1 ) i  j 0 i > j , @x j @d i = ⇢ ûi 1 ⇥ (x j x i 1 ) i  j 0 i > j .

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 ni centered at x i 1 , and moving a dihedral angle d i drives circular motion of downstream coordinate x j around bond vector ûi 1 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 bi and ãi via the transformations b i = log

⇣ 1 + e bi

⌘ and a i = ⇡ 1+e ãi . All references to the internal coordinates z and Jacobians @x @z will refer to the use of fully unconstrained representations (Table 2 ).

A.2 ENERGY 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 

z i = h b i ãi di i T [3,1] Unit vector from x i 1 to x i ûi [3,1] Unit vector normal to bond plane at x i 1 ni [3,1] Bond length ||x i x i 1 || b i [1] Bond angle \( ûi , ûi 1 ) a i [1] Dihedral angle \( ni , ni 1 ) d i [1] Unconstrained bond length bi [1] Unconstrained bond angle ãi [1] Jacobian matrix @x @z [3L,3L]

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.3 INTERNAL 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 quantification, 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 defined as a continuous-time stochastic differential equation, and are simulated in discrete time with the first 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 infinity, 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 efficiency 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

z (t+✏) z (t) ✏ 2 r z U (t) + p ✏p, p ⇠ N (0, I).

The advantages and disadvantages of the two coordinate systems are complementary: Cartesian dynamics efficiently sample local structural rearrangements and inefficiently sample global chain motions, while internal coordinate dynamics efficiently 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

)

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-fly geometry modification.

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 x (t+✏) = F ⇣ z (t) + z (t) ⌘ = x (t) + Z t+✏ t 1 ✏ @x @z (u) z (t) du.

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 first substitute an Euler approximation to predict x (t+✏) as t) , and then substitute the Jacobian evaluated at the predicted state x(t+✏) to form trapezoidal approximation t) . The comparison of this algorithm with naive integration is given in Figure 8 . 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 8 ). 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.

EQUATION

EQUATION

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 3L degrees of freedom, since only 3L 6 degrees of freedom are necessary to encode a centered and un-oriented structure5 . As a consequence, a significant 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 define the non-translational motion as xi = x i

x Trans i and the centered Cartesian coordinates as xi = x i hx i i. The point-wise angular momentum is then l i = xi ⇥ xi and we define a total angular velocity of the structure ! by summing these and dividing by the moment of inertia as ! = (

P i l i ) / P i || xi || 2 2 .

We convert the angular velocity ! into Cartesian displacements with an unrolled Heun integration as

x Rot i = 1 2 ! ⇥ ( xi + ! ⇥ xi )

, which leaves the isolated structural motions as

x Struct i = x i x Trans i x Rot i .

Algorithm 3: Mixed Integrator

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 modified 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., 2017) might accomplish a similar goal in a more principled manner. The final integrator combining these ideas is presented in Figure 3 .

B APPENDIX B: TRAINING:

B.1 DATA For a training and validation set, we downloaded all protein domains of length L  200 from Classes ↵, , 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 difficulty: 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,381 test domains. We further stratified this test set into C, A, T, and H categories based on their nearest CATH classification in the training set.

We also analyzed test set stratifications based on nearest neighbors in both training and validation in figure Figure 12 . We note that the validation set was not explicitly used to tune hyperparameters due to the large cost of training ( 2 months on 2 M40 GPUs), but we did keep track of validation statistics during training.

B.2 SGD:

We optimized all models for 200,000 iterations with Adam (Kingma & Ba, 2014).

B.3 LOSS:

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

X i<j w ij D (Model) ij D (Data) ij

, where the contact-focusing weights are

EQUATION

))

EQUATION

))

⌘ and (u) = 1 1+exp( u) is the sigmoid function.

Angle loss We use the loss

L angles = X i ||H(z (T ) i ) H(z (Data) i )||,

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 ) ). Taking z a as fixed 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.

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 vij (Figure 2C ) and doing the same contact-weighted averaging as the distance loss.

Template Modelling (TM) Score The TM-score (Zhang & Skolnick, 2005) ,

X i 1 1 + ⇣ Di D0 ⌘ 2 ,

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

U h-bond ij (X) = ✓ q N q O D NO + q H q C D HC + q H q O D HO + q N q C D NC ◆ 332 = 0.084 ✓ 1 D NO + 1 D HC 1 D HO 1 D NC ◆ 332

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

y data ij = 1 ⇣ U h-bond ij (X (data) ) < 0.5 ⌘ .

We 'predict' the probabilities of hydrogen bonds of the data given the model via logisitic regression of soft model assignments as

y model ij = ⇣ a ⇣ b ⇣ U h-bond ij (X (model) ) + 0.5 ⌘⌘ + c ⌘ ,

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 final hydrogen bond loss is the cross-entropy between these predictions and the data,

L h-bond = X |i j|>2 y data ij log y model ij + 1 y data ij log 1 y model ij .

Secondary Structure Prediction We output standard 8-class predictions of secondary structure and score them with a cross-entropy loss.

B.4 STABILIZING BACKPROPAGATION THROUGH TIME:

The combination of energy function, simulator, and refinement 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 specifics of the loss function, however, we will discuss a challenges and solutions for computing gradients of unrolled simulations in the face of chaos.

B.5 CHAOS 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 find 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) and are explored in detail in a concurrent work (Parmas et al., 2018) .

In Figure 9 , 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 (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 confined 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. t) dx (0) (bottom). When the step size ↵ is small, these dynamics converge to a periodic orbit over 2 k values where 0  k < 1. After some critical step size, the dynamics undergo a period-doubling bifurcation (Strogatz, 2018), become chaotic, and the gradients regularly diverge to huge numbers.

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

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, 10 4 I)6 . We regularize the sensitivity by adding are stable. To this end we introduce a damping factor to backpropagation that can adaptively combat exponentially diverging gradients with exponential discounting (Algorithm 4).

EQUATION

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 fixed 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 @ x (t) @x (t k) = ✓✓✓ @x (t) @x (t 1) ◆ @x (t 1) @x (t 2)

EQUATION

B.6 MULTIPLE SEQUENCE ALIGNMENT GENERATION:

We use multiple sequence alignments of evolutionarily related sequences for both profile construction ( § B.7) and (ii) data augmentation ( § B.8). 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 hhfilter (Remmert et al., 2012) such that all sequences are at least a normalized Hamming distance of 0.8 away from one another.

B.7 PROFILE GENERATION:

We briefly describe how we construct evolutionary profiles, or position-specific 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

f (i) a = ↵ + P j S (i) ja ↵ + M ,

where ↵ is a pseudocount that we set to 10. We compute our PSSM features for letter a at site i as

w (i) a = log f (i) a B a !

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) . , 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 profile and 1-seq based training.

C APPENDIX C: RESULTS

C.1 STRUCTURE GENERATION AND PROCESSING:

For each sequence from the CATH release 4.2 dataset, 100 structures were generated from both the profile 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.

NEMO RNN, 300 hidden units Figure 10: Sampling speed. Per-protein sampling times for various batch sizes across NEMO and one of the RNN baselines on a single Tesla M40 GPU with 12GB memory and 20 cores. For all results in the main paper, 100 models were sampled per protein followed by consensus clustering with 3D-jury, adding an additional factor of 10 2 cost between NEMO and the RNN.

:

We leverage a traditional estimator of the likelihood gradient (Boltzmann learning) rather than ratio-based estimators

Since our representation x is coarse-grained at one point per position, these are virtual bonds.

Specifically, the six-degrees of freedom parameterizing the rigid body placement of the structure are the'virtual' {b1, a1, a2, d1, d2, d3}.

In TensorFlow this operation is stop gradient.

In Internal coordinates the rotation and translation of the structure are encoded in b1, a1, a2, d1, d,2 , d3 while in Cartesian coordinates they are distributed across all coordinates.

For stochastic processes such as Langevin dynamics that depend on external noise, this must be cached and re-applied.

Figures and Tables:

Figure 2: A neural energy function models coarse grained structure and is sampled by internal coordinate dynamics. (A) The energy function is formulated as a Markov Random Field with structure-based features and sequence-based weights computed by neural networks (Figure 6). (B) To rapidly sample low-energy configurations, the Langevin dynamics simulator leverages both (i) an internal coordinate parameterization, which is more effective for global rearrangements, and (ii) a Cartesian parameterization, which is more effective for localized structural refinement. (C) The base features of the structure network are rotationally and translationally invariant internal coordinates (not shown), pairwise distances, and pairwise orientations.

Figure3: A transform integrator simulates Langevin dynamics in a more favorable coordinate system (e.g. internal coordinates z) directly in terms of the untransformed state variables (e.g. Cartesian x). This exchanges the cost of an inner-loop transformation step (e.g. geometry construction F(z)) for an extra Jacobian evaluation, which is fully parallelizable on modern hardware (e.g. GPUs).

Figure 4: Model generalizes and outperforms end-to-end baseline for unseen fold topologies. Colors indicate varying difficulty levels of protein domains in the test set, with the C (cyan) and A (magenta) subsets containing corresponding to test-set domains with topologies (folds) and superfamilies that were not represented in the training set. (Left) As the model exhibits higher confidence (reduced structural diversity), it becomes more accurate. (Center) The model occasionally achieves TM scores greater than 0.5 even for difficult C and A level generalization tasks. (Right) NEMO outperforms a strong RNN baseline for difficult generalization problems. All results for NEMO and RNN baselines are conditioned on profiles.

Figure 5: Examples of fold generalization at topology and architecture level. These predicted structures show a range of prediction accuracy for structural generalization (C and A) tasks, with the TM-score comparing the top ranked 3D-Jury pick against the target. The largest clusters are the three most-populated clusters derived from 100 models per domain with a within-cluster cutoff of TM > 0.5. CATH IDs: 2oy8A03; 5c3uA02; 2y6xA00; 3cimB00; 4ykaC00; 2f09A00; 3i5qA02; 2ayxA01.

Figure6: Model schematic. The model generates an atomic structure X (top right) from sequence information s (top left) via 3 steps: First, a sequence network takes in the sequence information s, processes it with a combination of 1D, 2D, and graph convolutions (MPNN, bottom left), and outputs energy function weights l as well as simulator hyperparameters (top center). Second, the simulator iteratively modifies the structure via Langevin dynamics based on the gradient of the energy landscape (Forces, bottom center). Third, the imputation network constructs predicted atomic coordinates X from the final simulator time step x(T ) . During training, the true atomic coordinates X(Data) , predicted atomic coordinates X, simulator trajectory x (1) , . . . , x(T ) , and secondary structure predictions SS(Model) feed into a composite loss function (Loss, bottom right), which is then optimized via backpropagation.

Figure 7: Component architectures. (Left) The energy function is the inner product of sequencebased weights and structure-based features. A combination of low-and high-level features capture multi-scale constraints on structure. (Center) The structure network is a lightweight convolutional network operating on both 1D (backbone) and 2D (interaction) features. (Right) Convolutional neural network modules used for sequence processing are composed of residual blocks that interleave spatial convolutions with 1x1 convolutions.

Figure 8: Accounting for second order errors is essential for internal coordinate dynamics. (Top) Discarding the corrector step rapidly accumulates errors due to the curvilinear motions of internal coordinate dynamics. (Bottom) Heun integration with a corrector step accounts for curvature in curvilinear motion.

• 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 )

Figure 9: Chaos impedes meta-learning for gradient descent in a well. (a) Gradient descent of a particle in a well with initial conditions x (0) and step size ↵. (b) Orbit diagrams visualize long-term dynamics from iterations 1000 to 2000 of the position x (top) and the gradient dx(t)  dx (0) (bottom). When the step size ↵ is small, these dynamics converge to a periodic orbit over 2 k values where 0  k < 1. After some critical step size, the dynamics undergo a period-doubling bifurcation (Strogatz, 2018), become chaotic, and the gradients regularly diverge to huge numbers.

Figure 9: Chaos impedes meta-learning for gradient descent in a well. (a) Gradient descent of a particle in a well with initial conditions x (0) and step size ↵. (b) Orbit diagrams visualize long-term dynamics from iterations 1000 to 2000 of the position x (top) and the gradient dx(t)  dx (0) (bottom). When the step size ↵ is small, these dynamics converge to a periodic orbit over 2 k values where 0  k < 1. After some critical step size, the dynamics undergo a period-doubling bifurcation (Strogatz, 2018), become chaotic, and the gradients regularly diverge to huge numbers.

Test set performance across different levels of generalization

Table content: <table><tr><td>Model</td><td colspan="2"># params Total</td><td>C</td><td>A</td><td>T</td><td>H</td></tr><tr><td>NEMO (ours, profile)</td><td>21.3m</td><td colspan="5">0.366 0.274 0.361 0.331 0.431</td></tr><tr><td>NEMO (ours, sequence-only)</td><td>19.1m</td><td colspan="5">0.248 0.198 0.245 0.254 0.263</td></tr><tr><td>RNN baseline model (profile)</td><td/><td/><td/><td/><td/></tr><tr><td>2x100</td><td>5.9m</td><td colspan="5">0.293 0.213 0.230 0.247 0.388</td></tr><tr><td>2x300 (avg. of 3)</td><td>8.8m</td><td colspan="5">0.335 0.229 0.282 0.278 0.446</td></tr><tr><td>2x500</td><td>13.7m</td><td colspan="5">0.347 0.222 0.272 0.286 0.477</td></tr><tr><td>2x700</td><td>21.4m</td><td colspan="5">0.309 0.223 0.259 0.261 0.403</td></tr><tr><td>Number of structures</td><td/><td colspan="2">10381 1537</td><td>1705</td><td>3198</td><td>3941</td></tr></table>

Coordinate systems and representations for protein structure.

Table content: <table><tr><td>Variable</td><td>Notation</td><td>Shape</td></tr><tr><td>Sequence</td><td>s</td><td>[L, 20]</td></tr><tr><td>Cartesian coordinates (coarse)</td><td>x</td><td>[3L,1]</td></tr><tr><td>Internal coordinates</td><td>z</td><td>[3L,1]</td></tr><tr><td>Cartesian coordinates (atomic)</td><td>X</td><td>[3L,A]</td></tr><tr><td>Cartesian coordinates for position i</td><td>x i</td><td>[3,1]</td></tr><tr><td>Internal coordinate for position i</td><td/><td/></tr></table>

Model architecture. Input number of channels q = 20 for sequence-only and q = 40 for profiles.

Table content: <table><tr><td>Location</td><td colspan="5">Type Channels # Blocks Width Dilation</td><td>Stride</td></tr><tr><td>Pre-MPNN MPNN</td><td>1D 1D</td><td>128 128</td><td>12 4</td><td>3 3</td><td colspan="2">[1, 2, 4, 8] ⇥ 3 1 [1, 2, 4, 8] 1</td></tr><tr><td>MPNN</td><td>2D</td><td>50</td><td>1</td><td>7</td><td>1</td><td>1</td></tr><tr><td colspan="2">Post-MPNN Post-MPNN* 2D 1D</td><td>q+256 100</td><td>12 1</td><td>3 9</td><td colspan="2">[1, 2, 4, 8] ⇥ 3 1 1 1</td></tr><tr><td>Imputation</td><td>1D</td><td>q+256</td><td>12</td><td>3</td><td colspan="2">[1, 2, 4, 8] ⇥ 3 1</td></tr></table>

Damped Backpropagation Through TimeInput :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 << < 1 Output :Exponentially damped gradient r ✓ L Compute loss L(x 1 , . . . , x T ) ; Compute gradient r ✓ L AutoDiff(L, ✓) ; † where ?(•) is the stop gradient function.

Table content: <table><tr><td>Initialize x (0) for t 2, . . . , T do F(z (0) );</td><td/><td/></tr><tr><td>Compute time step xt</td><td colspan="2">F (x t 1 , s t , ✓);</td></tr><tr><td>Decay the gradient  † x t</td><td>(1</td><td>)?(x t ) + xt ;</td></tr><tr><td>end</td><td/><td/></tr><tr><td/><td/><td>12)</td></tr><tr><td colspan="3">to the loss. Interestingly, the stochastic nature of this approximate regularizer is likely a good thing -a</td></tr><tr><td colspan="3">truly non-expansive map is quite limited in what it can model. However, being 'almost' non-expansive</td></tr><tr><td colspan="3">seems to be incredibly helpful for learning.</td></tr></table>

Qualitative timings. † Results on CATH dataset and 2 M40 GPUs.To reduce our reliance on alignments and the generation of profiles 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

Table content: <table><tr><td>Method</td><td>Generation time</td><td>Training time</td></tr><tr><td colspan="3">RNN baseline  † NEMO  † Coevolution-based methods minutes to hours Coupled to generation milliseconds ⇠ 1 week seconds ⇠ 2 months</td></tr><tr><td>Physical simulations</td><td>days to weeks</td><td>N/A</td></tr><tr><td colspan="2">B.8 EVOLUTIONARY DATA AUGMENTATION</td><td/></tr><tr><td>p t =</td><td>e EDAdt P M s=1 e EDAds</td><td/></tr></table>

Peer review (iclr2019):
 Post-rebuttal revision: The authors have adressed my concerns sufficiently. The paper still has issues with presentation, and weak comparisons to earlier methods. However, the field is currently rapidly developing, and comparing to earlier works is often difficult. I believe the Langevin-based prediction is a significant and clever contribution. I*m raising my score to 6. ------ The paper proposes an end-to-end neural architecture for learning protein structures from sequences. The problem is highly important. The method proposes to use a Langevin simulator to fold the protein ‘in silico’ from some initial state, proposes numerous tricks for the optimisation, and proposes neural networks to extract information from both the sequence and the fold state (energy function). The system works on internal coordinates, which are conditioned and integrated on the fly. The method seems to perform very well, improving upon their baseline model considerably. In spite of the paper being an outstanding work, I have two criticisms about the accessibility and impact of the paper on the broader ICLR audience. In its current form and complexity, the paper feels accessible mostly to a narrow audience. First, the framework proposed in the paper is massive, containing a large amount of components, neural networks, simulators, integrators, optimisation tricks, alignments, profiles, stabilizations, etc. The amount of work done in the manuscript is staggering, but the method is also difficult to understand from reading the main manuscript alone. The 10+ page appendix is critical for understanding (for instance, the appendix reveals that MSA is used to generate more data), and even with it the method is difficult to grasp as a whole. This paper should be presented in a journal form with a presentation not hindered by page limits, while currently one needs to jump between the main text and appendix to get the whole picture. I also wonder if some parts of the system have already been published, and perhaps the presentation could be condensed that way. Second, the introduction lists numerous competing methods both on the protein modelling side and on the MCMC vs optimisation side. The paper does not compare to any of these, which is strange, and makes it difficult to assess how much this paper improves upon state-of-the-art. Right now its unclear what is state-of-the-art in general. No bigger context of protein folding is given either, for instance, how well the method fares against purely alignment based approaches, or against purely physics-based simulators. Finally, the experimental section poorly describes how all the pieces of the system affect the final predictions. The discussion on the exploding gradients and dampening is excellent however. The only baseline is one with the simulator replaced by an RNN. There does not seem to be any running time analyses. As such, it is hard to interpret the current system, and it feels like a black box.