Original: G EO D IFF : A G EOMETRIC D IFFUSION M ODEL FOR M OLECULAR C ONFORMATION G ENERATION
Revised: G EO D IFF : A G EOMETRIC D IFFUSION M ODEL FOR M OLECULAR C ONFORMATION G ENERATION

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery.
Revised: Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery.

Original: Recently, signiﬁcant progress has been achieved with machine learning approaches, especially with deep generative models.
Revised: Recently, signiﬁcant progress has been achieved with machine learning approaches, especially with deep generative models.

Original: 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 G EO D IFF for molecular conformation prediction.
Revised: 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 G EO D IFF for molecular conformation prediction.

Original: G EO D IFF 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.
Revised: G EO D IFF 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.

Original: Modeling such a generation process is however very challenging as the likelihood of conformations should be rototranslational invariant.
Revised: Modeling such a generation process is however very challenging as the likelihood of conformations should be rototranslational invariant.

Original: 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.
Revised: 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.

Original: The whole framework can be efﬁciently trained in an end-toend fashion by optimizing a weighted variational lower bound to the (conditional) likelihood.
Revised: The whole framework can be efﬁciently trained in an end-toend fashion by optimizing a weighted variational lower bound to the (conditional) likelihood.

Original: Experiments on multiple benchmarks show that G EO D IFF is superior or comparable to existing state-of-the-art approaches, especially on large molecules
Revised: Experiments on multiple benchmarks show that G EO D IFF is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Graph representation learning has achieved huge success for molecule modeling in various tasks ranging from property prediction (Gilmer et al., 2017;
Revised: Graph representation learning has achieved huge success for molecule modeling in various tasks ranging from property prediction (Gilmer et al., 2017;

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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;
Revised: 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;

Original: Gebauer et al., 2021;
Revised: Gebauer et al., 2021;

Original: Jing et al., 2021; Batzner et al., 2021).
Revised: Jing et al., 2021; Batzner et al., 2021).

Original: Unfortunately, how to predict stable molecular conformation remains a challenging problem.
Revised: Unfortunately, how to predict stable molecular conformation remains a challenging problem.

Original: Traditional methods based on molecular dynamics (MD) or Markov chain Monte Carlo (MCMC) are very computationally expensive, especially for large molecules (Hawkins, 2017).
Revised: Traditional methods based on molecular dynamics (MD) or Markov chain Monte Carlo (MCMC) are very computationally expensive, especially for large molecules (Hawkins, 2017).

Original: Recently, signiﬁcant progress has been made with machine learning approaches, especially with deep generative models.
Revised: Recently, signiﬁcant progress has been made with machine learning approaches, especially with deep generative models.

Original: For example, Simm & Hernandez-Lobato (2020); Xu et al.
Revised: For example, Simm & Hernandez-Lobato (2020); Xu et al.

Original: (2021b)
Revised: (2021b)

Original: studied predicting atomic distances with variational autoencoders (VAEs) (Kingma & Welling, 2013) and ﬂow-based models (Dinh et al., 2017) respectively.
Revised: studied predicting atomic distances with variational autoencoders (VAEs) (Kingma & Welling, 2013) and ﬂow-based models (Dinh et al., 2017) respectively.

Original: Shi et al.
Revised: Shi et al.

Original: (2021) proposed to use denoising score matching (Song & Ermon, 2019; 2020) to estimate the gradient ﬁelds over atomic distances, through which the gradient ﬁelds over atomic coordinates can be calculated.
Revised: (2021) proposed to use denoising score matching (Song & Ermon, 2019; 2020) to estimate the gradient ﬁelds over atomic distances, through which the gradient ﬁelds over atomic coordinates can be calculated.

Original: Ganea et al.
Revised: Ganea et al.

Original: (2021) studied generating conformations by predicting both bond lengths and angles.
Revised: (2021) studied generating conformations by predicting both bond lengths and angles.

Original: 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.
Revised: 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.

Original: As a result, they are able to achieve very compelling performance.
Revised: As a result, they are able to achieve very compelling performance.

Original: 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.
Revised: 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.

Original: 2 for a detailed description).
Revised: 2 for a detailed description).

Original: 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.
Revised: 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.

Original: In this paper, we propose such a solution called G EO D IFF , a principled probabilistic framework based on denoising diffusion models (Sohl-Dickstein et al., 2015).
Revised: In this paper, we propose such a solution called G EO D IFF , a principled probabilistic framework based on denoising diffusion models (Sohl-Dickstein et al., 2015).

Original: Our approach is inspired by the diffusion process in nonequilibrium thermodynamics (De Groot & Mazur, 2013).
Revised: Our approach is inspired by the diffusion process in nonequilibrium thermodynamics (De Groot & Mazur, 2013).

Original: 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.
Revised: 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.

Original: At each time step, stochastic noises are added to the atomic positions.
Revised: At each time step, stochastic noises are added to the atomic positions.

Original: Our high-level idea is learning to reverse the diffusion process, which recovers the target geometric distribution from the noisy distribution.
Revised: Our high-level idea is learning to reverse the diffusion process, which recovers the target geometric distribution from the noisy distribution.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: To this end, we ﬁrst 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.
Revised: To this end, we ﬁrst 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.

Original: We further provide practical parameterization to deﬁne a roto-translational invariant prior distribution and a
Revised: We further provide practical parameterization to deﬁne a roto-translational invariant prior distribution and a

Original: Markov kernel imposing the equivariance constraints.
Revised: Markov kernel imposing the equivariance constraints.

Original: 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 efﬁciently optimized.
Revised: 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 efﬁciently optimized.

Original: A unique strength of G EO D IFF is that it directly acts on the atomic coordinates and entirely bypasses the usage of intermediate elements for both training and inference.
Revised: A unique strength of G EO D IFF is that it directly acts on the atomic coordinates and entirely bypasses the usage of intermediate elements for both training and inference.

Original: This general formulation enjoys several crucial advantages.
Revised: This general formulation enjoys several crucial advantages.

Original: First, the model can be naturally trained end-to-end without involving any sophisticated techniques like bilevel programming (Xu et al., 2021b), which beneﬁts from small optimization variances.
Revised: First, the model can be naturally trained end-to-end without involving any sophisticated techniques like bilevel programming (Xu et al., 2021b), which beneﬁts from small optimization variances.

Original: 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.
Revised: 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.

Original: Moreover, G EO D IFF enjoys a high model capacity to approximate the complex distribution of conformations.
Revised: Moreover, G EO D IFF enjoys a high model capacity to approximate the complex distribution of conformations.

Original: Thus, the model can better estimate the highly multi-modal distribution and generate structures with high quality and diversity.
Revised: Thus, the model can better estimate the highly multi-modal distribution and generate structures with high quality and diversity.

Original: We conduct comprehensive experiments on multiple benchmarks, including conformation generation and property prediction tasks.
Revised: We conduct comprehensive experiments on multiple benchmarks, including conformation generation and property prediction tasks.

Original: Numerical results show that G EO D IFF consistently outperforms existing state-of-the-art machine learning approaches, and by a large margin on the more challenging large molecules.
Revised: Numerical results show that G EO D IFF consistently outperforms existing state-of-the-art machine learning approaches, and by a large margin on the more challenging large molecules.

Original: The signiﬁcantly superior performance demonstrate the high capacity to model the complex distribution of molecular conformations and generate both diverse and accurate molecules.
Revised: The signiﬁcantly superior performance demonstrate the high capacity to model the complex distribution of molecular conformations and generate both diverse and accurate molecules.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Recently, various deep generative models have been proposed for conformation generation.
Revised: Recently, various deep generative models have been proposed for conformation generation.

Original: Among them, CVGAE (Mansimov et al., 2019) ﬁrst 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.
Revised: Among them, CVGAE (Mansimov et al., 2019) ﬁrst 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.

Original: To address this problem, the majority of subsequent models are based on intermediate geometric elements such as atomic distances and torsion angles.
Revised: To address this problem, the majority of subsequent models are based on intermediate geometric elements such as atomic distances and torsion angles.

Original: 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).
Revised: 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).

Original: However, such a decomposition suffers from several drawbacks for either training or sampling.
Revised: However, such a decomposition suffers from several drawbacks for either training or sampling.

Original: For example, G RAPH DG (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.
Revised: For example, G RAPH DG (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.

Original: C
Revised: C

Original: ONF VAE further improves this pipeline by designing an end-to-end framework via bilevel optimization (Xu et al., 2021b).
Revised: ONF VAE further improves this pipeline by designing an end-to-end framework via bilevel optimization (Xu et al., 2021b).

Original: 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.
Revised: 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.

Original: To overcome this problem, C ONF GF (Shi et al.,  2021) proposed to learn the gradient of the log-likelihood w.r.t coordinates.
Revised: To overcome this problem, C ONF GF (Shi et al., 2021; Luo et al., 2021) proposed to learn the gradient of the log-likelihood w.r.t coordinates.

Original: However, in practice the model is still aided by intermediate geometric elements, in that it
Revised: However, in practice the model is still aided by intermediate geometric elements, in that it

Original: ﬁrst 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.
Revised: ﬁrst 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Most recently, Ganea et al. (2021) proposed a highly systematic (rule-based) pipeline named G EO M OL , 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.
Revised: Most recently, another concurrent work (Ganea et al., 2021) proposed a highly systematic (rule-based) pipeline named G EO M OL , which learns to predict a minimal set of geometric quantities ( i.e. length and angles) and then reconstruct the local and global structures of the conformation in a sophisticated procedure.

Original: Despite its effectiveness, this method can be only adopted to molecules with a speciﬁc pattern, which limits it from wide applications.

Original: In our experiments, the pattern requirement ﬁltered out nearly one third of molecular graphs, making it incomparable in our setting.

Original: Besides, there has also been efforts to use reinforcement learning for conformation search Gogineni et al.
Revised: Besides, there has also been efforts to use reinforcement learning for conformation search Gogineni et al.

Original: Nevertheless, this method relies on rigid rotor approximation and can only model the torsion angles, and thus fundamentally differs from other approaches.
Revised: Nevertheless, this method relies on rigid rotor approximation and can only model the torsion angles, and thus fundamentally differs from other approaches.

Original:  P RELIMINARIES 3. N OTATIONS AND P ROBLEM D EFINITION
Revised:  P RELIMINARIES 3. N OTATIONS AND P ROBLEM D EFINITION

Original: Notations.
Revised: Notations.

Original: In this paper each molecule with n atoms is represented as an undirected graph G =
Revised: In this paper each molecule with n atoms is represented as an undirected graph G =

Original: (cid:104)V , E(cid:105) , where
Revised: (cid:104)V , E(cid:105) , where

Original: V = { v i } ni =1 is the set of vertices representing atoms and E =
Revised: V = { v i } ni =1 is the set of vertices representing atoms and E =

Original: { e ij | ( i, j ) ⊆ |V| × |V|} is the set of edges representing inter-atomic bonds.
Revised: { e ij | ( i, j ) ⊆ |V| × |V|} is the set of edges representing inter-atomic bonds.

Original: Each node v i ∈ V describes the atomic attributes, e.g. , the element type.
Revised: Each node v i ∈ V describes the atomic attributes, e.g. , the element type.

Original: Each edge e ij ∈ E describes the corresponding connection between v i and v j , and is labeled with its chemical type.
Revised: Each edge e ij ∈ E describes the corresponding connection between v i and v j , and is labeled with its chemical type.

Original: In addition, we also assign the unconnected edges with a virtual type.
Revised: In addition, we also assign the unconnected edges with a virtual type.

Original: For the geometry, each atom in V is embedded by a coordinate vector c ∈ R 3 into the 3-dimensional space, and the full set of positions ( i.e. , the conformation) can be represented as a matrix C
Revised: For the geometry, each atom in V is embedded by a coordinate vector c ∈ R 3 into the 3-dimensional space, and the full set of positions ( i.e. , the conformation) can be represented as a matrix C

Original: =
Revised: =

Original: [ c 1 , c 2 , · · · , c n ] ∈ R n × 3 .
Revised: [ c 1 , c 2 , · · · , c n ] ∈ R n × 3 .

Original: Problem Deﬁnition.
Revised: Problem Deﬁnition.

Original: The task of molecular conformation generation is a conditional generative problem, where we are interested in generating stable conformations for a provided graph G .
Revised: The task of molecular conformation generation is a conditional generative problem, where we are interested in generating stable conformations for a provided graph G .

Original: Given multiple graphs G , and for each G given its conformations C as i.i.d samples from an underlying
Revised: Given multiple graphs G , and for each G given its conformations C as i.i.d samples from an underlying

Original: 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.
Revised: 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.

Original: 3. E QUIVARIANCE
Revised: 3. E QUIVARIANCE

Original: Equivariance is ubiquitous in machine learning for atomic systems, e.g. , the vectors of atomic dipoles or forces should rotate accordingly w.r.t.
Revised: Equivariance is ubiquitous in machine learning for atomic systems, e.g. , the vectors of atomic dipoles or forces should rotate accordingly w.r.t.

Original: 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).
Revised: 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).

Original: 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;
Revised: 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;

Original: Satorras et al., 2021a).
Revised: Satorras et al., 2021a).

Original: Formally, a function F : X → Y is equivariant w.r.t a group G if:
Revised: Formally, a function F : X → Y is equivariant w.r.t a group G if:

Original: [Equation]
Revised: [Equation]

Original: where T g and S g are transformations for an element g ∈ G , acting on the vector spaces X and
Revised: where T g and S g are transformations for an element g ∈ G , acting on the vector spaces X and

Original: Y , respectively.
Revised: Y , respectively.

Original: In this work, we consider the SE(3) group, i.e. , the group of rotation, translation in 3D space.
Revised: In this work, we consider the SE(3) group, i.e. , the group of rotation, translation in 3D space.

Original: 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.
Revised: 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.

Original:  G EO D IFF M ETHOD
Revised:  G EO D IFF M ETHOD

Original: In this section, we elaborate on the proposed equivariant diffusion framework.
Revised: In this section, we elaborate on the proposed equivariant diffusion framework.

Original: We ﬁrst 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).
Revised: We ﬁrst 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).

Original: Then we emphasize several  non-trivial challenges of building diffusion models for geometry generation scenario, and show how we technically tackle these issues.
Revised: Then we emphasize several [Figure] non-trivial challenges of building diffusion models for geometry generation scenario, and show how we technically tackle these issues.

Original: Speciﬁcally, in Sec. 4.2, we present how we parameterize p θ ( C|G ) so that the conditional likelihood is roto-translational invariant, and in Sec.
Revised: Speciﬁcally, in Sec. 4.2, we present how we parameterize p θ ( C|G ) so that the conditional likelihood is roto-translational invariant, and in Sec.

Original: 4.3, we introduce our surgery of the training objective to make the optimization also invariant of translation and rotation.
Revised: 4.3, we introduce our surgery of the training objective to make the optimization also invariant of translation and rotation.

Original: Finally, we brieﬂy show how to draw samples from our model in Sec. 4.4.[Figure]
Revised: Finally, we brieﬂy show how to draw samples from our model in Sec. 4.4.

Original: 4. F ORMULATION
Revised: 4. F ORMULATION

Original: Let C 0 denotes the ground truth conformations and let C t for t = 1 , · · · , T be a sequence of latent variables with the same dimension, where t is the index for diffusion steps.
Revised: Let C 0 denotes the ground truth conformations and let C t for t = 1 , · · · , T be a sequence of latent variables with the same dimension, where t is the index for diffusion steps.

Original: 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.
Revised: 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.

Original: Intuitively, the diffusion process progressively injects small noises to the data C 0 , while the generative process learns to revert the diffusion process by gradually eliminating the noise to recover the ground truth.
Revised: Intuitively, the diffusion process progressively injects small noises to the data C 0 , while the generative process learns to revert the diffusion process by gradually eliminating the noise to recover the ground truth.

Original: We provide a high-level schematic of the processes in Fig.
Revised: We provide a high-level schematic of the processes in Fig.

Original: Diffusion process.
Revised: Diffusion process.

Original: Following the physical insight, we model the particles C as an evolving thermodynamic system.
Revised: Following the physical insight, we model the particles C as an evolving thermodynamic system.

Original: With time going by, the equilibrium conformation C 0 will gradually diffuse to the next chaotic states C t , and ﬁnally converge into a white noise distribution after T iterations.
Revised: With time going by, the equilibrium conformation C 0 will gradually diffuse to the next chaotic states C t , and ﬁnally converge into a white noise distribution after T iterations.

Original: Different from typical latent variable models, in diffusion model this forward process is deﬁned as a ﬁxed (rather than trainable) posterior distribution q ( C 1: T |C 0 ) .
Revised: Different from typical latent variable models, in diffusion model this forward process is deﬁned as a ﬁxed (rather than trainable) posterior distribution q ( C 1: T |C 0 ) .

Original: Speciﬁcally, we deﬁne it as a Markov chain according to a ﬁxed variance schedule β 1 , . .
Revised: Speciﬁcally, we deﬁne it as a Markov chain according to a ﬁxed variance schedule β 1 , . .

Original: . , β T :
Revised: . , β T :

Original: [Equation]
Revised: [Equation]

Original: Note that, in this work we do not impose speciﬁc (invariance) requirement upon the diffusion process, as long as it can efﬁciently draw noisy samples for training the generative process p θ ( C 0 ) .
Revised: Note that, in this work we do not impose speciﬁc (invariance) requirement upon the diffusion process, as long as it can efﬁciently draw noisy samples for training the generative process p θ ( C 0 ) .

Original: Let α t = 1 − β t
Revised: Let α t = 1 − β t

Original: and ¯ α t = (cid:81) ts =1{+t s =
+} α s , a special property of the forward process is that q ( C t |C 0 ) of arbitrary timestep t can be calculated in closed form q ( C t |C 0 ) =
Revised: and ¯ α t = (cid:81) {+t s =
+} α s , a special property of the forward process is that q ( C t |C 0 ) of arbitrary timestep t can be calculated in closed form q ( C t |C 0 ) =

Original: N ( C t ; √ ¯ α t C 0 , (1 − ¯ α t ) I ) 1 .
Revised: N ( C t ; √ ¯ α t C 0 , (1 − ¯ α t ) I ) 2 .

Original: This indicates with sufﬁciently large T , the whole forward process will convert C 0 to whitened isotropic
Revised: This indicates with sufﬁciently large T , the whole forward process will convert C 0 to whitened isotropic

Original: Gaussian, and thus it is natural to set p ( C T ) as a standard Gaussian distribution.
Revised: Gaussian, and thus it is natural to set p ( C T ) as a standard Gaussian distribution.

Original: Reverse Process.
Revised: Reverse Process.

Original: Our goal is learning to recover conformations C 0 from the white noise C T , given speciﬁed molecular graphs G .
Revised: Our goal is learning to recover conformations C 0 from the white noise C T , given speciﬁed molecular graphs G .

Original: We consider this generative procedure as a reverse dynamics of the above diffusion process, starting from the noisy particles C T ∼ p ( C T ) .
Revised: We consider this generative procedure as a reverse dynamics of the above diffusion process, starting from the noisy particles C T ∼ p ( C T ) .

Original: We formulate this reverse dynamics as a conditional Markov chain with learnable transitions:
Revised: We formulate this reverse dynamics as a conditional Markov chain with learnable transitions:

Original: [Equation]
Revised: [Equation]

Original: Herein µ θ are parameterized neural networks to estimate the means, and σ t can be any user-deﬁned variance.
Revised: Herein µ θ are parameterized neural networks to estimate the means, and σ t can be any user-deﬁned variance.

Original: The initial distribution p ( C T ) is set as a standard Gaussian.
Revised: The initial distribution p ( C T ) is set as a standard Gaussian.

Original: Given a graph G , its 3D structure is generated by ﬁrst drawing chaotic particles C T from p ( C T ) , and then iteratively reﬁned through the reverse Markov kernels p θ ( C t − 1 |G , C t ) .
Revised: Given a graph G , its 3D structure is generated by ﬁrst drawing chaotic particles C T from p ( C T ) , and then iteratively reﬁned through the reverse Markov kernels p θ ( C t − 1 |G , C t ) .

Original: Having formulated the reverse dynamics, the marginal likelihood can be calculated by p θ ( C 0 |G )
Revised: Having formulated the reverse dynamics, the marginal likelihood can be calculated by p θ ( C 0 |G )

Original: = (cid:82) p ( C T ) p θ ( C 0: T − 1 |G , C T )
Revised: = (cid:82) p ( C T ) p θ ( C 0: T − 1 |G , C T )

Original: d C 1: T .
Revised: d C 1: T .

Original: 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).
Revised: 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).

Original: In the following subsections, we will elaborate on how we parameterize the Markov kernels p θ ( C t − 1 |G , C t ) to achieve this desired property, and also how to maximize this likelihood by taking the invariance into account.
Revised: In the following subsections, we will elaborate on how we parameterize the Markov kernels p θ ( C t − 1 |G , C t ) to achieve this desired property, and also how to maximize this likelihood by taking the invariance into account.

Original: 4. E QUIVARIANT R EVERSE G ENERATIVE P ROCESS
Revised: 4. E QUIVARIANT R EVERSE G ENERATIVE P ROCESS

Original: Instead of directly leveraging existing methods, we consider building the density p θ ( C 0 ) that is invariant to rotation and translation transformations.
Revised: Instead of directly leveraging existing methods, we consider building the density p θ ( C 0 ) that is invariant to rotation and translation transformations.

Original: Intuitively, this requires the likelihood to be unaffected by translations and rotations.
Revised: Intuitively, this requires the likelihood to be unaffected by translations and rotations.

Original: Formally, let T g be some roto-translational transformations of a group element g ∈ SE(3), then we have the following statement:
Revised: Formally, let T g be some roto-translational transformations of a group element g ∈ SE(3), then we have the following statement:

Original: Proposition 1.
Revised: Proposition 1.

Original: Let p ( x T ) be an SE(3)-invariant density function, i.e., p ( x T )
Revised: Let p ( x T ) be an SE(3)-invariant density function, i.e., p ( x T )

Original: =
Revised: =

Original: p ( T g ( x T )) .
Revised: p ( T g ( x T )) .

Original: If
Revised: If

Original: Markov transitions p ( x t − 1 | x t ) are SE(3)-equivariant, i.e., p ( x t − 1 | x t ) =
Revised: Markov transitions p ( x t − 1 | x t ) are SE(3)-equivariant, i.e., p ( x t − 1 | x t ) =

Original: p ( T g ( x t − 1 ) | T g ( x t )) , then we have that the density p θ ( x 0 ) =
Revised: p ( T g ( x t − 1 ) | T g ( x t )) , then we have that the density p θ ( x 0 ) =

Original: (cid:82) p ( x T ) p θ ( x 0: T − 1 | x T )d x 1: T is also SE(3)-invariant.
Revised: (cid:82) p ( x T ) p θ ( x 0: T − 1 | x T )d x 1: T is also SE(3)-invariant.

Original: This proposition indicates that the dynamics starting from an invariant standard density along an equivariant Gaussian Markov kernel can result in an invariant density.
Revised: This proposition indicates that the dynamics starting from an invariant standard density along an equivariant Gaussian Markov kernel can result in an invariant density.

Original: Now we provide a practical implementation of G EO D IFF based on the recent denoising diffusion framework (Ho et al., 2020).
Revised: Now we provide a practical implementation of G EO D IFF based on the recent denoising diffusion framework (Ho et al., 2020).

Original: Invariant Initial Density p ( C T ) .
Revised: Invariant Initial Density p ( C T ) .

Original: We ﬁrst introduce the invariant distribution p ( C T ) , which will also be employed in the equivariant Markov chain.
Revised: We ﬁrst introduce the invariant distribution p ( C T ) , which will also be employed in the equivariant Markov chain.

Original: We borrow the idea from Köhler et al.
Revised: We borrow the idea from Köhler et al.

Original: (2020) to consider systems with zero center of mass (CoM), termed CoM-free systems.
Revised: (2020) to consider systems with zero center of mass (CoM), termed CoM-free systems.

Original: We deﬁne p ( C T ) as a “CoM-free standard density” ˆ ρ ( C ) , built upon an isotropic normal density ρ ( C ) : for evaluating the likelihood ˆ ρ ( C ) we can ﬁrstly translate C to zero CoM and then calculate ρ ( C ) , and for sampling from ˆ ρ ( C ) we can ﬁrst sample from ρ ( C ) and then move the CoM to zero.
Revised: We deﬁne p ( C T ) as a “CoM-free standard density” ˆ ρ ( C ) , built upon an isotropic normal density ρ ( C ) : for evaluating the likelihood ˆ ρ ( C ) we can ﬁrstly translate C to zero CoM and then calculate ρ ( C ) , and for sampling from ˆ ρ ( C ) we can ﬁrst sample from ρ ( C ) and then move the CoM to zero.

Original: 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.
Revised: 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.

Original: Consequently, ˆ ρ ( C ) is constructed as a roto-transitional invariant density.
Revised: Consequently, ˆ ρ ( C ) is constructed as a roto-transitional invariant density.

Original: Equivariant Markov Kernels p ( C t − 1 |G , C t ) .
Revised: Equivariant Markov Kernels p ( C t − 1 |G , C t ) .

Original: Similar to the prior density, we also consider equipping all intermediate structures C t as CoM-free systems.
Revised: Similar to the prior density, we also consider equipping all intermediate structures C t as CoM-free systems.

Original: Speciﬁcally, given mean µ θ ( G , C t , t ) and variance σ t , the likelihood of C t − 1 will be calculated by ˆ ρ ( C t − 1 − µ θ ( G , C t ,t ) σ t ) .
Revised: Speciﬁcally, given mean µ θ ( G , C t , t ) and variance σ t , the likelihood of C t − 1 will be calculated by ˆ ρ ( C t − 1 − µ θ ( G , C t ,t ) σ t ) .

Original: The CoM-free Gaussian ensures the translation invariance in the Markov kernels.
Revised: The CoM-free Gaussian ensures the translation invariance in the Markov kernels.

Original: Consequently, to achieve the equivariant property deﬁned in Proposition 1, we focus on the rotation equivariance.
Revised: Consequently, to achieve the equivariant property deﬁned in Proposition 1, we focus on the rotation equivariance.

Original: Then in general, the key requirement is to ensure the means µ θ ( G , C t , t ) to be roto-translation equivariant w.r.t C t .
Revised: Then in general, the key requirement is to ensure the means µ θ ( G , C t , t ) to be roto-translation equivariant w.r.t C t .

Original: Following Ho et al.
Revised: Following Ho et al.

Original: (2020), we consider the following parameterization of µ θ :
Revised: (2020), we consider the following parameterization of µ θ :

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) θ are neural networks with trainable parameters θ .
Revised: where (cid:15) θ are neural networks with trainable parameters θ .

Original: Intuitively, the model (cid:15) θ learns to predict the noise necessary to decorrupt the conformations.
Revised: Intuitively, the model (cid:15) θ learns to predict the noise necessary to decorrupt the conformations.

Original: This is analogous to the physical force ﬁelds (Schütt et al., 2017; Zhang et al., 2018;
Revised: This is analogous to the physical force ﬁelds (Schütt et al., 2017; Zhang et al., 2018;

Original: Hu et al., 2021; Shuaibi et al., 2021), which also gradually push particles towards convergence around the equilibrium states.
Revised: Hu et al., 2021; Shuaibi et al., 2021), which also gradually push particles towards convergence around the equilibrium states.

Original: Now the problem is transformed to constructing (cid:15) θ to be roto-translational equivariant.
Revised: Now the problem is transformed to constructing (cid:15) θ to be roto-translational equivariant.

Original: We draw inspirations from recent equivariant networks (Thomas et al., 2018;
Revised: We draw inspirations from recent equivariant networks (Thomas et al., 2018;

Original: Satorras et al., 2021b) to design an equivariant convolutional layer, named graph ﬁeld network (GFN).
Revised: Satorras et al., 2021b) to design an equivariant convolutional layer, named graph ﬁeld network (GFN).

Original: In the l -th layer, GFN takes node embeddings h l ∈ R n × b ( b denotes the feature dimension) and corresponding coordinate embeddings x l ∈ R n × 3 as inputs, and outputs h l +1 and x l +1 as follows:
Revised: In the l -th layer, GFN takes node embeddings h l ∈ R n × b ( b denotes the feature dimension) and corresponding coordinate embeddings x l ∈ R n × 3 as inputs, and outputs h l +1 and x l +1 as follows:

Original: [Equation]
Revised: [Equation]

Original: where Φ are feed-forward networks and d ij denotes interatomic distances.
Revised: where Φ are feed-forward networks and d ij denotes interatomic distances.

Original: N ( i ) denotes the neighborhood of i th 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.
Revised: N ( i ) denotes the neighborhood of i th 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.

Original: Initial embeddings h 0 are combinations of atom and timestep embeddings, and x 0 are atomic coordinates.
Revised: Initial embeddings h 0 are combinations of atom and timestep embeddings, and x 0 are atomic coordinates.

Original: 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 : R b
Revised: 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 : R b

Original: → R .
Revised: → R .

Original: Such vector ﬁeld x L enjoys the roto-translation equivariance property.
Revised: Such vector ﬁeld x L enjoys the roto-translation equivariance property.

Original: Formally, we have:
Revised: Formally, we have:

Original: Proposition 2.
Revised: Proposition 2.

Original: Parameterizing (cid:15) θ ( G , C , t ) as a composition of L GFN layers, and take the x L after
Revised: Parameterizing (cid:15) θ ( G , C , t ) as a composition of L GFN layers, and take the x L after

Original: L updates as the output.
Revised: L updates as the output.

Original: Then the noise vector ﬁeld (cid:15) θ is SE(3) equivariant w.r.t the 3D system C .
Revised: Then the noise vector ﬁeld (cid:15) θ is SE(3) equivariant w.r.t the 3D system C .

Original: Intuitively, given h l already invariant and x l equivariant, the message embedding m will also be invariant since it only depends on invariant features.
Revised: Intuitively, given h l already invariant and x l equivariant, the message embedding m will also be invariant since it only depends on invariant features.

Original: Since x is updated with the relative differences c
Revised: Since x is updated with the relative differences c

Original: i − c
Revised: i − c

Original: j weighted by invariant features, it will be translation-invariant and rotation-equivariant.
Revised: j weighted by invariant features, it will be translation-invariant and rotation-equivariant.

Original: Then inductively, composing (cid:15) θ with L GFN layers enables equivariance with C t .
Revised: Then inductively, composing (cid:15) θ with L GFN layers enables equivariance with C t .

Original: We provide the formal proof of equivariance properties in Appendix A.
Revised: We provide the formal proof of equivariance properties in Appendix A.

Original: 4. I MPROVED T RAINING O BJECTIVE
Revised: 4. I MPROVED T RAINING O BJECTIVE

Original: Having formulated the generative process and the model parameterization, now we consider the practical training objective for the reverse dynamics.
Revised: Having formulated the generative process and the model parameterization, now we consider the practical training objective for the reverse dynamics.

Original: Since directly optimizing the exact log-likelihood is intractable, we instead maximize the usual variational lower bound (ELBO) 2 :
Revised: Since directly optimizing the exact log-likelihood is intractable, we instead maximize the usual variational lower bound (ELBO) 3 :

Original: [Equation]
Revised: [Equation]

Original: 8) where q ( C t − 1 |C t , C 0 ) is analytically tractable as N ( 
Revised: where q ( C t − 1 |C t , C 0 ) is analytically tractable as N ( √ ¯ α

Original: [Equation]
Revised: [Equation]

Original: Most recently, Ho et al.
Revised: Most recently, Ho et al.

Original: (2020) showed that under the parameterization in equation 4, the ELBO of the diffusion model can be further simpliﬁed by calculating the KL divergences between Gaussians as weighted L 2 distances between the means (cid:15) θ and (cid:15) 2 .
Revised: (2020) showed that under the parameterization in equation 4, the ELBO of the diffusion model can be further simpliﬁed by calculating the KL divergences between Gaussians as weighted L 2 distances between the means (cid:15) θ and (cid:15) 3 .

Original: Formally, we have:
Revised: Formally, we have:

Original: Proposition 3.
Revised: Proposition 3.

Original: (Ho et al., 2020)
Revised: (Ho et al., 2020)

Original: Under the parameterization in equation 4, we have:
Revised: Under the parameterization in equation 4, we have:

Original: [Equation]
Revised: [Equation]

Original:  where C t = √ ¯ α t C 0 + √ 1 − ¯ α t (cid:15) .
Revised: (9) where C t = √ ¯ α t C 0 + √ 1 − ¯ α t (cid:15) .

Original: The weights γ t =
Revised: The weights γ t =

Original: [Equation]
Revised: [Equation]

Original: for [Equation] > 
Revised: for t > 1 , and γ 1 = 12 α 1 .

Original: [Equation] and γ [Equation] The intuition of this objective is to independently sample chaotic conformations of different timesteps from q ( C t − 1 |C t , C 0 ) , and use (cid:15) θ to model the noise vector (cid:15) .
Revised: The intuition of this objective is to independently sample chaotic conformations of different timesteps from q ( C t − 1 |C t , C 0 ) , and use (cid:15) θ to model the noise vector (cid:15) .

Original: To yield a better empirical performance,
Revised: To yield a better empirical performance,

Original: Ho et al.
Revised: Ho et al.

Original: (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).
Revised: (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).

Original: As (cid:15) θ is designed to be equivariant, it is natural to require its supervision signal (cid:15) to be equivariant with C t .
Revised: As (cid:15) θ is designed to be equivariant, it is natural to require its supervision signal (cid:15) to be equivariant with C t .

Original: Note that once this is achieved, the ELBO will also become invariant.
Revised: Note that once this is achieved, the ELBO will also become invariant.

Original: However, the (cid:15) in the forward diffusion process is not imposed with such equivariance, violating the above properties.
Revised: However, the (cid:15) in the forward diffusion process is not imposed with such equivariance, violating the above properties.

Original: Here we propose two approaches to obtain the modiﬁed noise vector ˆ (cid:15) , which, after replacing (cid:15) in the
Revised: Here we propose two approaches to obtain the modiﬁed noise vector ˆ (cid:15) , which, after replacing (cid:15) in the

Original: L 2 distance calculation in equation 9, achieves the desired equivariance:
Revised: L 2 distance calculation in equation 9, achieves the desired equivariance:

Original: Alignment approach .
Revised: Alignment approach .

Original: Considering the fact that (cid:15) can be calculated by C t −√ ¯ α t C 0 √ 1 − ¯ α t , we can ﬁrst rotate and translate C 0 to ˆ C 0 by aligning w.r.t C t , and then compute ˆ (cid:15) as C t −√ ¯ α t ˆ C 0 √ 1 − ¯ α t .
Revised: Considering the fact that (cid:15) can be calculated by C t −√ ¯ α t C 0 √ 1 − ¯ α t , we can ﬁrst rotate and translate C 0 to ˆ C 0 by aligning w.r.t C t , and then compute ˆ (cid:15) as C t − √ ¯ α t ˆ C 0 √ 1 − ¯ α t .

Original: Since the aligned conformation ˆ C 0 is equivariant with C t , the processed ˆ (cid:15) will also enjoy the equivariance.
Revised: Since the aligned conformation ˆ C 0 is equivariant with C t , the processed ˆ (cid:15) will also enjoy the equivariance.

Original: Speciﬁcally, the alignment is implemented by ﬁrst translating C 0 to the same CoM of C t and then solve the optimal rotation matrix by Kabsch alignment algorithm (Kabsch, 1976).
Revised: Speciﬁcally, the alignment is implemented by ﬁrst translating C 0 to the same CoM of C t and then solve the optimal rotation matrix by Kabsch alignment algorithm (Kabsch, 1976).

Original: Chain-rule approach .
Revised: Chain-rule approach .

Original: Another meaningful observation is that by reparameterizing the Gaussian distribution q ( C t |C 0 ) as C t = √ ¯ α t C 0 + √ 1 − ¯ α t (cid:15) , (cid:15) can be viewed as a weighted score function √ 1 − ¯ α t ∇ C t q ( C t |C 0 ) .
Revised: Another meaningful observation is that by reparameterizing the Gaussian distribution q ( C t |C 0 ) as C t = √ ¯ α t C 0 + √ 1 − ¯ α t (cid:15) , (cid:15) can be viewed as a weighted score function √ 1 − ¯ α t ∇ C t q ( C t |C 0 ) .

Original: Shi et al.
Revised: Shi et al.

Original: (2021) recently shows that generally this score function ∇ C t q ( C t |· ) can be designed to be equivariant by decomposing it into ∂ C t d t ∇ d t q ( C t |· ) with the chain rule, where d t can be any invariant features of the structures C t such as the inter-atomic distances.
Revised: (2021) recently shows that generally this score function ∇ C t q ( C t |· ) can be designed to be equivariant by decomposing it into ∂ C t d t ∇ d t q ( C t |· ) with the chain rule, where d t can be any invariant features of the structures C t such as the inter-atomic distances.

Original: We refer readers to Shi et al.
Revised: We refer readers to Shi et al.

Original: (2021) for more details.
Revised: (2021) for more details.

Original: The insight is that as gradient of invariant variables w.r.t equivariant variables, the partial derivative ∂ C t d t will always be equivalent with C t .
Revised: The insight is that as gradient of invariant variables w.r.t equivariant variables, the partial derivative ∂ C t d t will always be equivalent with C t .

Original: In this work, under the common assumption that d also follows a Gaussian distribution (Kingma & Welling, 2013), our practical implementation is to ﬁrst approximately calculate ∇ d t q ( C t |C 0 ) as d t −√ ¯ α t d 0 1 − ¯ α t , and then compute the modiﬁed noise vector ˆ (cid:15) as √ 1 − ¯ α t ∂ C t d t ( d t −√ ¯ α t d 0 1 − ¯ α t)
Revised: In this work, under the common assumption that d also follows a Gaussian distribution (Kingma & Welling, 2013), our practical implementation is to ﬁrst approximately calculate ∇ d t q ( C t |C 0 ) as d t −√ ¯ α t d 0 1 − ¯ α t , and then compute the modiﬁed noise vector ˆ (cid:15) as √ 1 − ¯ α t ∂ C t d t ( d t −√ ¯ α t d 0 1 − ¯ α t

Original: [Equation]
Revised: [Equation]

Original: t .

Original: 4. S AMPLING
Revised: 4. S AMPLING

Original: Input : the molecular graph G , the learned reverse model (cid:15) θ .
Revised: Input : the molecular graph G , the learned reverse model (cid:15) θ .

Original: Output : the molecular conformation C .
Revised: Output : the molecular conformation C .

Original: 1: Sample C T ∼ p ( C T ) =
Revised: 1: Sample C T ∼ p ( C T ) =

Original: N (0 , I ) 2: for s = T, T − 1 , · · · , 1 do 3:
Revised: N (0 , I ) 2: for s = T, T − 1 , · · · , 1 do 3:

Original: Shift C s to zero CoM 4:
Revised: Shift C s to zero CoM 4:

Original: Compute µ θ ( C s , G , s ) from (cid:15) θ ( C s , G , s ) using equation5:
Revised: Compute µ θ ( C s , G , s ) from (cid:15) θ ( C s , G , s ) using equation5:

Original: Sample C s − 1 ∼ N ( C s − 1 ; µ θ ( C s , G , s ) , σ 2 t I ) 6: end for 7: return C 0 as C
Revised: Sample C s − 1 ∼ N ( C s − 1 ; µ θ ( C s , G , s ) , σ 2 t I ) 6: end for 7: return C 0 as C

Original: With a learned reverse dynamics (cid:15) θ ( G , C t , t ) , the transition means µ θ ( G , C t , t ) can be calculated by equation 4.
Revised: With a learned reverse dynamics (cid:15) θ ( G , C t , t ) , the transition means µ θ ( G , C t , t ) can be calculated by equation 4.

Original: Thus, given a graph
Revised: Thus, given a graph

Original: G , its geometry C 0 is generated by ﬁrst sampling chaotic particles C T ∼ p ( C T ) , and then progressively sample C t − 1 ∼ p θ ( C t − 1 |G , C t ) for t =
Revised: G , its geometry C 0 is generated by ﬁrst sampling chaotic particles C T ∼ p ( C T ) , and then progressively sample C t − 1 ∼ p θ ( C t − 1 |G , C t ) for t =

Original: T, T − 1 , · · · , 1 .
Revised: T, T − 1 , · · · , 1 .

Original: This process is Markovian, which gradually shifts the previous noisy positions towards equilibrium states.
Revised: This process is Markovian, which gradually shifts the previous noisy positions towards equilibrium states.

Original: We provide the pseudo code of the whole sampling process in Algorithm 1.
Revised: We provide the pseudo code of the whole sampling process in Algorithm 1.

Original:  E XPERIMENT
Revised:  E XPERIMENT

Original: In this section, we empirically evaluate G EO D IFF on the task of equilibrium conformation generation for both small and drug-like molecules.
Revised: In this section, we empirically evaluate G EO D IFF on the task of equilibrium conformation generation for both small and drug-like molecules.

Original: 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:
Revised: 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:

Original: Conformation Generation (Sec. 5.2) and Property Prediction (Sec. 5.3).
Revised: Conformation Generation (Sec. 5.2) and Property Prediction (Sec. 5.3).

Original: We ﬁrst present the general experiment setups, and then describe task-speciﬁc evaluation protocols and discuss the results in each section.
Revised: We ﬁrst present the general experiment setups, and then describe task-speciﬁc evaluation protocols and discuss the results in each section.

Original: The implementation details are provided in Appendix C.
Revised: The implementation details are provided in Appendix C.

Original: 5. E XPERIMENT S ETUP
Revised: 5. E XPERIMENT S ETUP

Original: Datasets.
Revised: Datasets.

Original: 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.
Revised: 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.

Original: The former one contains small molecules while the latter one are medium-sized organic compounds.
Revised: The former one contains small molecules while the latter one are medium-sized organic compounds.

Original: We borrow the data split produced by Shi et al.
Revised: We borrow the data split produced by Shi et al.

Original: For both datasets, the training split consists of 40 ,molecules with 5 conformations for each, resulting in 200 , 000 conformations in total.
Revised: For both datasets, the training split consists of 40 ,molecules with 5 conformations for each, resulting in 200 , 000 conformations in total.

Original: The valid split share the same size as training split.
Revised: The valid split share the same size as training split.

Original: The test split contains 200 distinct molecules, with 22 ,conformations for QM9 and 14 , 324 ones for Drugs.
Revised: The test split contains 200 distinct molecules, with 22 ,conformations for QM9 and 14 , 324 ones for Drugs.

Original: Baselines.
Revised: Baselines.

Original: We compare G EO D IFF with 6 recent or established state-of-the-art baselines.
Revised: We compare G EO D IFF with 6 recent or established state-of-the-art baselines.

Original: For the ML approaches, we test the following models with highest reported performance: CVGAE (Mansimov et al., 2019), G RAPH DG (Simm & Hernandez-Lobato, 2020), CGCF (Xu et al., 2021a), C ONFVAE (Xu et al., 2021b) and C ONF GF (Shi et al., 2021).
Revised: For the ML approaches, we test the following models with highest reported performance: CVGAE (Mansimov et al., 2019), G RAPH DG (Simm & Hernandez-Lobato, 2020), CGCF (Xu et al., 2021a), C ONFVAE (Xu et al., 2021b) and C ONF GF (Shi et al., 2021).

Original: We also test the classic RDK IT (Riniker &
Revised: We also test the classic RDK IT (Riniker &

Original: Landrum, 2015) method, which is arguably the most popular open-source software for conformation generation.
Revised: Landrum, 2015) method, which is arguably the most popular open-source software for conformation generation.

Original: Some other recent works (TorsionNet (Gogineni et al., 2020) and G EO M OL (Ganea et al., 2021)) are omitted since these approaches are not directly comparable in our setting.

Original: We refer readers to Sec.
Revised: We refer readers to Sec.

Original: 2 for a detailed discussion of these models.
Revised: 2 for a detailed discussion of these models.

Original: 5. C ONFORMATION G ENERATION
Revised: 5. C ONFORMATION G ENERATION

Original: Evaluation metrics.
Revised: Evaluation metrics.

Original: The task aims to measure both quality and diversity of generated conformations by different models.
Revised: The task aims to measure both quality and diversity of generated conformations by different models.

Original: We follow Ganea et al.
Revised: We follow Ganea et al.

Original: (2021) to evaluate 4 metrics built upon root-mean-squaredeviation (RMSD), which is deﬁned as the normalized Frobenius norm of two atomic coordinates matrices, after alignment by Kabsch algorithm (Kabsch, 1976). Formally, let S g and S r denote the sets of generated and reference conformers respectively, then the Cov erage and Mat ching metrics (Xu [Table]
Revised: (2021) to evaluate 4 metrics built upon root-mean-square

Original: [Table]
Revised: [Table]

Original: A

Original: [Table]
Revised: [Table]

Original: G EO D IFF

Original: [Table]
Revised: [Table]

Original: C
Revised: (RMSD), which is deﬁned as the normalized Frobenius norm of two atomic coordinates matrices, after alignment by Kabsch algorithm (Kabsch, 1976).

Original: [Table] et al., 2021a) following the conventional Recall measurement can be deﬁned as:[Equation] [Equation]
Revised: Formally, let S g and S r denote the sets of generated and reference conformers respectively, then the Cov erage and Mat ching metrics (Xu et al., 2021a) following the conventional Recall measurement can be deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: where δ is a pre-deﬁned threshold.
Revised: where δ is a pre-deﬁned threshold.

Original: The other two metrics COV-P and MAT-P inspired by Precision can be deﬁned similarly but with the generated and reference sets exchanged.
Revised: The other two metrics COV-P and MAT-P inspired by Precision can be deﬁned similarly but with the generated and reference sets exchanged.

Original: In practice, S g is set as twice of the size of S r for each molecule.
Revised: In practice, S g is set as twice of the size of S r for each molecule.

Original: 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 δ .
Revised: 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 δ .

Original: By contrast, the MAT scores measure the average
Revised: By contrast, the MAT scores measure the average

Original: RMSD of conformers in one set with its closest neighbor in another set.
Revised: RMSD of conformers in one set with its closest neighbor in another set.

Original: In general, higher COV rates or lower MAT score suggest that more realistic conformations are generated.
Revised: In general, higher COV rates or lower MAT score suggest that more realistic conformations are generated.

Original: Besides, the Precision metrics depend more on the quality, while the Recall metrics concentrate more on the diversity.
Revised: Besides, the Precision metrics depend more on the quality, while the Recall metrics concentrate more on the diversity.

Original: Either metrics can be more appealing considering the speciﬁc scenario.
Revised: Either metrics can be more appealing considering the speciﬁc scenario.

Original: Following previous works (Xu et al., 2021a; Ganea et al., 2021), δ is set as 0 .
Revised: Following previous works (Xu et al., 2021a; Ganea et al., 2021), δ is set as 0 .

Original: 5 Å and 1 .
Revised: 5 Å and 1 .

Original: 25 Å for QM9 and Drugs datasets respectively.
Revised: 25 Å for QM9 and Drugs datasets respectively.

Original: Results & discussion.
Revised: Results & discussion.

Original: The results are reported in Tab. 1 and Tab. 2.
Revised: The results are summarized in Tab. 1 and Tab. 5 (left in Appendix. D).

Original: As noted in Sec.
Revised: As noted in Sec.

Original: 4.3, G EO D IFF can be trained with two types of modiﬁed ELBO, named alignment and chain-rule approaches.
Revised: 4.3, G EO D IFF can be trained with two types of modiﬁed ELBO, named alignment and chain-rule approaches.

Original: We denote models learned by these two objectives as G EO D IFF -A and G EO D IFF -C respectively.
Revised: We denote models learned by these two objectives as G EO D IFF -A and G EO D IFF -C respectively.

Original: As shown in the tables, G EO D IFF consistently outperform the state-of-the-art ML models on all datasets and metrics, especially by a signiﬁcant margin for more challenging large molecules (Drugs dataset).
Revised: As shown in the tables, G EO D IFF consistently outperform the state-of-theart ML models on all datasets and metrics, especially by a signiﬁcant margin for more challenging large molecules (Drugs dataset).

Original: The results demonstrate the superior capacity of G EO D IFF to model the multi modal distribution, and generative both accurate and diverse conformations.
Revised: The results demonstrate the superior capacity of G EO D IFF to model the multi modal distribution, and generative both accurate and diverse conformations.

Original: We also notice that in general G EO D IFF +C performs slightly better than G EO D IFF +A, which suggests that chain-rule approach leads to a better optimization procedure.
Revised: We also notice that in general G EO D IFF -C performs slightly better than G EO D IFF -A, which suggests that chain-rule approach leads to a better optimization procedure.

Original: We thus take G EO D IFF +C as the representative in the following comparisons.
Revised: We thus take G EO D IFF -C as the representative in the following comparisons.

Original: We visualize samples generated by different models in Fig.
Revised: We visualize samples generated by different models in Fig.

Original: 2 to provide a qualitative comparison, where G EO D IFF is shown to capture better both local and global structures.
Revised: 2 to provide a qualitative comparison, where G EO D IFF is shown to capture better both local and global structures.

Original: On the more challenging Drugs dataset, we further test RDK IT .
Revised: On the more challenging Drugs dataset, we further test RDK IT .

Original: 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.
Revised: As shown in Tab. 2, 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.

Original: However, for the new Precision -based metrics we found that
Revised: However, for the new Precision -based metrics we found that

Original: ML models are still not comparable.
Revised: ML models are still not comparable.

Original: This indicates that ML models tend to explore more possible representatives while RDK IT concentrates on a few most common ones, prioritizes quality over [Figure] [Table] G EO D IFF + FF
Revised: This indicates that ML models tend to explore more possible representatives while RDK IT concentrates on a few most common ones, prioritizes quality over diversity.

Original: [Table] diversity. Previous works (Mansimov et al., 2019; Xu et al., 2021b) suggest that this is because
Revised: Previous works (Mansimov et al., 2019; Xu et al., 2021b) suggest that this is because

Original: RDK IT involves an additional empirical force ﬁeld (FF)
Revised: RDK IT involves an additional empirical force ﬁeld (FF)

Original: (Halgren, 1996) to optimize the structure, and we follow them to also combine G EO D IFF with FF to yield a more fair comparison.
Revised: (Halgren, 1996) to optimize the structure, and we follow them to also combine G EO D IFF with FF to yield a more fair comparison.

Original: Results in
Revised: Results in

Revised: [Figure]

Revised: [Table]

Revised: G

Original:  Tab. 3 demonstrate that G EO D IFF +FF can keep the superior diversity ( Recall metrics) while also enjoy signiﬁcantly improved accuracy (( Precision metrics))
Revised: [Equation] IFF + FF [Table] Tab. 2 demonstrate that G EO D IFF +FF can keep the superior diversity ( Recall metrics) while also enjoy signiﬁcantly improved accuracy (( Precision metrics))

Original: [Table]
Revised: [Table]

Original: 5. P ROPERTY P REDICTION
Revised: 5. P ROPERTY P REDICTION

Original: [Table]
Revised: [Table]

Original: Evaluation metrics.
Revised: Evaluation metrics.

Original: This task estimates the molecular ensemble properties (Axelrod & Gomez-Bombarelli, 2020) over a set of generated conformations.
Revised: This task estimates the molecular ensemble properties (Axelrod & Gomez-Bombarelli, 2020) over a set of generated conformations.

Original: This can provide an direct assessment on the quality of generated samples.
Revised: This can provide an direct assessment on the quality of generated samples.

Original: In speciﬁc, we follow (Shi et al., 2021) to extract a split from
Revised: In speciﬁc, we follow Shi et al. (2021) to extract a split from

Original: GEOM-QM9 covering 30 molecules, and generate 50 samples for each.
Revised: GEOM-QM9 covering 30 molecules, and generate 50 samples for each.

Original: Then we use the chemical toolkit P SI 4 (Smith et al., 2020) to calculate each conformer’s energy E and HOMOLUMO gap (cid:15) , and compare the average energy E , lowest energy E min , average gap ∆ (cid:15) , minimum gap ∆ (cid:15) min , and maximum gap ∆ (cid:15) max with the ground truth.
Revised: Then we use the chemical toolkit P SI 4 (Smith et al., 2020) to calculate each conformer’s energy E and HOMOLUMO gap (cid:15) , and compare the average energy E , lowest energy E min , average gap ∆ (cid:15) , minimum gap ∆ (cid:15) min , and maximum gap ∆ (cid:15) max with the ground truth.

Original: Results & discussions.
Revised: Results & discussions.

Original: The mean absolute errors (MAE) between calculated properties and the ground truth are reported in Tab. 4.
Revised: The mean absolute errors (MAE) between calculated properties and the ground truth are reported in Tab. 3.

Original: CVGAE is excluded due to the poor performance, which is also reported in Simm & Hernandez-Lobato (2020); Shi et al.
Revised: CVGAE is excluded due to the poor performance, which is also reported in Simm & Hernandez-Lobato (2020); Shi et al.

Original: The properties are highly sensitive to geometric structure, and thus the superior performance demonstrate that G EO D IFF can consistently predict more accurate conformations across different molecules.
Revised: The properties are highly sensitive to geometric structure, and thus the superior performance demonstrate that G EO D IFF can consistently predict more accurate conformations across different molecules.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We propose G EO D IFF , a novel probabilistic model for generating molecular conformations.
Revised: We propose G EO D IFF , a novel probabilistic model for generating molecular conformations.

Original: G EO D IFF 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.
Revised: G EO D IFF 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.

Original: We derive a tractable invariant objective from the variational lower bound to optimize the likelihood.
Revised: We derive a tractable invariant objective from the variational lower bound to optimize the likelihood.

Original: Comprehensive experiments over multiple tasks demonstrate that G EO D IFF is competitive with the existing state-of-the-art models.
Revised: Comprehensive experiments over multiple tasks demonstrate that G EO D IFF is competitive with the existing state-of-the-art models.

Original: 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. 
Revised: 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. A CKNOWLEDGEMENT

Revised: Minkai thanks Huiyu Cai, David Wipf, Zuobai Zhang and Zhaocheng Zhu for their helpful discussions and comments.

Revised: This project is supported by the Natural Sciences and Engineering Research

Revised: Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award,

Original:  A P ROOFS A. P ROPERTIES OF THE D IFFUSION M ODEL
Revised: Tencent AI Lab Rhino-Bird Gift Fund and a NRC Collaborative R&D Project (AI4D-CORE-06) This project was also partially funded by IVADO Fundamental Research Project grant PRF-20193583139727. The Stanford team is supported by NSF(#1651565, #1522054, #1733686), ONR (N000141912145), AFOSR (FA95501910024), ARO (W911NF-21-1-0125) and Sloan Fellowship. A P ROOFS A. P ROPERTIES OF THE D IFFUSION M ODEL

Original: We include proofs for several key properties of the probabilistic diffusion model here to be selfcontained.
Revised: We include proofs for several key properties of the probabilistic diffusion model here to be selfcontained.

Original: For more detailed discussions, please refer to Ho et al.
Revised: For more detailed discussions, please refer to Ho et al.

Original: Let { β 0 , ..., β T } be a sequence of variances, and α t = 1 − β t and ¯ α t = (cid:81) ts =1 α s .
Revised: Let { β 0 , ..., β T } be a sequence of variances, and α t = 1 − β t and ¯ α t = (cid:81) ts =1 α s .

Original: The two following properties are crucial for deriving the ﬁnal tractable objective in equation 9.
Revised: The two following properties are crucial for deriving the ﬁnal tractable objective in equation 9.

Original: Property 1.
Revised: Property 1.

Original: Tractable marginal of the forward process:
Revised: Tractable marginal of the forward process:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Let (cid:15)
Revised: Let (cid:15)

Original: i ’s be independent standard Gaussian random variables.
Revised: i ’s be independent standard Gaussian random variables.

Original: Then, by deﬁnition of the
Revised: Then, by deﬁnition of the

Original: Markov kernels q ( C t |C t − 1 ) in equation 2, we have
Revised: Markov kernels q ( C t |C t − 1 ) in equation 2, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore q ( C t |C 0 ) is still Gaussian, and the mean of C t is √ ¯ α t C 0 , and the variance matrix is ( α t α t − 1 · · · α 2 β 1 + · · · + α t β t − 1 + β t )
Revised: Therefore q ( C t |C 0 ) is still Gaussian, and the mean of C t is √ ¯ α t C 0 , and the variance matrix is ( α t α t − 1 · · · α 2 β 1 + · · · + α t β t − 1 + β t )

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: 1 − ¯ α t ) I .
Revised: 1 − ¯ α t ) I .

Original: Then we have:
Revised: Then we have:

Original: [Equation]
Revised: [Equation]

Original: This property provides convenient closed-form evaluation of C t knowing C 0 :
Revised: This property provides convenient closed-form evaluation of C t knowing C 0 :

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) ∼ N (0 , I ) .
Revised: where (cid:15) ∼ N (0 , I ) .

Original: Besides, it is worth noting that,
Revised: Besides, it is worth noting that,

Original: [Equation]
Revised: [Equation]

Original: where ¯ α T = (cid:81) Tt =1 (1 − β t ) approaches zero with large T , which indicates the diffusion process can ﬁnally converge into a whitened noisy distribution [Equation]
Revised: where ¯ α T = (cid:81) Tt =1 (1 − β t ) approaches zero with large T , which indicates the diffusion process can ﬁnally converge into a whitened noisy distribution.

Original: Property 2.
Revised: Property 2.

Original: Tractable posterior of the forward process:
Revised: Tractable posterior of the forward process:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Let ˜ β t = 1 − ¯ α t − 1 1 − ¯ α t β t , then we can derive the posterior by Bayes rule:
Revised: Let ˜ β t = 1 − ¯ α t − 1 1 − ¯ α t β t , then we can derive the posterior by Bayes rule:

Original: [Equation]
Revised: [Equation]

Original: A. P ROOF OF P ROPOSITION
Revised: A. P ROOF OF P ROPOSITION

Original: Let T g be some roto-translational transformations of a group element g ∈ SE(3), and let p ( x T ) be a density which is SE(3)-invariant,
Revised: Let T g be some roto-translational transformations of a group element g ∈ SE(3), and let p ( x T ) be a density which is SE(3)-invariant,

Original: i.e. , p ( x T ) =
Revised: i.e. , p ( x T ) =

Original: p ( T g ( x T )) .
Revised: p ( T g ( x T )) .

Original: If the Markov transitions p ( x t − 1 | x t ) are
Revised: If the Markov transitions p ( x t − 1 | x t ) are

Original: SE(3)-equivariant, i.e. , p ( x t − 1 | x t ) =
Revised: SE(3)-equivariant, i.e. , p ( x t − 1 | x t ) =

Original: p ( T g ( x t − 1 ) | T g ( x t )) , then we have that the density p θ ( x 0 ) =
Revised: p ( T g ( x t − 1 ) | T g ( x t )) , then we have that the density p θ ( x 0 ) =

Original: (cid:82) p ( x T ) p θ ( x 0: T − 1 | x T )d x 1: T is also SE(3)-invariant.
Revised: (cid:82) p ( x T ) p θ ( x 0: T − 1 | x T )d x 1: T is also SE(3)-invariant.

Original: Proof.
Revised: Proof.

Original: [Equation]

Original: invariant

Original: [Equation]
Revised: [Equation]

Original: equivariant kernels
Revised: equivariant kernels

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: A. P ROOF OF P ROPOSITION
Revised: A. P ROOF OF P ROPOSITION

Original: In this section we prove that the output x of GFN deﬁned in equation 5, 6 and 7 is translationally invariant and rotationally equivariant with the input C .
Revised: In this section we prove that the output x of GFN deﬁned in equation 5, 6 and 7 is translationally invariant and rotationally equivariant with the input C .

Original: Let g ∈ R 3 denote any translation transformations and orthogonal matrices R ∈ R 3 × 3 denote any rotation transformations.
Revised: Let g ∈ R 3 denote any translation transformations and orthogonal matrices R ∈ R 3 × 3 denote any rotation transformations.

Original: let R x be shorthand for ( R x 1 , · · · , R x N ) .
Revised: let R x be shorthand for ( R x 1 , · · · , R x N ) .

Original: Formally, we aim to prove that the model satisﬁes:
Revised: Formally, we aim to prove that the model satisﬁes:

Original: [Equation]
Revised: [Equation]

Original: This equation indicates that, given x l already rotationally equivalent with C , and h l already invariant, then such property can propagate through a single GFN layer to x l +1 and h l +1 .
Revised: This equation indicates that, given x l already rotationally equivalent with C , and h l already invariant, then such property can propagate through a single GFN layer to x l +1 and h l +1 .

Original: Proof.
Revised: Proof.

Original: Firstly, given that h l already invariant to SE(3) transformations, we have that the messages m ij calculated from equation 5 will also be invariant.
Revised: Firstly, given that h l already invariant to SE(3) transformations, we have that the messages m ij calculated from equation 5 will also be invariant.

Original: This is because it sorely relies on the distance between two atoms, which are manifestly invariant to rotations (cid:107) R x li − R x lj (cid:107) 2 =
Revised: This is because it sorely relies on the distance between two atoms, which are manifestly invariant to rotations (cid:107) R x li − R x lj (cid:107) 2 =

Original: ( x li − x lj ) (cid:62) R
Revised: ( x li − x [Equation] [Equation] ) =

Original: [Equation]
Revised: [Equation]

Original: R ( x li −x lj ) = ( x li − x
Revised: −

Original: [Equation]
Revised: [Equation]

Original: ) (cid:62) I ( xli − x lj )
Revised: ) (cid:62) I [Equation] x

Revised: [Equation]

Original:  = (cid:107) x li − x lj (cid:107) 2 .
Revised: ) = (cid:107) x li − x lj (cid:107) 2 .

Original: Formally, the invariance of messages in equation 5 can be written as:
Revised: Formally, the invariance of messages in equation 5 can be written as:

Original: [Equation]
Revised: [Equation]

Original: And similarly, the h t +1 updated from equation 6 will also be invariant.
Revised: And similarly, the h t +1 updated from equation 6 will also be invariant.

Original: Next, we prove that the vector x updated from equation 7 preserves rotational equivariance and translational invariance.
Revised: Next, we prove that the vector x updated from equation 7 preserves rotational equivariance and translational invariance.

Original: Given m ij already invariant as proven above, we have that:
Revised: Given m ij already invariant as proven above, we have that:

Original: [Equation]
Revised: [Equation]

Original: Therefore, we have that rotating and translating c results in the same rotation and no translation on x l +1 by updating through equation 7.
Revised: Therefore, we have that rotating and translating c results in the same rotation and no translation on x l +1 by updating through equation 7.

Original: Thus we can conclude that the property deﬁned in equation 15 is satisﬁed.
Revised: Thus we can conclude that the property deﬁned in equation 15 is satisﬁed.

Original: 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.
Revised: 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.

Original: A. P ROOF OF P ROPOSITION
Revised: A. P ROOF OF P ROPOSITION

Original: We ﬁrst derive the variational lower bound (ELBO) objective in equation 8.
Revised: We ﬁrst derive the variational lower bound (ELBO) objective in equation 8.

Original: The ELBO can be calculated as follows:
Revised: The ELBO can be calculated as follows:

Original: [Equation]
Revised: [Equation]

Revised: −

Original: [Equation]
Revised: [Equation]

Original: It can be noted that the ﬁrst term KL (cid:0) q ( C T |C 0 ) (cid:107) p ( C T ) (cid:1) is a constant, which can be omitted in the objective.
Revised: It can be noted that the ﬁrst term KL (cid:0) q ( C T |C 0 ) (cid:107) p ( C T ) (cid:1) is a constant, which can be omitted in the objective.

Original: Furthermore, for brevity, we also merge the ﬁnal term log p θ ( C 0 |G , C 1 ) into the second term (sum over KL divergences), and ﬁnally derive that L ELBO =
Revised: Furthermore, for brevity, we also merge the ﬁnal term log p θ ( C 0 |G , C 1 ) into the second term (sum over KL divergences), and ﬁnally derive that L ELBO =

Original: (cid:80) Tt =1 D KL ( q ( C t − 1 |C t , C 0 ) (cid:107) p
Revised: (cid:80) Tt =1 D KL ( q ( C t − 1 |C t , C 0 ) (cid:107) p

Original: θ ( C t − 1 |G , C t ))
Revised: θ ( C t − 1 |G , C t ))

Original: as in equation 8.
Revised: as in equation 8.

Original: Now we consider how to compute the KL divergences as the proposition 3.
Revised: Now we consider how to compute the KL divergences as the proposition 3.

Original: Since both q ( C t − 1 |C t , C 0 ) and p θ ( C t − 1 |G , C t ) are Gaussian share the same covariance matrix ˜
Revised: Since both q ( C t − 1 |C t , C 0 ) and p θ ( C t − 1 |G , C t ) are Gaussian share the same covariance matrix ˜

Original: β t I , the KL divergence between them can be calculated by the squared (cid:96) 2 distance between their means weighed by a certain weights 12˜ β t .
Revised: β t I , the KL divergence between them can be calculated by the squared (cid:96) 2 distance between their means weighed by a certain weights 12˜ β t .

Original: By the expression of q ( C t |C 0 ) , we have the reparameterization that C t = √ ¯ α
Revised: By the expression of q ( C t |C 0 ) , we have the reparameterization that C t = √ ¯ α

Original: [Equation]
Revised: [Equation]

Original: (cid:15) .
Revised: (cid:15) .

Original: Then we can derive:
Revised: Then we can derive:

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where γ t represent the wights βt 2 α t
Revised: where γ t represent the wights β

Original: [Equation]
Revised: [Equation]

Original: 1 −

Original: [Equation]

Original: α t − 1 ) .

Original: And we ﬁnish the proof.
Revised: And we ﬁnish the proof.

Original: A. A NALYSIS OF THE INVARIANT DENSITY IN S EC . 4. 
Revised: A. A NALYSIS OF THE INVARIANT DENSITY IN S EC . 4.

Original: Given a geometric system x ∈ R N · 3 , we obtain the CoM-free ˆ x by subtracting its CoM. This can be considered as a linear transformation:
Revised: Given a geometric system x ∈ R N · 3 , we obtain the CoM-free ˆ x by subtracting its CoM. This can be considered as a linear transformation:

Original: [Equation]
Revised: [Equation]

Original: (20) where I k denotes the k × k identity matrix and 1 k denotes the k -dimensional vector ﬁlled with ones.
Revised: (20) where I [Table]

Original: It can be noted that Q is a symmetric projection operator, i.e. , Q 2 = Q and Q T = Q .
Revised: It can be noted that Q is a symmetric projection operator, i.e. , Q 2 = Q and Q T = Q .

Original: And we also have that rank[ Q ] =
Revised: And we also have that rank[ Q ] =

Original: ( N − 1) · 3 .
Revised: ( N − 1) · 3 .

Original: 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.
Revised: 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.

Original: Formally, let n = N · 3 and set R n with an isotropic normal distribution ρ = N (0 , I n ) , then the
Revised: Formally, let n = N · 3 and set R n with an isotropic normal distribution ρ = N (0 , I n ) , then the

Original: CoM-free density can be formally written as ˆ ρ = N (0 , QI n Q T ) =
Revised: CoM-free density can be formally written as ˆ ρ = N (0 , QI n Q T ) =

Original: N (0 , QQ T ) .
Revised: N (0 , QQ T ) .

Original: Thus, sampling from ˆ ρ can be trivially achieved by sampling from ρ and then projecting with Q .
Revised: Thus, sampling from ˆ ρ can be trivially achieved by sampling from ρ and then projecting with Q .

Original: And ˆ ρ ( y ) can be calculated by ρ ( y ) since for any y ∈ U we have (cid:107) y (cid:107) 

Revised: And ˆ ρ ( y ) can be calculated by ρ ( y ) since for any y ∈ U we have (cid:107) y (cid:107)

Original: = (cid:107) Qy (cid:107) 22 , and thus ρ ( y )
Revised: = (cid:107) Qy (cid:107) 22 , and thus ρ ( y )

Original: =
Revised: =

Original: ˆ ρ ( y ) .
Revised: ˆ ρ ( y ) .

Original: 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.
Revised: 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.

Original: And thus we can induce a well-deﬁned Markov chain on the subspace spanned by Q .
Revised: And thus we can induce a well-deﬁned Markov chain on the subspace spanned by Q .

Original: B O THER RELATED WORK
Revised: B O THER RELATED WORK

Original: Protein structure generation.
Revised: Protein structure generation.

Original: There has also been many recent works working on protein structure folding.
Revised: There has also been many recent works working on protein structure folding.

Original: For example, Boltzmann generators Noé et al.
Revised: For example, Boltzmann generators Noé et al.

Original: (2019) use ﬂow-based models to generate the structure of protein main chains.
Revised: (2019) use ﬂow-based models to generate the structure of protein main chains.

Original: AlQuraishi (2019) uses recurrent networks to model the amino acid sequences.
Revised: AlQuraishi (2019) uses recurrent networks to model the amino acid sequences.

Original: Ingraham et al.
Revised: Ingraham et al.

Original: (2019) proposed neural networks to learn an energy simulator to infer the protein structures.
Revised: (2019) proposed neural networks to learn an energy simulator to infer the protein structures.

Original: Most recently, AlphaFold Senior et al.
Revised: Most recently, AlphaFold Senior et al.

Original: ; Jumper et al.
Revised: ; Jumper et al.

Original: (2021) has signiﬁcantly improved the performance of protein structure generation.
Revised: (2021) has signiﬁcantly improved the performance of protein structure generation.

Original: Nevertheless, proteins are mainly linear backbone structures while general molecules are highly branched with various rings, making protein folding approaches unsuitable for our setting.
Revised: Nevertheless, proteins are mainly linear backbone structures while general molecules are highly branched with various rings, making protein folding approaches unsuitable for our setting.

Original: Point cloud generation.
Revised: Point cloud generation.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: C E XPERIMENT DETAILS
Revised: C E XPERIMENT DETAILS

Original: In this section, we introduce the details of our experiments.
Revised: In this section, we introduce the details of our experiments.

Revised: In practice, the means (cid:15) θ are parameterized as compositions of both typical invariant MPNNs (Schütt et al., 2017) and the proposed equivariant

Revised: GFNs in Sec.

Original: As a default setup, the MPNNs for parameterizing the means (cid:15) θ are all implemented with 4 layers, and the hidden embedding dimension is set as 128 .
Revised: As a default setup, the MPNNs for parameterizing the means (cid:15) θ are all implemented with 4 layers, and the hidden embedding dimension is set as 128 .

Revised: After the MPNNs, we can obtain the informative invariant atom embeddings, which we denote as h 0 .

Revised: Then the embeddings h 0 are fed into equivariant layers and updated with equation 5, equation 6, and equation 7 to obtain the equivariant output.

Original: For the training of G EO D IFF , we train the model on a single Tesla V100 GPU with a learning rate of 0 .
Revised: For the training of G EO D IFF , we train the model on a single Tesla V100 GPU with a learning rate of 0 .

Original: 001 until convergence and Adam (Kingma & Welling, 2013) as the optimizer.
Revised: 001 until convergence and Adam (Kingma & Welling, 2013) as the optimizer.

Original: The  other hyper-parameters of G EO D IFF 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.
Revised: The practical training time is ~48 hours. The other hyper-parameters of G EO D IFF are summarized in Tab. 4, 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.

Original: For sampling, the decoding time is in line with the most competitive baseline model C ONF GF (Shi et al., 2021) and requires around 8500 and 11500 seconds to decode the entire QM9 and Drugs test sets respectively.
Revised: [Table] D A DDITIONAL EXPERIMENTS D. R ESULTS FOR GEOM-QM  The results on the GEOM-QM9 dataset are reported in Tab. 5. [Table]

Original: [Table]
Revised: [Table]

Original: D E XTRA EXPERIMENTS
Revised: D. A BLATION STUDY WITH FEWER DIFFUSION STEPS

Original: We also test our method with fewer diffusion steps.
Revised: We also test our method with fewer diffusion steps.

Original: Speciﬁcally, we test the setting with T = 1000 , β 1 = 1e-7 and β T =
Revised: Speciﬁcally, we test the setting with T = 1000 , β 1 = 1e-7 and β T =

Original: 9e-3.
Revised: 9e-3.

Original: The results on the more challenging Drugs dataset are shown in Tab. 6.
Revised: The results on the more challenging Drugs dataset are shown in Tab. 6.

Original: 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.
Revised: 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.

Original: Note that, the most competitive baseline C ONF GF (Shi et al., 2021) also requires 5000 sampling steps, which indicates that our model can achieve better performance with fewer computational costs [Table]
Revised: Note that, the most competitive baseline C ONF GF (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.

Original: E M ORE V ISUALIZATIONS
Revised: E M ORE V ISUALIZATIONS

Original: We provide more visualization of generated structures in Fig.
Revised: We provide more visualization of generated structures in Fig.

Original: The molecules are chosen from the test split of GEOM-Drugs dataset.
Revised: The molecules are chosen from the test split of GEOM-Drugs dataset.

Original: [Figure]
Revised: [Figure]

Original: Conformations