Generating Realistic 3D Molecules with an Equivariant Conditional Likelihood Model
Authors: James P. Roney
Paul Maragakis
Peter Skopp
David E. Shaw
Venue: ICLR 2022 Submitted (2022)
Keywords: Generative Models
Molecular Graphs
3D Molecules
Drug Discovery
Equivariance

Abstract: The number of drug-like molecules that could potentially exist is thought to be above $10^{33}$, precluding exhaustive computational or experimental screens for molecules with desirable pharmaceutical properties. Machine learning models that can propose novel molecules with specific characteristics are powerful new tools to break through the intractability of searching chemical space. Most of these models generate molecular graphs—representations that describe the topology of covalently bonded atoms in a molecule—because the bonding information in the graphs is required for many downstream applications, such as virtual screening and molecular dynamics simulation. These models, however, do not themselves generate 3D coordinates for the atoms within a molecule (which are also required for these applications), and thus they cannot easily incorporate information about 3D geometry when optimizing molecular properties. In this paper, we present GEN3D, a model that concurrently generates molecular graphs and 3D geometries, and is equivariant to rotations, translations, and atom permutations. The model extends a partially generated molecule by computing a conditional distribution over atom types, bonds, and spatial locations, and then sampling from that distribution to update the molecular graph and geometries, one atom at a time. We found that GEN3D proposes molecules that have much higher rates of chemical validity, and much better atom-distance distributions, than those generated with previous models. In addition, we validated our model’s geometric accuracy by forcing it to predict geometries for benchmark molecular graph inputs, and found that it also advances the state of the art on this test. We believe that the advantages that GEN3D provides over other models will enable it to contribute substantially to structure-based drug discovery efforts.

Under review as a conference paper at ICLR 2022
GENERATING REALISTIC 3D MOLECULES WITH AN EQUIVARIANT CONDITIONAL LIKELIHOOD MODEL
Anonymous authors Paper under double-blind review
ABSTRACT
The number of drug-like molecules that could potentially exist is thought be above 10 precluding exhaustive computational or experimental screens for molecules with desirable pharmaceutical properties. Machine learning models that can propose novel molecules with specific characteristics are powerful new tools to break through the intractability searching chemical space. Most of these models generate molecular teaphss-repreeettttionnn that describe the topology of covalently bonded atoms in a molecule-because the bonding information in the graphs is required for many downstream applications, such virtual screening and molecular dynamics simulation. These models, however, do not themselves generate 3D coordinates for the atoms within molecule (which are also required for these applications), and thus they cannot easily incorporate information about 3D geometry when optimizing molecular properties. In this paper, we present GEN3D, a model that concurrently generates molecular graphs and 3D geome- tries, and equivariant to rotations, translations, and atom permutations. The model extends partially generated molecule by computing a conditional distri- bution over atom types, bonds, and spatial locations, and then sampling from that distribution to update the molecular graph and geometries, one atom at time. We found that GEN3D proposes molecules that have much higher rates of chemical validity, and much better atom disttnee distributions, than those generated with previous models. In addition, we validated our model's geometric accuracy by forcing it to predict geometries for benchmark molecular graph inputs, and found that also advances the state of the art on this test. We believe that the advantages that GEN3D provides over other models will enable contribute substantially to structure-basee drug discovery efforts.
INTRODUCTION
Identifying molecules with desirable characteristies is of fundamental importance in many fields, including drug discovery. The astronomically large number of possible drug-like compounds, how- ever, makes exhaustive experimental or virtual screening intractable (Polishchuk et al., (013). As result, large body of work has used machine learning to explore chemical space and propose molecules with specific characteristics.
Organic molecules can be represented as graphs in which nodes are individual atoms, and edges are covalent bonds through which atoms share electrons. Each node in molecular graph is labelled with the atomic number of its corresponding atom, and edges are labelled with the number of electrons shared in that covalent bond. Because each atom has predetermined number of electrons with which to form bonds, only a subset of the possible edge labels result in chemically valid molecules.
The connectivity- description of a molecule provided by molecular graphs is important for many applications, including chemical synthesis, virtual screening, and molecular dynamics sim- ulation, and has thus motivated extensive research into generative models for molecular graphs GombzrBoo et al.) 2018 Maragakis et al., 2020 Liu et al.) 2018 Shi et al. 2020 Jii al.) 2018 2020). The functional characteristics of a given molecule, however, arise not only from its connectivity, but from its configuration in 3D space. Physical limitations
the lengths and an-
gles of covalent bonds among specific types of atoms impose constraints on the set of geometric configurations that compatible with given molecular graph. Because most organic molecules

---

Under review as a conference paper at ICLR 2022
contain rotatable bonds, however, these constraints are generally not sufficient to unambiguously construct the 3D coordinates of all atoms in the molecule. The probability of observing particular 3D geometry for given molecule is in general a function of the quantum mechanical energy of that geometry, with low-energy geometries being more likely.
Graph-based generative models are capable of producing wide range of molecules whose atoms have the correct number and type of bonds, but their outputs do not contain any geometric informa- tion, and may include molecular graphs for which stable 3D geometries (i.e., low-energy geometries with physically realistic inter-atomic distances and bond angles) do not exist. Although it is possi- ble that model could learn to assess the geometric feasibility of molecular graph without training on any geometric data, model that is trained explicitly on 3D molecules would likely be better at generating molecular graphs with corresponding low-energy 3D geometries. The lack of geometric information is also problematic because many downstream applications require 3D information. For instance, many drug discovery efforts screen for molecules with high predicted affinity for target protein pocket using computational docking process to score molecule poses. It would be ex- tremely valuable if a machine learning model could directly generate such molecules, as this would expedite the screening process. A model that is explicitly trained to recreate the geometric poses of bound molecules, rather than just molecular graphs, has the potential to produce molecules and geometries that are more conducive binding.
To address the above shortcomings, we created GEN3D-a a graph-generative model that proposes molecules with 3D coordinates. The model is rotationally and translationally equivariant, providing an inductive bias that exploits the symmetries of chemical space. GEN3D creates molecular graphs through a sequential sampling process, like other graph-based generative models (Shi et al. 2020 Liu et al. 2018). but at each step GEN3D also calculates a likelihood function over positions in 3D space, and uses that function to sample coordinates for each new atom. We show that GEN3D generates novel, chemically valid molecular graphs that have realistic low-energy geometries. In particular, we found that our model outperforms existing graph generators in its ability to create chemically valid molecular graphs, and outperforms existing 3DD-generative models in the ability to create realistic geometries. In Appendix E, we also show that our pre-trained model can be tuned to generate novel molecules in geometries that score well in virtual screening tool, which illustrates the potential application of our method in drug discovery. Finally, we demonstrate the geometric accuracy of our model by using it to sample geometries for fixed molecular graphs, and show that it achieves state-of the-art art results on the previously established benchmark of Xu et al. (2021b).
RELATED WORK
Different models construct molecular graphs using Graph Neural Networks (GNNs) in a number of different ways. CGVAE, for example, uses a GNN to create VAE latent representations for each node in a molecular graph (Liu et al., 2018). The nodes and edges of the graph are then reconstructed in sequential decision process, guided by another GNN, in which atoms are connected one at time to a growing molecular graph. Another model, GraphAF, also generates molecular graphs one atom at a time, but it samples molecules using Gaussian Autoregressive Flow (Shi et al. 2020). Other models, like JT-VAE (Jin et al.) 2018) and EierVAE (Jin et al., 2020), generate molecules by sequentially adding multi- aaoti motifs rather than individual atoms.
Because of their importance in classification tasks on point clouds and 3D molecules, there has been a great deal of interest in machine learning problems involving 3D structured data. When working with such data, it is often desirable that models be invariant or equivariant to rotations or translations of the input data, as these transformations are only an arbitrary change of coordinate systems. Ap- proaches like SchNet (Schütt et al. 2017) and EGNN (Satorras et al. 2021b) achieve equivariance by using invariant features like pairwise distances as inputs to the model. These approaches have proven effective in many domains, but they are reflection- reflection-invariant, invaiiantt so they cannot distinguish be- tween mirror-image isomers. Other approaches, like DimeNet (Klicpera et al. (020) and SphereNet (Liu et al.; 2021), achieve equivariance through message- passing schemes that encode distance and angular information, while Tensor Field Networks (Thomas et al.) 2018) and Transsorrassrrmrrs (Fuchs et al. 2020) process data using basis of equivariant spherical harmonic functions.
There are several prior works on generative models for 3D molecular data. Two models from the same group generate 3D molecules by outputting a voxelized grid of atomic densities, which is
2

---

Under review as a conference paper at ICLR 2022
then converted to set of atoms and bonds in a secondary optimization step (Ragoza et al. 2020 Masuda et al. 2020). While efficient to implement using convolutions, the voxelized grid approach imposes practical limit on the size of generated molecules, and is not equivariant. The process of reconstructing a molecule from a density field also suffers from relatively high rates of chemical and geometric invalidity, and the one-shot nature of the generative process does not allow for valence constraints to be enforced. Another pair of recent papers aims to generate stable 3D molecules using reinforcement learning (Simm et al. 2019a b.. These models, however, are only capable of generating molecules by placing atoms from a set corresponding to a pre-determined stoichiometry, have only been demonstrated on very small molecules, and do not generate bonding information.
Two other 3D -D-generative models are closely related to GEN3D: E-NF (Satorras et al. 202 and G-SchNet (Gebauer et al. 2019). E-NF uses an equivariant, EGNN-based normalizing flow to con- vert random initial atom positions into realistic molecular geometries. This model, however, has only been demonstrated on relatively small molecules, very expensive to train, and produces molecules with low rates of chemical validity. G-SchNet is an equivariant model that autoregressively gener- ates 3D molecules, and has recently been extended to conditionally generate molecules with desired electronic properties and molecular fingerprints (Gebauer et al., (02T). G-SchNet, however, exclu- sively generates atomic positions, and does not generate molecular graphs that provide the chemical bonding information required by many downstream applications. In addition, without explicitly generating bonds, one cannot use simple rules to constrain chemical properties, like atomic valences and ring sizes, during the generative process.
There are also machine learning methods that predict molecular geometries for a given molecu- lar graph. A number of works have approached this problem by predicting inter-atomic distances from the molecular graph, and then using the predictions to generate a 3D geometry in a secondary optimization step (Xu et al. 2021a Simm Hernández-Lobate 2019). Most recently, Xu et al. (221b) proposed the Conf architecture, which estimates inter-atomic distances and then opti- mizes molecular geometry using an end-to-end differentiable optimization procedure.
PROBABILISTIC MODEL
We will represent molecule as a 3D-dimensional graph G (X) For a molecule with n atoms, V E Rgxd is a list of d-dimensional atom features, A E Raxnxb is an adjacency matrix with b-dimensional edge features, and X RB*3 is list of 3D atomic coordinates for each atom. In practice, V simply encodes the atomic number of each atom, and A encodes the number of shared electrons in each covalent bond. To model a chemical space interest, we consider the distribution p(V, A, X). Previous works have aimed at calculating various marginal and conditional densities with respect to this joint distribution. For example, graph-basee generative models learn the marginal distribution p(V. A) Jx PVV, X)dX, molecular geometry prediction amounts to learning the conditional distribution p(X[V, A), and 3D generative models like G-SchNet learn the distribution pVV = PVV A, X ddA. To learn the joint distribution p(V, A, X), it is necessary to factorize the density. We choose to use the following factorization:
P(V, = II p(V X., And V A X.-1,, - p(stop|V,
T1
II p(X., A.,, X. )p( A A4-1, 1,
Xa- p(stopp V, A, X)
Here, is the number of atoms in the input graph, and Vn,, A. and X,, indicate the graph (V, A.) restricted to the first atoms. Computing p( VV amounts predicting single atom type based on 3D graph (V.21,, A, X,--)). Calculating p(A,i) Vi, A,-1,, adjacency matrix. More concretely, computing the conditional density of A. E I Rxxixb amounts to
A,-11, X,+-1) ) is relatively simple because
is more complex because it involves a prediction over a new row of the
computing a joint density over the new entries of the adjacency matrix A. A E R To solve this problem, we further decompose this distribution:
p(A,, Va, A X P(A.. A Via A
II p( (A -1 VA, X
3

---

Under review as a conference paper at ICLR 2022
Intuitively, At, A,,-1 represent the edges from atom to atoms 1... 1.
Finally, estimating the density p(X,, Vas Ay, X.-1) involves modeling continuous distribution over positions X; E R3 for atom To accomplish this, we assume X, belongs finite set of points X, and model its probability mass as a product distributions over angles and distances:
P(X, Au, X
MIIX, X,ll Va, Au, Xái1)
C
II p(Angle(X, X, - X&) Vi, Ai, X3-1)
GRReZ
Vi,, A., -1) predicts the distances from each existing atom to the new atom, and p(Angle(X; XR,X, Xx) Vi A. 11) predicts the bond angles of connected triplets of atoms involving atom i. is set of pairs (j. k) where atom k is connected to atom i, and atom j is connected to atom k. "Angle" denotes the angle between two vectors.
Intuitively,
normalizing constant derived from summing this density over all of X. To increase the computational tractability molecular graph (V, A, X) are
of estimating this factorized density, we assume that the nodes listed in the order of breadth-firs' traversal over the molecular graph.
4 ARCHITECTURE AND TRAINING
Atom Network
Edge Network
Distance Network
Angle Network
Figure 1: The process by which GEN3D samples 3D molecule.
The GEN3D architecture uses collection of four equivariant neural networks: the atom network (denoted FA). the edge network (denoted Fg), the distance network (denoted Fp), and the angle network (denoted Fo). Each of these networks is implemented as 7-layer EGNN with a hidden dimension of 128, as described in Satorras et al.) (2021b). We found that adding batch normalization to the atom, distance, and angle networks improved stability (Toffe & Szegedy 2015). An EGNN network takes in a 3D graph as input, and outputs vector embeddings for each node in the input graph. GEN3D also uses four simple MLPs, DA, Dg. DD, and Dg, to decode the output embed- dings of each EGNN into softmax probabilities. GEN3D's subnetworks used to compute the components of the factorized density above as follows:
P(V.VVa-1. A 1 Softmax (P Xxx)))))
MA,glArss- Va, A4- X Sottma((D Va, A4-1, X4-1);))
4

---

Under review as a conference paper at ICLR 2022
P(|IX X,1l Via, A4, 11) = Softmax(D A4, 1)))
h = Fo(V..; Aii xa-) p(Angle(X: Nx, x, - xx) Va, As, Xa-1) = Softmax Do (hy. hw))
Note that the predicted distance and angle distributions are discrete softmax probabilities. These discrete distributions correspond to predictions over equal-width distance and angle bins. Because all of the comm-ucoptted densities are insensitive to translations and rotations of the input graph, the full product density is also insensitive to these transformations.
At training time, we compute a breadth-first decomposition of a graph XX). The subnetworks of GEN3D are trained to autoregressively predict the next atom types, edges, distances, and angles in this decomposition according to the model described above. We use cross-entropy losses to pe- nalize the model for making predictions that deviate from the actual next tokens in the breadth-first decomposition. While the model's 's density is not invariant across different breath-fing decompo- sitions of the same molecule, we resample each molecule's decomposition at every epoch, so the model should learn to ascribe equal densities to different rollouts of the same molecule.
The training algorithm is provided in detail in Appendix F. Our experiments use the Adam optimizer with a base learning rate of 0.001 (Kingma & Bal 2014). All models were able to train in approx- imately one day on a single NVIDIA A100 GPU. The model is trained using teacher forcing, so only learns to make accurate predictions when given well-formed structures as autoregressive inputs. This could make the model more brittle at generation time because will struggle recover from its own mistakes in previous iterations. To increase geometric robustness and mitigate this issue, we add uniform random noise of up to .05 Á the atomic coordinates during training for all datasets.
To sample a 3D molecule from a trained GEN3D model, we can start with a single initial atom or larger molecular fragment. First, the atom network computes a discrete distribution over new atom types to add, from which a new atom type can be sampled multinomially. The edge network is then used to sequentially sample the edge types joining the new atom to each of the previously generated atoms. The distance and angle networks compute distributions over interatomic distances and bond angles involving the newly sampled atom. To sample the new atom's position, we construct the discrete set of points X' as a fine grid surrounding the previously-genette point a probability according to the model's 's distance and angle predictions. Finally, the new atom's 's position is sampled multinomially from the set X. The resulting molecular graph, which has been extended by one atom, is then fed back into the autoregressive sampling procedure until a stop token generated. This sampling process is illustrated in Figure 1, and described in detail in Appendix H.
atoms, and assign each
5 RESULTS
We trained GEN3D to generate 3D molecules from three datasets: QM9, GEOM-QM9 and GEOM-
Drugs (Ramakrishnan et al.) 2014, Axelrod & K66mez-Bombarelll p020). QM9 contains 134,000 small molecules with up nine heavy atoms (i.e., not including hydrogen) of the chemical elements C. N, O and F. Each molecule has a single set of 3D coordinates obtained via Density Functional Theory calculations, which approximately compute the quantum mechanical energy of a set of atoms in 3D space. GEOM-QM9 contains the same set of compounds as QM9, but with multiple geome- tries for each molecule. BEOM--rugs also has multiple geometries for each molecule, and contains over 300,000 drug-like compounds with more heavy atoms and atomic species than QM9,
On QM9, we trained one version of the model with heavy atoms only, and one version with hydro- gens. To ensure the quality of our geometric data, we used OpenBabel (O'ooyle et al. 201T) to convert the coordinates from the QM9 source files into SDF files, which contain both coordinates and connectivity information inferred from inter-atomic distances. We discarded all molecules for which the inferred connectivity did not match the intended SMILES string (Weininger) 1988) from the QM9 source data, leaving approximately 124,000 molecules with SDF-formatted bonding in- formation. We used 100,000 of these molecules for training, with the remaining molecules used for validation. For GEOM-QM9 we trained on 200,000 molecule-geomete pairs, and excluded all SMILES strings from the test set of Xu et al. 2021a).. For GEOM- -Duggs we trained with heavy atoms only, using 50,000 randomly chosen molecule- molceule-geometry pairs for training. After 60 epochs of training, GEN3D was able to generate highly realistic 3D molecules from all of these datasets. Visualizations of GEN3D samples from QM9 and GEOM-Drugs are shown in Figure 2.
5

---

Under review as a conference paper at ICLR 2022
QM9 w/o Hydrogens
QM9 w/ Hydrogens
GEOM-Drugs
Samples of molecules produced by GEN3D when trained on various training sets listed the first column. These examples have been curated to showcase the qualitative diversity of molecules created by GEN3D. For larger number of raw samples, see Appendix A.
Figure 2:
5.1 CHARACTERISTICS OF GENERATED MOLECULAR GRAPHS
To assess the quality of our generated molecules, we began by analyzing the characteristics of our generated molecular graphs with QM9. We assessed the percentages novel and unique molec- ular graphs generated by the heavy atom QM9 model in sample of 10,000 molecules. A novel graph not present in the training data. The uniqueness rate is defined
molecular graph is defined by the number of distinct molecular graphs generated, divided by the total number of molecules generated. We focused on the heavy atom model for these metrics because most existing molecular graph generators use only heavy atoms. As mentioned previously, selections at each step of the generative process such that 100% of generated molecular graphs have the correct number of bonds for each atom. We also assess the percentage of molecules that are chemically valid, even when no masking performed, and doing so gives an indication of the model's understanding of fundamental chemical constraints. Using the results for novelty, valid- ity, and uniqueness metrics, we compared our approach against GraphAF and CGVAE, which are two recent molecular graph generators that also add one atom time. We also compared against GEN3D-, a geometry-unaware baseline that was created by removing the geometric networks from GEN3D and setting all positional inputs to These results are reported in Table
mask-out atom and edge
Table 1: Properties of QM9 Molecular Graphs (Heavy Atoms)
Model
Validity
Validity (w/o check)
Uniqueness
Novelty
Strain
CGVAE
100%
49.55%
97.09% 94.51 %
88.18%
70.84%
GraphAF
100%
67%
88,93 %
GEN3D-
100%
99.79%
95.59%
25.93%
90.97%
GEN3D (ours)
100%
98.80%
94.33%
33.18%
93.36%
Q.99 (truth) 100%
100%
92.41%
We found that, even without imposing checks at generation time, GEN3D produces molecules that obey valence constraints 98.8% of the time after training QM9, This far exceeds the unchecked validity rate of 67% achieved by GraphAF, suggesting that GEN3D has a better understanding of the basic rules of chemistry. Interestingly, ing that our improvements in chemical validity come from architectural differences that are unrelated to the generation of 3D geometries. GEN3D achieves a uniqueness of 94.3%, which is simi-
geometry-fre baseline achieves 99.8% validity, suggest-

---

Under review as a conference paper at ICLR 2022
lar to the rates for GraphAF and CGVAE. We also assessed the geometric feasibility of generated graphs by converting them into 3D coordinates using CORINA (Sadowski & Gasteiger 1993), and then computing the volume of the tetrahedron enclosed by each sp tetrahedral center, with vertices located Á along each tetrahedral bond. Graphs that could not be converted with CORINA, or contained tetrahedral centers with volumes less than 0.345 cubic Å, were classified as being overly strained. This test was adapted from Ruddigkeit et al. (0012), where it was used to filter overly strained graphs from the GDB-17 dataset. We found that GEN3D produced fewer overly strained molecules than other models, including GEN3D-, suggesting that explicitly generating molecular geometries helps bias the model towards stable compounds.
Interestingly, GEN3D showed novelty score of only 33.18%. which is far lower than the scores reported for previously published models. We are confident, however, that our model is not overfit to its training data. Of molecules created by GEN3D that matched a molecule in our QM9 dataset, 18.6% matched a molecule in the held-out validation set. Based on the relative sizes of our validation and training sets, an unbiased model would generate molecules matching the validation set 19.5% of the time. Our model thus suffers from minimal overfitting, despite its lower novelty score.
When considering the composition of the QM9 dataset, it is not particularly surprising that GEN3D achieved a relatively low novelty score. QM9 contains all physically plausible compounds with up to nine heavy atoms of C, O, N, and F (Ramakrishnan et al., 2014). It is derived from the massive GDB-17 dataset, which enumerates the physically plausible compounds with up to 17 heavy atoms (Ruddigkeit et al.) 2012). For model to generate a molecular graph outside of the QM9 dataset, that molecule must thus have more than nine heavy atoms, or violate basic chemical constraint that is satisfied by all compounds in the GDB-17 database. Models like CGVAE and GraphAF report >90% novelty rates when trained on the majority of compounds in QM9, meaning that their generated molecules violate the constraints of the QM9 data distribution very frequently. To further explain the discrepancy in novelty scores between GEN3D and other models, we analyzed why molecules generated by CGVAE and GEN3D are present or absent from QM9, and present the results in Appendix B. We found that CGVAE produces far more molecules with over nine heavy atoms, which artificially increases its novelty score. In contrast, GEN3D almost always generates molecules with 9 fewer heavy atoms. This is not because of any hard cap on the number of atoms in molecules produced by GEN3D. Rather, GEN3D has learned to obey the limit of nine heavy atoms found in the training data, whereas CGVAE has not. Because GEN3D learns to stay within this limit so accurately, and QM9 is an exhaustive dataset, this leads to lower novelty scores that approach the proportion of the dataset that was withheld during training-thhe optimal outcome for a model that aims to generate the region of chemical space enumerated by QM9.
5 .2 ACCURACY OF MOLECULAR GEOMETRIES
After assessing the quality of our generated graphs, we sought to verify the quality of the 3D ge- ometries produced by GEN3D. We compared our model to E-NF, a non-equivariant version of E-NF called GNF-attention, and G-SchNet, which are the only other published models that generate sam- ples from the distribution of 3D QM9 molecules. Both E-NF and G-SchNet produce the positions of heavy atoms and hydrogens as the output of their generative process. We compared these models to our all-atom QM9 model. The E-NF paper reports atomic stability as the percentage of atoms that have correct number of bonds, and molecular stability as the fraction of all molecules with the correct number of bonds for every atom. We report these metrics in Table and compare GEN3D to E-NF, G-SchNet, and related baselines.
GEN3D outperformed all other models, achieving 97.5% molecular stability without any valence masking, compared to 77% for G-SchNet and 4.3% for E-NF. In order to assess the geometric realism of the generated molecules, the authors of E-NF computed the Jensen-Shannon divergence between a normalized histogram of inter-atomic distances and the true distribution of pairwise dis- tances from the QM9 dataset. We also computed this metric, and found that GEN3D advances the state of the art, reducing the JS divergence by a factor of two over G-SchNet and a factor of four over E-NF. The fact that GEN3D substantially outperforms E-NF and G-SchNet, both of which only generate coordinates and do no generate bonding information, suggests that generating bonds as well coordinates significantly increases the quality of generated molecules.
7

---

Under review as a conference paper at ICLR 2022
To confirm this, we conducted a systematic ablation study (see Appendix G) in which we succes- sively removed the angle and edge networks of GEN3D to produce a baseline model that very similar G-SchNet. We found that performance in both geometric and chemical accuracy metrics dropped continuously as we removed these features, and that our baseline model performed very similarly to G-SchNet. In addition, GEN3D is much cheaper to train than E-NF's flow-based gen- erative process, and it is applicable to larger, drug-like molecules. These comparisons are reported in Table 2, and the true and learned histograms of pairwise distances are plotted in Appendix C. In order to be consistent with the E-NF paper, the Jensen- -hasennShannon divergence was only computed between generated and QM9 molecules with exactly 19 total atoms.
Table 2: Properties of QM9 Molecules (3D models with hydrogens)
Model
Atom Stability
Mol Stability
Distance JS
GNF-attention
72%
.3%
.007
E-NF
84%
4.2%
.006
G-SchNet
98.7%
77%
.0031
GEN3D (w/o check)
99.7%
97.5% %
.0014
GEN3D (w/ check)
99.87 %
99.1% %
.0014
QM9 (truth)
99.99%
99.9%
0
The Jensen-Shannon Shannon divergence metric provides confidence that GEN3D is, on average, generating accurate molecular geometries. This metric, however, is relatively insensitive to the correctness of individual molecular geometries because it only compares the aggregate distributions of distances. In order to further validate the accuracy of GEN3D's generated geometries, used GEN3D to pre- dict the geometries of specific molecular graphs, and compared its accuracy with purpose-buie built tools designed for molecular geometry prediction, such the model described in Xu al. (2021b). This evaluation amounts to verifying the accuracy of the conditional distribution p PV A) when the joint distribution p(V, A, X) is learned by GEN3D. We approximate this conditional distribution by using a search algorithm to identify geometries X that give high value to pVV, A, X), as calculated by GEN3D when V and A are known inputs. This method is described in detail in Appendix D.
To evaluate the ability of GEN3D to predict molecules from FEOM-QMS (Axelrod & molecular geometries, we trained GEN3D to generate 2020). We then followed the evalua- tion protocol described in Xu et al. (2021a) with the same set of 150 molecular graphs, which were excluded from our training set. As in prior works, we predicted an ensemble of geometries, and then computed COV and MAT scores with respect to the test set. The COV score measures what fraction of reference geometries have a "close" neighbor in the set of generated geometries, where closeness is measured with aligned RMSD threshold; we used threshold of 0.5 Å, following Xu et al. (021a). The MAT score summarizes the aligned RMSD of each reference geometry to its closest neighbor in the set of generated geometries (for the full evaluation protocol, see Xu et al. (2011)).
GEN3D achieved results that are among the best for published models on both metrics. In particu- lar, its MAT scores outperformed all prior methods that do not refine geometries using rules-based force field. We compared GEN3D with previous machine learning models for molecular geome- try prediction, as well as the ETKDG algorithm implemented in RDKit (which predicts molecular geometries using database of preferred torsional angles and bond lengths (Riniker & Landrum, 2013)). Table shows the results of these comparisons, and Figure visualizes representative geom- etry predictions. These data indicate that GEN3D is accurately sampling from the joint distribution of molecular graphs and molecular geometries.
CONCLUSIONS
In this paper, we have introduced GEN3D: an autoregressive model for generating molecules in 3D space that is rotationally and translationally equivariant. We have demonstrated that GEN3D produces 3D molecules that are chemically stable and geometrically realistic. and achieves of ahe-art results multiple benchmarks in molecular machine learning.
We believe models like GEN3D will have many practical applications in chemistry and drug discov- ery. A natural extension of GEN3D is model that generates molecules in 3D space occupied by
8

---

Under review as a conference paper at ICLR 2022
RDKit
GEN3D
Reference
KV50 .56
RXb0 .55
RMS0 + .27
RA55
RVGD
RM60
3/50 .22
RM50
RV6D - 41
RVSD
RMSD 00
R
ARA
RV600
RAB0
RMSD
RMGG
RMS0
Figure 3: Predicted geometries for GEOM-QM9, The right columns contains reference geome- tries, and the left two columns show the nearest neighbor to the reference geometries among the geometries generated by RDKit and GEN3D.
Table 3: Geometry Prediction on GEOM-QM9
COV (%) Median
MAT (A)
Metric
Mean
Mean
Median
GraphDG
55.09 %
56.47 %
0.4649
0.4298
CGCF
69.60 %
70.64 %
0.3915
0.3986
ConfVAE
77.98 %
82.82 %
0.3778
0.3770
RDKit
80.68 %
87.50 %
0.3349
0.3245
GEN3D (ours)
73.62 %
77.14 %
0.3168
0.3049
features of a protein pocket, such as pharmacophore features or amino acid side chain atoms. Such model could learn to generate potential binding molecules directly inside protein pocket. The generated molecules and geometries could then move directly into downstream stability screenings using molecular dynamics simulations docking programs. This approach could greatly decrease the amount of computation needed to select and verify promising drug molecules. Quantum mechan- ical energy calculations, docking scores, and molecular dynamics simulation results from generated poses could also be used to train the model via reinforcement learning.
Conditional likelihood models like GEN3D can also work together with other drug discovery meth- ods. For instance, if virtual screening campaign identifies several small molecules that each bind tightly to a different region in a protein pocket, GEN3D could be used to propose large molecules that retain the respective poses of the small molecules, and join them together in a chemically and geometrically feasible way. Alternatively, molecules that each bind to the same region of protein pocket and share moiety in the con- text their bound poses, GEN3D could also be used to extend this common molecular scaffold in new ways. Such applications of GEN3D could help with the enumeration of possible side chain substitutions and the closure of rings, which are common tasks is drug design.
screening campaign identifies a family of small
While GEN3D already demonstrates impressive performance, it only stands benefit from recent architectural developments 3D machine learning. For example, GEN3D is currently reflection- invariant due its use of EGNN subunits, but this could be fixed by substituting the EGNN with reflection-sensitive networks like DimeNet or SphereNet (Klicpera et al. 2020 Liu et al. 2021).
Beyond the fields of chemistry and drug discovery, GEN3D provides simple equivariant approach for generating 3D graphs, and is much cheaper to train than the flow- based generative models used in previous works. We hope GEN3D and similar subsequent models will find diverse uses in generating and optimizing structures in 3D space.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
In order to facilitate reproducibility, we plan to release the source code for the GEN3D model prior publication, together with pre-trained GEN3D model weights and our curated version of the QM9 dataset.
10

---

Under review as a conference paper at ICLR 2022
REFERENCES
Simon Axelrod and Rafael for property prediction and molecular generation. arXiv preprint arXiv:2006.0555. 05000, 2020.
Geom: annotated molecular conformations
Edsger W. Dijkstra. A note on two problems in connexion with graphs. Numerische mathematik, (1):269-271, 1959.
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Stanssrranssormmrs:: 3D) noto-translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Niklas W. A. Gebauer, Michael Gastegger, and Kristof T. Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. arXiv preprint arXiv:1906.0093. 2019.
Niklas W. A. Gebauer, Michael Gastegger, Stefaan S. P. Hessmann, Klaus-Robert Müller, and Kristof T. Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. arXiv preprint arXiv:1906.00995 2021.
J. Andrew Grant, Maria A. Gallardo, and Barry T. Pickup. A fast method of molecular shape comparison: A simple application gaussian description of molecular shape. Journal of Com- putational Chemistry, 17(14)11653-1666, 1996.
Rafael Gömez-Bombarelli, Jennifer Wei, David Duvenaud, José Miguel Hernández-Lobato, Timothy D.. Hirzel,
Benjamín Sánchez-l Sánchez-Lengelingg engelingg Dennis Sheberla, Jorge Ryan Adams, and Alán Aspuru-Guzik. Automatic chemical design using data-driven contin- uous representation of molecules. ACS Central Science, 4(2):268-276, 2018.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448 148-456. PMLR, 2015.
Paul Jaccard. The distribution of the flora in the alpine zone. New Phytologist, 11(2):37-50, 1912.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junetion tree variational autoencoder for molecular graph generation. In International Conference on Machine Learning, pp. 2323-2332. PMLR, 2018
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Hierarchical generation of molecular graphs using structural motifs. In International Conference Machine Learning, pp. 9839-4844, PMLR, 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14/2.6900, 2014.
Johannes Klicpera, Janck Groß, and Stephan Günnemann. Directional message passing for molec- ular graphs. arXiv preprint arXiv:2003.03123, 2020.
Greg Landrum et al. Rdkit: Open-source cheminformatics software. http://www.rdkit.org YVij6EbMJZU, 2021.
Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander L. Gaunt. Constrained graph variational autoencoders for molecule design. arXiv preprint arXiv: 00000009766, 2018.
Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3D graph networks. arXiv preprint arXiv:2102.0500 2021.
Paul Maragakis, Hunter Nisonoff, Brian Cole, and David E. Shaw. A deep- deep-learning view of chemical space designed facilitate drug discovery. Journal Chemical Information and Modeling, 60 (10):4487-4496, 2020.
Tomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3D molecular struc- tures conditional on a receptor binding site with deep generative models. arXiv preprint arXiv:2010. 14442, 2020.
11

---

Under review as a conference paper at ICLR 2022
Noel O'Boyle, Michael Banck, Craig James, Chris Morely, Tim Vandermeersch, and Geoffrey Hutchison. Open Babel: An open chemical toolbox. Journal of Cheminformatics, 3(33), 2011.
Emanuele Perola and Paul S. Charifson. Conformational analysis of drug-like molecules bound to proteins: an extensive study of ligand reorganization upon binding. Journal of Medicinal Chemistry, 47(10):2499-255 2510, 2004.
P. G. Polishchuk, T. I. Madzhidov, and A. Varnek. Estimation of the size of drug-like chemical space based on GDB-17 data. Journal of Computer-Aided Molecular Design, 27(8))675-679, 2013.
Matthew Ragoza, Tomohide Masuda, and David Ryan Koes. Learning a continuous representation of 3D molecular structures with deep generative models. arXiv preprint arXiv:2010.08687, 2020.
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(140022), 2014.
Sereina Riniker and Gregory A. Landrum. Better informed distance geometry: Using what we know to improve conformation generation. Journal of Chemical Information and Modeling, 55 (12):2562-257)) 2015.
David Rogers and Mathew Hahn. Extendeed-comnectvityy fingerprints. Journal of Chemical Infor- mation and Modeling, 50(5):742-754, 2010.
David J. Rogers and Taffee T. Tanimoto. A computer program for classifying plants. Science, 132 (3454): 1960.
Lars Ruddigkeit, Ruud van Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumcration of 166 billion organic small molecules in the chemical universe database GDB-17. Journal of Chemical Information and Modeling, (2 1):2864 (2(11):28642287,, 2875, 2012.
Jens Sadowski and Johann Gasteiger. From atoms and bonds to thrce-dimensional atomic coordi- nates: automatic model builders. Chemical Reviews, 93(7):256772588 1993.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows for molecule generation in 3D. arXiv preprint arXi::2105.09016 2021a.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- works. arXiv preprint 102:992102.09944,, 2021b.
Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller, Schnet: A continuous- filter convolutional neural network for modeling quantum interactions. arXiv preprint arXiv: 1706008666 2017.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. GraphAF: flow-based autoregressive model for molecular graph generation. arXiv preprint arXiv:2001.09388 2020.
Gregor N. C. Simm and José Miguel Loraande A generative model for molecular distance geometry. arXiv preprint arXiv: 2019.
Gregor N. C. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hemáánezz-Lobato Symmetry- aware actor-critic for 3d molecular design. arXiv preprint arXiv:2002.07711,,00 2019a.
Gregor N. C. Simm, Robert Pinsler, and José Miguel Hernández-LLobtto Reinforcement learning for molecular design guided by quantum mechanics. arXiv preprint 07X1v:2000.07717,220000 2019b.
Nathaniel Thomas, Tess Smidt, Steven Kearnes. Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and vranslaton-eeuuivrrinn neural networks for 3D point clouds. arXiv preprint 1802.082//0 2018.
David Weininger. Smiles, chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Journal of Chemical Information and Computer Sciences, 28(1):31-36 1988.
12

---

Under review as a conference paper at ICLR 2022
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint 102::2002..12240, 2021a.
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael
and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel program- ming. arXiv preprint arXiv:2105.07246 2021b.
13

---

Under review as a conference paper at ICLR 2022
Appendices
A
ADDITIONAL SAMPLE MOLECULES
Figure S1: Unfiltered, independent random samples from GEN3D models trained on QM9 without
hydrogens (top), QM9 with hydrogens (middle), and GEOM-Drugs (bottom).
14

---

Under review as a conference paper at ICLR 2022
B COMPOSITION OF GENERATED MOLECULES
1.0
In Training Set
In Held-Out QM9
> Heavy Atoms Geometry Error Other
0.8
Proportion
0.6
0.4
0.2
0.0
GEN3D
CGVAE
Figure S2: The composition of molecules generated from GEN3D and CGVAE.
15

---

Under review as a conference paper at ICLR 2022
C DISTRIBUTION OF INTER-ATOMIC DISTANCES
True
Learned
0.06
0.05
0.04
Probability
0.03
0.02
0.01
0.00
0
2
4
10
12
Inter-atomic Distance in Angstroms
Figure S3: Histograms of inter-atom distances for GEN3D -ggnerage and QM9 molecules with 19 total atoms.
16

---

Under review as a conference paper at ICLR 2022
D MOLECULAR GEOMETRY PREDICTION
In order to predict the geometry of a specific molecular graph, we utilized Dijkstra's algorithm to search for geometries of those molecules that assigned high likelihood (Dijkstra) 1959). The given molecular graph is unrolled in breadth-first order, so predicting the molecule's 's geometry amounts determining sequence of positions for each atom during the rollout. atomic positions are discretized, then the space of possible molecular geometries forms tree. Each edge in the tree can be assigned a likelihood by GEN3D. Predicting a plausible geometry thus amounts to finding path where the sum of the log-likelihoods of the edges is large. This can be accomplished using standard graph search algorithms like A Dijkstra's 's algorithm. The geometry prediction algorithm is presented in Algorithm I below.
Algorithm Molecular Geometry Search using GEN3D
Input: molecular graph G (V, A) with n and number of iterations T. Assume G has been processed such that the nodes of V and A and in breadth-first order.
Q EmpttyPriorityQuueell Q.push( 1, [0.0, 0.0, 0.0]))
Queue entries have the form (Negative Log Likelihood, Number of Atoms, Position Matrix)
1 + 0
results <
while < T do
nll, i, X f Q dedequee
if = n then
results. cat.saapapnnnllll X))
else
for j in {1..i} do if argmax(AA-1.) neigh + end if end for points MakeGrid(((neigg))
Find an atom that a neighbor of the next atom in the rollout
then
D Create fine grid of potential positions around the neighbor atom
for x in points do
dik pp (x: V 3+1 A441
angliik = pe (z) V A.41,
atm lik PA (V.+22 V 5+1 A StackCols(X, x))
edg lik = PE (2 (A 1 A.41,, (ta,kCols(Xxx))
nll nll log(dst_lik) log(ang log(atm_ log(edg
Q
end for
end if end while return results
We found this procedure be effective and computational feasible for molecules in GEOM-QM9. The tree of potential geometries expands rapidly with increased numbers atoms, however, lim- iting its efficiency for large molecules. Future work could identify search heuristics that make this approach more efficient.
17

---

Under review as a conference paper at ICLR 2022
E DEMONSTRATION OF PROPERTY OPTIMIZATION
We evaluated the ability GEN3D to generate 3D molecules in poses that have favorable predicted interactions with target protein pocket, evaluated by the ROCS virtual screening algorithm (Grant et al. 1996). We started from GEN3D model trained on GEOM-drugs (we call this model GEN3D-gd) -d). We curated a large pre-existing library of 62.9-million compounds, containing up to 250 molecular geometries generated with OpenEye Omega (Perola & Charifson, (004), for each compound, and screened the resulting 13.8-billion conformations against our target pocket using ROCS. We then selected the top 1000 scoring geometries belonging to distinct molecular graphs from the library, and fine-tuned GEN3D-gd on these 1000 3D molecules for 100 epochs (we call the resulting model GEN3D-ft).
To evaluate the ability of GEN3D learn chemical and geometrie features that are conducive to binding the pocket, we generated 10,000 molecules with 3D coordinates from GEN3D-gd and GEN3D-ft. In addition, we took the molecular graphs generated by GEN3D-ft and recalculated molecular geometries for them using OpenEye Omega. We excluded molecules generated by GEN3D-ft the molecular graph overlapped with the fine-tuning set (2.07% of the total), and scored the remainder using ROCS. We found that fine-tuning significantly increased the scores of generated compounds. Because GEN3D-ft was fine-tuned on high-scoring molecular geometries, the molecular geometries it generated implicitly include information about the target geometry that were unavailable to GEN3D-gd and OpenEye Omega. As a result, the scores for GEN3D-ft geome- tries were, average, better than those generated by other methods. These results shown in Figures $4.
3.5
GEN3D-gd, GEN3D-gd Geometry
GEN3D-ft, Omega Geometry
GEN3D-ft, GEN3D-ft Geometry
3.0
2.5
Probability Density
2.0
1.5
1.0
0.5
0.0
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
ROCS Score
Figure S4: The probability densities of ROCS scores for molecular graphs and geometries gener- ated by GEN3D-gd (blue), molecular graphs generated by GEN3D-ft with the Omega geometries (orange), and molecular graphs and geometries generated by GEN3D-ft (green). Higher ROCS scores indicate that a model is better able to produce desirable molecules.
Ideally, this training procedure would allow models to generate strong binders that signifi- cantly different from those in the fine-tuning set. To compare each model's ability to produce both high-quality and novel compounds, picked the top 2% of molecules (by ROCS score) generated
18

---

Under review as a conference paper at ICLR 2022
by each model, and plotted their ROCS scores against their maximum Tanimoto similarity coeffi- cient to an element the set used for fine-tuning. A Tanimoto similarity coefficient (also called a Jaccard coefficient of community) ranges from fully dissimilar 0.0 to identical at 1.0, and can be used as a measure of the structural closeness of two molecular graphs (Rogers & 'Tanimot, 1960) Jaccard, 1912). It is computed by representing prints, which are essentially lists of activated bits corresponding to substructures present in each molecule (Rogers Hahn 2010), (Here, we used RDKit's implementation of Morgan fingerprints with 2048 bit, radius 2, and without chirality (Landrum et al. 202T).
molecules with Extended-Connecttoe
Finger-
We found that GEN3D-ft generated molecules with high ROCS scores across wide range of Tani- moto similarities to the finettuning set. Molecules generated by GEN3D-ft had sigmificantly higher scores than those generated by GEN3D-gd, even when comparing molecules from each model with comparable similarities to the fine-tuning set. In this particular instance, the highest ROCS scor- ing molecule generated by GEN3D-ft had Tanimoto similarity the fine-tuning set of about 0.4. These results are shown in Figure S5.
These experiments indicate that GEN3D able to shift generative distribution into specific re- gions of chemical and geometric space.
1.5
GEN3D-gd GEN3D-ft
1.4
1.3
ROCS Score
1.2
1.1
1.0
0.9
0.8
0.0
0.2 Tanimoto Similarity Coefficient to Molecules Used for Fine-tuning
0.4
0.6
0.8
1.0
Figure S5: The similarity to the fine-tuning dataset the molecules with the top 2% of ROCS scores. For clarity, include those molecules that exactly match member of the fine-tuning training set (i.e., they have a Tanimoto similarity coefficient 1.0).
19

---

Under review as a conference paper at ICLR 2022
F DETAILED TRAINING PROCEDURE
In order train GEN3D, we decompose a 3D molecule into a sequence of partially completed molecular graphs, to which new atoms, edges, and spatial locations added sequentially. We use GEN3D to autoregressively predict the next atom types, edge types, interatomic distances, and bond angles in this sequence of partially completed graphs, and penalize incorrect predictions using cross-entropy loss. Algorithm describes this process in detail.
Algorithm 2 Training GEN3D
Initial: Parameters @ of GEN3D subnetworks FA. DA, FE, DE, FD. DD, Fe, Da- while @ not converged do
Sample 3D molecule from the training set. Let denote the number of atoms.
Perform Breadth Freatth-First search, producing graph G = (V,A,X) with the nodes. in breadth-first order.
Let G.; = VV A.,) denote the partially complete graph restricted to the atoms encountered by step of the BFS.
for i in {1..m} do atom-probs +- if = then
loss + CrosEntoopy(atonmprros StopToken)
else
loss + CrossEEnropp (alom probs, V4+1)
edge.probs Zeros(i, b)
edge_accum <- Zeros(i, b)
edge_accume [ 0] for j in {1..i} do edge_accum |: 5l A V + Stack VI+1, edge.acoum)
Initialize with all unbonded edge types
Vi R (2d+5)xx Row j contains type of atom j. type of atom i + 1, type of edge A++1j
+ +V AV,, X.;)
h_edge & D (FF(((
edge probs|] + Softmax(h end for
loss loss CrossEnnropy(edgge-roooo A141,0)
G¹2 + (StackRows(V., ( A.,, X.;)
dists + |X - X
D dists E RR. Distances to the new atom.
dist.probs Sofmax (F
loss + loss CrossEntropy -probs, make.bins bin((ditss))
h (C',)) angle_loss -
angle count for in {1..i) do
for l in (1..j - 1) do if argmax(A),t)
and argmax(A4+1.t) > 0 then angle-probs + Doltmax(Dy(Connattthaoggg/ hanglle(k)))
angle +- Angle(X Xx,, X, Xx)
angle +- angle + brossEntropy(angle.poooo make count +- angle_count count +
mak.. bins(angle))
end if
end for end for
loss loss + angle loss angle count
end if
< Φ n sloss
end for end while
20

---

Under review as a conference paper at ICLR 2022
G ABLATION STUDIES
In order to determine which features of GEN3D allow for to exceed the chemical and geometric accuracy of previous models like G-SchNet and E-NF, we performed a systematic ablation study of various GEN3D subnetworks. We hypothesized that the ability of GEN3D to generate and analyze chemical bonding information (rather than exclusively generating and analyzing atom types and positions) substantially contributes to its ability to generate realistic molecules. In order to ablate this feature, we replaced GEN3D's edge network with neural network that selects single previously generated atom as a "focus atom," which will be nearby the next atom added to the molecular graph. We also deleted all edge labels from the input data, leaving only the atom types and positions for analysis. This setup is very similar to G-SchNet. We found that this model (which Focus Atom") performed similarly to G-SchNet, and generated molccules with the correct valence 82% of the time after training QM9. Allowing the network generate and analyze bonding information significantly improved performance. Adding angle predictions and enforcing valence constrains during the generative process further improved chemical and geometric realism. It should be noted that, while both of these features improve performance, they are difficult to implement in models like G-SchNet and GEN3D Focus Atom, neither of which generate bonding information.
call "GEN3D
Because G-SchNet and GEN3D Focus Atom do not generate bonding information, secondary tool must be used to infer the bonds between atoms and assess whether or not atoms have the correct valences. To accomplish this, we used the same OpenBabel-based procedure contained in the G-SchNet reference implementation. In order guarantee a fair comparison, we also applied this procedure to molecules for which GEN3D generated its own edges, in order to eliminate the possibility that the validity gap between GEN3D and G-SchNet caused by errors OpenBabel's determination of bonding information. Even when inferring bonds using OpenBabel, the GEN3D models that generated edges achieved much higher validity rates that G-SchNet and GEN3D Focus Atom. The results of the ablation studies are shown in Table 4.
Table Performance of Ablated Models on QM9 Molecules (3D models with hydrogens)
Model
Atom Stability
Mol Stability
Distance JS
G-SchNet
98.7%
77.0%
.0031
GEN3D Focus Atom
98.7%
82.0%
.0037
GEN3D No Angles(w/o check, OpenBabel edges)
99.4%
93.4%
.0030
GEN3D No Angles (w/o check, GEN3D edges)
99.6%
95.7%
.0030
GEN3D No Angles (w/ check, OpenBabel edges)
99.5%
95.5%
.0030
GEN3D No Angles (w/ check, GEN3D edges)
99.7%
98.0%
.0030
GEN3D (w/o check, OpenBabel edges)
99.6%
96.4%
.0014
GEN3D (wlo check, GEN3D edges)
99.7%
97.5%
.0014
GEN3D (w/ check, OpenBabel edges)
99.75%
97.6%
.0014
GEN3D (w/ check, GEN3D edges)
99.87%
99.1%
.0014
QM9 (truth)
99.99%
99.9%
21

---

Under review as a conference paper at ICLR 2022
H
GENERATION PROCEDURE
GEN3D is an autoregressive model that augments a partially completed molecular graph. We denote partially completed graph with atoms as G, (Va, A., X..) V₁2 Rixd is list of one-hot encoded atom types (i.c., the different chemical elements appearing in the dataset), and d is the number of possible atom types. A, Rixixb is an adjacency matrix recording the one-hot encoded bond type between each pair of atoms, with b representing the number of bond types. X; Réxa list of atom positions. For the adjacency matrix Ag, we include an extra bond type indicating that the atoms are not chemically bonded (unbonded atoms are still connected in the sense that information can propagate between them during the EGNN computation).
The addition of a new atom proceeds in three steps. First, a new atom type is selected as follows:
SumPool(Fx(G..)
V1+1 Categorica((Sofmax) (Pa())))))HH))))))))
where V is the one-hot encoded type of the atom, and DA is neural network that decodes the EGNN graph embedding into a set of softmax probabilities. We implement Da as 3-layer MLP. Note that, addition all of the atom species in the training set, allow Vi+1 take on an extra "stop token" value. If this value is generated, the molecule is complete, and generation terminates.
The next step in the generation procedure is to connect the new atom to the existing graph with edges. We do this similar manner GraphAF, and query every atom sequentially determine its new bond type, updating the adjacency list as needed (Shi et al. 2020). More formally, this procedure works as follows:
Initialize E Rexb matrix containing each atom's edge type the new atom V4+1- At initialization, let E contain all unbonded edge types.
for j in 1.4 do:
V Concat(V.,, E, Vi+1) is a (2d + b) matrix of modified atom features. Row j contains the one-hot encoded type of atom j. the one-hot encoded type of atom j's current edge atom + 1, and the one-hot encoded type of atom i i+ 1.
G² = (V1, X,))
- = FF(F((
Ai41, ~
A411, is a sampled bond type between atom j and atom + 1. DE another MLP decoder that acts on the node-specific
embedding of atom j.
E, Ait1.j
Through this procedure, a set of bonds is sampled for the new atom.
In the final step, the new atom is given 3D position. We accomplish this by predicting a discrete distribution of distances from each atom in the graph to the new atom, and a discrete distribution of bond angles between edges that contain the new atom and all adjacent edges. These predictions induce distribution over 3D coordinates. We compute this distribution by calculating how much each point in 3D space would violate the distance and angle constraints generated by the model, and giving higher probability density to points that violate the constraints less. In a secondary step, we approximately sample from this spatial distribution by drawing points from a fine, stochastic 3D grid using the likelihood function given by the distance and angle predictions. More formally, the positions of the atoms are predicted as follows:
V StackRows( V.,, E, V(+1)
G/2, (V2,, A5,
hi,.,, h; = FD(G'))
M,..,,K Fo())))
Pj = Softmax(DD((()) for j j= 1..
Ojk Softmax(DD(())) M)) for j, k s.t. (onymaa(A44... > 0 and argmar (A,s) 0.
22

---

Under review as a conference paper at ICLR 2022
where Do and D3 are MLP decoders as before. Note that reusing the matrix E from the edge prediction step, which has accumulated all of the new edges to atom 1. The probability vectors Pi...., Pt now define discrete distributions new atom, and the vectors Gjs define distributions over bond angles. If we treat these distributions as the product rule compute the likelihood of any point in 3D space:
the distances between each atom in the graph and the
heing independent,
LD(X,-1))
II 1 X,Il) 1=1
Lo(X+41)
II (J.k)EI
(Angleeee( - Xy -xx)))
L(X++1)= = LD(( x Lo(Xi+)),
where I is the set of incident edges to the neighbors of atom i + 1, and "Angle" denotes the angle between two vectors. To sample a point from the likelihood L(X;+1), to every point in a fine, stochastic grid surrounding the atoms that are bonded to atom + 1, and sample from it categorical distribution produce
simply assign likelihood
spatial location.
By repeating this procedure until termination, GEN3D produces a 3D molecule from a single starting atom. Note that, because the generation process is sequential, it is simple to mask out atom edge selections that would violate valence constraints, thereby guaranteeing that generated molecules follow basic chemical rules. It also possible for the model predict a non-terminating atom, but then predict that no edges connect to that atom. In this case, we rerun the edge sampling procedure until at least one edge the new atom generated. If no edge to the new atom is produced after 10 resampling attempts, the new atom is discarded and the generation process is said to have terminated.
23


Reviews:
[{"review_id": "iuH6UE5wUH", "reviewer": "Reviewer_biHr", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. The related work is pretty detailed but neglects some relevant related papers. \n   1) Simm, G. N. C.; Pinsler, R.; Hernández-Lobato, J. M. Reinforcement Learning for Molecular Design Guided by Quantum Mechanics. ICML 2020.\n   1) Gebauer, N. W. A.; Gastegger, M.; Hessmann, S. S. P.; Müller, K.-R.; Schütt, K. T. Inverse Design of 3d Molecular Structures with Conditional Generative Neural Networks. arXiv:2109.04824 [physics, stat] 2021.\n   1) Simm, G. N. C.; Pinsler, R.; Csányi, G.; Hernández-Lobato, J. M. Symmetry-Aware Actor-Critic for 3D Molecular Design. ICLR 2020.\n\n2. It is not clear how working with a graph is beneficial. 1) The authors claim that organic chemistry rules can easily be enforced through a graph representation. However, it's unclear to what extent these relatively simple rules are holding ML models back from proposing more chemically meaningful molecules. 2) Working with a graph restricts the application of this model to organic chemistry, and metals or multi-molecular clusters become unattainable.", "summary_of_the_paper": "The authors propose GEN3D, an ML-based approach for 3D molecular structure generation. GEN3D places atoms in space to build molecules in an atom-by-atom fashion and constructs a corresponding molecular graph that can be employed to enforce organic chemistry rules.", "summary_of_the_review": "GEN3D appears to be a competitive approach for molecular structure generation. The main weakness of this paper is, however, that it lacks methodological novelty. Methods, such as G-Schnet, MolGym, and ENF, are closely related to this work, and the originality of this work is, in my view, insufficient for this venue.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XPlZZw-UFtg", "reviewer": "Reviewer_Qzyh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n* The proposed model is a nice extension of previous autoregressive graph generative models, e.g., GCPN. Inspired by the mechanism of distance geometry and internal coordinates (e.g., pairwise distance and angles), the authors propose to autoregressively predict the next atom/edge type, next atom distances/angles, thus enabling the model to concurrently generate molecular graphs and 3D geometries. The whole methodology is classic and easy to follow.\n* They experimentally validate the effectiveness of the proposed model. The experiment design seems good as it covers several research topics ranging from 2D molecular graph generation, 3D molecule design, to molecular conformation generation.\n* They plan to release the code in the future.\n\n**Weaknesses**\n* The novelty of the work is limited. Compared to the previous autoregressive graph generative models (e.g., GCPN, GraphAF, MolecularRNN[1,2,3]), the key difference is that they incorporate the atomic distance prediction and angle prediction into their model. However, distance geometry modeling and angle prediction are not something new. So the model is more like a combination of techniques in 2D graph generation and 3D structure modeling.\n* Since the paper introduces a **new generative model** for 3D molecule design, I recommend authors introduce the generative model in a more formal (mathematical) way. For example, we can introduce the probabilistic distribution we want to model, the way we factorize the distribution (autoregressively), and how we estimate the parameters. This is a minor suggestion, but I think it will improve the quality of the draft. \n\n**Reference**\n1. Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation\n2. GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation\n3. MolecularRNN: Generating realistic molecular graphs with optimized properties\n\nPost rebuttal: The authors partially addressed my concerns and I would like to update the score to 5", "summary_of_the_paper": "* The paper proposes an autoregressive generative model for 3D molecular graph generation, by predicting the next atom type, next edge type, next atom distances, and next atom angles sequentially based on the existing graph.\n* The authors claim that the proposed model can generate 3D molecules with higher rates of chemical validity, better atom-distance distributions, and is very cheap to train compared with previous models.", "summary_of_the_review": "My main concerns of the paper are its novelty and the way it introduces the generative model.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "E0wfQDNe6JO", "reviewer": "Reviewer_uzvM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n - Problem is relevant and clearly defined\n - Method is intuitive and accessible - generate atom type, generate connections to existing atoms, generate positions, repeat\n - Some results are positive, generally better or similar to other methods. The model's tendency to construct valid molecules without checks is impressive.\n - Writing is clear, organized, and developed, discrepancies are discussed\n\nCons:\n - The result in Table 1 make me question the overall value of the technique.\n   - The method w/o geometry has higher validity and uniqueness than the model w/ geometry.\n   - All of the methods (w/ checks) produce valid molecules and have similar uniqueness and higher novelty. The paper discusses this apparent shortcoming and largely concludes that the comparison is misleading due to intentional limitations (number of atoms) in the proposed method vs the other methods. If this is true, is it possible to reproduce Table 1 where only \"fair\" comparisons are made (maybe by dropping large molecules from the other methods)?\n - The paper discusses how the conditional model could be used to retain features from one molecule when constructing new molecules but doesn't demonstrate this. It would strengthen the paper to see a simple demonstration of this application.", "summary_of_the_paper": "The paper proposes a model to generate molecules using graph and geometric representations. The method is tested against several relevant baselines on standard and custom datasets, with a primary focus on chemical validity.", "summary_of_the_review": "The paper takes a positive step forward in constructing generative models over molecules by including graphical and geometrical information in a single model. However, some of the empirical evaluations muddle the overall story and could be improved/clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7KBgZbXL3zF", "reviewer": "Reviewer_7FsH", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Pros\n    - Highlights the importance of generating 3d structure for molecule generation.\n    - The paper provides new insight of the QM9 dataset.\n\n- Cons\n    - Insufficient description. The paper seems not-self contained. Difficult to understand some of the details.\n        - Since ICLR is not chemistry focused conference, it would be good to provide some details of domain-specific information\n    - The 3D structure generation is unclear.\n        - To generate 3D coordinates, the model places a probability mass over the grids of distances and angles.\n        - The provided process is likely to place a mass over the impossible configuration of distances and angles.\n        - From the grid, the model sample a new location. However, sampling is not back-propagatable in general. It is unclear how the model handles this.\n    - A mismatch between motivation and experiments\n        - None of the experiments is related to the proposed downstream tasks in the introduction. Although generating 3D structure seems plausible, the benefit of it is unclear without direct application of 3D structure\n            - JS divergence on distance is the only metric that is directly connected to the molecule geometry. However, it is not sufficient to show the usefulness of 3D structure.\n        - For example, it would be more interesting to show that the generated graphs have relatively low-energy conformations than those from the other models.\n    - A proper ablation study on model selection and training procedures would improve the strength of the proposed approach.\n\n- Questions\n    - What are the losses used to train the model? Please provide the detailed training procedure.\n    - Is the equivariance of the subnetworks sufficient to impose the equivariance of the entire network?\n        - Is permutation invariant unnecessary? Since the generation is autoregressive, to put the same distribution to the same molecule with different sequences, it needs to be permutation invariant.\n    - If the description of QM9 dataset is correct, how do we understand the reported performance of some previous models which achieve 90%+ uniqueness and novelty such as GraphDF[1] and MoFlow[2]. This seems impossible based on the description provided in this paper.\n    - In the Introduction, a significant portion has been used to criticize the approaches with SMILES. What problem can be arisen from converting SMILES string to 3D structure? and how the proposed solution can avoid this? - e.g., w. OpenBabel\n    - What is the GEOM-QM9 dataset, and how does it differ from QM9?\n        - In the last sentence of Section 3, a random uniform noise is added to each coordinate. Is this applied to all datasets?\n\n1. Luo, Youzhi, Keqiang Yan, and Shuiwang Ji. \"GraphDF: A discrete flow model for molecular graph generation.\" arXiv preprint arXiv:2102.01189 (2021).\n1. Liu, Meng, et al. \"GraphEBM: Molecular graph generation with energy-based models.\" arXiv preprint arXiv:2102.00546 (2021).", "summary_of_the_paper": "This paper proposes a generative model to sample a new molecule with its 3d coordinates. The generative process is divided into three different steps: 1) the atom network generates a new atom 2) the edge network connects the new atom to the existing atoms 3) the distance and angle network generates the pairwise distances and angles from existing atoms to the new atoms. EGNN is used for each of these sub-networks to impose the equivariance to translations and rotations. Experiments with QM9 datasets show the geometric correctness of the proposed generative process.\n", "summary_of_the_review": "Although the proposed method seems an important research direction for molecule generation, the experiments did not show the usefulness of such approaches to compare with the previously studied methods. Some of the descriptions are insufficient to understand the model details as well.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to generate 3D molecules using a step by step approach. The reviewers raised major concerns on the experiments, novelty, writing and technical details. The authors also were not aware of many of the important references, part (but not all) of which have been included during discussions. It is clear that this work is not ready to be accepted by ICLR.
Decision: Reject