Year: 2021

Abstract:
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 (G), 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. G 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 G learns more accurately the distribution of given molecules and helps identify novel molecules with better properties of interest.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 GEOMETRIC GRAPH GENERATOR
text: 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.

heading: 3.1 NOTATION
text: Let n be the number of nodes, df be the number of node features, de be the number of edge types, and dc be the dimension of the node coordinates (typically two or three). We denote a geometric graph as G = (A,X), where A ∈ {0, 1}n×n×de is the adjacency tensor and X = [C,F ] ∈ Rn×(dc+df ) is the node matrix, where each note vector xi = [ci, fi] is a concatenation of coordinates ci ∈ R1×dc and features fi ∈ R1×df . For molecules, the node features are one-hot encoding of the atom type, and the edge types correspond to bond types.

heading: 3.2 ENCODER
text: The purpose of the encoder is to use a latent vector to represent both the geometry and the structure of G. It contains two modules: a point encoder that maps the point cloud C to a geometry descriptor zp; and a graph encoder that takes A and F as input and maps them to a structure descriptor zg . Both descriptors are concatenated to form the latent representation z = [zp, zg] ∈ Rdz . A usual graph encoder processes the structure information A and node features F by using a GNN. When the coordinate information C is available, one could have treated C as part of the node features and feed them into the GNN, indiscriminately of F . However, the geometry information of a graph is often complementary to the structure. Hence, a drawback of such an approach is that one heavily relies on the local neighborhoods from graph structures to parameterize the geometry, resulting in a strong bias toward the structure. A notable feature of our design is that we separate the geometry
from the structure when encoding a geometric graph; we will show an ablation study supporting this design subsequently.
This idea offers a means to process all node information, including features F and coordinates C, independently of the combinatorial structure, as an additional parameterization over the pure use of a GNN. In what follows, we describe the two encoder modules and defer details of the specific network architecture in the appendix.
Point Encoder. In this module, we treat each row of C as a point and use a point cloud architecture to obtain the geometry descriptor zp. First, a multilayer perceptron (MLP) is applied point-wise to the point cloud to amplify the feature dimensions. The point-wise features are then passed through a series of multi-headed attention blocks (Vaswani et al., 2017) to further featurize the points. These attention-based blocks are more effective than simple point-wise MLPs used in PointNet (Qi et al., 2017a) as they are able to leverage global information, i.e. features on all other points, to produce rich and expressive point-wise features. At last, we use a global max-pooling over the points followed by an additional MLP to yield the geometry descriptor zp.
Graph Encoder. This module takes the graph adjacency tensor and the node feature matrix as inputs and outputs a structure descriptor zg. Virtually all GNNs serve the purpose; some are even specialized to molecules (e.g., Duvenaud et al. (2015); Kearnes et al. (2016); Gilmer et al. (2017)). Our empirical investigation finds that GINE, Graph Isomorphism Network with Edge features (Xu et al., 2018), obtains a good and stable performance for our purpose.

heading: 3.3 DECODER
text: Decoding a graph is substantially harder than encoding one. A common idea is to employ a sequence model and decode nodes and edges one by one Popova et al. (2019). Training such a model is increasingly more challenging for long sequences (i.e., large graphs). On the other hand, approaches decoding the adjacency tensor also face scalability challenges owing to theO(n2) degrees of freedom Samanta et al. (2020). In this work, we exploit the unique opportunity offered by a geometric graph to decode first the geometry, then the point-wise features, and lastly the structure.
Specifically, we propose a three-step decoding process that decodes a graph G′ = (A′, X ′) from z, where X ′ = [C ′, F ′]. The first step is to decode C ′, the coordinates for the point cloud, which involves using the latent representation z to morph a template of points to the correct geometry. By isolating the decoding of coordinates, we can leverage specialized architectures to achieve accurate reconstructions of the geometry, which is crucial for the subsequent decoding of F ′ and A′. Secondly, the reconstructed geometry C ′ is recycled along with the latent code z to predict the point-wise features F ′. Finally, X ′ = [C ′, F ′] is used to infer the graph structure A′ through a link predictor. This procedure sidesteps the scalability challenge of both the sequential generation approach and the adjacency decoding approach.
In what follows, we describe each component of decoder modules and defer details of the specific network architecture in the appendix.
Coordinate Decoder The idea of decoding a point cloud from a latent vector z is to successively fold a template of points until it forms the desired geometry. Our coordinate decoder is inspired by TearingNet (Pang et al., 2021), which is a refinement over FoldingNet (Yang et al., 2018; Tao, 2020), a popular archetypal architecture for point cloud reconstruction. Note that the template design is application-dependent; an effective template is guided by the geometry to be modeled. For example, a spherical template with roughly equidistant points is more suitable for the three-dimensional structure of molecules. This design is different from the original TearingNet architecture that uses a regular 2D grid to model surfaces.
Feature Decoder To decode the point-wise features F ′, we apply the folding operations with the decoded coordinates C ′ from the aforementioned coordinate decoder as the starting template. The latent code z is replicated and concatenated to C ′, which is then forwarded through point-wise MLPs to complete one set of folding. We repeat two sets of folding to arrive at the decoded features. The architecture of the feature decoder is illustrated in Figure 1.
Link Predictor (Graph Decoder) The link predictor determines the existence of an edge (and its type if so) for every pair of decoded nodes. A typical link predictor takes the inner product between two nodes to indicate the strength of connection (Trouillon et al., 2016). However, such a link predictor is oblivious to the geometry because a large inner product in 3D does not necessarily indicate closeness. Rather, the displacement of two points, their respective neighborhoods, and their neighbors’ features may all play a role in determining the existence and type of an edge. For example, in molecules, the bond between a pair of atoms depends on their distance as well as the neighboring atoms. Therefore, we augment the features of a decoded node xi from [ci, fi] to x̂i as
x̂i = [|ci − cn1 |, ..., |ci − cnk−1 |, fi, fn1 , ..., fnk−1 ],
where n1, . . . , nk−1 are the node indices of the sorted k-nearest neighbors of xi; fi ∈ Rdf and ci ∈ Rdc are decoded features and coordinates, respectively; and the absolute sign indicates element-wise absolute values. We then use these augmented features as well as the coordinate-wise displacement |ci − cj | to define the link predictor. Formally,
l(i, j) = MLP([x̂i, x̂j , |ci − cj |]) ∈ Rde , LinkPredictor(i, j) = l(i, j) + l(j, i) ∈ Rde .
We highlight that our link predictor works remarkably well as a standalone module for molecules, which can be characterized as a procedure for inferring inter-atomic bonds based on molecular conformation and atom types. To the best of our knowledge, there are no reliable public software available for this task, especially for large molecules with noise in coordinates or atom types. Therefore, our proposed link predictor has its own value other than being a part of G3. We refer appendix G for illustrating the outstanding performance of the standalone link predictor.

heading: 3.4 LATENT SPACE SAMPLING: NORMALIZING FLOW
text: To reliably sample from G3’s latent space, we make use of a normalizing flow model (Kobyzev et al., 2019), which consists of a sequence of invertible transformations that map between an easily sampled base distribution (e.g., standard normal) and an unknown distribution (e.g., the latent distribution of our autoencoder). The densities of the two distributions are related by the Jacobian of the transformations, so that log-likelihood of the unknown distribution can be computed and optimized for training. In practice, we find that RealNVP (Dinh et al., 2016) to be a simple yet effective model for our needs. After the training is finished, samples from the base distribution are forwarded through the normalizing flow and subsequently decoded into the generated molecules.

heading: 3.5 OVERALL MODEL AND TRAINING LOSS
text: The overall architecture of G3 is outlined in Figure 2, which summarizes all the components described in the preceding subsections. To train this model, we use a loss consisting of multiple terms with a tunable weighting scheme: L = λ1 · LC + λ2 · LF + λ3 · LE + λ4 · LNF + λ5 · LR. Node coordinate loss LC . In this loss, we use only the coordinate part C contained in the node matrix X . In our experiments, we use the Sliced Wasserstein-2 distance (SW2) as our loss for robust point cloud reconstructions (Rabin et al., 2011). Note that in the calculation of LC , every point c ∈ C is identified with the nearest match c′ ∈ C ′ and vice versa. This correspondence allows us to feasibly define the following losses LF and LE with differentiability. We remark that another commonly used loss for point-cloud matching, Chamfer distance (Yang et al., 2018), can also be used here.
Node feature loss LF . Since node correspondence is established, the difference between F and F ′ can be meaningfully quantified. When the rows are one-hot vectors, a cross-entropy loss is suitable; otherwise, mean absolute error (MAE) or other losses may be appropriate.
Edge loss LE . The edge loss quantifies whether edge types are correctly inferred. Hence, a crossentropy loss between A and A′ is natural, where node correspondence is needed.
Normalizing flow loss LNF . The normalizing flow model is equipped with a standard negative log-likelihood loss (Dinh et al., 2016).
Additional loss LR. This term enforces structural and other chemical constraints innate to the specific application. Some constraints may be easily violated even if the decoding of nodes and edges is reasonably accurate (as long as error exists, however small). Hence, a regularization that penalizes the violation of the constraints helps improve semantic validity. For example, for molecules, one incorrect decoding of a carbon atom to oxygen is enough to create an invalid molecule, because the valence of oxygen is much smaller. In this case, we follow the work by Ma et al. (2018) and include a soft valence regularization term LR to encourage higher validity of the generated molecules.

heading: 4 EXPERIMENTS

heading: 4.1 DATASETS AND BASELINE MODELS
text: We perform experiments on two benchmark datasets that come with atomic coordinate information: QM9 (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014) and ChEMBL (Mendez et al., 2019; Davies et al., 2015). The QM9 dataset contains 134k molecules with 4 atom types and up to 9 heavy atoms per molecule. The ChEMBL dataset contains 1.5 million molecules, and the largest one has hundreds of atoms. For feasible experimentation, we take a subset of it that contains 250k molecules with 19 atom types and up to 39 heavy atoms per molecule. We name this new dataset ChEMBL250k. Details of dataset proprocessing are provided in the appendix.
We compare G3 with several representative architectures that feature a variety of generative approaches. SDVAE (Dai et al., 2018) is a syntax-directed VAE model that processes SMILES representation of molecules and generates syntactically correct ones. MolGAN (De Cao & Kipf, 2018) combines GAN and reinforcement learning to generate molecules with optimized properties. JT-VAE (Jin et al., 2018) is a VAE-based model that assembles chemical substructures, boasting high validity of the generated molecules. GraphNVP (Madhawa et al., 2019) is a flow-based model with a one-shot generation procedure. GraphAF (Shi et al., 2020) is another flow-based autoregressive model that advertises high flexibility in density modeling. G-SchNet (Gebauer et al., 2019) is a quantum-chemistry-inspired model that learns accurate conformations and sequentially generates molecules close to the trained distribution.

heading: 4.2 GENERATION QUALITY
text: Metrics and processing. To evaluate the generative capability of different models, we measure the validity, uniqueness, and novelty (V, U, N) on their generated molecules. Validity is the percentage of
molecules that pass RDKit’s sanitization check. Uniqueness is the portion of valid molecules that are distinct. Novelty is the portion of valid molecules that do not appear in the training set.
In addition, we include metrics computed between the test and the generated molecules via MOSES (Polykovskiy et al., 2020): Frechet ChemNet Distance (FCD), which measures the distance between the two sets according to ChemNet’s penultimate layer’s activation; Fragment similarity (FRAG) measures the similarity in fragments present in both sets of molecules; Internal Diversity (IntDiv2) assesses the diversity within the generated molecules; LogP, QED, and Weight are the 1D Wasserstein distances between the two sets on different chemical properties: LogP measures the lipophilicity of a compound (Bhal, 2007); QED is the quantitative estimate of druglike-ness (Bickerton et al., 2012); Weight is the sum of atomic weights in a molecule. All scores are computed on 10k generated molecules.
In our experiments, G3 sometimes generates disconnected molecules, for which we treat the largest connected component as the final generated molecule (breaking ties arbitrarily) for evaluation. Additionally, we observe the normalizing flow cannot capture the autoencoder’s latent distribution perfectly when trained on ChEMBL250k. Therefore, we interpolate between samples from the normalizing flow and the nearest encoded training data before decoding.
Results and discussions. Results reported in Tables 1 and 2 invite a few observations. Overall, we see that G3 has strong performances across all metrics in both QM9 and ChEMBL250k. Notably, G3’s low FCD score indicates accurate modeling of the training distribution, while a high novelty rate means G3 has not simply memorized the training set. This is important because a model can always reproduce the training set to get perfect scores on e.g. FCD. In contrast, although JT-VAE has better FCD and LogP scores on QM9, its low novelty is characteristic of overfitting. In addition, G3’s validity scores are consistently the highest among the models with no hard valency checks (the others have 100% validity), and G3’s high uniqueness and molecular diversity suggest it is able to avoid mode collapse, which is evident in MolGAN. Furthermore, G3’s robust performance in LogP, QED, and Weight implies that G3’s generated molecules are not only faithful to the training set in terms of (sub)structure, but also with respect to various chemical properties.
Efficiency. G3’s performance on large molecular datasets is especially impressive considering its one-shot generation approach and is far superior to other models in the same camp (GraphNVP, MolGAN). Comparing to the best sequential models, e.g. JT-VAE, G3 can generate molecules as good in quality in a much faster fashion, thanks to its one-shot nature. When benchmarked on a machine with i7-8700K CPU, GeForce 1080 RTX GPU, and 16GB RAM, GraphAF and JT-VAE took 4804.36s and 2642.97s to generate 10k molecules when trained on ChEMBL250k, respectively. In contrast, G3 took 47.79s, which is approximately 55 times faster than JT-VAE and 100 times faster
than GraphAF. As discussed, the inefficiency of sequential generation models is mainly due to the challenge of parallelization.
Latent space interpolation. Lastly, we present examples of G3’s generated molecules as well as their local latent space in Figure 3. The surrounding molecules are decoded along two random orthogonal axes starting from the center molecule highlighted by the blue box. The transition between molecules are gradual and chemically sound, which illustrates G3’s great capability to smoothly parameterize discrete molecules in its latent space. More qualitative examples can be found in the appendix.

heading: 4.3 ABLATION STUDY AND PROPERTY PREDICTION
text: To process the complementary geometry and structure information, we use a pair of point encoder and graph encoder for G3. A natural question asks if two encoders are really necessary. Here, we conduct an ablation study and compare G3 with two simpler models. The first one uses only the point encoder and removes the graph encoder. Because of a lack of graph structure encoding, the edge decoding quality entirely depends on node information. The second simplified model is to remove the point encoder but retain the graph encoder. Since the graph encoder takes coordinates as part of the input, the geometry information is still encoded, albeit through a GNN module that was not designed to handle geometry in the first place.
We train the models on QM9 until convergence using the same hyperparameters and latent dimensions. From Table 3, G3’s encoder achieves the best results in two-thirds of the metrics, and the remaining third is split between the other approaches. We further highlight the combined encoder’s convincing lead in FCD, which is the best single indicator of the quality of generated molecules in the absence of overfitting. As a result, we find G3’s combined approach to be clearly favored.
To further indicate that the combination of geometry and graph information used in G3 learn a meaningful latent representation for the molecules, we show downstream tasks for predicting chemical properties. In the subsequent experiment, we show that a G3 encoder pretrained in a unsupervised manner can be used to predict chemical properties of QM9 molecules with an additional MLP to fit
desired properties. Alternatively, one can also train G3’s encoder in a supervised manner to predict the desired properties, which yield better results. The pretrained models are first trained on the QM9 dataset in an unsupervised manner for 80 epochs. The weights on the encoders are then frozen, and an additional MLP is trained on the encoded vectors until convergence. The supervised models are trained on the QM9 dataset until convergence.
We predict four properties that come with the QM9 dataset: Homo-Lumo energy (H-L), internal energy, free energy, and heat capacity. In all cases, G3’s graph-point encoder significantly outperforms the standalone graph encoder used in section 3.2. This further substantiates G3’s ability to leverage the geometric and combinatorial features in a geometric graph more effectively by separating the graph and the point encoder.

heading: 4.4 PROPERTY OPTIMIZATION
text: One use of generative models with latent spaces is to search for better molecules under metrics of interest. Following the approach outlined by Kusner et al. (2017), we apply Bayesian optimization (BO) on the G3 model pretrained from Table 1 and 2 to optimize two metrics: penalized logP and QED. Penalized logP is logP less the synthetic accessibility score and the number of long cycles (Bhal, 2007), and QED measures how likely a compound is to be a drug (Bickerton et al., 2012).
The results for QM9 and ChEMBL250k are reported in Table 5. One sees that G3 consistently finds the best molecule under both metrics. We remark that BO works better for low dimensional search spaces (Frazier, 2018). The latent spaces of SDVAE and JT-VAE are of a similar nature to that of G3, but GraphAF must use a latent space of the same dimension as the input space. Hence, BO may not be the best method for GraphAF. Indeed, Shi et al. (2020) employs reinforcement learning to chase for optimal metric values, which gives more promising results on ZINC250k. However, for fairness of comparison, we opt to use BO for all methods. More details are provided in the appendix.

heading: 5 CONCLUSION
text: We have presented a generative model G3 for geometric graphs that are equipped with additional geometry information. Geometric graphs appear in many practical applications, a notable one being molecules, whose three-dimensional structure is rarely captured by existing graph generative models. Our model is an autoencoder that employs both GNN architectures and point cloud architectures to handle the complementary structure and geometry information. Moreover, a normalizing flow is inserted into the latent space so that one can effectively sample new graphs. Experiment results show that G3 is able to more accurately learn the distribution of given graphs and better predict molecular properties. It helps discover novel molecules with better properties than do many representative graph generative models.

heading: A EXPERIMENTAL DETAILS
text: Datasets We preprocess each molecule by centering it based on the 3D coordinates and aligning its orientation with the principal axes via PCA. The molecules in the datasets have different number of atoms. Hence, we pad each graph with “phantom nodes” to reach the maximum molecule size in the dataset. The phantom nodes are indicated by one of the node features, and we also treat nonexistence of edges as an edge type.
Importantly, the phantom nodes need artificial coordinates. We cannot naively allocate all the phantom nodes to the origin, because overlapping coordinates cause non-unique correspondences in the point-cloud registration, making the learning of atom and edge types impossible. In practice, we use an extra coordinate dimension and assign the phantom coordinates for the i-th phantom node as (0, 0, i) or (0, 0, 0, i) if the original point cloud is in 2D or 3D, respectively.
In addition, we make a distinction between each atom type and its aromatic counterpart whenever possible. For example, a carbon atom may be part of an aromatic ring, in which case we mark its atom type as aromatic carbon to distinguish it from an ordinary carbon atom. This facilitates the reconstruction of ring structures in the link predictor, as it can predict aromatic bonds easily by checking if the atom types are aromatic.
Note that another popularly used benchmark, ZINC250k (Sterling & Irwin, 2015), does not contain coordinate information and hence we omit it. However, the scales of ChEMBL250k and ZINC250k are rather close, and their results unsurprisingly share similar trends when cross-checked against the performance of baseline models.
Models For G3, we use the same architecture and identical number of layers and hidden dimensions on both QM9 and ChEMBL250k, without separate fine tuning. Such a practice helps corroborate robustness of the model under different data size and complexity. The latent dimension used for QM9 is [dp, dg] = [10, 30], and [35, 25] for ChEMBL250k. The additional graph latent dimensions is required to accurately decode the diverse atom types present in ChEMBL250k, which is more than twice than that in QM9.
We did not check G-SchNet’s performance on ChEMBL250k for two reasons. First, G-SchNet excels at reconstructing and generating conformers, which are the 3D structures of molecules that are crucial for various computations in quantum chemistry. The coordinates of ChEMBL250k, however, come from molecular graphs and are all in 2D. Second, G-SchNet’s sequential generation process utilizes all previously placed atoms to predict the position of the new atom at each step, which scales poorly with respect to the size of molecules. As many ChEMBL250k molecules have more than 100 atoms after adding Hydrogens, the computational resource required to train a G-SchNet model on ChEMBL250k is likely to be prohibitively high.
All baseline results are obtained by running the implementation provided by the authors, with default hyperparameters. The baseline results generally match the reported ones from the original papers. One discrepancy occurs on GraphAF with QM9, which is caused by the disabling of a filter that ignores small molecules. Turning off the filter amounts to a more fair comparison because small molecules are still sensible results. The V, U, N metrics for MolGAN on QM9 are also somewhat off, which is likely a symptom of GAN’s training instability.
Code G3 is implemented with the Pytorch framework. All relevant code as well as the trained models will be made available upon publication.

heading: B ARCHITECTURE DETAILS
text: In this section, we specify the architecture used to produce the reported experiment results. Let MLP(n, h, o, l) denote a (point-wise) multi-layer perceptron that has input dimension n, hidden dimension h, output dimension o, and number of hidden layers l. Let MHA(n, h) be the multiheaded attention block with embedding dimension n and h heads. The implementation of multiheaded attention is borrowed from Paszke et al. (2019).
Point Encoder As illustrated in Figure 4, the G3 point encoder is :
MLP(n, 256, 768, 1)→MHA(256, 3)→MHA(256, 3)→ GlobalMaxPooling →MLP(768, 256, dp, 1)
Graph Encoder Let GINE(h, l) denote a GINE kernel (Xu et al., 2018) with embedding dimension h and l hidden layers in the MLP. Let Set2Set be a Set2Set global pooling layer (Vinyals et al., 2015). We first use linear layers to embed the node and edge features to 256 dimensional vectors. The graph encoder is
GINE(256, 3)→ GINE(256, 3)→ GINE(256, 3)→ Set2Set→MLP(512, 256, dg, 2)
A Batchnorm (Ioffe & Szegedy, 2015) layer is applied after every GINE convolution.
Coordinate Decoder This module is inspired by (Pang et al., 2021). Let U be the template with the i-th row denoded as ui, z = [zp, zg] be the latent vector, and Y ′ be the output with the i-th row denoded as y′i. Define the FoldingNet F, where F(U, z) = Y
′. It performs the following operation for each point i:
y1i ←MLP(ui, z) y2i ←MLP(y1i , z) y′i ←MLP(y2i , z)
Define the TearingNet T, where T(U,X, z) = Y ′. It performs the following operation for each point i:
y1i ←MLP(ui, xi, z) y′i ←MLP(ui, y1i , xi, z)
The coordinate decoder does the following operations:
C1 ← F(U, z) U1 ← T(U,X1, z) + U C ′ ← F(U1, z)
Where C ′ with the i-th row as c′i are the decoded coordinates. The MLPs all have two hidden layers with hidden dimensions 1024.
Feature Decoder Let C ′ be the decoded coordinates with the i-th row denoded as c′i, and z = [zp, zg] be the latent vector. Using the same FoldingNet notation as above, the feature decoder does the following:
F ′ ← F(C ′, z)
Where F ′ with the i-th row as f ′i are the decoded features. The MLPs all have two hidden layers with hidden dimensions 1024.
Link Predictor Let dx = 2((da + dc) · k) − dc. The MLP inside the link predictor is MLP(dx,1024,de,3).

heading: C HYPERPARAMETERS
text: Experiment on generation quality G3: For QM9, the latent vector z = [zp, zg] has dimension 10 + 30 and the loss weighting is (λ1, λ2, λ3, λ4, λ5) = (10, 1, 1, 10−3, 10). For ChEMBL250k, the latent dimension is 35 + 25 and the loss weighting is (10, 0.1, 9, 10−4, 5).
GraphAF: In this model, there is an option to filter out generated molecules with fewer than m atoms. Since no other models use filtering, we set m = 0 for a fair comparison. Because larger molecules are more likely to be novel and unique, the scores reported by the original authors are higher than those we obtain. Moreover, we have tried best efforts to tune hyperparameters for QM9 as the authors did not release their setting.
Other models: We use the default hyperparameters released by the authors for QM9. For ChEMBL250k, we borrow those used for ZINC250k. We find this to be a reasonable approach because ZINC250k and ChEMBL250k have the same number of molecules, and the maximum number of atoms per molecule is also quite close: 38 and 39, respectively. As such, it is expected that the resulting performance on these two datasets is similar, which is indeed the case.
Experiment on property optimization In each iteration of the Bayesian optimization, a Gaussian process is used to fit the explored data, and a batch of new candidates are suggested based on the use of an acquisition function. For all models, we use training molecules as the initial set of points and perform five iterations of Bayesian optimization, wherein the acquisition function is EI (expected improvement). In each iteration, 60 candidates are generated.
The implementation is provided by Kusner et al. (2017) via a custom build of Theano. To conduct Bayesian optimization on the baseline models, we use the default hyperparameters released by the authors whenever available, and follow the practice above otherwise.

heading: D LOSS DETAILS
text: For completeness, in what follows we give the detailed expressions for all the training losses. They are generally straightforward or have been defined in the literature.
D.1 COORDINATE LOSS
The Wasserstein distance is intimately connected with the study of Optimal Transport. Let (Ω,F) be a measurable space, and P(Ω) be the set of all probability measures on (Ω,F). Define two probability measures µ, ν ∈ P2(Ω), where P2(Ω) = {µ ∈ P(Ω) : ∫ Ω ‖x‖2µ(dx) < ∞}. The Wasserstein-2 distance W2(µ, ν) is defined as:
W 22 (µ, ν) = inf γ∈C(µ,ν) ∫ Ω×Ω ‖x− y‖2γ(dx, dy)
Where C(µ, ν) is the set of all couplings between µ, ν that satisfies: γ(A × Ω) = µ(A),∀A ∈ F , γ(Ω×B) = µ(B),∀B ∈ F , and F is the sigma algebra associated with Ω. It is known thatW2(µ, ν) defines a metric on P(Ω) (Ambrosio et al., 2008).
In general, W2(µ, ν) is difficult to compute, especially in high dimensions. However, it has a simple form in one-dimension:
W2(µ, ν) = ∫ 1 0 |F−1µ (τ)− F−1ν (τ)|2dτ
Where Fµ, Fν denote the cumulative distribution functions associated with µ, ν, respectively. The simple one-dimensional form of W2(µ, ν) inspires the usage of sliced Wasserstein distance, which is the expected value of the projected one-dimensional Wasserstein distance. Define the push-forward operation ∗ as: f∗µ(A) = µ(f−1(A)),∀A ∈ F , and the projection map as Θ(x) = 〈θ, x〉 for a normalized direction θ. Formally, the sliced Wasserstein-2 distance SW2(µ, ν) in d-dimension is:
SW2(µ, ν) = ∫ Sd−1 W2(Θ∗(µ),Θ∗(ν))dθ
where dθ is the uniform measure on Sd−1. Bonnotte (2013) has shown that SW2(µ, ν) also defines a metric on P(Ω), and it induces the same topology as W2(µ, ν) when Ω is compact. In practice, we compute SW2(µ, ν) by approximating the integral via Monte Carlo sampling directions on the sphere. We then compute the projection of the point clouds onto the sampled directions, sort the one-dimensional projections, and average their differences. Through experimentation, we found that 200 is an appropriate number of directions to sample.
D.2 NODE FEATURE LOSS LF AND EDGE LOSS LE
Let C ∈ Rn×dc be the original point cloud, where each row ci := Ci,: gives the coordinates of a point. Similarly, let C ′ ∈ Rn×dc be the reconstructed point cloud with c′i := C ′i,:. Define the permutation matrix σ as:
σij = { 1, if j = arg mink ‖ci − c′k‖ 0, otherwise.
Let F, F be the original and the reconstructed node feature matrices, respectively. The node feature loss is defined as:
LF = 1
n n∑ i=1 H(Fi,:, F ′ i,:, w N ),
where F ′ = σF is the reconstructed node feature matrix that is aligned with F ; wN is the weighting vector for each coordinate; and H(x, y, w) is the weighted cross entropy between probability vectors x and y:
H(x, y, w) = − ∑ i wixi log(yi).
Let A,A be the original and the reconstructed adjacency tensors, respectively. The edge loss is defined as:
LE = 1
n(n− 1) n∑ i=1 n∑ j=1,j 6=i H(Ai,j,:, A ′ i,j,:, w E),
where A′ is the reconstructed adjacency tensor that is aligned with A, satisfying A′:,:,k = σA:,:,kσ T for k = 1, 2, ..., de; and wE is the weighting vector for the edge classes. Note the exclusion of j = i in the inner summation; self-loops are not counted.
We use the following weighting for the cross entropies:
wN = 1, wE = ln(1 D),
where D is a probability vector that gives the empirical distribution of classes computed from data; and the operations and ln are element-wise.
D.3 NORMALIZNG FLOW LOSS LNF
Let z ∈ RD be a random vector with density pz(z) and let x be a vector resulting from a sequence of K invertible transformations from z:
x = fK ◦ fK−1 ◦ ... ◦ f2 ◦ f1(z). Define z0 = z, zi = fi(zi−1) for i = 1, ...,K, and zK = x. By the change-of-variable formula, the density of x can be expressed as
log px(x) = log pz(z)− K∑ i=1 log ∣∣∣∣det( ∂fi∂zi−1 )∣∣∣∣ . (1)
If z is standard normal, then
log pz(z) = − 1 2 ‖z‖2 − D 2 log 2π. (2)
In our work, we use RealNVP Dinh et al. (2016) to parameterize the mappings. Specifically, a mapping f(u) = v is defined as
v1:d = u1:d
vd+1:D = ud+1:D exp(MLPs(u1:d)) + MLPt(u1:d), where d = 12D. One easily calculates that the Jacobian of f is
∂f ∂u =
( I 0
∂vd+1:D ∂u1:d diag(exp(MLPs(u1:d)))
) .
Therefore,
det
( ∂f
∂u
) = exp D−d∑ j=1 MLPs(u1:d)j  . (3) Substituting equation 2 and equation 3 into equation 1, we obtain the negative log-likelilhood loss
− log px(x) = K∑ i=1 D−d∑ j=1 MLPis((zi−1)1:d)j + 1 2 ‖z0‖2 + D 2 log 2π. (4)
D.4 ADDITIONAL LOSS LR
We follow the approach outlined in Ma et al. (2018) to introduce a soft valence regularization to penalize the excess of valence capacity for each atom, thereby encouraging the decoder to output chemically valid molecules. In particular, define LR = ∑n i=1 σ(V (i) − U(i)) for each molecule, where σ is the ReLU function and U(i) and V (i) are the valence and capacity of atom i, respectively. The soft (thus differentiable) version is:
U(i) = ∑ r u(r)F ′i,r,
V (i) = ∑ j 6=i ∑ k h(k)A′i,j,k,
where F ′ and A′ are the reconstructed and aligned node feature matrix and adjacency tensor, respectively. Here, u(r) is the valence of node type r, determined from data; and h(k) is the capacity for edge type k. Specifically, h(k) = 1, 2, 3, 1.5 for single, double, triple, and aromatic bonds, respectively. Moreover, u(r) and h(k) are both 0 for phantom nodes and edges.

heading: E G3 PROPERTIES
text: A typical pitfall for many graph-based models is the ambiguity under node permutations. Reordering the nodes does not change the graph itself, but the latent representation of the graph may change if not properly modeled. Moreover, a single latent vector may not be able to decode an adjacency tensor with an arbitrary node permutation. Thanks to the geometry, our model G3 is able to identify node correspondence and maintain desirable permutation and translation invariance.
Proposition E.1. The G3 encoder and link predictor are invariant to input node permutations as well as translations on the coordinates.
Proof. The G3 encoder is composed of a point encoder and a graph encoder. The point encoder is invariant to node permutations because the MLP is weight sharing for all points and the global pooling is permutation invariant. The graph encoder is also invariant to node permutations because the graph convolution layers of MoNet are. The overall latent vector is a concatenation of the outputs of these two encoders and thus is also invariant to node permutations. Additionally, the point clouds are preprocessed so that each is centered at the origin, invariant to translations.
The link predictor takes pairwise nodes as input and is therefore invariant to node permutations. In addition, node coordinates are used always in the form of displacements. Hence, the link prediction is always unchanged regardless of translation.

heading: F QUALITATIVE EXAMPLES
text: In Figure 5, we show the generated examples for QM9 as well as results of latent space interpolation. They correspond to the results reported in Table 1. The generated samples and interpolation results for ChEMBL250k are in Figure 6, which correspond to the results reported in Table 2.

heading: G EXPERIMENTS ON THE STANDALONE LINK PREDICTOR
text: In this section, we empirically demonstrate the effectiveness of our link predictor as a standalone model. The task is to predict the edge types for all edges in the graph given the ground truth node types and coordinates for all nodes. We train the link predictor on QM9 and ChEMBL250k until convergence, then report the accuracy and the F-1 score on the test set. The classification reports are computed by using scikit-learn, where the "ZERO" type corresponds to the non-existence of a bond. In addition, we provide visualized comparisons between the ground truth molecules and the reconstructed molecules with ground truth nodes and predicted edges.
We remark that while link prediction is pairwise on nodes, the inclusion of neighbor information allows it to faithfully reconstruct complex structures like aromatic rings. Overall, G3’s link predictor attains near perfect aggregate F-1 score and accuracy on both QM9 and ChEMBL250k. It therefore serves as a viable standalone model for bond prediction in molecules.
H VISUALIZATION OF GEOMETRIC FEATURES
In this section, we provide qualitative examples for the learned geometric features in the point encoder. Similar to Wang et al. (2021), we present the channel-wise activation before the global pooling step in the point encoder of G3 pretrained on ChEMBL250k. Intuitively, the difference in the level of activation reveals the geometric primitives that the neurons focus on.
In general, we observe that G3’s point encoder can identify aromatic rings in molecules based purely on geometry; i.e. no bond information is provided. For example, Figure 9 shows that the 63rd neuron identifies rings to the right end of the molecules (see the red atoms).
In addition, different neurons can partition a molecule into meaningful sub-components. In Figure 10, we show neurons that are excited by the five-member ring on the left, the benzene ring in the middle, and the benzene ring on the right, respectively. In Figure 11, the top figure presents a neuron that identifies two rings on the right as a group, whereas the bottom figures showcase more discerning neurons that distinguish the two rings separately.

Review summaries:
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. 

This paper studies the problem of geometric graph generation (mainly focusing on molecule graph generation). Specifically, the authors propose a new method,  namely Geometric Graph Generator (G3), which generates three-dimensional geometric graphs. Different from others, G3 can capture both the combinatorial and the geometric information of graphs. The proposed algorithm can successfully generate molecule graphs and runs faster than other baseline methods. The key parts of G3 are the point encoder and graph encoder, and the three-step decoding process, that is the coordinate decode, feature decoder, and finally graph decoder.  The experimental results show that G3 has better performance than other generators in terms of validity, uniqueness, and novelty.