GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation
Authors: Minkai Xu
Lantao Yu
Yang Song
Chence Shi
Stefano Ermon
Jian Tang
Venue: ICLR 2022 Oral (2022)
Keywords: molecular conformation generation
deep generative models
diffusion probabilistic models

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 roto-translational 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-to-end 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. 

Published as a conference paper at ICLR 2022
GEODIFF: A GEOMETRIC DIFFUSION MODEL FOR MOLECULAR CONFORMATION GENERATION
Minkai Xu! Lantao Yu?, Yang Song", Chence Shi 2 Stefano Ermon J Tang Mila Québec Al Institute, Canada Université de Montréal, Canada Stanford University, USA *HEC Montréal, Canada 'CIFAR Al Research Chair (minkai xu, chence. shij @ @umontreal. ca (lantaoyu, yangsonn ermon | 8cs. stanford. edu Maan.ang@hec.cco
ABSTRACT
Predicting molecular conformations from molecular graphs is fundamental prob- lem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative mod- els. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to noise distribution, in this paper, we propose a novel generative model named GEODIFF for molecular conformation prediction. GEODIFF treats each atom particle and learns directly reverse the diffusion process (i.e., transforming from noise distribution stable conformations) as Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto- translational 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-to- end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments multiple benchmarks show that GEODIFF is superior comparable to existing state-of-the-art approaches, especially on large molecules.
INTRODUCTION
Graph representation learning has achieved huge success for molecule modeling in various tasks rang- ing from property prediction (Gilmer et al., 2017: Duvenaud et al., 2015) to molecule generation (Jin et al., 2018; Shi et al., 2020), where typically molecule is represented as atom-bond graph. Despite its effectiveness in various applications, a more intrinsic and informative representation for molecules is the 3D geometry, also known conformation, where atoms represented as their Cartesian coordinates. The 3D structures determine the biological and physical properties of molecules and hence play key role in many applications such as computational drug and material de- sign (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 al. (2021b) studied predicting atomic distances with variational autoencoders (VAEs) (Kingma & Welling, 2013) and flow-based models (Dinh al., 2017) respectively. Shi 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. Ganca et al. (2021) studied generating conformations by predicting both bond lengths and angles. As molecular conformations are aao-tralslational invariant, these approaches circumvent directly modeling atomic coordinates by leveraging intermediate geometric variables such atomic distances, bond and torsion angles, which
CCode is available https://giii com nkai Xu/Geali ff.

---

Published as a conference paper at ICLR 2022
noto-translationa invariant. As 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 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 al., 2015). Our approach is inspired by the diffision 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 latent variables, and formulate both the forward diffusion and reverse denoising process as Markov chains. Our goal to learn the transition kernels such that the reverse process recover realistic conformations from the chaotic positions sampled from a noise distribution. However, extending existing methods to geometric generation is highly non-trivial: direct application of diffusion models on the conformation generation task lead 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 ooo-translatiionl invariant density function. We further provide practical parameterization to define a -rotsraanllaiinal invariant prior distribution and a Markov kernel imposing the equivariance constraints. In addition, we derive weighted variational lower bound of the conditional likelihood of molecular conformations, which also enjoys the roto- translational 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 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.
RELATED WORK
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 ootottranslation 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. 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 decomposition suffers from several trawbacks for either training or sampling. For example, GRAPHDG (Simm & Hernandez- Lobato, 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
2

---

Published as a conference paper at ICLR 2022
distances. CONEVAE further improves this pipeline by designing 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; Luo et al., 2021) proposed to learn the gradient of the aggllikeihhod w.r.f coordinates. However, in practice the model is still aided by intermediate geometric elements, in that it first estimates the gradient w.F. 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 even contain negative values. As consequence, the model actually learned over invalid distance matrices but tested with valid ones calculated from coordinates, making it suffer from serious out-of out-of-distribution (Hendrycks & Gimpel, 2016) problem. Most recently, another concurrent work (Ganea et al., 2021) proposed a highly systematic (rule-based) pipeline named GEOMOL, which Iearns to predict minimal set of geometric quantities (i.e. length and angles) and then reconstruct the local and global structures of the conformation in sophisticated procedure. 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.
PRELIMINARIES
3.1 NOTATIONS AND PROBLEM DEFINITION
Notations. In this paper each molecule with n atoms is represented as an undirected graph 9 EV, where V [t=) is the set of vertices representing atoms and & = {esy (i.j) |V] [V]] is the set of edges representing inter-atomic bonds, Each node 22 E V describes the atomic attributes, e.g., the element type. Each edge ety E & describes the corresponding connection between o and Uy. and is labeled with its chemical type. In addition, we also assign the unconnected edges with virtual type. For the geometry, each atom in V is embedded by coordinate vector e E IR into the 3-dimensional space, and the full set of positions (i.e., the conformation) can be represented matrix C = [c1, C2, c.l E R=
Problem Definition. The task of molecular conformation generation conditional generative problem, where we are interested in generating stable conformations for a provided graph G. Given multiple graphs G. and for each 9 given its conformations as i.i.d samples from an underlying Boltzmann distribution (Noé et al., 2019), our goal is learning generative model po (ClG)) which is easy to draw samples from, to approximate the Boltzmann function.
3.2 EQUIVARIANCE
Equivariance is ubiquitous in machine learning for atomic systems, e.g., the vectors of atomic dipoles or forces should rotate accordingly w.r.f. 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. function F X' + y is equivariant w.r.f group G if:
F (x) S F(x),
where Ty and Sy are transformations for an element g E G, acting on the vector spaces X' and D, 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 will elaborate on how our method satisfy this property in Sec. 4.
GEODIFF METHOD
In this section, 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 al., 2015; Ho et al., 2020). Then we emphasize several
3

---

Published as a conference paper at ICLR 2022
Po (c.
CT
C
ct-1
CO
q(C4)(-1)
Figure 1: Illustration of the diffusion and reverse process of GEODIFF. For diffusion process, noise from fixed posterior distributions q(C*)(²)1 is gradually added until the conformation is de- stroyed. Symmetrically, for generative process, an initial state C* is sampled from standard Gaussian distribution, and the conformation is progressively refined via the Markov kernels pa [G C.).
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 pa (Clg) that the conditional likelihood is noto-translational invariant, and in Sec. 4.3, we introduce surgery of the training objective to make the optimization also invariant of translation and rotation. Finally, we briefly show how draw samples from our model in Sec. 4.4.
4.1 FORMULATION
Let C denotes the ground truth conformations and let C' for t = =1 T be sequence of latent variables with the same dimension, where is the index for diffusion steps. Then diffusion probabilistic model (Sohl-Dickstein al., 2015) be described processes: the forward diffusionn process, and the reverse generative process. Intuitively, the diffission process progressively injects small noises to the data C while the generative process learns to revert the diffusion process by gradually eliminating the noise to recover the ground truth. We provide high-level schematic of the processes Fig.
latent variable model with two
Diffusion process. Following the physical insight, we model the particles C as an evolving thermo- dynamic system. With time going by, the equilibrium conformation C will gradually diffuse the next chaotic states CS, and finally converge into white noise distribution after I iterations. Different from typical latent variable models, in diffusion model this forward process is defined as a fixed (rather than trainable) posterior distribution (CII:T(C)) Specifically, we define it a Markov chain according to a fixed variance schedule $1....,By:
T a
(CL:TICO)
II ct-+), [=1
((c))--)==NC VI (.0--133T)
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 (C").
Let at = 11- 3e and a; = II Cg., special property the forward process is that qC (*) of arbitrary timestep can calculated closed form ((C)(C)) AC',; va,Co (1- T)) This indicates with sufficiently large T, the whole forward process will convert C to whitened isotropic Gaussian, and thus it is natural to set p(C"") as standard Gaussian distribution.
Reverse Process. Our goal is learning to recover conformations C 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 pCC)). ). We formulate this reverse dynamics as a conditional Markov chain with learnable transitions:
Pe "G = II
Po
(3)
Herein μ4 are parameterized neural networks to estimate the means, and 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
Detailed derivations are provided in the Appendix A.
4

---

Published as a conference paper at ICLR 2022
Having formulated the revense dynamics, the marginal likelihood can be calculated by po (C (pc"))(oCCTT-11 CT " Herein non-trivial problem is that the likelihood should be invariant w.E. 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 pe pg(C C*) to achieve this desired property, and also how to maximize this likelihood by taking the invariance into account.
4.2 EQUIVARIANT REVERSE GENERATIVE PROCESS
Instead of directly leveraging existing methods, we consider building the density pp(C") 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- tro-stranlltional transformations of group element g E SE(3). then we have the following statement:
Proposition 1. Let p(IT) be an SE(3)-invariant density function, i.e., p(TT) = ((T((TT) If Markov transitions p(xx-1|(x)) are -equivvaaan i.e., p(x+-1|24) = p(T2((t111 ( then we have that the density po( = f p(2T) (20 1 is also SE(3)-invariant
This proposition indicates that the dynamics starting from 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(C") "CoM-free standard density" p(C), built upon an isotropic normal density p(C): for evaluating the likelihood ((C) we can firstly translate C' to zero CoM and then calculate p(C), and for sampling from p(C)) we can first sample from p(C') and then move the CoM to zero.
We provide a formal theoretical analysis of P(C) in Appendix A. Intuitively, the isotropic Gaussian 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, p(C) constructed as a roto-transitional invariant density.
Equivariant Markov Kernels p(C*-1 )*). Similar to the prior density, we also consider equip- ping all intermediate structures C' as CoM-free systems. Specifically, given mean [o t) and variance ar the likelihood of Ct-1 will be calculated by p0 -6 e .)). The CoM-free Gaussian ai 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 وو Ct t) to be noto-translation equivariant w.r.t Ct Following Ho et al. (2020), we consider the following parameterization of pg:
(C t) va
1
VI- (0
Be Be
where eg are neural networks with trainable parameters 0. Intuitively, the model eg 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 eg 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 (-th layer, GFN takes node embeddings h' e Bx (b denotes the feature dimension) and corresponding coordinate embeddings x' [gaxb as inputs, and outputs +1+ and x as follows:
my = 4m (h1, n2., ||x - x111 eeys 0m)
(5)
1 = Pn (1/s JEN(()
Σ miyi 0n)
(6)
x Σ dij
1
(c: -c.=
tr. (m.:: 0,)
JEN(()
5

---

Published as a conference paper at ICLR 2022
where @ are foed-forward networks and dı denotes interatomic distances. (1)) denotes the neighbor- hood of 7th node, including both connected atoms and other ones within radius threshold which enables the model to explicitly capture long-range interactions and support molecular graphs with disconnected components. Initial embeddings h are combinations of atom and timestep embeddings, and x" are atomic coordinates. The main difference between proposed GFN and other GNNs lies in equation 7, where x is updated Such vector field x- enjoys the roto-translation equivariance property. Formally, we have:
combination of radial directions weighted by D RR I.
Proposition 2. Parameterizing ea (g, C, 1) as a composition of L GFN layers, and take the x" after L updates as the oulput. Then the noise vector field eo is SE(3) equivariant w.r.t the 3D system C.
Intuitively, given h' already invariant and x' equivariant, the message embedding will also be invariant since it only depends on invariant features. Since x is updated with the relative differences C Cj weighted by invariant features, it will be translation-invariint and rotation-equivariant, Then inductively, composing 6e with I GFN layers enables equivariance with C We provide the formal proof of equivariance properties in Appendix A.
4.3 IMPROVED TRAINING OBJECTIVE
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- likellhoehood intractable, we instead maximize the usual variational lower bound (ELBO)":
Po (C (G)
E (C = E log Eyterreoy
(CI.T(C))
T
Eg Dea (c = LELBO
(8)
where qCC" (O) is analytically tractable as NV 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 KI. divergences between Gaussians as weighted L2 distances between the means co and é². Formally, have:
v 18 Co
- 1-û,
1 B.) Most
1-û.
Proposition 3. (Ho et al., 2020) Under the parameterization in equation 4, we have:
T
LELBO Σ
T g)wglt)))))...o.ooo 9).e~((0.1) ( - (o 11
where Ct = vorCa VI- āje. The weights Nt 2a,(1-a.-1) for and 71
2011
The intuition of this objective is to independently sample chaotic conformations of different timesteps from q(C'-1CC', (C)) and use eg to model the noise vector e. To yield a better empirical performance, Ho et al. (2020) suggests to set all weights 7/ as 1, which is in line with the the objectives of recent noise conditional score networks (Song & Ermon, 2019; 2020).
As ég is designed to be equivariant, it is natural to require its supervision signal e to be equivariant with C. Note that once this is achieved, the ELBO will also become invariant. However, the in the forward diffusson 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 c the La distance calculation in equation achieves the desired equivariance:
Alignment approach. Considering the fact that c can be calculated by Varco
we can first rotate
and translate C to ( by aligning w.r. C³, and then compute as conformation Ç equivariant with C²2 the processed é will also enjoy the equivariance. Specifically, the alignment is implemented by first translating C to the same CoM of Ct and then solve the optimal rotation matrix by Kabsch alignment algorithm (Kabsch, 1976).
VI-ai Since the aligned
"The detailed derivations and full proofs provided in Appendix

---

Published as a conference paper at ICLR 2022
Chain-rule approach. Another meaningful observation is that by reparameterizing the Gaussian distribution q(C Ca) as C = vā,Ca + vI-ā,,, e can be viewed as a weighted score function vI- āe Ve ct (CC Shi et al. (2021) recently shows that generally this score function Ve (CC"))) can be designed to be equivariant by decomposing it into d' Va.g(C().) with the chain rule, where d' can be any invariant features of the structures C such the inter-atomic distances. We refer 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 a d' will always be equivalent with C³. In this work, under the common assumption that also follows a Gaussian distribution (Kingma & Welling, 2013), d' d'-vard" and then our practical implementation is first approximately calculate Va.g(C((C)) as 1-o,
compute the modified noise vector as v1 ãs d d'' -vaad" derd'-(d'- (d'--a,d")
Vi-ai
4.4 SAMPLING
With a learned reverse dynamics Eo(G, Ct. t), the transition means HolG, C., f) can be calculated by equation 4. Thus, given a graph G, its geometry C is generated by first sampling chaotic parti- cles C² ~ P(CT), and then progressively sample Cr-1 Pa(C*= g )*) for = T.T- 1. 1. This process is Marko- vian, which gradually shifts the previous noisy positions towards
Algorithm 1 Sampling Algorithm of GEODIFF.
Input: the molecular graph G. the learned reverse model eg. Output: the molecular conformation C.
1: Sample CH p(C²) NN
2: for 8 11.
do
3: Shift C* to zero CoM
4: Compute [a(C", ,G. from eo (0 "CC, s) using equation 4
5: Sample Ca ~ NNC* ; (C o)
6: end for
7: return Co
equilibrium states. We provide the pseudo code of the whole sampling process in Algorithm 1.
EXPERIMENT
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.
5.1 EXPERIMENT SETUP
Datasets. Following prior works (Xu et al., 2021a;b), we also the recent GEOM-QM9 (Ramakr- ishnan 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 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-arr baselines. For the ML approaches, we test the following models with highest reported performance: CVGAE (Mansimov et al., 2019), GRAPIIDG (Simm& & Hernandez-Lobato, 2020), CGCF (Xu et al., 2021a), CONF- VAE (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. We refer readers to Sec. 2 for a detailed discussion of these models.
5.2 CONFORMATION GENERATION
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 moot-mean-square
7

---

Published as a conference paper at ICLR 2022
Table 1: Results on the GEOM-Drugs dataset, without FF optimization.
COV-R (%)
MAT-R (A)
COV-P (%) t
MAT-P (A) 4
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
CVGAE
0.00
0.00
3,0702
2.9937
GRAPHDG
8.27
0.00
1.9722
1.9845
2.08
0.00
2,4340
2,4100
CGCF
53.96
57.06
1.2487
1.2247
21.68
13.72
1.8571
1,8066
CONFVAE
55.20
59.43
1,2380
1.1417
22.96
14.05
1,8287
1.8159
GEOMOL
67.16
71.71
1.0875
1,0586
CONFGF
62.15
70.93
1,1629
1,1596
23.42
15.52
1,7219
1.6863
GEODIEF-A
88.36
96.09
0,8704
0.8628
60.14
61.25
1.1864
1.1391
GEODIFF-C 89.13
97.88
0.8629
0.8529
61.47
64.55
1.1712
1.1232
The COV-R and MAT-R results of CVGAE, GRAPHDG, CGCF, and CONFGF are borrowed from Shi et al. (2021). The results of GEOMOL are borrowed from a most recent study Zhu et al. (2022). Other results are obtained by our own experiments. The results of all models for the GEOM-QM9 dataset (summarized Tab. 5) are collected in the same way.
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 S, denote the sets of generated and reference conformers respectively, then the Coverage and Matching metrics (Xu et al., 2021 a) following the conventional Recall measurement can be defined as:
S,, [[ee E S,l RMSD(C, $ σ E.
S.)
(10)
min RMSD(C,C), MA, |S CES. CES,
(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 S, 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 certain threshold 8. By contrast, the MAT scores measure the average RMSD conformers in one set with its closest neighbor 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., 202 Ganea et al., 2021), ō is set as 0.5A and 1.25A for QM9 and Drugs datasets respectively.
Results & discussion. The results are summarized Tab. and Tab. (left in Appendix. D). 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 GEODIFT-A and GRODIFF-C respectively. As shown in the tables, GEODIFF consistently outperform the state-of-the- art MI. 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 model the multi modal distribution, and generative both accurate and diverse conformations. We also notice that in general GEODIFF-C performs slightly better than GEODITT-,, which suggests that chain-rule approach leads to 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 qualitative comparison, where GEODIFF is shown to capture better both local and global structures.
On the more challenging Drugs dataset, we further test RDKIT. shown in Tab. 2, our observation is in line with previous studies (Shi et al., 2021) that the state-of-the-art MI 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 more fair comparison. Results in
8

---

Published as a conference paper at ICLR 2022
A Q)
Graph
Referenee
GesDur
CoutiF
GmphPN
Figure 2: Examples of generated structures from Drugs dataset. For every model, show the
conformation best-aligned with the ground truth. More examples are provided in Appendix E.
Table 2: Results on the GEOM-Drugs dataset, with FF optimization.
COV-R (%)
MAT-R (Å)
COV-P (%) 1
MAT-P (Å)
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
RDKIT
60.91
65.70
1.2026
1.1252
72.22
88.72
1.0976
0.9539
GEODIFF + FF
92.27
100.00
0.7618
0.7340
84.51
95.86
0.9834
0.9221
Tab. 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 (Axel- rod & Gomez- Bombar-ll,,,,lli,, 2020) over a set of generated conformations. This can pro- vide an direct assessment on the quality of generated samples. In specific, low 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 Ps14 (Smith al., 2020) calculate each conformer's ' energy and HOMO- LUMO gap and compare the average energy E, lowest energy Emin, average gap Ac, minimum gap Armin and maximum gap A with the ground truth.
Table MAE of predicted ensemble properties in eV.
Method
E
Eyin
Ac
Armi
Afran
RDKIT
0.9233
0.6585
0.3698
0.8021
0,2359
GRAPHDG
9.1027
0.8882
1.7973
4.1743
0.4776
CGCF
28.9661
2.8410
2.8356
10.6361
0.5954
fol-
CONFVAE
8.2080
0.6100
1.6080
3.9111
CON+GF
2.7886
0.1765
0.4688
2.1843
0.1433
GEODIFF
0.25974
0.1551
0.3091
0.7033
0.1909
Results & discussions. The mean absolute errors (MAE) between calculated properties and the ground truth are reported Tab. 3. CVGAE is excluded due 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.
CONCLUSION
We propose GEODIFF, novel probabilistic model for generating molecular conformations, GEODIFF marries denoising diffusion models with geometric representations, where we parameterize the reverse generative dynamics density with equivariant Markov kernels. We derive a tractable invariant objective from the variational lower bound to optimize the likelihood. Comprehensive experiments that GEODIFF is competitive with the existing state-of-the-ars 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.
Markov chain, and novelly impose roto-translan noto-translational invariance into the
multiple tasks demonstrate

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
Minkai thanks Huiyu Cai, David Wipf, Zuobai Zhang, and Zhaocheng Zhu for their helpful discus- sions and comments. This project supported by the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR Al Chair Program, collaboration grants be- tween Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino- Baid Gift Fund and a NRC Collaborative R&D Project AI4D-CORE, This project was also partially funded by IVADO Fundamental Research Project grant PRF-2019- 1583139727. The Stanford team is supported by NSF(#1651565, #1 #1522054, #1 #1733686), ONR (N000141912145), AFOSR (FA95501910024), ARO (W91 (W91INF-21-1-011) and Sloan Fellowship.
REFERENCES
Mohammed AlQuraishi. End-to-end differentiable learning of protein structure. Cell systems, 8(4): 292-301, 2019.
Simon Axelrod and Rafael nomezz-ombarrlli. Geom: for property prediction and molecular generation. arXiv preprint M0Xiv:200600553/,, 2020.
molecular conformations
Simon Batzner, Tess E Smidt, Lixin Sun, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky. Se (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXiv:2101.000 164, 2021.
Julian Chibane, Thiemo Alldieck, and Gerard Pons-Moll. Implicit functions in feature space 3d shape reconstruction and completion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 6970-6981, 2020.
Sybren Ruurds De Groot and Peter Mazur. Non-equilibrium thermodynamics. Courier Corporation, 2013.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In ICLR, 2017.
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molecular fingerprints. In Advances in neural information processing systems, pp. 2224-2232, 2015.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Sc(3)-transformmrs:: 3d roto-
translation equivariant attention networks. NeurIPS, 2020.
Outevian-Fugen Ganea, Lagnajit Pattanaik, Connor W Coley, Regina Barzilay, Klavs F Jensen, William H Green, and Tommi S Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. arXiv preprint arXiv:2106.00 078000 2021.
Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. arXiv preprint 1XXi:22199008824 2021.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In Proceedings of the 34th International Conference on Machine Learning-
SooninggVoolmee Volume 70, pp. 1272. JMLR. org, 2017
T. Gogineni, Ziping Xu, Exequiel Punzalan, Runxuan Jiang, Joshua A Kammeraad, Ambuj Tewari, and P. Zimmerman. Torsionnet: A reinforcement learning approach to sequential conformer search. ArXiv, abs/2006.07078,. 2020.
Thomas A Halgren. Merck molecular force field. v. extension of mmff94 using experimental data, additional computational data, and empirical rules. Journal of Computational Chemistry, 17(5-)): 616-641, 1996.
Paul CD Hawkins. Conformation generation: the state of the art. Journal of Chemical Information and Modeling, 1777-1777 2017.
10

---

Published as a conference paper at ICLR 2022
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distribution examples in neural networks. arXiv preprint arXiv:1610.02136, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006. 1239, 2020.
Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram, Jure Leskovec, Devi Parikh, and Larry Zitnick. Forcenet: A graph neural network for large-scale quantum chemistry simulation. 2021.
John Ingraham, Adam Riesselman, Chris Sander, and Debora S Marks. Learning protein structure with a differentiable simulator. In International Conference on Learning Representations, 2019.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for 2018.
molecular graph generation. arXiv preprint
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. International Conference on Learning Representations, 2021.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-566)
2021.
Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors, Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallographyy 32(5):922-92 1976.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2nd International Conference on Learning Representations, 2013.
Jonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning. 2020.
Leo Liberti, Carlile Lavor, Nelson Maculan, and Antonio Mucherino. Euclidean distance geometry and applications. SIAM review, 56(1):3-69, 2014.
Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. ArXiv, abs/2103.01458, 2021.
Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predicting molecular conformation via dynamic graph score matching. Advances in Neural Information Processing Systems, 34, 2021.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using a deep generative graph neural network. arXiv preprint 1rXiv:1904.00314, 2019.
B. Miller, M. Geiger, T. Smidt, and F. Noé. Relevance of rotationally equivariant convolutions for predicting molecular properties. ArXiv, abs/2008.0844 084461, 2020.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457). 2019.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Sereina Riniker and Gregory A. Landrum. Better informed distance geometry: Using what we know to improve conformation generation. Journal Chemical Information and Modeling, 55(12): 2562-2574, 2015.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling, E (n) equivariant normalizing flows for molecule generation in 3d. arXiv preprint arXiv:2105.0900 2021a.
11

---

Published as a conference paper at ICLR 2022
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks, 2021b.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Mlaur-Robert Müller. Schnet: continuous fuutter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems, pp.
991-1001. Curran Associates, Inc., 2017.
Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek, Alexander WR Nelson, Alex Bridgland, et al. Improved protein structure prediction using potentials from deep learning. Nature, 77767792):706-711, 10, 2020.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: flow-based autoregressive model for molecular graph generation. arXiv preprint a7X1v:2001.09382.. 2020.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molccular conformation generation. ArXiv, 2021.
Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C LLwrrncce Zitmick. Rotation invariant graph neural networks using spin convolutions. arXiv preprint 1rXiv:2/00.09555 2021.
Gregor Simm and Jose Miguel Hernandez-Lobato A generative model for molecular distance geometry. Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119, pp. 8949 8958. PMLR, 2020.
Gregor N. C. Simm, Robert Pinsler, Gábor Csányi, and José Miguel Hermández-Lobal Symmetry- aware actor-critic for 3d molecular design. In International Conference Learning Representa- tions, 2021.
Daniel G. A. Smith, L. Burns, A. Simmonett, R. Parrish, M. C. Schieber, Raimondas Galvelis, P. Kraus, H. Kruse, Roberto Remigio, Asem Alenaizan, A. M. James, S. Lehtola, Jonathon Misiewicz, et al. Psi4 1.4: Open- software for thgh-throughput quantum chemistry. The Journal of chemical physics, 2020.
Jascha Sohl-Dickstein Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv: 1503155033303 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02550 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.
In Advances in Neural Information Processing Systems, pp. 11918-11930, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score- soore-baed generative models. NeurIPS, 2020.
N. Thomas, T. Smidt, Steven M. Kearnes, Lusann Yang, L. Li, Kai Kohlhoff, and P. Riley. Tensor field networks: Rotation- and rrausation-eeuivvarian neural networks for 3d point clouds. ArXiv, 2018.
M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. In NeurIPS, 2018.
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. In International Conference Learning Representations, 2021a.
Minkai Xu, Wujie Wang, Shitong Luo, Chence Shi, Yoshua Bengio, Rafael and Jian Tang. An end-to-end framework for molecular conformation generation via bilevel programming. arXiv preprint 11X:v:2105.07246, 2021b.
12

---

Published as a conference paper at ICLR 2022
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep Potential Molecular Dynamics: Scalable Model with the Accuracy of Quantum Mechanics, Physical Review Letters, 120(14):143001,, 2018.
Jinhua Zhu, Yingee Xia, Chang Liu, Lijun Wu, Shufang Xie, Tong Wang, Yusong Wang, Wengang Zhou, Tao Qin, Hougiang Li, et al. Direct molecular conformation generation. arXiv preprint 35.. 2022.
13

---

Published as a conference paper at ICLR 2022
A PROOFS
A. PROPERTIES OF THE DIFFUSION MODEL
We include proof's for several key properties of the probabilistic diffusion model here to self- contained. For more detailed discussions, please refer to Ho et al. (2020). Let (.,,..,, BT) be sequence of variances, and 1 1- Br and a; II as. The two following properties are crucial for deriving the final tractable objective in equation 9. Property 1. Tractable marginal of the forward process:
a(c*)c))= = f q(cl"*()) act)tttt)-)))))))))))))))))))))) =c =N(C'; =
(1-c)))
Proof. Let c,'s be independent standard Gaussian random variables. Then, by definition of the Markov kernels ((t) in equation 2, we have
C
Va,Ct-1 Bic
V0 v 0.31-184 + VBtes
Vatal-1ar-1(" + 1 V Va + VBer
(12)
va,Co +
+ ... + Va + VB,e.
(²) Therefore q(C (²) is still Gaussian, and the mean of C is va,C", and the variance matrix is (at0--1 -. 01 + ... + a18i-1 + B.) (1 - ā. Then we have:
q(C'IC") = -N(C"; vā,Co, (1-6:)) T))
This property provides convenient closed-form evaluation of C' knowing CO:
C VorC v1 āre,
where 1).
Besides, it is worth noting that,
[ CarCo, (1- (1)))).,
where âT = -I (1 3.) approaches zero with large T, which indicates the diffusion process can finally converge into whitened noisy distribution.
Property 2. Tractable posterior of the forward process:
vāl-1& cD +
Var(1 0n-1) Ct, (1 āx-1) 1 āe
NN
Br1).
1- at
ae.
Proof. Let 3 I
1-ā I-a4
3r, then derive the posterior by Bayes rule:
((C'(c-1) 9(c
(()
NiC ca
3.1) N(C--1 Var (1- (1-
varCo
(1-0x;)]]
-(2T8)) {2T((1 -
))-1)) $(2T(1 a. 5x
||C varC"
-13
||t-1
000100111
| - 112
exp
2
2(1 ā(-1)
2(1 - ar)
0 co
Var( (1 - Ō(-1) ctc
(2) -2 exp
ct-1
2
1 ai
(13)
Then have the posterior (CC C)) the given form.
14

---

Published as a conference paper at ICLR 2022
A.2 PROOF OF PROPOSITION 1
Let Tg be some roto-translational transformations of group element g E SE(3), and let p(ær) be density which is SE(3)-invariant, i.e., p(rr) ((Tg(TT)) (xT)). If the Markov transitions p -quiiiiqui,, i.e., p(zc-1(2;) ( (2t-1) (z)), then we have that the density po = (F
also SE(3)-invariant
Proof.
po (Tg
/
(invariant prior p(ær))
(14)
/ ((0 (equivariant kernels p( 1
/ pler))o((007-1133))dd...
(a()))
A.3 PROOF OF PROPOSITION 2
In this section we prove that the output of GFN defined in equation 5, 6 and 7 is translationally invariant and rotationally equivariant with the input C. Let g E IR' denote any translation transforma- tions and orthogonal matrices R 3 denote any rotation transformations. let Rx be shorthand for (Rx))-... Rxx)). Formally, we aim to prove that the model satisfies:
Rx h GFN( RC + gg,h".
(15)
This equation indicates that, given x' already rotationally equivalent with C. and h' already invariant, then such property can propagate through single GFN layer to +!+1 and h [+]
Firstly, given that h' already invariant to SE(3) transformations, we have that the messages my calculated from equation will also be invariant. This is because it sorely relies on
Proof.
distance between two atoms, which are manifestly invariant to rotations |[x Rx2|11 (x!
xx)) T R.T. R(x x)) = (x xx)) T [x - xx) = ||x -xxIll x21 1|2. Formally, the invariance of messages
in equation can be written as:
m, 4m [1/, h/. Rx! Rx| 1 eis) ;()) pm (n). b), || x311 .cu)
(16)
And similarly, the + 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 proven above, we have that:
Σ (EN(()
(Re, + - Rej- - g) t. (ma,) = R Σ d:y
(c+ c,) t2 (m.,g) = Rxx+1 (17)
di
JEN(()
Therefore, we have that rotating and translating e results in the same rotation and no translation on x/+1 by updating through equation? 7.
Thus we can conclude that the property defined in equation 15 satisfied.
Having proved the equivariance property of a single GFN layer, then inductively, we can draw conclusion that composition of T. GFN layers will also preserve the same equivariance.
15

---

Published as a conference paper at ICLR 2022
A.4 PROOF OF PROPOSITION 3
We first derive the variational lower bound (ELBO) objective in equation 8. The ELBO can be calculated follows:
Po ( p(ct)
log po (g) = E. log (0))
(CC
2 Eg log Pe (CO.T-1 CT)) x p(CT)
((C.T(C))
T
po (c CO)
- EEg logp(c
log ((c((ct-1) Σ
t=1
(C g,c
po 1o
log (CC
- By logg(CT) log
Σ
log (c---10',
q(=2) (C)
((C1(CU)
t=2
1g log p(c")
T
po (C -9
(C) (C g,Cl)
Σ
log ((cr-1C Cu)
T
+ = -Bg KL (9(c= (C (l)()(())
Σ KL (g(C* c*)
- log lo gccl
(18)
can be noted that the first term KL (P) is a constant, which can be omit- ted in the objective. Furthermore, for brevity, we also merge the final term log (C into the second term (sum over KL divergences), and finally derive that LELBO as in equation 8.
E Dxu(g(C=-1c' (C)
Now consider how compute the KL. divergences the proposition 3. Since both (C'' CC) and Po "g.""GGC"" are Gaussian share the same covariance matrix 3 the KL divergence between them can be calculated by the squared &g distance between their means weighed by certain weights 23, By the expression of (C") we have the reparameterization that C vanCo vI Oie.
Then we can derive:
E, KL (a(c--1
Ct, ())
1
E
00- Br co + var(L
at-1)
C Ct_ 1 1((e
ct- Be
co co(C",gtt)
)|²
2
1 ō
Vos
-1-
1
V Br - -I- diE
Vaill -a(-1) Ci
1
3r
)p²
Eco,
C
CC (CC,,,t) g. -1-ay a
23
ai
Vāl
ā.
Vay
1
87
Ec 10.1000 +c (,,,,)'
28
an(1 at)
87
Eco,e Coo,ele-Eeo('' (C(C", 1)
21-M--1
2
3,01(1 (1 ōz))
= =ne ||e- ||--<(C'', })
(19)
where "t represent the wights 20.(1-a.-11
8. And we finish the proof.
A.5 ANALYSIS OF THE INVARIANT DENSITY IN SEC. 4.2
Given geometric system & E N-3 we obtain the CoM-free á by subtracting its CoM. This can be considered as linear transformation:
f b= Qu, where Q Q=19 3 IN
1
(20)
where Is denotes the k x identity matrix and 1, denotes the k-diimensional vector filled with ones. It can be noted that Q is symmetric projection operator, iee., Q² Q and Q Q. And also
16

---

Published as a conference paper at ICLR 2022
have that rank[ = (N - Furthermore, let U represent the space of CoM-free systems, can easily have that Qy = for any y E U since the CoM of y is already zero.
( In). then the CoM-free density can be formally written as و = N (0, QI, = (0, QQ"). Thus, sampling from be trivially achieved by sampling from p and then projecting with Q. And p(y) can be calculated by p(y) since for any y E U we have llyll = |Qy l and thus p(y) p(y)-
Formally, let .N 3 and set I with isotropic normal distribution
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.
B OTHER RELATED WORK
Protein structure generation. There has also been many recent works working on protein structure folding. For example, Boltzmann generators Noć et al. (2019) flow-bused models to generate the structure protein main chains. AlQuraishi (2019) uses recurrent networks to model the amino acid sequences. Ingraham et al. (2019) proposed neural networks to learn 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 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.
C EXPERIMENT DETAILS
In this section, introduce the details of our experiments. In practice, the means 6a are parameterized as compositions of both typical invariant MPNNs (Schütt et al., 2017) and the proposed equivariant GFNs in Sec. 4.2. As default setup, the MPNNs for parameterizing the means eg all implemented with 4 layers, and the hidden embedding dimension is set as 128. After the MPNNs, we can obtain the informative invariant atom embeddings. which we denote as h". Then the embeddings he are fed into equivariant layers and updated with equation 5, equation 6, and equation 7 to obtain the equivariant output. For the training of GEODIFF, we train the model on a single Tesla V100 GPU with learning rate of 0.001 until convergence and Adam (Kingma & Welling, 2013) as the optimizer. parameters GEODIFF are summarized in
The practical training time is -48 hours. The other Tab. 4, including highest variance level Br, lowest variance level By, the variance schedule, number of diffusion timesteps T, radius threshold for determining the neighbor of atoms т, batch size, and number of training iterations.
Table 4: Additional hyperparameters of our GEODIFF,
Task
B1
Br
3 scheduler
T
Batch Size
Train Iter.
QM9
1e-7
2e-3
sigmoid
5000
10Å
64
1M
Drugs
le-7
2e-3
sigmoid
5000
10Å
32
IM
D ADDITIONAL EXPERIMENTS
D. RESULTS FOR GEOM-QM9
The results the GEOM-QM9 dataset are reported in Tab. 5.
17

---

Published as a conference paper at ICLR 2022
Table Results on the GEOM-QM9 dataset, without FF optimization.
COV-R (%) f
MAT-R (Å) 4
COV-P (%) f
MAT-P (Å)
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
CVGAE
0.09
0.00
1.6713
1.6088
GRAPHDGG
73.33
84.21
0,4245
0.3973
43.90
35.33
0,5809
0.5823
CGCF
78.05
82,48
0,4219
0.3900
36.49
33.57
0.6615
0.6427
CONFVAE
77.84
88.20
0.4154
0.3739
38.02
34.67
0.6215
0.6091
GROMOL
71.26
72.00
0.3731
0.3731
CONFGF
88.49
94.31
0.2673
0.2685
46.43
43,41
0.5224
0,5124
GEODIFF-A
90.54
94.61
0,2104
0.2021
52.35
50.10
0,4539
0,4399
GEODIFF-C
90.07
93.39
0.2090
0.1988
52.79
50.29
0.4448
0.4267
Table 6: Additional results on the GEOM-Drugs dataset, without FF optimization.
COV-R (
MAT-R (Å)
COV-P (%)
MAT-P (Å)
Models
Mean
Median
Mean
Median
Mean
Median
Mean
Median
GEODIEF (T=1000)
82.96
96.29
0.9525
0.9334
48.27
46.03
1.3205
1.2724
D.2 ABLATION STUDY WITH FEWER DIFFUSION STEPS
We also method with fewer diffusion steps. Specifically, test the setting with T 1000, 31 = and ôT =9e-3. The results on the more challenging Drugs dataset are shown in Tab. 6. Compared with the results in Tab. I, 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-e method.
E MORE VISUALIZATIONS
We provide more visualization of generated structures in Fig. 3. The molecules are chosen from the test split of GEOM-Drugs dataset.
18

---

Published as a conference paper at ICLR 2022
Graph
Conformations
taolot
Figure 3: Visualization of drug-like conformations generated by GEODIFF.
19


Reviews:
[{"review_id": "vL48helhaFK", "reviewer": "Reviewer_s7Z3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper brings together recent ideas and methods (e.g. diffusion, SE(3) equivariance) to the established task of molecular conformation generation with impressive empirical results. A few more experiments are necessary. The paper is well written. I am only tangentially in this area so I am not intimately familiar with the prior works and benchmarks, and may have missed something in the empirical evaluation.\n\nComments:\n* The authors state that roto-translation invariance is a critical inductive bias for this problem. But I am not convinced that it is critical for this setup and model class. It would be nice to include a baseline of the diffusion model without SE(3) invariance.\n* I am curious about the sensitivity of the performance of the model to the hyperparameters of the diffusion process. Can the authors describe how they reached the current settings given in Table 5. In particular, T=5000 steps seems quite high and I was curious about the performance with lower T say around 1000. As a minor note, there appears to be an inconsistency in the reported batch size (128 in the text vs. 32/64 in Table 5)\n* What is the generation time with 5000 diffusion steps? Does that become a bottleneck in conformation generation?\n* What happened to the other methods on the geom-drugs dataset for Table 3?\n* The authors should mention recent point cloud diffusion papers in their related works as it is addressing a very similar problem\n* What are the reverse process sigmas set to?\n\nMinor comments:\n* Typo: translatoinal\n* Typo: sophisticate\n* I noticed a reference that seemed incorrect: Xu et al 2021a as the reference for molecular generation with VAEs in the second paragraph. The authors should double check their references\n", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall a nice paper with technical novelty and good empirical results. I have some questions about the evaluation since some apt baselines are missing. I will raise my score if the authors can address my comments.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aQ5VdomRqq6", "reviewer": "Reviewer_sVuD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality:** The idea of the work is novel - **GeoDiff** is the first generative model for molecular conformation generation based on a diffusion framework.\n\n**Clarity and Quality:** The authors clearly provide motivation and state the differences of the proposed model from other approaches. The architecture is described in an easy-to-follow and intuitive way. The paper provides good **Related Work** and **Results & Discussion** sections. Still, it would be great if authors add a clear description of contributions and additional figures of architecture and an aligned cloud of conformations to visually evaluate the diversity of conformations.\n\n**Significance:** The contribution of the paper to the field is significant. The experiments show a significant metrics gap between the proposed method and baselines, still, it would be great if authors will add standard deviation of metrics to evaluate the stability of the method.\n\n**Drawbacks / questions**:\n\n1. In 3.1 Authors say, they are interested in generating stable conformations. Authors provide energy metrics for the generated ensemble on GEOM-QM9 split, which covers only 30 molecules. Still, they don't provide energy metrics for medium size molecules from GEOM-DRUGS dataset.\n2. Despite the proposed model being compared against a variety of baselines and that the authors mentioned the reason why they are not compared with the GeoMol, it requires a more detailed explanation of GeoMol limitations, since the GeoMol is the current state of the art approach in conformation generation task.\n3. Diffusion models suffer from high computational complexity. Based on the parameters provided by the authors in appendix B, it's logical to assume that 5000 iterations per molecule are very computationally expensive. The article lacks an estimate of the generation time for one molecule.", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "The idea of the paper is novel. The text of the paper is clear and easy to follow. The experiments show a significant gap in metrics values between the proposed model and baselines. Still, there are several minor drawbacks that can be fixed in the next revision of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "QDqqqk7H3C", "reviewer": "Reviewer_PLjF", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Pros\n\n* GeoDiff is an end-to-end model for generating Cartesian coordinates from a molecular graph. \n\n* GeoDiff shows very good performance on standard test sets.\n\n* GeoDiff outperforms various state-of-the-art methods for generating molecular conformations. \n\n## Cons\n\n* GeoDiff's model seems to involve a large number of parameters. \n\n## Comments and Questions\n\n* What is the computation time for generating a structure with GeoDiff? How does this compare to other ML approaches or RDKit? \n\n* What is the accuracy of the structures at intermediate stages of the diffusion process? Perhaps this could provide some insight into whether $T=5000$ diffusion steps are needed to reach the reported performance. \n\n* Have you also run tests with a smaller number of diffusion steps? \n\n* You show that GeoDiff's model and objective functions are invariant under rigid transformations. However, what if the training data are not invariant? Doesn't this still fix some arbitrary orientation in the trained model? \n\n* Which threshold $\\delta$ did you use in the computation of the precision and recall measures? \n\n* What is the total number of model parameters that are trained? \n\n* How is GeoDiff +FF implemented? What is GeoDiff's performance in property prediction without FF? \n\n## Language\n\n* Page 3: \"It has been shown effective\"\n\n* Page 9: \"can stay the superior diversity\"\n\n* Page 9: \"yeild\" instead of \"yield\"\n\n", "summary_of_the_paper": "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. \n", "summary_of_the_review": " GeoDiff is an end-to-end model for generating Cartesian coordinates from a molecular graph. GeoDiff outperforms various state-of-the-art methods for generating molecular conformations. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors focus on the conditional generation of molecular conformations (i.e. 3D cartesian atom positions) from a given molecular graph. They formulate the generation via diffusion probabilistic models.  Conformations are generated by a reverse diffusion process from isotropic Gaussian noise to molecular conformations. This diffusion process is learned from data using a SE(3) invariant formulation of the diffusion process. The authors work directly with atomic positions (i.e. a point cloud) instead of interatomic distances or an intermediate bond geometry representation. Experimental evaluations show state-of-the-art results according to COV/MAT metrics on GEOM-Drugs and GEOM-QM9 datasets.

Strengths:

- High technical novelty: first generative model for molecular conformation generation based on a diffusion framework
- Very clearly written paper.
- Impressive empirical results with state-of-the-art results on GEOM-Drugs and GEOM-QM9 datasets.

Weaknesses:

- Most of the weaknesses reported by the reviewers seem to have been addressed in the rebuttal.

The idea of the work is highly novel. The authors propose the first generative model for molecular conformation generation based on a diffusion framework. This paper brings together recent ideas and methods (e.g. diffusion, SE(3) equivariance) to the established task of molecular conformation generation with impressive empirical results. All the reviewers agree on acceptance with high scores. I recommend the authors to look at the reviewers' comments to improve the paper for the camera-ready version
Decision: Accept (Oral)