Direct Molecular Conformation Generation^††thanks: This work was done when Jinhua Zhu and Yusong Wang were interns at Microsoft Research AI4Science. Correspondence to Yingce Xia and Chang Liu.

Roto-translational and permutation invariance of conformations: Let $R\in\mathbb{R}^{|V|\times 3}$ and $\hat{R}\in\mathbb{R}^{|V|\times 3}$ denote the groundtruth conformation and the generated conformation. The roto-translation and permutation invariant loss is defined as follows:

$$\ell_{\rm RTP}(R,\hat{R})=\min_{\rho;\;\sigma\in\mathcal{S}}\|R-\rho(\sigma(\hat{R}))\|^{2}_{F}.$$

(1)

In Eqn.(1), (i) $\rho$ denotes a roto-translational operation, which means to rotate and translate a conformation rigidly; (ii) $\mathcal{S}$ denotes the collection of the permutation operations on symmetric atoms. For example, in Figure 1, $\mathcal{S}$ contains two elements $\sigma_{1}$ and $\sigma_{2}$, where $\sigma_{1}$ is an identical mapping, i.e. $\sigma_{1}(i)=i$ for any $i\in\{1,2,\cdots,18\}$, and $\sigma_{2}$ is the mapping on symmetric atoms of the pyrimidine: $\sigma_{2}(13)=17,\sigma_{2}(17)=13,\sigma_{2}(14)=16,\sigma_{2}(16)=14$ and $\sigma_{2}(i)=i$ for the remaining atom $i$’s. (iii) $\|A\|_{F}^{2}$ is defined as $\sum_{i,j}|A_{i,j}|^{2}$. In all, Eqn.(1) defines a loss between $R$ and $\hat{R}$ as the minimal achievable distance under any roto-translation operation and any permutation operation of symmetric atoms, hence is invariant to these operations. Eqn.(1) can be solved via quaternions (Karney, 2007; Hamilton, 1840) and graph isomorphism (Meli & Biggin, 2020).

To solve Eqn. (1), the optimization can be decomposed into two sub problems: (S1) $\ell_{\rm RT}=\min_{\rho}\|\rho(\hat{R})-R\|^{2}_{F}$; (S2) $\ell_{\rm P}=\min_{\sigma\in{\mathcal{S}}}\|\sigma(\hat{R})-R\|^{2}_{F}$.

Karney (2007) propose to use quaternions (Hamilton, 1840) to solve (S1). A quaternion $q$ is an extension of complex numbers, $q=q_{0}\bm{\mathsfit{u}}+q_{1}\bm{\mathsfit{i}}+q_{2}\bm{\mathsfit{j}}+q_{3}\bm{\mathsfit{k}}$, where $q_{0},q_{1},q_{2},q_{3}$ are real scalars and $\bm{\mathsfit{u}}$, $\bm{\mathsfit{i}}$, $\bm{\mathsfit{j}}$, $\bm{\mathsfit{k}}$ are orientation vectors. With quaternions, any rotation operation is specified by a $3\times 3$ matrix, where each element in the matrix is the summation/multiplication of $q_{0}$ to $q_{3}$. The solution to (S1) is the minimal eigenvalue of a $4\times 4$ matrix obtained by algebraic operations on $R$ and $\hat{R}$. To stabilize training, we stop gradient back-propagation through $\rho$ (see Appendix B.3 for the ablation study).

To solve (S2), we need to find all elements in ${\mathcal{S}}$, and then enumerate them to get the minimal value. ${\mathcal{S}}$ can be mathematically described as follows: (1) $\forall i\in V$, atom $i$ and atom $\sigma(i)$ have the same label, which is defined as the union of the atom type itself and also the types of all the bonds connected to it¹¹1For example, in Figure 1, the label of atom $11$ is “S-2Single”, and the label of atom $17$ is “N-2Aromatic”.. (2) There exists a bond between atoms $i$ and $j$ if and only if there exists a bond between atoms $\sigma(i)$ and $\sigma(j)$ in the same molecular graph. Therefore, we convert finding ${\mathcal{S}}$ into a graph isomorphism problem on molecular graphs. Inspired by Meli & Biggin (2020), we use the graph_tool toolkit²²2https://graph-tool.skewed.de to find all permutations in ${\mathcal{S}}$. By combining the above two strategies, we are able to solve Eqn.(1). We provide several examples in the online supplementary material to show how our method works.

Hopcroft & Wong (1974) proposed an algorithm whose complexity of testing planar graphs for isomorphism is $O(|E|)$ , where $|E|$ is the number of edges in a graph. A planar graph can be regarded as a type of graph that no edges cross each other (see Wiki for a quick introduction). For the widely used GEOM-QM9 and GEOM-Drugs datasets (Shi et al., 2021; Xu et al., 2022) of conformation generation, all the molecules are planar graphs. We also randomly sample 30M compounds from PubChem, and only 4.5k of them are not planar graphs (0.015%). This shows that although our method needs to test graph isomorphism, the time complexity could still be controlled. In addition, the $|{\mathcal{S}}|$’s of $99.8\%$ molecules in GEOM-Drugs are smaller than $100$ and efficient to enumerate them all in GPU. A limitation is that, when stepping from small molecules to proteins with a long chain, $|{\mathcal{S}}|$ will significantly increase, resulting in large computation cost of obtaining $\ell_{\rm RTP}$. We will improve it in the future.

One-to-many mapping of conformations: A molecule might correspond to multiple conformations. Thus, we introduce a random variable $z$ to our model for diverse conformation generation. Given a molecular graph $G$, different $z$ could result in different conformations (denoted as $\hat{R}(z,G)$). Inspired by the variational auto-encoder (VAE) (Kingma & Welling, 2014; Rezende et al., 2014; Sohn et al., 2015), we introduce a (conditional) inference model $q(z|R,G)$ to describe the posterior distribution of $z$, reform the reconstruction loss in a probabilistic style $\mathbb{E}_{q(z|R,G)}\left[\ell_{\rm RTP}(R,\hat{R}(z,G))\right]$, and append a regularization term in the form of the Kullback-Leibler (KL) divergence w.r.t. a prior distribution $p(z)$, i.e. $D_{\mathrm{KL}}(q(z|R,G)\|p(z))$. In this way, the aggregated (i.e. averaged/marginalized) posterior $\int q(z|R,G)p_{\text{data}}(R)\,\mathrm{d}R$ is driven towards the prior $p(z)$, which in turn allows generating a new conformation from $p_{\text{data}}(R)$ by passing through the decoder with a $p(z)$ sample. It is easy to draw a random variable $z$ from $p(z)$ and encourages diversity.

By properly choosing $q(z|R,G)$, the loss is tractable to optimize. We specify $q(z|R,G):=\mathcal{N}(z|\mu_{R,G},\Sigma_{R,G})$ as a multivariate Gaussian with a diagonal covariance matrix, where the $\mu_{R,G}$ and $\Sigma_{R,G}$ are outputs from an encoder. It enables tractable loss optimization via reparameterization (Kingma & Welling, 2014): $z\sim q(z|R,G)$ is equivalent to $z^{(i)}=\mu_{R,G}^{(i)}+\sqrt{\Sigma_{R,G}^{(i,i)}}\epsilon,\forall i$, where $\epsilon\sim\mathcal{N}(0,1)$, $z^{(i)}$ and $\mu^{(i)}_{R,G}$ are the $i$-th element of $z$ and $\mu_{R,G}$, and $\Sigma^{(i,i)}_{R,G}$ denotes the $i$-th diagonal element. The KL divergence loss is specialized as $D_{\mathrm{KL}}(\mathcal{N}(\mu_{R,G},\Sigma_{R,G})\|\mathcal{N}(0,{\bm{I}}))$, which has closed form solution.

Overall, the overall training objective function is defined as follows:

$\displaystyle\min\;\;\mathbb{E}_{\epsilon\sim\mathcal{N}(0,{\bm{I}})}\ell_{\rm RTP}(R,\hat{R}(z,G))+\beta D_{\mathrm{KL}}(\mathcal{N}(\mu_{R,G},\Sigma_{R,G})\|\mathcal{N}(0,{\bm{I}})),$

(2)

where $\beta>0$ is a hyperparameter. The minimization in Eqn.(2) is taken over all the network parameters (including the conformation generator and auxiliary model $q$).

Now we show the training and inference workflow. The training process involves three modules, $\varphi_{\rm 2D}$, $\varphi_{\rm 3D}$ and $\varphi_{\rm dec}$. The workflow is illustrated in Figure 2(a). Specifically,

(1) The encoder $\varphi_{\textrm{2D}}$ takes the molecular graph $G$ as its input, and outputs several representations: $H_{V}^{(0)}\in\mathbb{R}^{|V|\times d}$ for all atoms, $H_{E}^{(0)}\in\mathbb{R}^{|E|\times d}$ for all bonds, a global graph feature $U^{(0)}\in\mathbb{R}^{d}$, and initial conformation $\hat{R}^{(0)}\in\mathbb{R}^{|V|\times 3}$. Note $d$ is the dimension of the representations. Formally, $(H_{V}^{(0)},H_{E}^{(0)},U^{(0)},\hat{R}^{(0)})=\varphi_{\textrm{2D}}(G)$.

(2) The encoder $\varphi_{\textrm{3D}}$ extracts features of the conformation $R$ for constructing the conditional inference module $q(z|R,G)$. According to the above specification, $\varphi_{\rm 3D}$ only needs to output the mean and covariance of the Gaussian, or formally, $(\mu_{R,G},\Sigma_{R,G})=\varphi_{\textrm{3D}}(R,G)$.

(3) We randomly sample a variable $z$ from the Gaussian distribution $\mathcal{N}(\mu_{R,G},\Sigma_{R,G})$, and then feed $H_{V}^{(0)}$, $H_{E}^{(0)}$, $U^{(0)}$, $\hat{R}^{(0)}$, $z$ into the decoder $\varphi_{\textrm{dec}}$ to obtain the conformation $\hat{R}(z,G)$. That is, $\hat{R}(z,G)=\varphi_{\textrm{dec}}(\varphi_{\textrm{2d}}(G),z)=\varphi_{\textrm{dec}}(H_{V}^{(0)},H_{E}^{(0)},U^{(0)},\hat{R}^{(0)},z)$. Note that sampling $z\sim{\mathcal{N}}(\mu_{R,G},\Sigma_{R,G})$ is equivalent to sampling $\epsilon\sim{\mathcal{N}}(0,{\bm{I}})$ and then setting $z^{(i)}=\mu_{R,G}^{(i)}+\sqrt{\Sigma_{R,G}^{(i,i)}}\epsilon$.

(4) After obtaining $\hat{R}(z,G)$ and ${\mathcal{N}}(\mu_{R,G},\Sigma_{R,G})$, we optimize Eqn.(2) for training. Recall that $\hat{R}(z,G)$ is related to $\varphi_{\textrm{2D}}$, $\varphi_{\textrm{3D}}$, $\varphi_{\textrm{dec}}$, and $\mu_{R,G}$, $\Sigma_{R,G}$ are related to $\varphi_{\textrm{3D}}$.

The inference workflow is shown in Figure 2(b), where the well-trained $\varphi_{\rm 2D}$ and $\varphi_{\rm dec}$ are leveraged: (1) Given a molecular graph $G$, we use $\varphi_{\textrm{2D}}$ to encode $G$ and obtain $\hat{R}^{(0)}$, $H_{V}^{(0)}$, $H_{E}^{(0)}$, $U^{(0)}$; (2) we sample a random variable $z$ from Gaussian $\mathcal{N}(0,{\bm{I}})$; (3) we feed $\hat{R}^{(0)}$, $H_{V}^{(0)}$, $H_{E}^{(0)}$, $U^{(0)}$, $z$ into $\varphi_{\textrm{dec}}$ and obtain the eventual conformation $\hat{R}(z,G)$. Note that $\varphi_{\textrm{3D}}$ is not used in inference phase.

The encoders $\varphi_{\textrm{2D}}$, $\varphi_{\textrm{3D}}$ and the decoder $\varphi_{\textrm{dec}}$ share the same architecture. They all stack $L$ identical blocks. We take the decoder $\varphi_{\textrm{dec}}$ as an example to introduce its $l$-th block, and leave the details of $\varphi_{\textrm{2D}}$ and $\varphi_{\textrm{3D}}$ to Appendix A.1.

Figure 3 shows the architecture of the $l$-th block of $\varphi_{\textrm{dec}}$. Roughly speaking, this block takes the outputs from its preceding block (including the conformation $\hat{R}^{(l-1)}$, atom representations $H_{V}^{(l-1)}$, edge representations $H_{E}^{(l-1)}$ and the global representation $U^{(l-1)}$ of the whole molecule) and outputs refined conformation and representations of atoms, bonds, the whole graph. The process is repeated until the eventual output $\hat{R}^{(L)}$ is obtained. For the input of the first block (i.e., $l=1$), the $H^{(0)}_{V}$, $H^{(0)}_{E}$, $U^{(0)}$ and $\hat{R}^{(0)}$ are the outputs of $\varphi_{\textrm{2D}}$.

For ease of reference, let $h^{(l)}_{i}$ denote the representation of atom $i$ output by the $l$-th block, and $h^{(l)}_{ij}$ the representation of the bond between atom $i$ and $j$. Also, let MLP denote a feed-forward network.

Mathematically, the $l$-th block takes following operations:

(1) Update bond representations: We first incorporate the coordinate information into the representations by

	$\displaystyle\bar{h}^{(l)}_{i}=h^{(l-1)}_{i}+\texttt{MLP}(\hat{R}^{(l-1)}_{i})+z,\,\forall i\in V,$		(3)
	$\displaystyle\bar{h}^{(l)}_{ij}=h^{(l-1)}_{ij}+\texttt{MLP}(\|\hat{R}^{(l-1)}_{i}-\hat{R}^{(l-1)}_{j}\|),\,\forall(i,j)\in E,$

where $z\sim{\mathcal{N}}(\mu_{R,G},\Sigma_{R,G})$. After that, the bond representations are updated as follows: $\forall(i,j)\in E$,

$$h^{(l)}_{ij}=h^{(l-1)}_{ij}+\texttt{MLP}(\bar{h}^{(l-1)}_{i},\bar{h}^{(l-1)}_{j},\bar{h}^{(l-1)}_{ij},U^{(l-1)}).$$

(4)

(2) Update atom representations: for any atom $i\in V$,

		$\displaystyle\tilde{h}_{i}^{(l)}=\sum_{j\in N(i)}\alpha_{j}W_{v}\texttt{concat}(\bar{h}^{(l)}_{ij},\bar{h}_{j}^{(l-1)})\;{\rm where}\;\alpha_{j}\propto\exp(\bm{a}^{\top}\zeta(W_{q}\bar{h}^{(l-1)}_{i}+W_{k}\texttt{concat}(\bar{h}^{(l-1)}_{j},\bar{h}^{l}_{ij})));$		(5)
		$\displaystyle h^{(l)}_{i}=h_{i}^{(l-1)}+\texttt{MLP}\Big{(}\bar{h}^{(l-1)}_{i},\tilde{h}_{i}^{(l)},U^{(l-1)}\Big{)}.$

In Eqn.(5), $\bm{a}$, $W_{q}$, $W_{v}$ and $W_{k}$ are the parameters to be learned, concat$(\cdot,\cdot)$ is the concatenation of two vectors and $\zeta$ is the leaky ReLU activation. For atom $v_{i}$, we first use GATv2 (Brody et al., 2021) to aggregate the representations from its connected bonds to obtain $\tilde{h}_{i}$, and then update $v_{i}$ based on $\tilde{h}_{i}^{(l)}$, $\bar{h}^{(l-1)}_{i}$ and $U^{(l-1)}$.

(4) Update global molecule representation:

$$U^{(l)}=U^{(l-1)}+\texttt{MLP}\Big{(}\frac{1}{|V|}\sum\nolimits_{i=1}^{|V|}h^{(l)}_{i},\frac{1}{|E|}\sum_{i,j}h^{(l)}_{ij},U^{(l-1)}\Big{)}.$$

(6)

(5) Update the conformation: $\forall i\in V$,

$\displaystyle\bar{R}^{(l)}_{i}=\texttt{MLP}(h^{(l)}_{i}),\quad m^{(l)}=\frac{1}{|V|}\sum\nolimits_{j=1}^{|V|}\bar{R}^{(l)}_{j},\quad\hat{R}^{(l)}_{i}=\bar{R}^{(l)}_{i}-m^{(l)}+\hat{R}^{(l-1)}_{i}.$

(7)

An important step in Eqn.(7) is that, after making initial prediction $\bar{R}^{(l)}_{i}$, we calculate its center and normalize their coordinates by moving the center to the origin. This normalization ensures that the coordinates generated by each block are in reasonable numeric ranges.

We use $\hat{R}^{(L)}$ output by the last block in $\varphi_{\textrm{dec}}$ as the final prediction of the conformation.

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. For the small-scale setting, 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 GEOM-Drugs) molecule-conformation pairs respectively. After that, we work on the large-scale setting by sampling 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.37$M, $165$K and $174$K, and those for larger GEOM-Drugs are $2$M, $100$K and $100$K.

Model configuration: All of $\varphi_{\textrm{2D}}$, $\varphi_{\textrm{3D}}$ and $\varphi_{\textrm{dec}}$ have 6 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. Considering that our method outputs a conformation $\hat{R}^{(l)}$ at each block $l$, we also require that each $\hat{R}^{(l)}$ should try to be similar to the groundtruth $R$. Therefore, the $\ell_{\rm RTP}$ is Eqn.(2) is implemented as

$$\ell_{\rm RTP}(\hat{R}^{(L)},R)+\lambda\sum_{l=0}^{L-1}\ell_{\rm RTP}(\hat{R}^{(l)},R),$$

(8)

where $L$ is the number of blocks in the decoder, $\hat{R}^{(0)}$ is the output from $\varphi_{\rm 2D}$, and $\lambda$ is determined according to validation performance. More details are summarized in Appendix A.2.

Evaluation: Assuming 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 $2N_{x}$ conformations. Let $\mathbb{S}_{g}$ and $\mathbb{S}_{r}$ denote all generated and groundtruth conformations respectively. We use coverage score (COV) and matching score (MAT) to evaluate the generation quality. To measure the difference between $R$ and $\hat{R}$, we use the GetBestRMS in the RDKit package and denote the root-mean-square deviation as $\operatorname{RMSD}(R,\hat{R})$. The recall-based coverage and matching scores are defined as follows:

		$\displaystyle\operatorname{COV}(\mathbb{S}_{g},\mathbb{S}_{r})=\frac{1}{\left|\mathbb{S}_{r}\right|}\left|\left\{R\in\mathbb{S}_{r}\mid\operatorname{RMSD}(R,\hat{R})<\delta,\exists\hat{R}\in\mathbb{S}_{g}\right\}\right|;$		(9)
		$\displaystyle\operatorname{MAT}\left(\mathbb{S}_{g},\mathbb{S}_{r}\right)=\frac{1}{\left|\mathbb{S}_{r}\right|}\sum_{R\in\mathbb{S}_{r}}\min_{\hat{R}\in\mathbb{S}_{g}}\operatorname{RMSD}(R,\hat{R}).$

A good method should have a high COV score and a low MAT score. Following (Shi et al., 2021; Xu et al., 2022), the $\delta$’s are set as $0.5$ and $1.25$ for GEOM-QM9 and GEOM-Drugs, respectively. The COV-$\delta$ curves are left in Figure 9 of the appendix. There are also precision-based COV and MAT scores by switching the $\mathbb{S}_{r}$ and $\mathbb{S}_{g}$ in Eqn.(9). We leave the precision-based results in Appendix B.1.

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 & Herná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) and DGSM (Luo et al., 2021b), which uses score matching to generate the gradients w.r.t distances and then recover the conformation; (7) GeoDiff (Xu et al., 2022), which uses diffusion model to generate conformations; (8) GeoMol (Ganea et al., 2021), which predicts local atomic 3D structures and torsion angles. Considering Ganea et al. (2021) use a different data split from previous work, we reproduce their method following the more commonly used data split (Xu et al., 2021a; Shi et al., 2021).

The recall-based results are shown in Table 1. For small-scale datasets, we independently train our models with five different random seeds, and report the mean and standard derivations. We have the following observations:

(1) On the four settings in Table 1, our method achieves state-of-the-art results on all of them. The median COV(%) being 100% means that for more than half of the groundtruth conformations, there exist generated conformations that are close to them within a predefined threshold. These results show the effectiveness and scalability of our method.

(2) For the molecules in GEOM-QM9 and GEOM-Drugs, our method achieves more improvement on molecules with more heavy atoms. Take the small-scale results in Table 1 as an example. On average, GEOM-QM9 and GEOM-Drugs have $8.8$ and $24.9$ heavy atoms respectively. In terms of MAT mean values, on GEOM-QM9, our method improves ConfGF and GeoDiff by $22.7\%$ and $1.2\%$, while on GEOM-Drugs, the improvements are $37.9\%$ and $16.3\%$. The results demonstrate the effectiveness of our method on large molecules.

More analysis is in Appendix B.5.

(3) Our method is much more sample-efficient than methods based on Langevin dynamics like ConfGF, since we can generate IID samples free of the auto-correlation in a Markov chain. ConfGF requires $5000$ sequential forward steps, while we only need to sample once from ${\mathcal{N}}(0,{\bm{I}})$ and forward through the model. For a fair comparison, following the official implementation of ConfGF, we split the test sets of small-scale 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. Our method speeds up the decoding more than $200$ times. Our method is also much more efficient than the recent GeoMol algorithm , which takes $99.34$s and $668.95$s to decode the above two datasets.

(4) As shown in Table 1, the standard derivations of our method are significantly smaller than the gain compared to the previously best results. This shows the effectiveness and robustness of our method. In addition, considering that our method takes a random conformation as input, to test the confidence interval, we run decoding with $10$ different initial conformations. The mean COV and MAT scores on small-scale GEOM-QM9 are $96.24\pm 0.12$ and $0.2079\pm 0.0010$, and those two numbers on small-scale GEOM-Drugs are $96.38\pm 0.19$ and $0.7239\pm 0.0025$. Our method is not sensitive to the choice of initial conformations.

The number of rotatable bonds is an important metric of how flexible a molecule is. We report coverage score w.r.t. the number of rotatable bonds in Figure 4 based on small-scale GEOM-Drugs. More rotatable bonds indicate harder generation. Our method outperforms previous baselines.

In Figure 5, we visualize the conformation of different methods. We randomly select three molecules from the small-scale GEOM-drug dataset, generate several conformations, and visualize the best-aligned 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.

Computation cost analysis: We use PyTorch profiler³³3https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html to analyze the training time of the following components: (1) model forward time, which denotes the time of calculating the hidden representations from the input layer to output layer; (2) transformation time, which denotes the time of calculating the optimal roto-translation operation $\rho^{*}$; (3) permutation time, which denotes the time of enumerating all possible permutations in $\mathcal{S}$ and find the optimal one $\sigma^{*}\in{\mathcal{S}}$; note that we can use torch.no_grad to reduce time and memory; (4) loss forward time, which is the total of calculating the loss after obtaining $\rho^{*}$ and $\sigma^{*}$; (5) loss backward time, which denotes the time of gradient backpropagation.

The time is summarized in Table 2. We can see that model forward and loss forward/backward takes about $71.4\%$ of the total computation time. The transformation and permutation takes $20.4\%$ and $8.2\%$ of the total time. Note that there are $7$ transformation operations in the experiments (see Eqn.(8)). For the full training pipeline where data loading, model forwarding, loss forwarding, gradient backpropagation, metric calculation and CPU/GPU communications are all considered, DMCG takes $20\%$ more time than that without roto-translation and permutation.

We use graph isomorphism algorithms to find all ${\mathcal{S}}$. Although the general graph isomorphism problem is NP-hard, the size of drug-like molecules is largely limited, otherwise the molecule’s druggability is limited (one can refer to Lipinski’s rule of five). Therefore, our method does not need scalability to a large scale. In our experiments, it takes $4.9$ seconds to process $10k$ molecules in GEOM-QM9, and $6.6$ seconds to process $10k$ molecules in GEOM-Drugs. This is negligible compared to the training time, and we only need to process the data for one time in data preparing stage.

Molecular docking (Roy et al., 2015) is a widely used technique in drug discovery, which aims to find the optimal binding conformation of a drug (i.e., the small molecule) in the pocket of a given target protein and the corresponding binding affinity. In most cases, the molecular docking algorithms treat proteins as rigid bodies and take one conformation of the small molecules as the initial structure inputs. The algorithms then search for the optimal conformation in the conformation space of the small molecules guided by the scoring function. However, due to the complexity of the conformation space, it is difficult for the algorithm to converge to a global minimum. Therefore, the choice of the initial structure often leads to different binding conformations and needs to be taken seriously.

Previously, RDKit was often used to generate initial conformations of small molecules, which usually got reasonable but not optimal results after docking. To verify the effectiveness of our method, we compared the docked poses which take initial conformations generated by our method (DMCG), ConfGF, GeoMol, GeoDiff and RDKit as the initial conformations for docking respectively.

We use Smina (Koes et al., 2013) for molecular docking and make evaluation on PDBbind refined set (Liu et al., 2017) which is a comprehensive collection of experimentally measured binding affinity for all biomolecular complexes deposited in the Protein Data Bank⁴⁴4https://www.rcsb.org/. We randomly select $100$ protein-ligand pairs for evaluation. Appendix A.3 shows detailed optimization hyper-parameters.

Two metrics were used to evaluate the results of docking. One is the docking score (roughly, the estimation of binding affinity), which measures how well a molecule fits the binding site. A smaller value indicates better binding affinity. The other is the root-mean-square deviation (RMSD, the smaller, the better) compared to the crystal complex structure. As shown in Figure 6(a), the distributions for the three methods have the similar shape but our method is much more left-shifted than the others. This shows that for the same small molecule, our method tends to help docking to find conformations with higher binding affinities. Furthermore, docking tends to find lower RMSD binding conformations using the conformation generated by our method as the initial conformation, suggesting that our method can help docking to find binding conformations that are closer to the native crystal structures (Figure 6(b)). We also summarize the mean values of the docking scores and RMSD of different algorithms in the legends of Figure 6. All these results show that our method provides more proper initial conformations for molecular docking and thus facilitates the real application in computer-aided drug discovery.

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 first randomly choose $30$ molecules from GEOM-QM9 test sets, and then sample $50$ conformations for each molecule using RDKit, ConfGF and our method. We use the quantum chemical calculation package Psi4 (Smith et al., 2020) to calculate the energy, HOMO and LUMO for each generated conformation and groundtruth conformation. Next, we calculate the ensemble properties of average energy $\overline{E}$, lowest energy $E_{\rm min}$, average HOMO-LUMO gap $\overline{\Delta\epsilon}$, minimum gap $\Delta\epsilon_{\rm min}$ and maximum gap $\Delta\epsilon_{\rm max}$ based on the conformational properties of each molecule⁵⁵5From a physics perspective, using the Boltzmann-weighted average of the energies of the molecules is a better choice, but the distribution is missing from the dataset. Following (Simm & Hernández-Lobato, 2020; Shi et al., 2021; Luo et al., 2021b), we use the average number here instead of the weighted version.. We use mean absolute error to measure the property differences between the generated conformations and groundtruth conformations.

The results are shown in Table 3. Our method significantly outperforms GraphDG, ConfGF and the recent GeoMol, which shows the effectiveness of our method. We can observe that RDKit achieves the best results on $\Delta\epsilon_{\rm min}$, and we will combine our method with RDKit in the future.

We conduct ablation study on the small-scale GEOM-Drugs dataset. The results are shown in Table 4.

(1) We remove the permutation invariant loss and use the roto-translation invariant loss only, i.e., the $\ell_{\rm RTP}$ in Eqn.(2) is replaced with $\ell_{\rm RT}$ defined in Section.(2.1). The results are denoted as “No $\ell_{\rm P}$” in Table 4.

(2) We replace attentive node aggregation by a simple MLP network. That is, Eqn.(5) is replaced by

$$h^{(l)}_{i}=h^{(l-1)}_{i}+\texttt{MLP}(h^{(l-1)}_{i},U^{(l-1)},\frac{1}{|N(i)|}\sum_{j\in N(i)}h^{(l-1)}_{j}).$$

The results are denoted as “No attention” in Table 4.

(3) We remove the normalization step in Eqn.(7), i.e., the $m^{(l)}$ is not used. Denote the results as “No normalization”.

We can see that: (1) The permutation invariant loss is extremely important, without which the mean COV drops 18.91 while MAT increases 0.3434. We also visualize several cases in Appendix B.2 to compare the results with or without $\ell_{\rm P}$. (2) Without attentively aggregating the atom features, the mean COV drops $1.70$ points and MAT score increases $0.0345$ points. (3) Without the conformation normalization, the performance is also hurt. These results demonstrate the importance of the components in our method.

Finally, we compute the COV and MAT scores of $\hat{R}^{(l)}$ against the groundtruth, which is the output conformation of the $l$-th block in the decoder. $\hat{R}^{(0)}$ is the output of $\varphi_{\textrm{2D}}$. The results are shown in Figure 7. We can see that iteratively refining the conformations can improve the performances, which shows the effectiveness of our design. This phenomenon is consistency with the discovery in machine translation (Xia et al., 2017), image synthesis (Chen & Koltun, 2017) and protein structure prediction (Jumper et al., 2021).

We leave the discussion about additional constraints on loss functions, the comparison of model sizes and more discussions in Appendix B.