Year: 2021

Abstract:
Modeling many-body systems has been a long-standing challenge in science, from classical and quantum physics to computational biology. Equivariance is a critical physical symmetry for many-body dynamic systems, which enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which significantly boosts the prediction performance of down-streaming tasks. Some general equivariant models which are computationally efficient have been proposed, however, these models have no guarantee on the approximation power and may have information loss. In this paper, we leverage insights from the scalarization technique in differential geometry to model many-body systems by learning the gradient vector fields, which are SE(3) and permutation equivariant. Specifically, we propose the Equivariant Vector Field Network (EVFN), which is built on a novel tuple of equivariant basis and the associated scalarization and vectorization layers. Since our tuple equivariant basis forms a complete basis, learning the dynamics with our EVFN has no information loss and no tensor operations are involved before the final vectorization, which reduces the complex optimization on tensors to a minimum. We evaluate our method on predicting trajectories of simulated Newton mechanics systems with both full and partially observed data, as well as the equilibrium state of small molecules (molecular conformation) evolving as a statistical mechanics system. Experimental results across multiple tasks demonstrate that our model achieves best or competitive performance on baseline models in various types of datasets.

heading: 1 INTRODUCTION
text: Modeling many-body systems has been a long-standing challenge in scientific fields from classical and quantum physics (Carleo & Troyer, 2017; Zhang et al., 2018; Satorras et al., 2021c) to structural biology (Senior et al., 2020; Shi et al., 2021), due to its high numerical complexity and complicated evolving mechanism. Graph neural network (GNN), which is superior to model the high-dimensional structured data with permutation equivariance, brings a new opportunity to model the many-body systems in an end-to-end manner. Since many-body physical systems follow many physical constraints like SE(3) symmetry, pure black-box GNN models show limitations on generalization in this scenario and symmetry-preserving GNN models have become a hot research direction.
The core question to be solved for developing general equivariant GNN models is how to conduct nonlinear operations on tensors in a reference-free way. To represent and manipulate equivariant tensor of arbitrary orders, some approaches resort to equivariant function spaces such as spherical harmonics (Thomas et al., 2018; Fuchs et al., 2020; Bogatskiy et al., 2020; Fuchs et al., 2021) or lifting the spatial space to high-dimensional spaces such as Lie group space (Cohen & Welling, 2016; Cohen et al., 2018; 2019b; Finzi et al., 2020; Hutchinson et al., 2021). Since no restriction on the order of tensors is imposed on these methods, sufficient expressive power of these models is guaranteed. Unfortunately, transforming a many-body system into those high-dimensional spaces or calculating equivariant functions usually brings excessive computational cost and great optimization difficulty, which is unacceptable in some real-world scenarios. To remedy this issue, Satorras et al. (2021c) proposed EGNN to directly implement equivariant operations in the original space, providing an efficient way to preserve equivariance without performing complex space transformations.
Detailed experiments in (Satorras et al., 2021c) have shown that preserving equivariance without transforming the space is theoretically possible and computationally efficient in practice.
However, one trade-off of EGNN is abandoning a certain amount of tensor information 1, which causes the equivariance function class EGNN can approximate to be restricted. This drawback may become serious when modeling complex dynamical scenarios such as molecular simulation, where geometric information (e.g., angular potentials and rotatable bonds) plays an important role in inducing conformation changes (Klicpera et al. (2020); Xu et al. (2021)). To mitigate this issue, we propose a new model called Equivariant Vector Field Network (EVFN) to fit the gradient vector fields (Song & Ermon, 2019; Shi et al., 2021) of many-body systems. With a scalarization block and a vectorization block, EVFN is able to represent tensor information losslessly in the original space and outputs equivariant vector fields with no restriction on the direction.
Inspired by the scalarization technique from differential geometry (Kobayashi, 1963; Hsu, 2002), EVFN first introduces a tuple of complete basis associated with each particle pair that preserves permutation and SE(3) symmetry under global rotation and translation transformation on the manybody system. Based on this basis, the scalarization block losslessly transforms the geometric information into SO(3)-invariant scalar representations. In principle, the scalar representations can be fed into any permutation-equivariant network to implement complex nonlinear transformations. Moreover, the vectorization block could reverse the scalars back to the vector field without sacrificing geometric information with the complete basis. Once obtained the estimated gradient field, we can predict a certain state or the whole dynamical trajectory of a 3D many-body system via an integration procedure.
We evaluate the proposed framework on two many-body scenarios that require equivariance: (1) the simulated Newtonian many-body dynamics trajectory prediction and (2) the real-world molecular conformation generation. Our model achieves best or competitive results in various types of datasets.

heading: 2 BACKGROUND
text: In this section, we first introduce some basic concepts on the notion of equivariance and tensor field and then describe the scalarization technique from differential geometry. Finally, we define the ‘vector field’ as the differential of a many-body system. Let X = (x1, . . . ,xN ) ∈ RN×3 be a many-body system living in R3, where N is the number of particles. We use xi(t) to denote the position of the particle i at time t.
SE(3) group and Equivariance In the Euclidean space R3 we can consider affine transformations that preserve the distance between any two points, i.e., the isometric group SE(3). We call it the symmetry group w.r.t. the Euclidean metric, and it turns out that SE(3) can be generated by the translation group and the rotation group SO(3).
Once we have a symmetry group, it’s valid to define quantities that are “equivariant” under the symmetry group. Given a function f : Rm → Rn, assuming the symmetry group G acts on Rm and Rn, then f is G-equivariant if
f(gx) = gf(x), ∀x ∈ Rm and g ∈ G. For SO(3) group, if n = 1, i.e., the output of f is a scalar, then the group action on R1 is the identity map, in this case f should be SO(3)-invariant (Thomas et al., 2018): f(gx) = f(x).
The notion of tensor field can be defined for general Riemannian manifold (see Definition 2.1.10 of (Jost & Jost, 2008)). Let { ∂
∂xi }3i=1 and {dxi}3i=1 be the tangent vectors and dual vectors in R3
respectively, and ⊗ denotes the tensor product. Then recall the definition of tensor field for R3 w.r.t. the SO(3) group: Definition 2.1. A (r, s)- tensor field θ is a multi-linear map from a collection of r vectors and s dual vectors in R3 to R: θ(x) = θi1···irj1···js ∂ ∂xi1 ⊗ · · · ⊗ ∂ ∂xir
⊗ dxj1 ⊗ · · · ⊗ dxjs . It implies that under SO(3) coordinate transformation g := {gij}1≤i,j≤n, the tensor field θ transforms equivariantly: θ i′1···i ′ r
j′1···j′s = gi′1i1 · · · gi′rirg T j1j ′ 1 · · · gTjsj′sθ i1···ir j1···js , where g T is the inverse of g. 1For example, dihedral angle is a function of position vectors rather than the position’s norm, which are the
input of EGNN. Detailed definition is given in appendix A.2.4
Equivariant Vector field To model X(t), a natural way is to estimate its differential dX(t)dt and apply an ODE solver to integrate the differential to obtain the dynamic trajectory or a state at a given time. Due to the SO(3) symmetry, we define such differential as an equivariant vector field. Most 3D real-world scenarios adopt first-order and second-order equivariant vector fields to depict their dynamic evolving mechanisms, which are also the modeling targets in this paper. A typical secondorder vector field is the acceleration field of Newtonian systems.
Gradient field is a widely used terminology meaning the first-order derivative w.r.t. a scalar function (Jost & Jost, 2008; Song & Ermon, 2019). To generate molecular conformations (i.e. equilibrium state) with a single stage, (Shi et al., 2021) define “gradient field” to serve as pseudo-forces acting on each particle. By evolving the particles following the direction of the gradient field, the nonequilibrium system will finally converge to an equilibrium state. Gradient field is a special case of first-order vector field.

heading: 3 METHODOLOGY
text: Given a many-body system X , we aim at modeling its vector field to predict the long-term dynamic trajectory or the equilibrium state within a single stage. To preserve the physical symmetries of the system, the estimated vector field should be equivariant for permutation and SE(3) group. To achieve this goal, we represent the system as a spatial graph and construct EVFN based on it with three key components: (1) a Scalarization block to encode the geometric tensors into SO(3)invariant scalar representations attached to each node; (2) a Graph Transformer block to learn SO(3)-invariant edgewise embeddings by propagating and aggregating information on the graph; and (3) a Vectorization block to reverse scalar representations back to geometric tensors to estimate the vector field. A brief overview of EVFN is illustrated in Figure 1. Once the vector field network (EVFN) is optimized, an Evolving block is incorporated to integrate the vector field for predicting the dynamics. The translation symmetry can be easily preserved by moving the particle system’s centroid at t = 0 to the origin (the Centralization operation in Figure 1). Permutation equivariance is automatically guaranteed for the message-passing based Graph Transformer. We provide detailed proof about these symmetries in Appendix A.2.1. Now we concentrate on SO(3) symmetry in the following sections.

heading: 3.1 SCALARIZATION BLOCK
text: The scalarization block is designed to transform geometric tensors into edgewise SO(3)-invariant scalar features by introducing a novel tuple of complete basis. Given a particle xi(t), define neighbor
N (xi(t)) as the particles that react with xi(t). Then we can consider a particle pair (xi(t),xj(t)), where xj(t) ∈ N (xi(t)). Suppose we take the positions of the two particles as the relevant geometric information of edge ⟨i, j⟩, then the edgewise SO(3)-invariant scalars could be defined as tij := Scalarize(xi(t),xj(t),Fij), where Scalarize is the scalarization operation under an edgewise dynamical basis Fij defined below.
Equivariant basis construction For the particle pair (xi(t),xj(t)), let a(t) = xi(t)−xj(t)
∥xi(t)−xj(t)∥ and × denotes the cross product of two vectors, define
b(t) = xi(t)× xj(t) ∥xi(t)× xj(t)∥ and c(t) = a(t)× b(t), (3.1)
Then we build a SO(3)-equivariant basis Fij := (a(t), b(t), c(t)). In practice we add a small constant ϵ to the normalization factor in case that xi and xj collapse. Under the condition that the matrix (a(t), b(t), c(t)) is non-degenerate, Fij formulates a complete orthonormal basis (frame) of the tangent space at xi(t). Note that this is a dynamical basis w.r.t. t and the construction process of such basis is permutation-equivariant. Since the Euclidean metric is flat, the dual basis (living in the cotangent space of xi (Hsu, 2002)) of Fij is just its transpose: (aT (t), bT (t), cT (t)). The ‘bad’ event that Fij is degenerate for all neighbors happens only when all particles are restricted to a straight line, which is a measure zero set in R3. Therefore, we assume Fij is non-degenerate from now on. Proof for SO(3)-equivariance of Fij is provided in proposition A.1.
Invariant scalarization of geometric tensors With the complete equivariant basis, we can realize the scalarization operation (Kobayashi, 1963) in an elementary way. First of all, notice that under the basis Fij ((a(t), b(t), c(t))), the position vector of xk naturally owns a ‘coefficient’ or ‘scalar’ representation: (xk · a(t),xk · b(t),xk · c(t)). (3.2) We define the process obtaining such scalars as Scalarize operation. Here we demonstrate that the set of obtained coefficients (3.2) is actually a SO(3)-invariant scalar tuple. Let g ∈ SO(3) be an arbitrary orthogonal transformation, then
xk → gxk and (a(t), b(t), c(t))→ (g · a(t), g · b(t), g · c(t)). Therefore (3.2) undergoes
(xk · a(t),xk · b(t),xk · c(t))→ (gxk · ga(t), gxk · gb(t), gxk · gc(t)) = (xk · a(t),xk · b(t),xk · c(t)), (3.3)
where we use the fact that gT g = I to get the last line.
It is easy to prove that the Scalarize operation could transform arbitrary geometric tensors into SO(3)-invariant scalars. Taking (2,0)-type tensors as an example, by extending the complete basis (a, b, c) through tensor product, it’s easy to check that
{a⊗ a, b⊗ b, c⊗ c,a⊗ b, b⊗ a,a⊗ c, c⊗ a, b⊗ c, c⊗ b} forms an equivariant basis of the (2, 0)-type tensor space. Then the scalarization of a tensor is just the linear combination coefficients under this basis. In the same way as (3.3), we can prove that the coefficients are also SO(3)-invariant scalars. Given a (2, 0)-symmetric tensor θ (e.g., energymomentum tensor), under the complete basis Fij = (a, b, c), θ can be expressed as:
θ =θaaa⊗ a+ θbbb⊗ b+ θccc⊗ c+ θab(a⊗ b+ b⊗ a) + θac(a⊗ c+ c⊗ a) + θbc(b⊗ c+ c⊗ b). (3.4)
The scalars tuple tij := {θaa, θab, . . . } are the scalarization of θ under the equivariant basis Fij , which are SO(3)-invariant. Since any nonlinear transformations acting on SO(3)-invariant scalars are still SO(3)-invariant, therefore the scalars tuple can be fed into any neural network architectures without concerns about breaking the equivariance symmetry. Our scalarization is inspired by the scalarization technique on the frame bundle in differential geometry (Hsu, 2002), the equivalence of the scalarization technique and (3.4) is given in proposition A.2.
In practice, we focus on the scalarization of (1, 0)-type tensors (i.e., vectors). Since the most common geometric information of input is vectors in real-world scenarios. We define the Scalarize operation as:
Scalarize(xi, xj , Fij) = (xi · a(t), xi · b(t), xi · c(t), xj · a(t), xj · b(t), xj · c(t)) (3.5)

heading: 3.2 GRAPH TRANSFORMER BLOCK
text: Algorithm 1: EVFN Input: X = (x1, . . . ,xN ) ∈ RN×3,
hi ∈ Rh, eij ∈ Re,GX // Centralization X← Centralize(X); // Scalarization Block for xi ∈ X do
for xj ∈ N (xi) do Fij = EquiBasis(xj ,xj); tij = Scalarize(Fij ,xi,xj);
end end // Graph Transformer Block mij = GraphTransformer(tij , hi, eij); // Vectorization Block Vi = 1 N ∑ j∈N (xi) Vectorize(mij ,Fij); Output: Vi
After encoding the geometric tensors into SO(3)invariant scalars tij , we first embed them alone with other pre-defined node/edge attributes (hj , eij) into high-dimensional representations, and leverage an attention-based Graph Transformer architecture (Shi et al., 2020) to learn the SO(3)invariant edgewise message embeddings mij by propagating and aggregating information on the graph GX . The attention mechanism is introduced due to its powerful capacity in modeling those graphs with unknown topology. We provide further design details in Appendix A.2.2.

heading: 3.3 VECTORIZATION BLOCK
text: Given the refined edgewise message mij , the vectorization block is designed to transform these scalars back to equivariant vectors, which requires pairing mij with the corresponding complete basis2 Fij := (a, b, c). More precisely, we first project mij into a scalar triple {x1, x2, x3}, then define the vectorization process as:
(x1, x2, x3) Pairing−−−→ x1a+ x2b+ x3c. (3.6)
We encapsulate the pairing process as Vij = Vectorize(mij ,Fij). It’s easy to check that the output follows the transformation rule of vectors. Finally, we aggregate all vectors Vij associated with xi to estimate the ground-truth vector field Vi.
So far, we have achieved permutation and SE(3) equivariance by employing the three blocks. The evolving block is introduced to simulate dynamics of the system with the optimized vector field network, acting like an ODE solver. We will discuss the evolving block designed specifically for each scenario in Section 4. The workflow of our method is summarized in Algorithm 1.
Discussion: First, the SO(3)-equivariance of EVFN can be proved from the following facts: (1) The scalarization block is SO(3)-invariant because its inputs are SO(3)-invariant scalars; (2) The SO(3)invariant variables are still SO(3)-invariant after being transformed by the graph transformer block (or any nonlinear neural network); (3) The vectorization block is SO(3)-equivariant if the input is SO(3)-invariant. We put the detailed proof in appendix A.2.3. Second, note that by pairing the scalars obtained from the scalarization block with the equivariant basis, we can recover back the tensor fields in the data. Therefore the transformation is invertible, which implies that there is no information loss when going through the scalarization block.

heading: 4 EXPERIMENTS
text: Our method (EVFN together with the evolving block) is a general framework for modeling manybody systems. To validate the effectiveness of the method, we conduct extensive experiments on two scenarios: (1) simulated Newtonian many-body systems trajectory prediction (second-order vector field task) and (2) the real-world molecular conformation prediction (first-order vector field task).

heading: 4.1 NEWTONIAN MANY-BODY SYSTEM
text: In this experiment, we apply our model to predict the long-term motion trajectory of an unknown many-body Newton system given its initial position and velocity. To highlight the strength of EVFN on including complete geometric information, we place particular emphasis on non-radical Newtonian forces under three settings.
2Since every vector is a linear combination of a basis, the output of the vectorization block has no restriction on the direction. We will discuss how our basis changes under reflection in remark A.3.
Partially observed system (POS). This system consists of six particles under Newton’s gravitation (radical force) but only four of them could be observed, i.e., for each trajectory, we are provided with positions X(t) ∈ R4×3 and velocities V (t) ∈ R4×3. The two unobserved particles act as a ‘virtual’ external field for the system.
Static external force field (SEFF). This system consists of six particles under both Newton’s gravitation and an external static force : fη = (0, 0, η), where η > 0 is the magnitude of the external field along the z-axis.
Dynamical external force field (DEFF). This system consists of three particles governed by both Newton’s gravitation and a Lorentz-like dynamical external force field, which means there exists a force field perpendicular to the current direction of each particle’s velocity: fl(v(t)) = qv(t)×B, where q is a positive constant that mimics the charge of particles and B denotes the pseudo-vector of the external field.
Problem definition. Following (Zhuang et al., 2020; Li et al., 2021), we formulate trajectory prediction as two tasks: Interpolation and Extrapolation in the original and rotated reference.
The experimental setting is as follows. For each trajectory, given the initial condition, we use observations xi(t), t ∈ {∆t, 2∆t . . . , T1} as the training labels and observations xi(t), t ∈ {T1 + ∆t, T1 + ∆t, . . . , T2} as the validation set. To evaluate the interpolation and extrapolation capacity of all methods, the observations xi(t), t ∈ { 12∆t, 3 2∆t, . . . , T1 + 1 2∆t} and xi(t), t ∈ {T2 + ∆t, T2 + 2∆t, . . . , T3} are used as interpolation and extrapolation test sets respectively. We measure the mean square error (MSE) between the predicted trajectory and ground truth for both tasks. To measure the exactness of equivariance, we follow (Fuchs et al., 2020) to apply uniformly sampled SO(3)-transformations on the input and output. The MSE between the predicted trajectory with rotated input and rotated ground truth could reflect the transformation robustness of method. The normalized distance between the rotated prediction with original input and the original prediction with the rotated input defines the equivariance error ∆EQ:
∆EQ = ∥LsΦ(x)− ΦLs(x)∥/∥LsΦ(x)∥, (4.1) where Ls and Φ denote SO(3) transformations and equivariant neural networks, respectively.
Learning Framework. Inspired by (Norcliffe et al., 2020), for a Newtonian system, we utilize EVFN to parameterize its acceleration vector field and adopt a second-order neural ODE (SNODE) as the evolving block (see Appendix A.2.5) to integrate both the position and velocity trajectories. Only the MSE between predicted position trajectory and ground truth is taken as the loss penalty:
L(θ) = 1 n n∑ i=1 L2(xti ,SNODE(xt0 ,vt0 , t0, ti,Θ)), t0 < t1 < · · · < tn, (4.2)
where (xt0 ,vt0) and Θ denote the initial condition of the system and the parameters of EVFN Φ.
Implementation Details. Following (Zhuang et al., 2020), all trajectories are simulated using the Dopri5 solver (Dormand & Prince, 1980) with the tolerance to 10−7 and the modified physical rules. In this experiment, we set η to 0.98, q to 1 and B to [0.5, 0.5, 0.5]⊤, respectively. We sample 1, 100 and 100 trajectories for the three systems as our evaluation platform, where the data is split into training, validation and test sets by the time span: T1 = {1.5, 0.5, 0.5}, T2 = {2, 0.55, 0.55} and T3 = {2.5, 0.6, 0.6}. We compare our method to the non-equivariant graph convolutional network (GCN) (Kipf & Welling, 2016) without and with SO(3) data augmentation (denoted as GCN (Aug)), as well as two equivariant models designed for vector field modeling: Radial Field (Köhler et al., 2019) and EGNN (Satorras et al., 2021a;c). Further details about data generation and model implementation are provided in Appendix A.3.1.
Results. All the baselines and EVFN implement the feature transformations in the original space. Notice that GNNs are manifestly permutation-equivariant and the centroid of data is reduced for all models to preserve translation equivariance. Thus we only need to evaluate the generalization capacity of all models for SO(3) transformations.
From Table 1, we have the following conclusions: (1) With the original input, EVFN outperforms all other equivariant methods in the interpolation and extrapolation tasks, which demonstrates that EVFN exhibits stronger expressive power by representing geometric information losslessly3;
3We put more discussions especially the performance of the non-equivariant GCN in Appendix A.3.1.
(2) With the rotated input, EVFN still performs best and all equivariant networks exhibit better equivariance-preserving capacity than GCNs even with data augmentation, implying that it is difficult to achieve equivariance with simple data augmentation. We leave a detailed analysis about different augmentation degrees in Appendix A.3.1. 4

heading: 4.2 MOLECULAR CONFORMATION GENERATION
text: In this experiment, we leverage EVFN to fit the first-order vector field of small molecules’ equilibrium distribution and generate reasonable conformations for a given molecular graph. Note that in this case, the SO(3) and translation symmetry are in the distribution level, rather than a specific trajectory. Previous works demonstrate that geometric information (e.g., angular potentials and rotatable bonds) plays a critical role in inducing conformation changes (Xu et al., 2021; Klicpera et al., 2020), which provides a natural platform to evaluate the strength of EVFN.
Evaluation Tasks. We conduct experiments on two tasks: (1) Conformation Generation evaluates the capacity of EVFN to learn the conformation distribution by measuring the diversity and accuracy of generated conformations. (2) Distributions Over Distances evaluates the discrepancy of distance geometry between the generated conformations and the reference conformations.
Learning Framework. Following (Shi et al., 2021), for this first-order statistical ensemble system, we leverage a score-based generative modeling framework to estimate the gradient field of atomic positions (See more details about score-based networks in (Shi et al., 2021; Song et al., 2020) or Appendix A.2.5, A.3.2). The optimization objective of EVFN Φ can be summarized as:
L(θ) = 1 n n∑ i=1 EX(0){λ(ti)∥∇Hti(X(ti))− Φ(X(ti), ti,Θ)∥ 2 2, t0 < t1 < · · · < tn, (4.3)
where λ(t) : [0, T ] → R+ is a positive weighting function and ∇Hti is the pre-computed gradient field of noisy atomic positions. Once the score network is optimized, we can use an annealed Langevin dynamics (ALD) sampler or an ODE-based PC sampler as the evolving block to generate conformations (Song et al., 2020).
4Note that the division operation in the basis construction procedure will bring considerable numerical error into EVFN which increases ∆EQ. Therefore, we tune the constant C in a(t) =
xi(t)−xj(t) ∥xi(t)−xj(t)∥+C
. The model variant EVFN* with the suitable C can achieve a comparable equivariant error with other baselines.
Datasets. Following (Xu et al., 2021; Shi et al., 2021) we evaluate the proposed model on the GEOM-QM9 and GEOM-Drugs datasets (Axelrod & Gomez-Bombarelli, 2020) as well as the ISO17 dataset (Schütt et al., 2017). To keep a fair comparison with the existing state-of-the-art (SOTA) method ConfGF (Shi et al., 2021), we reproduce its data collection and split settings rigorously. Further details are described in Appendix A.3.2.
Metrics. Given the RMSD of heavy atoms that measures the distance between generated conformation and the reference, Coverage (COV) and Matching (MAT) scores are defined to measure the diversity and accuracy for a given RMSD threshold δ respectively.
COV(Sg, Sr) = 1
|Sr| |{R ∈ Sr|RMSD(R, R̂) < δ, R̂ ∈ Sg}|, (4.4)
MAT(Sg, Sr) = 1 |Sr| ∑ R∈Sr min R̂∈Sg RMSD(R, R̂), (4.5)
where Sg and Sr denote generated and reference conformations, respectively.
Implementation Details. The EVFN is equipped with 4 Graph Transformer blocks and the hidden dimensions are set to 288. All models are trained with Adam optimizer via the loss function (4.3) for 400 epochs. For each molecule in the test set, we follow (Shi et al., 2021) to sample twice as many conformations as the reference ones from each model. We provide all hyperparameters of the score-based framework in Appendix A.3.2.
Baselines. We compare EVFN to four classic methods for conformation generation. Specifically, both RDKit (Landrum, 2013) and CGCF(Xu et al., 2021) are distance-based approaches. ConfGF (Shi et al., 2021) is most close to us, attempting to generate conformations by learning the gradient field of the data distribution in an equivariant manner. However, ConfGF only utilizes the distance matrix as the geometric input. We also reproduce EGNN on this task as our baseline.
Results. We summarize the mean and median COV and MAT scores on two benchmarks for all methods. As shown in Table 2, EVFN achieves the best performance on almost all metrics and datasets, demonstrating the effectiveness of our proposed method. Compared with ConfGF which employs a similar learning strategy with us, EVFN significantly increases 26.5% COV-mean and 26.6% COV-median scores on the GEOM-Drugs dataset. A potential interpretation is that molecules in Drugs usually contain more atoms and complex chemical functional groups (e.g., Benzene rings) than those of QM9, thus distance-based geometry is not sufficient to model the gradient field of this complex distribution. EVFN also achieves significant improvement in the Distributions Over Distances task, and we provide the empirical results and further discussions in Appendix A.3.2.
Ablations. Although the superior performance on multiple tasks verifies the effectiveness of EVFN, it remains unclear whether the proposed strategies make a critical contribution. In light of this, we set up several ablative configurations and list the empirical results in Table 3. For the scalarization block, we conduct a variant of EVFN without the scalarization block, named EVFN w/o Sca., which only takes the distance matrix of molecules as the input geometric feature. The results show that including scalarization block plays an important role in the model, as the COV-mean and COV-median scores on the Drugs dataset increase by 23.0% and 21.9%, respectively, which implies that including all geometric information will boost the performance of the model. For the graph transformer block,
we replace the block with the GIN network (Xu et al., 2018) that is employed in ConfGF, getting the variant named EVFN w/o GT. The results indicate that introducing the attention mechanism will also contribute to the gradient field modeling. We cannot conduct the ablative study for the vectorization block because it guarantees the output of EVFN is an equivariant vector field.

heading: 5 RELATED WORK
text: Equivariant neural network. Existing equivariant networks with a theoretical guarantee can be roughly classified into two categories by whether conducting all operations in the original space or not. The first category of methods lift the data into high-dimensional spaces (e.g., lie group) or introduce equivariant functions (e.g., spherical harmonics) to preserve equivariance (Worrall et al., 2017; Thomas et al., 2018; Kondor et al., 2018; Weiler et al., 2018b;a; Weiler & Cesa, 2019; Esteves et al., 2020; Romero et al., 2020; Klicpera et al., 2020; Anderson et al., 2019). They exhibit sufficient expressive power but usually bring expensive computational costs. Our work follows methods of the second category that operates on the original space in a computationally efficient way (Schütt et al., 2018; Köhler et al., 2019; Shi et al., 2021). However, most of these approaches (e.g., EGNN (Satorras et al., 2021c)) abandon a certain amount of geometric (tensor) information, causing their expressive power to be restricted. Different from existing methods, we propose a novel architecture that avoids complex vector-level transformations while preserving complete geometric information (see the discussion at the end of section 3). Essentially, the scalarization is also adopted by EGNN, which is only computed from one basis and EVFN extends the one basis (i.e., radical direction xi − xj) in EGNN into a tuple of bases for information lossless. Gradient fields modeling. Gradient fields modeling is one of the popular tools for modeling many-body systems from predicting motion trajectories of physical systems (Greydanus et al., 2019; Norcliffe et al., 2020; Li et al., 2020) to estimating probabilistic densities of complex systems (Song & Ermon, 2019; Cai et al., 2020; Shi et al., 2021). Neural ODEs (Chen et al., 2018; Zhuang et al., 2020) are built for learning the gradient of a system by the adjoint method. With the estimated gradient, the NODEs can make predictions for irregular time series by integrating to any given time. Score-based methods attempt to model the data distribution by learning the gradient of its probabilistic density (Song & Ermon, 2020; Song et al., 2020). More precisely, Sohl-Dickstein et al. (2015); Shi et al. (2021); Wu et al. (2021) model the equilibrium states by going through the process determined by the gradient of the density.

heading: 6 CONCLUSION AND FUTURE WORK
text: To learn the gradient fields for many-body system modeling, we propose an equivariant vector field neural network (EVFN) aiming at lossless utilization of tensors without incorporating highdimensional spaces or equivariant functions. With the proposed scalarization technique, EVFN could cooperate with any neural networks without concerns about breaking the equivariance symmetry. Theoretical analyses and extensive empirical results verify the effectiveness of the proposed method. In the future, we will investigate the performance of EVFN in large-scale many-body systems and extend the strategy to other symmetry groups or local symmetry group such as gauge group (Cohen et al., 2019a).

heading: A APPENDIX

heading: A.1 PROOF AND DISCUSSION ON SCALARIZATION AND VECTORIZATION
text: Proposition A.1. The complete basis (a(t), b(t), c(t)) defined by (3.1) is equivariant under SO(3) transformation of the spatial space.
Proof. Let g ∈ SO(3), then under the action of g, the positions of the many-body system X(t) transform equivariantly:
(x1(t), . . . ,xn(t)) g−→ (gx1(t), . . . , gxn(t)).
Then from the definition of a(t), we know that
a(t) g−→ ga(t).
For b(t), since
(gxi(t))× (gxj(t)) = det(g)(gT )−1(xi(t)× xj(t)) (A.1) = g(xi(t)× xj(t)), (A.2)
where we have used g−1 = gT for orthogonal matrix g to get the last line. Therefore, b(t) g−→ gb(t). Applying (A.1) once again, we have c(t) g−→ gc(t).
Frame bundle and scalarization technique A frame at x ∈ M is a linear isomorphism u from Rd to the tangent space at x:TxM . We use F (M)x to denote the space of all frames at x. Then GL(d,R) acts on F (M)x by Rd g−→ Rd u−→ TxM. Then the frame bundle F (M) = ∪x∈MF (M)x can be made into a differential manifold. From the principle bundle point of view, each differential manifold M is a quotient of its frame bundle F (M) by the general linear group GL(d,R): M = F (M)/GL(d,R). We denote the quotient map by π: F (M) π−→ M. Then, each point of u ∈ F (M) is a reference frame located at x := π(u) ∈ M . On the other hand, Rd can be seen as a d-dimensional differential manifold with a global coordinates chart and SO(3) is the structurepreserving group of the Euclidean metric with a fixed orientation.
Following Hsu (2002), let {ei, 1 ≤ i ≤ d} be the canonical basis of Rd, and {ei} the corresponding dual basis. At each frame u, the vectors Yi := uei form a basis of TxM . Let {Y i} be the dual frame of T ∗xM , then a (r,s)-tensor θ can be expressed as
θ = θi1···irj1···jsYi1 ⊗ · · · ⊗ Yir ⊗ Y j1 ⊗ · · · ⊗ Y js .
The scalarization of θ at u is
θ̃(u) := θi1···irj1···jsei1 ⊗ · · · ⊗ eir ⊗ e j1 ⊗ · · · ⊗ ejs . (A.3)
Through scalarization, a tensor field θ becomes an ordinary vector space valued function on F (M):
θ̃ : F (M)→ Rr ⊗ Rs.
Therefore, geometric operations such as covariant derivative and tensor product on manifolds can be realized as directional derivative and tensor product of ordinary vector spaces (Hsu, 2002). Proposition A.2. There is a one-to-one correspondence between the scalarization of tensor fields on the orthonormal frame bundle O(R3) (A.3) and the SO(3)-invariant scalars tuple obtained by the scalarization block (3.4) under an equivariant basis.
Proof. Since we are working in R3 with a fixed orientation, the GL(3,R) group action is reduced to SO(3) global group action acting on the orthonormal frame bundle O(R3). A frame u ∈ O(R3) at π(u) ∈ R3 can be transported to another point in R3 by translation. Therefore, we neglect the
origin of the frame in the proof. Let ue = (e1, e2, e3) be an equivariant basis of R3, then given a scalars tuple {θi1,...,ir}, we construct a vector-valued function on O(R3) by:
θ̃i1,...,ir (u) = gi1i′1 · · · giri′rθ i′1,...,i ′ r ,
where g is the SO(3) transformation from ue to another frame u. Moreover, θ̃i1,...,ir (u) is SO(3)equivariant, since
θ̃(gu) = gθ̃(u),
where the g on the right side means the usual extension of the action of SO(3) from R to the tensor space R⊗r. We have constructed the (r,0) tensor field from the scalars tuple. It’s easy to check that θ̃i1,...,ir (u) induces a (r,0) tensor field on R3 by following the definition of (2.1).
From scalarization θ̃i1,...,ir (u) on O(R3) to scalarization is obvious. Note that any equivariant basis ue is also a point on O(R3), therefore the scalars tuple is just the values of θ̃i1,...,ir at ue:
θi1,...,ir = θ̃i1,...,ir (ue).
For general (r,s)-type tensors, the proof is the same by adding the dual basis of ue.
Remark A.3. For molecular structures, one common obstruction for distance-based modeling (Shi et al., 2021) is the chirality problem (Bruice, 2000). Therefore it’s meaningful to investigate how the equivariant basis (3.1) transforms under reflection x→ −x. Notice that a→ −a, and
b = xi × xj → b.

heading: It implies that c → −c. In conclusion, the orientation of the equivariant basis remains unchanged under reflection.

heading: A.2 METHODOLOGY

heading: A.2.1 PERMUTATION AND TRANSLATION EQUIVARIANCE
text: For a many-body system X(t) = (x1(t), . . . ,xn(t), the centroid C(t) is defined by
c(t) = x1(t) + · · ·xn(t)
n .
Translating the reference by h, then
X(t) + h→ c(t) + h.
Therefore, recentering the reference’s origin to the centroid at the input’s time t = 0, we have
X(t)− c(0) translation by h at t=0−−−−−−−−−−−−−→X(t)− c(0).
That is, the system is translation-invariant under the recentered reference if the translation is done at the input’s time t = 0, which is exactly the scenario considered in predicting the future trajectory or state. Note that although the gradient field is supposed to be translation-invariant, the state of the many-body system after integrating should be translation-equivariant. Therefore, we add c(0) back for the output of the evolving block.
To emphasis the permutation symmetry, we quote the Kolmogorov–Arnold representation theorem: if f(x1, . . . ,xn) is a permutation invariant multivariate continuous function, then
f(x1, . . . ,xn) = g( n∑ i=1 ϕ(xi)). (A.4)
The crucial point is the function ϕ is shared among all the points, which exactly fits the definition of the so-called message-passing scheme. For a many-body system X , let v = (v1, . . . vn) be its vector field, then vi ∈ R3 corresponds the equivariant vector for xi ∈ {x1, . . . ,xn}. Denote EVFN with parameters θ by Φθ,t = {Φθi }ni=1, then
vi(t) = Φ θ i (X(t), t),
for a fixed particle xi. Suppose (i1, . . . , ik) (neighbors of xi) indicate indexes of particles which have interaction with xi, then obviously 1 ≤ k ≤ n− 1. By (A.4), Φi(X(t), t) is an aggregation of message from xi’s neighbors, therefore we have:
Φi(X(t), t) = 1
k k∑ j=1 ϕ(xi(t),xij (t), t), (A.5)
and ϕ is a SO(3)-equivariant network with vector output. Note that (A.5) performs aggregation at the level of vectors, therefore we choose g in (A.4) to be the arithmetic mean to preserve SO(3) symmetry. A shared neural network ϕ guarantees the permutation equivariance of EVFN.

heading: A.2.2 GRAPH TRANSFORMER BLOCK
text: For a many-body system X , we first represent it as a spatial graph and utilize an attention-based to learn the SO(3)-invariant edgewise embeddings mij from the graph. Considering that there does not exist any graph topology in most real-world scenarios, we introduce the attention mechanism due to its powerful capacity in learning the correlations between inter-instances (Vaswani et al., 2017). Now we discuss the workflow of the GTB block. Note that here we omit the time index of the SO(3)-invariant scalars for the simplicity of description.
Feature embedding Given geometric scalars tij , node features hi and edge features mij , GTB first embeds them into high-dimensional representations:
hi = MLP(hi), (A.6) eij = MLP(eij), (A.7) tij = Fourier(tij), (A.8) eij = eij + tij , (A.9)
where MLP denotes a fully connected network and Fourier denotes a Fourier transformation with a tuple of learnable frequencies.
Transformer block The overall architecture of our GTB block is inspired by (Shi et al., 2020). For each GTB block, we first compute the edge-wise message with ϕ1m (A.7), then leverage the message embeddings and node embeddings with a transformer encoder-like architecture to refine the node embeddings. After that, we update edgewise messages with a residue block (A.13). The whole pipeline is expressed as (A.9)-(A.15).
mij = ϕ 1 m(hi, hj , eij), (A.10)
qi = ϕq(hi), kij = ϕk(hi,mij), vij = ϕv(mij), (A.11)
αij = ⟨qi, kij⟩∑
j′∈N (i) ⟨qi, kij′⟩ , (A.12)
Mi = LayerNorm( ∑
j∈N (i)
αijvij), (A.13)
hi = ϕh(hi,Mi), (A.14) hi = hi + LayerNorm(hi), (A.15)
mij = ϕ 2 m(hi, hj , eij) +mij , (A.16)
where ϕ1m, ϕq , ϕk, ϕv , ϕh and ϕ 2 m are fully connected networks. αij andMi denote the attention weights and the refined nodewise embeddings, respectively. LayerNorm refers to the normalization layer adopted in (Vaswani et al., 2017).

heading: A.2.3 EQUIVARIANCE OF EVFN
text: We need the following lemma:
Lemma A.4. Suppose h is an invariant function, f1, . . . , fk are arbitrary nonlinear functions. Then, the composition fk ◦ · · · ◦ f1 ◦ h is an invariant function.
Proof. Consider the group action g acting on x ∈ R3, then f is invariant means that f ◦ g(x) := f(gx) = f(x). We have fk ◦ · · · ◦ f1 ◦ g = fk ◦ · · · ◦ f1, hence the composition fk ◦ · · · ◦ f1 is invariant.
Let f1 be the scalarization block. Because the output of f1 are scalars (from (3.3)) and SO(3)scalars must be invariant under SO(3) group action, we can conclude that f1 is invariant. Denote each graph-transformer layer by fi, i ∈ {2, . . . , k}, then by Lemma A.4, we can conclude that the output of the graph-transformer block is also invariant. Finally, from the standard fact that scalars multiply vectors yield equivariant vectors, we conclude that the output of the vectorization block is SO(3)-equivariant. Combining with the translation equivariance in A.2.1, our model is SE(3)-equivariant.

heading: A.2.4 DIHEDRAL ANGLE
text: For a given node xl with three neighbors xi, xj and xk. The dihedral angle of the plane spanned by (xl − xi,xl − xj) and the plane spanned by (xl − xi,xl − xk) is given by the inner product of normal vectors of the two planes:
(xl − xi)× (xl − xj) ∥(xl − xi)× (xl − xj)∥ ,
and (xl − xi)× (xl − xk) ∥(xl − xi)× (xl − xk)∥ .
Therefore, the dihedral angle is a function of the position vectors xl, xi, xj and xk, which cannot be determined purely by the positions’ norm: ∥xl − xi∥,∥xl − xj∥ and ∥xl − xk∥. However, by the complete equivariant basis, the position vectors are fully expressed (only the coordinates are transformed from the natural coordinate system to the equivariant coordinate system), then all functions of the position vectors can also be expressed.

heading: A.2.5 EVOLVING BLOCK
text: According to the order of vector field in dynamic systems, we design different evolving blocks to integrate the estimated vector field to obtain the dynamics.
Recall that given a multi-particle Newton dynamic system X(t) = {xi(t) ∈ R3, i = 1, 2, · · · , n} whose underlying physical rules are unknown, let vi denote the velocity of particle xi. We attempt to predict its position or trajectory by integrating a second-order equivariant vector field. In this case, the EVFN network Φ is implemented for modeling the acceleration vector v̇(t). More precisely, the hidden differential unit of NODE (Chen et al. (2018)) has the following form:[
Ẋ(t) v̇(t)
] = [ v(t)
Φ(X(t),v(t), t)
] . (A.17)
For the i-th particle,
v̇i(t) = Φi(X(t),v(t)) := k∑ j=1 ϕ(xi(t),xij (t), t). (A.18)
The ordinary differential equation (A.17) is then solved by black-box ODE solver as in (Chen et al., 2018). We denote the evolving block in the second-order vector field case by SNODE.
As to the statistical ensemble system, we try to predict the reverse evolving process from a random state to equilibrium by integrating a first-order equivariant vector field. For example, all physicalallowable molecule conformations are located in an equilibrium state determined by the energy function. Suppose the forward process from equilibrium to non-equilibrium of the system satisfies:
dX(t) = f(X(t), t)dt+ g(t)dWt, 0 ≤ t ≤ T,
where Wt is the Brownian motion and the initial state X(0) follows an unknown equilibrium distribution p0. Denote the marginal distribution at time t by pt, then
pt(X) = exp{−βHt(X)}, so the Hamiltonian function Ht at time t is entangled with p0. According to the Liouville equation (the probability flow in Song et al. (2020)), the reverse evolving process satisfies the following ODE:
dX(T − t) = f(X(T − t), T − t)dt− 1 2 g2(T − t)∇HT−t(xT−t)dt.
The gradient field of the Hamiltonian function are also called the force field. Therefore the NODE evolving block has the following form:
Ẋ(T − t) = Φ(X(T − t)), 0 ≤ t ≤ T. and for the i-th particle, following form:
Φi(X(T − t)) = fi(X(T − t), T − t)− 1
2 g2(T − t)[ k∑ j=1 ϕ(xi(T − t),xij (T − t), T − t)]. (A.19)
In this case, we will use the PC scheme proposed by (Song et al., 2020) for correcting the numerical integration error. Note that the vector-valued function f(x, t) and the scalar function g(t) is prior knowledge and set to be fixed, so the only learnable module in EVFN is the vector field network ϕ. In Shi et al. (2021) and our molecular experiment, we use the discretization of VP SDE (Song et al., 2020), where f ≡ 0 and
g(t) =
√ [dσ2(t)]
dt .

heading: A.2.6 NEURAL ODE AS A CONTINUOUS LIMIT
text: The idea of neural ODE is to fit dynamics by modelling its infinitesimal rates of change (gradient) and implementing numerical integration. The whole process can be seen as taking the continuous limit of the residual network as the depth goes to infinity. Consider a residual network where all hidden layers have the same dimension, denote the neural network’s parameters at t-th layer by θt, then ht+1 = ht + f(ht, θt), t ∈ N+. Taking the discrete layer index t to its continuous limit, we get
dh(t)
dt = f(h(t), θt), t ∈ R+.
The discrete back-propagation method also has a continuous limit: the adjoint method Chen et al. (2018); Norcliffe et al. (2020).

heading: A.2.7 SCALABILITY
text: Compared to EGNN, the extra cost of EVFN is to calculate the values of the scalars in tij ,∀i, j. Compared to the computational cost of back-propagation in the neural networks, the cost of calculating the O(3 ∗ 3 ∗N2) (N is the number of nodes in the graph) scalars is much less. For example, for the molecular conformation generation task, transforming the tensors into SO(3)-invariant scalars only brings 9.6% extra real-time computational cost and 17.4% extra memory cost.

heading: A.3 EXPERIMENT

heading: A.3.1 NEWTONIAN MANY-BODY SYSTEM
text: Partially observed system This system consists of six particles under Newton’s gravitation but only four of them could be observed, i.e., for each trajectory we are provided with positions X(t) ∈ R4×3 and velocities V (t) ∈ R4×3. The time evolution of the particles is given by
ẍi = ∑
j∈{1,...,i−1,i+1,...6}
−mj xi − xj ∥xi − xj∥3 , 1 ≤ i ≤ 4. (A.20)
Gravity field This system consists of six particles under both the mutual newton’s gravitation and an external static force of the form: fg = (0, 0, η). The time evolution of the particles are given by
ẍi = ∑
j∈{1,...,i−1,i+1,...6}
−mj xi − xj ∥xi − xj∥3 + fη, 1 ≤ i ≤ 6. (A.21)
Lorentz force field This system consists of three particles and controlled by Newton’s law of motion and a Lorentz field, which means there exists a force field perpendicular to the direction of velocity v, i.e., fl(v) = qv × B, where q and B denote the charge of particles and the direction vector of the electromagnetic field respectively. The time evolution of the particles is given by:
ẍi = ∑
j∈{1,...,i−1,i+1,...3}
−mj xi − xj ∥xi − xj∥3 + f il (vi), 1 ≤ i ≤ 3. (A.22)
Implementation Details Following (Zhuang et al., 2020), all trajectories are simulated using the Dopri5 solver (Dormand & Prince, 1980) with the tolerance to 10−7 and the modified physical rules. The trajectory points are uniformly sampled with ∆t = 10−3 for the POS dataset or ∆t = 5 · 10−4 for the other two datasets. The number of training labels (i.e., trajectory points) for three tasks is 1.5k, 100k and 100k, respectively. We implement all baselines and our method with Pytorch (Paszke et al., 2019). All models use the same ODE solver (Dopri5) as the evolving blocks and are trained with Adam optimizer (Kingma & Ba, 2014) via an MSE loss for 800 epochs. We set the number of layers to 2 for all models and adjust the hidden dimensions of each model separately to keep the parameters in the same level. We adopt EGNN from (Satorras et al., 2021a) for outputting vectors. For all datasets, we take the time t and the Ł2 norm of velocities as the node features and take the relative distances as the edge features.
Further Analysis of Results As shown in Table 1, the non-equivariant GCN performs the best with the original input on SEFF task. The reason is that the equivariant models are superior to model the force that is a function of the radial direction, however, the external force in SEFF is a constant vector that is independent with the radial direction.
We also compare EVFN with the non-computational efficient equivariant models such as SE(3)Transformer (Fuchs et al., 2020), which leverages spherical harmonics to preserve equivariance. In practice, we find that it is hard to jointly optimize the neural ODE and the SE(3)-Transformer model by directly combining the two models together. And the empirical results in Table 4 show that SE(3)-Transformer achieves an extremely poor performance compared with EVFN. To the best of our knowledge, there is no previous work studying the combination of SE(3)-Transformer and neural ODEs. We would like to further explore this problem in future work. For the results in Table 4, all SE(3)-Transformer models are implemented according to the official codebase 5, where the number of hidden layers is set to 2 and each layer has representation degrees {0, 1} and channels {8, 4}. This setting makes the number of parameters of SE(3)-Transformer (i.e., 40k) comparable with EVFN and other baselines. We also implemented another variant with the default setting of SE(3)Transformer, where each layer has representation degrees {0, 1, 2, 3} and 8 channels per degree. We find it would bring much more parameters (i.e., 118k) and computational efforts, but only marginal improvement (e.g., 0.405 vs. 0.422 for MSE on the POS dataset).
Data Augmentation For all Newtonian many-body systems, we propose to force a naive GCN to be SO(3) equivariant by introducing numerous SO(3) augmentation samples. The augmentation strategy is that, for each iteration, we rotate the initial condition and the ground-truth trajectory with a random probability p. We visualize four MSE metrics of GCN (Aug) with different augmentation probabilities on the POS dataset in Figure 2. As shown in Figure 2, with the original input, the interpolation and extrapolation capacity of GCN is damaged as p grows. with the rotated input, the MSE of GCN (Aug) gradually decreases as p grows, but it is still higher several magnitudes than that of real equivariant networks, implying that it is difficult to achieve equivariance with simple data augmentation.
5https://github.com/FabianFuchsML/se3-transformer-public
The Impact of the Number of Training samples To further explore the impact of the number of training data for both non-equivariant and equivariant models, we re-train GCN and EVFN with the different number of training samples on the POS dataset and summarize the results in Figure 3. In particular, under the same data splitting setting described in Section 4.1, we choose to sample the trajectories with different time intervals {10−2, 1.5 · 10−3, 5 · 10−4, 3 · 10−4}, corresponding to 150, 1000, 3000 and 5000 training labels, respectively. The time interval of the test set remains the same as before (i.e., 10−3). Note that we avoid introducing new trajectories into the dataset because it may influence the problem complexity of the original interpolation and extrapolation tasks. As shown in Figure 3 (a) and (b), with the number of samples growing, the interpolation performance of GCN becomes better but its extrapolation becomes worse, indicating that the GCNbased neural ODE may easily overfit to the interpolation task. Compared to GCN, EVFN achieves more robust performance to the different number of training samples. As shown in Figure 3 (c) and (d), introducing more training samples in the original coordinate frame cannot enhance the equivariant capacity of GCN efficiently, which demonstrates that SO(3) augmentation is a more reasonable choice to realize equivariance for non-equivariant methods.
Equivariant Errors As discussed in Section 4.1, the bigger equivariance error of EVFN (Table 1, col 7) is caused by numerical errors introduced by the basis construction procedure and the neural
ODE training framework. More precisely, to build the equivariant basis, we incorporated the normalization operation x∥x∥ → x ∥x∥+C (x denotes vectors and C denotes the normalization constant) to obtain unit vectors, where the division operation would introduce numerical errors and neural ODE would lead to a larger cumulative error with the integration process. The similar phenomenon is also observed in SE(3)-Transformer (Fuchs et al., 2020) and EGNN (Satorras et al., 2021b). As shown in Table 5, with some suitable normalization constants, EVFN can achieve comparable or even better equivariant errors than other equivariant models, validating the equivariance of the pro-
posed method. In most cases, EVFN with larger normalization constants can still outperform other baselines significantly.

heading: A.3.2 MOLECULAR CONFORMATION GENERATION
text: Dataset For each dataset, 40, 000 molecules are randomly drawn and 5 most likely conformations (sorted by energy) are selected for each molecule, and 200 molecules are drawn from the remaining data, which results in 200, 000 conformations in the training set, 22, 408 and 14, 324 conformations in the test set for GEOM-QM9 and GEOM-Drugs datasets, respectively. The distances over distributions task are evaluated on the ISO17 dataset, where we follow the setup in (Simm & Hernández-Lobato, 2019).
Implementation Details Besides the geometric input, we feed the node type, edge type and relative distances as extra node/edge attributes into the graph transformer block. Our score-based training framework is adapted from (Shi et al., 2021). The maximum and minimum noise scales are set to 10 and 0.01. Let {σi}Li=1 be a positive geometric progression scheme with a common ratio, we split the noise range into 50 levels. For the reverse process, we find the performance difference between the ALD sampler and the PC sampler is marginal, so we do not compare this point with quantitative results. We choose the PC sampler as our evolving block and respectively set the iteration steps of predictor and corrector to 10 and 100. The sample step size ηs is chosen according to (Song & Ermon, 2020). We keep all hyper-parameters mentioned in the forward and reverse process the same as (Shi et al., 2021). The results reported in Table 2 are copied from (Shi et al., 2021) considering that we rigorously evaluate EVFN on the same benchmark and data split setting.
Conformation Generation Here we introduce the calculation equation of RMSD:
RMSD(R, R̂) = min( 1
n n∑ i=1 ||Ri − R̂i||2) 1 2 , (A.23)
where n denotes the number of heavy atoms.
We visualize several conformations in the Drugs dataset in Figure 4 that are best aligned with the reference ones generated by different methods, illustrating EVFN’s superior capacity on generating high-quality drug molecular conformation.
Distributions over Distances To evaluate the distribution of the generated conformations, we utilize maximum mean discrepancy (MMD) (Gretton et al., 2012) to measure the discrepancy between the generated distributions and the reference distributions. As shown in Table 6, EVFN dramatically
outperforms the previous SOTA (ConfGF), demonstrating the strong capacity of the proposed model in modeling molecular dynamics data. In particular, EVFN reduces the MMD by a magnitude in both Single-median and Pair-Median metrics.

Review summaries:
The authors introduce a model to predict the time evolution of Newton mechanical systems and small molecules.  The model takes a graphical representation as input and converts it to an SE(3) and permutation equivariant representation using physical principles (i.e., white-box model). This representation is passed through a learned graph transformer module to produce a vector field which is used to predict the time evolution of the system/molecule. 

This work (EVFM) aims to to improve predictions of n-body system dynamics by combining continuous lie symmetries with permutation symmetry. The authors propose to do this by encoding SO(3) invariant representations of each node, followed by the use of a graph transformer and a vectorization block to estimate the vector field. An evolving block is subsequently used to predict dynamics.