Original: E QUIVARIANT E NERGY -G UIDED SDE FOR I NVERSE M OLECULAR D ESIGN
Revised: E QUIVARIANT E NERGY -G UIDED SDE FOR I NVERSE M OLECULAR D ESIGN

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Inverse molecular design is critical in material science and drug discovery, where the generated molecules should satisfy certain desirable properties.
Revised: Inverse molecular design is critical in material science and drug discovery, where the generated molecules should satisfy certain desirable properties.

Original: In this paper, we propose equivariant energy-guided stochastic differential equations (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.
Revised: In this paper, we propose equivariant energy-guided stochastic differential equations (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.

Original: Formally, we show that
Revised: Formally, we show that

Original: EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.
Revised: EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.

Original: Empirically, under the guidance of designed energy functions, EEGSDE significantly improves the baseline on QM9, in inverse molecular design targeted to quantum properties and molecular structures.
Revised: Empirically, under the guidance of designed energy functions, EEGSDE significantly improves the baseline on QM9, in inverse molecular design targeted to quantum properties and molecular structures.

Original: Furthermore, EEGSDE is able to generate molecules with multiple target properties by combining the corresponding energy functions linearly.
Revised: Furthermore, EEGSDE is able to generate molecules with multiple target properties by combining the corresponding energy functions linearly.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The discovery of new molecules with desired properties is critical in many fields, such as the drug and material design (Hajduk & Greer, 2007; Mandal et al., 2009; Kang et al., 2006;
Revised: The discovery of new molecules with desired properties is critical in many fields, such as the drug and material design (Hajduk & Greer, 2007; Mandal et al., 2009; Kang et al., 2006;

Original: Pyzer-Knapp et al., 2015).
Revised: Pyzer-Knapp et al., 2015).

Original: However, brute-force search in the overwhelming molecular space is extremely challenging.
Revised: However, brute-force search in the overwhelming molecular space is extremely challenging.

Original: Recently, inverse molecular design (Zunger, 2018) provides an efficient way to explore the molecular space, which directly predicts promising molecules that exhibit desired properties.
Revised: Recently, inverse molecular design (Zunger, 2018) provides an efficient way to explore the molecular space, which directly predicts promising molecules that exhibit desired properties.

Original: A natural way of inverse molecular design is to train a conditional generative model (SanchezLengeling & Aspuru-Guzik, 2018).
Revised: A natural way of inverse molecular design is to train a conditional generative model (SanchezLengeling & Aspuru-Guzik, 2018).

Original: Formally, it learns a distribution of molecules conditioned on certain properties from data, and new molecules are predicted by sampling from the distribution with the condition set to desired properties.
Revised: Formally, it learns a distribution of molecules conditioned on certain properties from data, and new molecules are predicted by sampling from the distribution with the condition set to desired properties.

Original: Among them, equivariant diffusion models (EDM) (Hoogeboom et al., 2022) leverage the current state-of-art diffusion models (Ho et al., 2020), which involves a forward process to perturb data and a reverse process to generate 3D molecules conditionally or unconditionally.
Revised: Among them, equivariant diffusion models (EDM) (Hoogeboom et al., 2022) leverage the current state-of-art diffusion models (Ho et al., 2020), which involves a forward process to perturb data and a reverse process to generate 3D molecules conditionally or unconditionally.

Original: While EDM generates stable and valid 3D molecules, we argue that a single conditional generative model is insufficient for generating accurate molecules that exhibit desired properties (see Table 1 and Table 2 for an empirical verification).
Revised: While EDM generates stable and valid 3D molecules, we argue that a single conditional generative model is insufficient for generating accurate molecules that exhibit desired properties (see Table 1 and Table 2 for an empirical verification).

Original: In this work, we propose equivariant energy-guided stochastic differential equations (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.
Revised: In this work, we propose equivariant energy-guided stochastic differential equations (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.

Original: EEGSDE formalizes the generation process as an equivariant stochastic differential equation, and plug in energy functions to improve the controllability of generation.
Revised: EEGSDE formalizes the generation process as an equivariant stochastic differential equation, and plugs in energy functions to improve the controllability of generation.

Original: Formally, we show that EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.
Revised: Formally, we show that EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.

Original: We apply EEGSDE to various applications by carefully designing task-specific energy functions.
Revised: We apply EEGSDE to various applications by carefully designing task-specific energy functions.

Original: When targeted to quantum properties, EEGSDE is able to generate more accurate molecules than
Revised: When targeted to quantum properties, EEGSDE is able to generate more accurate molecules than

Original: EDM, e.g., reducing the mean absolute error by more than 30% on the dipole moment property.
Revised: EDM, e.g., reducing the mean absolute error by more than 30% on the dipole moment property.

Original: When targeted to specific molecular structures, EEGSDE better capture the structure information in molecules than EDM, e.g, improving the similarity to target structures by more than 10%.
Revised: When targeted to specific molecular structures, EEGSDE better capture the structure information in molecules than EDM, e.g, improving the similarity to target structures by more than 10%.

Original: Furthermore, EEGSDE is able to generate molecules targeted to multiple properties by combining the corresponding energy functions linearly.
Revised: Furthermore, EEGSDE is able to generate molecules targeted to multiple properties by combining the corresponding energy functions linearly.

Original: These demonstrate that our EEGSDE enables a flexible and controllable generation of molecules, providing a smart way to explore the chemical space.
Revised: These demonstrate that our EEGSDE enables a flexible and controllable generation of molecules, providing a smart way to explore the chemical space.

Original: [Figure]
Revised: [Figure]

Original: As the energy function is invariant to rotational transformation R , its gradient (i.e., the energy guidance) is equivariant to R , and therefore the distribution of generated samples is invariant to R .
Revised: As the energy function is invariant to rotational transformation R , its gradient (i.e., the energy guidance) is equivariant to R , and therefore the distribution of generated samples is invariant to R .

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

Original: Diffusion models are initially proposed by Sohl-Dickstein et al. (2015).
Revised: Diffusion models are initially proposed by Sohl-Dickstein et al. (2015).

Original: Recently, they are better understood in theory by connecting it to score matching and stochastic differential equations (SDE) (Ho et al., 2020; Song et al., 2020).
Revised: Recently, they are better understood in theory by connecting it to score matching and stochastic differential equations (SDE) (Ho et al., 2020; Song et al., 2020).

Original: After that, diffusion models have shown strong empirical performance in many applications Dhariwal & Nichol (2021); Ramesh et al.
Revised: After that, diffusion models have shown strong empirical performance in many applications Dhariwal & Nichol (2021); Ramesh et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2020); Kong et al.
Revised: (2020); Kong et al.

Original: There are also variants proposed to improve or accelerate diffusion models (Nichol & Dhariwal, 2021; Vahdat et al., 2021; Dockhorn et al., 2021;
Revised: There are also variants proposed to improve or accelerate diffusion models (Nichol & Dhariwal, 2021; Vahdat et al., 2021; Dockhorn et al., 2021;

Original: Bao et al., 2022b;a;
Revised: Bao et al., 2022b;a;

Original: Salimans & Ho, 2022; Lu et al., 2022).
Revised: Salimans & Ho, 2022; Lu et al., 2022).

Original: Guidance is a technique to control the generation process of diffusion models.
Revised: Guidance is a technique to control the generation process of diffusion models.

Original: Initially, Song et al.
Revised: Initially, Song et al.

Original: (2020); Dhariwal & Nichol (2021) use classifier guidance to generate samples belonging to a class.
Revised: (2020); Dhariwal & Nichol (2021) use classifier guidance to generate samples belonging to a class.

Original: Then, the guidance is extended to CLIP (Radford et al., 2021) for text to image generation, and semantic-aware energy
Revised: Then, the guidance is extended to CLIP (Radford et al., 2021) for text to image generation, and semantic-aware energy

Original: (Zhao et al., 2022) for image-to-image translation.
Revised: (Zhao et al., 2022) for image-to-image translation.

Original: Prior guidance methods focus on image data, and are nontrivial to apply to molecules, since they do not consider the geometric symmetry.
Revised: Prior guidance methods focus on image data, and are nontrivial to apply to molecules, since they do not consider the geometric symmetry.

Original: In contrast, our work proposes a general guidance framework for 3D molecules, which leverage the geometric symmetry of molecules.
Revised: In contrast, our work proposes a general guidance framework for 3D molecules, where an invariant energy function is employed to leverage the geometric symmetry of molecules.

Original: Molecule generation.
Revised: Molecule generation.

Original: Several works attempt to model molecules as 3D objects via deep generative models Nesterov et al.
Revised: Several works attempt to model molecules as 3D objects via deep generative models Nesterov et al.

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

Original: ; Satorras et al.
Revised: ; Satorras et al.

Original: (2021a)
Revised: (2021a)

Original: ; Hoffmann & No´e (2019); Hoogeboom et al.
Revised: Hoffmann & No ´ e (2019); Hoogeboom et al.

Original: Among them, the most relevant one is the equivariant diffusion model (EDM) (Hoogeboom et al., 2022), which generates molecules in an iterative denoising manner.
Revised: Among them, the most relevant one is the equivariant diffusion model (EDM) (Hoogeboom et al., 2022), which generates molecules in an iterative denoising manner.

Original: Benefiting from recent advances of diffusion models, EDM is stable to train and is able to generate high quality molecules.
Revised: Benefiting from recent advances of diffusion models, EDM is stable to train and is able to generate high quality molecules.

Original: We provide a formal description of EDM in Section 3.
Revised: We provide a formal description of EDM in Section 3.

Original: Some other methods generate simplified representations of molecules, such as 1D SMILES strings (Weininger, 1988) and 2D graphs of molecules.
Revised: Some other methods generate simplified representations of molecules, such as 1D SMILES strings (Weininger, 1988) and 2D graphs of molecules.

Original: These include variational autoencoders (Kusner et al., 2017;Dai et al., 2018; Jin et al., 2018; Simonovsky & Komodakis, 2018;
Revised: These include variational autoencoders (Kusner et al., 2017;

Original:  Liu et al., 2018), normalizing flows (Madhawa et al., 2019; Zang & Wang, 2020; Luo et al., 2021), generative adversarial networks (Bian et al., 2019; Assouel et al., 2018), and autoregressive models (Popova et al., 2019; Flam-Shepherd et al., 2021).
Revised: Dai et al., 2018; Jin et al., 2018; Simonovsky & Komodakis, 2018; Liu et al., 2018), normalizing flows (Madhawa et al., 2019; Zang & Wang, 2020; Luo et al., 2021), generative adversarial networks (Bian et al., 2019; Assouel et al., 2018), and autoregressive models (Popova et al., 2019; Flam-Shepherd et al., 2021).

Revised: There are also methods on generating torsion angles in molecules.

Revised: For instance, Torsional Diffusion (Jing et al., 2022) employs the SDE formulation of diffusion models to model torsion angles in a given 2D molecular graph for the conformation generation task.

Original: Inverse molecular design.
Revised: Inverse molecular design.

Original: Generative models have been applied to inverse molecular design.
Revised: Generative models have been applied to inverse molecular design.

Original: For example, conditional autoregressive models (Gebauer et al., 2022) and EDM (Hoogeboom et al., 2022) directly generate 3D molecules with desired quantum properties.
Revised: For example, conditional autoregressive models (Gebauer et al., 2022) and EDM (Hoogeboom et al., 2022) directly generate 3D molecules with desired quantum properties.

Original: Gebauer et al.
Revised: Gebauer et al.

Original: (2019) also finetune pretrained generative models on a biased subset to generate 3D molecules with small
Revised: (2019) also finetune pretrained generative models on a biased subset to generate 3D molecules with small

Original: HOMO-LUMO gaps.
Revised: HOMO-LUMO gaps.

Original: In contrast to these conditional generative models, our work further proposes a guidance method, a flexible way to control the generation process of molecules.
Revised: In contrast to these conditional generative models, our work further proposes a guidance method, a flexible way to control the generation process of molecules.

Original: Some other methods apply optimization methods to search molecules with desired properties, such as reinforcement learning (Zhou et al., 2019;
Revised: Some other methods apply optimization methods to search molecules with desired properties, such as reinforcement learning (Zhou et al., 2019;

Original: You et al., 2018) and genetic algorithms (Jensen, 2019; Nigam et al., 2019).
Revised: You et al., 2018) and genetic algorithms (Jensen, 2019; Nigam et al., 2019).

Original: These optimization methods generally consider the 1D SMILES strings or 2D graphs of molecules, and the 3D information is not provided.
Revised: These optimization methods generally consider the 1D SMILES strings or 2D graphs of molecules, and the 3D information is not provided.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: 3D representation of molecules.
Revised: 3D representation of molecules.

Original: Suppose a molecule has M atoms and let x i ∈ R n ( n = 3 in general) be the coordinate of the i th atom.
Revised: Suppose a molecule has M atoms and let x i ∈ R n ( n = 3 in general) be the coordinate of the i th atom.

Original: The collection of coordinates x = ( x 1 , . . .
Revised: The collection of coordinates x = ( x 1 , . . .

Original: x M ) ∈
Revised: x M ) ∈

Original: R
Revised: R

Original: Mn determines the conformation of the molecule.
Revised: Mn determines the conformation of the molecule.

Original: In addition to the coordinate, each atom is also associated with an atom feature, e.g., the atom type.
Revised: In addition to the coordinate, each atom is also associated with an atom feature, e.g., the atom type.

Original: We use h i ∈ R d to represent the atom feature of the i th atom, and use h = ( h 1 , . .
Revised: We use h i ∈ R d to represent the atom feature of the i th atom, and use h = ( h 1 , . .

Original: , h M ) ∈
Revised: , h M ) ∈

Original: R
Revised: R

Original: Md to represent the collection of atom features in a molecule.
Revised: Md to represent the collection of atom features in a molecule.

Original: We use a tuple z = ( x , h ) to represent a molecule, which contains both the 3D geometry information and the atom feature information.
Revised: We use a tuple z = ( x , h ) to represent a molecule, which contains both the 3D geometry information and the atom feature information.

Original: Equivariance and invariance.
Revised: Equivariance and invariance.

Original: Suppose R is a transformation.
Revised: Suppose R is a transformation.

Original: A distribution p ( x , h ) is said to be invariant to R , if p ( x , h ) =
Revised: A distribution p ( x , h ) is said to be invariant to R , if p ( x , h ) =

Original: p ( Rx , h ) holds for all x and h .
Revised: p ( Rx , h ) holds for all x and h .

Original: Here Rx =
Revised: Here Rx =

Original: ( Rx 1 , . . .
Revised: ( Rx 1 , . . .

Original: , Rx M ) is applied to each coordinate.
Revised: , Rx M ) is applied to each coordinate.

Original: A function ( a x , a h ) =
Revised: A function ( a x , a h ) =

Original: f ( x , h ) that have two components a x , a h in its output is said to be equivariant to R , if f ( Rx , h ) =
Revised: f ( x , h ) that have two components a x , a h in its output is said to be equivariant to R , if f ( Rx , h ) =

Original: ( Ra x , a h ) holds for all x and h .
Revised: ( Ra x , a h ) holds for all x and h .

Original: A function f ( x , h ) is said to be invariant to R , if f ( Rx , h ) =
Revised: A function f ( x , h ) is said to be invariant to R , if f ( Rx , h ) =

Original: f ( x , h ) holds for all x and h .
Revised: f ( x , h ) holds for all x and h .

Original: Zero CoM subspace.
Revised: Zero CoM subspace.

Original: It has been shown that the invariance to translational and rotational transformations is an important factor for the success of 3D molecule modeling (K¨ohler et al., 2020;
Revised: It has been shown that the invariance to translational and rotational transformations is an important factor for the success of 3D molecule modeling (K ¨ ohler et al., 2020;

Original: Xu et al., 2022).
Revised: Xu et al., 2022).

Original: However, the translational invariance is impossible for a distribution in the full space R Mn (Satorras et al., 2021a).
Revised: However, the translational invariance is impossible for a distribution in the full space R Mn (Satorras et al., 2021a).

Original: Nevertheless, we can view two collections of coordinates x and y as equivalent if x can be translated from y , since the translation doesn’t change the identity of a molecule.
Revised: Nevertheless, we can view two collections of coordinates x and y as equivalent if x can be translated from y , since the translation doesn’t change the identity of a molecule.

Original: Such an equivalence relation partitions the whole space R Mn into disjoint equivalence classes.
Revised: Such an equivalence relation partitions the whole space R Mn into disjoint equivalence classes.

Original: Indeed, all elements in the same equivalence classes represent the same conformation, and we can use the element with zero center of mass (CoM), i.e., 1 M (cid:80) Mi
Revised: Indeed, all elements in the same equivalence classes represent the same conformation, and we can use the element with zero center of mass (CoM), i.e., 1 M (cid:80) Mi

Original: =1 x i = 0 , as the specific representation.
Revised: =1 x i = 0 , as the specific representation.

Original: These elements collectively form the zero CoM linear subspace X
Revised: These elements collectively form the zero CoM linear subspace X

Original: (Xu et al., 2022; Hoogeboom et al., 2022), and the rest of the paper always uses elements in X to represent conformations.
Revised: (Xu et al., 2022; Hoogeboom et al., 2022), and the rest of the paper always uses elements in X to represent conformations.

Original: Equivariant graph neural network.
Revised: Equivariant graph neural network.

Original: Satorras et al.
Revised: Satorras et al.

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

Original: propose equivariant graph neural networks (EGNNs), which incorporate the equivariance inductive bias into neural networks.
Revised: propose equivariant graph neural networks (EGNNs), which incorporate the equivariance inductive bias into neural networks.

Original: Specifically, ( a x , a h ) =
Revised: Specifically, ( a x , a h ) =

Original: EGNN( x , h ) is a composition of L equivariant convolutional layers .
Revised: EGNN( x , h ) is a composition of L equivariant convolutional layers .

Original: The l -th layer takes the tuple ( x l , h l ) as the input and outputs an updated version ( x l +1 , h l +1 ) , as follows:
Revised: The l -th layer takes the tuple ( x l , h l ) as the input and outputs an updated version ( x l +1 , h l +1 ) , as follows:

Original: [Equation]
Revised: [Equation]

Original: where Φ m , Φ w , Φ h , Φ x are parameterized by fully connected neural networks with parameters θ m , θ w , θ h , θ x respectively, and e ij are optional feature attributes.
Revised: where Φ m , Φ w , Φ h , Φ x are parameterized by fully connected neural networks with parameters θ m , θ w , θ h , θ x respectively, and e ij are optional feature attributes.

Original: We can verify that these layers are equivariant to orthogonal transformations, which include rotational transformations as special cases.
Revised: We can verify that these layers are equivariant to orthogonal transformations, which include rotational transformations as special cases.

Original: As their composition, the EGNN is also equivariant to orthogonal transformations.
Revised: As their composition, the EGNN is also equivariant to orthogonal transformations.

Original: Furthermore, let EGNN h ( x , h )
Revised: Furthermore, let EGNN h ( x , h )

Original: = a h , i.e., the second component in the output of the EGNN.
Revised: = a h , i.e., the second component in the output of the EGNN.

Original: Then
Revised: Then

Original: EGNN h ( x , h ) is invariant to orthogonal transformations.
Revised: EGNN h ( x , h ) is invariant to orthogonal transformations.

Original: Equivariant diffusion models (EDM) (Hoogeboom et al., 2022) are a variant of diffusion models (Ho et al., 2020) for molecule data.
Revised: Equivariant diffusion models (EDM) (Hoogeboom et al., 2022) are a variant of diffusion models (Ho et al., 2020) for molecule data.

Original: EDMs gradually inject noise to the molecule data z
Revised: EDMs gradually inject noise to the molecule data z

Original: = ( x , h ) via a forward process
Revised: = ( x , h ) via a forward process

Original: [Equation]
Revised: [Equation]

Original: where α n and β n represent the noise schedule and satisfy α n + β n = 1 , and N X represent the
Revised: where α n and β n represent the noise schedule and satisfy α n + β n = 1 , and N X represent the

Original: Gaussian distribution in the zero CoM subspace X (see its formal definition in Appendix A.2).
Revised: Gaussian distribution in the zero CoM subspace X (see its formal definition in Appendix A.2).

Original: Let α n = α 1 α 2 · · · α n , β n = 1 − α n and ˜ β n = β n β n − 1 /β n .
Revised: Let α n = α 1 α 2 · · · α n , β n = 1 − α n and ˜ β n = β n β n − 1 /β n .

Original: To generate samples, the forward process is reversed using a Markov chain:p (
Revised: To generate samples, the forward process is reversed using a Markov chain:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Here p ( z N ) =
Revised: Here p ( z N ) =

Original: N X ( x N | 0 , 1) N ( h N | 0 , 1) .
Revised: N X ( x N | 0 , 1) N ( h N | 0 , 1) .

Original: The mean µ n ( z n ) =
Revised: The mean µ n ( z n ) =

Original: ( µ xn ( z n ) , µ hn ( z n )) is parameterized by a noise prediction network ϵ θ ( z n , n ) , and is trained using a MSE loss, as follows:µ n (
Revised: ( µ xn ( z n ) , µ hn ( z n )) is parameterized by a noise prediction network ϵ θ ( z n , n ) , and is trained using a MSE loss, as follows:

Original: [Equation]
Revised: [Equation]

Original: where ϵ n = z n −√ α n z 0 √ β n is the standard Gaussian noise injected to z 0 and w ( t ) is the weight term.
Revised: where ϵ n = z n −√ α n z 0 √ β n is the standard Gaussian noise injected to z 0 and w ( t ) is the weight term.

Original: Hoogeboom et al.
Revised: Hoogeboom et al.

Original: (2022) show that the distribution of generated samples p ( z 0 ) is invariant to rotational transformations if the noise prediction network is equivariant to orthogonal transformations.
Revised: (2022) show that the distribution of generated samples p ( z 0 ) is invariant to rotational transformations if the noise prediction network is equivariant to orthogonal transformations.

Original: In Section 4.2, we extend this proposition to the SDE formulation of molecular diffusion modelling.
Revised: In Section 4.2, we extend this proposition to the SDE formulation of molecular diffusion modelling.

Original: Hoogeboom et al.
Revised: Hoogeboom et al.

Original: (2022) also present a conditional version of EDM for inverse molecular design by adding an extra input of the condition  to the noise prediction network.
Revised: (2022) also present a conditional version of EDM for inverse molecular design by adding an extra input of the condition c to the noise prediction network as ϵ θ ( z n , c, n ) .

Original:  E QUIVARIANT E NERGY -G UIDED SDE
Revised:  E QUIVARIANT E NERGY -G UIDED SDE

Original: In this part, we introduce our equivariant energy-guided SDE (EEGSDE), which is illustrated in Figure 1.
Revised: In this part, we introduce our equivariant energy-guided SDE (EEGSDE), as illustrated in Figure 1.

Original: EEGSDE is based on the SDE formulation of molecular diffusion modeling, which is described in Section 4.1 and Section 4.2.
Revised: EEGSDE is based on the SDE formulation of molecular diffusion modeling, which is described in Section 4.1 and Section 4.2.

Original: Then, we formally present our EEGSDE  in Section 4.3.
Revised: Then, we formally present our EEGSDE that incorporates an energy function to guide the molecular generation in Section 4.3.

Original: We provideproof and derivations in Appendix A.
Revised: We provide derivations in Appendix A.

Original: 4. SDE IN THE P RODUCT S PACE
Revised: 4. SDE IN THE P RODUCT S PACE

Original: Recall that a molecule is represented as a tuple z = ( x , h ) , where x = ( x 1 , . . .
Revised: Recall that a molecule is represented as a tuple z = ( x , h ) , where x = ( x 1 , . . .

Original: , x M ) ∈ X represents the conformation and h =
Revised: , x M ) ∈ X represents the conformation and h =

Original: ( h 1 , . . .
Revised: ( h 1 , . . .

Original: , h M ) ∈ R
Revised: , h M ) ∈ R

Original: Md represents atom features.
Revised: Md represents atom features.

Original: Here X =
Revised: Here X =

Original: { x ∈ R Mn : 1 M (cid:80) Mi
Revised: { x ∈ R Mn : 1 M (cid:80) Mi

Original: =1 x i = 0 } is the zero CoM subspace as mentioned in Section 3, and d is the dimension of features.
Revised: =1 x i = 0 } is the zero CoM subspace as mentioned in Section 3, and d is the dimension of features.

Original: We first introduce a continuous-time diffusion process { z t } 0 ≤ t ≤ T in the product space
Revised: We first introduce a continuous-time diffusion process { z t } 0 ≤ t ≤ T in the product space

Original: X × R Md , which gradually adds noise to both x and h .
Revised: X × R Md , which gradually adds noise to both x and h .

Original: This can be described by the following forward SDE:
Revised: This can be described by the following forward SDE:

Original: [Equation]
Revised: [Equation]

Original: where f ( t ) and g ( t ) are two scalar functions, w x and w h are independent standard Wiener processes in X and R Md respectively, and the SDE starts from the data distribution q ( z 0 ) .
Revised: where f ( t ) and g ( t ) are two scalar functions, w x and w h are independent standard Wiener processes in X and R Md respectively, and the SDE starts from the data distribution q ( z 0 ) .

Original: Note that w x can be constructed by subtracting the CoM of a standard Wiener process w in R Mn , i.e., w x = w − w , where w = 1 M (cid:80)
Revised: Note that w x can be constructed by subtracting the CoM of a standard Wiener process w in R Mn , i.e., w x = w − w , where w = 1 M (cid:80)

Original: Mi
Revised: Mi

Original: =1 w i is the CoM of w = ( w 1 , . .
Revised: =1 w i is the CoM of w = ( w 1 , . .

Original: . , w M ) .
Revised: . , w M ) .

Original: It can be shown that the SDE has a linear Gaussian transition kernel q ( z t | z s ) =
Revised: It can be shown that the SDE has a linear Gaussian transition kernel q ( z t | z s ) =

Original: q ( x t | x s ) q ( h t | h s ) from x s to x t , where 0 ≤ s < t ≤
Revised: q ( x t | x s ) q ( h t | h s ) from x s to x t , where 0 ≤ s < t ≤

Original: T .
Revised: T .

Original: Specifically, there exists two scalars α t | s and β t | s , s.t., q ( x t | x s ) =
Revised: Specifically, there exists two scalars α t | s and β t | s , s.t., q ( x t | x s ) =

Original: N X ( x t |√
Revised: N X ( x t |√

Original: α t | s x s , β t | s ) and q ( h t | h s ) =
Revised: α t | s x s , β t | s ) and q ( h t | h s ) =

Original: N ( h t |√
Revised: N ( h t |√

Original: α t | s h s , β t | s ) .
Revised: α t | s h s , β t | s ) .

Original: Here N X denotes the Gaussian distribution in the subspace X , and see Appendix A.2 for its formal definition.
Revised: Here N X denotes the Gaussian distribution in the subspace X , and see Appendix A.2 for its formal definition.

Original: Indeed, the forward process of EDM in Eq.
Revised: Indeed, the forward process of EDM in Eq.

Original: (1) is a discretization of the forward SDE in Eq.
Revised: (1) is a discretization of the forward SDE in Eq.

Original: To generate molecules, we reversethe diffusion process in Eq.
Revised: To generate molecules, we reverse Eq.

Original: (3) from T to 0 .
Revised: (3) from T to 0 .

Original: Such a time reversal forms another a SDE, which can be represented by both the score function form and the noise prediction form :
Revised: Such a time reversal forms another a SDE, which can be represented by both the score function form and the noise prediction form :

Original: [Equation]
Revised: [Equation]

Original: noise prediction
Revised: [Table] [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Original: Here q t ( z ) is the marginal distribution of z t , ∇ x log q t ( z ) is the gradient of log q t ( z ) w.r.t.
Revised: Here q t ( z ) is the marginal distribution of z t , ∇ x log q t ( z ) is the gradient of log q t ( z ) w.r.t.

Original: x 1 , ∇ x log q t ( z ) = 1 M (cid:80) Mi
Revised: x 1 , ∇ x log q t ( z ) = 1 M (cid:80) Mi

Original: =1 ∇ x i log q t ( z ) is the CoM of ∇ x log q t ( z ) , d t is the infinitesimal negative timestep, ˜ w x and ˜ w h are independent reverse-time standard Wiener processes in X and R Md respectively, and ϵ t =
Revised: =1 ∇ x i log q t ( z ) is the CoM of ∇ x log q t ( z ) , d t is the infinitesimal negative timestep, ˜ w x and ˜ w h are independent reverse-time standard Wiener processes in X and R Md respectively, and ϵ t =

Original: z t −√ α t | 0 z 0 √ β t |is
Revised: z t −√ α t | 0 z 0 √ β t |is

Original: the standard Gaussian noise injected to z 0 .
Revised: the standard Gaussian noise injected to z 0 .

Original: Compared to the original SDE introduced by Song et al.
Revised: Compared to the original SDE introduced by Song et al.

Original: (2020), our reverse SDE in Eq.
Revised: (2020), our reverse SDE in Eq.

Original: (4) additionally subtracts the
Revised: (4) additionally subtracts the

Original: CoM of ∇ x log q t ( z ) .
Revised: CoM of ∇ x log q t ( z ) .

Original: This ensures x t always stays in the zero CoM subspace as time flows back.
Revised: This ensures x t always stays in the zero CoM subspace as time flows back.

Original: To sample from the reverse SDE in Eq.
Revised: To sample from the reverse SDE in Eq.

Original: (4), we use a noise prediction network ϵ θ ( z t , t ) to estimate E q ( z 0 | z t ) ϵ t , through minimizing thefollowing MSE loss 
Revised: (4), we use a noise prediction network ϵ θ ( z t , t ) to estimate E q ( z 0 | z t ) ϵ t , through minimizing the MSE loss min θ E t E q ( z 0 , z t ) w ( t ) ∥

Original: [Equation] where t is uniformly sampled from [0 , T ] , and w ( t ) controls the weight of the loss term at time t .
Revised: ϵ θ ( z t , t ) − ϵ t ∥ 2 , where t is uniformly sampled from [0 , T ] , and w ( t ) controls the weight of the loss term at time t .

Original: Note that the noise ϵ t is in the product space X × R Md , so we subtract the CoM of the predicted noise of x t to ensure ϵ θ
Revised: Note that the noise ϵ t is in the product space X × R Md , so we subtract the CoM of the predicted noise of x t to ensure ϵ θ

Original: ( z t , t ) is also in the product space.
Revised: ( z t , t ) is also in the product space.

Original: Substituting ϵ θ ( z t , t ) into Eq.
Revised: Substituting ϵ θ ( z t , t ) into Eq.

Original: (4), we get an approximate reverse-time SDE parameterized by θ :
Revised: (4), we get an approximate reverse-time SDE parameterized by θ :

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where p T ( z T ) =
Revised: where p T ( z T ) =

Original: N X ( x T | 0 , 1) N ( h T | 0 , 1) is a Gaussian prior in the product space that approximates q T ( z T ) .
Revised: N X ( x T | 0 , 1) N ( h T | 0 , 1) is a Gaussian prior in the product space that approximates q T ( z T ) .

Original: We define p θ ( z 0 ) as the marginal distribution of Eq.
Revised: We define p θ ( z 0 ) as the marginal distribution of Eq.

Original: (5) at time t = 0 , which is the distribution of our generated samples.
Revised: (5) at time t = 0 , which is the distribution of our generated samples.

Original: Similarly to the forward process, the reverse process of EDM in Eq.
Revised: Similarly to the forward process, the reverse process of EDM in Eq.

Original: (2) is a discretization of the reverse SDE in Eq.
Revised: (2) is a discretization of the reverse SDE in Eq.

Original: 4. E QUIVARIANT SDE
Revised: 4. E QUIVARIANT SDE

Original: To leverage the geometric symmetry in 3D molecular conformation, p θ ( z 0 ) should be invariant to translational and rotational transformations.
Revised: To leverage the geometric symmetry in 3D molecular conformation, p θ ( z 0 ) should be invariant to translational and rotational transformations.

Original: As mentioned in Section 3, the translational invariance of p θ ( z 0 ) is already satisfied by considering the zero CoM subspace.
Revised: As mentioned in Section 3, the translational invariance of p θ ( z 0 ) is already satisfied by considering the zero CoM subspace.

Original: The rotational invariance can be satisfied if the noise prediction network is equivariant to orthogonal transformations, as summarized in the following theorem:
Revised: The rotational invariance can be satisfied if the noise prediction network is equivariant to orthogonal transformations, as summarized in the following theorem:

Original: Theorem 1.
Revised: Theorem 1.

Original: Let ( ϵ x θ ( z t , t ) , ϵ h
Revised: Let ( ϵ x θ ( z t , t ) , ϵ h

Original: θ ( z t , t ))
Revised: θ ( z t , t ))

Original: =
Revised: =

Original: ϵ θ ( z t , t ) , where ϵ x θ ( z t , t ) and ϵ h
Revised: ϵ θ ( z t , t ) , where ϵ x θ ( z t , t ) and ϵ h

Original: θ ( z t , t ) are the predicted noise of x t and h t respectively.
Revised: θ ( z t , t ) are the predicted noise of x t and h t respectively.

Original: If for any orthogonal transformation R ∈ R n × n , ϵ θ ( z t , t ) is equivariant to R , i.e., ϵ θ
Revised: If for any orthogonal transformation R ∈ R n × n , ϵ θ ( z t , t ) is equivariant to R , i.e., ϵ θ

Original: ( Rx t , h t , t ) =
Revised: ( Rx t , h t , t ) =

Original: ( Rϵ x θ ( x t , h t , t ) , ϵ h θ ( x t , h t , t )) , and p T ( z T ) is invariant
Revised: ( Rϵ x θ ( x t , h t , t ) , ϵ h θ ( x t , h t , t )) , and p T ( z T ) is invariant

Original: to
Revised: to

Original: R , i.e., p T ( Rx T , h T ) =
Revised: R , i.e., p T ( Rx T , h T ) =

Original: p T ( x T , h T ) , then p θ ( z 0 ) is invariant to any rotational transformation.
Revised: p T ( x T , h T ) , then p θ ( z 0 ) is invariant to any rotational transformation.

Original: As mentioned in Section 3, the EGNN satisfies the equivariance constraint, and we parameterize ϵ θ ( z t , t ) using an EGNN following Hoogeboom et al.
Revised: As mentioned in Section 3, the EGNN satisfies the equivariance constraint, and we parameterize ϵ θ ( z t , t ) using an EGNN following Hoogeboom et al.

Original: See details in Appendix C.
Revised: See details in Appendix D.

Original: 4. E QUIVARIANT E NERGY -G UIDED SDE
Revised: 4. E QUIVARIANT E NERGY -G UIDED SDE

Original: Now we describe equivariant energy-guided SDE (EEGSDE), which guides the generated molecules of Eq.
Revised: Now we describe equivariant energy-guided SDE (EEGSDE), which guides the generated molecules of Eq.

Original: (5) towards desired properties c by leveraging a time-dependent energy function
Revised: (5) towards desired properties c by leveraging a time-dependent energy function

Original: E
Revised: E

Original: ( z , c, t ) :
Revised: ( z , c, t ) :

Original: [Equation]
Revised: [Equation]

Original: energy gradient taken in the product space
Revised: energy gradient taken in the product space

Original: [Equation]
Revised: [Equation]

Original: which defines a distribution p θ ( z 0 | c ) conditioned on the property c .
Revised: which defines a distribution p θ ( z 0 | c ) conditioned on the property c .

Original: Here the CoM ∇ x E ( z , c, t ) of the gradient is subtracted to keep the SDE in the product space, which ensures the translational invariance of p θ ( z 0 | c ) .
Revised: Here the CoM ∇ x E ( z , c, t ) of the gradient is subtracted to keep the SDE in the product space, which ensures the translational invariance of p θ ( z 0 | c ) .

Original: Besides, the rotational invariance is satisfied by using energy invariant to orthogonal transformations, as summarized in the following theorem:
Revised: Besides, the rotational invariance is satisfied by using energy invariant to orthogonal transformations, as summarized in the following theorem:

Original: Require: Number of steps N ∆ t = TN z ← ( x − x , h ) , where x ∼ N ( 0 , 1) , h ∼ N ( 0 , 1) { Sample from the prior p T ( z T ) } for i

Original: = N to 1 do t ← i ∆ t g x ← ∇ x E ( z , c, t ) , g h ← ∇ h E ( z , c, t ) { Calculate the gradient of the energy function } g ← ( g x − g x , g h ) { Subtract the CoM of the gradient }

Original: F ← f ( t ) z

Original: [-+ g ( t ) 
-]
Revised: [-+ g ( t ) 
-]

Original: [-( 1 √ β t | 
-]
Revised: [-( 1 √ β t | 
-]

Original: ϵ θ

Original: ( z , t ) + g ) ϵ ←

Original: ( ϵ x − ϵ x , ϵ h ) , where ϵ x ∼ N ( 0 , 1) , ϵ h ∼ N ( 0 , 1) z ← z − F ∆ t + g ( t ) √ ∆ t ϵ { Update z according to Eq.

Original: (6) } end for return z

Original: Theorem 2.
Revised: Theorem 2.

Original: Suppose the assumptions in Theorem 1 hold and E ( z , c, t ) is invariant to any orthogonal transformation R , i.e., E ( Rx , h , c, t )
Revised: Suppose the assumptions in Theorem 1 hold and E ( z , c, t ) is invariant to any orthogonal transformation R , i.e., E ( Rx , h , c, t )

Original: =
Revised: =

Original: E ( x , h , c, t ) .
Revised: E ( x , h , c, t ) .

Original: Then p θ ( z 0 | c ) is invariant to any rotational transformation.
Revised: Then p θ ( z 0 | c ) is invariant to any rotational transformation.

Original: Note that we can also use a conditional model ϵ θ ( z , c, t ) in Eq.
Revised: Note that we can also use a conditional model ϵ θ ( z , c, t ) in Eq.

Original: See Appendix B for details.
Revised: See Appendix C for details.

Original: To sample from p θ ( z 0 | c ) , various solvers can be used for Eq.
Revised: To sample from p θ ( z 0 | c ) , various solvers can be used for Eq.

Original: (6), such as the Euler-Maruyama method (Song et al., 2020) and the Analytic-DPM sampler (Bao et al., 2022b;a).
Revised: (6), such as the Euler-Maruyama method (Song et al., 2020) and the Analytic-DPM sampler (Bao et al., 2022b;a).

Original: In Algorithm 1, we present the Euler-Maruyama method as an example.
Revised: We present the Euler-Maruyama method as an example in Algorithm 1 at Appendix B.

Original: Compared with Eq.
Revised: Our EEGSDE is a general framework, which can be applied to various applications by specifying different energy functions

Original: (5), Eq.
Revised: E ( z , c, t ) .

Revised: For example, we can design the energy function according to consistency between the molecule z and the property c , where a low energy represents a well consistency.

Revised: As the generation process in Eq.

Original: (6) additionally performs a gradient descent on the energy function. This encourages generated molecules to have a low energy.
Revised: (6) proceeds, the gradient of the energy function encourages generated molecules to have a low energy, and consequently a well consistency.

Original: Thus, we can design the energy function according to the consistency between the  molecule z and the property c .
Revised: Thus, we can expect the generated molecule z aligns well with the property c .

Original: In the rest of the paper, we specify the choice of energy functions, and show these energies improve controllable molecule generation targeted to quantum properties, molecular structures, and even a combination of them. 
Revised: In the rest of the paper, we specify the choice of energy functions, and show these energies improve controllable molecule generation targeted to quantum properties, molecular structures, and even a combination of them. Remark.

Revised: The term “energy” in this paper refers to a general notion in statistical machine learning, which is a scalar function that captures dependencies between input variables (LeCun et al., 2006).

Revised: Thus, the “energy” in this paper can be set to a MSE loss when we want to capture how the molecule align with the property (as done in Section 5).

Revised: Also, the “energy” in this paper does not exclude potential energy or free energy in chemistry, and they might be applicable when we want to generate molecules with small potential energy or free energy.

Original:  G ENERATING M OLECULES WITH D ESIRED Q UANTUM P ROPERTIES
Revised:  G ENERATING M OLECULES WITH D ESIRED Q UANTUM P ROPERTIES

Original: Let c ∈ R be a certain quantum property.
Revised: Let c ∈ R be a certain quantum property.

Original: To generate molecules with the desired property, we set the energy function as the squared error between the predicted property and the desired property:
Revised: To generate molecules with the desired property, we set the energy function as the squared error between the predicted property and the desired property:

Original: [Equation]
Revised: [Equation]

Original: where g ( z t , t ) is a time-dependent property prediction model, and s is the scaling factor controlling the strength of the guidance.
Revised: where g ( z t , t ) is a time-dependent property prediction model, and s is the scaling factor controlling the strength of the guidance.

Original: Specifically, we parameterize g ( z t , t ) using the second component in the output of EGNN (see Section 3) followed by a decoder (Dec):
Revised: Specifically, we parameterize g ( z t , t ) using the second component in the output of EGNN (see Section 3) followed by a decoder (Dec):

Original: [Equation]
Revised: [Equation]

Original: where the concatenation is performed on each atom feature, and the decoder consists of a node-wise multilayer perceptron (MLP), a sum pooling layer and an MLP for the final property prediction based on Satorras et al.
Revised: where the concatenation is performed on each atom feature, and the decoder is a small neural network based on Satorras et al.

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

Original: The training objective of g ( z t , t ) is available in Appendix D.
Revised: The training objective of g ( z t , t ) is available in Appendix F.

Original: This parameterization ensures that the energy function E ( z t , c, t ) is invariant to orthogonal transformations, and thus the distribution of generated samples is also invariant according to Theorem 2. 
Revised: This parameterization ensures that the energy function E ( z t , c, t ) is invariant to orthogonal transformations, and thus the distribution of generated samples is also invariant according to Theorem 2. Note that we can also generate molecules targeted to multiple quantum properties by combining energy functions linearly. We present details in Appendix G.1. 5. S ETUP

Original: We evaluate on the QM9 dataset (Ramakrishnan et al., 2014), which contains quantum properties and coordinates of ∼ 130k organic molecules with up to nine heavy atoms from (C, N, O, F).
Revised: We evaluate on the QM9 dataset (Ramakrishnan et al., 2014), which contains quantum properties and coordinates of ∼ 130k organic molecules with up to nine heavy atoms from (C, N, O, F).

Original: We follow Hoogeboom et al. for training and evaluation settings. Data split: We split QM9 into training, validation and test sets, which include 100K, 18K and 13K samples respectively.
Revised: Following EDM (Hoogeboom et al., 2022), we split QM9 into training, validation and test sets, which include 100K, 18K and 13K samples respectively.

Original: The training set is further divided into two non-overlapping halves D a , D b equally.
Revised: The training set is further divided into two nonoverlapping halves D a , D b equally.

Original: Training setting: We train the noise prediction network and the  energy function on setting: We report the mean absolute error (MAE), which is evaluated by a property classifier network ϕ c (Satorras et al., 2021b) trained onthe other half D a . We report the loss of ϕ c on D b as the lower bound (Lbound) (Hoogeboom et al., 2022) of the MAE.
Revised: The noise prediction model and the time-dependent property prediction model g ( z t , t ) used in energy function are trained on D b .

Original: We also report the novelty of generated molecules, as well as the atom stability (AS) and molecule stability (MS), which measure the validity of generated molecules.
Revised: Note that they are trained separately on their own objectives.

Original: Since the noise prediction network and the energy function are trained on D b ,
Revised: By default, EEGSDE uses a conditional noise prediction model ϵ θ

Original: [Table] [Table]
Revised: ( z , c, t ) in Eq.

Original: [Table] the novelty is also evaluated on D b , which leads to a higher value than that evaluated on the whole dataset (Hoogeboom et al., 2022).
Revised: (6), since we find it generate more accurate molecules than using an unconditional one ϵ θ ( z , t ) (see Appendix H.1 for an ablation study).

Original: We provide additional experimental settings in Appendix E.2.
Revised: Additional details such as training hyperparameters are provided in Appendix G.3.

Original: By default, EEGSDE uses a conditional noise prediction network in Eq.
Revised: Evaluation metric: Following EDM (Hoogeboom et al., 2022), we use the mean absolute error (MAE) to evaluate how generated molecules align with the desired property.

Original: (6), since we find it generates more accurate molecules than using an unconditional one (see Appendix F.1 for an ablation study).
Revised: Specifically, we train another property prediction model ϕ p (Satorras et al., 2021b) on D a , and the MAE is calculated as 1 K (cid:80) Ki =1 | ϕ p ( z i ) − c i | , where z i is a generated molecule, c i is its desired property and K is the number of generated samples.

Original: We compare our EEGSDE to EDM, which only adopts a conditional noise prediction network.
Revised: We generate K =10,000 samples for evaluation. Note that the property [Table] [Table]

Revised: [Table]

Revised: prediction model ϕ p is different from g ( z t , t ) used in the energy function, and they are trained on the two non-overlapping training subsets D a , D b respectively.

Revised: This ensures no information leak occurs when evaluating our EEGSDE, and therefore ensures the fairness.

Revised: For fairness and completeness, we also report the novelty, the atom stability and the molecule stability following Hoogeboom et al.

Revised: (2022) in Appendix H.2, although they are not our main focus.

Revised: Baseline: The most direct baseline is conditional EDM (Hoogeboom et al., 2022), which only adopts a conditional noise prediction network without the guidance of an energy function.

Original: We also compare two additional baselines “U-bound” and “#Atoms” from Hoogeboom et al.
Revised: We also compare two additional baselines “U-bound” and “#Atoms” from Hoogeboom et al.

Original: (2022) (see Appendix E.1 for details). 
Revised: (2022) (see Appendix G.2 for details). 5. R ESULTS

Original: Following Hoogeboom et al.
Revised: Following Hoogeboom et al.

Original: (2022), we consider six quantum properties in QM9: polarizability α , highest occupied molecular orbital energy ε HOMO , lowest unoccupied molecular orbital energy [Figure] Compared to EDM, the molecular structures of EEGSDE align better with the target structures.
Revised: (2022), we consider six quantum properties in QM9: polarizability α , highest occupied molecular orbital energy ε HOMO , lowest unoccupied molecular orbital energy ε LUMO , HOMO-LUMO gap ∆ ε , dipole moment µ and heat capacity C v .

Original: [Figure]

Original: ε LUMO , HOMO-LUMO gap ∆ ε , dipole moment µ and heat capacity C v .

Original: Firstly, we generate molecules targeted to one of these six properties.
Revised: Firstly, we generate molecules targeted to one of these six properties.

Original: As shown in Table 1, with the energy guidance, our EEGSDE has a significantly better MAE than the  EDM on all properties.
Revised: As shown in Table 1, with the energy guidance, our EEGSDE has a significantly better MAE than the conditional EDM on all properties.

Original: Remarkably, with a proper scaling factor s , the MAE of EEGSDE is reduced by more than 25% compared to  EDM on properties ∆ ε , ε HOMO , and more than 30% on µ .
Revised: Remarkably, with a proper scaling factor s , the MAE of EEGSDE is reduced by more than 25% compared to conditional EDM on properties ∆ ε , ε LUMO , and more than 30% on µ .

Original: Note that the energy guidance does not affect the novelty, the atom stability and molecule stability much (always within 6% compared to EDMs).

Original: We further generate molecules targeted to multiple quantum properties by combining energies linearly.
Revised: We further generate molecules targeted to multiple quantum properties by combining energy functions linearly.

Original: As shown in Table 2, our EEGSDE still has a significantly better MAE than the EDM, and does not affect the novelty, the atom stability and  molecule stability  much.
Revised: As shown in Table 2, our EEGSDE still has a significantly better MAE than the conditional EDM. [Figure] For completeness, we also report the novelty, the atom stability and the molecule stability of generated molecules in Appendix H.2, and the energy guidance does not affect these metrics much.

Revised: Conclusion: These results suggest that molecules generated by our EEGSDE align better with the desired properties than molecules generated by the conditional EDM baseline (for both the singleproperty and multiple-properties cases).

Original: As a result, our EEGSDE is able to explore the chemical space in a guided way to generate promising molecules, which may benefit applications such as  drug and material discovery.
Revised: As a consequence, EEGSDE is able to explore the chemical space in a guided way to generate promising molecules for downstream applications such as the virtual screening, which may benefit drug and material discovery.

Original:  G ENERATING M OLECULES WITH T ARGET S TRUCTURES
Revised:  G ENERATING M OLECULES WITH T ARGET S TRUCTURES

Original: Following Gebauer et al.
Revised: Following Gebauer et al.

Original: (2022), we use the molecular fingerprint to encode the structure information of a molecule.
Revised: (2022), we use the molecular fingerprint to encode the structure information of a molecule.

Original: The molecular fingerprint c = ( c 1 , . . .
Revised: The molecular fingerprint c = ( c 1 , . . .

Original: , c L ) is a series of bits that capture the presence or absence of substructures in the molecule.
Revised: , c L ) is a series of bits that capture the presence or absence of substructures in the molecule.

Revised: Specifically, a substructure is mapped to a specific position l in the bitmap, and the corresponding bit c l will be 1 if the substructure exists in the molecule and will be 0 otherwise.

Original: To generate molecules with a specific structure, we set the energy function as the squared error E ( z t , c, t ) =
Revised: To generate molecules with a specific structure (encoded by the fingerprint c ), we set the energy function as the squared error E ( z t , c, t ) =

Original: s
Revised: s

Original: ∥ m ( z t , t ) − c ∥ 2 between a time-dependent multi-label classifier m ( z t , t ) andthe corresponding fingerprint c .
Revised: ∥ m ( z t , t ) − c ∥ 2 between a time-dependent multi-label classifier m ( z t , t ) and c .

Original: Here s is the scaling factor.
Revised: Here s is the scaling factor, and m ( z t , t ) is trained with binary cross entropy loss to predict the fingerprint as detailed in Appendix F. Note that the choice of the energy function is flexible and can be different to the training loss of m ( z t , t ) .

Original: In initial experiments, we also try binary cross entropy loss for the energy, but we find it causes the generation process unstable.
Revised: In initial experiments, we also try binary cross entropy loss for the energy function, but we find it causes the generation process unstable.

Original: The multi-label classifier is parameterized by EGNN as

Original: [Equation]

Original: The multi-label classifier has the same backbone to the property prediction model in Eq.
Revised: The multi-label classifier is parameterized in a similar way to the property prediction model in Eq.

Original: (8), except that the decoder outputs a vector of dimension L , and thesigmoid function σ is adopted for multi-label classification.
Revised: (8), and we present details in Appendix E. 6. S ETUP We evaluate on QM9 and GEOM-Drug in this experiment. We train our method (including the noise prediction network and the multi-label classifier) on the whole training set.

Original: The training objective of m ( zt ,  t )is available inAppendix D. Similarly to Eq.
Revised: By default, we use a conditional noise prediction network ϵ θ ( z , c, t ) in Eq.

Original: (8), the EGNN in the multi-label classifier guarantees the invariance of the distribution of generated samples according to Theorem 2.
Revised: (6), since we find it generally performs better than an unconditional one ϵ θ

Original: [Figure]
Revised: ( z , t ) .

Original: We train the noise prediction network and the energy function on the whole training set of QM9.
Revised: More details including baselines are provided in Appendix G.4.

Original:  We evaluate the Tanimoto similarity (Gebauer et al., 2022) betweenunseen target structures from test set and structuresof generated molecules, which is computed by comparing fingerprints of two structures.
Revised: Evaluation metric: We want to measure how the structure of a generated molecule aligns with the target structure. Recall that the structure information is encoded by the molecular fingerprint, which is a bitmap representing the presence or absence of substructures. To this end, we use the Tanimoto similarity (Gebauer et al., 2022), which captures similarity between structures by comparing their fingerprints.

Original:  We provide more experimental details in Appendix E.3.
Revised: Specifically, let S g be the set of bits that are set to 1 in the fingerprint of a generated molecule, and S t be the set of bits that are set to 1 in the target structure. The Tanimoto similarity is defined as | S g ∩ S t | / | S g ∪ S t | , where | · | denotes the number of elements in a set. Baseline: The most direct baseline is conditional EDM (Hoogeboom et al., 2022), which only adopts a conditional noise prediction network without the guidance of an energy model. We also consider cG-SchNet(Gebauer et al., 2022), which generates molecules in an autoregressive manner.

Revised: 6. R ESULTS

Original: As shown in Table 3, EEGSDE significantly improves the similarity between target structures and generated structures compared to EDM.
Revised: As shown in Table 3, EEGSDE significantly improves the similarity between target structures and generated structures compared to conditional EDM and cG-SchNet on QM9.

Original: Also note in a proper range, a larger scaling factor results in a better similarity, and EEGSDE with s =1 improves the similarity by more than 10% compared to EDM.
Revised: Also note in a proper range, a larger scaling factor results in a better similarity, and EEGSDE with s =1 improves the sim- S 1 = S 2 =𝛼 = 100 Bohr S 2 =

Revised: T a r g e t

Original:  In Figure 2, we plot generated molecules of  EDM and EEGSDE ( s =1) targeted to specific structures, where our EEGSDE aligns better with them.
Revised: S t r u c t u r e S 1 = 0 .= ilarity by more than 10% compared to conditional EDM. In Figure 2, we plot generated molecules of conditional EDM and EEGSDE ( s =1) targeted to specific structures, where our EEGSDE aligns better with them.

Original: We further visualize the effect of the scaling factor in Table 3, where the generated structures align better as the scaling factor grows.
Revised: We further visualize the effect of the scaling factor in Appendix H.3, where the generated structures align better as the scaling factor grows.

Original: These results demonstrate that our EEGSDE captures the structure information in molecules well.
Revised: These results demonstrate that our EEGSDE captures the structure information in molecules well.

Original: We also perform experiments on the more challenging GEOM-Drug (Axelrod & Gomez-Bombarelli, 2022) dataset, and we train the conditional EDM baseline using the default setting of Hoogeboom et al.(2022).
Revised: We also perform experiments on the more challenging GEOM-Drug (Axelrod & Gomez-Bombarelli, 2022) dataset, and we train the conditional EDM baseline following the default setting of Hoogeboom et al.

Original: We find the  EDM baseline has a similarity of 0.166, which is much lower than  QM9.
Revised: As shown in Table 3, we find the conditional EDM baseline has a similarity of 0.165, which is much lower than the value on QM9.

Original: We hypothesize this is because molecules in GEOM-Drug has much more atoms than QM9 with a more complex structure, and the default setting in Hoogeboom et al.
Revised: We hypothesize this is because molecules in GEOM-Drug has much more atoms than QM9 with a more complex structure, and the default setting in Hoogeboom et al.

Original: (2022) is suboptimal.
Revised: (2022) is suboptimal.

Original: For example, the  EDM on GEOM-Drug has a smaller number of parameters than the  EDM on QM9 (15M v.s. 26M), which is insufficient to capture the structure information.
Revised: For example, the conditional EDM on GEOM-Drug has a smaller number of parameters than the conditional EDM on QM9 (15M v.s. 26M), which is insufficient to capture the structure information.

Original: Nevertheless, our EEGSDE still improves the similarity by ∼ %17.
Revised: Nevertheless, our EEGSDE still improves the similarity by ∼ %17.

Original: We provide more results on GEOM-Drug in Appendix F.3.
Revised: We provide generated molecules on GEOM-Drug in Appendix H.5.

Original: Finally, we demonstrate that our EEGSDE is a flexible framework to generate molecules targeted to multiple properties, which is often the practical case.
Revised: Finally, we demonstrate that our EEGSDE is a flexible framework to generate molecules targeted to multiple properties, which is often the practical case.

Original: We additionally target to the quantum property α (polarizability) by combining the energy function for molecular structures in this section and the energy function for quantum properties in Section 5.
Revised: We additionally target to the quantum property α (polarizability) by combining the energy function for molecular structures in this section and the energy function for quantum properties in Section 5.

Original: Here we choose α = 100 Bohr 3 , which is a relatively large value, and we expect it to encourage less isometrically shaped structures.
Revised: Here we choose α = 100 Bohr 3 , which is a relatively large value, and we expect it to encourage less isometrically shaped structures.

Original: As shown in Figure 4, the generated molecule aligns better with the target structure as the scaling factor sgrows, and meanwhile a ring substructure in the generated molecule vanishes as the scaling factor for polarizability s 2 grows, leading to a less isometrically shaped structure, which is as expected. 
Revised: As shown in Figure 3, the generated molecule aligns better with the target structure as the scaling factor sgrows, and meanwhile a ring substructure in the generated molecule vanishes as the scaling factor for polarizability s 2 grows, leading to a less isometrically shaped structure, which is as expected. Conclusion: These results suggest that molecules generated by our EEGSDE align better with the target structure than molecules generated by the conditional EDM baseline. Besides, EEGSDE can generate molecules targeted to both specific structures and desired quantum properties by combining energy functions linearly. As a result, EEGSDE may benefit practical cases in molecular design when multiple properties should be considered at the same time.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This work presents equivariant energy-guided SDE (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.
Revised: This work presents equivariant energy-guided SDE (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models.

Original: EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.
Revised: EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations.

Original: EEGSDE significantly improves the  EDM baseline in inverse molecular design targeted to quantum properties and molecular structures.
Revised: EEGSDE significantly improves the conditional EDM baseline in inverse molecular design targeted to quantum properties and molecular structures.

Original: Furthermore, EEGSDE is able to generate molecules with multiple target properties by combining the corresponding energy functions linearly.
Revised: Furthermore, EEGSDE is able to generate molecules with multiple target properties by combining the corresponding energy functions linearly.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: Inverse molecular design is critical in fields like material science and drug discovery.
Revised: Inverse molecular design is critical in fields like material science and drug discovery.

Original: Our EEGSDE is a flexible framework for inverse molecular design and thus might benefit these fields.
Revised: Our EEGSDE is a flexible framework for inverse molecular design and thus might benefit these fields.

Original: Currently the negative consequences are not obvious.
Revised: Currently the negative consequences are not obvious.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: Our code is included in the supplementary material.
Revised: Our code is included in the supplementary material.

Original: The implementation of our experiment is described in Section 5 and Section 6.
Revised: The implementation of our experiment is described in Section 5 and Section 6.

Original: Further details such as the hyperparameters of training and model backbones are provided in 
Revised: Further details such as the hyperparameters of training and model backbones are provided in Appendix G.

Original: Appendix E. We provide complete proofs and derivations of all theoretical results in Appendix A
Revised: We provide complete proofs and derivations of all theoretical results in Appendix A

Original: A D ERIVATIONS A. Z ERO C O M S UBSPACE
Revised: A D ERIVATIONS A. Z ERO C O M S UBSPACE

Original: The zero CoM subspace X = { x ∈ R
Revised: The zero CoM subspace X = { x ∈ R

Original: Mn : (cid:80) Mi
Revised: Mn : (cid:80) Mi

Original: =1 x i = 0 } is a ( M − 1) n dimensional subspace of R Mn .
Revised: =1 x i = 0 } is a ( M − 1) n dimensional subspace of R Mn .

Original: Therefore, there exists an isometric isomorphism ϕ
Revised: Therefore, there exists an isometric isomorphism ϕ

Original: from R ( M − 1) n to X , i.e., ϕ is a linear bijection from R ( M − 1) n to X , and ∥ ϕ
Revised: from R ( M − 1) n to X , i.e., ϕ is a linear bijection from R ( M − 1) n to X , and ∥ ϕ

Original: (ˆ x ) ∥ 2 = ∥
Revised: (ˆ x ) ∥ 2 = ∥

Original: ˆ x ∥ 2 for all ˆ x ∈ R ( M − 1) n .
Revised: ˆ x ∥ 2 for all ˆ x ∈ R ( M − 1) n .

Original: We use A ϕ ∈
Revised: We use A ϕ ∈

Original: R Mn × ( M − 1) n represent the matrix corresponding to ϕ , so we have ϕ (ˆ x ) =
Revised: R Mn × ( M − 1) n represent the matrix corresponding to ϕ , so we have ϕ (ˆ x ) =

Original: A ϕ ˆ x .
Revised: A ϕ ˆ x .

Original: An important property of the isometric isomorphism is that A ϕ
Revised: An important property of the isometric isomorphism is that A ϕ

Original: A ⊤ ϕ x = x − x for all x ∈ R Mn .
Revised: A ⊤ ϕ x = x − x for all x ∈ R Mn .

Original: We show the proof as following.
Revised: We show the proof as following.

Original: Proposition 1.
Revised: Proposition 1.

Original: Suppose ϕ is an isometric isomorphism from R ( M − 1) n to X , and let A ϕ ∈
Revised: Suppose ϕ is an isometric isomorphism from R ( M − 1) n to X , and let A ϕ ∈

Original: R Mn × ( M − 1) n be the matrix corresponding to ϕ .
Revised: R Mn × ( M − 1) n be the matrix corresponding to ϕ .

Original: Then we have A ϕ
Revised: Then we have A ϕ

Original: A ⊤ ϕ x = x − x for all x ∈ R Mn , where x = 1 M (cid:80) Mi
Revised: A ⊤ ϕ x = x − x for all x ∈ R Mn , where x = 1 M (cid:80) Mi

Original: =1 x i .
Revised: =1 x i .

Original: Proof.
Revised: Proof.

Original: We consider a new subspace of R Mn , X ⊥ =
Revised: We consider a new subspace of R Mn , X ⊥ =

Original: { x ∈ R Mn : x ⊥ X } , i.e., the orthogonal component of X .
Revised: { x ∈ R Mn : x ⊥ X } , i.e., the orthogonal component of X .

Original: We can verify that
Revised: We can verify that

Original: X ⊥ = { x ∈ R Mn : x 1 = x 2 = · · · = x M } , and X ⊥ is n dimensional.
Revised: X ⊥ = { x ∈ R Mn : x 1 = x 2 = · · · = x M } , and X ⊥ is n dimensional.

Original: Thus, there exists an isometric isomorphism ψ from R n to X ⊥ .
Revised: Thus, there exists an isometric isomorphism ψ from R n to X ⊥ .

Original: Let A ψ ∈ R Mn × n represent the matrix corresponding to ψ .
Revised: Let A ψ ∈ R Mn × n represent the matrix corresponding to ψ .

Original: Then we define λ (ˆ x , ˆ y )
Revised: Then we define λ (ˆ x , ˆ y )

Original: = ϕ
Revised: = ϕ

Original: (ˆ x )+ ψ
Revised: (ˆ x )+ ψ

Original: (ˆ y ) , where ˆ x ∈ R ( M − 1) n and ˆ y ∈ R n .
Revised: (ˆ y ) , where ˆ x ∈ R ( M − 1) n and ˆ y ∈ R n .

Original: The image of λ is { x + y : x ∈ X,
Revised: The image of λ is { x + y : x ∈ X,

Original: y ∈
Revised: y ∈

Original: X ⊥ } =
Revised: X ⊥ } =

Original: R Mn .
Revised: R Mn .

Original: Therefore, λ is a linear bijection from R Mn to R Mn , and the matrix corresponding to λ is A λ
Revised: Therefore, λ is a linear bijection from R Mn to R Mn , and the matrix corresponding to λ is A λ

Original: =
Revised: =

Original: [ A ϕ , A ψ ] .
Revised: [ A ϕ , A ψ ] .

Original: Furthermore, λ is an isometric isomorphism, since ∥ λ (ˆ x , ˆ y )
Revised: Furthermore, λ is an isometric isomorphism, since ∥ λ (ˆ x , ˆ y )

Original: ∥ 2 = ∥
Revised: ∥ 2 = ∥

Original: ϕ (ˆ x ) ∥ 2 + ∥ ψ (ˆ y ) ∥ 2 + 2 ⟨ ϕ (ˆ x ) , ψ (ˆ y ) ⟩ = ∥
Revised: ϕ (ˆ x ) ∥ 2 + ∥ ψ (ˆ y ) ∥ 2 + 2 ⟨ ϕ (ˆ x ) , ψ (ˆ y ) ⟩ = ∥

Original: ˆ x ∥ 2 + ∥ ˆ y ∥ 2 = ∥
Revised: ˆ x ∥ 2 + ∥ ˆ y ∥ 2 = ∥

Original: (ˆ x ⊤ , ˆ y ⊤ )
Revised: (ˆ x ⊤ , ˆ y ⊤ )

Original: ⊤ ∥ 2 .
Revised: ⊤ ∥ 2 .

Original: This means A λ is orthogonal transformation.
Revised: This means A λ is orthogonal transformation.

Original: Therefore,
Revised: Therefore,

Original: A λ A ⊤ λ =
Revised: A λ A ⊤ λ =

Original: A ϕ
Revised: A ϕ

Original: A ⊤ ϕ x + A ψ A ⊤ ψ =
Revised: A ⊤ ϕ x + A ψ A ⊤ ψ =

Original: I and A ϕ
Revised: I and A ϕ

Original: A ⊤ ϕ x + A ψ A ⊤ ψ x = x .
Revised: A ⊤ ϕ x + A ψ A ⊤ ψ x = x .

Original: Since A ϕ
Revised: Since A ϕ

Original: A ⊤ ϕ x ∈ X
Revised: A ⊤ ϕ x ∈ X

Original: and
Revised: and

Original: A ψ A ⊤ ψ x ∈ X ⊤ , we can conclude that A ϕ
Revised: A ψ A ⊤ ψ x ∈ X ⊤ , we can conclude that A ϕ

Original: A ⊤ ϕ x is the orthogonal projection of x to X , which is exactly x − x .
Revised: A ⊤ ϕ x is the orthogonal projection of x to X , which is exactly x − x .

Original: Since R ( M − 1) n and X are two intrinsically equivalent spaces, an equivalence of distributions in these two spaces can also be established, as shown in the following propositions.
Revised: Since R ( M − 1) n and X are two intrinsically equivalent spaces, an equivalence of distributions in these two spaces can also be established, as shown in the following propositions.

Original: Proposition 2.
Revised: Proposition 2.

Original: Suppose x is a random vector distributed in X , and ˆ x = ϕ − 1 ( x ) is its equivalent representation in R ( M − 1) n .
Revised: Suppose x is a random vector distributed in X , and ˆ x = ϕ − 1 ( x ) is its equivalent representation in R ( M − 1) n .

Original: If x ∼ q ( x ) , then ˆ x ∼ ˆ q (ˆ x ) , where ˆ q (ˆ x ) =
Revised: If x ∼ q ( x ) , then ˆ x ∼ ˆ q (ˆ x ) , where ˆ q (ˆ x ) =

Original: q ( ϕ (ˆ x )) .
Revised: q ( ϕ (ˆ x )) .

Original: Proposition 3.
Revised: Proposition 3.

Original: Suppose x , y is are two random vectors distributed in X , and
Revised: Suppose x , y is are two random vectors distributed in X , and

Original: ˆ x = ϕ − 1 ( x ) , ˆ y = ϕ − 1 ( y ) are their equivalent representations in R ( M − 1) n .
Revised: ˆ x = ϕ − 1 ( x ) , ˆ y = ϕ − 1 ( y ) are their equivalent representations in R ( M − 1) n .

Original: If x | y ∼ q ( x | y ) , then ˆ x | ˆ y ∼ ˆ q (ˆ x | ˆ y ) , where ˆ q (ˆ x | ˆ y )
Revised: If x | y ∼ q ( x | y ) , then ˆ x | ˆ y ∼ ˆ q (ˆ x | ˆ y ) , where ˆ q (ˆ x | ˆ y )

Original: =
Revised: =

Original: q ( ϕ (ˆ x ) | ϕ
Revised: q ( ϕ (ˆ x ) | ϕ

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

Original: A. SDE IN THE P RODUCT S PACE
Revised: A. SDE IN THE P RODUCT S PACE

Original: Definition 1 (Gaussian distributions in the zero CoM subspace) .
Revised: Definition 1 (Gaussian distributions in the zero CoM subspace) .

Original: Suppose µ ∈ X .
Revised: Suppose µ ∈ X .

Original: Let
Revised: Let

Original: N X ( x | µ , σ 2 ) :
Revised: N X ( x | µ , σ 2 ) :

Original: = (2 πσ 2 ) − ( M − 1) n/
Revised: = (2 πσ 2 ) − ( M − 1) n/

Original: 2 exp( − 1 2 σ 2 ∥ x − µ ∥ 22 ) , which is the isotropic Gaussian distribution with mean µ and variance σ 2 in the zero CoM subspace X .
Revised: 2 exp( − 1 2 σ 2 ∥ x − µ ∥ 22 ) , which is the isotropic Gaussian distribution with mean µ and variance σ 2 in the zero CoM subspace X .

Original: Proposition 4 (Transition kernels of a SDE in the zero CoM subspace) .
Revised: Proposition 4 (Transition kernels of a SDE in the zero CoM subspace) .

Original: Suppose d x = f ( t ) x
Revised: Suppose d x = f ( t ) x

Original: d t + g ( t )d w x , x 0 ∼ q ( x 0 ) is a SDE in the zero CoM subspace.
Revised: d t + g ( t )d w x , x 0 ∼ q ( x 0 ) is a SDE in the zero CoM subspace.

Original: Then the transition kernel from x s to x t ( 0 ≤ s < t ≤ T ) can be expressed as q ( x t | x s ) =
Revised: Then the transition kernel from x s to x t ( 0 ≤ s < t ≤ T ) can be expressed as q ( x t | x s ) =

Original: N X ( x t |√
Revised: N X ( x t |√

Original: α t | s x s , β t | s ) , where α t | s = exp(2 (cid:82) t s f ( τ )d τ ) and β t | s = α t | s (cid:82) t s g ( τ ) 2 /α τ | s d τ are two scalars determined by f ( · ) and g ( · ) .
Revised: α t | s x s , β t | s ) , where α t | s = exp(2 (cid:82) t s f ( τ )d τ ) and β t | s = α t | s (cid:82) t s g ( τ ) 2 /α τ | s d τ are two scalars determined by f ( · ) and g ( · ) .

Original: Proof.
Revised: Proof.

Original: Firstly, we map the process { x t } Tt =0 in the zero CoM subspace to the equivalent space R ( M − 1) n through the isometric isomorphism ϕ introduced in Appendix A.1.
Revised: Firstly, we map the process { x t } Tt =0 in the zero CoM subspace to the equivalent space R ( M − 1) n through the isometric isomorphism ϕ introduced in Appendix A.1.

Original: This produces a new process { ˆ x } Tt =0 , where ˆ x t = ϕ − 1 ( x t ) .
Revised: This produces a new process { ˆ x } Tt =0 , where ˆ x t = ϕ − 1 ( x t ) .

Original: By applying ϕ − 1 to the SDE in the zero CoM subspace, we know ˆ x t = ϕ − 1 ( x t ) satisfies the following SDE in R ( M − 1) n :
Revised: By applying ϕ − 1 to the SDE in the zero CoM subspace, we know ˆ x t = ϕ − 1 ( x t ) satisfies the following SDE in R ( M − 1) n :

Original: [Equation]
Revised: [Equation]

Original: where ˆ w is the standard Wiener process in R ( M − 1) n and
Revised: where ˆ w is the standard Wiener process in R ( M − 1) n and

Original: ˆ q (ˆ x 0 ) =
Revised: ˆ q (ˆ x 0 ) =

Original: q ( ϕ (ˆ x 0 )) .
Revised: q ( ϕ (ˆ x 0 )) .

Original: According to Song et al.
Revised: According to Song et al.

Original: (2020), the transition of Eq.
Revised: (2020), the transition of Eq.

Original: (9) from ˆ x s to ˆ x t ( s < t ) can be expressed as ˆ q (ˆ x t | ˆ x s ) =
Revised: (9) from ˆ x s to ˆ x t ( s < t ) can be expressed as ˆ q (ˆ x t | ˆ x s ) =

Original: N ( ˆ x t | √ α t | s ˆ x s , β t | s I ) , where α t | s = exp(2 (cid:82) t s f ( τ )d τ ) and β t | s = α t | s (cid:82) t s g ( τ ) 2 /α τ | s d τ are two scalars determined by f ( · ) and g ( · ) .
Revised: N ( ˆ x t | √ α t | s ˆ x s , β t | s I ) , where α t | s = exp(2 (cid:82) t s f ( τ )d τ ) and β t | s = α t | s (cid:82) t s g ( τ ) 2 /α τ | s d τ are two scalars determined by f ( · ) and g ( · ) .

Original: According to Proposition 3, we know the transition kernel q ( x t | x s ) from x s to x t satisfies
Revised: According to Proposition 3, we know the transition kernel q ( x t | x s ) from x s to x t satisfies

Original: [Equation]
Revised: [Equation]

Original: Proposition 5 (Transition kernels of the SDE in the product space) .
Revised: Proposition 5 (Transition kernels of the SDE in the product space) .

Original: Suppose d z = f ( t )
Revised: Suppose d z = f ( t )

Original: z d t + g ( t )d( w x , w h ) , z 0 ∼ q ( z 0 ) is the SDE in the product space
Revised: z d t + g ( t )d( w x , w h ) , z 0 ∼ q ( z 0 ) is the SDE in the product space

Original: X × R Md , as introduced in
Revised: X × R Md , as introduced in

Original: Eq.
Revised: Eq.

Original: Then the transition kernel from z s to z t ( 0 ≤ s < t ≤ T ) can be expressed as q ( z t | z s )
Revised: Then the transition kernel from z s to z t ( 0 ≤ s < t ≤ T ) can be expressed as q ( z t | z s )

Original: =
Revised: =

Original: q ( x t | x s ) q ( h t | h s ) , where q ( x t | x s )
Revised: q ( x t | x s ) q ( h t | h s ) , where q ( x t | x s )

Original: =
Revised: =

Original: N X ( x t |√
Revised: N X ( x t |√

Original: α t | s x s , β t | s ) and q ( h t | h s ) =
Revised: α t | s x s , β t | s ) and q ( h t | h s ) =

Original: N ( h t | √ α t | s h s , β t | s ) .
Revised: N ( h t | √ α t | s h s , β t | s ) .

Original: Here α t | s and β t | s are defined as in Proposition 4.
Revised: Here α t | s and β t | s are defined as in Proposition 4.

Original: Proof.
Revised: Proof.

Original: Since w x and w h are independent to each other, the transition kernel from z s to z t can be factorized as q ( z t | z s )
Revised: Since w x and w h are independent to each other, the transition kernel from z s to z t can be factorized as q ( z t | z s )

Original: =
Revised: =

Original: q ( x t | x s ) q ( h t | h s ) , where q ( x t | x s ) is the transition kernel of d x = f ( t ) x
Revised: q ( x t | x s ) q ( h t | h s ) , where q ( x t | x s ) is the transition kernel of d x = f ( t ) x

Original: d t + g ( t )d w x and q ( h t | h s ) is the transition kernel of d h = f ( t ) h d t + g ( t )d w h .
Revised: d t + g ( t )d w x and q ( h t | h s ) is the transition kernel of d h = f ( t ) h d t + g ( t )d w h .

Original: According to Proposition 4 and Song et al.
Revised: According to Proposition 4 and Song et al.

Original: (2020), we have q ( x t | x s ) =
Revised: (2020), we have q ( x t | x s ) =

Original: N X ( x t |√
Revised: N X ( x t |√

Original: α t | s x s , β t | s ) and q ( h t | h s ) =
Revised: α t | s x s , β t | s ) and q ( h t | h s ) =

Original: N ( h t |√
Revised: N ( h t |√

Original: α t | s h s , β t | s ) .
Revised: α t | s h s , β t | s ) .

Original: Remark 1.
Revised: Remark 1.

Original: The marginal distribution of z t is
Revised: The marginal distribution of z t is

Original: [Equation]
Revised: [Equation]

Original: where λ is the Lebesgue measure in the zero CoM subspace X .
Revised: where λ is the Lebesgue measure in the zero CoM subspace X .

Original: While q t ( z t ) is a distribution in X × R Md , it has a natural differentiable extension to R Mn × R Md , since q ( x t | x 0 ) =
Revised: While q t ( z t ) is a distribution in X × R Md , it has a natural differentiable extension to R Mn × R Md , since q ( x t | x 0 ) =

Original: N X ( x t |√
Revised: N X ( x t |√

Original: α t | s x s , β t | s ) has a differentiable extension to R Mn according to Definition 1.
Revised: α t | s x s , β t | s ) has a differentiable extension to R Mn according to Definition 1.

Original: Thus, we can take gradient of q t ( z t ) w.r.t.
Revised: Thus, we can take gradient of q t ( z t ) w.r.t.

Original: x t in the whole space R Mn .
Revised: x t in the whole space R Mn .

Original: Proposition 6 (Time reversal of the SDE in the product space) .
Revised: Proposition 6 (Time reversal of the SDE in the product space) .

Original: Suppose d z = f ( t )
Revised: Suppose d z = f ( t )

Original: z d t + g ( t )d( w x , w h ) , z 0 ∼ q ( z 0 ) is the SDE in the product space
Revised: z d t + g ( t )d( w x , w h ) , z 0 ∼ q ( z 0 ) is the SDE in the product space

Original: X × R Md , as introduced in Eq.
Revised: X × R Md , as introduced in Eq.

Original: Then its time reversal satisfies the following reverse-time SDE, which can be represented by both the score function form and the noise prediction form:
Revised: Then its time reversal satisfies the following reverse-time SDE, which can be represented by both the score function form and the noise prediction form:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: prediction
Revised: prediction

Original: [Equation]
Revised: [Equation]

Original: where ˜ w x and ˜ w h are reverse-time standard Wiener processes in X and R Md respectively, and ϵ t =
Revised: where ˜ w x and ˜ w h are reverse-time standard Wiener processes in X and R Md respectively, and ϵ t =

Original: z t −√ α t | 0 z 0 √ β t |is
Revised: z t −√ α t | 0 z 0 √ β t |is

Original: the standard Gaussian noise in X × R Md injected to z 0 .
Revised: the standard Gaussian noise in X × R Md injected to z 0 .

Original: Furthermore, we have ∇ x log q t ( x ) −∇ x log q t ( x ) = −
Revised: Furthermore, we have ∇ x log q t ( x ) −∇ x log q t ( x ) = −

Original: [Equation]
Revised: [Equation]

Original: ϵ t , where ϵ t = x t −√ α t | 0 x 0 √ β t |is the standard Gaussian noise in the zero CoM subspace injected to x 0 .
Revised: ϵ t , where ϵ t = x t −√ α t | 0 x 0 √ β t |is the standard Gaussian noise in the zero CoM subspace injected to x 0 .

Original: Proof.
Revised: Proof.

Original: Let ˆ z t =
Revised: Let ˆ z t =

Original: (ˆ x t , h t ) , where ˆ x t = ϕ − 1 ( x t ) , introduced in the proof of Proposition 4.
Revised: (ˆ x t , h t ) , where ˆ x t = ϕ − 1 ( x t ) , introduced in the proof of Proposition 4.

Original: Then { ˆ z t } Tt =0 is a process in R ( M − 1) n × R Md determined by the following SDE
Revised: Then { ˆ z t } Tt =0 is a process in R ( M − 1) n × R Md determined by the following SDE

Original: [Equation]
Revised: [Equation]

Original: where ˆ w is the standard Wiener process in R ( M − 1) n × R Md and ˆ q ( ˆ z 0 )
Revised: where ˆ w is the standard Wiener process in R ( M − 1) n × R Md and ˆ q ( ˆ z 0 )

Original: =
Revised: =

Original: q ( ϕ ( ˆ x 0 ) , h ) .
Revised: q ( ϕ ( ˆ x 0 ) , h ) .

Original: According to Song
Revised: According to Song

Original: (10) has reversal:
Revised: (10) has reversal:

Original: [Equation]
Revised: [Equation]

Original: where ˆ q t ( z ) is the marginal distribution of ˆ z t , which satisfies ˆ q t (ˆ z t )
Revised: where ˆ q t ( z ) is the marginal distribution of ˆ z t , which satisfies ˆ q t (ˆ z t )

Original: =
Revised: =

Original: q t ( ϕ (ˆ x t ) , h t ) according to
Revised: q t ( ϕ (ˆ x t ) , h t ) according to

Original: Proposition 2, and ˜ w is the reverse-time standard Wiener process in R ( M − 1) n × R Md .
Revised: Proposition 2, and ˜ w is the reverse-time standard Wiener process in R ( M − 1) n × R Md .

Original: Then we apply the linear transformation T ˆ z
Revised: Then we apply the linear transformation T ˆ z

Original: = ( ϕ (ˆ x ) , h ) to Eq.
Revised: = ( ϕ (ˆ x ) , h ) to Eq.

Original: (11), which maps ˆ z t back to z t .
Revised: (11), which maps ˆ z t back to z t .

Original: This yields
Revised: This yields

Original: [Equation]
Revised: [Equation]

Original: Here ˜ w x and ˜ w h are reverse-time standard Wiener processes in X and R Md respectively, and
Revised: Here ˜ w x and ˜ w h are reverse-time standard Wiener processes in X and R Md respectively, and

Original: T ( ∇ ˆ z log ˆ q t (ˆ z )) can be expressed as ( ϕ ( ∇ ˆ x log ˆ q t (ˆ z )) , ∇ h log ˆ q t (ˆ z )) .
Revised: T ( ∇ ˆ z log ˆ q t (ˆ z )) can be expressed as ( ϕ ( ∇ ˆ x log ˆ q t (ˆ z )) , ∇ h log ˆ q t (ˆ z )) .

Original: Since ∇ ˆ x log ˆ q t (ˆ z ) =
Revised: Since ∇ ˆ x log ˆ q t (ˆ z ) =

Original: A ⊤ ϕ ∇ x log q t ( x , h ) =
Revised: A ⊤ ϕ ∇ x log q t ( x , h ) =

Original: A ⊤ ϕ ∇ x log q t ( z ) , we have ϕ ( ∇ ˆ x log ˆ q t (ˆ z ))
Revised: A ⊤ ϕ ∇ x log q t ( z ) , we have ϕ ( ∇ ˆ x log ˆ q t (ˆ z ))

Original: =
Revised: =

Original: A ϕ A ⊤ ϕ ∇ x log q t ( z ) , where A ϕ represents the matrix corresponding to ϕ .
Revised: A ϕ A ⊤ ϕ ∇ x log q t ( z ) , where A ϕ represents the matrix corresponding to ϕ .

Original: According to Proposition 1, we have ϕ ( ∇ ˆ x log ˆ q t (ˆ z ))
Revised: According to Proposition 1, we have ϕ ( ∇ ˆ x log ˆ q t (ˆ z ))

Original: = ∇ x log q t ( z ) − ∇ x log q t ( z ) .
Revised: = ∇ x log q t ( z ) − ∇ x log q t ( z ) .

Original: Besides, ∇ h log ˆ q t (ˆ z ) =
Revised: Besides, ∇ h log ˆ q t (ˆ z ) =

Original: ∇ h log q t ( z ) .
Revised: ∇ h log q t ( z ) .

Original: Thus, Eq.
Revised: Thus, Eq.

Original: (12) can be written as
Revised: (12) can be written as

Original: [Equation]
Revised: [Equation]

Original: ˜
Revised: ˜

Original: [Equation]
Revised: [Equation]

Original: ) , which is the score function form of the reverse-time SDE.
Revised: ) , which is the score function form of the reverse-time SDE.

Original: We can also write the score function in Eq.
Revised: We can also write the score function in Eq.

Original: (11) as ∇ ˆ z log ˆ q t (ˆ z ) =
Revised: (11) as ∇ ˆ z log ˆ q t (ˆ z ) =

Original: E ˆ q
Revised: E ˆ q

Original: (ˆ
Revised: (ˆ

Original: [Equation]
Revised: [Equation]

Original: | ˆ z t ) ∇ ˆ z log ˆ q (ˆ z
Revised: | ˆ z t ) ∇ ˆ z log ˆ q (ˆ z

Original: [Equation]
Revised: [Equation]

Original: | ˆ z
Revised: | ˆ z

Original: [Equation]
Revised: [Equation]

Original: ˆ
Revised: ˆ

Original: [Equation]
Revised: [Equation]

Original: | ˆ
Revised: | ˆ

Original: [Equation]
Revised: [Equation]

Original: t ) ˆ ϵ t , where ˆ
Revised: t ) ˆ ϵ t , where ˆ

Original: [Equation]
Revised: [Equation]

Original: is the standard Gaussian noise injected to ˆ z 0 .
Revised: is the standard Gaussian noise injected to ˆ z 0 .

Original: With this expression, we have T ( ∇ ˆ z log ˆ q t (ˆ z ))
Revised: With this expression, we have T ( ∇ ˆ z log ˆ q t (ˆ z ))

Original: = −
Revised: = −

Original: [Equation]
Revised: [Equation]

Original: E q (
Revised: E q (

Original: [Equation]
Revised: [Equation]

Original: | z
Revised: | z

Original: [Equation]
Revised: [Equation]

Original: ϵ t , where ϵ
Revised: ϵ t , where ϵ

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: t |is the standard Gaussian noise in X × R Md injected to z 0 .
Revised: t |is the standard Gaussian noise in X × R Md injected to z 0 .

Original: Thus, Eq.
Revised: Thus, Eq.

Original: (12) can also be written as
Revised: (12) can also be written as

Original: [Equation]
Revised: [Equation]

Original: which is the noise prediction form of the reverse-time SDE.
Revised: which is the noise prediction form of the reverse-time SDE.

Original: A. E QUIVARIANCE
Revised: A. E QUIVARIANCE

Original: Theorem 1.
Revised: Theorem 1.

Original: Let ( ϵ x θ ( z t , t ) , ϵ h
Revised: Let ( ϵ x θ ( z t , t ) , ϵ h

Original: θ ( z t , t ))
Revised: θ ( z t , t ))

Original: =
Revised: =

Original: ϵ θ ( z t , t ) , where ϵ x θ ( z t , t ) and ϵ h
Revised: ϵ θ ( z t , t ) , where ϵ x θ ( z t , t ) and ϵ h

Original: θ ( z t , t ) are the predicted noise of x t and h t respectively.
Revised: θ ( z t , t ) are the predicted noise of x t and h t respectively.

Original: If for any orthogonal transformation R ∈ R n × n , ϵ θ ( z t , t ) is equivariant to R , i.e., ϵ θ
Revised: If for any orthogonal transformation R ∈ R n × n , ϵ θ ( z t , t ) is equivariant to R , i.e., ϵ θ

Original: ( Rx t , h t , t ) =
Revised: ( Rx t , h t , t ) =

Original: ( Rϵ x θ ( x t , h t , t ) , ϵ h θ ( x t , h t , t )) , and p T ( z T ) is invariant
Revised: ( Rϵ x θ ( x t , h t , t ) , ϵ h θ ( x t , h t , t )) , and p T ( z T ) is invariant

Original: to
Revised: to

Original: R , i.e., p T ( Rx T , h T ) =
Revised: R , i.e., p T ( Rx T , h T ) =

Original: p T ( x T , h T ) , then p θ ( z 0 ) is invariant to any rotational transformation.
Revised: p T ( x T , h T ) , then p θ ( z 0 ) is invariant to any rotational transformation.

Original: Proof.
Revised: Proof.

Original: Suppose R ∈ R n × n is an orthogonal transformation.
Revised: Suppose R ∈ R n × n is an orthogonal transformation.

Original: Let z θ t =
Revised: Let z θ t =

Original: ( x θ t , h θ t ) ( 0 ≤ t ≤ T ) be the process determined by Eq.
Revised: ( x θ t , h θ t ) ( 0 ≤ t ≤ T ) be the process determined by Eq.

Original: Let y θ t =
Revised: Let y θ t =

Original: Rx θ t and u θ t =
Revised: Rx θ t and u θ t =

Original: ( y θ t , h θ t ) .
Revised: ( y θ t , h θ t ) .

Original: We use p u t ( u t ) and p z t ( z t ) to denote the distributions of u θ t and z θ t respectively, and they satisfy p u t ( y t , h t ) =
Revised: We use p u t ( u t ) and p z t ( z t ) to denote the distributions of u θ t and z θ t respectively, and they satisfy p u t ( y t , h t ) =

Original: p z t ( R − 1 y t , h t ) .
Revised: p z t ( R − 1 y t , h t ) .

Original: By applying the transformation T z
Revised: By applying the transformation T z

Original: =
Revised: =

Original: ( Rx , h ) to Eq.
Revised: ( Rx , h ) to Eq.

Original: (5), we know the new process { u θ t } Tt = satisfies the following SDE:
Revised: (5), we know the new process { u θ t } Tt = satisfies the following SDE:

Original: [Equation]
Revised: [Equation]

Original: / by equivariance =
Revised: / by equivariance =

Original: [Equation]
Revised: [Equation]

Original: // Winner process is isotropic and its initial distribution is p u T
Revised: // Winner process is isotropic and its initial distribution is p u T

Original: [Equation]
Revised: [Equation]

Original: u
Revised: u

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: This indicates that the distribution of u θ 0 is the same to the distribution of z θ 0 , i.e., p u 0 ( u 0 ) =
Revised: This indicates that the distribution of u θ 0 is the same to the distribution of z θ 0 , i.e., p u 0 ( u 0 ) =

Original: p z 0 ( u 0 ) =
Revised: p z 0 ( u 0 ) =

Original: p θ ( u 0 ) .
Revised: p θ ( u 0 ) .

Original: Also note that p u 0 ( u 0 ) =
Revised: Also note that p u 0 ( u 0 ) =

Original: p z 

Revised: p z

Original: ( R − 1 y 0 , h 0 )
Revised: ( R − 1 y 0 , h 0 )

Original: =
Revised: =

Original: p θ ( R − 1 y 0 , h 0 ) .
Revised: p θ ( R − 1 y 0 , h 0 ) .

Original: Thus, p θ ( u 0 ) =
Revised: Thus, p θ ( u 0 ) =

Original: p θ ( R − 1 y 0 , h 0 ) , and consequently p θ ( Rx 0 , h 0 ) =
Revised: p θ ( R − 1 y 0 , h 0 ) , and consequently p θ ( Rx 0 , h 0 ) =

Original: p θ ( x 0 , h 0 ) .
Revised: p θ ( x 0 , h 0 ) .

Original: This means p θ ( z 0 ) is invariant to any orthogonal transformation, which includes rotational transformations as special cases.
Revised: This means p θ ( z 0 ) is invariant to any orthogonal transformation, which includes rotational transformations as special cases.

Original: Theorem 2.
Revised: Theorem 2.

Original: Suppose the assumptions in Theorem 1 hold and E ( z , c, t ) is invariant to any orthogonal transformation R , i.e., E ( Rx , h , c, t )
Revised: Suppose the assumptions in Theorem 1 hold and E ( z , c, t ) is invariant to any orthogonal transformation R , i.e., E ( Rx , h , c, t )

Original: =
Revised: =

Original: E ( x , h , c, t ) .
Revised: E ( x , h , c, t ) .

Original: Then p θ ( z 0 | c ) is invariant to any rotational transformation.
Revised: Then p θ ( z 0 | c ) is invariant to any rotational transformation.

Original: Proof.
Revised: Proof.

Original: Suppose R ∈ R n × n is an orthogonal transformation.
Revised: Suppose R ∈ R n × n is an orthogonal transformation.

Original: Taking gradient to both side of
Revised: Taking gradient to both side of

Original: E ( Rx , h , c, t ) =
Revised: E ( Rx , h , c, t ) =

Original: E ( x , h , c, t ) w.r.t.
Revised: E ( x , h , c, t ) w.r.t.

Original: x , we get
Revised: x , we get

Original: [Equation]
Revised: [Equation]

Original: Multiplying R to both sides, we get
Revised: Multiplying R to both sides, we get

Original: [Equation]
Revised: [Equation]

Original: Let ϕ ( z , c, t )
Revised: Let ϕ ( z , c, t )

Original: =
Revised: =

Original: ( ∇ x E ( z , c, t ) − ∇ x E ( z , c, t ) , ∇ h E ( z , c, t )) .
Revised: ( ∇ x E ( z , c, t ) − ∇ x E ( z , c, t ) , ∇ h E ( z , c, t )) .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: x , h , c, t) , ∇ h E ( Rx , h , c, t )) =( R ( [Equation] E ( x , h , c, t )
Revised: x , h , c, t

Original: [Equation]

Original: E ( x , h , c, t ))

Original: [Equation]
Revised: [Equation]

Original: E ( x , h , c, t )) .

Original: Thus, ϕ ( z , c, t ) is equivariant to R .
Revised: Thus, ϕ ( z , c, t ) is equivariant to R .

Original: Let ˆ ϵ θ
Revised: Let ˆ ϵ θ

Original: ( z , c, t ) =
Revised: ( z , c, t ) =

Original: ϵ θ ( z , t ) + (cid:112) β t | 0 ϕ
Revised: ϵ θ ( z , t ) + (cid:112) β t | 0 ϕ

Original: ( z , c, t ) , which is a linear combination of two equivariant functions and is also equivariant to R .
Revised: ( z , c, t ) , which is a linear combination of two equivariant functions and is also equivariant to R .

Original: Then, Eq.
Revised: Then, Eq.

Original: (6) can be written as
Revised: (6) can be written as

Original: [Equation]
Revised: [Equation]

Original: According to Theorem 1, we know its marginal distribution at time t = 0 , i.e., p θ ( z 0 | c ) , is invariant to any rotational transformation. 
Revised: According to Theorem 1, we know its marginal distribution at time t = 0 , i.e., p θ ( z 0 | c ) , is invariant to any rotational transformation. B S AMPLING

Revised: In Algorithm 1, we present the Euler-Maruyama method to sample from EEGSDE in Eq. 6.

Revised: Algorithm 1 Sample from EEGSDE using the Euler-Maruyama method

Original: B C ONDITIONAL N OISE P REDICTION N ETWORKS
Revised: Require: Number of steps N ∆ t = TN z ← ( x − x , h ) , where x ∼ N ( 0 , 1) , h ∼ N ( 0 , 1) { Sample from the prior p T ( z T ) } for i = N to 1 do t ← i ∆ t g [Equation] ← ∇

Revised: [Equation]

Revised: ( z , c,

Revised: [Equation]

Revised: ← ∇

Revised: [Equation]

Revised: E ( z , c, t ) { Calculate the gradient of the energy function } g ← ( g x − g x , g h ) { Subtract the CoM of the gradient }

Revised: F ← f ( t )

Revised: [Equation]

Revised: ( t )

Revised: [Equation]

Revised: [Equation]

Revised: ( z ,

Revised: [Equation]

Revised: [Equation]

Revised: g ) ϵ ← ( ϵ x − ϵ x , ϵ h ) , where ϵ x ∼ N ( 0 , 1) , ϵ h ∼ N ( 0 , 1) z ← z − F ∆ t + g ( t ) √ ∆ t ϵ { Update z according to Eq.

Revised: (6) } end for return z

Revised: C C ONDITIONAL N OISE P REDICTION N ETWORKS

Original: In Eq.
Revised: In Eq.

Original: (6), we can alternatively use a conditional noise prediction network ϵ θ
Revised: (6), we can alternatively use a conditional noise prediction network ϵ θ

Original: ( z , c, t ) for a stronger guidance, as follows
Revised: ( z , c, t ) for a stronger guidance, as follows

Original: [Equation]
Revised: [Equation]

Original: The conditional noise prediction network is trained similarly to the unconditional one, using the following MSE loss
Revised: The conditional noise prediction network is trained similarly to the unconditional one, using the following MSE loss

Original: [Equation]
Revised: [Equation]

Original: C P ARAMETERIZATION OF N OISE P REDICTION N ETWORKS
Revised: D P ARAMETERIZATION OF N OISE P REDICTION N ETWORKS

Original: We parameterize the noise prediction network following Hoogeboom et al.
Revised: We parameterize the noise prediction network following Hoogeboom et al.

Original: , and we provide the specific parameterization for completeness.
Revised: , and we provide the specific parameterization for completeness.

Original: For the unconditional model ϵ θ ( z , t ) , we first concatenate each atom feature
Revised: For the unconditional model ϵ θ ( z , t ) , we first concatenate each atom feature

Original: h i and t , which gives h i ′ =
Revised: h i and t , which gives h i ′ =

Original: ( h i , t ) .
Revised: ( h i , t ) .

Original: Then we input x and h ′ =
Revised: Then we input x and h ′ =

Original: ( h 1 ′ , . .
Revised: ( h 1 ′ , . .

Original: , h M ′ ) to the EGNN as follows
Revised: , h M ′ ) to the EGNN as follows

Original: [Equation]
Revised: [Equation]

Original: Finally, we subtract the CoM of a x , and gets the parameterization of ϵ θ ( z , t ) :
Revised: Finally, we subtract the CoM of a x , and gets the parameterization of ϵ θ ( z , t ) :

Original: [Equation]
Revised: [Equation]

Original: ) , where a h comes from discarding the last component of a h ′ that corresponds to the time.
Revised: where a h comes from discarding the last component of a h ′ that corresponds to the time.

Original: For the conditional model ϵ θ ( z , c, t ) , we additionally concatenate c to the atom feature
Revised: For the conditional model ϵ θ ( z , c, t ) , we additionally concatenate c to the atom feature

Original: h
Revised: h

Original: i , i.e., h i ′ =
Revised: i , i.e., h i ′ =

Original: ( h i , t, c ) , and other parts in the parameterization remain the same.
Revised: ( h i , t, c ) , and other parts in the parameterization remain the same.

Revised: E P ARAMETERIZATION OF THE M ULTI -L ABEL C LASSIFIER

Revised: The multi-label classifier is parameterized by EGNN as

Revised: [Equation]

Revised: The multi-label classifier has the same backbone to the property prediction model in Eq.

Original: D T RAINING O BJECTIVES OF E NERGY F UNCTIONS
Revised: (8), except that the decoder outputs a vector of dimension L , and the sigmoid function σ is adopted for multi-label classification. Similarly to Eq. (8), the EGNN in the multi-label classifier guarantees the invariance of the distribution of generated samples according to Theorem 2. F T RAINING O BJECTIVES OF E NERGY F UNCTIONS

Original: Time-dependent property prediction model.
Revised: Time-dependent property prediction model.

Original: Since the quantum property is a scalar, we train the time-dependent property prediction model g ( z t , t ) using the ℓ 1 loss
Revised: Since the quantum property is a scalar, we train the time-dependent property prediction model g ( z t , t ) using the ℓ 1 loss

Original: [Equation]
Revised: [Equation]

Original: where t is uniformly sampled from [0 , T ] .
Revised: where t is uniformly sampled from [0 , T ] .

Original: Time-dependent multi-label classifier.
Revised: Time-dependent multi-label classifier.

Original: Since the fingerprint is a bit map, predicting it can be viewed as a multi-label classification task.
Revised: Since the fingerprint is a bit map, predicting it can be viewed as a multi-label classification task.

Original: Thus, we use a time dependent multi-label classifiermodel m ( z t , t ) , and train it using the binary cross entropy loss
Revised: Thus, we use a time dependent multi-label classifier m ( z t , t ) , and train it using the binary cross entropy loss

Original: [Equation]
Revised: [Equation]

Original: where t is uniformly sampled from [0 , T ] , and m l ( z t , t ) is the l -th component of m ( z t , t ) .
Revised: where t is uniformly sampled from [0 , T ] , and m l ( z t , t ) is the l -th component of m ( z t , t ) .

Original: E E XPERIMENTAL D ETAILS E. T HE “U- BOUND ” AND “#A TOMS ” B ASELINES
Revised: G E XPERIMENTAL D ETAILS G. H OW TO G ENERATE M OLECULES T ARGETED TO M ULTIPLE Q UANTUM P ROPERTIES

Revised: When we want to generate molecules with K quantum properties c = ( c 1 , c 2 , . . . , c K ) , we combine energy functions for single properties linearly as E ( z t , c , t ) =

Revised: (cid:80)

Revised: Kk =1 E k ( z t , c k , t ) , where

Revised: E k ( z t , c k , t ) =

Revised: s k | g k ( z t , t ) − c k | 2 is the energy function for the k -th property, s k is the scaling factor and g k ( z t , t ) is the time-dependent property prediction model for the k -th property.

Revised: Then we use the gradient of E ( z t , c , t ) to guide the reverse SDE as described in Eq.

Revised: G. T HE “U- BOUND ” AND “#A TOMS ” B ASELINES

Original: The “U-bound” and “#Atoms” baselines are from Hoogeboom et al.
Revised: The “U-bound” and “#Atoms” baselines are from Hoogeboom et al.

Original: The “U-bound” baseline shuffles the labels in D b and then calculate the loss of ϕ c on it, which can be regarded as an upper bound of the MAE.
Revised: The “U-bound” baseline shuffles the labels in D b and then calculate the loss of ϕ c on it, which can be regarded as an upper bound of the MAE.

Original: The “#Atoms” baseline predicts a quantum property c using only the number of atoms in a molecule.
Revised: The “#Atoms” baseline predicts a quantum property c using only the number of atoms in a molecule.

Original: E. G ENERATING M OLECULES WITH D ESIRED Q UANTUM P ROPERTIES
Revised: G. G ENERATING M OLECULES WITH D ESIRED Q UANTUM P ROPERTIES

Original: For the noise prediction network, we use the same setting with EDM Hoogeboom et al.
Revised: For the noise prediction network, we use the same setting with EDM Hoogeboom et al.

Original: for a fair comparison, where the models is trained ∼ 2000 epochs with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.
Revised: for a fair comparison, where the models is trained ∼ 2000 epochs with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.

Original: The EGNN used in the energy function has 192 hidden features and 7 layers.
Revised: The EGNN used in the energy function has 192 hidden features and 7 layers.

Original: We train 2000 epochs with a batch size of 128, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.
Revised: We train 2000 epochs with a batch size of 128, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.

Original: E. G ENERATING M OLECULES WITH T ARGET S TRUCTURES
Revised: G. G ENERATING M OLECULES WITH T ARGET S TRUCTURES

Original: QM9.
Revised: QM9.

Original: For the backbone of the noise prediction network, we use a three-layer MLP with 768,and 192 hidden nodes as embedding to encode the fingerprint and add the output of it with the embedding of atom features h , which is then fed into the following EGNN.
Revised: For the backbone of the noise prediction network, we use a three-layer MLP with 768,and 192 hidden nodes as embedding to encode the fingerprint and add the output of it with the embedding of atom features h , which is then fed into the following EGNN.

Original: The EGNN hashidden features and 9 layers.
Revised: The EGNN hashidden features and 9 layers.

Original: We train it 1500 epoch with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.
Revised: We train it 1500 epoch with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.

Original: The energy function is trained with 1750 epoch with a batch size of 128, a learning rate of 0.
Revised: The energy function is trained with 1750 epoch with a batch size of 128, a learning rate of 0.

Original: with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.
Revised: with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.

Original: Its
Revised: Its

Original: EGNN has 192 hidden features and 7 layers.
Revised: EGNN has 192 hidden features and 7 layers.

Revised: For the baseline cG-SchNet, we reproduce it using the public code.

Revised: Since the default data split in cG-SchNet is different with ours, we train the cG-SchNet under the same data split with ours for a fair comparison.

Revised: We report results at 200 epochs (there is no gain on similarity metric after 150 epochs).

Revised: We evaluate the Tanimoto similarity on the whole test set.

Original: GEOM-Drug.
Revised: GEOM-Drug.

Original: We use the same data split of GEOM-Drug with Hoogeboom et al.
Revised: We use the same data split of GEOM-Drug with Hoogeboom et al.

Original: (2022), where training, validation and test set include 554K, 70K and 70K samples respectively.
Revised: (2022), where training, validation and test set include 554K, 70K and 70K samples respectively.

Original: We train the noise prediction network and the energy function on the trainingset and evaluate the Tanimoto similarity on the randomly selected 10k molecules from the test set.
Revised: We train the noise prediction network and the energy function on the training set.

Original: For the noise prediction network, we use the recommended hyperparameters of EDM (Hoogeboom et al., 2022), where the EGNN hashidden features and 4 layers, and the other part is the same as that in QM9.
Revised: For the noise prediction network, we use the recommended hyperparameters of EDM (Hoogeboom et al., 2022), where the EGNN has 256 hidden features and 4 layers, and the other part is the same as that in QM9.

Original: We train 10 epoch with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.
Revised: We trainepoch with a batch size of 64, a learning rate of 0.0001 with the Adam optimizer and an exponential moving average (EMA) with a rate of 0.9999.

Original: The backbone of the energy function is the same as that in QM9 and we train the energy function for 14 epoch.
Revised: The backbone of the energy function is the same as that in QM9 and we train the energy function for 14 epoch.

Revised: We evaluate the Tanimoto similarity on randomly selected 10K molecules from the test set.

Original: F A DDITIONAL R ESULTS F. A BLATION S TUDY ON N OISE P REDICTION N ETWORKS AND E NERGY G UIDANCE
Revised: H A DDITIONAL R ESULTS H. A BLATION S TUDY ON N OISE P REDICTION N ETWORKS AND E NERGY G UIDANCE

Original: We perform an ablation study on the conditioning, i.e., use conditional or unconditional noise prediction networks, and the energy guidance.
Revised: We perform an ablation study on the conditioning, i.e., use conditional or unconditional noise prediction networks, and the energy guidance.

Original: When neither the conditioning nor the energy guidance is adopted, a single unconditional model can’t perform conditional generation, and therefore we only report its atom stability and the molecule stability.
Revised: When neither the conditioning nor the energy guidance is adopted, a single unconditional model can’t perform conditional generation, and therefore we only report its atom stability and the molecule stability.

Original: As shown in Table 4, Table 5 and Table 6, the conditional noise prediction network improves the MAE compared to the unconditional one, and the energy guidance improves the MAE compared to a sole conditional model.
Revised: As shown in Table 4, Table 5 and Table 6, the conditional noise prediction network improves the MAE compared to the unconditional one, and the energy guidance improves the MAE compared to a sole conditional model.

Original: Both the conditioning and the energy guidance do not affect the atom stability and the molecule stability much.
Revised: Both the conditioning and the energy guidance do not affect the atom stability and the molecule stability much.

Original: [Table] [Table] 
Revised: [Table] [Table] [Table]

Revised: H. R ESULTS ON N OVELTY , A TOM S TABILITY AND M OLECULE S TABILITY

Original: [Table]
Revised: For completeness, we also report novelty, atom stability and molecule stability on 10K generated molecules. Below we briefly introduce these metrics. • Novelty (Simonovsky & Komodakis, 2018) is the proportion of generated molecules that do not appear in the training set. Specifically, let G be the set of generated molecules, the novelty is calculated as 1 − | G ∩ D b | | G | .

Original: F. R EPRODUCE
Revised: Note that the novelty is evaluated on D b , since the reproduced conditional EDM and our method are trained on D b . This leads to an inflated value compared to the one evaluated on the whole dataset (Hoogeboom et al., 2022). • Atom stability (AS) (Hoogeboom et al., 2022) is the proportion of atoms that have the right valency. Molecule stability (MS) (Hoogeboom et al., 2022) is the proportion of gener-

Revised: ated molecules where all atoms are stable.

Revised: We use the official implementation for the two metrics from the EDM paper (Hoogeboom et al., 2022).

Revised: As shown in Table 7 and Table

Revised: [Table]

Revised: conditional EDM and EEGSDE with a small scaling factor have a slightly better stability, and EEGSDE with a large scaling factor has a slightly better novelty in general.

Revised: The additional energy changes the distribution of generated molecules, which improves the novelty of generated molecules in general.

Revised: Since there is a tradeoff between novelty and stability (see the caption of Table 5 in the EDM paper (Hoogeboom et al., 2022)), a slight decrease on the stability is possible when the scaling factor is large

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: H.

Original: We compare our reproduced results and the original results of EDM (Hoogeboom et al., 2022) in 
Revised: We visualize the effect of the scaling factor in Figure [Table] where the generated structures align better as the scaling factor grows

Original: Table 7.
Revised: [Table] [Figure] S = S = 0.

Revised: H. R EPRODUCE

Revised: We compare our reproduced results and the original results of conditional EDM (Hoogeboom et al., 2022) in Table 9.

Original: The results are consistent.
Revised: The results are consistent.

Original:  [Table]
Revised: [Figure] [Table]

Original: F. G ENERATING M OLECULES WITH T ARGET S TRUCTURES ON GEOM-D RUG
Revised: H. G ENERATING M OLECULES WITH T ARGET S TRUCTURES ON GEOM-D RUG

Original: We provide the Tanimoto similarity between target structures and generated structures in Table 8.

Original: The conditional EDM baseline is trained following the default setting of Hoogeboom et al.

Original: We plot generated molecules on GEOM-Drug in Figure
Revised: We plot generated molecules on GEOM-Drug in Figure

Original: [Table]
Revised: [Table]

Original: and it can be observed that the atom types of generated molecules with EEGSDE often match the target better than EDM.
Revised: and it can be observed that the atom types of generated molecules with EEGSDE often match the target better than conditional EDM

Original: [Table] [Figure]
Revised: [Table] [Figure]