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-to-end-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. 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 DISPLAYFORM0 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 BID0 BID17 . 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 An unrolled simulator as a model for protein structure. NEMO combines a neural energy function for coarse protein structure, a stochastic simulator based on Langevin dynamics with learned (amortized) initialization, and an atomic imputation network to build atomic coordinate output from sequence information. It is trained end-to-end by backpropagating through the unrolled folding simulation.protein folding (Dill et al., 2017) , in which the folds that natural protein sequences adopt are those that minimize free energy. Without the availability of external information such as coevolutionary information (Marks et al., 2012) or homologous structures (Martí-Renom et al., 2000) to constrain the energy function, however, contemporary simulations are challenged to generate globally favorable low-energy structures in available time.How can we get the representational benefits of energy-based models with the sampling efficiency 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 fixed amount of time. Leveraging this idea, we construct an end-to-end differentiable model of protein structure that is trained by backpropagtion through folding ( FIG0 ). 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• Efficient sampling algorithm that is based on a transform integrator for efficient 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 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. (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.
We use an unrolled simulator as an end-to-end differentiable model of protein structure and show it can (sometimes) hierarchically generalize to unseen fold topologies.
{"src": "scitldr-AIC_train", "paper_id": "Byg3y3C9Km"}