Year: 2021

Abstract:
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 first predict interatomic distances and then generate conformations based on them. This twostage approach has a potential limitation that the predicted distances may conflict 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 refined 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.

heading: 1 INTRODUCTION
text: 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 fields, like UFF (Rappe et al., 1992) or MMFF (Halgren, 1996), to estimate conformations, which is efficient but unacceptably inaccurate (Kanal et al., 2018).
Recently, people start to explore machine learning methods to generate conformation. Simm & Hernández-Lobato (2020), Shi et al. (2020) and Shi et al. (2021) leveraged variational auto-encoder (VAE), flow-based models and score-based methods for conformation generation, respectively. With learned models, molecule conformations can be sampled independently (Simm & HernándezLobato, 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). Specifically, Simm & Hernández-Lobato (2020); Shi et al. (2020); 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 first 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 conflict 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 3N − 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 & Hernández-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 refined 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 fine-grained loss function for training, that is invariant to the roto-translation of the conformation.
(3) We propose a new model that iteratively refines the conformations. Our model is inspired by multiple advanced architectures like GATv2 (Brody et al., 2021), graph network (GN) block (Battaglia et al., 2018), that can effectively model molecules.

heading: 2 PRELIMINARIES
text: In this section, we introduce the notations used in this work, give a formal definition of the molecular conformation generation problem, and briefly 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. Specifically, V = {v1, v2, · · · , v|V |} with the i-th atom vi. Let eij denote the bond between atom vi and vj . LetN(i) denote the neighbors of atom i, i.e.,N(i) = {vj | eij ∈ E}. We useR to represent the conformation ofG, whereR ∈ R|V |×3. The i-th row ofR is the coordinate of atom vi.
Let ρ(·) denote a roto-translation operation, i.e. an affine transformation ρ(R) = Rt + b where the orientation-preserving orthogonal transformation t ∈ SO(3) ⊂ R3×3 represents a rotation and the vector b ∈ R1×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 definition: 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. 1
Matching loss: Let R1 and R2 denote two conformations, both of which are N × 3 matrices (N is the number of rows in R1 and R2). The matching loss between R1 and R2 is defined as follows:
`M (R1, R2) = min ρ ‖ρ(R1)−R2‖2F , (1)
1A molecule corresponds to different conformations at different energy level. To model such cases, the problem is to generate R based on G and a random noise z.
where ‖ · ‖F denotes the Frobenius norm, i.e., ‖A‖2F = ∑ i,j |Aij |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 R1 is obtained by a roto-translational operation of R2.
Karney (2007) proposed to use quaternions to solve Eqn.(1). A quaternion q is an extension of complex numbers, q = q0u + q1i + q2j + q3k, where q0, q1, q2, q3 are scalars and u, i, j, k are orientation vectors. With quaternions, any rotation operation is specified by a 3 × 3 matrix, where each element in the matrix is the basic algebra of q0 to q3 and some constants. The solution to Eqn.(1) is the minimal eigenvalue of a 4× 4 matrix obtained by the basic algebra on R1 and R2, the corresponding optimal rotation and translation ρ∗ can be obtained from the eigenvectors. Note that the best ρ∗ is related to both R1 and R2. To stabilize training, we set the gradients of ρ∗ w.r.t. the model parameter as zero. More details can be found in Karney (2007).

heading: 3 OUR METHOD
text: We introduce our proposed method in this section.

heading: 3.1 CONFORMATION GENERATION FRAMEWORK
text: 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 flexible model architecture (vs. flow-based models), efficient (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) = ∫ p(z)p(R|z,G) dz is however intractable to evaluate and maximize to fit data. VAE handles this by introducing a (conditional) inference model q(z|R,G) which leads to the following identity:
log p(R|G) = Eq(z|R,G) [log p(R|z,G)]−DKL(q(z|R,G)‖p(z)) +DKL(q(z|R,G)‖p(z|R,G)), (2)
where DKL denotes the Kullback-Leibler (KL) divergence, and p(z|R,G) = p(z)p(R|z,G)∫ p(z)p(R|z,G) dz is the true posterior. Due to the nonnegativity of the third term, the first 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:
Eq(z|R,G)[− log p(R|z,G)] +DKL(q(z|R,G)‖p(z)). (3)
Now we develop our specification for the conformation generation task. We construct the likelihood model with regard to a decoded reference conformation R̂(z,G):
p(R|z,G) ∝ exp { − 1 γ `M ( R̂(z,G), R )} , (4)
where γ > 0 is a variance parameter that we choose to fix as 1 (see Sec. 3.4 (1) for explanation). Note that the usage of `M guarantees all roto-translation equivalent conformations have the same probability. For the inference model, we choose
q(z|R,G) = N (z|µR,G,ΣR,G), (5)
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 where ∼ N (0, I). With the above specification, the loss function becomes:
minE ∼N (0,I)`M (R̂(µR,G + ΣR,G ,G), R) + βDKL(N (µR,G,ΣR,G)‖N (0, I)), (6)
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.

heading: 3.2 WORKFLOW
text: 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 workflow of how they work together (see Figure 1). Details of the ϕ··· modules will be covered in the next subsection.
The training workflow 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) ∈ Rd, and initial conformation R̂(0) ∈ R|V |×3. Formally, (H(0)V , H (0) E , U (0), R̂(0)) = ϕ2D(G).
(2) 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 specification, this encoder only needs to output the mean and variance of the Gaussian, or formally, (µR,G,ΣR,G) = ϕ3D(R,G).
(3) We randomly sample a variable z from the Gaussian distributionN (µ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 ∼ N (0, I) and then setting z = µR,G + ΣR,G .
(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 workflow, 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.

heading: 3.3 DETAILED MODEL ARCHITECTURE
text: 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 finer level. (1) The conformation is iteratively refined by each decoding block. Specifically, 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 first 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 ← H (l−1) V + F
(l) + z, where z ∼ N (µR,G,ΣR,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 workflow 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
h (l) ij = h (l−1) ij + MLP(h (l−1) i , h (l−1) j , h (l−1) ij , U (l−1)). (7)
(2) Update atom features: The atom features are updated using an attentive way. For any i ∈ [|V |],
h̄ (l) i = ∑ j∈N (i) αjWvconcat(h (l) ij , h (l−1) j );
αj ∝ exp(a>LeakyReLU(Wqh(l−1)i +Wkconcat(h (l−1) j , h l ij)));
h (l) i = h (l−1) i + MLP ( h (l−1) i , h̄ (l) i , U (l−1) ) .
(8)
In Eqn.(8), a, Wq , Wv and Wk are the parameters to be learned and concat(·, ·) is the concatenation of two vectors. For atom vi, we first use GATv2 (Brody et al., 2021) to aggregate the features from its connected bonds and obtain h̄i, and then update vi based on h̄ (l) i , h (l−1) i and U (l−1).
(3) Update global features: The global feature is updated as follows:
U (l) = U (l−1) + MLP  1 |V | |V |∑ i=1 h (l) i , 1 |E| ∑ i,j h (l) ij , U (l−1)  . (9) (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:
R̄ (l) i = MLP(h (l) i ), m
(l) = 1
|V | |V |∑ j=1 R̄ (l) j , R̂ (l) i = R̄ (l) i −m (l) + R̂ (l−1) i . (10)
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 final prediction of the conformation.

heading: 3.4 DISCUSSIONS
text: 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 & Hernández-Lobato, 2020). Our method, pursuing the same spirit, makes several finer 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 & Wipf, 2019,
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 fix γ. When γ is fixed, 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 fixing γ 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 refine 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.

heading: 4 EXPERIMENT

heading: 4.1 SETTINGS
text: 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 200K, 2.5K 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.2M molecule-conformation pairs for GEOM-Drugs. The numbers of training, validation and test sets for the larger GEOM-QM9 setting are 1.37M, 165K and 174K, and those for larger GEOM-Drugs are 2M, 100K and 100K.
Model configuration: 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 first one maps the input features from dimension 256 to 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 Nx conformations. Following Shi et al. (2020; 2021), for each molecule x in the test set, we generate 2Nx conformations. Let Sg(G) and Sr(G) denote all generated and groundtruth conformations respectively. We use coverage score (COV) and matching score (MAT) to evaluate the generation quality. Mathematically,
COV(Sg(G),Sr(G)) = 1
|Sr| |{R ∈ Sr | `M (R,R′) < δ, ∃R′ ∈ Sg}| ; (11)
MAT (Sg(G),Sr(G)) = 1 |Sr| ∑
R′∈Sr
min R∈Sg
`M (R,R ′) , (12)
where `M is defined 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 fields; (2) CVGAE (Mansimov et al., 2019), which is an early attempt to generate raw coordinates; (3) GraphDG (Simm & Hernández-Lobato, 2020), a representative distance-based method with VAE; (4) CGCF (Xu et al., 2021a), which is another distance-based method leveraging continuous normalizing flow; (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.

heading: 4.2 RESULTS
text: The results on the small-scale and large-scale datasets are in Table 1 and Table 2 respectively.
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 predefined 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 1 as an example. On average, QM9 and Drugs have 8.8 and 24.9 heavy atoms respectively. In terms of MAT median, on QM9, 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 [10i+1, 10(i+1)]. We compare our method against ConfGF and GraphDG. The results are in Table 3.
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 significantly 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 efficient 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 official implementation of ConfGF, we split the test sets of smallscale GEOM-QM9 and GEOM-Drugs into 200 batches. ConfGF requires 8511.60 and 11830.42 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.

heading: 4.3 PROPERTY PREDICTION
text: We conduct experiments on the property prediction task, which is to predict molecular properties based on an ensemble of generated conformations (Axelrod & Gomez-Bombarelli, 2021). We first 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. Next, we use the package Psi4 (Smith et al., 2020) to calculate the energy, HOMO and LUMO for each generated conformation and groundtruth conformation. After that, we calculate ensemble properties including average energy Ē, lowest energy Emin, average HOMO-LUMO gap ∆ , minimum gap ∆ min, and maximum gap ∆ max for the conformations. Finally, we use the mean absolute error to measure the gap of the above properties between the groundtruth and generated conformations. Although a better choice is to use Boltzmann distribution to re-weight each conformation, due to the lack of such the distribution in the dataset, we are not able to use it now and use the average value instead.
The results are shown in Table 4. On QM9, our method achieves slightly better results than ConfGF, while on Drugs, the results are significantly better than ConfGF and GraphDG. This shows that our method can provide good conformations for property prediction tasks.

heading: 4.4 ABLATION STUDY
text: 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
h (l) i = h (l−1) i + MLP(h (l−1) i , U (l−1), 1 |N (i)| ∑
j∈N (i)
h (l−1) j ); (13)
(2) We remove the normalization step in Eqn.(10), i.e., the m(l) is not used; (3) We further use the force field in RDKit to refine 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 refined 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 refining the conformations can improve the performances, which shows the effectiveness of our design.

heading: 5 CONCLUSIONS AND FUTURE WORK
text: 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 refine 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 flow-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.

heading: B MORE EXPERIMENT DETAILS
text: We use AdamW optimizer (Loshchilov & Hutter, 2019) with initial learning rate η0 = 2 × 10−4 and weight decay 0.01. In the first 4000 iterations, the learning rate is linearly increased from 10−6 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 fixed 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.

heading: C MORE EXPERIMENTAL RESULTS
text: C.1 COMBINATION WITH DISTANCE-BASED AND ANGLE-BASED LOSS FUNCTIONS
In addition to the matching loss defined 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 Ri denote the groundtruth coordinate of atom vi and R̂i denote the predicted coordinate of atom vi. Remind in Section 2, we use E to denote the collection of all bonds. We also define E2 as {(i, j, k)|(i, j) ∈ E, (i, k) ∈ E, k 6= j}. Following (Winter et al., 2021) we use the following two functions:
`angle = 1 |E2| ∑
(i,j,k)∈E2,
‖ cosine(Rj −Ri, Rk −Ri)− cosine(R̂j − R̂i, R̂k − R̂i)‖2F , (14)
`bond = 1 |E| ∑
(i,j)∈E
( distance(Rj , Ri)− distance(R̂j − R̂i) )2 , (15)
where cosine(a, b) = a >b
‖a‖‖b‖ and distance(a, b) = ‖a − b‖, 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.
C.2 PROPERTY PREDICTION
Following Shi et al. (2021), besides energy, we also calculate HOMO and LUMO of all the conformation using Psi4. Note that all the generated conformations are refined by MMFF. Then we will calculate ensemble property average energy E, lowest energy Emin, average HOMO-LUMO gap ∆ , minimum gap ∆ min and maximum gap ∆ 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. As the results shown, our method achieves the best accuracy in general.
C.3 FAILURE 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.
C.4 STUDY 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 10.
We can see that: (i) when we reduce our model size to 0.98M, 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 parameters are similar. (iii) When we increase the model size of ConfGF, the performance becomes worse, which shows that ConfGF cannot benefit 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.5 MORE DISCUSSIONS ON THE CONFORMATION WITH MORE HEAVY ATOMS
In Table 3, 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 sufficient 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.

Review summaries:
The authors present a method that is able to accurately generate molecule confirmations, achieving competitive results across four datasets.  Unlike previous methods that often rely on directly predicting various inter-atomic distances (through a distance matrix) (which in some cases results in physically impossible molecules), the presented work predicts the location of all atoms. The loss function is invariant to rotation, an important property of many molecules. 

 The paper presents a direct method for generating a molecular conformation conditional on the molecular graph.  Inspired by the success of AlphaFold2, the authors demonstrated that one could directly write out 3D coordinates for a small molecule rather than use the previous ways of generating intermediate interatomic distance matrices or iterate over forces. The method displays a decent performance on the standard benchmark set of QM9.  It performs better than past baseline methods for larger molecules. 

The paper proposed a new VAE-based generative model for generating molecular conformations from graphs. The whole formulation and training objective are similar to existing VAE-based models like CVGAE, and the main contrition lies in the architecture of the decoder. The architecture borrows the idea from the recent advance of AlphaFold: the proposed decoder is composed of several identical updating blocks, where they iteratively update the node/edge embeddings as well as the coordinates in each block. And the blocks are implemented with advanced graph neural networks.