L EARNING N EURAL G ENERATIVE D YNAMICS FOR M OLECULAR C ONFORMATION G ENERATION
A BSTRACT
We study how to generate molecule conformations ( i.e. , 3D structures) from a molecular graph.
Traditional methods, such as molecular dynamics, sample conformations via computationally expensive simulations.
Recently, machine learning methods have shown great potential by training on a large collection of conformation data.
Challenges arise from the limited model capacity for capturing complex distributions of conformations and the difﬁculty in modeling long-range dependencies between atoms.
Inspired by the recent progress in deep generative models, in this paper, we propose a novel probabilistic framework to generate valid and diverse conformations given a molecular graph.
We propose a method combining the advantages of both ﬂow-based and energy-based models, enjoying: (1) a high model capacity to estimate the multimodal conformation distribution; (2) explicitly capturing the complex long-range dependencies between atoms in the observation space.
Extensive experiments demonstrate the superior performance of the proposed method on several benchmarks, including conformation generation and distance modeling tasks, with a signiﬁcant improvement over existing generative models for molecular conformation sampling.
I NTRODUCTION
Recently, we have witnessed the success of graph-based representations for molecular modeling in a variety of tasks such as property prediction (Gilmer et al., 2017) and molecule generation (You et al., 2018; Shi et al., 2020).
However, a more natural and intrinsic representation of a molecule is its 3D structure, commonly known as the molecular geometry or conformation , which represents each atom by its 3D coordinate.
The conformation of a molecule determines its biological and physical properties such as charge distribution, steric constraints, as well as interactions with other molecules.
Furthermore, large molecules tend to comprise a number of rotatable bonds, which may induce ﬂexible conformation changes and a large number of feasible conformations in nature.
Generating valid and stable conformations of a given molecule remains very challenging.
Experimentally, such structures are determined by expensive and time-consuming crystallography.
Computational approaches based on Markov chain Monte Carlo (MCMC) or molecular dynamics (MD) (De Vivo et al., 2016) are computationally expensive, especially for large molecules (Ballard et al., 2015).
Machine learning methods have recently shown great potential for molecular conformation generation by training on a large collection of data to model the probability distribution of potential conformations R based on the molecular graph G , i.e. , p ( R |G ) .
For example, Mansimov et al.
(2019) proposed a Conditional Variational Graph Autoencoders (CVGAE) for molecular conformation generation.
A graph neural network (Gilmer et al., 2017) is ﬁrst applied to the molecular graph to get the atom representations, based on which 3D coordinates are further generated.
One limitation of such an approach is that by directly generating the 3D coordinates of atoms it fails to model the rotational and translational invariance of molecular conformations.
To address this issue, instead of generating the 3D coordinates directly, Simm & Hern´andez-Lobato (2020) recently proposed to ﬁrst model the molecule’s distance geometry ( i.e. , the distances between atoms)—which are rotationally and translationally invariant—and then generate the molecular conformation based on the distance geometry through a post-processing algorithm (Liberti et al., 2014).
Similar to Mansimov et al.
(2019), a few layers of graph neural networks are applied to the molecular graph to learn the representations of different edges , which are further used to generate the distances of different edges independently.
This approach is capable of more often generating valid molecular conformations.
Although these new approaches have made tremendous progress, the problem remains very challenging and far from solved.
First, each molecule may have multiple stable conformations around a number of states which are thermodynamically stable.
In other words, the distribution p ( R |G ) is very complex and multi-modal.
Models with high capacity are required to model such complex distributions.
Second, existing approaches usually apply a few layers of graph neural networks to learn the representations of nodes (or edges) and then generate the 3D coordinates (or distances) based on their representations independently.
Such approaches are necessarily limited to capturing a single mode of p ( R |G )
(since the coordinates or distances are sampled independently) and are incapable of modeling multimodal joint distributions and the form of the graph neural net computation makes it difﬁcult to capture long-range dependencies between atoms, especially in large molecules.
Inspired by the recent progress with deep generative models, this paper proposes a novel and principled probabilistic framework for molecular geometry generation, which addresses the above two limitations.
Our framework combines the advantages of normalizing ﬂows (Dinh et al., 2014) and energy-based approaches (LeCun et al., 2006), which have a strong model capacity for modeling complex distributions, are ﬂexible to model long-range dependency between atoms, and enjoy efﬁcient sampling and training procedures.
Similar to the work of Simm & Hern ´ andez-Lobato (2020), we also ﬁrst learn the distribution of distances d given the graph G , i.e. , p ( d |G ) , and deﬁne another distribution of conformations R given the distances d , i.e. , p ( R | d , G ) .
Speciﬁcally, we propose a novel Conditional Graph Continuous Flow (CGCF) for distance geometry ( d ) generation conditioned on the molecular graph G .
Given a molecular graph G , CGCF deﬁnes an invertible mapping between a base distribution ( e.g. , a multivariate normal distribution) and the molecular distance geometry, using a virtually inﬁnite number of graph transformation layers on atoms represented by a
Neural Ordinary Differential Equations architecture (Chen et al., 2018).
Such an approach enjoys very high ﬂexibility to model complex distributions of distance geometry.
Once the molecular distance geometry d is generated, we further generate the 3D coordinates R by searching from the probability p
( R | d , G ) .
Though the CGCF has a high capacity for modeling complex distributions, the distances of different edges are still independently updated in the transformations, which limits its capacity for modeling long-range dependency between atoms in the sampling process.
Therefore, we further propose another unnormalized probability function, i.e. , an energy-based model (EBM) (Hinton & Salakhutdinov, 2006; LeCun et al., 2006; Ngiam et al., 2011), which acts as a tilting term of the ﬂow-based distribution and directly models the joint distribution of R .
Speciﬁcally, the EBM trains an energy function E ( R , G ) , which is approximated by a neural network.
The ﬂow- and energy-based models are combined in a novel way for joint training and mutual enhancement.
First, energy-based methods are usually difﬁcult to train due to the slow sampling process.
In addition, the distribution of conformations is usually highly multi-modal, and the sampling procedures based on Gibbs sampling or Langevin Dynamics (Bengio et al., 2013a;b) tend to get trapped around modes, making it difﬁcult to mix between different modes (Bengio et al., 2013a).
Here we use the ﬂow-based model as a proposal distribution for the energy model, which is capable to generate diverse samples for training energy models.
Second, the ﬂow-based model lacks the capacity to explicitly model the long-range dependencies between atoms, which we ﬁnd can however be effectively modeled by an energy function E ( R , G ) .
Our sampling process can be therefore viewed as a two-stage dynamic system, where we ﬁrst take the ﬂow-based model to quickly synthesize realistic conformations and then used the learned energy E ( R , G ) to reﬁne the generated conformations through Langevin Dynamics.
We conduct comprehensive experiments on several recently proposed benchmarks, including GEOM-QM9, GEOM-Drugs (Axelrod & Gomez-Bombarelli, 2020) and ISO17 (Simm &
Hern´andez-Lobato, 2020).
Numerical evaluations show that our proposed framework consistently outperforms the previous state-of-the-art (GraphDG) on both conformation generation and distance modeling tasks, with a clear margin.
P ROBLEM D EFINITION AND P RELIMINARIES 2. P ROBLEM D EFINITION
Notations.
Following existing work (Simm & Hern´andez-Lobato, 2020), each molecule is represented as an undirected graph G = (cid:104)V , E(cid:105) , where V is the set of nodes representing atoms and E is the set of edges representing inter-atomic bonds.
Each node v in
V is labeled with atomic properties such as element type.
The edge in
E connecting u and v is denoted as e uv , and is labeled with its bond type.
We also follow the previous work (Simm & Hern´andez-Lobato, 2020) to expand the molecular graph with auxiliary bonds, which is elaborated in Appendix B. For the molecular 3D representation, each atom in V is assigned with a 3D position vector r ∈ R 3 .
We denote d uv = (cid:107) r u − r v (cid:107) 2 as the Euclidean distance between the u th and v th atom.
Therefore, we can represent all the positions { r v } v ∈V as a matrix R ∈ R |V|× 3 and all the distances between connected nodes { d uv } e uv ∈E as a vector d ∈ R |E| .
Problem Deﬁnition.
The problem of molecular conformation generation is deﬁned as a conditional generation process.
More speciﬁcally, our goal is to model the conditional distribution of atomic positions R given the molecular graph G , i.e. , p ( R |G ) .
2. P RELIMINARIES
Continuous Normalizing Flow.
A normalizing ﬂow (Dinh et al., 2014; Rezende & Mohamed, 2015) deﬁnes a series of invertible deterministic transformations from an initial known distribution p ( z ) to a more complicated one p ( x ) .
Recently, normalizing ﬂows have been generalized from discrete number of layers to continuous (Chen et al., 2018; Grathwohl et al., 2018) by deﬁning the transformation f θ as a continuous-time dynamic ∂z ( t ) ∂t = f θ ( z ( t ) , t ) .
Formally, with the latent variable z ( t 0 ) ∼ p ( z ) at the start time, the continuous normalizing ﬂow (CNF) deﬁnes the transformation x
= z ( t 0 )
+ (cid:82) t 1 t
f
θ
( z ( t ) , t ) dt .
Then the exact density for p θ ( x ) can be computed by:
[Equation]
where z ( t 0 ) can be obtained by inverting the continuous dynamic z ( t 0 )
=
x + (cid:82) t 0 t 1 f θ ( z ( t ) , t ) dt .
A black-box ordinary differential equation (ODE) solver can be applied to estimate the outputs and inputs gradients and optimize the CNF model (Chen et al., 2018; Grathwohl et al., 2018).
Energy-based Models.
Energy-based models (EBMs) (Dayan et al., 1995; Hinton & Salakhutdinov, 2006; LeCun et al., 2006) use a scalar parametric energy function E φ ( x ) to ﬁt the data distribution.
Formally, the energy function induces a density function with the Boltzmann distribution p φ ( x ) = exp( − E φ ( x )) /Z
( φ ) , where Z = (cid:82) exp( − E φ ( x ))
dx denotes the partition function.
EBM can be learned with Noise contrastive estimation (NCE) (Gutmann & Hyv¨arinen, 2010) by treating the normalizing constant as a free parameter.
Given the training examples from both the dataset and a noise distribution q ( x ) , φ can be estimated by maximizing the following objective function:
[Equation]
which turns the estimation of EBM into a discriminative learning problem.
Sampling from E φ can be done with a variety of methods such as Markov chain Monte Carlo (MCMC) or Gibbs sampling (Hinton & Salakhutdinov, 2006), possibly accelerated using Langevin dynamics (Du &
Mordatch, 2019; Song et al., 2020), which leverages the gradient of the EBM to conduct sampling:
[Equation]
where (cid:15) refers to the step size.
x 0 are the samples drawn from a random initial distribution and we take the x K with K Langevin dynamics steps as the generated samples of the stationary distribution.
[Table]
/YKu9ndj6W2z9JW6SK0zig4HHezPMzAtizrRx3W+nsLa+sblV3C7t7O7tH5QPj1o6ShShTRLxSHUCrClnkjYNM5x2YkWxCDhtB+Obqd9+okqzSD6YSUx9gYeShYxgY6VWLxDpfdYvV9yaOwNaJV5OKpCj0S//9AYRSQSVhnCsdddzY+OnWBlGOM1KvUTTGJMxHtKupRILqv10dm2GzqwyQGGkbEmDZurfiRQLrScisJ0Cm5Fe9qbif143MeGVnzIZJ4ZKMl8UJhyZCE1fRwOmKDF8YgkmitlbERlhhYmxAS1sCURWsqF4yxGsktZ5zbuoXd9dVOrVPJ4inMApVMGDS6jDLTSgCQQe4QVe4c15dt6dD+dz3lpw8pljWIDz9QvyzpXD</latexit> p <latexit sha1_base64="ORiRlP3d8Ck3qImxG8IgZjn4WBQ=">AAAB+HicbVA9SwNBEJ2LXzF+RS1tFoNgIeFOAmoXsLFMwHxAcoS9zSRZsnt37O4J8cgvsNXeTmz9N7b+EjfJFSbxwcDjvRlm5gWx4Nq47reT29jc2t7J7xb29g8Oj4rHJ00dJYphg0UiUu2AahQ8xIbhRmA7VkhlILAVjO9nfusJleZR+GgmMfqSDkM+4IwaK9XjXrHklt05yDrxMlKCDLVe8afbj1giMTRMUK07nhsbP6XKcCZwWugmGmPKxnSIHUtDKlH76fzQKbmwSp8MImUrNGSu/p1IqdR6IgPbKakZ6VVvJv7ndRIzuPVTHsaJwZAtFg0SQUxEZl+TPlfIjJhYQpni9lbCRlRRZmw2S1sCOS3YULzVCNZJ87rsVcp39UqpepXFk4czOIdL8OAGqvAANWgAA4QXeIU359l5dz6cz0VrzslmTmEJztcvBTaTkA==</latexit>
GradientDescent
MCMC
Predict distances for the input graph.
Search 3D coordinates given the distances.
[Figure]
Further optimize the generated structures.
[Figure]
Continuous Flow (CGCF); 2) search the possible 3D coordinates according to the generated distances and 3) further optimize the generated conformation via a MCMC procedure with the Energy-based Tilting Model (ETM).
M ETHOD 3. O VERVIEW
We ﬁrst present a high-level description of our model.
Directly learning a generative model on Cartesian coordinates heavily depends on the (arbitrary) rotation and translation (Mansimov et al., 2019).
Therefore, in this paper we take the atomic pairwise distances as intermediate variables to generate conformations, which are invariant to rotation and translation.
More precisely, the cornerstone of our method is to factorize the conditional distribution p θ ( R |G ) into the following formulation:
[Equation]
where p θ ( d |G ) models the distribution of inter-atomic distances given the graph G and p ( R | d , G ) models the distribution of conformations given the distances d .
In particular, the conditional generative model p θ ( d |G ) is parameterized as a conditional graph continuous ﬂow, which can be seen as a continuous dynamics system to transform the random initial noise to meaningful distances.
This ﬂow model enables us to capture the long-range dependencies between atoms in the hidden space during the dynamic steps.
Though CGCF can capture the dependency between atoms in the hidden space, the distances of different edges are still independently updated in the transformations, which limits the capacity of modeling the dependency between atoms in the sampling process.
Therefore we further propose to correct p θ ( R |G ) with an energy-based tilting term E φ ( R , G ) :
[Equation]
The tilting term is directly deﬁned on the joint distribution of R and G , which explicitly captures the long-range interaction directly in observation space.
The tilted distribution p θ,φ ( R |G ) can be used to provide reﬁnement or optimization for the conformations generated from p
θ ( R
[Equation]
This energy function is also designed to be invariant to rotation and translation.
In the following parts, we will ﬁrstly describe our ﬂow-based generative model p θ ( R |G ) in Section 3.2 and elaborate the energy-based tilting model E φ ( R , G ) in Section 3.3.
Then we introduce the two-stage sampling process with both deterministic and stochastic dynamics in Section 3.4.
An illustration of the whole framework is given in Fig.
3. F LOW - BASED G ENERATIVE M ODEL
Conditional Graph Continuous Flows p θ ( d |G ) .
We parameterize the conditional distribution of distances p θ ( d |G ) with the continuous normalizing ﬂow, named C onditional G raph C ontinuous F low (CGCF).
CGCF deﬁnes the distribution through the following dynamics system:
[Equation]
where the dynamic f θ is implemented by Message Passing Neural Networks (MPNN) (Gilmer et al., 2017), which is a widely used architecture for representation learning on molecular graphs.
MPNN takes node attributes, edge attributes and the bonds lengths d ( t ) as input to compute the node and edge embeddings.
Each message passing layer updates the node embeddings by aggregating the information from neighboring nodes according to its hidden vectors of respective nodes and edges.
Final features are fed into a neural network to compute the value of the dynamic f θ for all distances independently.
As t 1 → ∞ , our dynamic can have an inﬁnite number of steps and is capable to model long-range dependencies.
The invertibility of F θ allows us to not only conduct fast sampling, but also easily optimize the parameter set θ by minimizing the exact negative log-likelihood:
[Equation]
Closed-form p ( R | d , G ) .
The generated pair-wise distances can be converted into 3D structures through postprocessing methods such as the classic Euclidean Distance Geometry (EDG) algorithm.
In this paper, we adopt an alternative way by deﬁning the conformations as a conditional distribution:
[Equation]
where Z is the partition function to normalize the probability and { α uv } are parameters that control the variance of desired Cartesian coordinates, which can be either learned or manually designed according to the graph structure G .
With the probabilistic formulation, we can conduct either sampling via MCMC or searching the local optimum with optimization methods.
This simple function is fast to calculate, making the generation procedure very efﬁcient with a negligible computational cost.
Compared with the conventional EDG algorithm adopted in GraphDG (Simm & Hern´andez-Lobato, 2020), our probabilistic solution enjoys following advantages: 1) p ( R | d , G ) enables the calculation for the likelihood p θ ( R |G ) of Eq.
4 by approximation methods, and thus can be further combined with the energy-based tilting term E φ ( R , G ) to induce a superior distribution; 2) GraphDG suffers the drawback that when invalid sets of distances are generated, EDG will fail to construct 3D structure.
By contrast, our method can always be successful to generate conformations by sampling from the distribution p
( R | d , G ) .
3. E NERGY - BASED T ILTING M ODEL
The last part of our framework is the E nergy-based T iling M odel (ETM) E φ ( R , G ) , which helps model the long-range interactions between atoms explicitly in the observation space.
E φ ( R , G ) takes the form of SchNet (Sch ¨ utt et al., 2017), which is widely used to model the potential-energy surfaces and energy-conserving force ﬁelds for molecules.
The continuous-ﬁlter convolutional layers in SchNet allow each atom to aggregate the representations of all single, pairwise, and higherorder interactions between the atoms through non-linear functions.
The ﬁnal atomic representations are pooled to a single vector and then passed into a network to produce the scalar output.
Typically the EBMs can be learned by maximum likelihood, which usually requires the lengthy
MCMC procedure and is time-consuming for training.
In this work, we learn the ETM by Noise
Contrastive Estimation (Gutmann & Hyv¨arinen, 2010), which is much more efﬁcient.
In practice, the noise distribution is required to be close to data distribution, otherwise the classiﬁcation problem would be too easy and would not guide E φ to learn much about the modality of the data.
We propose to take the pre-trained CGCF to serve as a strong noise distribution, leading to the following discriminative learning objective for the ETM 1 :
[Equation]
3. S AMPLING
We employ a two-stage dynamic system to synthesize a possible conformation given the molecular graph representation G .
In the ﬁrst stage, we ﬁrst draw a latent variable ˆ z from the Gaussian prior
N (0 , I ) , and then pass it through the continuous deterministic dynamics model F θ deﬁned in Eq.to get ˆ d 0 =
F θ (ˆ z 0 , G ) .
Then an optimization procedure such as stochastic gradient descent is employed to search the realistic conformations R with local maximum probability of p ( R | d , G ) (deﬁned in Eq.
By doing this, an initial conformation R (0) can be generated.
In the second stage, we further reﬁne the initial conformation R (0) with the energy-based model deﬁned in Eq.with K steps of Langevin dynamics:
[Equation]
where E θ,φ ( R |G )
= − log p θ,φ ( R |G ) =
E φ ( R , G
[Equation]
−
[Equation]
10) where (cid:15) denotes the step size.
The second integration term in E θ,φ can be estimated through approximate methods.
In practice, we use Monte Carlo Integration to conduct the approximation, which is simple yet effective with just a few distance samples from the CGCF model p θ ( d |G ) .
E XPERIMENTS 4. E XPERIMENT S ETUP
Evaluation Tasks.
To evaluate the performance of proposed model, we conduct experiments by comparing with the counterparts on: (1) Conformation Generation evaluates the model’s capacity to learn the distribution of conformations by measuring the diversity and accuracy of generated samples (section 4.2); (2) Distribution over distances is ﬁrst proposed in Simm & Hern ´ andez-Lobato (2020), which concentrate on the distance geometry of generated conformations (section 4.2).
Benchmarks.
We use the recent proposed GEOM-QM9 and GEOM-Drugs (Axelrod & GomezBombarelli, 2020) datasets for conformation generation task and ISO17 dataset (Simm &
Hern ´ andez-Lobato, 2020) for distances modeling task.
The choice of different datasets is because of their distinct properties.
Speciﬁcally, GEOM datasets consist of stable conformations, which is suitable to evaluate the conformation generation task.
By contrast, ISO17 contains snapshots of molecular dynamics simulations, where the structures are not equilibrium conformations but can reﬂect the density around the equilibrium state.
Therefore, it is more suitable for the assessment of similarity between the model distribution and the data distribution around equilibrium states.
More speciﬁcally, GEOM-QM9 is an extension to the QM9 (Ramakrishnan et al., 2014) dataset: it contains multiple conformations for most molecules while the original QM9 only contains one.
This dataset is limited to 9 heavy atoms (29 total atoms), with small molecular mass and few rotatable bonds.
We randomly draw 50000 conformation-molecule pairs from GEOM-QM9 to be the training set, and take another 17813 conformations covering 150 molecular graphs as the test set.
GEOM-Drugs dataset consists of much larger drug molecules, up to a maximum of 181 atoms ( heavy atoms).
It also contains multiple conformations for each molecule, with a larger variance in structures, e.g. , there are the 6 . 5 rotatable bonds in average.
We randomly take 50000 conformationmolecule pairs from GEOM-Drugs as the training set, and another 9161 conformations (covering 100 molecular graphs) as the test split.
ISO17 dataset is also built upon QM9 datasets, which consists of 197 molecules, each with 5000 conformations.
Following Simm & Hern´andez-Lobato (2020), we also split ISO17 into the training set with 167 molecules and the test set with anothermolecules.
Baselines .
We compared our proposed method with the following state-of-the-art conformation generation methods.
CVGAE (Mansimov et al., 2019) uses a conditional version of VAE to directly generate the 3D coordinates of atoms given the molecular graph.
GraphDG (Simm & Hern´andezLobato, 2020) also employs the conditional VAE framework.
Instead of directly modeling the 3D structure, they propose to learn the distribution over distances.
Then the distances are converted into conformations with an EDG algorithm.
Furthermore, we also take RDKit (Riniker & Landrum, 2015) as a baseline model, which is a classical EDG approach built upon extensive calculation collections in computational chemistry.
[Table]
, the threshold δ is set as 0 .
5 ˚ A for QM9 and 1 .
25 ˚ A for Drugs.
More results of COV scores with different threshold δ are given in Appendix H.
4. C ONFORMATION G ENERATION
In this section, we evaluate the ability of the proposed method to model the equilibrium conformations.
We focus on both the diversity and accuracy of the generated samples.
More speciﬁcally, diversity measures the model’s capacity to generate multi-modal conformations, which is essential for discovering new conformations, while accuracy concentrates on the similarity between generated conformations and the equilibrium conformations.
Evaluation.
For numerical evaluations, we follow previous work (Hawkins, 2017; Mansimov et al., 2019) to calculate the Root-Mean-Square Deviation (RMSD) of the heavy atoms between generated samples and reference ones.
Precisely, given the generated conformation R and the reference R ∗ , we obtain ˆ R by translating and rotating R ∗ to minimize the following predeﬁned RMSD metric:
[Equation]
(11) where n is the number of heavy atoms.
Then the smallest distance is taken as the evaluation metric.
Built upon the RMSD metric, we deﬁne Cov erage (COV) and Mat ching (MAT) score to measure the diversity and quality respectively.
Intuitively, COV measures the fraction of conformations in the reference set that are matched by at least one conformation in the generated set.
For each conformation in the generated set, its neighbors in the reference set within a given RMSD threshold
[Table]
δ are marked as matched:
[Equation]
where S g ( G ) denotes the generated conformations set for molecular graph G
[Table]
and S r ( G ) denotes the reference set.
In practice, the number of samples in the generated set is two times of the reference set.
Typically, a higher COV score means the a better diversity performance.
The COV score is able to evaluate whether the generated conformations are diverse enough to cover the ground-truth.
While COV is effective to measure the diversity and detect the mode-collapse case, it is still possible for the model to achieve high COV with a high threshold tolerance.
Here we deﬁne the MAT score as a complement to measure the quality of generated samples.
For each conformation in the reference set, the RMSD distance to its nearest neighbor in the generated set is computed and averaged:
[Equation]
This metric concentrate on the accuracy of generated conformations.
More realistic generated samples lead to a lower matching score.
Results.
Tab.
1 shows that compared with the existing state-of-the-art baselines, our CGCF model can already achieve superior performance on all four metrics (top 4 rows).
As a CNF-based model,
CGCF holds much the higher generative capacity for both diversity and quality compared than VAE approaches.
The results are further improved when combined with ETM to explicitly incorporate the long-range correlations.
We visualize several representative examples in Fig.
2, and leave more examples in Appendix G. A meaningful observation is that though competitive over other neural models, the rule-based RDKit method occasionally shows better performance than our model, which indicates that RDKit can generate more realistic structures.
We argue that this is because after generating the initial coordinates, RDKit involves additional hand-designed molecular force ﬁeld (FF) energy functions (Rapp´e et al., 1992; Halgren, 1996a) to ﬁnd the stable conformations with local minimal energy.
By contrast, instead of ﬁnding the local minimums, our deep generative models aim to model and sample from the potential distribution of structures.
To yield a better comparison, we further test our model by taking the generated structures as initial states and utilize the Merck Molecular Force Field (MMFF)
(Halgren, 1996a) to ﬁnd the local stable points.
A more precise description of about the MMFF Force Field algorithms in RDKit is given in
Appendix I.
This postprocessing procedure is also employed in the previous work (Mansimov et al., 2019).
Additional results in Tab.
1 verify our conjecture that FF plays a vital role in generating more realistic structures, and demonstrate the capacity of our method to generate high-quality initial coordinates.
4. D ISTRIBUTIONS O VER D ISTANCES
Tough primarily designed for 3D coordinates, we also following Simm & Hern´andez-Lobato (2020) to evaluate the generated distributions of pairwise distance, which can be viewed as a representative element of the model capacity to model the inter-atomic interactions.
Evaluation.
Let p ( d uv |G ) denote the conditional distribution of distances on each edge e uv given a molecular graph
G .
The set of distances are computed from the generated conformations R .
We
[Table]
calculate maximum mean discrepancy (MMD) (Gretton et al., 2012) to compare the generated distributions and the ground-truth distributions.
Speciﬁcally, we evaluate the distribution of individual distances p
( d uv |G ) , pair distances p ( d uv , d ij |G ) and all distances p ( d |G ) .
For this benchmark, the number of samples in the generated set is the same as the reference set.
Results.
The results of MMD are summarized in Tab. 2.
The statistics show that RDKit suffers the worst performance, which is because it just aims to generate the most stable structures as illustrated in Section 4.2.
For CGCF, the generated samples are signiﬁcantly closer to the ground-truth distribution than baseline methods, where we consistently achieve the best numerical results.
Besides, we notice that ETM will slightly hurt the performance in this task.
However, one should note that this phenomenon is natural because typically ETM will sharpen the generated distribution towards the stable conformations with local minimal energy.
By contrast, the ISO17 dataset consists of snapshots of molecular dynamics where the structures are not equilibrium conformations but samples from the density around the equilibrium state.
Therefore, ETM will slightly hurt the results.
This phenomenon is also consistent with the observations for RDKit.
Instead of generating unbiased samples from the underlying distribution, RDKit will only generate the stable ones with local minimal energy by involving the hand-designed molecular force ﬁeld (Simm & Hern ´ andez-Lobato, 2020).
And as shown in the results, though highly competitive in Tab. 1, RDKit also suffers much weaker results in Tab. 2.
The marginal distributions P ( d uv |G ) for pairwise distances in visualized in Appendix K, which further demonstrate the superior capacity of our proposed method.
We also follow Mansimov et al.
(2019) to calculate the diversity of conformations generated by all compared methods, which is measured by calculating the mean and standard deviation of the pairwise RMSD between each pair of generated conformations per molecule.
The results shown in Tab.
3 demonstrate that while our method can achieve the lowest MMD, it does not collapse to generating extremely similar conformations.
Besides, we observe that ETM will slightly hurt the diversity of CGCF, which veriﬁes our statement that ETM will sharpen the generated distribution towards the stable conformations with local minimal energy.
C ONCLUSION AND F UTURE W ORK
In this paper, we propose a novel probabilistic framework for molecular conformation generation.
Our generative model combines the advantage of both ﬂow-based and energy-based models, which is capable of modeling the complex multi-modal geometric distribution and highly branched atomic correlations.
Experimental results show that our method outperforms all previous state-of-the-art baselines on the standard benchmarks.
Future work includes applying our framework on much larger datasets and extending it to more challenging structures ( e.g. , proteins).
A CKNOWLEDGMENTS
This project is supported by the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ldt., Amazon Faculty Research Award,
Tencent AI Lab
Rhino-Bird Gift Fund and a NRC Collaborative R&D Project (AI4D-CORE-06).
This project was also partially funded by IVADO Fundamental Research Project grant PRF-2019-3583139727.
A R ELATED W ORKS
Conformation Generation.
There have been results showing deep learning speeding up molecular dynamics simulation by learning efﬁcient alternatives to quantum mechanics-based energy calculations (Sch¨utt et al., 2017; Smith et al., 2017).
However, though accelerated by neural networks, these approaches are still time-consuming due to the lengthy MCMC process.
Recently, Gebauer et al.
and Hoffmann & No´e (2019) propose to directly generate 3D structures with deep generative models.
However, these models can hardly capture graph- or bond-based structure, which is typically complex and highly branched.
Some other works (Lemke & Peter, 2019; AlQuraishi, 2019; Ingraham et al., 2019; No´e et al., 2019; Senior et al., 2020) also focus on learning models to directly generate 3D structure, but focus on the protein folding problem.
Unfortunately, proteins are linear structures while general molecules are highly branched, making these methods not naturally transferable to general molecular conformation generation tasks.
Energy-based Generative Model.
There has been a long history for energy-based generative models.
Xie et al.
(2016) proposes to train an energy-based model parameterized by modern deep neural network and learned it by Langevin based MLE.
The model is called generative ConvNet since it can be derived from the discriminative ConvNet.
In particular, this paper is the ﬁrst to formulate modern ConvNet-parametrized EBM as exponential tilting of a reference distribution, and connect it to discriminative ConvNet classiﬁer.
More recently, Du & Mordatch (2019) implemented the deep
EBMs with ConvNet as energy function and achieved impressive results on image generation.
Different from the previous works, we concentrate on molecular geometry generation, and propose a novel and principled probabilistic framework to address the domain-speciﬁc problems.
More speciﬁcally, we ﬁrst predict the atomic distances through the continuous normalizing ﬂow, and then convert them to the desired 3D conformation and optimize it with the energy-based model.
This procedure enables us to keep the rotational and translational invariance property.
Besides, to the best of our knowledge, we are the ﬁrst one to combine neural ODE with EBMs.
We take the ODE model to improve the training of EBM, and combine both to conduct the two-stage sampling dynamics.
B D ATA P REPROCESS
Inspired by classic molecular distance geometry (Crippen et al., 1988), we also generate the conﬁrmations by ﬁrstly predicting all the pairwise distances, which enjoys the invariant property to rotation and translation.
Since the bonds existing in the molecular graph are not sufﬁcient to characterize a conformation, we pre-process the graphs by extending auxiliary edges.
Speciﬁcally, the atoms that are 2 or 3 hops away are connected with virtual bonds , labeled differently from the real bonds of the original graph.
These extra edges contribute to reducing the degrees of freedom in the 3D coordinates, with the edges between second neighbors helping to ﬁx the angles between atoms, and those between third neighbors ﬁxing dihedral angles.
C N ETWORK A RCHITECTURE
In this section, we elaborate on the network architecture details of CGCF and ETM.
C. C ONTINUOUS G RAPH F LOW
In CGCF, the dynamic function f θ deﬁned in Eq.
6 is instanced with a message passing neural networks.
Given the node attributes, edge attributes and intermediate edge lengths as input, we ﬁrst embed them into the feature space through feedforward networks:
[Equation]
Then, the node and edge features along are passed sequentially into L layers message passing networks with the graph structure G :
[Equation]
where N G ( v ) denotes the ﬁrst neighbors in the graph G and σ is the activation function.
After L message passing layers, we use the ﬁnal hidden representation h ( L ) as the node representations.
Then for each bond, the corresponding node features are aggregated along with the edge feature to be fed into a neural network to compute the value of the dynamic f θ :
[Equation]
The ETM is implemented with SchNet.
It takes both the graph and conformation information as input and output a scalar to indicate the energy level.
Let the atoms are described by a tuple of features X l = ( x l 1 , . . .
, x ln ) , where n denote the number of atoms and l denote the layer.
Then given the positions R , the node embeddings are updated by the convolution with all surrounding atoms:
[Equation]
where ”o” represents the element-wise multiplication.
It is straightforward that the above function enables to include translational and rotational invariance by computing pairwise distances instead of using relative positions.
After L convolutional layers, we perform a sum-pooling operator over the node embeddings to calculate the global embedding for the whole molecular structure.
Then the global embedding is fed into a feedforward network to compute the scalar of the energy level.
D T WO - STAGE D YNAMIC S YSTEM FOR S AMPLING
Input : molecular graph G , CGCF model with parameter θ , ETM with parameter φ , the number of optimization steps for p ( R | d , G ) M and its step size r , the number of MCMC steps for E θ,φ N and its step size (cid:15) molecular
(0 , I ) 2:
d = F θ ( d ( t 0 ) ,
G ) 3: for m = 1 , ..., M do 4:
R m =
R m − 1 + r ∇ R log p
( R | d , G ) 5: end for 6: for n = 1 , ..., N do 7:
R n = R n − 1 − (cid:15) 2 ∇ R E
θ,φ ( R |G ) + √ (cid:15)ω, ω ∼ N (0 , I ) , 8: end for
E I MPLEMENTATION D ETAILS
Our model is implemented in PyTorch (Paszke et al., 2017).
The MPNN in CGCF is implemented with 3 layers, and the embedding dimension is set as 128.
And the SchNet in ETM is implemented with 6 layers with the embedding dimension set as 128.
We train our CGCF with a batch size ofand a learning rate of 0.001 until convergence.
After obtaining the CGCF, we train the ETM with a batch size of 384 and a learning rate of 0.001 until convergence.
For all experimental settings, we use Adam (Kingma & Ba, 2014) to optimize our model.
F D ETAILED D ERIVATIONS OF E NERGY - BASED M ODEL
Here we present the detailed derivations of the training objective function of Energy-based Tilting
Model (ETM) in Eq. 9:
[Equation]
G M ORE G ENERATED S AMPLES
We present more visualizations of generated 3D structures in Fig.
3, which are generated from our model (CGCF + ETM) learned on both GEOM-QM9 and GEOM-Drugs datasets.
The visualizations demonstrate that our proposed framework holds the high capacity to model the chemical structures in the 3D coordinates.
H M ORE R ESULTS OF C OVERAGE S CORE
We give more results of the coverage (COV) score with different threshold δ in Fig. 4.
As shown in the ﬁgure, our proposed method can consistently outperform the previous state-of-the-art baselines
CVGAE and GraphDG, which demonstrate the effectiveness of our model.
[Figure]
I I MPLEMENTATION FOR MMFF
In this section, we give a more precise description of the MMFF Force Field implementation in the
RDKit toolkit (Riniker & Landrum, 2015).
In MMFF, the energy expression is constituted by seven terms: bond stretching, angle bending, stretch-bend, out-of-plane bending, torsional, van der Waals and electrostatic.
The detailed functional form of individual terms can be found in the original literature (Halgren, 1996a).
To build the force ﬁeld for a given molecular system, the ﬁrst step is to assign correct types to each atom.
At the second step, atom-centered partial charges are computed according to the MMFF charge model (Halgren, 1996b).
Then, all bonded and non-bonded interactions in the molecular system under study, depending on its structure and connectivity, are loaded into the energy expression.
Optionally, external restraining terms can be added to the MMFF energy expression, with the purpose of constraining selected internal coordinates during geometry optimizations.
Once all bonded and non-bonded interactions, plus optional restraints, have been loaded into the MMFF energy expression, potential gradients of the system under study can be computed to minimize the energy.
J M ORE E VALUATIONS FOR C ONFORMATION G ENERATION
Junk Rate.
The COV and MAT score in Section 4.2 do not appear to explicitly measure the generated false samples.
Here we additionally deﬁne Junk rate measurement.
Intuitively, JUNK measures the fraction of generated conformations that are far away from all the conformations in the reference set.
For each conformation in the generated set, it will be marked as a false sample if its RMSD to all the conformations of reference set are above a given threshold δ :
[Equation] [Equation]
Typically, a lower JUNK rate means better generated quality.
The results are shown in Tab. 4.
As shown in the table, our CGCF model can already outperform the existing state-of-the-art baselines with an obvious margin.
The results are further improved when combined with ETM to explicitly incorporate the long-range correlations.
K D ISTANCE D ISTRIBUTION V ISUALIZATION
In Fig.
5, we plot the marginal distributions p
( d uv |G ) for all pairwise distances between C and O atoms of a molecular graph in the ISO17 test set.
As shown in the ﬁgure, though primarily designed
[Table]
reported JUNK
[Table]
threshold δ
[Table]
Drugs
[Table]
for 3D structure generation, our method can make much better estimation of the distances than
GraphDG, which is the state-of-the-art model for molecular geometry prediction.
As a representative element of the pairwise property between atoms, the inter-atomic distances demonstrate the capacity of our model to capture the inter-atomic interactions.
ISO17 GraphDG
Ours 0 8 5 0 1 84 5 6