GEODIFF: A GEOMETRIC DIFFUSION MODEL FOR MOLECULAR CONFORMATION GENERATION
Year: 2021

Abstract:
Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GEODIFF for molecular conformation prediction. GEODIFF treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be rototranslational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-toend fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GEODIFF is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

heading: 1 INTRODUCTION
text: Graph representation learning has achieved huge success for molecule modeling in various tasks ranging from property prediction (Gilmer et al., 2017; Duvenaud et al., 2015) to molecule generation (Jin et al., 2018; Shi et al., 2020), where typically a molecule is represented as an atom-bond graph. Despite its effectiveness in various applications, a more intrinsic and informative representation for molecules is the 3D geometry, also known as conformation, where atoms are represented as their Cartesian coordinates. The 3D structures determine the biological and physical properties of molecules and hence play a key role in many applications such as computational drug and material design (Thomas et al., 2018; Gebauer et al., 2021; Jing et al., 2021; Batzner et al., 2021). Unfortunately, how to predict stable molecular conformation remains a challenging problem. Traditional methods based on molecular dynamics (MD) or Markov chain Monte Carlo (MCMC) are very computationally expensive, especially for large molecules (Hawkins, 2017).
Recently, significant progress has been made with machine learning approaches, especially with deep generative models. For example, Simm & Hernandez-Lobato (2020); Xu et al. (2021b) studied predicting atomic distances with variational autoencoders (VAEs) (Kingma & Welling, 2013) and flow-based models (Dinh et al., 2017) respectively. Shi et al. (2021) proposed to use denoising score matching (Song & Ermon, 2019; 2020) to estimate the gradient fields over atomic distances, through which the gradient fields over atomic coordinates can be calculated. Ganea et al. (2021) studied generating conformations by predicting both bond lengths and angles. As molecular conformations are roto-translational invariant, these approaches circumvent directly modeling atomic coordinates by leveraging intermediate geometric variables such as atomic distances, bond and torsion angles, which are roto-translational invariant. As a result, they are able to achieve very compelling performance. However, as all these approaches seek to indirectly model the intermediate geometric variables, they have inherent limitations in either training or inference process (see Sec. 2 for a detailed description). Therefore, an ideal solution would still be directly modeling the atomic coordinates and at the same time taking the roto-translational invariance property into account.
In this paper, we propose such a solution called GEODIFF, a principled probabilistic framework based on denoising diffusion models (Sohl-Dickstein et al., 2015). Our approach is inspired by the diffusion process in nonequilibrium thermodynamics (De Groot & Mazur, 2013). We view atoms as particles in a thermodynamic system, which gradually diffuse from the original states to a noisy distribution in contact with a heat bath. At each time step, stochastic noises are added to the atomic positions. Our high-level idea is learning to reverse the diffusion process, which recovers the target geometric distribution from the noisy distribution. In particular, inspired by recent progress of denoising diffusion models on image generation (Ho et al., 2020; Song et al., 2020), we view the noisy geometries at different timesteps as latent variables, and formulate both the forward diffusion and reverse denoising process as Markov chains. Our goal is to learn the transition kernels such that the reverse process can recover realistic conformations from the chaotic positions sampled from a noise distribution. However, extending existing methods to geometric generation is highly non-trivial: a direct application of diffusion models on the conformation generation task lead to poor generation quality. As mentioned above, molecular conformations are roto-translational invariant, i.e., the estimated (conditional) likelihood should be unaffected by translational and rotational transformations (Köhler et al., 2020). To this end, we first theoretically show that a Markov process starting from an roto-translational invariant prior distribution and evolving with roto-translational equivariant Markov kernels can induce an roto-translational invariant density function. We further provide practical parameterization to define a roto-translational invariant prior distribution and a Markov kernel imposing the equivariance constraints. In addition, we derive a weighted variational lower bound of the conditional likelihood of molecular conformations, which also enjoys the rototranslational invariance and can be efficiently optimized.
A unique strength of GEODIFF is that it directly acts on the atomic coordinates and entirely bypasses the usage of intermediate elements for both training and inference. This general formulation enjoys several crucial advantages. First, the model can be naturally trained end-to-end without involving any sophisticated techniques like bilevel programming (Xu et al., 2021b), which benefits from small optimization variances. Besides, instead of solving geometries from bond lengths or angles, the one-stage sampling fashion avoids accumulating any intermediate error, and therefore leads to more accurate predicted structures. Moreover, GEODIFF enjoys a high model capacity to approximate the complex distribution of conformations. Thus, the model can better estimate the highly multi-modal distribution and generate structures with high quality and diversity.
We conduct comprehensive experiments on multiple benchmarks, including conformation generation and property prediction tasks. Numerical results show that GEODIFF consistently outperforms existing state-of-the-art machine learning approaches, and by a large margin on the more challenging large molecules. The significantly superior performance demonstrate the high capacity to model the complex distribution of molecular conformations and generate both diverse and accurate molecules.

heading: 2 RELATED WORK
text: Recently, various deep generative models have been proposed for conformation generation. Among them, CVGAE (Mansimov et al., 2019) first proposed a VAE model to directly generate 3D atomic coordinates, which fails to preserve the roto-translation equivariance property of conformations and suffers from poor performance. To address this problem, the majority of subsequent models are based on intermediate geometric elements such as atomic distances and torsion angles. A favorable property of these elements is the roto-translational invariance, (e.g. atomic distances does not change when rotating the molecule), which has been shown to be an important inductive bias for molecular geometry modeling (Köhler et al., 2020). However, such a decomposition suffers from several drawbacks for either training or sampling. For example, GRAPHDG (Simm & HernandezLobato, 2020) and CGCF (Xu et al., 2021a) proposed to predict the interatomic distance matrix by VAE and Flow respectively, and then solve the geometry through the Distance Geometry (DG) technique (Liberti et al., 2014), which searches reasonable coordinates that matches with the predicted distances. CONFVAE further improves this pipeline by designing an end-to-end framework via bilevel optimization (Xu et al., 2021b). However, all these approaches suffer from the accumulated error problem, meaning that the noise in the predicted distances will misguide the coordinate searching process and lead to inaccurate or even erroneous structures. To overcome this problem, CONFGF (Shi et al., 2021) proposed to learn the gradient of the log-likelihood w.r.t coordinates. However, in practice the model is still aided by intermediate geometric elements, in that it first estimates the gradient w.r.t
interatomic distances via denoising score matching (DSM) (Song & Ermon, 2019; 2020), and then derives the gradient of coordinates using the chain rule. The problem is, by learning the distance gradient via DSM, the model is fed with perturbed distance matrices, which may violate the triangular inequality or even contain negative values. As a consequence, the model is actually learned over invalid distance matrices but tested with valid ones calculated from coordinates, making it suffer from serious out-of-distribution (Hendrycks & Gimpel, 2016) problem. Most recently, Ganea et al. (2021) proposed a highly systematic (rule-based) pipeline named GEOMOL, which learns to predict a set of geometric quantities (i.e. length and angles) and then reconstruct the local and global structures of the conformation in a sophisticated procedure. Despite its effectiveness, this method can be only adopted to molecules with a specific pattern, which limits it from wide applications. In our experiments, the pattern requirement filtered out nearly one third of molecular graphs, making it incomparable in our setting. Besides, there has also been efforts to use reinforcement learning for conformation search Gogineni et al. (2020). Nevertheless, this method relies on rigid rotor approximation and can only model the torsion angles, and thus fundamentally differs from other approaches.

heading: 3 PRELIMINARIES

heading: 3.1 NOTATIONS AND PROBLEM DEFINITION
text: Notations. In this paper each molecule with n atoms is represented as an undirected graph G = 〈V, E〉, where V = {vi}ni=1 is the set of vertices representing atoms and E = {eij | (i, j) ⊆ |V| × |V|} is the set of edges representing inter-atomic bonds. Each node vi ∈ V describes the atomic attributes, e.g., the element type. Each edge eij ∈ E describes the corresponding connection between vi and vj , and is labeled with its chemical type. In addition, we also assign the unconnected edges with a virtual type. For the geometry, each atom in V is embedded by a coordinate vector c ∈ R3 into the 3-dimensional space, and the full set of positions (i.e., the conformation) can be represented as a matrix C = [c1, c2, · · · , cn] ∈ Rn×3. Problem Definition. The task of molecular conformation generation is a conditional generative problem, where we are interested in generating stable conformations for a provided graph G. Given multiple graphs G, and for each G given its conformations C as i.i.d samples from an underlying Boltzmann distribution (Noé et al., 2019), our goal is learning a generative model pθ(C|G), which is easy to draw samples from, to approximate the Boltzmann function.

heading: 3.2 EQUIVARIANCE
text: Equivariance is ubiquitous in machine learning for atomic systems, e.g., the vectors of atomic dipoles or forces should rotate accordingly w.r.t. the conformation coordinates (Thomas et al., 2018; Weiler et al., 2018; Fuchs et al., 2020; Miller et al., 2020; Simm et al., 2021; Batzner et al., 2021). It has shown effectiveness to integrate such inductive bias into model parameterization for modeling 3D geometry, which is critical for the generalization capacity (Köhler et al., 2020; Satorras et al., 2021a). Formally, a function F : X → Y is equivariant w.r.t a group G if:
F ◦ Tg(x) = Sg ◦ F(x), (1) where Tg and Sg are transformations for an element g ∈ G, acting on the vector spaces X and Y , respectively. In this work, we consider the SE(3) group, i.e., the group of rotation, translation in 3D space. This requires the estimated likelihood unaffected with translational and rotational transformations, and we will elaborate on how our method satisfy this property in Sec. 4.

heading: 4 GEODIFF METHOD
text: In this section, we elaborate on the proposed equivariant diffusion framework. We first present a high level description of our 3D diffusion formulation in Sec. 4.1, based on recent progress of denoising diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020). Then we emphasize several non-trivial challenges of building diffusion models for geometry generation scenario, and show how we technically tackle these issues. Specifically, in Sec. 4.2, we present how we parameterize pθ(C|G) so that the conditional likelihood is roto-translational invariant, and in Sec. 4.3, we introduce our surgery of the training objective to make the optimization also invariant of translation and rotation. Finally, we briefly show how to draw samples from our model in Sec. 4.4.

heading: 4.1 FORMULATION
text: Let C0 denotes the ground truth conformations and let Ct for t = 1, · · · , T be a sequence of latent variables with the same dimension, where t is the index for diffusion steps. Then a diffusion probabilistic model (Sohl-Dickstein et al., 2015) can be described as a latent variable model with two processes: the forward diffusion process, and the reverse generative process. Intuitively, the diffusion process progressively injects small noises to the data C0, while the generative process learns to revert the diffusion process by gradually eliminating the noise to recover the ground truth. We provide a high-level schematic of the processes in Fig. 1.
Diffusion process. Following the physical insight, we model the particles C as an evolving thermodynamic system. With time going by, the equilibrium conformation C0 will gradually diffuse to the next chaotic states Ct, and finally converge into a white noise distribution after T iterations. Different from typical latent variable models, in diffusion model this forward process is defined as a fixed (rather than trainable) posterior distribution q(C1:T |C0). Specifically, we define it as a Markov chain according to a fixed variance schedule β1, . . . , βT :
q(C1:T |C0) = T∏ t=1 q(Ct|Ct−1), q(Ct|Ct−1) = N (Ct; √ 1− βtCt−1, βtI). (2)
Note that, in this work we do not impose specific (invariance) requirement upon the diffusion process, as long as it can efficiently draw noisy samples for training the generative process pθ(C0).
Let αt = 1 − βt and ᾱt = ∏t s=1 αs, a special property of the forward process is that q(Ct|C0) of
arbitrary timestep t can be calculated in closed form q(Ct|C0) = N (Ct; √ ᾱtC0, (1− ᾱt)I)1. This indicates with sufficiently large T , the whole forward process will convert C0 to whitened isotropic Gaussian, and thus it is natural to set p(CT ) as a standard Gaussian distribution. Reverse Process. Our goal is learning to recover conformations C0 from the white noise CT , given specified molecular graphs G. We consider this generative procedure as a reverse dynamics of the above diffusion process, starting from the noisy particles CT ∼ p(CT ). We formulate this reverse dynamics as a conditional Markov chain with learnable transitions:
pθ(C0:T−1|G, CT ) = T∏ t=1 pθ(Ct−1|G, Ct), pθ(Ct−1|G, Ct) = N (Ct−1;µθ(G, Ct, t), σ2t I). (3)
Herein µθ are parameterized neural networks to estimate the means, and σt can be any user-defined variance. The initial distribution p(CT ) is set as a standard Gaussian. Given a graph G, its 3D structure is generated by first drawing chaotic particles CT from p(CT ), and then iteratively refined through the reverse Markov kernels pθ(Ct−1|G, Ct). Having formulated the reverse dynamics, the marginal likelihood can be calculated by pθ(C0|G) =∫ p(CT )pθ(C0:T−1|G, CT )dC1:T . Herein a non-trivial problem is that the likelihood should be invariant w.r.t translation and rotation, which has proved to be a critical inductive bias for 3D object generation (Köhler et al., 2020; Satorras et al., 2021a). In the following subsections, we will elaborate on how we parameterize the Markov kernels pθ(Ct−1|G, Ct) to achieve this desired property, and also how to maximize this likelihood by taking the invariance into account.
1Detailed derivations are provided in the Appendix A.

heading: 4.2 EQUIVARIANT REVERSE GENERATIVE PROCESS
text: Instead of directly leveraging existing methods, we consider building the density pθ(C0) that is invariant to rotation and translation transformations. Intuitively, this requires the likelihood to be unaffected by translations and rotations. Formally, let Tg be some roto-translational transformations of a group element g ∈ SE(3), then we have the following statement: Proposition 1. Let p(xT ) be an SE(3)-invariant density function, i.e., p(xT ) = p(Tg(xT )). If Markov transitions p(xt−1|xt) are SE(3)-equivariant, i.e., p(xt−1|xt) = p(Tg(xt−1)|Tg(xt)), then we have that the density pθ(x0) = ∫ p(xT )pθ(x0:T−1|xT )dx1:T is also SE(3)-invariant.
This proposition indicates that the dynamics starting from an invariant standard density along an equivariant Gaussian Markov kernel can result in an invariant density. Now we provide a practical implementation of GEODIFF based on the recent denoising diffusion framework (Ho et al., 2020).
Invariant Initial Density p(CT ). We first introduce the invariant distribution p(CT ), which will also be employed in the equivariant Markov chain. We borrow the idea from Köhler et al. (2020) to consider systems with zero center of mass (CoM), termed CoM-free systems. We define p(CT ) as a “CoM-free standard density” ρ̂(C), built upon an isotropic normal density ρ(C): for evaluating the likelihood ρ̂(C) we can firstly translate C to zero CoM and then calculate ρ(C), and for sampling from ρ̂(C) we can first sample from ρ(C) and then move the CoM to zero. We provide a formal theoretical analysis of ρ̂(C) in Appendix A. Intuitively, the isotropic Gaussian is manifestly invariant to rotations around the zero CoM. And by considering CoM-free system, moving the particles to zero CoM can always ensure the translational invariance. Consequently, ρ̂(C) is constructed as a roto-transitional invariant density.
Equivariant Markov Kernels p(Ct−1|G, Ct). Similar to the prior density, we also consider equipping all intermediate structures Ct as CoM-free systems. Specifically, given mean µθ(G, Ct, t) and variance σt, the likelihood of Ct−1 will be calculated by ρ̂(C t−1−µθ(G,Ct,t) σt
). The CoM-free Gaussian ensures the translation invariance in the Markov kernels. Consequently, to achieve the equivariant property defined in Proposition 1, we focus on the rotation equivariance.
Then in general, the key requirement is to ensure the means µθ(G, Ct, t) to be roto-translation equivariant w.r.t Ct. Following Ho et al. (2020), we consider the following parameterization of µθ:
µθ(Ct, t) = 1 √ αt
( Ct − βt√
1− ᾱt θ(G, Ct, t)
) , (4)
where θ are neural networks with trainable parameters θ. Intuitively, the model θ learns to predict the noise necessary to decorrupt the conformations. This is analogous to the physical force fields (Schütt et al., 2017; Zhang et al., 2018; Hu et al., 2021; Shuaibi et al., 2021), which also gradually push particles towards convergence around the equilibrium states.
Now the problem is transformed to constructing θ to be roto-translational equivariant. We draw inspirations from recent equivariant networks (Thomas et al., 2018; Satorras et al., 2021b) to design an equivariant convolutional layer, named graph field network (GFN). In the l-th layer, GFN takes node embeddings hl ∈ Rn×b (b denotes the feature dimension) and corresponding coordinate embeddings xl ∈ Rn×3 as inputs, and outputs hl+1 and xl+1 as follows:
mij = Φm ( hli,h l j , ‖xli − xlj‖2, eij ; θm ) (5)
hl+1i = Φh ( hli, ∑ j∈N (i) mij ; θh ) (6)
xl+1i = ∑
j∈N (i)
1
dij (ci − cj) Φx (mij ; θx) (7)
where Φ are feed-forward networks and dij denotes interatomic distances. N (i) denotes the neighborhood of ith node, including both connected atoms and other ones within a radius threshold τ , which enables the model to explicitly capture long-range interactions and support molecular graphs with disconnected components. Initial embeddings h0 are combinations of atom and timestep embeddings, and x0 are atomic coordinates. The main difference between proposed GFN and other GNNs lies
in equation 7, where x is updated as a combination of radial directions weighted by Φx : Rb → R. Such vector field xL enjoys the roto-translation equivariance property. Formally, we have:
Proposition 2. Parameterizing θ(G, C, t) as a composition of L GFN layers, and take the xL after L updates as the output. Then the noise vector field θ is SE(3) equivariant w.r.t the 3D system C.
Intuitively, given hl already invariant and xl equivariant, the message embedding m will also be invariant since it only depends on invariant features. Since x is updated with the relative differences ci − cj weighted by invariant features, it will be translation-invariant and rotation-equivariant. Then inductively, composing θ with L GFN layers enables equivariance with Ct. We provide the formal proof of equivariance properties in Appendix A.

heading: 4.3 IMPROVED TRAINING OBJECTIVE
text: Having formulated the generative process and the model parameterization, now we consider the practical training objective for the reverse dynamics. Since directly optimizing the exact log-likelihood is intractable, we instead maximize the usual variational lower bound (ELBO)2:
E [ log pθ(C0|G) ] = E [ logEq(C1:T |C0)
pθ(C0:T |G) q(C1:T |C0) ] ≥ −Eq
[ T∑ t=1 DKL(q(Ct−1|Ct, C0)‖pθ(Ct−1|Ct,G)) ] := −LELBO (8)
where q(Ct−1|Ct, C0) is analytically tractable asN ( √ ᾱt−1βt 1−ᾱt C 0 + √ αt(1−ᾱt−1) 1−ᾱt C t, 1−ᾱt−11−ᾱt βt)
2. Most recently, Ho et al. (2020) showed that under the parameterization in equation 4, the ELBO of the diffusion model can be further simplified by calculating the KL divergences between Gaussians as weighted L2 distances between the means θ and 2. Formally, we have: Proposition 3. (Ho et al., 2020) Under the parameterization in equation 4, we have:
LELBO = T∑ t=1 γtE{C0,G}∼q(C0,G), ∼N (0,I) [ ‖ − θ(G, Ct, t)‖22 ] (9)
where Ct = √ ᾱtC0 + √ 1− ᾱt . The weights γt = βt2αt(1−ᾱt−1) for t > 1, and γ1 = 1 2α1 .
The intuition of this objective is to independently sample chaotic conformations of different timesteps from q(Ct−1|Ct, C0), and use θ to model the noise vector . To yield a better empirical performance, Ho et al. (2020) suggests to set all weights γt as 1, which is in line with the the objectives of recent noise conditional score networks (Song & Ermon, 2019; 2020).
As θ is designed to be equivariant, it is natural to require its supervision signal to be equivariant with Ct. Note that once this is achieved, the ELBO will also become invariant. However, the in the forward diffusion process is not imposed with such equivariance, violating the above properties. Here we propose two approaches to obtain the modified noise vector ̂, which, after replacing in the L2 distance calculation in equation 9, achieves the desired equivariance: Alignment approach. Considering the fact that can be calculated by C t− √ ᾱtC0√
1−ᾱt , we can first rotate
and translate C0 to Ĉ0 by aligning w.r.t Ct, and then compute ̂ as C t− √ ᾱtĈ0√
1−ᾱt . Since the aligned
conformation Ĉ0 is equivariant with Ct, the processed ̂ will also enjoy the equivariance. Specifically, the alignment is implemented by first translating C0 to the same CoM of Ct and then solve the optimal rotation matrix by Kabsch alignment algorithm (Kabsch, 1976).
Chain-rule approach. Another meaningful observation is that by reparameterizing the Gaussian distribution q(Ct|C0) as Ct = √ ᾱtC0 + √ 1− ᾱt , can be viewed as a weighted score function√
1− ᾱt∇Ct q(Ct|C0). Shi et al. (2021) recently shows that generally this score function ∇Ct q(Ct|·) can be designed to be equivariant by decomposing it into ∂Ctdt∇dtq(Ct|·) with the chain rule, where dt can be any invariant features of the structures Ct such as the inter-atomic distances. We refer
2The detailed derivations and full proofs are provided in Appendix A.
readers to Shi et al. (2021) for more details. The insight is that as gradient of invariant variables w.r.t equivariant variables, the partial derivative ∂Ctdt will always be equivalent with Ct. In this work, under the common assumption that d also follows a Gaussian distribution (Kingma & Welling, 2013), our practical implementation is to first approximately calculate ∇dtq(Ct|C0) as d t− √ ᾱtd 0
1−ᾱt , and then
compute the modified noise vector ̂ as √ 1− ᾱt ∂Ctdt(d t− √ ᾱtd 0
1−ᾱt ) = ∂Ctd
t·(dt− √ ᾱtd
0)√ 1−ᾱt .
4.4 SAMPLING
Algorithm 1 Sampling Algorithm of GEODIFF. Input: the molecular graph G, the learned reverse model θ. Output: the molecular conformation C.
1: Sample CT ∼ p(CT ) = N (0, I) 2: for s = T, T − 1, · · · , 1 do 3: Shift Cs to zero CoM 4: Compute µθ(Cs,G, s) from θ(Cs,G, s) using equation 4 5: Sample Cs−1 ∼ N (Cs−1;µθ(Cs,G, s), σ2t I) 6: end for 7: return C0 as C With a learned reverse dynamics θ(G, Ct, t), the transition means µθ(G, Ct, t) can be calculated by equation 4. Thus, given a graph G, its geometry C0 is generated by first sampling chaotic particles CT ∼ p(CT ), and then progressively sample Ct−1 ∼ pθ(Ct−1|G, Ct) for t = T, T − 1, · · · , 1. This process is Markovian, which gradually shifts the previous noisy positions towards equilibrium states. We provide the pseudo code of the whole sampling process in Algorithm 1.

heading: 5 EXPERIMENT
text: In this section, we empirically evaluate GEODIFF on the task of equilibrium conformation generation for both small and drug-like molecules. Following existing work (Shi et al., 2021; Ganea et al., 2021), we test the proposed method as well as the competitive baselines on two standard benchmarks: Conformation Generation (Sec. 5.2) and Property Prediction (Sec. 5.3). We first present the general experiment setups, and then describe task-specific evaluation protocols and discuss the results in each section. The implementation details are provided in Appendix C.

heading: 5.1 EXPERIMENT SETUP
text: Datasets. Following prior works (Xu et al., 2021a;b), we also use the recent GEOM-QM9 (Ramakrishnan et al., 2014) and GEOM-Drugs (Axelrod & Gomez-Bombarelli, 2020) datasets. The former one contains small molecules while the latter one are medium-sized organic compounds. We borrow the data split produced by Shi et al. (2021). For both datasets, the training split consists of 40, 000 molecules with 5 conformations for each, resulting in 200, 000 conformations in total. The valid split share the same size as training split. The test split contains 200 distinct molecules, with 22, 408 conformations for QM9 and 14, 324 ones for Drugs.
Baselines. We compare GEODIFF with 6 recent or established state-of-the-art baselines. For the ML approaches, we test the following models with highest reported performance: CVGAE (Mansimov et al., 2019), GRAPHDG (Simm & Hernandez-Lobato, 2020), CGCF (Xu et al., 2021a), CONFVAE (Xu et al., 2021b) and CONFGF (Shi et al., 2021). We also test the classic RDKIT (Riniker & Landrum, 2015) method, which is arguably the most popular open-source software for conformation generation. Some other recent works (TorsionNet (Gogineni et al., 2020) and GEOMOL (Ganea et al., 2021)) are omitted since these approaches are not directly comparable in our setting. We refer readers to Sec. 2 for a detailed discussion of these models.

heading: 5.2 CONFORMATION GENERATION
text: Evaluation metrics. The task aims to measure both quality and diversity of generated conformations by different models. We follow Ganea et al. (2021) to evaluate 4 metrics built upon root-mean-square deviation (RMSD), which is defined as the normalized Frobenius norm of two atomic coordinates matrices, after alignment by Kabsch algorithm (Kabsch, 1976). Formally, let Sg and Sr denote the sets of generated and reference conformers respectively, then the Coverage and Matching metrics (Xu
et al., 2021a) following the conventional Recall measurement can be defined as:
COV-R(Sg, Sr) = 1
|Sr| ∣∣∣{C ∈ Sr|RMSD(C, Ĉ) ≤ δ, Ĉ ∈ Sg}∣∣∣, (10) MAT-R(Sg, Sr) = 1
|Sr| ∑ C∈Sr min Ĉ∈Sg RMSD(C, Ĉ), (11)
where δ is a pre-defined threshold. The other two metrics COV-P and MAT-P inspired by Precision can be defined similarly but with the generated and reference sets exchanged. In practice, Sg is set as twice of the size of Sr for each molecule. Intuitively, the COV scores measure the percentage of structures in one set covered by another set, where covering means the RMSD between two conformations is within a certain threshold δ. By contrast, the MAT scores measure the average RMSD of conformers in one set with its closest neighbor in another set. In general, higher COV rates or lower MAT score suggest that more realistic conformations are generated. Besides, the Precision metrics depend more on the quality, while the Recall metrics concentrate more on the diversity. Either metrics can be more appealing considering the specific scenario. Following previous works (Xu et al., 2021a; Ganea et al., 2021), δ is set as 0.5Å and 1.25Å for QM9 and Drugs datasets respectively.
Results & discussion. The results are reported in Tab. 1 and Tab. 2. As noted in Sec. 4.3, GEODIFF can be trained with two types of modified ELBO, named alignment and chain-rule approaches. We denote models learned by these two objectives as GEODIFF-A and GEODIFF-C respectively. As shown in the tables, GEODIFF consistently outperform the state-of-the-art ML models on all datasets and metrics, especially by a significant margin for more challenging large molecules (Drugs dataset). The results demonstrate the superior capacity of GEODIFF to model the multi modal distribution, and generative both accurate and diverse conformations. We also notice that in general GEODIFF +C performs slightly better than GEODIFF +A, which suggests that chain-rule approach leads to a better optimization procedure. We thus take GEODIFF +C as the representative in the following comparisons. We visualize samples generated by different models in Fig. 2 to provide a qualitative comparison, where GEODIFF is shown to capture better both local and global structures.
On the more challenging Drugs dataset, we further test RDKIT. As shown in Tab. 3, our observation is in line with previous studies (Shi et al., 2021) that the state-of-the-art ML models (shown in Tab. 1) perform better on COV-R and MAT-R. However, for the new Precision-based metrics we found that ML models are still not comparable. This indicates that ML models tend to explore more possible representatives while RDKIT concentrates on a few most common ones, prioritizes quality over
diversity. Previous works (Mansimov et al., 2019; Xu et al., 2021b) suggest that this is because RDKIT involves an additional empirical force field (FF) (Halgren, 1996) to optimize the structure, and we follow them to also combine GEODIFF with FF to yield a more fair comparison. Results in Tab. 3 demonstrate that GEODIFF +FF can keep the superior diversity (Recall metrics) while also enjoy significantly improved accuracy ((Precision metrics)).
5.3 PROPERTY PREDICTION
Evaluation metrics. This task estimates the molecular ensemble properties (Axelrod & Gomez-Bombarelli, 2020) over a set of generated conformations. This can provide an direct assessment on the quality of generated samples. In specific, we follow (Shi et al., 2021) to extract a split from GEOM-QM9 covering 30 molecules, and generate 50 samples for each. Then we use
the chemical toolkit PSI4 (Smith et al., 2020) to calculate each conformer’s energy E and HOMOLUMO gap , and compare the average energy E, lowest energy Emin, average gap ∆ , minimum gap ∆ min, and maximum gap ∆ max with the ground truth.
Results & discussions. The mean absolute errors (MAE) between calculated properties and the ground truth are reported in Tab. 4. CVGAE is excluded due to the poor performance, which is also reported in Simm & Hernandez-Lobato (2020); Shi et al. (2021). The properties are highly sensitive to geometric structure, and thus the superior performance demonstrate that GEODIFF can consistently predict more accurate conformations across different molecules.

heading: 6 CONCLUSION
text: We propose GEODIFF, a novel probabilistic model for generating molecular conformations. GEODIFF marries denoising diffusion models with geometric representations, where we parameterize the reverse generative dynamics as a Markov chain, and novelly impose roto-translational invariance into the density with equivariant Markov kernels. We derive a tractable invariant objective from the variational lower bound to optimize the likelihood. Comprehensive experiments over multiple tasks demonstrate that GEODIFF is competitive with the existing state-of-the-art models. Future work includes further improving or accelerating the model with other recent progress of diffusion models, and extending our method to other challenging structures such as proteins.

heading: A PROOFS
text: A.1 PROPERTIES OF THE DIFFUSION MODEL
We include proofs for several key properties of the probabilistic diffusion model here to be selfcontained. For more detailed discussions, please refer to Ho et al. (2020). Let {β0, ..., βT } be a sequence of variances, and αt = 1− βt and ᾱt = ∏t s=1 αs. The two following properties are crucial for deriving the final tractable objective in equation 9.
Property 1. Tractable marginal of the forward process: q(Ct|C0) = ∫ q(C1:t|C0) dC1:(t−1) = N (Ct; √ ᾱtC0, (1− ᾱt)I).
Proof. Let i’s be independent standard Gaussian random variables. Then, by definition of the Markov kernels q(Ct|Ct−1) in equation 2, we have
Ct = √αtCt−1 + √ βt t
= √ αtαt−1Ct−2 + √ αtβt−1 t−1 + √ βt t = √ αtαt−1αt−1Ct−3 + √ αtαt−1βt−2 t−2 + √ αtβt−1 t−1 + √ βt t = · · · = √ ᾱtC0 + √ αtαt−1 · · ·α2β1 1 + · · ·+ √ αtβt−1 t−1 + √ βt t
(12)
Therefore q(Ct|C0) is still Gaussian, and the mean of Ct is √ ᾱtC0, and the variance matrix is (αtαt−1 · · ·α2β1 + · · ·+ αtβt−1 + βt)I = (1− ᾱt)I . Then we have:
q(Ct|C0) = N (Ct; √ ᾱtC0, (1− ᾱt)I).
This property provides convenient closed-form evaluation of Ct knowing C0:
Ct = √ ᾱtC0 + √ 1− ᾱt ,
where ∼ N (0, I). Besides, it is worth noting that,
q(CT |C0) = N (CT ; √ ᾱTC0, (1− ᾱT )I),
where ᾱT = ∏T t=1(1− βt) approaches zero with large T , which indicates the diffusion process can finally converge into a whitened noisy distribution.
Property 2. Tractable posterior of the forward process:
q(Ct−1|Ct, C0) = N (Ct−1; √ ᾱt−1βt
1− ᾱt C0 +
√ αt(1− ᾱt−1)
1− ᾱt Ct, (1− ᾱt−1) 1− ᾱt βtI).
Proof. Let β̃t = 1−ᾱt−1 1−ᾱt βt, then we can derive the posterior by Bayes rule:
q(Ct−1|Ct, C0) = q(C t|Ct−1) q(Ct−1|C0)
q(Ct|C0)
= N (Ct;√αtCt−1, βtI) N (Ct−1; √ ᾱt−1C0, (1− ᾱt−1)I)
N (Ct; √ ᾱtC0, (1− ᾱt)I)
= (2πβt) − d2 (2π(1− ᾱt−1))− d 2 (2π(1− ᾱt)) d 2×
exp ( − ‖Ct −√αtCt−1‖2
2βt − ‖C
t−1 −√ᾱt−1C0‖2
2(1− ᾱt−1) + ‖Ct −
√ ᾱtC0‖2
2(1− ᾱt) ) = (2πβ̃t) − d2 exp ( − 1
2β̃t
∥∥∥∥Ct−1 − √ᾱt−1βt1− ᾱt C0 − √ αt(1− ᾱt−1) 1− ᾱt Ct ∥∥∥∥2 )
(13)
Then we have the posterior q(Ct−1|Ct, C0) as the given form.
A.2 PROOF OF PROPOSITION 1
Let Tg be some roto-translational transformations of a group element g ∈ SE(3), and let p(xT ) be a density which is SE(3)-invariant, i.e., p(xT ) = p(Tg(xT )). If the Markov transitions p(xt−1|xt) are SE(3)-equivariant, i.e., p(xt−1|xt) = p(Tg(xt−1)|Tg(xt)), then we have that the density pθ(x0) =∫ p(xT )pθ(x0:T−1|xT )dx1:T is also SE(3)-invariant.
Proof.
pθ(Tg(x0)) = ∫ p(Tg(xT ))pθ(Tg(x0:T−1)|Tg(xT ))dx1:T
= ∫ p(Tg(xT ))Π T t=1pθ(Tg(xt−1)|Tg(xt))dx1:T
= ∫ p(xT )Π T t=1pθ(Tg(xt−1)|Tg(xt))dx1:T (invariant prior p(xT ))
= ∫ p(xT )Π T t=1pθ(xt−1|xt)dx1:T (equivariant kernels p(xt−1|xt))
= ∫ p(xT )pθ(x0:T−1|xT )dx1:T
= pθ(x0)
(14)
A.3 PROOF OF PROPOSITION 2
In this section we prove that the output x of GFN defined in equation 5, 6 and 7 is translationally invariant and rotationally equivariant with the input C. Let g ∈ R3 denote any translation transformations and orthogonal matrices R ∈ R3×3 denote any rotation transformations. let Rx be shorthand for (Rx1, · · · , RxN ). Formally, we aim to prove that the model satisfies:
Rxl+1,hl+1 = GFN(Rxl, RC + g,hl). (15)
This equation indicates that, given xl already rotationally equivalent with C, and hl already invariant, then such property can propagate through a single GFN layer to xl+1 and hl+1.
Proof. Firstly, given that hl already invariant to SE(3) transformations, we have that the messages mij calculated from equation 5 will also be invariant. This is because it sorely relies on the distance between two atoms, which are manifestly invariant to rotations ‖Rxli − Rxlj‖2 = (xli − xlj) >R>R(xli − xlj) = (xli − xlj)>I(xli − xlj) = ‖xli − xlj‖2. Formally, the invariance of messages in equation 5 can be written as:
mi,j = Φm ( hli,h l j , ∥∥Rxli −Rxlj∥∥2 , eij) = Φm (hli,hlj ,∥∥xli − xlj∥∥2 , eij) . (16)
And similarly, the ht+1 updated from equation 6 will also be invariant.
Next, we prove that the vector x updated from equation 7 preserves rotational equivariance and translational invariance. Given mij already invariant as proven above, we have that:∑ j∈N (i) 1 dij (Rci + g −Rcj − g) Φx (mi,j) = R ∑ j∈N (i) 1 dij (ci − cj) Φx (mi,j) = Rxl+1i . (17)
Therefore, we have that rotating and translating c results in the same rotation and no translation on xl+1 by updating through equation 7.
Thus we can conclude that the property defined in equation 15 is satisfied.
Having proved the equivariance property of a single GFN layer, then inductively, we can draw conclusion that a composition of L GFN layers will also preserve the same equivariance.
A.4 PROOF OF PROPOSITION 3
We first derive the variational lower bound (ELBO) objective in equation 8. The ELBO can be calculated as follows:
E log pθ(C0|G) = E logEq(C1:T |C0) [pθ(C0:T−1|G, CT )× p(CT )
q(C1:T |C0) ] ≥ Eq log
pθ(C0:T−1|G, CT )× p(CT ) q(C1:T |C0)
= Eq [ log p(CT )− T∑ t=1 log pθ(Ct−1|G, Ct) q(Ct|Ct−1) ] = Eq [ log p(CT )− log pθ(C
0|G, C1) q(C1|C0) − T∑ t=2 ( log pθ(Ct−1|G, Ct) q(Ct−1|Ct, C0) + log q(Ct−1|C0) q(Ct|C0) )] = Eq [ log
p(CT ) q(CT |C0) − log pθ(C0|G, C1)− T∑ t=2 log pθ(Ct−1|G, Ct) q(Ct−1|Ct, C0) ] = −Eq [ KL ( q(CT |C0)‖p(CT ) ) +
T∑ t=2 KL ( q(Ct−1|Ct, C0)‖pθ(Ct−1|G, Ct) ) − log pθ(C0|G, C1) ] .
(18) It can be noted that the first term KL ( q(CT |C0)‖p(CT ) ) is a constant, which can be omit-
ted in the objective. Furthermore, for brevity, we also merge the final term log pθ(C0|G, C1) into the second term (sum over KL divergences), and finally derive that LELBO =∑T t=1DKL(q(Ct−1|Ct, C0)‖pθ(Ct−1|G, Ct)) as in equation 8.
Now we consider how to compute the KL divergences as the proposition 3. Since both q(Ct−1|Ct, C0) and pθ(Ct−1|G, Ct) are Gaussian share the same covariance matrix β̃tI , the KL divergence between them can be calculated by the squared `2 distance between their means weighed by a certain weights
1 2β̃t
. By the expression of q(Ct|C0), we have the reparameterization that Ct = √ ᾱtC0 + √ 1− ᾱt .
Then we can derive: Eq KL ( q(Ct−1|Ct, C0)‖pθ(G, Ct−1|Ct) ) = 1
2β̃t EC0 ∥∥∥∥√ᾱt−1βt1− ᾱt C0 + √ αt(1− ᾱt−1) 1− ᾱt Ct − 1√ αt ( Ct − βt√ 1− ᾱt θ(Ct,G, t) )∥∥∥∥2 = 1
2β̃t EC0, ∥∥∥∥√ᾱt−1βt1− ᾱt · C t − √ 1− ᾱt √ ᾱt + √ αt(1− ᾱt−1) 1− ᾱt Ct − 1√ αt ( Ct − βt√ 1− ᾱt θ(Ct,G, t) )∥∥∥∥2 = 1
2β̃t · β
2 t
αt(1− ᾱt) EC0, ∥∥0 · Ct + − θ(Ct,G, t)∥∥2 = β2t
2 1−ᾱt−11−ᾱt βtαt(1− ᾱt) EC0, ∥∥ − θ(Ct,G, t)∥∥2 = γtEC0,
∥∥ − θ(Ct, t)∥∥2 , (19)
where γt represent the wights βt2αt(1−ᾱt−1) . And we finish the proof.
A.5 ANALYSIS OF THE INVARIANT DENSITY IN SEC. 4.2
Given a geometric system x ∈ RN ·3, we obtain the CoM-free x̂ by subtracting its CoM. This can be considered as a linear transformation:
x̂ = Qx, where Q = I3 ⊗ ( IN − 1N 1N1 T N ) (20)
where Ik denotes the k × k identity matrix and 1k denotes the k-dimensional vector filled with ones. It can be noted that Q is a symmetric projection operator, i.e., Q2 = Q and QT = Q. And we also
have that rank[Q] = (N − 1) · 3. Furthermore, let U represent the space of CoM-free systems, we can easily have that Qy = y for any y ∈ U since the CoM of y is already zero. Formally, let n = N · 3 and set Rn with an isotropic normal distribution ρ = N (0, In), then the CoM-free density can be formally written as ρ̂ = N (0, QInQT ) = N (0, QQT ). Thus, sampling from ρ̂ can be trivially achieved by sampling from ρ and then projecting with Q. And ρ̂(y) can be calculated by ρ(y) since for any y ∈ U we have ‖y‖22 = ‖Qy‖22, and thus ρ(y) = ρ̂(y). And in this paper, with the SE(3)-equivariant Markov kernels of the reverse process, any CoM-free system will transit to another CoM-free system. And thus we can induce a well-defined Markov chain on the subspace spanned by Q.

heading: B OTHER RELATED WORK
text: Protein structure generation. There has also been many recent works working on protein structure folding. For example, Boltzmann generators Noé et al. (2019) use flow-based models to generate the structure of protein main chains. AlQuraishi (2019) uses recurrent networks to model the amino acid sequences. Ingraham et al. (2019) proposed neural networks to learn an energy simulator to infer the protein structures. Most recently, AlphaFold Senior et al. (2020); Jumper et al. (2021) has significantly improved the performance of protein structure generation. Nevertheless, proteins are mainly linear backbone structures while general molecules are highly branched with various rings, making protein folding approaches unsuitable for our setting.
Point cloud generation. Recently, some other works (Luo & Hu, 2021; Chibane et al., 2020) has also been proposed for 3D structure generation with diffusion-based models, but focus on the point cloud problem. Unfortunately, in general, point clouds are not considered as graphs with various atom and bond information, and equivariance is also not widely considered, making these methods fundamentally different from our model.

heading: C EXPERIMENT DETAILS
text: In this section, we introduce the details of our experiments. As a default setup, the MPNNs for parameterizing the means θ are all implemented with 4 layers, and the hidden embedding dimension is set as 128. For the training of GEODIFF, we train the model on a single Tesla V100 GPU with a learning rate of 0.001 until convergence and Adam (Kingma & Welling, 2013) as the optimizer. The other hyper-parameters of GEODIFF are summarized in Tab. 5, including highest variance level βT , lowest variance level βT , the variance schedule, number of diffusion timesteps T , radius threshold for determining the neighbor of atoms τ , batch size, and number of training iterations. For sampling, the decoding time is in line with the most competitive baseline model CONFGF (Shi et al., 2021) and requires around 8500 and 11500 seconds to decode the entire QM9 and Drugs test sets respectively.

heading: D EXTRA EXPERIMENTS
text: We also test our method with fewer diffusion steps. Specifically, we test the setting with T = 1000, β1 =1e-7 and βT =9e-3. The results on the more challenging Drugs dataset are shown in Tab. 6. Compared with the results in Tab. 1, we can observe that when setting the diffusion steps as 1000, though slightly weaker than the performance with 5000 decoding steps, the model can already outperforms all existing baselines. Note that, the most competitive baseline CONFGF (Shi et al., 2021) also requires 5000 sampling steps, which indicates that our model can achieve better performance with fewer computational costs compared with the state-of-the-art method.

heading: E MORE VISUALIZATIONS
text: We provide more visualization of generated structures in Fig. 3. The molecules are chosen from the test split of GEOM-Drugs dataset.

Review summaries:
The work introduces a novel **GeoDiff** model for conformation generation task, based on a promising diffusion model approach that shows state-of-the-art results in other domains. The authors improved and adapted the concept of the diffusion model for the new domain so that it can work with rotation and translation invariant objects, such as conformations - authors introduced an architecture based on SE(3)-equivariant Graph Field Network (GFN) for parameterizing Markov kernels, also used roto-translational invariant starter density.  The main contributions of the paper are: the authors are the first who propose an architecture based on principal novel generative diffusion framework for molecular conformation generation and explore suitable roto-translational invariant architecture to parameterize the kernel of the method. Experiments show that the proposed model outperforms recent state-of-the-art approaches with a huge gap on **Conformation Generation** and **Property Prediction** tasks.

This paper tackles the problem of conditional generation of molecular conformations (i.e. 3D cartesian atom positions) given a molecular graph. The authors formulate the generation process via diffusion probabilistic models; Conformations are generated by learning a reverse diffusion process from isotropic gaussian noise to molecular conformations. They use a SE(3) invariant formulation of the diffusion process based on Kohler et al 2020, and they operate directly on atomic positions (i.e. a point cloud) instead of interatomic distances or an intermediate bond geometry representation. The authors show state of the art results evaluated by COV/MAT metrics on GEOM-Drugs and GEOM-QM9 datasets.

This paper describes and tests GeoDiff, an end-to-end denoising diffusion model for generating molecular conformations from a molecular graph. The idea is to learn how to denoise a random conformation so as to generate a realistic structure of a small molecule. To do so, a molecular structure is first transformed into a random configuration from an uncorrelated multivariate normal distribution via a diffusion process (forward process). By training a backwards model that inverts the diffusion process, molecular conformations can then be generated from white noise. The reverse process is modeled as a sequence of Markov kernels where each kernel is a Gaussian whose means are neural networks that are trained based on the forward  trajectories. Molecular conformations are directly parameterized in Cartesian coordinates. Therefore, an important property that should be imposed on the reverse process is *equivariance*. Equivariance refers to the following behavior of a mapping under the action of a group: Transforming the input and then mapping it to some output is the same as applying the mapping first and then transforming the output. Here, the relevant group, the special Euclidean group $SE(3)$, is formed by the rigid transformations in $\mathbb R^3$. Equivariance is implemented by drawing the initial conformation from a distribution that is invariant under rigid transformations (isotropic Gaussian that removes the center of mass) and by using equivariant Markov kernels for the reverse process. The equivariance of the Markov kernels is implemented with equivariant convolutional layers (graph field network (GFN)). Two invariant objective functions for training are derived ("alignment approach" and "chain-rule approach"). GeoDiff is tested on two standard benchmarks and shows a superior performance compared to other molecular configuration generators. Also property predictions made by GeoDiff are better than with concurrent methods.