Neural Energy Minimization for Molecular Conformation Optimization

1 INTRODUCTION . The 3D molecular conformation is one of the most important features in determining many physicochemical and biological properties of a molecule . The molecule ’ s 3D charge distribution and physical shape are crucial when considering the steric constraints or electronic effects for chemical reactions and interaction . Therefore , molecular conformers are widely adopted in quantitative structure-activity relationships ( QSAR ) prediction ( Verma et al. , 2010 ) and drug discovery ( Hawkins , 2017 ) . With a growing interest in virtual drug screening ( Zhavoronkov et al. , 2019 ; McCloskey et al. , 2020 ; Stokes et al. , 2020 ) and de novo drug design ( De Cao & Kipf , 2018 ; Gómez-Bombarelli et al. , 2018 ; Jin et al. , 2018 ; Zhou et al. , 2019 ; Mercado et al. , 2020 ) , it becomes more desirable to predict the molecular conformations both quickly and accurately such that 3D geometry features can be considered . Deep learning techniques have recently been introduced to address this problem and demonstrated promising results . For instance , generative models ( Hoffmann & Noé , 2019 ; Simm & HernandezLobato , 2020 ; Xu et al. , 2020 ) are proposed to first generate pairwise distances for atom pairs , and then infer the conformer from such matrix.Since the 3D structures are generated indirectly , the second step can be sensitive towards the error in estimated distances , and extra model capacity could be required to encode the redundancy ( Hoffmann & Noé , 2019 ) . In the meantime , many SE ( 3 ) -equivariant networks emerge to work with 3D roto-translation symmetrical inputs enabling direct manipulation of the 3D Cartesian coordinates . One of the most recent successful applications ∗Equal Contributions . of such models is the AlphaFold system for protein structure prediction ( Jumper et al. , 2021 ) . Despite the many benefits , however , existing SE ( 3 ) -equivariant networks are either derived from complicated mathematical theory involving expensive coefficient calculations ( Thomas et al. , 2018 ; Anderson et al. , 2019 ; Finzi et al. , 2020 ; Fuchs et al. , 2020 ) , or heuristically designed ( Satorras et al. , 2021 ) based on the intuition of message passing ( Gilmer et al. , 2017 ) . This work presents a new formulation that unrolls the molecular conformer prediction into a fixed number of optimization steps , where a neural network is parametrized to learn gradient fields of the implicit conformational energy landscape . Under such formulation , the model can refine a given conformer toward a more stable and energetically preferred state . Instead of relying on complex mathematical tools or designing the model heuristically , we propose a novel energy-inspired framework that explicitly connects an underlying enregy function with an SE ( 3 ) -equivariant model . We are able to not only provide new interpretations for some existing models but also derive new variants of SE ( 3 ) -equivariant model from the assumption of a conformational energy function , aligning a model ’ s architecture with its physical assumptions . Through extensive experiments , these new variants show superior performance in conformer optimization comparing to other SE ( 3 ) equivariant models . Moreover , the proposed optimization method is also suitable for conformer generation and capable of generating diverse and accurate conformers . In conclusion , the main contributions of this work include : • Formulating the conformer prediction problem as a multi-step unrolled optimization where a model learns the gradient field of the conformational energy landscape . • Proposing a novel framework that explicitly connects an underlying energy function with an SE ( 3 ) -equivariant neural network such that one can derive new models that align with a physical assumption , or interpret some of the existing models from an energy perspective . • Demonstrating competitive empirical performance for molecular conformation prediction in both the optimization and generation settings . 2 RELATED WORK . 2.1 MOLECUALR CONFORMATION GENERATION . In conventional computational approaches for conformer generation , we observe an accuracy and efficiency trade-off where highly accurate ab initio methods are very time-consuming while efficient methods leveraging empirical force fields are not accurate enough ( Unke et al. , 2021 ) . To efficiently and accurately estimates molecular conformations , many learning-based methods have been proposed . For instance , models , such as VAE ( Simm & Hernandez-Lobato , 2020 ) , GAN ( Hoffmann & Noé , 2019 ) , and continuous normalizing flow ( Xu et al. , 2020 ) , have all been developed to predict molecular conformers by first generating an intermediate pairwise atom distance matrix . However , the postprocessing step for 3D structure construction can be susceptible to the error , and model capacity could be wasted to encode the redundancy in such low-rank matrix ( Hoffmann & Noé , 2019 ) . Despite the many drawbacks of such methods , the distance matrix is difficult to circumvent because the models can not capture the 3D roto-translation symmetry nature of a conformer . Besides modeling the atom distances , another solution for conformer generation is to estimate the force field with machine learning methods . For instance , sGDML ( Chmiela et al. , 2019 ) is capable of reproducing the molecular potential energy surface given a few reference conformations and generating compelling results with MD simulation . Although such a model scales well with the molecule size , it requires retraining for every new molecules ( Unke et al. , 2021 ) . Similarly but more generally , methods including Behler & Parrinello ( 2007 ) ; Bartók et al . ( 2010 ; 2013 ) ; Smith et al . ( 2017 ) ; Schütt et al . ( 2017 ) ; Faber et al . ( 2018 ) can also be used to generate molecular conformers with their learned force fields . Recently , ConfGF ( Shi et al. , 2021 ) specifically tackled the conformer generation problem by estimating the gradient fields of inter-atomic distances ( similar to force fields ) , where the 3D coordinates of the conformer can be generated through steps of Langevin dynamics achieving state-of-the-art performance . While our model also tries to estimate the gradient field , we instead employ an unrolled optimization architecture ( Domke , 2012 ; Liu et al. , 2018 ) aiming to model the entire optimization process instead of a single force field optimization step . Such one-pass end-to-end design is proved to be more effective empirically as shown in the experiment section . In addition , our approach also enjoys a much faster run time . 2.2 SE ( 3 ) -EQUIVARIANT NEURAL NETWORKS . Since 3D roto-translation symmetry is a common property for many physical systems , various methods have been proposed to model such symmetry in problems like quantum mechanics ( Ramakrishnan et al. , 2014 ) , molecular design ( Simm et al. , 2020 ) , particle simulations ( Thomas et al. , 2018 ) , and point clouds ( Uy et al. , 2019 ) . The ability to capture the spatial symmetry makes these models quite suitable for conformer optimization . In general , there are two ways to introduce the equivariance into a model . One thread of research relies on building models with irreducible representations and spherical harmonics ( Thomas et al. , 2018 ; Anderson et al. , 2019 ; Finzi et al. , 2020 ; Fuchs et al. , 2020 ; Batzner et al. , 2021 ) . Different from this line of work , our approach aims to avoid the complexity of such mathematical tools , considering the higher-order features in these more theoretical framework might not be practical for our purpose . The other line of research relies on vector-based methods , where equivariant filters are constructed by taking the gradient of invariant filters ( Schütt et al. , 2021 ) or designed ( Satorras et al. , 2021 ) based on the intuition of message passing ( Gilmer et al. , 2017 ) . By explicitly connecting an underlying energy function with an SE ( 3 ) -equivariant neural network in our framework , we can not only reinterpret some of the existing models but also derive new variants that aligns the model parametrization with an underlying physical assumption as shown in Section 3.3 . 3 METHODS . 3.1 PROBLEM DEFINITION . We represent each molecule as a graph Gmol and use a set of 3D coordinates X to denote the conformer . The molecular graph is defined as Gmol = 〈V , E〉 where V = { vi } is the set of nodes representing atom features , and E = { ei } is the set of edges representing the relationship between atoms : including covalent interactions and distance of their shortest path . X consists of the 3D coordinate { xi|xi ∈ R3 } of each atom . Given a molecular graph Gmol and an initial conformer X0 , we want to learn an optimization model φ that refines the initial conformation towards an energetically preferred state X∗ as X∗ = arg minX E ( X|Gmol ) where E is some energy surface of nature . 3.2 CONFORMATIONAL ENERGY MINIMIZATION WITH A NEURAL NETWORK . Molecular conformation represents a 3D arrangement of the atoms and shares a close relationship with the conformational energy , as each equilibrium of a molecular conformation can be considered as a local minimum of the potential energy surface ( Axelrod & Gomez-Bombarelli , 2020 ) . One could therefore find these stable conformations through gradient descent as a process of energy minimization . Specifically , given a particular energy function E , we can optimize the conformation X by updating the 3D Cartesian coordinates with gradient descent such that E can be minimized : Xt+1 = Xt − γ∇XE ( Xt|Gmol ) ( 1 ) where γ is the learning rate , and we will observe X∗ = X∗− γ∇XE ( X∗|Gmol ) for an energetically stable conformer . Therefore , we can easily find a stable molecular conformation from an initial conformation by applying the gradient descent update iteratively if the energy function E ( X|Gmol ) is fully characterized and differentiable with respect to X . Unfortunately , this is often not the case . Instead of hand-designing the energy function and its gradient , we can take advantage of the approximation power of a neural network and use it to learn the gradient field from the data directly . Once the network is trained , we can then unroll the optimization ( Domke , 2012 ; Liu et al. , 2018 ) with the model and perform a neural version of the above optimization where the gradient term ∇E in Equation 1 is not derived from an explicit energy function but estimated by a neural network . More concretely , we parameterize the energy surface as a function of the molecular graph G = 〈V , E〉 and the 3D conformer X . We can consider X as an external 3-dimensional representation of the molecule , while V and E represent an internal high dimensional representation of the molecule . These internal embeddings V and E not only encode the initial atom/bond information but also embed information about their neighborhood related to the evolving conformation X . For each step , we want to minimize E ( X , V , E ) with the following gradient updates : xt+1i = x t i + φx ( Gt , Et ) vt+1i = v t i + φv ( Gt , Et ) et+1ij = e t ij + φe ( Gt , Et ) for all xi ∈ X , vi ∈ V , and eij ∈ E , where φ are neural networks estimating their respective gradient field , and the exact parametrization of the neural network will depend on the energy function as described in the following section . In this work , we view eij as a constant representation for simplicity , and therefore , it is not updated in our models .

Summary: This paper proposes a new method to generate molecular conformations (i.e., the spatial arrangements of atoms belonging to a given molecule). The method works by updating initial atom positions (which can be initialized either randomly or using an alternative conformer generation method) through an iterative process. Each iteration involves predicting the gradients associated with atoms' positions and then changing their coordinates by taking a step in this direction (Eqn. 1 & Fig. 1b).   The authors: 1. show how the parameterization of the gradients relates to how one models the conformational energy (Section 3.3); 2. demonstrate competitive performance on optimization (finding a single best conformer; Section 4.1) and generation (finding a set of relatively stable conformers; Section 4.2) tasks compared to previous benchmarks. This is judged in terms of (a) how well the generated conformers match up to ground truth conformers (e.g., Table 1) as well as (b) how well the generated conformers enable the prediction of molecular properties (Table 2).