GEOMOL: Torsional Geometric Generation of Molecular 3D Conformer Ensembles
Year: 2022

Abstract:
Prediction of a molecule’s 3D conformer ensemble from the molecular graph holds a key role in areas of cheminformatics and drug discovery. Existing generative models have several drawbacks including lack of modeling important molecular geometry elements (e.g., torsion angles), separate optimization stages prone to error accumulation, and the need for structure fine-tuning based on approximate classical force-fields or computationally expensive methods. We propose GEOMOL — an end-to-end, non-autoregressive, and SE(3)-invariant machine learning approach to generate distributions of low-energy molecular 3D conformers. Leveraging the power of message passing neural networks (MPNNs) to capture local and global graph information, we predict local atomic 3D structures and torsion angles, avoiding unnecessary over-parameterization of the geometric degrees of freedom (e.g., one angle per non-terminal bond). Such local predictions suffice both for both the training loss computation and for the full deterministic conformer assembly (at test time). We devise a non-adversarial optimal transport based loss function to promote diverse conformer generation. GEOMOL predominantly outperforms popular open-source, commercial, or state-of-the-art machine learning (ML) models, while achieving significant speed-ups. We expect such differentiable 3D structure generators to significantly impact molecular modeling and related applications. 4

text: 1 Overview
Problem & importance. We tackle the problem of molecular conformer generation (MCG), i.e., predicting the ensemble of low-energy 3D conformations of a small molecule solely based on the molecular graph (fig. 1). A single conformation is represented by the list of 3D coordinates for each atom in the respective molecule. In this work, we assume that the low-energy states are implicitly defined by the given dataset, i.e., our training data consist of molecular graphs and corresponding sets of energetically favorable 3D conformations. Low-energy structures are
the most stable configurations and, thus, expected to be observed most often experimentally.
⇤Equal first authorship and contribution. Correspondence to {oct,lagnajit}@mit.edu . †Department of Chemical Engineering, MIT, Cambridge, MA 02139 ‡Computer Science and Artificial Intelligence Laboratory, MIT, Cambridge, MA 02139 4Code is available at https://github.com/PattanaikL/GeoMol.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Dealing with molecules in their natural 3D structure is of great importance in areas such as cheminformatics or computational drug discovery because conformations determine biological, chemical, and physical properties [Guimaraes et al., 2012, Schütt et al., 2018, Klicpera et al., 2019, Axelrod and Gomez-Bombarelli, 2020b, Schütt et al., 2021, Liu et al., 2021] such as charge distribution, potential energy, docking poses [McGann, 2011], shape similarity [Kumar and Zhang, 2018], pharmacophore searching [Schwab, 2010], or descriptors for 3/4D QSAR [Verma et al., 2010]. For instance, in drug design it is crucial to understand how a molecule binds to a specific target protein; this process heavily depends on the 3D structures of the two components, both in terms of geometric (shape matching) and chemical (hydrophobic/hydrophilic) interactions [Gainza et al., 2020, Sverrisson et al., 2020].
Motivation & challenges of existing methods. The main challenge in MCG comes from the enormous size of the 3D structure space consisting of bond lengths, bond angles, and torsion angles. It is known that the molecular graph imposes specific constraints on possible 3D conformations, e.g., bond length ranges depend on the respective bond types, while tetrahedral centers dictate local spatial arrangement. However, the space of possible conformations grows exponentially with the graph size and number of rotatable bonds, thus hindering exhaustive brute force exploration even for relatively small molecules. Additionally, the number of plausibly-stable low-energy states is unknown a priori and can vary between one and several thousand conformations for a single molecule [Chan et al., 2021]. Nevertheless, various facets of the curse of dimensionality have been favorably tackled by ML models in different contexts, and our goal is to build on the recent ML efforts for MCG [Mansimov et al., 2019, Simm and Hernandez-Lobato, 2020, Lemm et al., 2021, Xu et al., 2021].
Molecular conformations can be determined experimentally, but existing techniques are very expensive. As a consequence, predictive computational models have been developed over the past few decades, traditionally being categorized as either stochastic or systematic (rule-based) methods [Hawkins, 2017]. Stochastic approaches have traditionally been based on molecular dynamics (MD) or Markov chain Monte Carlo (MCMC) techniques, potentially combined with genetic algorithms (GAs). They can do extensive explorations of the energy landscape and accurately sample equilibrium structures, but quickly become prohibitively slow for larger molecules [Shim and MacKerell Jr, 2011, Ballard et al., 2015, De Vivo et al., 2016, Hawkins, 2017], e.g., they require several CPU minutes for a single drug-like molecule. Moreover, stochastic methods have difficulties sampling diverse and representative conformers, prioritizing quantity over quality. On the other hand, rule-based systematic methods achieve state-of-the-art in commercial software [Friedrich et al., 2017] with OMEGA [Hawkins et al., 2010, Hawkins and Nicholls, 2012] being a popular example. They usually process a single drug-like molecule under a second. They address the aforementioned challenges of stochastic methods by relying on carefully curated torsion templates (torsion rules), rule-based generators, and knowledge bases of rigid 3D fragments, which are assembled together and combined with subsequent stability score ranking. However, torsion angles are mostly varied independently (based on their fragments), without explicitly capturing their global interactions, which results in difficulties for larger and more flexible molecules. Furthermore, the curated fragments and rules are inadequate for more challenging inputs (e.g., transition states or open-shell molecules).
Both types of methods can be combined with Distance Geometry (DG) techniques to generate the initial 3D conformation. First, the 3D atom distance matrix is generated based on a set of distance constraints or from a specialized model. Subsequently, the corresponding 3D atom coordinates are learned to approximately match these predicted distances [Havel et al., 1983b,a, Crippen et al., 1988, Havel, 1998, Lagorce et al., 2009, Riniker and Landrum, 2015]. Indeed, modern stochastic algorithms are entirely based on DG methods [Riniker and Landrum, 2015]. The inductive bias of rotational and translational invariance is guaranteed for DG, thus being appealing for ML models [Simm and Hernandez-Lobato, 2020, Xu et al., 2021, Pattanaik et al., 2020b]. However, several drawbacks weaken this important direction: i) the distance matrix is overparameterized compared to the actual number of degrees of freedom, ii) it is difficult to enforce 3D Euclidean distance constraints as well as geometric graph constraints (e.g., on torsion angles or rings [Riniker and Landrum, 2015]); iii) important aspects of molecular geometry are not explicitly modeled, e.g., torsion angles of rotatable bonds or tetrahedral centers; iv) expensive force-field energy fine-tuning of the generated conformers is vital for a reasonable quality [Xu et al., 2021, Simm and Hernandez-Lobato, 2020]; iv) the resulting multi-stage pipeline is prone to error accumulation as opposed to an end-to-end model.
Previous methods often rely on a force field (FF) energy function minimization to fine-tune the conformers. These are hand-designed energy models which use parameters estimated from experi-
ment and/or computed from quantum mechanics (e.g., Universal Force Field [Rappé et al., 1992], Merck Molecular Force Field [Halgren, 1996]). However, FFs are crude approximations of the true molecular potential energy surface [Kanal et al., 2018], limited in the interactions they can capture in biomolecules due to their strong assumptions [Barman et al., 2015]. In addition, FF energy optimization is relatively slow and increases error accumulation in a multi-pipeline method.
Relation to protein folding. There has been impressive recent progress on modeling protein folding dynamics [Ingraham et al., 2018, AlQuraishi, 2019, Noé et al., 2019, Senior et al., 2020], where crystallized 3D structures are predicted solely from the amino-acid sequence using ML methods. However, molecules pose unique challenges, being highly branched graphs containing cycles, different types of bonds, and chirality information. This makes protein folding approaches not readily transferable to general molecular data.
Our key contributions & model in a nutshell. In this work, we investigate the question:
Can we design a fast and generalizable deep learning model to predict high-quality, representative, and diverse 3D conformational ensembles from input molecular graphs?
To tackle this question, we propose GEOMOL (shown in fig. 2), exhibiting the following merits:
• It is end-to-end trainable, non-autoregressive, and does not rely on DG techniques (thus avoiding aforementioned drawbacks). More precisely, it outputs a minimal set of geometric quantities (i.e., angles and distances) sufficient for full deterministic reconstruction of the 3D conformer.
• It models conformers in an SE(3)-invariant (translation/rotation) manner by design. This desirable inductive bias was previously either achieved using multi-step DG methods [Simm and HernandezLobato, 2020] or not captured at all [Mansimov et al., 2019].
• It explicitly models and predicts essential molecular geometry elements: torsion angles and local 3D structures (bond distances and bond angles adjacent to each atom). Together with the input molecular graph, these are used for k-hop distance computation at train time and full deterministic conformation assembly at test time. Crucially, we do not over-parameterize these predictions, i.e., a single torsion angle is computed per each non-terminal bond, irrespective of the number and permutation of the neighboring atoms at each end-point of the respective bond.
• The above geometric elements (torsion angles, local structures) are SE(3)-invariant (by definition or usage) and we jointly predict them using MPNNs [Gilmer et al., 2017] and self-attention networks. Thus, unlike [Mansimov et al., 2019], we are not affected by MPNNs’ pitfalls that obstruct direct predictions of 3D atom coordinates from node embeddings, e.g., symmetric or locally isomorphic nodes would always have identical MPNN embeddings [Xu et al., 2019, Garg et al., 2020] and, as a consequence, would be inappropriately assigned identical 3D coordinates.
• To promote diverse conformer ensembles with good coverage, we devise a tailored generative loss that does not use slow or difficult-to-optimize adversarial training techniques. Using optimal transport, GEOMOL finds the best matching between generated and ground truth conformers based on their pairwise log-likelihood loss, requiring only minimization.
• It explicitly and deterministically distinguishes reflected structures (enantiomers) by solving tetrahedral stereocenters using oriented volumes and local chiral descriptors, bypassing the need for iterative optimization usually done in DG approaches.
• Empirically, we conduct experiments on two benchmarks: GEOM-QM9 (smaller molecules relevant to gas-phase chemistry) and GEOM-DRUGS (drug-like molecules) [Axelrod and GomezBombarelli, 2020a]. Our method often outperforms previous ML and two popular open-source or commercial methods in different metrics. Moreover, we show competitive quality even without the frequently-used computationally-demanding fine-tuning FF strategies.
• GEOMOL processes drug-like molecules in seconds or less, being orders of magnitude faster than popular baselines (e.g., ETKDG/RDKit[Riniker and Landrum, 2015]), without sacrificing quality.

heading: 2 Method
text: Problem setup & notations. Our input is any molecular graph G = (V,E) with node and edge features, xv 2 Rf , 8v 2 V and eu,v 2 Rf 0 , 8(u, v) 2 E representing atom types, formal charges, bond types, etc. For each molecular graph G, we have a variable-size set of low-energy ground truth 3D conformers {C⇤l }l that we predict with a model {Ck}k def = ⇣(G). A conformer is a map C : V ! R3 from graph nodes to 3D coordinates, but a simplified notation is cv 2 R3 for v 2 V . We use additional notations: d(X,Y ) = kcX cY k is the 3D distance between X and Y; \XY Z is the counter-clockwise (CCW) angle \cXcY cZ ; \(XY Z,XY T ) is the CCW dihedral angle of the 2D planes cXcY cZ and cXcY cT (formula is in appendix C). We use the corresponding c ⇤ v, d ⇤ (X,Y ),\⇤XY Z,\⇤(XY Z,XY T ) when manipulating a ground truth conformer.
Any conformer is defined up to a SE(3) transformation, i.e., any translation or rotation applied to the set {cv}v2V . A classic conformer distance function that satisfies this constraint is root-mean-square deviation of atomic positions (RMSD), computed by the Kabsch alignment algorithm [Kabsch, 1976].

heading: 2.1 GEOMOL high-level overview
text: Our approach, shown in fig. 2, comprises three steps. First, we predict the local 3D structure of each non-terminal atom, which we deem local structure (LS), by combining self-attention layers and MPNNs with deterministic corrections for tetrahedral centers. Bond distances and bond angles are computed from the predicted LS. Next, we assemble all neighboring pairs of LSs by predicting the torsion angles and aligning them. Importantly, since LSs are fixed, it suffices to only predict a single value for the dihedral angle of each bond. Towards this goal, we develop a canonical representation of torsion angles via a local coordinate system defined SE(3)-equivariantly w.r.t. the full structure, which allows us to predict exactly the number of degrees of freedom. Finally, at test time, we assemble all predicted pairs of neighboring LSs to construct the full conformer, applying deterministic ring corrections. In order to generate diverse conformers, we append random Gaussian noise vectors to each initial node feature vector and use an optimal transport-based loss function for training.

heading: 2.2 Message passing neural networks (MPNNs)
text: Given an input graph G, an MPNN [Gilmer et al., 2017, Battaglia et al., 2018, Yang et al., 2019] computes node embeddings hv 2 Rd, 8v 2 V using T layers of iterative message passing:
h (t+1) u =
h (t) u ,
X
v2Nu
(h (t) v ,h (t) u , eu,v) ! , where h(0)v def = concat[xv, zv], zv ⇠ N (0, sId)
(1) for each t 2 [0 .. T 1], where Nu = {v 2 V |(u, v) 2 E}, while and are generic functions, e.g., implemented using multilayer perceptrons (MLP) or attention [Veličković et al., 2017]. Final node embeddings are obtained by the embedding of the last layer: hv def = h (T ) v , 8v 2 V . Finally, we
also compute a molecular embedding: hmol def = MLP ( P v2V hv). We leave comparison with other MPNN variants for future work, e.g., Kipf and Welling [2017], Veličković et al. [2017], Hamilton et al. [2017], Xu et al. [2019].

heading: 2.3 Local structure (1-hop) prediction model
text: Following notations in fig. 3, for each non-terminal graph vertex X 2 V having n graph neighbors NX = {Ti}i2[1..n], we predict its local 3D structure (LS), i.e., the relative 3D positions of all Ti, when X is centered in the origin. The generic model is a function f(hT1 , . . . ,hTn ;hX) = (p1, . . . ,pn) 2 R3⇥n that, additionally, should satisfy permutation equivariance w.r.t. Ti’s, namely, the 3D position of each neighbor Ti should not change regardless of the ordering of the X’s neighbors:
f(hT⇡(1) , . . . ,hT⇡(n) ;hX) = (p⇡(1), . . . ,p⇡(n)), 8⇡ 2 Sn (2)
Figure 3: For each non-terminal atom X, we predict the relative 3D position of each of its graph neighbors, {Ti}i2[1..n], in a permutation equivariant manner.
Our choice is the encoder part of a transformer [Vaswani et al., 2017], without any positional encoding, thus satisfying permutation equivariance. This model takes as input the set {concat[hTi ,hX ]; i 2 [1..n]} in any order and synchronously updates the n embeddings based on several transformer layers. The final layer projects the embeddings to 3 dimensions, resulting in a list (p1, . . . ,pn) 2 R3⇥n having the exact same node order as the input list.
Enforcing local consistency. We desire the LS model f() to be distance-consistent, i.e., any bond distance d(X,Y ) is the same, no matter if it is computed from the LS of node X or of node Y. To achieve this, we use the above transformer
just to compute bond directions (which will be aligned using a separate approach described in section 2.4), while we obtain the bond distances with a separate symmetric model. Concretely, let the above transformer f() predict (p1, . . . ,pn) 2 R3⇥n, while the final local 3D coordinates are p0i def = pi kpikdGNN (hX ,hTi), 8i, where each bond distance is predicted with a symmetric model dGNN (hX ,hY ) def = softplus( (hX ,hY )+ (hY ,hX)), 8(X,Y ) 2 E, with the same shared (e.g., an MLP). For notation simplicity, we will just use pi instead of p0i.
Regarding SE(3) invariance. The above model is not SE(3)-invariant per se, but it is used as such. Namely, on one hand we compute SE(3)-invariant quantities: 1- hop distances d(Ti, X), 2-hop distances d(Ti, Tj), and bending angles \TiXTj . These will be compared to their ground-truth counterparts in the final loss, see section 2.5.
On the other hand, the LS of adjacent graph nodes are assembled together for computing torsion angles or for building the full conformer at test time. This process is explicitly defined to be SE(3)-invariant as described in section 2.4.
Tetrahedral chiral corrections. When embedding the local neighborhood of a node in 3D space, one has to carefully account for tetrahedral stereocenters (fig. 4). Tetrahedral chirality is a common form of stereochemistry which restricts the 3D location of neighboring substituents of a central atom with four distinct neighbors; molecules which differ by a single tetrahedral stereocenter, i.e., enantiomers, are mirror images of each other. Chirality heavily impacts some properties of small molecules–e.g., bioactivity. Existing MPNNs using only the molecular graph cannot distinguish chiral centers (fig. 4), but solutions exist [Pattanaik et al., 2020a]. Mathematically,
enantiomers can be differentiated based on the oriented volume around the tetrahedral center. That is, given the ordered set of neighbor 3D coordinates around the center, namely p1,p2,p3,p4 2 R3, the
sign of the volume of the tetrahedron formed by the neighbors is
OV (p1,p2,p3,p4) def = sign
0
B@
1 1 1 1
x1 x2 x3 x4
y1 y2 y3 y4
z1 z2 z3 z4
1
CA
Enantiomeric structures always have opposite signs for the oriented volume [Crippen et al., 1988]. Since we generate local 3D structures directly, we can also use local 3D chiral descriptors to ensure the correct generation of tetrahedral stereocenters. RDKit internally keeps track of these local chiral labels, denoted by CW/CCW labels (detailed in e.g., Pattanaik et al. [2020a]). Importantly, each local chiral label corresponds to a certain oriented volume (CW = +1 and CCW = -1). Thus, when generating an LS for a tetrahedral center, we calculate the oriented volume and check against the internal RDKit label. If it results in the incorrect oriented volume (i.e., the incorrect stereocenter was generated), we simply reflect the structure by flipping against the z-axis. This ensures that all tetrahedral stereocenters centers are generated exactly, and no iterative optimization is necessary as with traditional DG-based generators.

heading: 2.4 Torsion angle representation and local structure (LS) assembly
text: Once the LS of each atom/vertex is predicted, we assemble them in pairs corresponding to each non-terminal bond in the molecular graph. We describe this process for a bond connecting atoms X and Y, each having additional graph neighbors {Ti}i2[1..n] and, resp., {Zj}j2[1..m]. See fig. 5.
Torsion angle over-parameterization. We first note that, for any assembled bond XY (fig. 5, right), and 8i, k 2 [1..n], 8j, l 2 [1..m], the dihedral angles \(XY Ti, XY Tk) and \(XY Zl, XY Zj) are fully determined by the LS of nodes X and Y, respectively, so they do not depend on the torsion angle of bond XY. Next, observe that there is exactly one torsion angle for any bond XY, given unique indexing of the neighbors. This happens because of the following constraint:
\(XY Ti, XY Zj) = [\(XY Tk, XY Zl)+\(XY Ti, XY Tk)+\(XY Zl, XY Zj)](mod 2⇡) (3)
Thus, in order to avoid unnecessary over-parameterization, we predict a single torsion angle ↵ per each bond XY connecting non-terminal atoms.
Torsion angle formulation. However, it is still unclear at this point how to define this unique angle in a canonical way that is: i) permutation invariant w.r.t. the nodes in the set {Ti}i2[1..n] and, respectively, in the set {Zj}j2[1..m], ii) SE(3)-invariant w.r.t. the full 3D conformer, and iii) agrees with eq. (3).
Let ij def = \(XY Ti, XY Zj) and sij def=
 cos( ij)
sin( ij)
. Let
cij 2 R be real coefficients such that s def = P
i,j cijsij 2 R2 is not the null vector. Then, we define the torsion angle as5: ↵ def = atan2(
s ksk ). It is easy to see that this formulation satisfies
both invariances claimed above. We further state (and prove in appendix A) that our proposed formulation gives a torsion angle uniquely determined by all local angles \(Y XTi, Y XTk), \(Y XZj , Y XZl) and by the true underlying torsion angle:
Proposition 1. Given 3D coordinates of nodes X,Y, Ti, Zj and fixed weights cij 2 R such thatP
i,j cijsij 2 R2 is not the null vector, then ↵ def = atan2( s ksk ) is unique, i.e., if we change the torsion angle of bond XY, then ↵ will change. Formally, if we rotate the set of bonds {XTi}i jointly around the line XY with the same angle , then ↵ will be exactly shifted with .
5We define atan2 slightly different than standard: atan2(r cos(↵), r sin(↵)) def= ↵, 8↵ 2 [0, 2⇡), r 2 R⇤+.
How to set cij? Breaking symmetries. A simple solution is to choose cij = 1, 8i, j. However, in some important cases, local symmetries may result in s = 0. For example, this happens if, for some j, we have ij = 2i⇡n + ct., 8i 2 [1..n]. One solution is to use different cij to differentiate between the different subgraphs rooted at different Ti (and similarly for Zj). This is reminiscent of traditional group priorities used for distinguishing E/Z isomers. We devise a flexible solution to distinguish these subgraphs: a differentiable real valued function computed from the MPNN node embeddings as cij = MLP (hTi + hZj ) 2 R, with MLP being a neural network shared across all bonds and molecules. Note that we constrain cij = cji, thus guaranteeing that the same ↵ is obtained if we swap X and Y (and their neighbors, respectively).
Final LS assembly for a single bond. We now describe the assembly process depicted in fig. 5. We first predict the LS of node X as in section 2.3, obtaining several 3D coordinates: pX = 0,pY ,pTi 2 R3, 8i 2 [1..n], as well as the LS of node Y: qY = 0,qX ,qZj 2 R3, 8j 2 [1..m]. By design, we have that kqXk = kpY k. These two sets are currently not aligned. To achieve this, we first rotate the LS of X such that pY becomes [kpY k 0 0]
>, while pX remains 0. Next, we rotate and translate the LS of Y such that qY becomes pY and qX becomes pX = 0. These two rotations have one degree of freedom each, which we set randomly. Exact formulas are in appendix B. Thus, the bond XY is now matched, but the torsional rotation is still arbitrary/random. The remaining step is to rotate the LS of X with an angle such that all dihedrals \(XY Ti, XY Zj) match their true counterparts. This
is done by applying to all vectors pTi the same rotation of type: H :=
" 1 0 0
0 cos( ) sin( ) 0 sin( ) cos( )
# .
How to compute ? The current dihedrals curij def = \cur(XY Ti, XY Zj) depend on the random torsional rotations from the initial assembly step of LS of X and of Y. After applying the H rotation, we obtain the new dihedral angles: [ curij ] mod 2⇡ that should match the ground truth dihedral angles ⇤ij def = \⇤(XY Ti, XY Zj). This is equivalently written as s⇤ij = Acurij s ,
where s def =
 cos( )
sin( )
and Acurij
def =
 cos(
cur ij ) sin( cur ij )
sin( cur ij ) cos( curij )
. Let s⇤ def= P i,j cijs ⇤ ij and Acur def =
P i,j cijA cur ij . The necessary condition for becomes s = (Acur) > s ⇤, which is also sufficient due to proposition 1. This implies it is enough to predict only the normalized s ⇤
ks⇤k and, in practice, we do
that by predicting s↵ def =
 cos(↵)
sin(↵)
using a function commutative in X and Y (i.e., swapping X and Y
does not change ↵):
↵ = [ (hX ,hY ,hmol) + (hY ,hX ,hmol)]mod 2⇡ (4)
where is a neural network (e.g., MLP). Finally, s =  cos( )
sin( )
= 1 k(Acur)>s↵k (A cur ) > s↵.

heading: 2.5 An optimal transport (OT) loss function for diverse conformer generation
text: Loss per single conformer. Assume first that we predict a single conformer C. Based on all LS and torsion angle predictions, we deterministically compute all 1/2/3-hop distances and bond/torsion angles. If the corresponding ground truth conformer C⇤ is known, we feed those quantities into a negative log-likelihood loss, denote by L(C, C⇤) and detailed in appendix D. Similar to Senior et al. [2020], we fit distances using normal distributions and angles using von Mises distributions. This is a much faster approach compared to habitual RMSD losses that compare full conformers.
Node symmetries. Our current formulation has difficulties distinguishing pairs of symmetric graph nodes that are less than 3 hops away, e.g., hydrogen groups. We address this using a tailored matching loss detailed in appendix D and exemplified in fig. 6.
Total OT loss per ensemble of conformers. In practice, our model generates a set of conformers {Ck}k2[1..K] that needs to match a variable sized set of low-energy ground truth conformers, {C⇤l }l2[1..L]. However, we do not know a priori the number L of true conformers or the matching between generated and true conformers. We also wish to avoid expensive and problematic adversarial training. Our solution is an OT-based, minimization-only, loss function:
Lensemble def= EMDL(·,·)({Ck}k, {C⇤l }l) = min T2QK,L
X
k,l
TklL(Ck, C⇤l )
where EMD is the Earth Mover Distance, T is the transport plan satisfying QK,L def = {T 2 RK⇥L+ : T1L = 1 K1K ,T T 1K = 1 L1L}. The minimization w.r.t. T is computed quickly using Earth Mover Distance and the POT library [Flamary and Courty, 2017].

heading: 2.6 Full conformer assembly at test time
text: Knowing all true LSs and torsion angles is, in theory, enough for a deterministic unique SE(3)invariant reconstruction of the full conformer. However, in practice, these predictions might have small errors that accumulate, e.g., in rings. To mitigate this issue, we deterministically build the full conformer (only at test time) by first predicting a smoothed structure of (fused) rings separately, and then assembling the full conformer following any graph traversal order (any order gives the same conformer, so this procedure does not break the non-autoregressive behavior). We detail this step in appendix E.

heading: 3 Experiments
text: We empirically evaluate GEOMOL on the task of low-energy conformer ensemble generation for small and drug-like molecules. We largely follow the evaluation protocols of recent methods [Simm and Hernandez-Lobato, 2020, Xu et al., 2021], but also introduce new useful metrics.
Datasets & splits. We use two popular datasets: GEOM-QM9 [Ramakrishnan et al., 2014] and GEOM-DRUGS [Axelrod and Gomez-Bombarelli, 2020a]. Statistics and other details are in fig. 10 and in Mansimov et al. [2019]. Datasets are preprocessed as described in appendix G. We split them randomly based on molecules into train/validation/test (80%/10%/10%). At the end, for each dataset, we sample 1000 random test molecules as the final test set. Thus, the splits contain 106586/13323/1000 and 243473/30433/1000 molecules for GEOM-QM9 and GEOM-DRUGS, resp.
Baselines. We compare to established or recent baselines (discussed in section 1). ETKDG/RDKit [Riniker and Landrum, 2015] is likely the most popular open-source software, a stochastic DG-based method developed in the RDKit package. OMEGA [Hawkins et al., 2010, Hawkins and Nicholls, 2012, Friedrich et al., 2017], a rule-based method, is one of the most established commercial software, with more than a decade of continuous development. OMEGA and ETKDG are some of the fastest and best scaling existing approaches. Finally, we compare with the recent ML models of highest reported quality: GraphDG [Simm and Hernandez-Lobato, 2020] and CGCF [Xu et al., 2021].
Evaluation metrics. We follow prior work [Simm and Hernandez-Lobato, 2020, Xu et al., 2021] and use root-mean-square deviation of atomic positions (RMSD) to compare any two conformers. This is defined as the normalized Frobenius norm of the two corresponding matrices of 3D coordinates after being SE(3)-aligned a priori (using the Kabsch alignment algorithm [Kabsch, 1976]). Next, we introduce four types of metrics to compare two conformer ensembles, generated by a method, {Ck}k2[1..K], and ground truth, {C⇤l }l2[1..L]. These metrics follow the established classification metrics of Precision and Recall and are defined for a given threshold > 0 as:
COV - R (Recall) def= 1 L |{l 2 [1..L] : 9k 2 [1..K], RMSD(Ck, C⇤l ) < }|
AMR - R (Recall) def= 1
L
X
l2[1..L]
min k2[1..K]
RMSD(Ck, C⇤l ) (5)
where AMR is "Average Minimum RMSD", COV is "Coverage", and COV - P (Precision) and AMR - P (Precision) are defined as in eq. (5), but with the generated and ground truth conformer sets swapped. The recall metrics measure how many of the ground truth conformers are correctly predicted, while the precision metrics indicate how many generated structures are of high quality. Specifically, in terms of recall, COV measures the percentage of correct generated conformers from the ground truth set (where a correct conformer is defined as one within an RMSD threshold of the true conformer), while AMR measures the average RMSD of each generated conformer with its closest groun truth match. Depending on the application, either of the metrics might be of greater interest. We follow Xu et al. [2021] and set = 0.5Å for GEOM-QM9 and = 1.25Å for GEOM-DRUGS.
Training and test details. For each input molecule having K ground truth conformers, we generate exactly 2K conformers using any of the considered methods. For GEOMOL, this is done by sampling
different random noise vectors that are appended to node and edge features before the MPNN (eq. (1)). At train time, our model uses a standard deviation (std) s (see eq. (1)) of 5 for both GEOM-QM9 and GEOM-DRUGS. At test time, GEOMOL can use the same or different s values, depending on the downstream application, i.e., higher s results in more diverse conformers, while lower s gives more quality (better precision). For OMEGA, it is not possible to specify a desired number of conformers. So, we tune the RMSD threshold (which decides how many conformers to keep) such that the total generated conformers by OMEGA are approximately 2K. For GEOM-QM9, this corresponds to no RMSD cutoff (i.e., OMEGA generates all possible conformers), and for GEOM-DRUGS, this corresponds to a cutoff of 0.7Å (meaning no two generated conformers will have a distance smaller than this cutoff). We discuss hyper-parameters and additional training details in appendix H.
Results & discussion. Results are shown in table 1 and table 2, and confidence intervals are in appendix I. As noted above, GEOMOL can be run with different noise std at test time, depending on which metric the user is interested in. Even though OMEGA is an established commercial software with more than a decade of continuous development, our model remarkably frequently outperforms it. Note that OMEGA fails to generate any conformers for 7% of the QM9 test set (many of which include fused rings). Moreover, we also outperform the popular RDKit/ETKDG open-source model (except for AMR-P on QM9) and very recent ML models such as GraphDG and CGCF, sometimes by a large margin. For a qualitative insight, we show generated examples in fig. 7 and appendix N.
Additionally, we show in fig. 7 how COV Recall results are affected by the increasing number of rotatable bonds in the test molecule. As expected, having more rotatable bonds makes the problem harder, and this affects all baselines, but GEOMOL maintains a reasonable coverage even for more difficult molecules. Moreover, in appendix K and table 8 we show energy calculations of the generated conformers to support their plausibility. Additionally, results with energy-based relaxations are given in appendix J.
Running time. Fig. 7 shows conformer generation test running times. Our model is the fastest method from the considered baselines, being much faster than CGCF or ETKDG/RDKit. Moreover, GEOMOL scales favorably for molecules with increasing number of rotatable bonds.

heading: 4 Conclusion
text: We proposed GEOMOL, an end-to-end generative approach for molecular 3D conformation ensembles that explicitly models various molecular geometric aspects such as torsion angles or chirality. We expect that such differentiable structure generators will significantly impact small molecule conformer generation along with many related applications (e.g., protein-ligand binding), thus speeding up areas such as drug discovery. GEOMOL’s full source code will be made publicly available.
Limitations & future work. A few current limitations are highlighted and left for future extensions (see also discussion in appendix O). First, our model does not currently support disconnected molecular graphs, e.g., ionic salts, but it can be applied to each connected component, followed by a 3D alignment. Next, our approach would benefit from explicit modeling of long distance interactions, especially for macrocycles or large molecules. This remains to be addressed in an efficient manner. Third, explicitly using ground truth energy values could further improve GEOMOL. Last, we look forward to fine-tune GEOMOL on applications such as generating molecular docking poses or descriptors for 4D QSAR.
Acknowledgments and Disclosure of Funding
OEG thanks Bracha Laufer, Tian Xie, Xiang Fu, Peter Mikhael, and the rest of the RB and TJ group members for their helpful comments and suggestions. LP thanks Camille Bilodeau and the rest of the WHG, KFJ, and CWC research groups for their useful discussions. We also thank Pat Walters, Simon Axelrod, and Rafael Gomez-Bombarelli for their insightful feedback as well as Minkai Xu and Shitong Luo for helping run the CGCF model. Both OEG and LP are funded by the Machine Learning for Pharmaceutical Discovery and Synthesis consortium.

Review summaries:
The authors propose the generative model that produces 3d conformations from a molecular graph based on predictions of local structures. First, for each anchor atom, the model predicts relative 1-hop positions of neighboring atoms, then predicts torsion angles. Finally, assemble the molecule in the deterministic way given predicted components. The authors comprehensively evaluate the model on the recently proposed dataset GEOM and compare it with state-of-the-art neural-based conformer generators and special software.  The paper's main contributions are: the authors propose a principal novel scheme for conformation generation that avoids expensive molecular geometry optimization used in current state-of-the-art neural conformer generators, still being translation, rotation, and reflection equivariant. The proposed model outperforms other baselines in both speeds of conformation generation and diversity of generated objects. 

This paper introduces GeoMol, an end-to-end learning method for molecular conformation generation.  It is based on the construction and alignment of elements of local atomic structure, is efficient in its parameterization of geometric degrees of freedom, and is SE(3) invariant.  It achieves competitive performance on the GEOM-DRUGS and GEOM-QM9 tasks.