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 deep neural network (DNN) architecture to generate molecular 3D conformation. It derives its architecture starting from gradient-based updates for minimization of the molecule's energy function. Then the architecture becomes a repetition of vertex-wise aggregation layer of SE(3)-equivariant Transformer-based DNN. The architecture takes a molecule and 3D conformation pre-calculated from RDkit as input. The authors show how the proposed architecture generalize existing DNN architectures for molecular conformation generation. Empirical results demonstrate superiority of the proposed method.