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 addresses the problem of generating 3D molecular geometries. The proposed G-SphereNet takes SphereNet as the backbone to generate 3D molecular geometries from scratch. The generation is done in a step-by-step fashion, with each step generating a (distance, angle, torsion) triplet, conditioning on features extracted by an attention mechanism.