An Autoregressive Flow Model for 3D Molecular Geometry Generation from Scratch

We consider the problem of generating 3D molecular geometries from scratch . While multiple methods have been developed for generating molecular graphs , generating 3D molecular geometries from scratch is largely under-explored . In this work , we propose G-SphereNet , a novel autoregressive flow model for generating 3D molecular geometries . G-SphereNet employs a flexible sequential generation scheme by placing atoms in 3D space step-by-step . Instead of generating 3D coordinates directly , we propose to determine 3D positions of atoms by generating distances , angles and torsion angles , thereby ensuring both invariance and equivariance properties . In addition , we propose to use spherical message passing and attention mechanism for conditional information extraction . Experimental results show that G-SphereNet outperforms previous methods on random molecular geometry generation and targeted molecule discovery tasks . 1 INTRODUCTION . Designing and synthesizing novel molecules with desirable properties is a challenging task in drug discovery ( Wang et al. , 2020 ; Stokes et al. , 2020 ) and chemical science . The size of search space of all chemical molecules is estimated to be on the order of 1033 ( Polishchuk et al. , 2013 ) , thereby making exhaustive search infeasible . In recent years , advances in machine learning methods have greatly accelerated the progress in this filed . Many studies represent molecules as 2D molecular graphs and propose to automatically generate molecular graphs and optimize molecular properties with deep generative models , such as variational auto-encoders ( Kingma & Welling , 2014 ) . However , complete information of molecules can not be obtained from 2D molecular graphs because 3D structures , also known as 3D molecular geometries , are critical in determining many molecular properties . 3D molecular geometries represent the 3D coordinates of atoms and are important for the accurate prediction of quantum properties ( Schütt et al. , 2017 ) . Hence , we argue that generating 2D molecular graphs might not be the best way to identify novel molecules with certain desirable quantum properties . Instead , developing a generative model that can generate 3D molecular geometries from scratch is a promising solution to this problem . Currently , this area remains under-explored . Recently , a series of seminal studies ( Xu et al. , 2021b ; a ; Shi et al. , 2021 ) have proposed to generate 3D molecular geometries from given 2D molecular graphs . These methods themselves do not generate novel molecules once the 2D molecular graphs are given . In this work , we propose G-SphereNet , a Generative model for 3D molecular geometry generation from scratch inspired by SphereNet ( Liu et al. , 2021d ) . In G-SphereNet , 3D molecular geometries are generated by sequentially placing atoms in 3D space . The 3D positions of atoms are implicitly determined by generating distances , angles and torsion angles to ensure both invariance and equivariance properties . Our work is inspired by SphereNet , which use distances , angles and torsion angles to compute predictive representations of molecules . In addition , G-SphereNet employs SphereNet and attention mechanism to extract conditional information . Experimental results show that G-SphereNet can outperform prior methods on the 3D molecular geometry generation task . 2 BACKGROUND AND RELATED WORK . 2.1 2D MOLECULAR GRAPH GENERATION . Recently , thanks to the advances of deep generative models , significant progress has been made in the problem of molecule design and generation . Some methods ( Gómez-Bombarelli et al. , 2018 ; Kusner et al. , 2017 ; Dai et al. , 2018 ; Liu et al. , 2020 ) use sequence models to generate SMILES ( Weininger , 1988 ) string representations of molecules . Other studies consider molecules as graphs , in which atoms and chemical bonds of the molecule are represented by nodes and edges , respectively . These studies either generate the node type and adjacency matrix of the graph ( Simonovsky & Komodakis , 2018 ; Ma et al. , 2018 ; De Cao & Kipf , 2018 ; Liu et al. , 2021c ) , or form the molecular graph by sequentially adding nodes and edges ( You et al. , 2018 ; Shi * et al. , 2020 ) , or compose the molecule from a junction tree of molecular motifs ( Jin et al. , 2018 ) . However , these methods only generate the graph structures of molecules while crucial 3D molecular geometries are ignored . In other words , the 3D coordinates of atoms in the molecule are unknown . Hence , these generative models can not distinguish spatial isomers , i.e. , molecules with the same molecular graph but different 3D molecular geometries . In addition , 3D molecular geometries are required for the computation of some quantum properties of molecules , such as HOMO-LUMO gap ( Hu et al. , 2021 ; Liu et al. , 2021a ; Ying et al. , 2021 ; Addanki et al. , 2021 ; Xu et al. , 2021c ) . Hence , these generation methods can not be used when spatial isomers or quantum properties are needed . 2.2 3D MOLECULAR GEOMETRY GENERATION FROM 2D INFORMATION . Currently , some studies have proposed to generate 3D molecular geometries from the 2D conditional information of target molecules . Some methods ( Simm et al. , 2020 ; 2021 ) propose to generate 3D molecular geometries through minimizing the energy of atomic systems , in which the numbers and types of atoms are explicitly given in atom bags . Other methods ( Mansimov et al. , 2019 ; Simm & Hernandez-Lobato , 2020 ; Gogineni et al. , 2020 ; Xu et al. , 2021b ; a ; Shi et al. , 2021 ; Ganea et al. , 2021 ) propose to randomly sample multiple 3D molecular geometries from the molecular graph of the target molecule with deep generative models . These conditional generation methods assume that we are given the target molecules in the form of molecular graphs , but we do not know their 3D geometries . However , this assumption does not hold in the targeted molecule discovery problem . In this problem , we aim to discover novel molecules with good quantum properties , such as low HOMO-LUMO gaps . In other words , target molecules themselves are unknown and to be generated , and their 3D geometries are also needed because quantum properties can not be accurately estimated only from molecular graphs . Hence , we argue that developing a method to generate 3D molecular geometries from scratch is more useful to this problem . 2.3 3D MOLECULAR GEOMETRY GENERATION FROM SCRATCH . In this work , we consider the problem of generating 3D molecular geometries from scratch . Let G = { Gj } mj=1 be a set of 3D molecular geometries , and the function S ( G ) ∈ R computes a specific quantum property score of G. We consider the two generation tasks defined as below . • Learning a random generation model pθ ( · ) from G , so that the model can sample a valid 3D molecular geometry G with a high probability pθ ( G ) . • Learning a targeted molecule discovery model pθ ( · ) so as to maximize ( or minimize ) the expected quantum property score EG∼pθ [ S ( G ) ] . This problem is largely under explored and there are only a few studies attempting to solve it . GSchNet ( Gebauer et al. , 2019 ) uses an autoregressive model based on SchNet ( Schütt et al. , 2017 ) to sequentially generate the new atom and place it at the local grid point of a focal atom . On the other hand , EDMNet ( Hoffmann & Noé , 2019 ) and 3DMolNet ( Nesterov et al. , 2020 ) generate pairwise distances between atoms with generative adversarial networks ( GAN ) ( Goodfellow et al. , 2014 ) and variational auto-encoders ( VAE ) ( Kingma & Welling , 2014 ) , respectively . Besides , E-NFs ( Satorras et al. , 2021a ) proposes a geometry generation model by combining flow models ( Rezende & Mohamed , 2015 ) with E ( n ) equivariant graph neural networks ( Satorras et al. , 2021b ) . It generates 3D coordinates of all atoms in a one-shot fashion and defines the prior distribution in a subspace of the latent space to ensure translation invariance . Different from these methods , our method employs a flexible sequential generation pipeline based on autoregressive flow models , which can capture the density of 3D molecular geometries more effectively . 2.4 FLOW MODELS . A flow model defines a parametric invertible mapping fθ : z ∈ Rd → x ∈ Rd , where the data point x and the latent variable z are both random variables . Given that z is sampled from a known prior distribution pZ and fθ is invertible , we can compute the log-likelihood of x as log pX ( x ) = log pZ ( f−1θ ( x ) ) + log |detJ | , ( 1 ) where J = ∂f −1 θ ( x ) ∂x is the Jacobian matrix . To train the flow model on a given dataset X = { xi } m i=1 , the log-likelihoods of data points are computed from Eqn . ( 1 ) and maximized by gradient descent . Hence , tractable and cheap computation of detJ is needed for efficient training . A common choice of fθ in most flow models is the affine coupling mapping ( Dinh et al. , 2014 ; 2016 ) , in which case computing detJ is very easy because J is an upper triangular matrix . Flow models have been used in a variety of generation tasks ( Rezende & Mohamed , 2015 ; Tran et al. , 2019 ; Köhler et al. , 2020 ) . Compared with VAE and GAN , they allow for exact likelihood computation , and can model the density of data more accurately . Because of these advantages , many recent studies have used flow models in the molecule generation task . Some one-shot methods , including GraphNVP ( Madhawa et al. , 2019 ) , GRF ( Honda et al. , 2019 ) , and MoFlow ( Zang & Wang , 2020 ) , consider the node type and adjacency matrix as the generation targets . On the other hand , GraphAF ( Shi * et al. , 2020 ) and GraphDF ( Luo et al. , 2021 ) generate molecular graphs by generating nodes and edges sequentially through autoregressive flow models ( Papamakarios et al. , 2017 ) . These models have stronger capacity to model graph structures than one-shot methods , and achieve state-of-the-art performance in the molecular graph generation task . 3 METHODS . While autoregressive flow models have been successfully applied to the molecular graph generation task , it remains unclear whether they are sufficiently powerful to model more complicated 3D molecular geometries . In this section , we present G-SphereNet , a novel 3D molecular geometry generation method . It adopts a flexible , effective , and efficient sequential generation pipeline based on autoregressive flow models , which can ensure the equivariance property of coordinates and the invariance property of likelihood simultaneously . In addition , expressive spherical message passing based graph neural networks ( Liu et al. , 2021d ) and multi-head attention networks ( Vaswani et al. , 2017 ) are used in G-SphereNet to extract 3D conditional information for accurate generation . To the best of our knowledge , G-SphereNet is the first likelihood-based autoregressive generative model for 3D molecular geometry generation . 3.1 SEQUENTIAL GENERATION . Let k be the number of atom types . We use a 3D point cloud G = ( A , R ) to represent the 3D geometry of a molecule with n atoms , where A ∈ { 0 , 1 } n×k is the atom type matrix and R ∈ Rn×3 is the atom coordinate matrix . Each row in the matrix A is a one-hot vector , and A [ j , u ] = 1 represents that the j-th atom has type u . The 3-dimensional row vector at the j-th row of the matrix R represents the 3D Cartesian coordinate of the j-th atom . We consider the generation of 3D molecular geometries as a sequential decision process . We start from a molecular geometry G1 with one carbon atom at the origin point , and generate the complete geometry by adding a new atom at each step . Specifically , at the i-th step , let the intermediate 3D molecular geometry generated from the previous i − 1 steps be Gi = ( Ai , Ri ) , which has i atoms . The atom type ai ∈ { 0 , 1 } k of the new atom is generated by the generative model ga based on the latent variable zai . Afterwards , the generative model g r decides the 3D Cartesian coordinate ri ∈ R3 of the new atom based on the latent variable zri . Both ga and gr are autoregressive functions of intermediately generated geometries . The overall sequential generation process can be described by the following equations : ai = g a ( zai ; Ai , Ri ) , ri = g r ( zri ; Ai , Ri ) , i ≥ 1 . ( 2 )

Summary: The paper introduces a new generative model for 3D molecular geometry, taking the form of a graph embedded in 3D where each node is an atom, and each edge is a bond. Differently from existing methods, the proposed pipeline attempts to generate the 3D molecule from scratch, instead of starting, e.g., from a 2D molecular graph as other methods do. The pipeline is sequential and is based on autoregressive flow models, which have been previously used for the generation of 2D molecular graphs, but not for the present task. The paper convincingly shows that the proposed model works better than competing approaches for the same task.