$G^3$: Representation Learning and Generation for Geometric Graphs

A geometric graph is a graph equipped with geometric information ( i.e. , node coordinates ) . A notable example is molecular graphs , where the combinatorial bonding is supplement with atomic coordinates that determine the three-dimensional structure . This work proposes a generative model for geometric graphs , capitalizing on the complementary information of structure and geometry to learn the underlying distribution . The proposed model , Geometric Graph Generator ( G3 ) , orchestrates graph neural networks and point cloud models in a nontrivial manner under an autoencoding framework . Additionally , we augment this framework with a normalizing flow so that one can effectively sample from the otherwise intractable latent space . G3 can be used in computer-aided drug discovery , where seeking novel and optimal molecular structures is critical . As a representation learning approach , the interaction of the graph structure and the geometric point cloud also improve significantly the performance of downstream tasks , such as molecular property prediction . We conduct a comprehensive set of experiments to demonstrate that G3 learns more accurately the distribution of given molecules and helps identify novel molecules with better properties of interest . 1 INTRODUCTION . Geometric deep learning refers to the development of deep learning techniques for data from nonEuclidean domains , such as graphs and manifolds ( Bronstein et al. , 2017 ) . In recent years , graph neural networks ( GNN ) emerged to be a promising family of architectures that models the relational inductive bias ubiquitous to graph structured data ( Battaglia et al. , 2018 ; Zhou et al. , 2018 ) . With the rise of deep generative modeling , various generative models have been adapted for graphs and proven to be effective in learning the underlying distribution implicitly defined by a set of given graphs ( Goodfellow et al. , 2014 ; Kingma & Welling , 2013 ; Rezende & Mohamed , 2015 ; van den Oord et al. , 2016 ; De Cao & Kipf , 2018 ; You et al. , 2018b ; Jin et al. , 2018 ; Shi et al. , 2020 ) . In this work , we consider a special type of graphs—geometric graphs—which are equipped with geometric information in addition to the combinatorial structure ( Pach , 2013 ) . In practice , this additional information corresponds to low-dimensional geometric structures and typically appears as node coordinates in R2 or R3 . The coordinates find broad uses in molecules ( Van Aalten et al. , 1996 ) , meshes ( Alliez et al. , 2005 ) , and graph drawing ( Frishman & Tal , 2008 ) . This work aims to develop a generative model that captures both the combinatorial and the geometric characteristics exhibited in a given collection of geometric graphs . One driving application of geometric graphs is molecule generation . In pharmacology , drug discovery is the process of identifying new candidate medications in the form of molecules . The process is known to be difficult , costly , and time-consuming ( Paul et al. , 2010 ) , because of the discrete nature of the search space and its vast size ( estimated to contain at least 1033 molecules ) ( Polishchuk et al. , 2013 ) . Recently , deep learning techniques have been developed to represent molecules as continuous vectors and apply continuous optimization ( Jin et al. , 2018 ) or reinforcement learning ( You et al. , 2018a ) to conduct an efficient search . In almost all existing work based on the molecular-graph approach , the graph is treated as a combinatorial object and the three-dimensional geometric information serves as node features only . However , the atomic coordinates encode vital energy information of the molecule and they can be instrumental for the inference of structure and state . Therefore , a proper modeling of the geometric information is paramount for a more accurate representation of the molecule in downstream tasks . In this work , we develop an autoencoder for geometric graphs , orchestrating a GNN and a point cloud model in a nontrivial manner . Besides the straightforward use of a GNN ( Duvenaud et al. , 2015 ; Kearnes et al. , 2016 ; Gilmer et al. , 2017 ) to encode the combinatorial structure , we treat the nodes with spatial coordinates as low-dimensional point clouds and use a point cloud model to encode the geometry . Different from decoders merely based on graph structures in conventional graph autoencoders , on the other hand , we decode the graph structure by first reconstructing the geometry . Specifically , we use a folding-based technique ( Yang et al. , 2018 ; Pang et al. , 2021 ) to map a template of points to the correct geometry , from which the combinatorial structure is inferred . This way , both sources of information are organically fused and processed . Furthermore , to fulfill the generative capability , we augment the autoencoder with a normalizing flow ( Kobyzev et al. , 2019 ; Dinh et al. , 2016 ; 2014 ) , so that one is able to effectively sample from the otherwise intractable latent space . The proposed geometric graph generator , G3 , addresses geometry that is rarely considered in past graph generative models . A naive alternative to incorporate node coordinates is to treat them as part of features when consumed by a GNN . This approach , however , biases the geometry by the graph structure owing to local neighborhood aggregation—the global geometry may be compromised . For geometric graphs , the structure itself is combinatorial and discrete , while the geometry is continuous . These two sources of information are often complementary and benefit from collaborative processing . Our strategy is to exploit both graph techniques and point cloud techniques to maximally retain information from the two sources . To evaluate the effectiveness of G3 , we focus on molecule applications and use popular datasets ( QM9 and ChEMBL ) that contain atomic coordinates . We demonstrate that G3 outperforms a number of representative graph generative models in generating novel and valid molecules and it also significantly outperforms either a GNN or a point model alone for property prediction . Appealingly , G3 decodes molecules substantially faster than do many popular models . Furthermore , with the use of Bayesian optimization in concert with a trained model , we can identify better molecules under metrics of interest than do existing methods . 2 RELATED WORK . Generative models refer to a class of machine learning models that can learn the underlying distribution , implicitly defined by a given set of data , and to sample from it . Noteworthy generative models in the deep learning era include variational autoencoders ( VAE ) ( Kingma & Welling , 2013 ) , generative adversarial networks ( GAN ) ( Goodfellow et al. , 2014 ) and normalizing flows ( NF ) ( Kobyzev et al. , 2019 ; Dinh et al. , 2016 ; 2014 ) . Among them , NF bears a unique advantage in estimating the density ( likelihood ) . For training , NF directly optimizes the likelihood , whereas VAE maximizes a lower bound of it ( called the ELBO Kingma & Welling ( 2013 ) ) and GAN minimizes the discrepancy between the input and the transformed noise distribution . Under these generative models , a graph can be generated in a sequential or one-shot manner : the former samples nodes and edges incrementally while the latter samples a graph in its entirety . Combining GAN and reinforcement learning , MolGAN ( De Cao & Kipf , 2018 ) and GCPN ( You et al. , 2018a ) generate molecules in a one-shot and sequential fashion , respectively . MolecularRNN ( Popova et al. , 2019 ) is an autoregressive model based on recurrent neural networks ; it performs validity checks during the sequential generation process and rejects invalid structures . JT-VAE ( Jin et al. , 2018 ) is a VAE-based tree model that constructs the molecular graph sequentially from sub-components . Popova et al . ( 2019 ) point out that JT-VAE may suffer ambiguity in the conversion of the tree structure to a molecular graph , affecting property optimization . GraphNVP ( Madhawa et al. , 2019 ) and GraphAF ( Shi et al. , 2020 ) are recently proposed NF-based models that parameterize the mapping from a latent vector to a graph by using a flow . The former employs one-shot sampling while the latter is sequential . Molecules admit representations other than graphs . An alternative is the simplified molecular-input line-entry system ( SMILES ) ( Weininger , 1988 ) , a string notation that universally describes molecular structures . Grammar VAE ( GVAE ) ( Kusner et al. , 2017 ) and syntax-directed VAE ( SD-VAE ) ( Dai et al. , 2018 ) are VAE-based models that process SMILES strings and apply grammatical rules to sequentially reconstruct the strings . Another family of models , which are motivated by quantum chemistry principles , models molecules as point clouds and learns a physically robust representation to predict , e.g. , molecular energies and equilibrium conformations . Schütt et al . ( 2018 ) propose SchNet , a network with continuous-filter convolution layers that model interactions between atoms and learn a representation that is invariant to translation and rotation . Using the same filters , Gebauer et al . ( 2019 ) devise Generative-SchNet ( G-SchNet ) , which accurately models the target molecule distribution and generates novel 3D conformations that are relatively stable . In addition to the graph structure , our work treats nodes as a point cloud , which is a collection of points in R3 . Point clouds are a prominent subject in computer vision and computer graphics ; they are obtained through , e.g. , LIDAR scanning of object surfaces ( Yang et al. , 2018 ) . Deep learning with point clouds is faced with challenges in defining the convolution operator ( Bruna et al. , 2013 ; Bronstein et al. , 2017 ; Schonsheck et al. , 2018 ; Jin et al. , 2019 ) . Volumetric CNN ( Wu et al. , 2015 ; Qi et al. , 2016 ; Maturana & Scherer , 2015 ) applies convolution filters on voxels obtained from discretization of the three-dimensional space . Multiview CNN ( Su et al. , 2015 ; Qi et al. , 2016 ) reduces point clouds to a collection of 2D images and applies gridded convolutions on them . PointNet ( Qi et al. , 2017a ; b ) uses point-wise MLPs to featurize individual points , followed by a symmetric pooling that generates a permutation-invariant global description . 3 GEOMETRIC GRAPH GENERATOR . In this section , we present G3 . The model is an autoencoder that contains interacting modules to process and reconstruct the geometry ( i.e. , node coordinates ) , the structure ( i.e. , edges and edge types ) , and additional features ( e.g. , node features ) . Such a sophisticated autoencoder allows forming a more accurate representation of the geometric graph and learning better the input distribution . A proper processing of the geometric and combinatorial information is crucial to the success of G3 and downstream tasks . For example , an authentic representation of the node coordinates serves as the basis for decoding both the node and the edge types : the training hinges on an accurate registration between the reconstructed point cloud and the ground truth . The straightforward adaptation of GNNs for encoding the geometric information ( as node features ) results in poor coordinate reconstruction and worse generation quality overall . Moreover , one can naively reconstruct coordinates and node features with a single decoder . However , mixing two modes of information together introduces unnecessary bias through , e.g. , the difference in scales . Hence , separate but dependent decoders for features and geometry render a more accurate reconstruction .

Summary: The paper proposes a generative models for learning over the space of geometric graphs -- those graphs whose nodes are associated with geometric coordinates (point clouds). The point cloud serves as basis for decoding graph structure. This makes the entire system efficient. An extensive suite of experiments on chemical benchmarks (QM9 and ChEBML) show that the proposed method is competitive against state-of-the-art rivals on a variety of measures and tasks.