D IRECT M OLECULAR C ONFORMATION G ENERATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Molecular conformation generation, which is to generate 3 dimensional coordinates of all the atoms in a molecule, is an important task for bioinformatics and pharmacology.
Most existing machine learning based methods ﬁrst predict interatomic distances and then generate conformations based on them.
This twostage approach has a potential limitation that the predicted distances may conﬂict with each other, e.g., violating the triangle inequality.
In this work, we propose a method that directly outputs the coordinates of atoms, so that there is no violation of constraints.
The conformation generator of our method stacks multiple blocks, and each block outputs a conformation which is then reﬁned by the following block.
We adopt the variational auto-encoder (VAE) framework and use a latent variable to generate diverse conformations.
To handle the roto-translation equivariance, we adopt a loss that is invariant to rotation and translation of molecule coordinates, by computing the minimal achievable distance after any rotation and translation.
Our method outperforms strong baselines on four public datasets, which shows the effectiveness of our method and the great potential of the direct approach.
The code is released at https: //github.com/DirectMolecularConfGen/DMCG .
I NTRODUCTION
Molecular conformation generation aims to generate the 3D coordinates of all the atoms of molecules, which then can be used in molecular property prediction (Axelrod & Gomez-Bombarelli, 2021), docking (Roy et al., 2015), structure-based virtual screening (Kontoyianni, 2017), etc.
Molecular conformation can be physically observed using X-ray crystallography, but it is prohibitively costly for industry-scale tasks (Mansimov et al., 2019).
Ab initio methods, such as density functional theory (DFT) (Parr, 1980; Baseden & Tye, 2014), can accurately predict the molecular shapes, but take up to several hours per small molecule (Hu et al.,
2021).To handle large-scale molecules, people turn to leverage the classical force ﬁelds, like UFF (Rappe et al., 1992) or MMFF (Halgren, 1996), to estimate conformations, which is efﬁcient but unacceptably inaccurate (Kanal et al., 2018).
Recently, people start to explore machine learning methods to generate conformation.
Simm &
Hern ´ andez-Lobato (2020), Shi et al.
(2020) and Shi et al.
(2021) leveraged variational auto-encoder (VAE), ﬂow-based models and score-based methods for conformation generation, respectively.
With learned models, molecule conformations can be sampled independently (Simm & Hern´andezLobato, 2020; Shi et al., 2020) or using Langevin dynamics (Xu et al., 2021a; Shi et al., 2021).
The common part of these methods is that they are all built upon the interatomic distances among atoms (i.e., the distance between atom pairs).
Speciﬁcally, Simm & Hern´andez-Lobato (2020); Shi et al.
Xu et al.
(2021a;b) use various generative methods to model the distribution of interatomic distances, and then reconstruct conformations based on distances.
Shi et al.
(2021) leveraged a score-matching network to model the density gradient of interatomic distances.
Moreover, Winter et al.
(2021) use variational auto-encoder to ﬁrst predict bond length, bond angle and dihedral angle and then reconstruct the coordinates based on the intermediate results.
A major reason of using distance-based methods is that the interatomic distances are invariant to rotation and translation of the conformation.
However, a potential drawback is that the predicted distances might conﬂict with each other.
For example, they did not explicitly consider the triangle inequality for distances among any three atoms.
The underlying degree-of-freedom of these distances is only 3 N − 6 ( N is the number of atoms in a molecule) while practically, these methods often generate distances with the degree-of-freedom of roughly N ( N − 1) / 2 , which usually lead to violations of the triangle inequality.
For example, we found that in GraphDG (Simm & Hern´andez-Lobato, 2020), a prevailing representative of distance-based method, 8 .
65% of molecular graphs in the GEOM-Drugs test set will produce distance matrix that violates the triangle inequality.
In this work, we explore the possibility to directly generate the coordinates of all atoms.
This is a straightforward and more natural choice and avoids violations of the triangle inequality, and has demonstrated remarkable performance on protein structure prediction by the AlphaFold 2 (Jumper et al., 2021).
We design a model that generates atom coordinates directly.
The generator in our model stacks multiple blocks, and each block outputs a conformation which is then reﬁned by the following block.
A block consists of several modules that can encode the previous conformation as well as the features of bonds, atoms and global information.
At the end of each block, we add a normalization layer that centers the coordinates at the origin.
Since a molecule may have multiple conformations, we use the variational auto-encoder (VAE) framework which allows diverse generation.
To realize roto-translation equivariance, i.e., rotating and translating a set of atom coordinates do not change the conformation, we adopt a loss that is invariant to rotation and translation of atom coordinates.
We conduct experiments on four benchmark settings, which are GEOM-QM9 and GEOM-Drugs with small-scale setting (Shi et al., 2021) and large-scale setting (Axelrod & Gomez-Bombarelli, 2021).
Compared with previous methods, ours achieves state-of-the-art results on GEOM-Drugs (both small-scale and large-scale settings), and outperforms almost all previous baselines on GEOMQM9, demonstrating the effectiveness of our method.
Our contributions are summarized as follows:
(1) We explore a new direction of conformation generation, that directly generates the coordinates of a molecular conformation without generating interatomic distances.
We empirically show that directly generating conformation achieves state-of-the-art results on several tasks.
(2) We leverage a ﬁne-grained loss function for training, that is invariant to the roto-translation of the conformation.
(3) We propose a new model that iteratively reﬁnes the conformations.
Our model is inspired by mul - tiple advanced architectures like GATv2 (Brody et al., 2021), graph network (GN) block (Battaglia et al., 2018), that can effectively model molecules.
P RELIMINARIES
In this section, we introduce the notations used in this work, give a formal deﬁnition of the molecular conformation generation problem, and brieﬂy introduce how to align two conformations under rotation and translation.
Notations : Let G =
( V, E ) denote a molecular graph, where V and E are collections of atoms and bonds, respectively.
Speciﬁcally, V = { v 1 , v 2 , · · · , v | V | } with the i -th atom v i .
Let e ij denote the bond between atom v i and v j .
Let N ( i ) denote the neighbors of atom i , i.e., N ( i )
= { v j | e ij ∈ E } .
We use R to represent the conformation of G , where R ∈ R | V |× 3 .
The i -th row of R is the coordinate of atom v i .
Let ρ ( · ) denote a roto-translation operation, i.e. an afﬁne transformation ρ ( R )
=
R t + b where the orientation-preserving orthogonal transformation t ∈ SO(3) ⊂
R 3 × 3 represents a rotation and the vector b ∈ R 1 × 3 represents a translation.
Let MLP ( · · · ) denote the multi-layer perception network, where the inputs are concatenated together and then processed by linear mapping, Batch Normalization (Ioffe & Szegedy, 2015) and ReLU activation sequentially.
Problem deﬁnition :
Given a graph G =
( V, E ) , our task is to learn a mapping, that can output the coordinates R of all atoms in V , i.e., R ∈ R | V |× 3 .
Matching loss : Let R 1 and R 2 denote two conformations, both of which are N × 3 matrices
N is the number of rows in R 1 and R 2 ).
The matching loss between R 1 and R 2 is deﬁned as follows:
[Equation]
[Equation] where (cid:107) · (cid:107) F denotes the Frobenius norm, i.e., (cid:107) A (cid:107) 2 F = (cid:80) i,j | A ij | 2 .
The matching loss is invariant to the rotation and translation of either of the two input conformations.
It is obviously nonnegative, and is zero if R 1 is obtained by a roto-translational operation of R 2 .
Karney (2007) proposed to use quaternions to solve Eqn.(1).
A quaternion q is an extension of complex numbers, q = q 0 u + q 1 i + q 2 j + q 3 k , where q 0 , q 1 , q 2 , q 3 are scalars and u , i , j , k are orientation vectors.
With quaternions, any rotation operation is speciﬁed by a 3 × 3 matrix, where each element in the matrix is the basic algebra of q 0 to q 3 and some constants.
The solution to
Eqn.(1) is the minimal eigenvalue of a 4 × 4 matrix obtained by the basic algebra on R 1 and R 2 , the corresponding optimal rotation and translation ρ ∗ can be obtained from the eigenvectors.
Note that the best ρ ∗ is related to both R 1 and R 2 .
To stabilize training, we set the gradients of ρ ∗ w.r.t.
the model parameter as zero.
More details can be found in Karney (2007).
O UR METHOD
We introduce our proposed method in this section.
3. C ONFORMATION GENERATION FRAMEWORK
Our goal can be formalized as learning to generate samples from the conditional distribution p
( R | G ) .
We follow the variational auto-encoder (VAE) scheme (Kingma & Welling, 2014; Rezende et al., 2014; Sohn et al., 2015) as it enables ﬂexible model architecture (vs. ﬂow-based models), efﬁcient (i.e., i.i.d) generation (vs. score-matching methods), and stable training and diverse sample generation (vs. adversarial methods).
Given a molecule graph G , the conditional variant of VAE generates
R using a likelihood model p ( R | z, G ) where z is a latent code drawn from a prior p ( z ) =
N (0 , I ) .
The corresponding p ( R | G ) = (cid:82) p ( z ) p ( R | z, G )
d z is however intractable to evaluate and maximize to ﬁt data.
VAE handles this by introducing a (conditional) inference model q ( z | R, G ) which leads to the following identity:
[Equation]
where D KL denotes the Kullback-Leibler (KL) divergence, and p ( z | R, G )
=
[Equation]
z is the true posterior.
Due to the nonnegativity of the third term, the ﬁrst two terms form a lower bound of log p
( R | G ) .
When q ( z | R, G ) is properly chosen, the bound is tractable to optimize.
Moreover, as the left-hand-side is independent of q ( z | R, G ) , tightening the bound (i.e., minimizing the third term) can be achieved by maximizing the bound w.r.t q ( z | R, G ) .
Therefore, as a loss function to be minimized, the objective is:
[Equation]
Now we develop our speciﬁcation for the conformation generation task.
We construct the likelihood model with regard to a decoded reference conformation ˆ R ( z, G ) :
[Equation]
where γ > 0 is a variance parameter that we choose to ﬁx as 1 (see Sec.
3.4 (1) for explanation).
Note that the usage of (cid:96) M guarantees all roto-translation equivalent conformations have the same probability.
For the inference model, we choose
[Equation]
where the conditional mean and variance are outputs from some encoder.
It enables tractable loss optimization via reparameterization: z ∼ q ( z | R, G ) is equivalent to z = µ R,G + Σ R,G (cid:15) where (cid:15) ∼ N (0 , I ) .
With the above speciﬁcation, the loss function becomes:
[Equation]
where the minimization is applied over all model parameters.
Note that in Eqn.(6), an additional hyperparameter β is introduced in the spirit of β -VAE (Higgins et al., 2016) to handle the strength of the prior regularization.
The second term in Eqn.(6) has a closed-form expression for optimization.
3. W ORKFLOW
According to Eqn.(3), we need a module ϕ
3D to model q ( z | G, R ) , a module ϕ dec to model p ( R | z, G ) and a module ϕ 2D to model the conditional input G .
In this subsection, we introduce the general workﬂow of how they work together (see Figure 1).
Details of the ϕ ··· modules will be covered in the next subsection.
The training workﬂow is shown as follows:
(1) The encoder ϕ 2D takes the molecular graph G as its input, and outputs d -dimensional graph representations H (0) V ∈ R | V |× d for all atoms, H (0) E ∈ R | E |× d for all bonds, a global graph feature U (0)
∈ R d , and initial conformation
ˆ R (0) ∈ R | V |× 3 .
Formally, ( H (0) V , H (0) E , U (0) , ˆ R (0) )
=
ϕ 2D ( G ) .
We use another encoder ϕ
3D to extract features of the conformation R for constructing the conditional inference model q ( z | R, G ) .
According to the above speciﬁcation, this encoder only needs to output the mean and variance of the Gaussian, or formally, ( µ R,G , Σ R,G ) =
ϕ 3D ( R, G ) .
We randomly sample a variable z from the Gaussian distribution N ( µ R,G , Σ R,G ) , and then feed
H (0) V , H (0) E , U (0) , ˆ R (0) , z into the decoder ϕ dec to obtain the conformation ˆ R .
That is, ˆ R ( z, G ) =
ϕ dec ( ϕ 2d ( G ) , z ) =
ϕ
dec ( H (0) V , H (0) E , ˆ R (0) , U (0) , z ) .
Note that sampling z ∼ N ( µ R,G , Σ R,G ) is equivalent to sampling (cid:15) ∼ N (0 , I ) and then setting z = µ R,G + Σ R,G (cid:15) .
(4) After obtaining ˆ R ( z, G ) and and N ( µ R,G , Σ R,G ) , we use Eqn.(6) as the objective function for training.
Remind that ˆ R is related to ϕ 2D , ϕ 3D , ϕ dec , and µ R,G , Σ R,G are related to ϕ 3D .
For the inference workﬂow, there are three steps in total: (1) Given a molecular graph G , we use ϕ 2D to encode G and obtain ˆ R (0) , H (0) V , H (0) E , U (0) ; (2) sample a random variable z from Gaussian
N (0 , I ) ; (3) feed ˆ R (0) , H (0) V , H (0) E , U (0) , z into ϕ dec and obtain the eventual conformation
ˆ R ( z, G ) .
Note that ϕ 3D is not used in inference phase.
3. D ETAILED MODEL ARCHITECTURE
The encoders ϕ 2D , ϕ 3D and the decoder ϕ dec share the same architecture.
They are all stacks of
L identical blocks.
We introduce the l -th block in the decoder ϕ dec , and leave the introduction of ϕ 2D and ϕ 3D in Appendix A.
In contrast to previous work, we introduce some novel designs that handle the task at a ﬁner level.
The conformation is iteratively reﬁned by each decoding block.
Speciﬁcally, we use a normalization technique that moves the center of each decoded conformation to the origin.
(2) We use more advanced model components, like attention models (Brody et al., 2021) and the GN block (Battaglia et al., 2018), that can more effectively represent the molecules.
The detailed architecture is shown in Figure 2.
The inputs include the following information output by the ( l − 1)
-th block: the conformation ˆ R ( l − 1) , atom features H ( l − 1) V ∈ R | V |× 3 , edge features
H ( l − 1) E and the global feature U ( l − 1) .
Note H (0) V , H (0) E , U (0) and ˆ R (0) are the outputs of ϕ 2D .
ˆ R ( l − 1) is ﬁrst mapped by MLP and get F ( l ) =
MLP ( ˆ R ( l − 1) ) .
The atom features are then fused with the coordinate features by H ( l − 1) V ← [Equation] z , where z ∼ N [Equation]
µ R
[Equation]
G , Σ R [Equation] G ) . We use a variant of the GN block (Battaglia et al., 2018) as the backbone of the model due to its superior performance in molecular modeling.
In each block, we update the bond features, atom fea- tures and global features sequentially.
Each block will output its own prediction on the coordinates of atoms.
For ease of reference, we use h ( l ) i (single subscript) to denote the feature of atom i output by the l -th block, and use h ( l ) ij (two subscripts) to denote the bond features between atom i and j output by the l -th block.
The workﬂow of the l -th block is shown as follows: (
1) Update bond features : For each bond feature h ( l − 1) ij
in H ( l − 1) E ,
it
is updated by
[Equation]
Update atom features : The atom features are updated using an attentive way.
For any i ∈ [ | V | ]
[Equation]
In Eqn.(8), a , W q , W v and W k are the parameters to be learned and concat ( · , · ) is the concatenation of two vectors.
For atom v i , we ﬁrst use GATv2 (Brody et al., 2021) to aggregate the features from its connected bonds and obtain ¯ h
i , and then update v i based on ¯ h ( l ) i , h ( l − 1) i and U ( l − 1) .
(3) Update global features : The global feature is updated as follows:
[Equation]
(4) Output the prediction : After obtaining the new features of the l -th block, for any atom i ∈ [ | V | ] , it predicts the conformations ˆ R ( l ) i as follows:
[Equation]
An important step in Eqn.(10) is that, after making initial prediction ¯ R ( l ) i , we calculate its center and normalize their coordinates by moving the center to the origin.
We add this kind of normalization to ensure that input coordinates of each block are in reasonable numeric ranges.
We use ˆ R ( L ) output by the last block in ϕ dec as the ﬁnal prediction of the conformation.
3. D ISCUSSIONS
In this section, we discuss the relation between our method and two previous work.
Comparison with CVGAE : Mansimov et al.
(2019) developed an early attempt to generate conformation by raw coordinates, but its performance is not as desired as distance-based methods developed afterwards (Shi et al., 2020; Simm & Hern´andez-Lobato, 2020).
Our method, pursuing the same spirit, makes several ﬁner designs and implementations.
(1) Under the VAE formulation, Mansimov et al.
(2019) also leaned the variance parameter γ of the likelihood model.
i.e., the γ in Eqn.(4).
However, a decent VAE analysis (Dai
Thm. 3) revealed that the VAE objective prefers a vanishing γ .
This is desired in the ideal case when the encoder and decoder are optimized perfectly (Dai & Wipf, 2019, Thm. 4), but in practice it would distract the optimizer towards a “lazy way” to focus on γ and stagnate the optimization of the encoder and decoder.
Therefore, we choose to ﬁx γ .
When γ is ﬁxed, the loss function is equivalent to using γ = 1 and a β parameter as the formulation in Eqn.(6).
We hence choose a small value of β corresponding to ﬁxing γ to a small value in the original formulation (i.e. without β ) as advocated (Dai & Wipf, 2019).
(2) Our model consists of more advanced modules, including GATv2 (Brody et al., 2021) and GN block (Battaglia et al., 2018) to better model the input.
In comparison, Mansimov et al.
(2019) mainly leveraged GRU (Bahdanau et al., 2015) and its variant on graphs, which is outperformed by the modules used in our model.
(3) We iteratively reﬁne the output of each block, while (Mansimov et al., 2019) only outputs the conformation in the last layer.
Comparison with ConfGF : Shi et al.
(2021) used score matching to model the gradients w.r.t interatomic distances, and then recover the coordinates based on the gradients.
That is, (Shi et al., 2021) still tries to model the interatomic distances, while we completely abandon modeling the distances.
We propose a new direction for molecular conformation generation.
E XPERIMENT 4. S ETTINGS
Datasets : Following prior works (Xu et al., 2021a; Shi et al., 2021), we use the GEOM-QM9 and
GEOM-Drugs datasets (Axelrod & Gomez-Bombarelli, 2021) for conformation generation.
We verify our method on both small-scale setting and large-scale setting.
First, we use the same datasets provided by Shi et al.
(2021) for fair comparison with prior works.
The training, validation and test sets of the two datasets consist of 200 K, 2 .
5 K and 22408 (for GEOM-QM9)/ 14324 (for GEOMDrugs) molecule-conformation pairs respectively.
After that, we sample larger datasets from the original GEOM to validate the scalability of our method.
We use all data in GEOM-QM9 and 2 .
2 M molecule-conformation pairs for GEOM-Drugs.
The numbers of training, validation and test sets for the larger GEOM-QM9 setting are 1 .
37 M, 165 K and 174 K, and those for larger GEOM-Drugs are 2 M, 100 K and 100 K.
Model conﬁguration : All of ϕ 2D , ϕ 3D and ϕ dec have the same number of blocks.
The dimension d of the features is 256 .
Inspired by the feed-forward layer in Transformer (Vaswani et al., 2017), MLP also consists of two sub-layers, where the ﬁrst one maps the input features from dimensionto hidden states, followed by Batch Normalization and ReLU activation.
Then the hidden states is mapped to 256 again using linear mapping.
More parameters are summarized in Appendix B.
Evaluation : Assume in the test set, the molecule x has N x conformations.
Following Shi et al.
(2020; 2021), for each molecule x in the test set, we generate 2 N x conformations.
Let S g
( G ) and S r ( G ) denote
all generated and groundtruth conformations respectively.
We use coverage score (COV) and matching score (MAT) to evaluate the generation quality.
Mathematically,
[Equation]
(12) where (cid:96) M is deﬁned in Eqn.(1).
A good method should have a high COV score and a low MAT score.
Following (Shi et al., 2021), the δ ’s are set as 0 .
5 and 1 .
25 for QM9 and Drugs, respectively.
Baselines : (1) RDKit, which is a widely used toolkit and generates the conformation based on the force ﬁelds; (2) CVGAE (Mansimov et al., 2019), which is an early attempt to generate raw coordinates; (3) GraphDG (Simm & Hern´andez-Lobato, 2020), a representative distance-based method with VAE; (4) CGCF (Xu et al., 2021a), which is another distance-based method leveraging continuous normalizing ﬂow; (5) ConfVAE (Xu et al., 2021b), an end-to-end framework for molecular conformation generation, which still uses the pairwise distances among atoms as intermediate variables; (6) ConfGF (Shi et al., 2021), which uses score matching to generate the gradients w.r.t.
distances and then recover the conformation.
Note that ConfGF leverages Langevin dynamics, which requires much more additional decoding time.
4. R ESULTS
The results on the small-scale and large-scale datasets are in Table 1 and Table 2 respectively.
[Table]
[Table]
We have the following observations:
(1) On the four settings in Table 1 and Table 2, our method achieves state-of-the-art results on most of them, i.e., the two large-scale settings and small dataset setting on GEOM-drug.
The only exception is that our method is not so good as ConfGF in terms of the COV mean and MAT mean on small-scale QM9, and the mean COV score on large-scale QM9.
The median COV(%) being 100% means that for more than half of the reference molecules, these exist generated molecules that are close to them within a predeﬁned threshold.
These results already showed the effectiveness and scalability of our method.
(2) Our method achieves more improvement on molecules with more heavy atoms.
Take the results in Table as an example.
On average, QM9 and Drugs have 8 . 8 and 24 .
9 heavy atoms respectively.
In terms of MAT median,
on QM , our method improves ConfGF by 0 .
026 points, while on Drugs, the improvement is 0 . 103 .
To verify our conjecture, on Drugs, we categorize the molecules based on their numbers of heavy atoms.
The number of heavy atoms in the i -th group lie in [10 i +1 , 10( i +1)] .
We compare our method against ConfGF and GraphDG.
The results are in Table 3.
[Table]
The i indicates that the heavy atom number lies in range [10 i + 1 , 10( i + 1)] , i ∈ { 1 , 2 , 3 } . [Figure]
We can observe that among the three algorithms, ConfGF achieves the best results on the molecules whose numbers of heavy atoms lie in [11 , 20] , and our method is slightly worse than it.
However, when the numbers of heavy atoms are larger than 21 , our method signiﬁcantly outperforms ConfGF and GraphDG.
This is consistent with our conjecture, and demonstrates the potential of directly generating coordinates for molecules with more heavy atoms. More discussions are in Appendix C.5.
(3) Our method is much more sample efﬁcient than ConfGF.
In (Shi et al., 2021), generating a conformation requires 5000 sequential forward steps.
With our method, we only need to sample one variable from N (0 , I ) and then generate a corresponding conformation by forwarding once.
For a fair comparison, following the ofﬁcial implementation of ConfGF, we split the test sets of smallscale GEOM-QM9 and GEOM-Drugs into 200 batches.
ConfGF requires 8511 .
60 and 11830 .seconds to decode QM9 and Drugs test sets, while our method only requires 32 . 68 and 54 .
89 seconds respectively.
That is, our method speeds up the decoding more than 200 times.
In Figure 3, we visualize the conformation of different methods.
We randomly select four molecules from the small-scale GEOM-drug dataset, generate several conformations, and visualize the bestaligned ones with the groundtruth.
We can see that our method can generate high-quality conformations than previous methods, which are the most similar to the groundtruth.
4. P ROPERTY PREDICTION
In addition to conformation generation task, we also conduct experiments on property prediction task, which is to predict molecular property based on an ensemble of generated conformation (Axelrod & Gomez-Bombarelli, 2021).
We ﬁrst randomly choose 30 molecules from QM9 and Drugs test sets, and then sample 50 conformations for each molecule using RDkit, ConfGF and our method as initial coordinates.
Then we will use MMFF function in RDkit to optimize the generated coordinates and quantum chemical calculation package Psi4 (Smith et al., 2020) to calculate the energy for each conformation.
We also calculate energy for each groundtruth conformation of these molecules in the same way, as the groundtruth energy.
Finally, we use the mean absolute error (MAE) of the average energy (denoted as “Average E”) and the lowest energy (denoted as “Lowest E”) among these conformations to measure the accuracy of each method on ensemble property prediction task.
[Table]
[Table]
The results are shown in Table 4.
On QM9, our method achieves slightly better results than ConfGF, while on Drugs, the results are signiﬁcantly better than ConfGF and GraphDG.
This shows that our method can provide good conformations for property prediction tasks.
4. A BLATION STUDY
To further verify our algorithm, we conduct the following ablation study: (1) We remove the attentive node aggregation; instead, it is replaced by a simple MLP network.
That is, Eqn.(8) is replaced by
[Equation]
(2) We remove the normalization step in Eqn.(10), i.e., the m ( l ) is not used; (3) We further use the force ﬁeld in RDKit to reﬁne the output of our method.
The results are summarized in Table 5.
We can see that: (1) Without attentively aggregating the atom features, the performances drops: The mean COV drops 5 .
6 points and MAT score increases 0 .
06 points.
(2) Without the conformation normalization, the performance is greatly effected, which shows the importance of this step; (3) Our method can still be reﬁned by the RDKit, which shows that our method can be combined with the classic methods for more improvements.
Finally, we compute the COV and MAT scores of ˆ R ( l ) against the groundtruth, which is the output conformation of the l -th block in the decoder.
ˆ R (0) is the output of ϕ 2D .
The results are shown in Figure 4.
We can see that iteratively reﬁning the conformations can improve the performances, which shows the effectiveness of our design.
[Figure]
C ONCLUSIONS AND FUTURE WORK
In this work, we propose a new method, that directly generates the coordinates of conformations.
For this purpose, we leverage a loss function that is invariant to roto-translation, and design a new model with many advanced modules (i.e., GATv2, GN block) that can iteratively reﬁne the conformations.
Experimental results on both small-scale and large-scale GEOM-QM9 and GEOM-Drugs demonstrate the effectiveness of our method.
For future work, there are many interesting directions.
First, we will combine with more generative methods like ﬂow-based model, Langevin dynamics and so on.
Second, during case analyzing, we found that there are still some cases that current methods cannot successfully handle (e.g., rotatable rings) and we will improve them in the future (see Appendix C.3).
Third, current methods are mainly non-autoregressive, where all coordinates are generated simultaneously.
We will study the autoregressive setting so as to further improve the accuracy.
Fourth, the current loss function is not invariant to permutations of symmetric atoms, and we will improve it in the future.
A D ETAILS OF ϕ 2D AND ϕ 3D
The model architectures of ϕ 2D and ϕ 3D are similar to ϕ dec , with the following differences.
Comparing ϕ 2D with ϕ dec , the differences are the initial conformation ˆ R (0) and initial features (i.e., the H (0) V , H (0) E and U (0) ).
ϕ 2D takes a random conformation sampled from uniform distribution in [ − 1 , 1] as input.
The initial atom and edge features are the embeddings of the atoms and edges respectively.
ϕ 2D will also output a prediction of the conformation.
Note that the random variable z sampled from Gaussian N ( µ R,G , Σ R,G ) is not used in ϕ 2D .
Comparing ϕ 3D with ϕ dec , the differences are the initial conformation ˆ R (0) , initial features (i.e., the
H (0) V , H (0) E and U (0) ) too.
ϕ 3D takes the groundtruth conformation as input.
The initial atom and edge features are the embeddings of the atoms and edges respectively.
Another difference is that the fourth step of ϕ dec , i.e., Eqn.(10), is not used.
More details are available in our code.
B M ORE EXPERIMENT DETAILS
We use AdamW optimizer (Loshchilov & Hutter, 2019) with initial learning rate η 0 = 2 × 10 −and weight decay 0 .
In the ﬁrst 4000 iterations, the learning rate is linearly increased from 10 −to 2 × 10 − 4 .
After that, we use cosine learning rate scheduler (Loshchilov & Hutter, 2016), where the learning rate at the t -th iteration is η 0 (1 + cos( π tT )) / 2 , where T is the half of the period (i.e., the iteration numbers of 10 epochs in our setting).
Similarly, we also use the cosine scheduler to dynamically set the β at range [0 . 0001 , 0 . 008] .
The batch size is ﬁxed as 128 .
All models are trained for 100 epochs.
For the two small-scale settings, the experiments are conducted on a single
V100 GPU.
For the two large-scale settings, we use two V100 GPUs for experiments.
The detailed hyper-parameters are described in Table 6.
[Table]
C M ORE EXPERIMENTAL RESULTS C. C OMBINATION WITH DISTANCE - BASED AND ANGLE - BASED LOSS FUNCTIONS
In addition to the matching loss deﬁned in Eqn.(1) which is related to coordinates only, one may be curious about whether using distance-based loss and angle-based can further improve the performance, since the latter two are equivariant to the transformation of coordinates.
For ease of reference, let R i denote the groundtruth coordinate of atom v i and ˆ R
i denote the predicted coordinate of atom v i .
Remind in Section 2, we use E to denote the collection of all bonds.
We also deﬁne
E 2 as { ( i, j, k ) | ( i, j ) ∈ E, ( i, k )
∈
E, k (cid:54) = j } .
Following (Winter et al., 2021)
we use the following two functions:
[Equation]
where cosine( a, b ) =
a
[Equation]
b
[Equation]
a (cid:107)(cid:107) b (cid:107) and distance( a, b ) = (cid:107) a − b (cid:107) , a and b are two vectors.
That is, we apply additional constraints to bond length and bond angles.
Please note that with the above two auxiliary loss functions, our method still generates coordinates directly and does not need to generate intermediate distances and angles.
We set the weight of the above two loss functions as 0 . 1 , and add them to Eqn.(6).
We conduct experiments on small-scale QM9 and Drugs.
The results are reported in Table 7, which are in the row starting with “Ours + Aux”.
We can see after using those two auxiliary loss functions which are invariant to the roto-translational operations on the coordinates, the results can be further improved.
We will further explore this direction in the future.
[Table]
C. P ROPERTY P REDICTION
Following Shi et al.
(2021), besides energy, we also calculate HOMO and LUMO of all the conformation using Psi4.
Then we will calculate ensemble property average energy E , lowest energy
E min , average HOMO-LUMO gap ∆ (cid:15) , minimum gap ∆ (cid:15) min and maximum gap ∆ (cid:15) max for each method and use mean absolute error (MAE) and median absolute error with ground-truth conformations to measure the accuracy.
The results are showed in Table 8 and Table 9.
We also show the result of each method without MMFF reﬁnement in Table 10.
As the results shown, our method achieves the best accuracy in general.
[Table]
[Table]
[Table]
C. F AILURE CASES
We notice some failure cases in our model prediction.
Two examples are shown in Figure 5.
We can see that there exists a rotatable ring at the end of a molecule, where the ring is symmetric to the bond connecting itself to the rest of the molecule.
Our method fails to generate the coordinates of such rings, but simply puts them in a line.
This is because the model is trapped into local optimal.
By using the additional loss functions as we introduced Appendix C.1, we can successfully recover the conformations of those rings (see the row Ours + Aux).
We will keep exploring along this direction.
[Figure]
C. S TUDY OF MODEL PARAMETERS
In this section, we compare the performances of our method and ConfGF.
By default, our model has 13.29M parameters, and the ConfGF model has 0.81M parameters.
We explore the following two settings: (1) we reduce our model size to 0.98M by reducing hidden dimension sizes and layer number; (2) we increase the model size of ConfGF to 12.28M by increasing hidden dimension sizes, and re-run the experiments.
The results are shown in Table 11.
We can see that: (i) when we reduce our model size to 0 .
98 M, our method still outperforms ConfGF on GEOM-Drugs; on GEOM-QM9, our method outperforms ConfGF evaluated by the median values of COV and MAT.
(ii)
The performances of our model with 13.29M parameters and 0.98M
[Table]
parameters are similar.
(iii) When we increase the model size of ConfGF, the performance becomes worse, which shows that ConfGF cannot beneﬁt from more parameters.
We observe that a larger
ConfGF (12.28M) suffers from larger training loss than a smaller ConfGF (0.98M), which shows its limitation.
C. M ORE DISCUSSIONS ON THE CONFORMATION WITH MORE HEAVY ATOMS
In Table
[Table]
we observe that our method works better than distance-based methods on molecules with more heavy atoms.
Our conjecture is that for these distance-based works, they usually extend the molecular graph with 1,2,3-order neighbors, which is sufﬁcient to determine the 3D structure in principle.
For GEOM-QM9 dataset, considering the number of atoms is less than 10, this extended graph is nearly a complete graph and can provide enough signals to reconstruct the 3D structure.
Therefore, these distance-based performances are good on GEOM-QM9 dataset.
For GEOM-Drugs dataset, the numbers of atoms are much more than those in GEOM-QM9.
Although in theory, the distances in a third-order extended graph can reconstruct the 3D structure, practically the signals are still not enough.
Our method does not rely on the interatomic distances, and can achieve good results on large molecules.