Original: N EURAL E NERGY M INIMIZATION FOR M OLECULAR C ONFORMATION O PTIMIZATION
Revised: N EURAL E NERGY M INIMIZATION FOR M OLECULAR C ONFORMATION O PTIMIZATION

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: This paper studies an important problem in computational chemistry: predicting a molecule’s spatial atom arrangements, or a molecular conformation.
Revised: This paper studies an important problem in computational chemistry: predicting a molecule’s spatial atom arrangements, or a molecular conformation.

Original: We propose a neural energy minimization formulation that casts the prediction problem into an unrolled optimization process, where a neural network is parametrized to learn the gradient ﬁelds of a conformational energy landscape.
Revised: We propose a neural energy minimization formulation that casts the prediction problem into an unrolled optimization process, where a neural network is parametrized to learn the gradient ﬁelds of a conformational energy landscape.

Original: Assuming different forms of the underlying potential energy function, we can not only reinterpret and unify many of the existing models but also derive new variants of SE(3)-equivariant neural networks in a principled manner.
Revised: Assuming different forms of the underlying potential energy function, we can not only reinterpret and unify many of the existing models but also derive new variants of SE(3)-equivariant neural networks in a principled manner.

Original: In our experiments, these new variants show superior performance in molecular conformation optimization comparing to existing SE(3)-equivariant neural networks.
Revised: In our experiments, these new variants show superior performance in molecular conformation optimization comparing to existing SE(3)-equivariant neural networks.

Original: Moreover, our energy minimization formulation is also suitable for molecular conformation generation, where we can generate more diverse and accurate conformers comparing to existing baselines.
Revised: Moreover, our energy minimization formulation is also suitable for molecular conformation generation, where we can generate more diverse and accurate conformers comparing to existing baselines.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The 3D molecular conformation is one of the most important features in determining many physicochemical and biological properties of a molecule.
Revised: The 3D molecular conformation is one of the most important features in determining many physicochemical and biological properties of a molecule.

Original: The molecule’s 3D charge distribution and physical shape are crucial when considering the steric constraints or electronic effects for chemical reactions and interaction.
Revised: The molecule’s 3D charge distribution and physical shape are crucial when considering the steric constraints or electronic effects for chemical reactions and interaction.

Original: Therefore, molecular conformers are widely adopted in quantitative structure-activity relationships (QSAR) prediction (Verma et al., 2010) and drug discovery (Hawkins, 2017).
Revised: Therefore, molecular conformers are widely adopted in quantitative structure-activity relationships (QSAR) prediction (Verma et al., 2010) and drug discovery (Hawkins, 2017).

Original: With a growing interest in virtual drug screening (Zhavoronkov et al., 2019; McCloskey et al., 2020; Stokes et al., 2020) and de novo drug design (De Cao & Kipf, 2018; Gómez-Bombarelli et al., 2018; Jin et al., 2018;
Revised: With a growing interest in virtual drug screening (Zhavoronkov et al., 2019; McCloskey et al., 2020; Stokes et al., 2020) and de novo drug design (De Cao & Kipf, 2018; Gómez-Bombarelli et al., 2018; Jin et al., 2018;

Original: Zhou et al., 2019; Mercado et al., 2020), it becomes more desirable to predict the molecular conformations both quickly and accurately such that 3D geometry features can be considered.
Revised: Zhou et al., 2019; Mercado et al., 2020), it becomes more desirable to predict the molecular conformations both quickly and accurately such that 3D geometry features can be considered.

Original: Deep learning techniques have recently been introduced to address this problem and demonstrated promising results.
Revised: Deep learning techniques have recently been introduced to address this problem and demonstrated promising results.

Original: For instance, various generative models (Gebauer et al., 2019; Hoffmann & Noé, 2019; Simm & Hernandez-Lobato, 2020; Xu et al., 2020) are developed to ﬁrst generate a pairwise distance matrix for atom pairs, and then infer the conformer from the matrix.
Revised: For instance, various generative models (Hoffmann & Noé, 2019; Simm & Hernandez-Lobato, 2020; Xu et al., 2020) are developed to ﬁrst generate a pairwise distance matrix for atom pairs, and then infer the conformer from the matrix.

Original: Since the 3D structures are generated indirectly, the second step can be sensitive towards the error in the distance matrix, and extra model capacity could be required to encode the redundancy observed in such  matrix (Hoffmann & Noé, 2019).
Revised: Since the 3D structures are generated indirectly, the second step can be sensitive towards the error in the distance matrix, and extra model capacity could be required to encode the redundancy observed in such a matrix (Hoffmann & Noé, 2019).

Original: In the meantime, within a different thread of research, many SE(3)-equivariant networks have been proposed to work with 3D roto-translation symmetrical inputs (Finzi et al., 2020; 
Revised: In the meantime, within a different thread of research, many SE(3)-equivariant networks have been proposed to work with 3D roto-translation symmetrical inputs (Finzi et al., 2020; Fuchs et al., 2020;

Original: Fuchs et al., 2020; Horie et al., 2020;
Revised: Horie et al., 2020;

Original: Satorras et al., 2021; Thomas et al., 2018), enabling direct manipulation of the 3D Cartesian coordinates.
Revised: Satorras et al., 2021; Thomas et al., 2018), enabling direct manipulation of the 3D Cartesian coordinates.

Original: One of the most recent successful applications of such models is the AlphaFold system for protein structure prediction (Jumper et al., 2021).
Revised: One of the most recent successful applications of such models is the AlphaFold system for protein structure prediction (Jumper et al., 2021).

Original: Despite the many beneﬁts, however, existing SE(3)-equivariant networks are either derived from complicated mathematical theory involving expensive coefﬁcient calculations (Thomas et al., 2018; Anderson et al., 2019;
Revised: Despite the many beneﬁts, however, existing SE(3)-equivariant networks are either derived from complicated mathematical theory involving expensive coefﬁcient calculations (Thomas et al., 2018; Anderson et al., 2019;

Original: Finzi et al., 2020; Fuchs et al., 2020), or heuristically designed (Satorras et al., 2021)
Revised: Finzi et al., 2020; Fuchs et al., 2020), or heuristically designed (Satorras et al., 2021)

Original: based on the intuition of message passing (Gilmer et al., 2017).
Revised: based on the intuition of message passing (Gilmer et al., 2017).

Original: This work presents a new formulation that unrolls the molecular conformer prediction into a ﬁxed number of optimization steps, where a neural network is parametrized to learn gradient ﬁelds of the conformational energy landscape.
Revised: This work presents a new formulation that unrolls the molecular conformer prediction into a ﬁxed number of optimization steps, where a neural network is parametrized to learn gradient ﬁelds of the conformational energy landscape.

Original: Under such formulation, the model can reﬁne any conformer toward a more stable and energetically preferred state as its prediction.
Revised: Under such formulation, the model can reﬁne any conformer toward a more stable and energetically preferred state as its prediction.

Original: Instead of relying on complex mathematical tools or designing the model heuristically, we propose a novel framework based on the perspective of neural energy minimization.
Revised: Instead of relying on complex mathematical tools or designing the model heuristically, we propose a novel framework based on the perspective of neural energy minimization.

Original: Such perspective not only provides new interpretations for some existing models but also allows us to derive new variants of SE(3)-equivariant model from a designed conformational energy function, which clearly aligns a model’s architecture with the physical assumptions.
Revised: Such perspective not only provides new interpretations for some existing models but also allows us to derive new variants of SE(3)-equivariant model from a designed conformational energy function, which clearly aligns a model’s architecture with the physical assumptions.

Original: Through extensive experiments, these new variants show superior performance in conformer optimization comparing to existing SE(3)-equivariant neural networks.
Revised: Through extensive experiments, these new variants show superior performance in conformer optimization comparing to existing SE(3)-equivariant neural networks.

Original: Moreover, the proposed optimization method is also suitable for conformer generation and capable of generating more diverse and structurally accurate conformers comparing to existing baselines.
Revised: Moreover, the proposed optimization method is also suitable for conformer generation and capable of generating more diverse and structurally accurate conformers comparing to existing baselines.

Original: In conclusion, the main contributions of this work include:
Revised: In conclusion, the main contributions of this work include:

Original: • Formulating the conformer prediction problem as an unrolled optimization where a model learns the gradient ﬁeld of the conformational energy landscape.
Revised: • Formulating the conformer prediction problem as an unrolled optimization where a model learns the gradient ﬁeld of the conformational energy landscape.

Original: • Proposing a novel framework for deriving variants of SE(3)-equivariant neural networks and reinterpreting existing models from the perspective of neural energy minimization.
Revised: • Proposing a novel framework for deriving variants of SE(3)-equivariant neural networks and reinterpreting existing models from the perspective of neural energy minimization.

Original: • Demonstrating competitive empirical performance for molecular conformation prediction in both optimization and generation settings.
Revised: • Demonstrating competitive empirical performance for molecular conformation prediction in both optimization and generation settings.

Original:  R ELATED W ORK 2. M OLECUALR C ONFORMATION G ENERATION
Revised:  R ELATED W ORK 2. M OLECUALR C ONFORMATION G ENERATION

Original: The conventional computational approach for conformer generation  relies onab initio Molecular
Revised: The conventional computational approach for conformer generation mainly relies on Molecular

Original: Dynamics (MD) simulation (Shim & MacKerell Jr, 2011; De Vivo et al., 2016).
Revised: Dynamics (MD) simulation (Shim & MacKerell Jr, 2011; De Vivo et al., 2016).

Original: However, such simulation often involves Density Functional Theory (DFT) (Parr, 1980), which is prohibitively expensive for large molecules with high degrees of freedom, or hand-crafted force ﬁelds (Rappé et al., 1992; Halgren, 1996), which can be obtained quickly but are often inaccurate (Kanal et al., 2018).
Revised: However, such simulation often involves Density Functional Theory (DFT) (Parr, 1980), which is prohibitively expensive for large molecules with high degrees of freedom, or hand-crafted force ﬁelds (Rappé et al., 1992; Halgren, 1996), which can be obtained quickly but are often inaccurate (Kanal et al., 2018).

Original: To efﬁciently generate high-quality estimates for molecular conformation, many learning-based methods have been proposed.
Revised: To efﬁciently generate high-quality estimates for molecular conformation, many learningbased methods have been proposed.

Original: For instance, various generative models, such as variational autoencoder (Simm & Hernandez-Lobato, 2020),autoregressive models (Gebauer et al., 2019), generative adversarial networks (Hoffmann & Noé, 2019), and continuous normalizing ﬂow (Xu et al., 2020), have all been developed to predict molecular conformers by ﬁrst generating an intermediate pairwise atom distance matrix .
Revised: For instance, various generative models, such as variational autoencoder (Simm & Hernandez-Lobato, 2020), generative adversarial networks (Hoffmann & Noé, 2019), and continuous normalizing ﬂow (Xu et al., 2020), have all been developed to predict molecular conformers by ﬁrst generating an intermediate pairwise atom distance matrix .

Original: However, the post-processing step for 3D structure construction can be susceptible to the error in the distance matrix, and model capacity could be wasted to encode the redundancy in such low-rank matrix (Hoffmann & Noé, 2019).
Revised: However, the post-processing step for 3D structure construction can be susceptible to the error in the distance matrix, and model capacity could be wasted to encode the redundancy in such low-rank matrix (Hoffmann & Noé, 2019).

Original: Despite the many drawbacks of such methods, the distance matrix is difﬁcult to circumvent because the models cannot capture the 3D roto-translation symmetry nature of a conformer.
Revised: Despite the many drawbacks of such methods, the distance matrix is difﬁcult to circumvent because the models cannot capture the 3D roto-translation symmetry nature of a conformer.

Original: Another solution for conformer generation is to estimate the force ﬁeld with machine learning methods.
Revised: Another solution for conformer generation is to estimate the force ﬁeld with machine learning methods.

Original: For instance, sGDML (Chmiela et al., 2019) is capable of reproducing the molecular potential energy surface given a few reference conformations and producing compelling results with MD simulation.
Revised: For instance, sGDML (Chmiela et al., 2019) is capable of reproducing the molecular potential energy surface given a few reference conformations and producing compelling results with MD simulation.

Original: However, such a model needs to be trained separately for different molecules and does not scale well  for larger drug-like molecules.
Revised: Although such a model scales well with the molecule size, it requires retraining for every new molecules (Unke et al., 2021).

Original: Recently, ConfGF (Shi et al., 2021) achieved state-of-the-art performance for conformer generation by estimating the gradient ﬁelds of inter-atomic distances and therefore can be treated as a replacement for force-ﬁeld methods.
Revised: Recently, ConfGF (Shi et al., 2021) achieved state-of-the-art performance for conformer generation by estimating the gradient ﬁelds of inter-atomic distances and therefore can be treated as a replacement for force-ﬁeld methods.

Original: The 3D coordinates of the conformer can be generated through steps of Langevin dynamics in the end.
Revised: The 3D coordinates of the conformer can be generated through steps of Langevin dynamics in the end.

Original: While our model is also trying to estimate the gradient ﬁeld, we instead employ an unrolled optimization architecture (Domke, 2012; Liu et al., 2018) aiming to model the entire optimization process.
Revised: While our model is also trying to estimate the gradient ﬁeld, we instead employ an unrolled optimization architecture (Domke, 2012; Liu et al., 2018) aiming to model the entire optimization process.

Revised: Such design with one-pass end-to-end training and inference is proved to be more effective empirically as shown in Table3.

Original: By learning to minimize a parametric energy function in a ﬁxed number of steps, our models enjoy a much faster inference time.
Revised: By optimizing the conformer in a ﬁxed number of steps, our approach also enjoys a much faster run time

Original: 2. SE(3)-E QUIVARIANT N EURAL N ETWORKS
Revised: 2. SE(3)-E QUIVARIANT N EURAL N ETWORKS

Original: Since 3D roto-translation symmetry is a common property for many physical systems, various methods(Thomas et al., 2018; Finzi et al., 2020; Fuchs et al., 2020; Horie et al., 2020; Satorras et al., 2021; Dym & Maron, 2020) have been proposed to model such symmetry in problems like quantum mechanics (Ramakrishnan et al., 2014), molecular design (Simm et al., 2020), particle simulations (Thomas et al., 2018), and point clouds (Uy et al., 2019).
Revised: Since 3D roto-translation symmetry is a common property for many physical systems, various methods have been proposed to model such symmetry in problems like quantum mechanics (Ramakrishnan et al., 2014), molecular design (Simm et al., 2020), particle simulations (Thomas et al., 2018), and point clouds (Uy et al., 2019).

Original: The ability to implicitly capture the spatial symmetry makes these models a promising class for performing conformer optimization.
Revised: The ability to implicitly capture the spatial symmetry makes these models a promising class for performing conformer optimization.

Original: For instance,
Revised: In general, there are two ways to introduce the equivariance into a model.

Original:  Fuchs et al. (2021) shows one can apply a SE(3)-Transformers (Fuchs et al., 2020) in an iterative fashion and perform energy minimization involving ten particles with a proof-of-concept toy problem.
Revised: One thread of research relies on building models with irreducible representations and spherical harmonics (Thomas et al., 2018; Anderson et al., 2019; Finzi et al., 2020; Fuchs et al., 2020; Batzner et al., 2021).

Original: More recently, the AlphaFold system also invokes the SE(3)-equivariant neural network for accurate protein structure prediction (Jumper et al., 2021).
Revised: Different from the above works, our approach aims to avoid the complexity of such mathematical tools.

Original: Despite their promising applications in conformer
Revised: The other line of research relies

Original: [Figure]
Revised: [Figure]

Original: optimization, existing SE(3)-equivariant models could be difﬁcult to build upon as they either relied on spherical-harmonics-based techniques involving expensive coefﬁcient calculation (Thomas et al., 2018; Anderson et al., 2019; Finzi et al., 2020; Fuchs et al., 2020) or designed (Satorras et al., 2021) based on the intuition of message passing (Gilmer et al., 2017).
Revised: on vector-based methods, where equivariant ﬁlters are constructed by taking the gradient of invariant ﬁlters (Schütt et al., 2021) or designed (Satorras et al., 2021) based on the intuition of message passing (Gilmer et al., 2017).

Original: However, with our proposed neural energy minimization perspective, we can not only reinterpret some of the existing models but also derive new variants  as shown in Section 3.3.
Revised: However, with our proposed neural energy minimization perspective, we can not only reinterpret some of the existing models but also derive new variants that aligns the model parametrization with the underlying physical assumption as shown in Section 3.3.

Original:  M ETHODS 3. P ROBLEM D EFINITION
Revised:  M ETHODS 3. P ROBLEM D EFINITION

Original: We represent each molecule as a graph G mol and use a set of 3D coordinates X to denote the conformer.
Revised: We represent each molecule as a graph G mol and use a set of 3D coordinates X to denote the conformer.

Original: The molecular graph is deﬁned as G mol =
Revised: The molecular graph is deﬁned as G mol =

Original: (cid:104)V , E(cid:105) where V = { v i } is the set of  atom nodes, and E = { e i } is the set of edges representing the relationship between atoms: including covalent interactions and distance of their shortest path.
Revised: (cid:104)V , E(cid:105) where V = { v i } is the set of nodes representing atom features, and E = { e i } is the set of edges representing the relationship between atoms: including covalent interactions and distance of their shortest path.

Original: The conformer X consists of the 3DCartesian coordinate { x i | x i ∈ R 3 } of each atom.
Revised: X consists of the 3D coordinate { x i | x i ∈ R 3 } of each atom.

Original: Given a molecular graph G mol and an initial conformer X 0 , we want to learn an optimization model φ that reﬁnes the initial conformation towards an energetically preferred state X ∗ as X ∗ =
Revised: Given a molecular graph G mol and an initial conformer X 0 , we want to learn an optimization model φ that reﬁnes the initial conformation towards an energetically preferred state X ∗ as X ∗ =

Original: arg min X E ( X  ) where E is some energy surface of nature.
Revised: arg min X E ( X |G mol ) where E is some energy surface of nature.

Original: 3. N EURAL E NERGY M INIMIZATION
Revised: 3. N EURAL E NERGY M INIMIZATION

Original: Molecular conformation represents a 3D arrangement of the atoms and shares a close relationship with the conformational energy, as each equilibrium of a molecular conformation can be considered as a local minimum of the potential energy surface (Axelrod & Gomez-Bombarelli, 2020).
Revised: Molecular conformation represents a 3D arrangement of the atoms and shares a close relationship with the conformational energy, as each equilibrium of a molecular conformation can be considered as a local minimum of the potential energy surface (Axelrod & Gomez-Bombarelli, 2020).

Original: One could therefore ﬁnd these stable conformations through gradient descent as a process of energy minimization.
Revised: One could therefore ﬁnd these stable conformations through gradient descent as a process of energy minimization.

Original: Speciﬁcally, given a particular energy function E( X ) , we can optimize the conformation X by updating the 3D Cartesian coordinates with gradient descent against some energy function E :
Revised: Speciﬁcally, given a particular energy function E , we can optimize the conformation X by updating the 3D Cartesian coordinates with gradient descent against some energy function E :

Original: [Equation]
Revised: [Equation]

Original:  where γ is the learning rate, and we will observe X ∗ =
Revised: 1) where γ is the learning rate, and we will observe X ∗ =

Original: X ∗ − γ ∇  E ( X ∗  ) for an energetically stable conformer.
Revised: X ∗ − γ ∇ X E ( X ∗ |G mol ) for an energetically stable conformer.

Original: Therefore, we can easily ﬁnd a stable molecular conformation from an initial conformation by applying the gradient descent update iteratively if the energy function E ( X  ) is fully characterized and differentiable with respect to X .
Revised: Therefore, we can easily ﬁnd a stable molecular conformation from an initial conformation by applying the gradient descent update iteratively if the energy function E ( X |G mol ) is fully characterized and differentiable with respect to X .

Original: Unfortunately, this is often not the case, especially for large and complex molecules.
Revised: Unfortunately, this is often not the case, especially for large and complex molecules.

Original: Instead of hand-designing the energy function and its gradient, we can take advantage of the approximation power of a neural network and use it to learn the gradient ﬁeld from the data directly.
Revised: Instead of hand-designing the energy function and its gradient, we can take advantage of the approximation power of a neural network and use it to learn the gradient ﬁeld from the data directly.

Original: Once the network is trained, we can then unroll the optimization (Domke, 2012; Liu et al., 2018) with the model and perform a neural version of the energy minimization where the gradient term ∇ E in
Revised: Once the network is trained, we can then unroll the optimization (Domke, 2012; Liu et al., 2018) with the model and perform a neural version of the energy minimization where the gradient term ∇ E in

Original: Equation 1 is not derived from an explicit energy function but estimated by a neural network.
Revised: Equation 1 is not derived from an explicit energy function but estimated by a neural network.

Original: More concretely, we deﬁne the energy surface as a function of the molecular graph G =
Revised: More concretely, we deﬁne the energy surface as a function of the molecular graph G =

Original: (cid:104)V , E(cid:105) and the 3D conformer X .
Revised: (cid:104)V , E(cid:105) and the 3D conformer X .

Original: We can consider X as an external 3 -dimensional representation of the molecule, while V and E represent an internal high dimensional representation of the molecule.
Revised: We can consider X as an external 3 -dimensional representation of the molecule, while V and E represent an internal high dimensional representation of the molecule.

Original: These internal embeddings V and E not only encode the initial atom/bond information but also embed information about their neighborhood related to the evolving conformation X .
Revised: These internal embeddings V and E not only encode the initial atom/bond information but also embed information about their neighborhood related to the evolving conformation X .

Original: For each step, we want to minimize E ( X , V , E ) with the following gradient updates:
Revised: For each step, we want to minimize E ( X , V , E ) with the following gradient updates:

Original: [Equation]
Revised: [Equation]

Original: for all x i ∈ X , v i ∈ V , and e ij ∈ E , where φ are neural networks estimating their respective gradient ﬁeld, and the exact parametrization of the neural network will depend on the energy function as described in the following section.
Revised: for all x i ∈ X , v i ∈ V , and e ij ∈ E , where φ are neural networks estimating their respective gradient ﬁeld, and the exact parametrization of the neural network will depend on the energy function as described in the following section.

Original: In this work, we view e ij as a constant representation for simplicity, and therefore, it is not updated in our models.
Revised: In this work, we view e ij as a constant representation for simplicity, and therefore, it is not updated in our models.

Original: 3. R ECIPE FOR DERIVING NEURAL ENERGY MINIMIZATION MODELS
Revised: 3. R ECIPE FOR DERIVING NEURAL ENERGY MINIMIZATION MODELS

Original: Taking the perspective of neural energy minimization, we can derive various models by making different assumptions of the target energy function, resulting in a principle architectural design that aligns with the underlying physical assumption.
Revised: Taking the perspective of neural energy minimization, we can derive various models by making different assumptions of the target energy function, resulting in a principle architectural design that aligns with the underlying physical assumption.

Original: In Appendix A, we show that a model derived from an SE(3)-invariant energy function (as they should be) always enjoys the SE(3)-equivariance naturally.
Revised: In Appendix A, we show that a model derived from an SE(3)-invariant energy function (as they should be) always enjoys the SE(3)-equivariance naturally.

Original: This section instantiates three variants of SE(3)-equivariant networks as examples for leveraging such perspective for future model development.
Revised: This section instantiates three variants of SE(3)-equivariant networks as examples for leveraging such perspective for future model development.

Original: Basic Two-Atom Model We start with a simple energy formulation where the conformation energy considers atom pairs independently:
Revised: Basic Two-Atom Model We start with a simple energy formulation where the conformation energy considers atom pairs independently:

Original: [Equation]
Revised: [Equation]

Original: where d ij = (cid:107) x
Revised: where d ij = (cid:107) x

Original: i − x j (cid:107) is the euclidean distance between two atoms i and j while u ( · ) is some unknown the potential energy function in nature.
Revised: i − x j (cid:107) is the euclidean distance between two atoms i and j while u ( · ) is some unknown the potential energy function in nature.

Original: Following the formulation, we can then derive the gradient ﬁeld of atom coordinates:
Revised: Following the formulation [Equation] we can then derive the gradient ﬁeld of atom coordinates:

Original: [Equation]
Revised: [Equation]

Original: where f x is a Transformer-based neural network (Vaswani et al., 2017) approximating the gradient ﬁeld ( ∂u∂d 2 ij ) with a constant scalar.
Revised: where f x is a Transformer-based neural network (Vaswani et al., 2017) approximating the gradient ﬁeld ( ∂u∂d 2 ij ) with a constant scalar.

Original: The optimization step for a two-atom model therefore becomes:
Revised: The optimization step for a two-atom model therefore becomes:

Original: [Equation]
Revised: [Equation]

Original: and since f x is SE(3)-invariant and ( x ti − x tj ) is rotational equivariant, the optimization step Equationobserves SE(3)-equivariance (proof in Appendix A).
Revised: and since f x is SE(3)-invariant and ( x ti − x tj ) is rotational equivariant, the optimization step Equationobserves SE(3)-equivariance (proof in Appendix A).

Original: The equivariance here is consistent with the invariant nature of the energy function, as the gradient ﬁeld should rotate with the input coordinate.
Revised: The equivariance here is consistent with the invariant nature of the energy function, as the gradient ﬁeld should rotate with the input coordinate.

Original: Similarly, we can derive the following update for the internal atom representation:
Revised: Similarly, we can derive the following update for the internal atom representation:

Original: [Equation]
Revised: [Equation]

Original: High-Order / Three-Atom Model While we start with an energy function that only considers atom pairs, the energy function can also be extended to consider atom triplets allowing us to capture the potential energy related to different bond angles:
Revised: High-Order / Three-Atom Model While we start with an energy function that only considers atom pairs, the energy function can also be extended to consider atom triplets allowing us to capture the potential energy related to different bond angles:

Original: [Equation]
Revised: [Equation]

Original: where r ij = x i − x j and (cid:104)· , ·(cid:105) denotes the inner product of two vectors capturing the bond angle ∠ jik .
Revised: where r ij = x i − x j and (cid:104)· , ·(cid:105) denotes the inner product of two vectors capturing the bond angle ∠ jik .

Original: The dependency on internal atom and edge embeddings is also extended to atom triplets V ijk = { v i , v j , v k } and their corresponding relationship representations E ijk = { e ij , e ik , e jk } .
Revised: The dependency on internal atom and edge embeddings is also extended to atom triplets V ijk = { v i , v j , v k } and their corresponding relationship representations E ijk = { e ij , e ik , e jk } .

Original: Following the derivation in Appendix B, we observe the following updates for the three-atom model:
Revised: Following the derivation in Appendix B, we observe the following updates for the three-atom model:

Original: [Equation]
Revised: [Equation]

Original: where SE(3)-equivariance is also achieved.
Revised: where SE(3)-equivariance is also achieved.

Original: Similarly, we have the following update function for the internal atom representation: v t +

Revised: Similarly, we have the following update function for the internal atom representation: v t +

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: While we only extend the framework to atom triplets, one could extend the framework to an even higher-order model to capture the energy contribution for more complex geometric structures such as the torsional angles.
Revised: While we only extend the framework to atom triplets, one could extend the framework to an even higher-order model to capture the energy contribution for more complex geometric structures such as the torsional angles.

Original: Geometry of the Internal Representation In the basic two-atom model, we derive a direct v update formula as Equation 4.
Revised: Geometry of the Internal Representation In the basic two-atom model, we derive a direct v update formula as Equation 4.

Original: However, we can also view v as an internal spatial organization of the atoms, and therefore, the proximity ( P ) between two n -dim vectors also contributes to the energy:
Revised: However, we can also view v as an internal spatial organization of the atoms, and therefore, the proximity ( P ) between two n -dim vectors also contributes to the energy:

Original: [Equation]
Revised: [Equation]

Original: where P is some distance metric for v i and v j .
Revised: where P is some distance metric for v i and v j .

Original: For instance, using cosine similarity P ( v i , v j ) =
Revised: For instance, using cosine similarity P ( v i , v j ) =

Original: [Equation]
Revised: [Equation]

Original: as the distance metric, we arrive at the following update by taking the gradient of E w.r.t v i :
Revised: as the distance metric, we arrive at the following update by taking the gradient of E w.r.t v i :

Original: [Equation]
Revised: [Equation]

Original: Reinterpretation for existing models Besides new variants, the perspective of energy minimization also allows us to reinterpret some existing models.
Revised: Reinterpretation for existing models Besides new variants, the perspective of energy minimization also allows us to reinterpret some existing models.

Original: For instance, in EGNN (Satorras et al., 2021), the authors propose the following updates for atom "representations" v and "coordinates" x :
Revised: For instance, in EGNN (Satorras et al., 2021), the authors propose the following updates for atom "representations" v and "coordinates" x :

Original: [Equation]
Revised: [Equation]

Original: ij , v ti , v tj , e ij )) .
Revised: ij , v ti , v tj , e ij )) .

Original: From the perspective of neural energy minimization, these are equivalent to optimizing an energy function that depends on both the atom presentations and inter-atomic distances just like our basic
Revised: From the perspective of neural energy minimization, these are equivalent to optimizing an energy function that depends on both the atom representations and inter-atomic distances just like our basic

Original: Two-Atom model, i.e. E = (cid:80) i,j ∈V ,i (cid:54) =
Revised: Two-Atom model, i.e. E = (cid:80) i,j ∈V ,i (cid:54) =

Original: j u ( d 2 ij , v i , v j , e ij ) .
Revised: j u ( d 2 ij , v i , v j , e ij ) .

Original: On the other hand, the "gradient" for their atom representations are only aggregated from the corresponding neighbors (as message passing (Gilmer et al., 2017)) instead of all atoms (as ours).
Revised: On the other hand, the "gradient" for their atom representations are only aggregated from the corresponding neighbors (as message passing (Gilmer et al., 2017)) instead of all atoms (as ours).

Original: In addition, this formulation uses a learned function φ h to update the representation v ti instead of simply adding the negative gradient.
Revised: In addition, this formulation uses a learned function φ h to update the representation v ti instead of simply adding the negative gradient.

Original: In the context of optimization, the operation can also be interpreted as optimizing with a learned optimizer (Metz et al., 2019) instead of performing simple gradient descent.
Revised: In the context of optimization, the operation can also be interpreted as optimizing with a learned optimizer (Metz et al., 2019) instead of performing simple gradient descent.

Original: Other reinterpretations for existing models can also be found in Appendix.
Revised: Other reinterpretations for existing models can also be found in Appendix.

Original: C.
Revised: C.

Original: We believe neural energy minimization represents a new perspective for thinking about SE(3)- equivariant model and expect the community to derive more interesting variants by designing different energy functions with more suitable assumptions for their applications.
Revised: We believe neural energy minimization represents a new perspective for thinking about SE(3)- equivariant model and expect the community to derive more interesting variants by designing different energy functions with more suitable assumptions for their applications.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: To evaluate the proposed methods, we perform experiments in two separated settings:
Revised: To evaluate the proposed methods, we perform experiments in two separated settings:

Original: The optimization setting where the goal for the model is to take a 3D conformer as input and produce a single optimized conformer with the most stable energy.
Revised: The optimization setting where the goal for the model is to take a 3D conformer as input and produce a single optimized conformer with the most stable energy.

Original: Since our model is proposed to solve an optimization task, this setting is a direct measure of its performance against other SE(3)- equivariant models also capable of taking 3D Cartesian coordinates as input.
Revised: Since our model is proposed to solve an optimization task, this setting is a direct measure of its performance against other SE(3)- equivariant models also capable of taking 3D Cartesian coordinates as input.

Original: Our results suggest that models derived from higher-order energy function can achieve better performance, and we believe the perspective of neural energy minimization would also accelerate more powerful extensions of
Revised: Our results suggest that models derived from higher-order energy function can achieve better performance, and we believe the perspective of neural energy minimization would also accelerate more powerful extensions of

Original: SE(3)-equivariant models.
Revised: SE(3)-equivariant models.

Original: The generation setting where the goal is to generate multiple conformers that capture a diverse set of relatively stable conformers.
Revised: The generation setting where the goal is to generate multiple conformers that capture a diverse set of relatively stable conformers.

Original: This setting has been studied recently, even outside the context of
Revised: This setting has been studied recently, even outside the context of

Original: SE(3)-equivariant model (Mansimov et al., 2019; Simm & Hernandez-Lobato, 2020; Xu et al., 2020;
Revised: SE(3)-equivariant model (Mansimov et al., 2019; Simm & Hernandez-Lobato, 2020; Xu et al., 2020;

Original: Shi et al., 2021), and represents a more realistic application scenario for practitioners.
Revised: Shi et al., 2021), and represents a more realistic application scenario for practitioners.

Original: Our experiment suggests that the energy minimization formulation can be easily extended for the generative setting.
Revised: Our experiment suggests that the energy minimization formulation can be easily extended for the generative setting.

Original: Compared to existing baselines, our models can generate a more diverse and accurate ensemble of conformers, demonstrating a realistic application scenario for the proposed formulation.
Revised: Compared to existing baselines, our models can generate a more diverse and accurate ensemble of conformers, demonstrating a realistic application scenario for the proposed formulation.

Original: 4. M OLECULAR C ONFORMER O PTIMIZATION
Revised: 4. M OLECULAR C ONFORMER O PTIMIZATION

Original: Data We test molecular conformer optimization on the QM9 dataset with small molecules (up to 9 heavy atoms) (Ramakrishnan et al., 2014) as well as the GEOM-Drugs dataset with larger molecules (up to 91 heavy atoms)
Revised: Data We test molecular conformer optimization on the QM9 dataset with small molecules (up to 9 heavy atoms) (Ramakrishnan et al., 2014) as well as the GEOM-Drugs dataset with larger molecules (up to 91 heavy atoms)

Original: (Axelrod & Gomez-Bombarelli, 2020).
Revised: (Axelrod & Gomez-Bombarelli, 2020).

Original: For the QM9 dataset, only one reference conformer generated through DFT calculation is considered as the lowest-energy conformer.
Revised: For the QM9 dataset, only one reference conformer generated through DFT calculation is considered as the lowest-energy conformer.

Original: For the GEOM-Drugs dataset, since multiple low-energy conformers are given through a semi-empirical method and after realizing larger molecules could have multiple similarly optimal conformers, we select multiple conformers with high Boltzmann weights as reference conformers (see
Revised: For the GEOM-Drugs dataset, since multiple low-energy conformers are given through a semi-empirical method and after realizing larger molecules could have multiple similarly optimal conformers, we select multiple conformers with high Boltzmann weights as reference conformers (see

Original: Appendix D for details).
Revised: Appendix D for details).

Original: We randomly split the QM9 dataset into 123k, 5k, 5k, and the GEOM-Drugs dataset  testing.
Revised: We randomly split the QM9 dataset into 123k, 5k, 5k, and the GEOM-Drugs dataset into 270K, 10K, and 10K for respective training, validation, and testing.

Original: Baselines Since the goal for this experiment is to show the energy-derived model is indeed capable of optimizing an existing 3D conformer towards its most energy stable state, we compare our model with two state-of-the-art SE(3)-equivariant models: EGNN (Satorras et al., 2021) and SE(3)-
Revised: Baselines Since the goal for this experiment is to show the energy-derived model is indeed capable of optimizing an existing 3D conformer towards its most energy stable state, we compare our model with two state-of-the-art SE(3)-equivariant models: EGNN (Satorras et al., 2021) and SE(3)-

Original: Transformer (Fuchs et al., 2020) that can also take 3D Cartesian coordinates as input directly.
Revised: Transformer (Fuchs et al., 2020) that can also take 3D Cartesian coordinates as input directly.

Original: For non-machine-learning method, we also include one classical approach RDKit + MMFF (Riniker &
Revised: For non-machine-learning method, we also include one classical approach RDKit + MMFF (Riniker &

Original: Landrum, 2015), where the predict conformer is estimated through Euclidean Distance Geometry (EDG) and further optimized with Merck Molecular Force Field (MMFF) (Halgren, 1996).
Revised: Landrum, 2015), where the predict conformer is estimated through Euclidean Distance Geometry (EDG) and further optimized with Merck Molecular Force Field (MMFF) (Halgren, 1996).

Original: Setup Following Sec. 3.3, we include all three variants for this experiment: (1) Ours-TwoAtom , a basic two-atom model where we only consider pairwise atomic distances in the energy function.
Revised: Setup Following Sec. 3.3, we include all three variants for this experiment: (1) Ours-TwoAtom , a basic two-atom model where we only consider pairwise atomic distances in the energy function.

Original: (2) Ours-Ext v , which extends the two-atom model by considering the geometry of the internal representation v in the energy function.
Revised: (2) Ours-Ext v , which extends the two-atom model by considering the geometry of the internal representation v in the energy function.

Original: (3) Ours-ThreeAtom , which further extends the two variants above by considering the atom triplets in the energy function.
Revised: (3) Ours-ThreeAtom , which further extends the two variants above by considering the atom triplets in the energy function.

Original: To train our models and other machine learning baselines, we optimize model parameters by minimizing the L 2 loss between pairwise distance matrices of the generated conformer and its closest reference conformer (see
Revised: To train our models and other machine learning baselines, we optimize model parameters by minimizing the L 2 loss between pairwise distance matrices of the generated conformer and its closest reference conformer (see

Original: Appendix E for more model details).
Revised: Appendix E for more model details).

Original: Taking advantage of the SE(3)-equivariant property of the baseline and our models, we train  the model to take the RDKit + MMFF predicted conformers as input and allow the model to start with a reasonable estimation at a low computation cost. 
Revised: Taking advantage of the SE(3)-equivariant property of the baseline and our models, we train all the models to take the RDKit + MMFF predicted conformers as input and allow the models to start with a reasonable estimation at a low computation cost. We perform same number of optimization step (=9) for all models and both datasets.

Original: The results for random initialization can also be found in Appendix H.1.
Revised: The results for random initialization can also be found in Appendix H.1, and the results for effects of the number of optimization steps can be found in Appendix K.

Original: Results Following Mansimov et al.
Revised: Results Following Mansimov et al.

Original: (2019), we measure the difference between a predicted conformer and its reference by the Root-Mean-Square Deviation (RMSD) for all the heavy atoms.
Revised: (2019), we measure the difference between a predicted conformer and its reference by the Root-Mean-Square Deviation (RMSD) for all the heavy atoms.

Original: As shown in Table 1, all the optimization models can improve upon the initial conformation estimated by RDKit+MMFF, and new variants derived from the neural energy minimization formulation also outperform the two state-of-the-art SE(3)-equivariant models with a clear margin.
Revised: As shown in Table 1, all the optimization models can improve upon the initial conformation estimated by RDKit+MMFF, and new variants derived from the neural energy minimization formulation also outperform the two state-of-the-art SE(3)-equivariant models with a clear margin.

Original: More interestingly, model variants derived with more complex and higher-order energy deﬁnitions also outperform thebasic model that considers inter-atomic distance.
Revised: More interestingly, model variants derived with more complex and higher-order energy deﬁnitions also outperform the

Original: In the GEOM-Drugs dataset, where we have large molecules with many rotatable bonds, we hypothesize the higher-order model can outperform the

Original: [Table]
Revised: [Table]

Original: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: basic  two-atom model because the corresponding higher-order energy function can better approximate the actual potential energy surface by considering the bond-angle explicitly.
Revised: basic model that considers inter-atomic distance. In the GEOM-Drugs dataset, where we have large molecules with many rotatable bonds, we hypothesize the higher-order model can outperform the basic two-atom model because the corresponding higher-order energy function can better approximate the actual potential energy surface by considering the bond-angle explicitly.

Original: Since we only consider three variants in this work, and they all have shown great performance, we look forward to other more powerful variants the community can develop with such perspective.
Revised: Since we only consider three variants in this work and they all have shown great performance, we look forward to other more powerful variants the community can develop with such perspective.

Original: In Figure 2, examples of model-predicted conformers are visualized along with their initial conformers from RDKit+MMFF and the reference conformers in their optimal energy state.
Revised: In Figure 2, examples of model-predicted conformers are visualized along with their initial conformers from RDKit+MMFF and the reference conformers in their optimal energy state.

Original: The visualization shows that our models can produce conformers that are not only realistic but also energetically stable.
Revised: The visualization shows that our models can produce conformers that are not only realistic but also energetically stable.

Original: With no intentional design, our model-predicted conformers can also capture the structure of the basic components quite well (e.g., benzene forms a ring in the same plane).
Revised: With no intentional design, our model-predicted conformers can also capture the structure of the basic components quite well (e.g., benzene forms a ring in the same plane).

Original: For larger molecules where the initial conformers (RDKit+MMFF) are not folded/rotated in the most stable fashion, our model can also optimize the conformers towards the stable state as reference conformers.
Revised: For larger molecules where the initial conformers (RDKit+MMFF) are not folded/rotated in the most stable fashion, our model can also optimize the conformers towards the stable state as reference conformers.

Original: To measure the quality of the optimized conformers beyond the RMSD metric, we evaluate two downstream applications with the QM9 test set conformers predicted by different methods.
Revised: To measure the quality of the optimized conformers beyond the RMSD metric, we evaluate two downstream applications with the QM9 test set conformers predicted by different methods.

Original: In the ﬁrststudy, we compare the HOMO and LUMO energy (Smith et al., 2020) of a molecule when calculated with the reference conformers vs. the model-predicted conformers. As shown in Figure 3, we ﬁnd the
Revised: In the ﬁrst

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: study, we compare the HOMO and LUMO energy (Smith et al., 2020) of a molecule when calculated with the reference conformers vs. the model-predicted conformers.

Original: [Table]
Revised: As shown in Figure 3, we ﬁnd the

Original: HOMO-LUMO gap calculated from our model-predicted conformer match well with the one from reference, while the correlation between the reference and force ﬁeld optimized conformers is much lower.
Revised: HOMO-LUMO gap calculated from our model-predicted conformer match well with the one from reference, while the correlation between the reference and force ﬁeld optimized conformers is much lower.

Original: In the second study, we train a standard neural network model, SchNet (Schütt et al., 2017) for the QM9 benchmark (Ramakrishnan et al., 2014) with the reference or predicted conformers.
Revised: In the second study, we train a standard neural network model, SchNet (Schütt et al., 2017) for the QM9 benchmark (Ramakrishnan et al., 2014) with the reference or predicted conformers.

Original: As shown in Table 2, the model trained with our predicted conformers can outperform the ones trained with other predicted conformers.
Revised: As shown in Table 2, the model trained with our predicted conformers can outperform the ones trained with other predicted conformers.

Original: The result indicates better inductive bias in our predicted 3D structures compared to the ones estimated using other algorithms.
Revised: The result indicates better inductive bias in our predicted 3D structures compared to the ones estimated using other algorithms.

Original: Even though there is still a performance gap when comparing results with the model trained with the reference conformers, we believe the model-predicted conformer could still be helpful for datasets where no references are not available.
Revised: Even though there is still a performance gap when comparing results with the model trained with the reference conformers, we believe the model-predicted conformer could still be helpful for datasets where no references are not available.

Original: More experiment details can be found in Appendix F.
Revised: More experiment details can be found in Appendix F.

Original: 4. M OLECULAR C ONFORMER G ENERATION
Revised: 4. M OLECULAR C ONFORMER G ENERATION

Original: Data We test molecular conformer generation on both GEOM-QM9 and GEOM-Drugs datasets (Axelrod & Gomez-Bombarelli, 2020) following the same data split as Shi et al.
Revised: Data We test molecular conformer generation on both GEOM-QM9 and GEOM-Drugs datasets (Axelrod & Gomez-Bombarelli, 2020) following the same data split as Shi et al.

Original: (2021) including 40k training molecules and 200 test molecules for both datasets (see Appendix D for details).
Revised: (2021) including 40k training molecules and 200 test molecules for both datasets (see Appendix D for details).

Original: Baselines We compare our model against an EDG-based approach RDKit (Riniker & Landrum, 2015) as well as four deep generative models speciﬁcally studied for the generation setting: CVGAE (Mansimov et al., 2019) is a conditional VAE that directly generates atom coordinates based on molecular graphs.
Revised: Baselines We compare our model against an EDG-based approach RDKit (Riniker & Landrum, 2015) as well as four deep generative models speciﬁcally studied for the generation setting: CVGAE (Mansimov et al., 2019) is a conditional VAE that directly generates atom coordinates based on molecular graphs.

Original: GraphDG (Simm & Hernandez-Lobato, 2020) and CGCF (Xu et al., 2020) are
Revised: GraphDG (Simm & Hernandez-Lobato, 2020) and CGCF (Xu et al., 2020) are

Original: VAE- and ﬂow-based method that only generate the pairwise distances for 3D conformer conversion.
Revised: VAE- and ﬂow-based method that only generate the pairwise distances for 3D conformer conversion.

Original: ConfGF (Shi et al., 2021) is a recently-published state-of-the-art model that learns the gradient ﬁeld for molecular conformations and perform generation through iterative Langevin sampling.
Revised: ConfGF (Shi et al., 2021) is a recently-published state-of-the-art model that learns the gradient ﬁeld for molecular conformations and perform generation through iterative Langevin sampling.

Original: Setup In order to test our optimization formulation for the generative setting, we adopt the same two-atom model from the optimization setting here with the same hyper-parameters .
Revised: Setup In order to test our optimization formulation for the generative setting, we adopt the same two-atom model from the optimization setting here with the same hyper-parameters .

Original: However, instead of training the model to optimize a single initialization towards one reference, we train the model to optimize K initialization such that these K optimized conformers can be exactly matched to K different sampled references.
Revised: However, instead of training the model to optimize a single initialization towards one reference, we train the model to optimize K initialization such that these K optimized conformers can be exactly matched to K different sampled references.

Original: More concretely, we optimize the model parameters against an optimal-transport-like loss similar to other generative models (Genevay et al., 2018):
Revised: More concretely, we optimize the model parameters against an optimal-transport-like loss similar to other generative models (Genevay et al., 2018):

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: The

Original: [Table]
Revised: [Table]

Original: where C ( X ∗ i , X j ) is the same L 2 loss between the distance matrices of the reference ( X ∗ i ) and generated ( X j ) conformer, while the optimal transport plan π ∈ { 0 , 1 } K × K is realized as a minimum weight bipartite graph matching and Γ is the set of all permutations.
Revised: where C ( X ∗ i , X j ) is the same L 2 loss between the distance matrices of the reference ( X ∗ i ) and generated ( X j ) conformer, while the optimal transport plan π ∈ { 0 , 1 } K × K is realized as a minimum weight bipartite graph matching and Γ is the set of all permutations.

Original: For the initial conformer, we again take advantage of the SE(3)-equivariance of our method and evaluate the generation performance with initialization from either a random set of coordinates (i.e., Ours-Random model) or a rough
Revised: For the initial conformer, we again take advantage of the SE(3)-equivariance of our method and evaluate the generation performance with initialization from either a random set of coordinates (i.e., Ours-Random model) or a rough

Original: RDKit estimate (i.e., Ours-RDKit model, whose training inputs are perturbed references).
Revised: RDKit estimate (i.e., Ours-RDKit model, whose training inputs are perturbed references).

Original: Results To evaluate the diversity, precision, and structural quality of generated conformers, we adopt the corresponding COV erage, MIS matching, and MAT ching scores following Xu et al.
Revised: Results To evaluate the diversity, precision, and structural quality of generated conformers, we adopt the corresponding COV erage, MIS matching, and MAT ching scores following Xu et al.

Original: (2020) and Shi et al.
Revised: (2020) and Shi et al.

Original:  As shown in Table.
Revised: (2021) as deﬁned in Appendix G. As shown in Table.

Original: 3, our models achieve competitive performance in all three metrics, especially for the more challenging GEOM-Drugs dataset where we outperform existing machine learning methods by a clear margin.
Revised: 3, our models achieve competitive performance in all three metrics, especially for the more challenging GEOM-Drugs dataset where we outperform existing machine learning methods by a clear margin.

Original: These results suggest that our optimization formulation can be easily extended to the generation setting for a diverse and accurate ensemble of conformers.
Revised: These results suggest that our optimization formulation can be easily extended to the generation setting for a diverse and accurate ensemble of conformers.

Revised: Ablation study in Appendix J also suggests the improvement is mainly contributed by the SE(3)-equivariant setup and the optimal transport loss.

Original: Since no Langevin sampling is required in our method, the inference time (<1s per molecule) for our method is also much faster than ConfGF ( ∼ 170s per molecule).
Revised: Since no Langevin sampling is required in our method, the inference time (<1s per molecule) for our method is also much faster than ConfGF ( ∼ 170s per molecule).

Original: We also noted the performance difference between random initialization and RDKit initialization also represents a classical trade-off between diversity and accuracy, where a RDKit initialization can achieve better accuracy (MAT / MIS) by providing a more accurate starting point, while a random initialization can achieve better coverage (COV) by sufﬁcient sampling of the space.
Revised: We also noted the performance difference between random initialization and RDKit initialization also represents a classical trade-off between diversity and accuracy, where a RDKit initialization can achieve better accuracy (MAT / MIS) by providing a more accurate starting point, while a random initialization can achieve better coverage (COV) by sufﬁcient sampling of the space.

Original: We left more interesting initialization (e.g. mixture of the two) for future work, and please consult Appendix G for details of the implementations.
Revised: We left more interesting initialization (e.g. mixture of the two) for future work, and please consult Appendix G for details of the implementations.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: This paper proposes a new formulation for molecular conformer prediction by parametrizing an
Revised: This paper proposes a new formulation for molecular conformer prediction by parametrizing an

Original: SE(3)-equivariant network to model the gradient ﬁeld of the conformational energy landscape.
Revised: SE(3)-equivariant network to model the gradient ﬁeld of the conformational energy landscape.

Original: From the formulation of neural energy minimization, we also found a new perspective to principally derive new variants of SE(3)-equivariant models that align with an explicit set of underlying assumptions for the physical system.
Revised: From the formulation of neural energy minimization, we also found a new perspective to principally derive new variants of SE(3)-equivariant models that align with an explicit set of underlying assumptions for the physical system.

Original: Through an extensive set of experiments, we show that the proposed method is capable of optimizing a given conformer towards its most energetic-stable state and generating an ensemble of diverse and accurate conformers efﬁciently.
Revised: Through an extensive set of experiments, we show that the proposed method is capable of optimizing a given conformer towards its most energetic-stable state and generating an ensemble of diverse and accurate conformers efﬁciently.

Original: We believe the perspective of neural energy minimization also represents a new direction for thinking about SE(3)-equivariant model and expect the community to derive more interesting variants with energy assumptions that are more suitable for their applications.
Revised: We believe the perspective of neural energy minimization also represents a new direction for thinking about SE(3)-equivariant model and expect the community to derive more interesting variants with energy assumptions that are more suitable for their applications.

Original: For future work, our interest also goes beyond conformations of small molecules, as we can further scale up the system to model larger bio-molecular systems such as protein-ligand complexes.
Revised: For future work, our interest also goes beyond conformations of small molecules, as we can further scale up the system to model larger bio-molecular systems such as protein-ligand complexes.

Original: Since our training objective is blind to the chirality, our system will need to infer the chirality post model generation, and it remains an interesting question that whether we can build the system with chirality built-in (Pattanaik et al., 2020).
Revised: Since our training objective is blind to the chirality, our system will need to infer the chirality post model generation, and it remains an interesting question that whether we can build the system with chirality built-in (Pattanaik et al., 2020).

Original: Last but not least, with the ability to quickly generate high-quality ensembles of molecular conformers, we are also interested in molecular machine learning with conformer ensembles (Axelrod & Gómez-Bombarelli, 2020) especially for datasets where conformers are unavailable yet.
Revised: Last but not least, with the ability to quickly generate high-quality ensembles of molecular conformers, we are also interested in molecular machine learning with conformer ensembles (Axelrod & Gómez-Bombarelli, 2020) especially for datasets where conformers are unavailable yet.

Original: A P ROOF OF SE(3)- EQUIVARIANCE A. P ROOF OF SE(3)- EQUIVARIANCE FOR THE TWO - ATOM MODEL
Revised: A P ROOF OF SE(3)- EQUIVARIANCE A. P ROOF OF SE(3)- EQUIVARIANCE FOR THE TWO - ATOM MODEL

Original: One of the crucial challenges for modeling molecular conformers is to capture the 3D roto-translation symmetry, as rotating or moving the conformer in the 3D space on its own would not change the conformational energy.
Revised: One of the crucial challenges for modeling molecular conformers is to capture the 3D roto-translation symmetry, as rotating or moving the conformer in the 3D space on its own would not change the conformational energy.

Original: Therefore, to capture such symmetry for better generalization (Köhler et al., 2019), our model is required to change its output accordingly if the input conformer X undergoes a transformation of rotation or/and translation in the 3D space.
Revised: Therefore, to capture such symmetry for better generalization (Köhler et al., 2019), our model is required to change its output accordingly if the input conformer X undergoes a transformation of rotation or/and translation in the 3D space.

Original: More concretely, the model needs to satisfy SE(3)-equivariance:
Revised: More concretely, the model needs to satisfy SE(3)-equivariance:

Original: [Equation]
Revised: [Equation]

Original: where T : X → ˆ X is a transformation on X for the SE(3) group.
Revised: where T : X → ˆ X is a transformation on X for the SE(3) group.

Original: Speciﬁcally, for each 3D coordinate x , we have T ( x ) =
Revised: Speciﬁcally, for each 3D coordinate x , we have T ( x ) =

Original: Ax + b where A ∈ R 3 × 3 is the rotation matrix and b ∈ R 3 is the translation vector.
Revised: Ax + b where A ∈ R 3 × 3 is the rotation matrix and b ∈ R 3 is the translation vector.

Original: In the rest of the section, we simply denote ˆ x as T ( x ) .
Revised: In the rest of the section, we simply denote ˆ x as T ( x ) .

Original: Now that we have deﬁned SE(3)-equivariance, we want to prove that the following coordinate update in Eq.
Revised: Now that we have deﬁned SE(3)-equivariance, we want to prove that the following coordinate update in Eq.

Original: equation 4 satisﬁes SE(3)-equivariant:
Revised: equation 4 satisﬁes SE(3)-equivariant:

Original: [Equation]
Revised: [Equation]

Original: and we will do so by ﬁrst showing that f x ( · ) is SE(3)-invariant.
Revised: and we will do so by ﬁrst showing that f x ( · ) is SE(3)-invariant.

Original: To show f x ( · ) is SE(3)-invariant, it is sufﬁcient to show that its various inputs d ij , v i , and e ij are SE(3)-invariant.
Revised: To show f x ( · ) is SE(3)-invariant, it is sufﬁcient to show that its various inputs d ij , v i , and e ij are SE(3)-invariant.

Original: Intuitively, d ij should not change with the 3D roto-translation T , but formally we have:
Revised: Intuitively, d ij should not change with the 3D roto-translation T , but formally we have:

Original: [Equation]
Revised: [Equation]

Original: Since the initialization and update of v in Equation 4 only depend on variables satisfying SE(3)- invariant, v is also SE(3)-invariant.
Revised: Since the initialization and update of v in Equation 4 only depend on variables satisfying SE(3)- invariant, v is also SE(3)-invariant.

Original: Similarly, the relationship embedding e is also SE(3)-invariant.
Revised: Similarly, the relationship embedding e is also SE(3)-invariant.

Original: Therefore, function f x satisﬁes SE(3)-invariant.
Revised: Therefore, function f x satisﬁes SE(3)-invariant.

Original: Now, if we look at the update for φ ( T ( x t )) or φ (ˆ x t ) , we have:
Revised: Now, if we look at the update for φ ( T ( x t )) or φ (ˆ x t ) , we have:

Original: [Equation]
Revised: [Equation]

Original: Thus, the two-atom model update φ is SE(3)-equivariant.
Revised: Thus, the two-atom model update φ is SE(3)-equivariant.

Original: A. P ROOF OF SE(3)- EQUIVARIANT FOR THE GENERAL CASE
Revised: A. P ROOF OF SE(3)- EQUIVARIANT FOR THE GENERAL CASE

Original: For the more general case, we have the following equivariance theorem:
Revised: For the more general case, we have the following equivariance theorem:

Original: Theorem 1.
Revised: Theorem 1.

Original: Let G be an SE(3)-group which acts on R n × 3 .
Revised: Let G be an SE(3)-group which acts on R n × 3 .

Original: If the potential function Φ : R n × 3 → R is a G -invariant function (i.e. the assumed energy function E shown in Equation 1 is G -invariant), then the gradient vector ﬁeld ∇ x Φ will be G -equivariant.
Revised: If the potential function Φ : R n × 3 → R is a G -invariant function (i.e. the assumed energy function E shown in Equation 1 is G -invariant), then the gradient vector ﬁeld ∇ x Φ will be G -equivariant.

Original: Proof.
Revised: Proof.

Original: For a rotation transformation ∀ g ∈ G with a rotation matrix R g , we have the following if Φ is SE(3)-invariant: ∇ x ( R g ◦ Φ) =
Revised: For a rotation transformation ∀ g ∈ G with a rotation matrix R g , we have the following if Φ is SE(3)-invariant: ∇ x ( R g ◦ Φ) =

Original: ∇ x Φ .
Revised: ∇ x Φ .

Original: Now, according to the chain rule, we also have the following: ∇ x ( R g ◦ Φ) = ( ∇ x R g x ) ◦ ( ∇ R g x Φ)
Revised: Now, according to the chain rule, we also have the following: ∇ x ( R g ◦ Φ) = ( ∇ x R g x ) ◦ ( ∇ R g x Φ)

Original: =
Revised: =

Original: R Tg ◦ ( ∇ R g x Φ) .
Revised: R Tg ◦ ( ∇ R g x Φ) .

Original: Now putting the two equations together, and multiply the orthogonal matrix R g on both sides, we have: R g ◦ ∇ x Φ( x ) =
Revised: Now putting the two equations together, and multiply the orthogonal matrix R g on both sides, we have: R g ◦ ∇ x Φ( x ) =

Original: R g ◦ R
Revised: R g ◦ R

Original: Tg ◦∇ R g x Φ( R g x )
Revised: Tg ◦∇ R g x Φ( R g x )

Original: = ∇ R g x
Revised: = ∇ R g x

Original: Φ( R g x ) .
Revised: Φ( R g x ) .

Original: Similarly, for a transnational transformation g ( x ) = x + t , one could show that ∇ x + t Φ( x + t ) =
Revised: Similarly, for a transnational transformation g ( x ) = x + t , one could show that ∇ x + t Φ( x + t ) =

Original: ∇ x Φ( x ) if Φ is SE(3)-invariant.
Revised: ∇ x Φ( x ) if Φ is SE(3)-invariant.

Original: The combination of these achieve the SE(3)-equivariant of the update formula:
Revised: The combination of these achieve the SE(3)-equivariant of the update formula:

Original: [Equation]
Revised: [Equation]

Original: i.e. f ◦ T g =
Revised: i.e. f ◦ T g =

Original: T g ◦ f , where f is our neural network and T g is the SE(3) transformation.
Revised: T g ◦ f , where f is our neural network and T g is the SE(3) transformation.

Original: This shows that if the constructed energy function E is SE(3)-invariant as it should be, it is sufﬁcient to guarantee that the gradient ﬁeld ∇ E (and estimated neural network) are SE(3)-equivariant.
Revised: This shows that if the constructed energy function E is SE(3)-invariant as it should be, it is sufﬁcient to guarantee that the gradient ﬁeld ∇ E (and estimated neural network) are SE(3)-equivariant.

Original: The theorem can be generalized to any group G utilizing the Riemannian geometry.
Revised: The theorem can be generalized to any group G utilizing the Riemannian geometry.

Original: We refer Wasserman (1969) and Katsman et al.
Revised: We refer Wasserman (1969) and Katsman et al.

Original: (2021) to readers of interest.
Revised: (2021) to readers of interest.

Original: B D ERIVATION OF THE T HREE -A TOM M ODEL
Revised: B D ERIVATION OF THE T HREE -A TOM M ODEL

Original: For the three-atom model, we have the following energy formulation:
Revised: For the three-atom model, we have the following energy formulation:

Original: [Equation]
Revised: [Equation]

Original: and following the energy function, we can derive the following gradient update w.r
Revised: and following the energy function, we can derive the following gradient update w.r

Original: [Equation]
Revised: [Equation]

Original: t. x :
Revised: t. x :

Original: [Equation]
Revised: [Equation]

Original: ∈V
Revised: ∈V

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: x j
Revised: x j

Original: [Equation]
Revised: [Equation]

Original: g x
Revised: g x

Original: [Equation]
Revised: [Equation]

Original: · )] and therefore, we have the following x updates for the three-atom models:
Revised: · )] and therefore, we have the following x updates for the three-atom models:

Original: [Equation]
Revised: [Equation]

Original: C R EINTERPRETATIONS OF EXISTING MODELS C. E QUIVARIANT B OLTZMANN G ENERATOR
Revised: C R EINTERPRETATIONS OF EXISTING MODELS C. E QUIVARIANT B OLTZMANN G ENERATOR

Original: Equivariant Boltzmann Generator (Köhler et al., 2019) introduces an E(n)-equivariant update for x i in the following form:
Revised: Equivariant Boltzmann Generator (Köhler et al., 2019) introduces an E(n)-equivariant update for x i in the following form:

Original: [Equation]
Revised: [Equation]

Original: where φ is a Multilayer Perceptron (MLP).
Revised: where φ is a Multilayer Perceptron (MLP).

Original: It explicitly optimizes the energy function E
Revised: It explicitly optimizes the energy function E

Original: =
Revised: =

Original: (cid:80) i,j ∈V , which can be interpreted as a simpliﬁed case of the energy function we proposed in Equation 2 since it only concerns about the distance between two nodes and disregards the information associated with the nodes and edges.
Revised: (cid:80) i,j ∈V , which can be interpreted as a simpliﬁed case of the energy function we proposed in Equation 2 since it only concerns about the distance between two nodes and disregards the information associated with the nodes and edges.

Original: C. TFN AND SE(
Revised: C. TFN AND SE(

Original: -T RANSFORMER
Revised: -T RANSFORMER

Original: TFN (Thomas et al., 2018) and SE(3)-Transformer (Fuchs et al., 2020), the attention-based extension of TFN, are a family of SE(3)-equivariant model that performs the update for a representation ˆ x i with a learned kernel
Revised: TFN (Thomas et al., 2018) and SE(3)-Transformer (Fuchs et al., 2020), the attention-based extension of TFN, are a family of SE(3)-equivariant model that performs the update for a representation ˆ x i with a learned kernel

Original: φ utilizing spherical harmonics such that the model is equivariant with respect to x i :
Revised: φ utilizing spherical harmonics such that the model is equivariant with respect to x i :

Original: [Equation]
Revised: [Equation]

Original: where φ u is a simple projection matrix, and φ x j − x i is a kernel matrix that learns to project ˆ x j based on ( x j − x i ) .
Revised: where φ u is a simple projection matrix, and φ x j − x i is a kernel matrix that learns to project ˆ x j based on ( x j − x i ) .

Original: Since kernel φ x j − x i consists of a learned radial function for the distance d ij and an angular preserve basis for x j − x i d ij , this formulation is in close relationship with Equationas the angular preserve basis also captures the term ( x i − x j ) while the learnable radial function is equivalent to the gradient estimator f v .
Revised: Since kernel φ x j − x i consists of a learned radial function for the distance d ij and an angular preserve basis for x j − x i d ij , this formulation is in close relationship with Equationas the angular preserve basis also captures the term ( x i − x j ) while the learnable radial function is equivalent to the gradient estimator f v .

Original: Therefore, such operation can also be interpreted as optimizing the same energy function of our two-atom model, but with higher-order equivariant representations involving special designed basis.
Revised: Therefore, such operation can also be interpreted as optimizing the same energy function of our two-atom model, but with higher-order equivariant representations involving special designed basis.

Original: Last but not least, since the optimized x can only be obtained at the last step through a linear projection, the whole model is effectively performance only an one-step optimization.
Revised: Last but not least, since the optimized x can only be obtained at the last step through a linear projection, the whole model is effectively performance only an one-step optimization.

Original: D D ATA P REPROCESSING D ETAILS D. R EFERENCE S ELECTION
Revised: D D ATA P REPROCESSING D ETAILS D. R EFERENCE S ELECTION

Original: There are three datasets we used in the experiment section: QM9, GEOM-QM9, and GEOM-Drugs.
Revised: There are three datasets we used in the experiment section: QM9, GEOM-QM9, and GEOM-Drugs.

Original: For the molecular conformation optimization task, we choose to use QM9 instead of GEOM-QM because the conformers of it are computed with high ﬁdelity DFT method.
Revised: For the molecular conformation optimization task, we choose to use QM9 instead of GEOM-QM because the conformers of it are computed with high ﬁdelity DFT method.

Original: For the GEOM-Drugs dataset computed with a semi-empirical method, they include multiple low-energy stable conformers.
Revised: For the GEOM-Drugs dataset computed with a semi-empirical method, they include multiple low-energy stable conformers.

Original: If we compare the Boltzmann weights (proxy of energy stability) for the most likely conformers and the second likely conformers, we ﬁnd many of them are pretty closed, as shown in Figure S1a).
Revised: If we compare the Boltzmann weights (proxy of energy stability) for the most likely conformers and the second likely conformers, we ﬁnd many of them are pretty close, as shown in Figure S1a).

Original: In
Revised: In

Original: Figure S1b), we also showed an example where two conformers with very similar Boltzmann weight but drastically different conformation.
Revised: Figure S1b), we also showed an example where two conformers with very similar Boltzmann weight but drastically different conformation.

Original: Therefore, for GEOM-Drugs in the optimization setting, we use multiple reference conformers during the training and evaluation for GEOM-Drugs, but the model predicted conformers only need to match one of the references.
Revised: Therefore, for GEOM-Drugs in the optimization setting, we use multiple reference conformers during the training and evaluation for GEOM-Drugs, but the model predicted conformers only need to match one of the references.

Original: In order to select similarly energetically preferred conformers, we select the set of target conformers as { X | P X ≥ 0 .
Revised: In order to select similarly energetically preferred conformers, we select the set of target conformers as { X | P X ≥ 0 .

Original: 5 ∗ P X ∗ } where P X ∗ is the Boltzmann weight for the most likely conformers.
Revised: 5 ∗ P X ∗ } where P X ∗ is the Boltzmann weight for the most likely conformers.

Original: For the molecular conformation generation task, we use both GEOM-QM9 and GEOM-Drugs as the dataset since they contain multiple references for each molecule.
Revised: For the molecular conformation generation task, we use both GEOM-QM9 and GEOM-Drugs as the dataset since they contain multiple references for each molecule.

Original: However, to ﬁlter out some of the rare conformers, we ﬁrst sort reference conformers by their Boltzmann weights and set a threshold of the sum Boltzmann weight at 0.95 in the training set: { X 1: n | (cid:80) ni
Revised: However, to ﬁlter out some of the rare conformers, we ﬁrst sort reference conformers by their Boltzmann weights and set a threshold of the sum Boltzmann weight at 0.95 in the training set: { X 1: n | (cid:80) ni

Original: =1 P X
Revised: =1 P X

Original: i ≥ 0 .
Revised: i ≥ 0 .

Original: 95 , (cid:80) n − 1 i =1 P X
Revised: 95 , (cid:80) n − 1 i =1 P X

Original: i < 0 .
Revised: i < 0 .

Original: 95 , P X 1 ≥ · · · ≥ P X N } , but still report metrics on all reference conformers as previous work does.
Revised: 95 , P X 1 ≥ · · · ≥ P X N } , but still report metrics on all reference conformers as previous work does.

Original: D. F EATURIZATION
Revised: D. F EATURIZATION

Original: We construct the molecular graph G as a fully-connected graph to represent the molecules where the nodes are featured with atom type, aromatic, hybridization, and the number of hydrogens, while the edges are featured with the bond type (including no-bond), length of the shortest path in the graph, and whether the atom pair at both ends are in the same conjugated system.
Revised: We construct the molecular graph G as a fully-connected graph to represent the molecules where the nodes are featured with atom type, aromatic, hybridization, and the number of hydrogens, while the edges are featured with the bond type (including no-bond), length of the shortest path in the graph, and whether the atom pair at both ends are in the same conjugated system.

Original: We perform an initial feature projection with v -update in each model to get the atom embedding and relationship embedding from these features.
Revised: We perform an initial feature projection with v -update in each model to get the atom embedding and relationship embedding from these features.

Original: Similar to other works (Mansimov et al., 2019; Xu et al., 2020), we only model the heavy atoms for the molecular graph and conformer in this experiment.
Revised: Similar to other works (Mansimov et al., 2019; Xu et al., 2020), we only model the heavy atoms for the molecular graph and conformer in this experiment.

Original: E E XPERIMENT D ETAILS FOR M OLECULAR C ONFORMATION O PTIMIZATION E. 
Revised: E E XPERIMENT D ETAILS FOR M OLECULAR C ONFORMATION O PTIMIZATION E.

Original: M ODEL P ARAMETRIZATION
Revised: M ODEL P ARAMETRIZATION

Original: SE(3)-Transformer (Fuchs et al., 2020) is parametrized with the following hyper-parameters which maximize the GPU memory: num_degrees=2 , num_layers=9 , hidden_dim=64 , and ﬁnally n_heads=8 .
Revised: SE(3)-Transformer (Fuchs et al., 2020) is parametrized with the following hyper-parameters which maximize the GPU memory: num_degrees=2 , num_layers=9 , hidden_dim=64 , and ﬁnally n_heads=8 .

Original: [Figure]
Revised: [Figure]

Original: EGNN (Satorras et al., 2021) is parametrized with the following hyper-parameters: num_layers=9 , hidden_dim=256 , and the distance square features are expanded using an
Revised: EGNN (Satorras et al., 2021) is parametrized with the following hyper-parameters: num_layers=9 , hidden_dim=256 , and the distance square features are expanded using an

Original: RBF kernel with 50 basis between 0Å - 10Å.
Revised: RBF kernel with 50 basis between 0Å - 10Å.

Original: Ours -TwoAtom and -Ext v contain three optimization blocks sharing the parameters, and within each optimization block, there are three layers of x and v updates.
Revised: Ours -TwoAtom and -Ext v contain three optimization blocks sharing the parameters, and within each optimization block, there are three layers of x and v updates.

Original: The neural network in these updates are featured as a Transformer (Vaswani et al., 2017) with hidden_dim=256 and n_heads=16 .
Revised: The neural network in these updates are featured as a Transformer (Vaswani et al., 2017) with hidden_dim=256 and n_heads=16 .

Original: The key/value embedding and attention score are generated through a 2-layer MLP with LayerNorm and ReLU.
Revised: The key/value embedding and attention score are generated through a 2-layer MLP with LayerNorm and ReLU.

Original: In addition, an additional interaction weight is introduced to model different energy contribution of the atom pairs which is simply modeled as α ij = sigmoid ( MLP ( d ij )) .
Revised: In addition, an additional interaction weight is introduced to model different energy contribution of the atom pairs which is simply modeled as α ij = sigmoid ( MLP ( d ij )) .

Original: Last but not least, we implement a dynamic cutoff where we ignore the energy contribution between atom pairs whose euclidean distance is less than 10Å. Similar to the EGNN implementation above, all distance features are also expanded with the same RBF kernel.
Revised: Last but not least, we implement a dynamic cutoff where we ignore the energy contribution between atom pairs whose euclidean distance is less than 10Å. Similar to the EGNN implementation above, all distance features are also expanded with the same RBF kernel.

Original: Figure S2 shows the overall architecture.
Revised: Figure S2 shows the overall architecture.

Original: [Figure]
Revised: [Figure]

Original: Ours-ThreeAtom is a three-atom model extension of Ours-Ext v , and therefore follow the same model hyper-parameters as above.
Revised: Ours-ThreeAtom is a three-atom model extension of Ours-Ext v , and therefore follow the same model hyper-parameters as above.

Original: Since we consider the bond angles ∠ jik and ∠ kij the same, the neural networks f x , g x and f v need to satisfy permutation invariant between ij and ik .
Revised: Since we consider the bond angles ∠ jik and ∠ kij the same, the neural networks f x , g x and f v need to satisfy permutation invariant between ij and ik .

Original: Therefore, for these three functions, we construct them as the following:
Revised: Therefore, for these three functions, we construct them as the following:

Original: [Equation]
Revised: [Equation]

Original: In addition, we generalize the three-atom model to the two-atom case by including the triplet ( i, j, k ) where j = k .
Revised: In addition, we generalize the three-atom model to the two-atom case by including the triplet ( i, j, k ) where j = k .

Original: While treating the molecule as a fully connected graph allows us to model the non-local relationship, the challenge is that the number of atom triplets grows very quickly with the size of the molecule.
Revised: While treating the molecule as a fully connected graph allows us to model the non-local relationship, the challenge is that the number of atom triplets grows very quickly with the size of the molecule.

Original: Since the bond angle in a conjugated ring system with no rotatable bond is not super interesting, we trim down the number of atom triplets by only considering ( i, j, k ) triplet where e ij or e ik represents a rotatable bond.
Revised: Since the bond angle in a conjugated ring system with no rotatable bond is not super interesting, we trim down the number of atom triplets by only considering ( i, j, k ) triplet where e ij or e ik represents a rotatable bond.

Original: E. T RAINING P ROCEDURES
Revised: E. T RAINING P ROCEDURES

Original: Our proposed models and other machine learning baselines are all trained in the same procedure.
Revised: Our proposed models and other machine learning baselines are all trained in the same procedure.

Original: We sample a batch size of 128 molecules for the QM9 dataset and a batch size of 32 molecules for the GEOM-Drugs dataset during the training.
Revised: We sample a batch size of 128 molecules for the QM9 dataset and a batch size of 32 molecules for the GEOM-Drugs dataset during the training.

Original: In one training batch, one initial conformer is randomly sampled for each molecule from a pool of ten pre-initialized conformers from RDKit+MMFF.
Revised: In one training batch, one initial conformer is randomly sampled for each molecule from a pool of ten pre-initialized conformers from RDKit+MMFF.

Original: Since we expect the model to stop the optimization when the input conformer is already stable, we also sample reference conformer as the initial input conformer with a probability p = 0 .
Revised: Since we expect the model to stop the optimization when the input conformer is already stable, we also sample reference conformer as the initial input conformer with a probability p = 0 .

Original: The model is trained end-to-end by the L 2 loss between the distance matrices for the predicted conformer and the reference conformer via gradient descent method Adam Kingma & Ba (2015).
Revised: The model is trained end-to-end by the L 2 loss between the distance matrices for the predicted conformer and the reference conformer via gradient descent method Adam Kingma & Ba (2015).

Original: The hyper-parameter settings for the optimizer are init_learning_rate=0.001 , betas=(0.95, 0.999) , and clip_gradient_norm=8 .
Revised: The hyper-parameter settings for the optimizer are init_learning_rate=0.001 , betas=(0.95, 0.999) , and clip_gradient_norm=8 .

Original: We also schedule to decay the learning rate exponentially with a factor of 0.5 and a minimum learning rate of 0.00001.
Revised: We also schedule to decay the learning rate exponentially with a factor of 0.5 and a minimum learning rate of 0.00001.

Original: The learning rate is decayed if there is no improvement for the validation loss in 8 consecutive evaluations, and the training will be completely stopped if no improvement is found for 20 consecutive evaluations.
Revised: The learning rate is decayed if there is no improvement for the validation loss in 8 consecutive evaluations, and the training will be completely stopped if no improvement is found for 20 consecutive evaluations.

Original: The evaluation is done for every 2000 training steps/batches.
Revised: The evaluation is done for every 2000 training steps/batches.

Revised: We tried a small number of learning rate schedules based on the validation set to make sure the training is sufﬁcient.

Revised: Gradient norm is clipped to make the training stable.

Revised: Other hyper-parameters are set following the common choices.

Original: E. C OMPUTATION R ESOURCE AND T RAINING T IME
Revised: E. C OMPUTATION R ESOURCE AND T RAINING T IME

Original: We train all conformer optimization models for the QM9 dataset with one NVIDIA GeForce GTX 1080 GPU and use one GeForce GTX 3090 GPU for the GEOM-Drugs dataset.
Revised: We train all conformer optimization models for the QM9 dataset with one NVIDIA GeForce GTX 1080 GPU and use one GeForce GTX 3090 GPU for the GEOM-Drugs dataset.

Original: All the models can be trained within 36-48 hours except the Ours-ThreeAtom model, which takes a longer time.
Revised: All the models can be trained within 36-48 hours except the Ours-ThreeAtom model, which takes a longer time.

Original: E. C ODE AND D ATA A VAILABILITY
Revised: E. C ODE AND D ATA A VAILABILITY

Original: We train all conformer optimization models for the QM9 dataset with one NVIDIA GeForce GTX 1080 GPU and one GeForce GTX 3090 GPU for the GEOM-Drugs dataset.
Revised: We train all conformer optimization models for the QM9 dataset with one NVIDIA GeForce GTX 1080 GPU and one GeForce GTX 3090 GPU for the GEOM-Drugs dataset.

Original: All the models can be trained within 36-48 hours except the Ours-ThreeAtom model, which takes longer.
Revised: All the models can be trained within 36-48 hours except the Ours-ThreeAtom model, which takes longer.

Original: F E XPERIMENT D ETAILS FOR D OWNSTREAM A PPLICATIONS F. C ORRELATION FOR THE HOMO-LUMO G AP
Revised: F E XPERIMENT D ETAILS FOR D OWNSTREAM A PPLICATIONS F. C ORRELATION FOR THE HOMO-LUMO G AP

Original: For this study, we sample 400 molecules from the test dataset of the conformer optimization experiment (Sec. 4.1) for the quantum mechanic calculation.
Revised: For this study, we sample 400 molecules from the test dataset of the conformer optimization experiment (Sec. 4.1) for the quantum mechanic calculation.

Original: Since the dataset and the model is heavy atom only, we infer the coordinate of hydrogens via RDKit (Riniker & Landrum, 2015).
Revised: Since the dataset and the model is heavy atom only, we infer the coordinate of hydrogens via RDKit (Riniker & Landrum, 2015).

Original: The HOMO and
Revised: The HOMO and

Original: LUMO energy for the reference and two-atom model predicted conformers are ﬁnally calculated via
Revised: LUMO energy for the reference and two-atom model predicted conformers are ﬁnally calculated via

Original: Psi4 (Smith et al., 2020) with the MP2/6-311++G(d,p) basis set.
Revised: Psi4 (Smith et al., 2020) with the MP2/6-311++G(d,p) basis set

Original:  F. T RAINING N EURAL N ETWORKS WITH P REDICTED C ONFORMERS
Revised: . F. T RAINING N EURAL N ETWORKS WITH P REDICTED C ONFORMERS

Original: In this study, we train the standard SchNet (Schütt et al., 2017) on the common QM9 property prediction benchmark with the data and splits provided by Anderson et al.
Revised: In this study, we train the standard SchNet (Schütt et al., 2017) on the common QM9 property prediction benchmark with the data and splits provided by Anderson et al.

Original: However, we want to train and evaluate the model with the generated conformers instead of using the reference conformers.
Revised: However, we want to train and evaluate the model with the generated conformers instead of using the reference conformers.

Original: Since SchNet (Schütt et al., 2017) requires the hydrogen atom for its full potential and the benchmark dataset follows a different split as our conformer optimization task, we train a new two-atom model on the training split of this benchmark with the hydrogen atom included and use this new model to generate conformers.
Revised: Since SchNet (Schütt et al., 2017) requires the hydrogen atom for its full potential and the benchmark dataset follows a different split as our conformer optimization task, we train a new two-atom model on the training split of this benchmark with the hydrogen atom included and use this new model to generate conformers.

Original: In our study, SchNet contains 7 layers and a hidden dimension of 128.
Revised: In our study, SchNet contains 7 layers and a hidden dimension of 128.

Original: To expand the distance feature, we use 50 Gaussian basis, and we construct the molecular graph with a distance cutoff of 10 Å. We use
Revised: To expand the distance feature, we use 50 Gaussian basis, and we construct the molecular graph with a distance cutoff of 10 Å. We use

Original: L 1 Loss to train the model end to end with Adam (Kingma & Ba, 2015) ( weight_decay=5e-16 , batch_size=96 , and num_epochs=1000 ).
Revised: L 1 Loss to train the model end to end with Adam (Kingma & Ba, 2015) ( weight_decay=5e-16 , batch_size=96 , and num_epochs=1000 ).

Original: The learning rate starts between 5e-4 and 1efor different tasks and follows a cosine decay.
Revised: The learning rate starts between 5e-4 and 1efor different tasks and follows a cosine decay.

Original: Since the single reference conformer does not constrain us, we consider multiple conformers in this study.
Revised: Since the single reference conformer does not constrain us, we consider multiple conformers in this study.

Original: RDKit+MMFF ﬁrst estimates multiple initial conformers, and different predicted conformers are generated from these initial conformers with our two-atom model and the baseline EGNN model.
Revised: RDKit+MMFF ﬁrst estimates multiple initial conformers, and different predicted conformers are generated from these initial conformers with our two-atom model and the baseline EGNN model.

Original: To select for these initial conformers, we ﬁrst sample 10 conformers with RDKit+MMFF, and the top 5 conformers with the lowest energy are selected.
Revised: To select for these initial conformers, we ﬁrst sample 10 conformers with RDKit+MMFF, and the top 5 conformers with the lowest energy are selected.

Original: One of the ﬁve conformers is sampled randomly at each batch during training, but the average prediction from the ﬁve conformers is reported as the model estimate during inference.
Revised: One of the ﬁve conformers is sampled randomly at each batch during training, but the average prediction from the ﬁve conformers is reported as the model estimate during inference.

Original: G E XPERIMENT D ETAILS FOR M OLECULAR C ONFORMATION G ENERATION
Revised: G E XPERIMENT D ETAILS FOR M OLECULAR C ONFORMATION G ENERATION

Original: We follow the setting in Shi et al.
Revised: We follow the setting in Shi et al.

Original: and report baseline results from their paper.
Revised: and report baseline results from their paper.

Original: The metrics we used are deﬁned as follows:
Revised: The metrics we used are deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: where RMSD( X , X ∗ ) =
Revised: where RMSD( X , X [Equation] ) =

Original: min Φ ( 1 n (cid:80) ni =1 (cid:107) Φ( X i ) − X ∗ i (cid:107) 2 ) 12 .
Revised: [Equation] Φ [Equation] X i ) − X ∗ i (cid:107) 2 ) 12 .

Original: n is the number of heavy atoms and Φ is an alignment function which aligns two conformers with rotational and translational operations.
Revised: n is the number of heavy atoms and Φ is an alignment function which aligns two conformers with rotational and translational operations.

Revised: The thresholds of COV and MIS are δ = 0 .

Revised: 5 and δ = 1 . 25 for GEOM-QM9 and the larger GEOM-Drugs respectively, which are same to Shi et al.

Original: We used the same setup as the two-atom model in the conformation optimization setting with the same hyper-parameters for this experiment.
Revised: We used the same setup as the two-atom model in the conformation optimization setting with the same hyper-parameters for this experiment.

Original: Tor train the model, we use the optimal transport loss during the training as described in Eq.
Revised: Tor train the model, we use the optimal transport loss during the training as described in Eq.

Original: and choose K = 5 .
Revised: and choose K = 5 .

Original: In Table.
Revised: In Table.

Original: 3, the conformers in both training and evaluation phases of Ours-Random are initialized with random noise, which is drawn from a Gaussian N (0 , σ 2 ) and σ = 0 .
Revised: 3, the conformers in both training and evaluation phases of Ours-Random are initialized with random noise, which is drawn from a Gaussian N (0 , σ 2 ) and σ = 0 .

Original: 028 ∗ |V| .
Revised: 028 ∗ |V| .

Original: On the other hand, Ours-RDKit are initialized with permuted reference conformers during training, where the permutation is drawn from a Gaussian
Revised: On the other hand, Ours-RDKit are initialized with permuted reference conformers during training, where the permutation is drawn from a Gaussian

Original: N (0 , σ ) and σ = 0 . 5 for GEOM-QM9 and σ = 1 . 0 for GEOM-Drugs to account for the average volume difference.
Revised: N (0 , σ ) and σ = 0 . 5 for GEOM-QM9 and σ = 1 . 0 for GEOM-Drugs to account for the average volume difference.

Original: In evaluation, we apply the model with initialization from RDKit since ground truth reference is not available.
Revised: In evaluation, we apply the model with initialization from RDKit since ground truth reference is not available.

Original: H A DDITIONAL E XPERIMENTS H. C ONFORMER O PTIMIZATION WITH R ANDOM I NITIALIZATION
Revised: H A DDITIONAL E XPERIMENTS H. C ONFORMER O PTIMIZATION WITH R ANDOM I NITIALIZATION

Original: While our optimization experiments are done with an initial conformer estimated by RDKit+MMFF, the initial conformer can be generated elsewhere, and even from other generative models like CGCF
Revised: While our optimization experiments are done with an initial conformer estimated by RDKit+MMFF, the initial conformer can be generated elsewhere, and even from other generative models like CGCF

Original: Xu et al.
Revised: Xu et al.

Original: In the simplest case, we could also initiate the conformers from random points for the model training and evaluation where each coordinate value is drawn from a Gaussian N (0 , σ 2 ) and σ = 0 .
Revised: In the simplest case, we could also initiate the conformers from random points for the model training and evaluation where each coordinate value is drawn from a Gaussian N (0 , σ 2 ) and σ = 0 .

Original: 1 ∗ |V| .
Revised: 1 ∗ |V| .

Original: As shown in Table S1, while the model performance is less satisfactory with the randomly initialized conformers, the predicted conformers are still quite accurate and even more accurate than the ones optimized by the hand-designed force-ﬁeld method.
Revised: As shown in Table S1, while the model performance is less satisfactory with the randomly initialized conformers, the predicted conformers are still quite accurate and even more accurate than the ones optimized by the hand-designed force-ﬁeld method.

Original: This also shows the potential of our methods as the model prediction can be even more accurate if better initialization is possible. 
Revised: This also shows the potential of our methods as the model prediction can be even more accurate if better initialization is possible. [Table]

Original: H. R EPRESENTATION L EARNING
Revised: H. P REDICTING E NSEMBLE P ROPERTY

Revised: Here we perform the same ensemble property prediction as Table 3 of Shi et al.

Revised: [Table]

Revised: where they measure how well the generated conformers can capture the population statistics of electronic property of the reference set.

Revised: The authors are very kind and provide us with the code and themolecules used in their experiments.

Revised: We reproduce the results here for future references

Original:  While we have been focusing on optimizing molecular conformation, we can also apply our SE(3)- equivariant model for molecular representation learning, even when the native conformer is already given.
Revised: [Table] eV. H. R EPRESENTATION L EARNING While we have been focusing on optimizing molecular conformation, we can also apply our SE(3)- equivariant model for molecular representation learning, even when the native conformer is already given.

Original: In this experiment, we treat the conformer X as a constant and only update the internal atom representation V under the framework of neural energy minimization.
Revised: In this experiment, we treat the conformer X as a constant and only update the internal atom representation V under the framework of neural energy minimization.

Original: Following the same setup as
Revised: Following the same setup as

Original: [Table]

Original: [Table]

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

Original: (2021), we compare our model against other models on the QM9 property prediction benchmark.
Revised: (2021), we compare our model against other models on the QM9 property prediction benchmark.

Original: As shown in Table S2, our v -update-only model can achieve a comparable performance with state-of-the-art models as well.
Revised: As shown in Table S3, our v -update-only model can achieve a comparable performance with state-of-the-art models as well.

Original: The consistently good performance again demonstrates the generalizability of our framework where powerful SE(3)-equivariant and -invariant models can be derived with different deﬁnitions of the energy function. 
Revised: The consistently good performance again demonstrates the generalizability of our framework where powerful SE(3)-equivariant and -invariant models can be derived with different deﬁnitions of the energy function. [Table] [Table] et al.

Original: I S TUDY OF M ODEL A BLATION AND V ARIATION I. E FFECTS OF GRAPH CONNECTIVITY
Revised: [Table] [Table] (Fuchs et al., 2020) [Table]

Revised: Satorras et al., 2021)

Revised: [Table]

Revised: I S TUDY OF M ODEL A BLATION AND V ARIATION I. E FFECTS OF GRAPH CONNECTIVITY

Original: While our model is based on the attention mechanism on the whole graph, we can also deﬁne the graph connectivity to reduce the computational complexity.
Revised: While our model is based on the attention mechanism on the whole graph, we can also deﬁne the graph connectivity to reduce the computational complexity.

Original: Performing attention directly on the molecular graph may be inefﬁcient.
Revised: Performing attention directly on the molecular graph may be inefﬁcient.

Original: One possible solution is to perform attention on the augmented graph as (Xu et al., 2020; Simm & Hernandez-Lobato, 2020) where the molecular graph is augmented with virtual bonds connecting atoms that are 2 or 3 hops away in the original graph.
Revised: One possible solution is to perform attention on the augmented graph as (Xu et al., 2020; Simm & Hernandez-Lobato, 2020) where the molecular graph is augmented with virtual bonds connecting atoms that are 2 or 3 hops away in the original graph.

Original: However, such node distances are sometimes not informative enough since two atoms far apart in graph distance may be close in the 3D space.
Revised: However, such node distances are sometimes not informative enough since two atoms far apart in graph distance may be close in the 3D space.

Original: An alternative way is to construct a fully connected graph to capture all long-range dependencies, but it could introduce noise to the attention process.
Revised: An alternative way is to construct a fully connected graph to capture all long-range dependencies, but it could introduce noise to the attention process.

Original: Our model dynamically calculates the graph connectivity based on the euclidean distances between atom pairs/triplets at each intermediate optimization step.
Revised: Our model dynamically calculates the graph connectivity based on the euclidean distances between atom pairs/triplets at each intermediate optimization step.

Original: Speciﬁcally, we exclude the contributions for pairs/triplets if they are too far from each other in the 3D space and only consider interaction inside this radius cutoff.
Revised: Speciﬁcally, we exclude the contributions for pairs/triplets if they are too far from each other in the 3D space and only consider interaction inside this radius cutoff.

Original: Table S3 shows that dynamic connectivity is consistently better than other approaches for both the
Revised: Table S4 shows that dynamic connectivity is consistently better than other approaches for both the

Original: QM9 and GEOM-Drugs datasets.
Revised: QM9 and GEOM-Drugs datasets.

Original: Although fully connected attention works better than augmented attention on the QM9 dataset, it is worse than the augmented graph on GEOM-Drugs, even with more computational complexity.
Revised: Although fully connected attention works better than augmented attention on the QM9 dataset, it is worse than the augmented graph on GEOM-Drugs, even with more computational complexity.

Original: The performance degradation here could be caused by the noise from unrelated atom pairs/triplets.
Revised: The performance degradation here could be caused by the noise from unrelated atom pairs/triplets.

Original: In contrast, our dynamic connectivity can take advantage of the predicted conformers and construct more informative connectivity than the augmented approach based on the graph distance.
Revised: In contrast, our dynamic connectivity can take advantage of the predicted conformers and construct more informative connectivity than the augmented approach based on the graph distance.

Original: [Table]
Revised: [Table]

Original: I. E FFECTS OF THE PARAMETERIZATION FOR INTERACTION WEIGHT
Revised: I. E FFECTS OF THE PARAMETERIZATION FOR INTERACTION WEIGHT

Original: While the dynamic connection discussed in the above section treats the interaction between atoms in a binary fashion, we also consider their interaction weight fractionally.
Revised: While the dynamic connection discussed in the above section treats the interaction between atoms in a binary fashion, we also consider their interaction weight fractionally.

Original: In EGNNSatorras et al.
Revised: In EGNNSatorras et al.

Original: (2021), it estimates the interaction weights with a linear layer which takes v as inputs.
Revised: (2021), it estimates the interaction weights with a linear layer which takes v as inputs.

Original: In our model, we propose a prediction layer for weights of interaction, which considers the distance between an atom pair and their atom type
Revised: In our model, we propose a prediction layer for weights of interaction, which considers the distance between an atom pair and their atom type.

Original: [Table]

Original: In Table S
Revised: In Table S

Original: [Table]
Revised: [Table]

Original: we test different approaches to parameterize the interaction weights.
Revised: we test different approaches to parameterize the interaction weights.

Original: While we ﬁnd the Per-layer v (EGNN-like) performs worse than others, the effects of distance
Revised: While we ﬁnd the Per-layer v (EGNN-like) performs worse than others [Table]

Original:  d based prediction net vary from dataset to dataset.
Revised: the effects of distance d based prediction net vary from dataset to dataset.

Original: In the models reported in the main text [Table] we use Per-block d for all variants
Revised: In the models reported in the main text, we use Per-block d for all variants

Original: [Table]
Revised: [Table]

Original: but a Per-layer d setup could improve the performance of the GEOM-Drugs dataset.
Revised: but a Per-layer d setup could improve the performance of the GEOM-Drugs dataset.

Original: Thus, we could tune this design from dataset to dataset.
Revised: Thus, we could tune this design from dataset to dataset.

Original: We leave further exploration to future work
Revised: We leave further exploration to future work

Original: [Table]
Revised: [Table]

Original: I. E FFECTS OF PARAMETER SHARING
Revised: I. E FFECTS OF PARAMETER SHARING

Original: For our models reported in the main text, we implement three optimization blocks that share their parameters while each block contains three layers of x and v updates that do not share their parameters.
Revised: For our models reported in the main text, we implement three optimization blocks that share their parameters while each block contains three layers of x and v updates that do not share their parameters.

Original: However, one could push this to either extreme in terms of parameter sharing.
Revised: However, one could push this to either extreme in terms of parameter sharing.

Original: For the same number of total update layers, we could have either nine shared optimization blocks with only one layer per block or a single optimization block with a total of nine layers that do not share parameters at all.
Revised: For the same number of total update layers, we could have either nine shared optimization blocks with only one layer per block or a single optimization block with a total of nine layers that do not share parameters at all.

Original: It is worth mentioning that all these parameter sharing designs make sense since they correspond to the different hypotheses of the shape of the gradient ﬁeld.
Revised: It is worth mentioning that all these parameter sharing designs make sense since they correspond to the different hypotheses of the shape of the gradient ﬁeld.

Original: For example, if the gradient ﬁeld is relatively smooth, sharing parameters across all blocks could also perform well.
Revised: For example, if the gradient ﬁeld is relatively smooth, sharing parameters across all blocks could also perform well.

Original: We report the performance for these two extreme cases in Table S5 and found our hybrid design strikes the right balance between the two.[Table] Block × 1-Layer [Table]
Revised: We report the performance for these two extreme cases in Table S6 and found our hybrid design strikes the right balance between the two.

Original: Block ×

Original: [Table]

Original: -Layer

Original: [Table]
Revised: [Table]

Original: I. E FFECTS OF ITERATIVE UPDATE
Revised: I. E FFECTS OF ITERATIVE UPDATE

Original: For our models in the main text, the coordinate x and the internal representation v are updated in an iterative fashion, where the update for x t +1 depends on v t +1 .
Revised: For our models in the main text [Table] the coordinate x and the internal representation v are updated in an iterative fashion, where the update for x t +1 depends on v t +1 .

Original: However, we can also update them in a parallel fashion as EGNN does, where the update for x t +1 will depend on v t instead.
Revised: However, we can also update them in a parallel fashion as EGNN does, where the update for x t +1 will depend on v t instead.

Original: These are both reasonable approaches depending on the optimization strategies.
Revised: These are both reasonable approaches depending on the optimization strategies.

Original: Empirical results in Table S6 show that optimizing x and v iteratively is slightly better than optimizing them in a synchronized fashion.
Revised: Empirical results in Table S7 show that optimizing x and v iteratively is slightly better than optimizing them in a synchronized fashion [Table] I. C OMPARISON WITH EGNN [Table]

Revised: Finally

Original: [Table]
Revised: [Table]

Revised: we summarize our different designs and their effects comparing to EGNN in Table S

Revised: [Table]

Revised: There are a couple of differences between our methods and EGNN coming from the neural energy minimization formulation including per-block interaction weight (I.2), shared parameters (I.3)

Revised: [Table]

Revised: and iterative updates (I.4).

Revised: These differences make up some of the improvements together.

Revised: On the parameterization side, arguably the biggest differences are that EGNN uses an MLP architecture while ours uses a Transformer-like architecture.

Revised: However, switching their MLP architecture with a

Revised: Transformer-like architecture provides a little bit of lift, but it’s not the major contributor.

Original: J M ORE E XAMPLES OF M ODEL -O PTIMIZED C ONFORMERS
Revised: J A BLATION S TUDIES IN C ONFORMATION G ENERATION

Revised: [Table]

Revised: We introduce ablation studies to further study the importance of different components of our proposed method in the generative setting.

Revised: We ﬁnd the formulation of framing the generative task as a ﬁxed-step unroll optimization process using an SE(3)-equivariance model is the more important bit, as we only see a very small performance regression when switching our backbone optimization model to EGNN.

Revised: On the other hand, ConfGF is proposed to model one-step optimization and relies on

Revised: Langevin dynamics for performing the entire sequence of optimization.

Revised: Additionally, the optimal transport loss also helps quite a bit, as it encourages the model to optimize the conformers towards different directions instead of collapsing the mode.

Revised: K I NFLUENCE OF N UMBER OF U NROLLED O PTIMIZATION S TEPS

Revised: We test the optimization performance on the QM9 dataset with different numbers of optimization steps.

Revised: We observe a diminishing return for steps >4 since the “gradient” would be small if we are close to the solution.

Revised: Therefore, we believe the method is not sensitive to the number of steps as long as a sufﬁcient number of steps are taken.

Revised: In all our experiments, we perform a ﬁxed number ofoptimization steps.

Revised: [Figure]

Revised: L M ORE E XAMPLES OF M ODEL -O PTIMIZED C ONFORMERS

Original: Please see next page.
Revised: Please see next page.

Original: [Figure]
Revised: [Figure]

Original:  MMFF
Revised: RDKit + MMFF

Original: EGNN
Revised: EGNN

Original: Ours
Revised: Ours

Original: Reference
Revised: Reference