Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning

By Frank No'e, Jonas Köhler, Hao Wu. 2019

873 citations recorded by the source corpus.

Efficient sampling of equilibrium states Molecular dynamics or Monte Carlo methods can be used to sample equilibrium states, but these methods become computationally expensive for complex systems, where the transition from one equilibrium state to another may only occur through rare events. Noé et al. used neural networks and deep learning to generate distributions of independent soft condensed-matter samples at equilibrium (see the Perspective by Tuckerman). Supervised training is used to construct invertible transformations between the coordinates of the complex system of interest and simple Gaussian coordinates of the same dimensionality. Thus, configurations can be sampled in this simpler coordinate system and then transformed back into the complex one using the correct statistical weighting. Science, this issue p. eaaw1147; see also p. 982 By combining deep learning and statistical mechanics, neural networks sample the equilibrium distribution of many-body systems. INTRODUCTION Statistical mechanics aims to compute the average behavior of physical systems on the basis of their microscopic constituents. For example, what is the probability that a protein will be folded at a given temperature? If we could answer such questions efficiently, then we could not only comprehend the workings of molecules and materials, but we could also design drug molecules and materials with new properties in a principled way. To this end, we need to compute statistics of the equilibrium states of many-body systems. In the protein-folding example, this means to consider each of the astronomically many ways to place all protein atoms in space, to compute the probability of each such “configuration” in the equilibrium ensemble, and then to compare the total probability of unfolded and folded configurations. As enumeration of all configurations is infeasible, one instead must attempt to sample them from their equilibrium distribution. However, we currently have no way to generate equilibrium samples of many-body systems in “one shot.” The main approach is thus to start with one configuration, e.g., the folded protein state, and make tiny changes to it over time, e.g., by using Markov-chain Monte Carlo or molecular dynamics (MD). However, these simulations get trapped in metastable (long-lived) states: For example, sampling a single folding or unfolding event with atomistic MD may take a year on a supercomputer. RATIONALE Here, we combine deep machine learning and statistical mechanics to develop Boltzmann generators. Boltzmann generators are trained on the energy function of a many-body system and learn to provide unbiased, one-shot samples from its equilibrium state. This is achieved by training an invertible neural network to learn a coordinate transformation from a system’s configurations to a so-called latent space representation, in which the low-energy configurations of different states are close to each other and can be easily sampled. Because of the invertibility, every latent space sample can be back-transformed to a system configuration with high Boltzmann probability (Fig. 1). We then employ statistical mechanics, which offers a rich set of tools for reweighting the distribution generated by the neural network to the Boltzmann distribution. RESULTS Boltzmann generators can be trained to directly generate independent samples of low-energy structures of condensed-matter systems and protein molecules. When initialized with a few structures from different metastable states, Boltzmann generators can generate statistically independent samples from these states and efficiently compute the free-energy differences between them. This capability could be used to compute relative stabilities between different experimental structures of protein or other organic molecules, which is currently a very challenging problem. Boltzmann generators can also learn a notion of “reaction coordinates”: Simple linear interpolations between points in latent space have a high probability of corresponding to physically realistic, low-energy transition pathways. Finally, by using established sampling methods such as Metropolis Monte Carlo in the latent space variables, Boltzmann generators can discover new states and gradually explore state space. CONCLUSION Boltzmann generators can overcome rare event–sampling problems in many-body systems by learning to generate unbiased equilibrium samples from different metastable states in one shot. They differ conceptually from established enhanced sampling methods, as no reaction coordinates are needed to drive them between metastable states. However, by applying existing sampling methods in the latent spaces learned by Boltzmann generators, a plethora of new opportunities opens up to design efficient sampling methods for many-body systems. Boltzmann generators overcome sampling problems between long-lived states. The Boltzmann generator works as follows: 1. We sample from a simple (e.g., Gaussian) distribution. 2. An invertible deep neural network is trained to transform this simple distribution to a distribution pX(x) that is similar to the desired Boltzmann distribution of the system of interest. 3. To compute thermodynamics quantities, the samples are reweighted to the Boltzmann distribution using statistical mechanics methods. Computing equilibrium states in condensed-matter many-body systems, such as solvated proteins, is a long-standing challenge. Lacking methods for generating statistically independent equilibrium samples in “one shot,” vast computational effort is invested for simulating these systems in small steps, e.g., using molecular dynamics. Combining deep learning and statistical mechanics, we developed Boltzmann generators, which are shown to generate unbiased one-shot equilibrium samples of representative condensed-matter systems and proteins. Boltzmann generators use neural networks to learn a coordinate transformation of the complex configurational equilibrium distribution to a distribution that can be easily sampled. Accurate computation of free-energy differences and discovery of new configurations are demonstrated, providing a statistical mechanics tool that can avoid rare events during sampling without prior knowledge of reaction coordinates.

Subfields: Atomic and Molecular Physics, and Optics, Materials Chemistry, Molecular Biology, Generative Models, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Atomic and Molecular Physics, and Optics

In Materials Chemistry

In Molecular Biology

In Generative Models

In Probabilistic & Bayesian Methods

Citations in the corpus

Stochastic Normalizing Flows (2020, methodology): “∆ S t can be interpreted as the entropy in the variational free energy of the ﬂow transformation (No´e et al., 2019; Li & Wang, 2018).”
VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models (2020, background): “…(Kingma et al., 2014; Izmailov et al., 2020) to 3D point cloud generation (Yang et al., 2019), light-transport simulation (Müller et al., 2019), molecular modeling (Sanchez-Lengeling & AspuruGuzik, 2018; Noé et al., 2019) and equivariant sampling in theoretical physics (Kanwar et al., 2020).”
Equivariant Flows: exact likelihood generative learning for symmetric densities (2020, methodology): “, 2017), which can be made exact using importance sampling (Müller et al., 2018; Noé et al., 2019).”
Equivariant Flows: sampling configurations for multi-body systems with symmetric energies (2019, methodology): “Two obvious options are available in order to train the generator to match qX,θ ≈ pX [1, 12]: 1.”
GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation (2022, methodology): “Given multiple graphs G , and for each G given its conformations C as i.i.d samples from an underlying Boltzmann distribution (Noé et al., 2019), our goal is learning a generative model p θ ( C|G ) , which is easy to draw samples from, to approximate the Boltzmann function.”
Torsional Diffusion for Molecular Conformer Generation (2022, methodology): “This connects with the literature on Boltzmann generators —generative models which aim to sample the Boltzmann distribution of physical systems without expensive molecular dynamics or MCMC simulations [Noé et al., 2019, Köhler et al., 2021].”
Inverse design of 3d molecular structures with conditional generative neural networks (2021, background): “This includes specifically designed approaches to translate given molecular graphs to 3d conformations [32–38], map from coarse-grained to fine-grained structures [39], sample unbiased equilibrium configurations of a given system [40, 41], or focus on protein folding [42–46].”
Protein storytelling through physics (2020, background): “One new accelerator of molecular dynamics is Boltzmann generators, which use deep learning methods to learn where to find the deepest free-energy wells on landscapes ( 129 ) and improve sampling ( 130 ).”
Learning Neural Generative Dynamics for Molecular Conformation Generation (2021, background): “Some other works (Lemke & Peter, 2019; AlQuraishi, 2019; Ingraham et al., 2019; No´e et al., 2019; Senior et al., 2020) also focus on learning models to directly generate 3D structure, but focus on the protein folding problem.”
An introduction to deep generative modeling (2021, background): “Those include new scientiﬁc applications of DGMs, for example, in physics and computational chemistry [8, 33, 6].”
A Unifying Review of Deep and Shallow Anomaly Detection (2020, background): “Consequently, the layers of flow models are usually designed so that the Jacobian is guaranteed to be upper (or lower) triangular or have some other nice structure such that one does not need to compute the full Jacobian to evaluate its determinant [282], [285], [286] (see [287] for an application i”