Equivariant Vector Field Network for Many-body System Modeling
Authors: weitao Du
He Zhang
Yuanqi Du
Qi Meng
Wei Chen
Bin Shao
Tie-Yan Liu
Venue: ICLR 2022 Submitted (2022)
Keywords: equivariant neural network
gradient fields
many-body system
molecular conformation generation

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. 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.

Under review as a conference paper at ICLR 2022
EQUIVARIANT VECTOR FIELD NETWORK FOR MANY- BODY SYSTEM MODELING
Anonymous authors Paper under double-blind review
ABSTRACT
Modeling many-body systems has been long-standing challenge in science, from classical and quantum physics to computational biology. Equivariance is critical physical symmetry for many-body dynamic systems, which enables robust and ac- curate 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 equívariant models which are computationally efficient have been pro- posed, however, these models have guarantee the approximation power and may have information loss. In this paper, we leverage insights from the scalariza- tion technique in differentiall geometry to model many- systems by learning the gradient vector fields, which are SE(3) and permutation equivariant. Specifi- cally, propose the Equivariant Vector Field Network (EVFN), which is built a novel tuple of equivariant basis and the associated scalarization and vectoriza- tion layers. Since our tuple equivariant basis forms complete basis, learning the dynamics with our EVEN has no information loss and no tensor operations are in- volved before the final vectorization, which reduces the complex optimization on tensors to a minimum. We evaluate our method on predicting trajectories of sim- ulated Newton mechanics systems with both full and partially observed data, as well as the equilibrium state of small molecules (molecular conformation) evolv- ing as 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.
INTRODUCTION
Modeling many-body systems has been long-standing challenge in scientific fields from classi- cal 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 new opportunity model the many-body -ody systems in an end-to-end manner. Since many-body physical systems follow many physical constraints like SE(() symmetry, pure black-box GNN models show limitations on general- ization in this scenario and symmetry-preserviee 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 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 imposed on these methods, sufficient expressive power these models is guaranteed. Unfortunately, transforming 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, provid- ing an efficient way to preserve equivariance without performing complex space transformations.

---

Under review as a conference paper at ICLR 2022
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 certain amount of tensor information ' which causes the equivariance function class EGNN can approximate to be restricted. This drawback may become serious when modeling complex dynamical scenarios such 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 al., 2021) of many-body systems. With scalarization block and a vectorization block, EVEN 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), EVEN first introduces tuple of complete basis associated with each particle pair that preserves permutation and SE(3) symmetry under global rotation and translation transformation on the many- body system. Based on this basis, the scalarization block losslessly transforms the geometric in- formation into SO(3)-invariant scalar representations. In principle, the scalar representations can be fed into any permutation iaaittioonn-vvviiiant network to implement complex nonlinear transformations. Moreover, the vectorization block could reverse the scalars back to the vector field without sacri- ficing geometric information with the complete basis. Once obtained the estimated gradient field, can predict certain state the whole dynamical trajectory of 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.
BACKGROUND
In this section, 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' the differential of many-hody system. Let X (æ1..... EN ) E Nxx be a many-body system living in R3, where N is the number of particles. We use (t(t) to denote the position of the particle i time t.
SE(3) group and Equivariance In the Euclidean space R we can consider affine transformations that preserve the distance between any two points, i.c., the isometric group SE(3). We call 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 symmetry group, it's valid to define quantities that are "equivariant" under the symmetry group. Given function f R. + R", assuming the symmetry group acts on Rm and R", then is G-equivariant if
f(gx) = gf Vz E Rm and g e G.
For SO(3) group, if = 1, i.e., the output of ff is a scalar, then the group action on RR. is the identity map, in this case ff should be -3niaiin (Thomas et al., 2018): f(gz) = f(x).
The notion of tensor field can be defined for general Riemannian manifold (see Definition 1.1.10 of (Jost & Jost, 2008)). Let { a 18 and {dx' ) be the tangent vectors and dual vectors in R respectively, and @ denotes the tensor product. Then recall the definition of tensor field for IR3 w.n.. the SO(3) group:
8z,
Definition 2.1. A (r, tensor field 0 is a multi-linear map from a collection of r vectors and s dual vectors in IR3 IR: 0(x) ff -j 924, @ .... a2s, @ dx @... 2 dx² It implies that under SO(3) coordinate transformation == {gi)) JSi.gnn the tensor field 0 transforms equivariantly; -4.1 "94,,,99, ... of 8/1::55,, where g² is the inverse of g. For example, dihedral angle is a function of position vectors rather than the position's 's norm, which are the input of EGNN. Detailed definition is given in appendix A.2.4
2

---

Under review as a conference paper at ICLR 2022
Complete Basis Construction
Centralization
Scalariation
Embedding
Edge Feature
Node Feature
Ag Aegregation
EaJ
c-C
Graph Transformer
Vi
b.
Vectorization
An overview of EVEN. For many-body system X, we first centralize the positions to preserve translation equivariance. Then we introduce tuple of edge-level complete basis Fiy to transform the geometric tensors at into SO(3)-invariant scalars. Afterwards, the scalar embeddings lij- pre-defined node features h; and edge features eij are fed to the Graph Transformer to learn edgewise embeddings myy- Finally, vectorization block transforms the edgewise emheddings into nodewise vector fields V,.
Figure 1:
Equivariant Vector field To model X (1). a natural way is to estimate differential 4X(+) apply ODE solver to integrate the differential to obtain the dynamic trajectory state 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 second- order vector field the acceleration field of Newtonian systems.
de and
Gradient field is 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 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 non- equilibrium system will finally converge to an equilibrium state. Gradient field is special case of first-order vector field.
METHODOL.OGY
Given a many-body system X, we aim at modeling vector field to predict the long-term dy- namic 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, represent the system with three key components: (1) a Scalarization block to encode the geometric tensors into SO(3)- invariant scalar representations attached to each node; (2) Graph Transformer block learn edgewise embeddings by propagating and aggregating information on the graph; and (3) a Vectorization block reverse scalar representations back to geometric tensors estimate the vector field. A brief overview of EVEN is illustrated in Figure 1. Once the vector field network (EVFN) is optimized, an Evolving block incorporated to integrate the vector field for predicting the dynamics. The translation symmetry can be easily preserved by moving the partiele system's 's to the origin (the Centralization operation in Figure 1). Permutation equivariance is automatically guaranteed for the message- pagggeggpaingggggg based Graph Transformer. We provide detailed proof about these symmetries in Appendix 1. Now we concentrate on SO(3) symmetry in the following sections.
spatial graph and construct EVFN based on
centroid at
3.1 SCALARIZATION BLOCK
The scalarization block is designed to transform geometric tensors into edgewise SO(3)-invariant scalar features by introducing novel tuple of complete basis. Given particle (1), define neighbor
3

---

Under review as a conference paper at ICLR 2022
(t)) as the particles that react with (t). Then can consider particle pair (z) 2 (t)), where (t) E (t), Suppose we take the positions of the two particles as the relevant ge- ometric information of edge (i, j). then the edgewise bij == Scalarize( (4), (A). Fig), where Scalarize is the scalarization operation under an edge-
scalars could be defined as
wise dynamical basis Fig defined below.
Equivariant basis construction For the partiele pair (a;( let a(t) = [0,())_, (0]l and
==(()= (10)-a,00 (1)
X denotes the cross product of two vectors, define
b(t) xx(() (t) x xz (t)
(t. x x (t)]l and c(t) a(t) x b(t),
(3.1)
Then we build S0(3)-equivarrin basis F.g
(a(t), c(t). In practice add small constant e to the normalization factor in case that æ; and 2j collapse. Under the condition that the matrix (a()), b((), e(()) is non-degenerate, Fij formulates complete orthonormal basis (frame) of the tangent space at æz (t). Note that this is dynamical basis w.r.t. t and the construction process of such basis is permutation- -aquuaiiiinniivviiaatt..m Since the Euclidean metric is flat, the dual basis (living in the cotangent space of æ; (Hsu, 2002)) of Fig is just its transpose: (a" (t), 6" (t), c" (t)). The 'bad' event that Fij is degenerate for all neighbors happens only when all particles are restricted to straight line, which is a measure zero set in R$, Therefore, assume Fij is non-degenerate from now on. Proof for SO(3) (equivvrrvae of Fiy 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 æ naturally owns a 'coefficient" 'scalar' representation:
(3.2) We define the process obtaining such scalars Scalarize operation. Here we demonstrate that the set of obtained coefficients (3.2) is actually a SO(3)-invariant scalar tuple. Let g E SO(3) be an arbitrary orthogonal transformation, then
(at a()), 2 b(t), Zx c(t)).
a + gæz and (a(t), c(t), + (g a(t),g b(t), ((t)).
Therefore (3.2) undergoes
( a(t), xy b(t), xp c(t)) - (gDy ga((), gad gb(t), ga gc(t)) (ax (l),, ax b(I), ox e(l)),
(3.3)
where we use the fact that g" g- I to get the last line.
is easy to prove that the Scalarize operation could transform arbitrary geometric tensors into SO(3)-invariant scalars. Taking (2,0)-tppe tensors an example, by extending the complete basis (a, b, e) through tensor product, it's easy to check that
{a 6 a, b @ b, e c, @ b, b a, a @ c, a, b)
forms an equivariant basis of the (2, O)-tyee 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, O)-ymmeetric tensor 0 (e.g., energy- momentum tensor), under the complete basis Fij (a., b, c), 0 can be expressed as:
0 =0aa a @ a + 8 b b + 0 2 c 0rh (a @ b 0 a)
+ 0 & c c & ghe c c b).
(3.4)
The scalars tuple big := {0en, gab geb,.... are the scalarization of 9 under the equivariant basis Fijs scalars are
which are 50(3)-invariant Since any nonlinear transformations acting on still -nvaaiinn, 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 com- mon geometric information of input is vectors in real-world scenarios. We define the Scalarize operation as:
Scalarize(II, D.j; Fiz) = (az; a()), 0; b((), as c(l), Vj a()), 2, b(I), my e(l)) (3.5)
4

---

Under review as a conference paper at ICLR 2022
3.2 GRAPH TRANSFORMER BLOCK
Algorithm 1: EVEN
After encoding the geometric tensors into SO(3)- invariant scalars t.j. we first embed them alone with other pre-defined node/edge attributes (h egy) into high-dimensional representations, and
Input: X (.,,,.,,,,, E RNXS,
1y E ERER, eij E R 9x
// Centralization
X Centralize(X);
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 intro- duced due to its powerful capacity in modeling those graphs with unknown topology. We pro- vide further design details in Appendix A.2.2.
// Scalarization Block for æ; E X do
for 2; E (ax;) do
Fij Equi EquiBasis(m,
Li - scalarize Cú 2y);
end
end
// Graph Transformer Block
mij = GraphTrans h;, eiz);
3.3 VECTORIZATION BLOCK
Vectorization Block
Vi V N ESeN(E)) Vectorizetonnnn Fij): Output: V;
Given the refined edgewise message Mij. the vectorization block is designed to transform these scalars back to equivariant vectors, which re- quires pairing myy with the corresponding complete basis? Fij := (a, c). More precisely, we first project myy into scalar triple {a² . xx2x²²) then define the vectorization process as:
(x x² 2²) Pairing, x'a 2² x²c.
(3.6)
We encapsulate the pairing process Vij Vectorize(m (mi)z Fiy). ). It's easy to check that the output follows the transformation rule of vectors. Finally, we aggregate all vectors Vij associated with æ; to estimate the ground-truth vector field V,-
So far, we have achieved permutation and SE(3) equivariance by employing the three blocks. The evolving block is introduced 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 (0(3))-quivarianee of EVFN can be proved from the following facts: (1) The scalarization block SO(3)-invarian because its inputs are SO(3)-invariant scalars; (2) The SO(3)- invariant variables are still (0(3)-invarian after being transformed by the graph transformer block (or any nonlinear neural network); (3) The vectorization block is -O())-eqivvarian if the input is SO(3) 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.
4 EXPERIMENTS
Our method (EVFN together with the evolving block) is general framework for modeling many- body 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).
4.1 NEWTONIAN MANY-BODY SYSTEM
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 Newto- nian forces under three settings.
Sinnee every vector is a linear combination of a hasis, the output of the vectorization block has restriction on the direction. We will discuss how our basis changes under reflection in remark A.3.
5

---

Under review as a conference paper at ICLR 2022
Partially observed system (POS). This system consists of six particles under Newton's gravittatin (radical force) but only four of them could be observed, ie.. for each trajectory, we are provided with positions X (t) E R4x3 and velocities ((t) E R4x3, The two unobserved particles act 'virtual' external field for the system.
Static external force field (SEFF). This system consists of six particles under both Newton's grav- itation and an external static force: fn = (0, O, 7), where dj > O 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 's gravitation and Lorentz-like dynamical external force field, which means there exists force field perpendicular to the current direction of each particle's ' velocity: fi(v(t) = qu(() x B, where q is positive constant that mimics the charge of particles and denotes the pseudo-vector of the external field.
Problem definition. Following (Zhuang al., 2020; Li et al., 2021), we formulate trajectory pre- diction 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 ær; (t), E (At, 2 } as the training labels and observations az (t), f E {T + At, Ti + At,...;T:) as the validation set. To evaluate the interpolation and extrap- olation capacity of all methods, the observations x; (t), E ( A T + f and a,,t) E {72 + At, T2 + 2 ,,, Ta} 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 tranforroosoo predicted trajectory with rotated input and rotated ground truth could reflect the transformation ro- bustness of method. The normalized distance between the rotated prediction with original input and the original prediction with the rotated input defines the equivariance error AEQ:
the input and output. The MSE between the
AEQ = |4, (x) - (z))l) ||,,,,,
(4.1)
where L.2 and & denote SO(3) transformations and equivariant neural networks, respectively.
Learning Framework. Inspired by (Norcliffe et al., 2020), for Newtonian system, we utilize EVEN 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:
1 Σ T1
((0)
L2( SNODE Uta to, t, Q)), to < t1
to
(4.2)
where (2tt: (tp) 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-² and the modified physical rules. In this experiment, we set 7) to 0.98, g to 1 and B to [0.5, 0.5, 0.5] T 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: Ti- = {1.5, 0.5, 0.5). TV = (2, 0.35, 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 (Köhler al., 2019) and EGNN (Satorras et al., 2021a; Further details about data generation and model implementation are provided in Appendix 1.
Results. All the baselines and EVEN implement the feature transformations in the original space. Notice that GNNs are manifestly permuration-equivariae 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 we have the following conclusions: (1) With the original input, EVFN outper- forms all other equivariant methods in the interpolation and extrapolation tasks, which demonstrates that EVEN exhibits stronger expressive power by representing geometric information losslessly
We put more discussions especially the performance of the non-equivariant GCN in Appendix A.3.1.

---

Under review as a conference paper at ICLR 2022
Table 1: Comparison three simulated Newtonian systems. Inter. and Extra. denote the interpola- tion and extrapolation task respectively. Rot. denotes the results in the rotated reference.
Setting
Method
Inter.
Extra.
Rot. Inter.
Rot. Extra.
A AEQ
GCN
0.143
1.452
0.772
2.633
18.163
GCN (Aug)
1.177
6.696
1.821
14.310
18.163
Radial Field
1.996
7.665
1.996
7.665
POS
EGNN
0.726
6.449
0.726
6.462
9.19 10
EVEN
0.138
2.502
0.137
2.428
1.23 10-4
EVEN
0.206
2.432
0.206
2.442
4.13 10-6
GCN
1.24 10-22
0.119
0.462
3.824
0.154
GCN (Aug)
0.173
1.389
0.463
3.931
0.154
Radial Field
0.701
6.032
0.701
6.032
5.23 10-7
SEFF
EGNN
0.463
4.248
0.467
4.275
5.88 10-6
EVEN
0.139
1.289
0.142
1.310
1.05 10-6
EVEN#
0.278
2,406
0.289
2.556
8.17 10-7
GCN
2.99 10-3
4.75 10
0.203
1.524
8.86 10
GCN (Aug)
8.81 10
0.740
9.07 10-2
0.760
8.86 10-2
Radial Field
9.57 10
0.804
9.37 10-2
0.804
5.10 5.10-10-7
DEFF
EGNN
1.99 1
0.213
3.22 10
0.290
8.55 10-
EVEN
1.38 10
2.43
10.-10022
1.39 10
2.48
10
1.15 10
EVEN2
1.66 10
3.13
10-2
1.81 10-3
3.24 10
6.51 10-7
(2) With the rotated input, EVEN still performs best and all equivariant networks exhibit better equivariance-preserviin capacity than GCNs even with data augmentation, implying that it dif- ficult achieve equivariance with simple data augmentation. We leave detailed analysis about different augmentation degrees in Appendix A.3.1.
4.2 MOLECULAR CONFORMATION GENERATION
In this experiment, we leverage EVEN the first-order vector field of small molecules' equi- librium distribution and generate reasonable conformations for given molecular graph. Note that in this case, the SO(3) and translation symmetry are the distribution level. rather than a specific trajectory. Previous works demonstrate that geometric information (e.g., angular potentials and ro- tatable bonds) plays critical role inducing conformation changes (Xu et al., 2021; Klicpera al., 2020), which provides natural platform to evaluate the strength of EVEN.
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.
Evaluation Tasks.
(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), 1), for this first-order statistical ensemble system, leverage score-based generative modeling framework to estimate the gradient field of atomic positions (See more details about score- score-based networks (Shi et al., 2021: Song et al., 2020) or Appendix A.2.5, A.3.2). The optimization objective of EVEN & can be summarized as:
(())= Σ Ex HXL 2x 0)| to <
< Ln. (4.3)
where X(f) J0, T] + is positive weighting function and VH, 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 al., 2020).
*Not that the division operation in the basis construction procedure will bring considerable numerical
into EVFN which increases A BQ- Therefore, tune the constant C in a(I) [m,(() [000))=,;000-- [0 The model variant EVEN* with the suitable C can achieve comparable equivariant with other baselines.

---

Under review as a conference paper at ICLR 2022
Table 2: COV and MAT scores of different approaches GEOM-QM9 and GHOM-Drugs datasets. For the COV score, the threshold is set to 0.5A for QM9 and 1.25A for Drugs
Dataset
GEOM-QMS
GEOM-Drugs
COV (%)²
MAT (A)
COV (%)
MAT (A)
Metric
Mean
Median
Mean
Median
Mean
Median
Mean
Median
RDKit
83.26
90.78
0.3447
0.2935
60.91
65.70
1.2026
1.1252
CGCF
78.05
82.48
0.4219
0.3900
53.96
57.06
1.2487
1.2247
ConfGF
88.49
94.13
0.2673
0.2685
62.15
70.93
1.1629
1.1596
EGNN 80.93
86.27
0.3832
0.3898
40.71
33.01
1.3574
1.3346
EVEN
90.21
93.14
0.2430
0.2457
88.64
97.56
0.9040
0.9023
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 IS017 dataset (Schütt et al., 2017). To keep fair comparison with the existing state-of-the- (SOTA) method ConfGF (Shi et al., 2021), we reproduce its data collection and split settings rigor- ously. Further details described in Appendix A.3.2.
Datasets.
Metrics. Given the RMSD of heavy atoms that measures the distance between generated confor- mation and the reference, Coverage (COV) and Matching (MAT) scores defined to measure the diversity and accuracy for given RMSD threshold 8 respectively.
COV(S,, S.) = [..] (R E S.RMS(RRR it) < 6, R S.),,
(4.4)
MAT(S, S.)
Σ
min RMSD R,, ft).
(4.5)
S.
RES,
RES.
where Sg and S, denote generated and reference conformations, respectively.
Implementation Details. The EVFN 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 al., 2021) to sample twice as many conformations as the reference ones from each model. We provide all hyperparameters of the score-based framework Appendix A.3.2.
Baselines. We compare EVEN 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 attempting generate conformations by learning the gradient field 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 baseline.
Results. We summarize the mean and median COV and MAT scores on two benchmarks for all methods. As shown in Table 2, EVEN achieves the best performance almost all metrics and datasets, demonstrating the effectiveness of our proposed method. Compared with ConfGF which employs 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-basee this complex distribution. EVEN also achieves significant improvement the Distributions Over Distances task, and we provide the empirical results and further discussions in Appendix A.3.2.
geometry is not sufficient to model the gradient field of
A Ablations Although the superior performance on multiple tasks verifies the effectiveness of EVEN, 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 variant of EVIN 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 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,
8

---

Under review as a conference paper at ICLR 2022
Table 3: Ablations for EVEN two datasets.
Dataset
GEOM-QM9
GEOM-Drugs
COV (%)t
MAT (A)
COV (%)
MAT (A)L
Metric
Mean
Median
Mean
Median
Mean
Median
Mean
Median
EVEN w/o Sea
88.99
94.55
0.3050
0.3066
65.67
75.63
1.1410
1.1132
EVEN w/o GT EVEN
85.21
91.18
0.3057
0.3060
70.57
81.82
1.1075
1.1004
90.21
93.14
0.2430
0.2457
88.64
97.56
0.9040
0.9023
we replace the block with the GIN network (Xu al., 2018) that is employed in ConfGF, getting the variant named EVEN 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.
5 RELATED WORK
Equivariant neural network. Existing equivariant networks with 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., 201 8b;a; Weiler & Cesa, 2019; Esteves et al,,,2000 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 et al., 2018; Köhler et al., 2019; Shi et al., 2021). However, most of these approaches (e.g.. EGNN (Satorras et al., 2021c)) abandon 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 only computed from one basis and EVFN extends the one basis (i.c., radical direction 2; 2zy) in EGNN into tuple of bases for information lossless.
the original space in computationally efficient way (Schütt
Gradient fields modeling. Gradient fields modeling 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 al., 2020) 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 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.
6 CONCLUSION AND FUTURE WORK
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 high- dimensional spaces or equivariant functions. With the proposed scalarization technique, EVEN could cooperate with any neural networks without concerns about breaking the equivariance sym- metry. Theoretical analyses and extensive empirical results verify the effectiveness of the proposed method. In the future, will investigate the performance of EVEN in large-scale many-bode sys- tems and extend the strategy to other symmetry groups or local symmetry group such as gauge group (Cohen et al., 2019a).

---

Under review as a conference paper at ICLR 2022
REFERENCES
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In NeurlPS, 2019.
Simon Axelrod and Rafael Gomez-Bombarchee GEOM: Energy-annotater molecular conformations for property prediction and molecular generation. arXiv preprint arXiv:2006.0555 0, 2020.
Alexander Bogatskiy, Brandon Anderson, Jan Offermann, Marwah Roussi, David Miller, and Risi Kondor. Lorentz group equivariant neural network for particle physics. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 199-1002 PMLR, 13-18 Jul 2020. URL https: httpss:/ppcceeeidggeee wwt.mmlr..r..... / v1 htt.//ootttttteoooo s20. html.
Paula Yurkanis Bruice. Organic chemistry 4th edition, 2000.
Ruojin Cai, Guandao Yang. Hadar Averbuch-Elor Zekun Hao, Serge J. Belongie, Noah Snavely, and Bharath Hariharan. Learning gradient fields for shape generation. In European Conference on Computer Vision, pp. 364-381, 2020.
Giuseppe Carleo and Matthias Troyer. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325)x602-6666 2017.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differ- ential equations. arXiv preprint arXiv:1806.07366, 2018.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International confer- ence on machine learning, pp. 2990-2999. PMLR, 2016.
Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling, Gauge equivariant convolu- tional networks and the icosahedral cnn. In International Conference on Machine Learning, pp- 1321-1330. PMLR, 2019a.
Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical cnns. arXiv preprint 10130000.000.,... 2018.
Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cnns on homoge- neous spaces. Advances In Neural Information Processing Systems 32 (Nips 2019), 32(CONF), 2019b.
John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26, 1980.
Carlos Esteves, Ameesh Makadia, and Kostas Danillidis. Spin-weighted spherical CNNs. In NeurIPS, 2020.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for equivariance lie groups on arbitrary continuous data. In International Conference Machine Learning, pp. 3165-3176. PMLR, 2020.
Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. (3)-transformmrs: 3d roto- translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Fabian B Fuchs, Edward Wagstaff, Justas Dauparas, and Ingmar Posner. Transformers. arXiv preprint 102. 13419, 2021.
Iterative SE(3)-
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Berhaar Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Sam Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In NeurIPS, 2019.
Elton P Hsu. Stochastic analysis on manifolds. Number 38. American Mathematical Soc., 2002.
10

---

Under review as a conference paper at ICLR 2022
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups. In International Confer- on Machine Learning, pp. 4533-4543. PMLR, 2021.
Jürgen Jost and Jeurgen Jost. Riemannian geometry and geometric analysis, volume 42005.
Springer, 2008.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint TrXiv::442.6996 2014.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: 2016.
Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. arXiv preprint 2003v:2003.03) 1233, 2020.
Shoshichi Kobayashi. Foundations of differential geometry vol I (new york: Interscience) kobayashi S and nomizu 1969. Foundations of differential geometry, 2, 1963.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi- body systems with symmetric energies. arXiv preprint 1:X10:111..00753, 2019.
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-Gordan nets: fully Fourier space spher- ical convolutional neural network. In NewrIPS, 2018.
Greg Landrum. Rdkit: software suite for cheminformatics, computational chemistry, and predic- tive modeling, 2013.
Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky T. Q. Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 3870-3882, PMLR, 26-28 Aug 2020. URL https: http///ooeeddddeeee iigs mlr www.s.pprssss... v1 08 12011111000 html.
Ziming Li. Bohan Wang. Qi Meng, Wei Chen, Max Tegmark, and Tie- Yan Liu. Machine-learning non-conservative dynamics for new-physics detection. arXiv preprint arXiv:2106.0000 2021.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Liò. On second order behaviour in augmented neural odes. arXiv preprint arXiv:2006.07220 072200, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
David W Romero, Erik J Bekkers, Jakub M Tomczak, and Mark Hoogendoorn. Attentive group equivariant convolutional networks. In ICML, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows for molecule generation 3d. arXiv preprint arXiv:2105.0900 000119666, 2021a.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- works, 2021b.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) cquivariant graph neural net- works. arXiv preprint arXhv:2102.09844, 2021c.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: continuous- convolutional neural network for modeling quantum interactions. In NewrIPS, 2017.
Kristof T Schütt, Huziel E Sauceda, P- Kindermans, Alexandre Tkatchenko. and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
11

---

Under review as a conference paper at ICLR 2022
A.W Senior, R Evans, J Jumper, and et al. Improved protein structure prediction using potentials from deep learning. Nature, 2020.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular confor- mation generation. arXiv preprini arXiv:2 10X.::2005.0390,,, 2021.
Yunsheng Shi, Zhengjie Huang, Wenjin Wang, Hui Zhong, Shikun Feng, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. arXiv preprint arXiv:2009.03509 2020.
Gregor NC Simm and José Miguel Heráández-Lobato A generative model for molecular distance geometry. arXiv preprint arXiv:/9099 / 1459, 2019.
Jascha bohl-Dicksteen, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibriun thermodynamics.. In International Conference on Machine Learn- ing. pp. 2256-2265 PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Proceedings of the 33rd Annual Conference on Neural Information Processing Systems, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems, volume 33, pp. 12438-12448, 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-ane consslationequivaiinnn neural networks for 3d point clouds. arXiv preprint arXiv:1802.082/19 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, 2017.
Maurice Weiler and Gabriele Cesa. General E(2)-equivariant steerable steerbal CNNs. In NeurlPS, 2019.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco S Cohen, 3D stecrable CNNs: Learning rotationally equivariant features in volumetric data. In NeurlPS, 2018a.
Maurice Weiler, Fred A Hamprecht, and Martin Storath. Learning steerable filters for rotation equivariant CNNs. In CVPR, 2018b.
Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In CVPR, 2017.
Jiaxiang Wu, Tao Shen, Haidong Lan, Yatao Bian, and Junzhou Huang. Se (3)-equivariant energy- based models for end-to-end protein folding. bioRxiv, 2021.
Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint 1rXiv:1810.0060 2018
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2102. 10240, 2021.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and E Weinan. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Physical review letters, 123(14) 2018.
Juntang Zhuang, Nicha Dvomek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Duncan. Adaptive checkpoint adjoint method for gradient estimation in neural ode. In International Conference on Machine Learning, pp. 11639-116e 11649. PMLR, 2020.
12

---

Under review as a conference paper at ICLR 2022
A APPENDIX
A. PROOF AND DISCUSSION ON SCALARIZATION AND VECTORIZATION
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 E SO(3), then under the action of g. the positions of the many-body system X transform equivariantly:
(.)( nn (1) & (gan (ga....)..gg .0r. (1)).
Then from the definition of (t), we know that
a(t) - ga(().
For b(t), since
(gxx(t)) (gz;((t) = (o" =
2y(())
(A.1)
= g(, (t) x xz (t))
(A.2)
where we have used g - g" for orthogonal matrix g to get the last line. Therefore, b(t) 2 gb(t). Applying (A.1) once again, we have e(t) 2 ge(t).
Frame bundle and scalarization technique A frame at x E M is a linear isomorphism u from Rd the tangent space at x:T,M. We F(M)z to denote the space of all frames at Then GL(d, IR) acts on F(M), by
R 2 R - T.M.
Then the frame bundle
F(M) UsEM ((M),
can be made into a differential manifold. From the principle bundle point of view, each differential manifold M is quotient of its frame bundle F 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 F(M) is a reference frame located at æ == (u) E M. On the other hand, Rd can be seen as d-dimensional differential manifold with global coordinates chart and SO(3) is the structure- preserving group of the Euclidean metric with a fixed orientation.
Following Hsu (2002), let {ei, 1 i < d} be the canonical basis of IRd, and {ee) the corresponding dual basis. At each frame u, the vectors Y = ne; form a basis of T. M. Let {Y') be the dual frame of T M, then (r,s)-tensor 0 can be expressed
Yn
2 Y, ② yi 8....... SY.
The scalarization of f at u is
((u) = 0 j.
ei
2 cde
(A.3)
Through scalarization, a tensor field 0 becomes an ordinary vector space valued function on F(M):
Ø F F(M) + R R
Therefore, geometric operations such covariant derivative and tensor product 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 frame bundle (R3) (A.3) and the SO(3)-invariant scalars tuple obtained by the scalarization block (3.4) under an equivariant basis.
Proof. Since to SO(3) global group action acting on the orthonormal frame bundle O ). A frame u E ( (u) I can be transported to another point in R by translation. Therefore, neglect the
working in 3 with fixed orientation, the GL R) group action is reduced
13

---

Under review as a conference paper at ICLR 2022
origin of the frame in the proof. Let ue (e1, E2, e3) be equivariant basis of R then given scalars tuple {011m)t)) ). we construct a vector-valued function on ( by:
8
(u)
9 Pi,1, 8
where g is the SO(3) transformation from ue to another frame u. Moreover, ( is SO(3)-
equivariant, since
(gu) = gf(w),
where the g on the right side means the usual extension of the action of SO(3) from R the tensor space Rer We have constructed the (r,0) tensor field from the scalars tuple. It's easy to check that (e (u) induces a (r,0) tensor field on R2 by following the definition of (2.1).
From scalarization 0f.....f f (u) on ( ) to scalarization is obvious. Note that any equivariant basis We is also a point (RR)), therefore the scalars tuple is just the values of (11.... at Ue:
P @fr.... (we).
For general (r,s)-type tensors, the proof is the same by adding the dual basis of the
Remark A.3. For molecular structures, one common obstruction for distance-based modeling (Shi al., 2021) the chirality problem (Bruice, 2000). Therefore it's meaningful to investigate how the equivariant basis (3.1) transforms under reflection - =z. Notice that a - aa, and
b= x x b. implies that - cc In conclusion, the orientation of the equivariant basis remains unchanged under reflection.
A.2 METHODOLOGY
A.2.1
PERMUTATION AND TRANSLATION EQUIVARIANCE
For a many-body system (1) = (t), (t), the centroid C(t) is defined by
xn c(t)= = x1 (1(() + xz (t)
n
Translating the reference by h, then
X(1) h c(t) th.
Therefore, recentering the reference's 's origin to the centroid at the input's time / = 0, we have
X(t) c(0)
translation
h t=
X c(0).
That is, the system is translation-invarian under the recentered reference if the translation is done the input's time f 0, which is exactly the scenario considered in predicting the future trajectory or state. Note that although the gradient field is supposed to be cransltion-invariinn, the state of the many-body system after integrating should be inisslatio-equivarian Therefore, add c(0) back for the output of the evolving block.
To emphasis the permutation symmetry, we quote the Kolmogorov- representation theorem: if ( ...) is permutation invariant multivariate continuous function, then
f( ,t) = g of (x:)).
(A.4)
The crucial point is the function shared among all the points, which exactly fits the definition of the so-called message-passing scheme. For many-body system X, let = (U).... u,, be its vector field, then v; R corresponds the equivariant vector for æ; [#1,... .2,,...,200 Fr). Denote EVEN with parameters 0 by क,, = [ - then
", (1)
(1),
14

---

Under review as a conference paper at ICLR 2022
for fixed particle z;. Suppose (d1,..., (x) (neighbors of æ;)) indicate indexes of particles which have interaction with 2,, then obviously I I n 1. By (A.4), (x) is an aggregation of message from æ; 's neighbors, therefore we have:
- (t) k
(x((()),))
Σ j=1
(A.5)
and & is a -quivvaaaa network with vector output. Note that (A.5) performs aggregation at the level of vectors, therefore choose g (A.4) to be the arithmetic mean to preserve SO(3)
symmetry. A shared neural network á guarantees the permutation equivariance of EVEN.
A.2.2 GRAPH TRANSFORMER BLOCK
For many-body system X, we first represent it as spatial graph and utilize an attention srrention-bbsed to learn the SO(33)invariant edgewise embeddings my from the graph. Considering that there does not exist any graph topology in most real-world scenarios, we introduce the attention mechanism due its powerful capacity in learning the correlations between inter-instances (Vaswani et al., 2017). Now discuss the workflow of the GTB block. Note that here we omit the time index of the SO(3)-invarian) scalars for the simplicity of description.
Feature embedding Given geometric scalars lij. node features h; and edge features maj. GTB first embeds them into high-dimensional representations:
h; = MLP(h;):
(A.6)
Gj MLP((e;);;
(A.7)
4j = Fourier(4:j).
(A.8)
Cij Cij taj.
(A.9)
where MLP denotes fully connected network and Fourier denotes Fourier transformation with a tuple of learnable frequencies.
Transformer block The overall architecture of our GTB block inspired by (Shi et al., 2020).
For each GTB block, we first compute the edge-wise message with o, (A.7). then leverage the message embeddings and node embeddings with transformer encoder-like architecture refine the node embeddings. After that, we update edgewise messages with residue block (A.13). The whole pipeline is expressed as (A.5).
mij (n hi..
(A.10)
== kij = dx dn(hs, (lssssmms)) mis), Vij d (m
(A.11)
(g:n kiz)
Oij EPEN() (gis kig))
(A.12)
M, LayerNorm( Σ aijoiy)
(A.13)
JEN(1)
h; = (h,,
(A.14)
h h, Laye
(A.15)
My = 0z (l,, (in) + maj:
(A.16)
where d1, Φy. 6x. 6v, Φn and oh are fully connected networks. di and M, denote the attention weights and the refined nodewise embeddings, respectively. LayerNorm refers to the normalization layer adopted in (Vaswani et al., 2017).
A.2.3 EQUIVARIANCE OF EVEN
We need the following lemma:
Lemma A.4. Suppose h an invariant function, fi...., fx. are arbitrary nonlinear functions. Then, the composition ft. .... fú h an invariant function.
15

---

Under review as a conference paper at ICLR 2022
Proof. Consider the group action g acting on R2, then f is invariant means that o g(z) == f(gz) = f(z). We have
fa 0..00 fi g fa o...o fi,
hence the composition fa o ..... fi is invariant.
Let fi be the scalarization block. Because the output of fi are scalars (from (3.3)) and SO(3)- scalars must be invariant under SO(3) group action, each graph- traphff-rrnffrrmer layer by fi, i (2...., kj. then by Lemma A.4, we can conclude that the output of the graph-transformer block is also invariant.
conclude that is invariant. Denote
Finally, from the standard fact that scalars multiply vectors yield equivariant vectors, we conclude that the output of the vectorization block is -O(3)-qquivarian. Combining with the translation equivariance in 1. our model is SE(3)-equivarian.
A.2.4 DIHEDRAL ANGLE
For given node 2y with three neighbors æ; zy and æg. The dihedral angle of the plane spanned by (2) - 2j, 2 xy) and the plane spanned by (x) - Zz zz Ix)) is given by the inner product of normal vectors of the two planes:
(0) x;)) x (x)) - 1.))
[[x: x3) x (x) zzy)
and
(x) x. (z) - 2²k)
(( œ;) x (x) - xx)
Therefore, the dihedral angle is a function of the position vectors wy Ty. By and ægs which cannot be determined purely by the positions' norm: || =;||.||æ;- æ, and | -xxll-
However, by the complete equivariant basis, the position vectors are fully expressed (only the co- ordinates are transformed from the natural coordinate system to the equivariant coordinate system), then all functions of the position vectors can also be expressed.
A.2.5 EVOLVING BLOCK
According 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) (t) e R 1,2, n3 whose underlying physical rules are unknown, let v; denote the velocity of particle æ.. We attempt predict its position or trajectory by integrating second-order equivariant vector field. In this case, the EVFN network & is implemented for modeling the acceleration vector é(f). More precisely, the hidden differential unit of NODE (Chen et al. 22))8) has the following form:
(t) b(t)
v(t)
(())))
(A.17)
For the i-th particle,
(+) = (x(()))))()) = Σ (t),
(A.18)
The ordinary differential equation (A.17) is then solved by block-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 physical- allowable molecule conformations are located in an equilibrium state determined by the energy function. Suppose the forward process from equilibrium to conlequilirrium of the system satisfies:
dx f(x (1), tydt + g())W) OSLST,
16

---

Under review as a conference paper at ICLR 2022
where W2 is the Brownian motion and the initial state (() follows an unknown equilibrium distri- bution Po- Denote the marginal distribution at time t by pt. then
Pr (X) exp(-
so the Hamiltonian function H, at time is entangled with po- According to the Liouville equation (the probability flow in Song et al. (2020)), the reverse evolving process satisfies the following ODE:
1 2 gT -
dX -) f( t), T t)dt
V
The gradient field of the Hamiltonian function are also called the force field. Therefore the NODE evolving block has the following form:
()-
= (( t)), OStST.
and for the i-th particle, following form:
&
+ oT (T- 1 - g
(T- (A.19)
In this case, we will use the PC scheme proposed by (Song al., 2020) for correcting the numerical integration error. Note that the vector-valued function f(r, t) and the scalar function g(t) is prior knowledge and set 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 al., 2020), where f 0 and
ddo2 dt
g(t)
A.2.6 NEURAL ODE AS A CONTINUOUS LIMIT
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 seen as taking the continuous limit of the residual network the depth goes to infinity. Consider residual network where all hidden layers have the same dimension, denote the neural network's parameters at t-th layer by O., then
h++1 = J1 + f(,,, 0e), t e N Taking the discrete layer index to its continuous limit, we get
dh(t) f (h((t, 0.), EER
dt
The discrete back- proppapppantioon method also has continuous limit: the adjoint method Chen et al. (2018); Norcliffe et al. (2020).
A.2.7 SCALABILITY
Compared to EGNN, the extra cost of EVEN is to calculate the values of the scalars in lij, Vi,j. Compared to the computational cost of back- sock-propagation
the neural networks, the cost of calcu- lating the O(3 * 3 * N² (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 sealars only brings 9.6% extra real-time computational cost and 17.4% extra memory cost.
A.3 EXPERIMENT
A.3.1 NEWTONIAN MANY-BODY SYSTEM
Partially observed system This system consists of six particles under Newton's gravitation but only four them could be observed, i.e. for each trajectory we are provided with positions X(t) 1443 and velocities (tt E 1x The time evolution of the particles is given by
5; Σ
mm
[z
C;
xz 3
23
1 ii5 4.
(A.20)
JE(L 1.++1...66
17

---

Under review as a conference paper at ICLR 2022
Gravity field This system consists of six particles under both the mutual newton's 'oon's gravitation and an external static force of the form: fg = (0, 0, 7). The time evolution of the particles are given by
my 2; Dj x; 3 13 + fn 1sis6. 6.
(A.21)
JE(
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., fi()) = qu 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:
Σ JE(L 1.
mi ||,,
zjj 2 + fi(ss, 15is3.
(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 and the modified physical rules. The trajectory points are uniformly sampled with At = 10 -3 for the POS dataset or Åt tt=3 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 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 for all models and adjust the hidden dimensions of each model separately 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 L2 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 EVEN 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 hard to jointly optimize the neural ODE and the SE(3)- Transformer model by directly combining the SE(3)-Transfomm achieves an extremely poor performance compared with EVEN. To the best of Transformer and neural
models together. And the empirical results in Table 4 show that
our knowledge, there is no previous work studying the combination of (E(3)-ETTann(tn ODEs. We would like to further explore this problem in future work. For the results in Table 4, all Srans-Transfommer 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 Transs-rrnse (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, 2,3 and channels per degree. We find it would bring much more parameters (i.e., 118k) and computational efforts, but only marginal improvement (c.g., 0.405 vs. 0.422 for MSE on the POS dataset).
Data Augmentation For all Newtonian many-body systems, propose to force 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 We visualize four MSE metrics of GCN (Aug) with different augmentation
a random probability probabilities on the POS dataset in Figure As shown in Figure 2, with the original input, the interpolation and extrapolation capacity of GCN is damaged as 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.
18

---

Under review as a conference paper at ICLR 2022
Table 4: Comparison with SE(3)-Transforme
Setting
Method
Inter.
Extra.
Rot. Inter.
Rot. Extra.
A FQ
(B(3))Trrnsformes
0.422
4.555
0.442
4.559
7.82 10
POS
EVFN
0.138
2.502
0.137
2.428
1.23 10
(oan)-Transf
12.623
53.358
12.623
53.258
3.67
10
SEFF
EVFN
0.139
1.289
0.142
1.310
1.05
10
SE(3)-Transfed Transforerr
9.119
37.313
9.119
37.312
3.31
10
DEFF
EVFN
1.38 10
2.43 10
1.39 10
2.48 10
1.15
10
1.2
Inter.
1.0
MSE
0.8
MSE
4-
0.6
3-
0.4
2
Extra.
0.2
0.0
0.2
0.4
0.6
0.8
0.0
0.6 Augmentation probability
0.2
0.4
0.8
Augmentation probability
(a)
(b)
2.50
Rot. Inter:
Rot. Extra.
2.00
20
MSE
1.75
MSE
15
1.50
1.25
10
1.00
0.75
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
Augmentation probability
Augmentation probability
(c)
(d)
Figure 2: Results of GCN with various data augmentation degrees on the POS dataset
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 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 10-3,5 5 10-1,3 10-"), corresponding to 150, 1000, 3000 and 5000 training labels, respectively. The time interval of the test set remains the same as before (i.c., 10 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 GCN- based 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 cquivariant 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 enror of EVEN (Table 1, col 7) is caused by numerical emrors introduced by the basis construction procedure and the neural
19

---

Under review as a conference paper at ICLR 2022
GCN
EVEN
0.19
0.18
0.17
Inter.
Extra
2.5
0.15
0.14
0.13
0.11
1.0
150
1000
3000
5000
150
1000
3000
5000
Number of trajectories (a)
Number of trajectories
(b)
3.5
30
3.0
25
Rot. Inter.
2.5
Rot. Extra.
20
2.0
15-
10
0.5
0.0
150
1000
3000
5000
150
1000
3000
5000
Number of trajectories
Number of trajectories (d)
(c)
Figure 3: Results for different number of training trajectories on the POS dataset
Table 5: The impact different normalization constants. SE(3-T denotes the SE(3)- Transformer
model.
Setting
Method
C
Inter.
Extra.
Rot. Inter.
Rot. Extra.
AEq
EGNN
(())
0.726
6.449
0.726
6.462
9.19 10-7
SE(3)-T
0.422
4.555
0.442
4.559
7.82 10-4
EVEN 0(10~))
0.138
2.502
0.137
2.428
1.23 10-4
POS
EVEN O(1)
0.206
2.432
0.206
2.442
4.13 10-6
EVEN
0(10)
0.608
2.344
0.608
2.344
3.60 10
EGNN
0(1)
0.463
4.248
0.467
4.275
5.88 10-6
SE(3)-T
12.623
53,358
12.623
53.258
3.67 10-5
EVEN 0(10-))
0.139
1.289
0.142
1.310
1.05 10-6
SEFF
EVEN O(1)
0.213
1.888
0.213
1.888
3.38 10
EVEN
O(10)
0.498
4.417
0.496
4.414
5.41 10
EGNN
O(1)
1.99 10-2
0.213
3.22 10-2
0.290
8.55 10-7
SE(3)-T
9.119
37.313
9.119
37.312
-
EVEN 0(10~))
1.38 10
2.43 10
1.39 10
10-8.10-2
1.15 10
DEFF
EVEN 0(1)
3.17 1.88 10-3
1.88
10-3
8029-10-2
2.11 10-6
EVEN O(10)
1.66 10
3.13-10-2
1.81 10-3
3.24 10
6.51 -0
ODE training framework. More precisely, to build the equivariant basis, we incorporated the nor- Φ
malization operation || T | (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 shown in Table 5, with some suitable normalization constants, EVEN can achieve comparable or even better equivariant emoor than other equivariant models, validating the equivariance of the pro-
(Fuchs et al., 2020) and EGNN (Satorras al., 2021b). As
20

---

Under review as a conference paper at ICLR 2022
Input Graph
Ref.
Conf.
Ours
ConfGF
Figure 4: Visualizations of generated conformations. For each molecule randomly selected from Drugg dataset, we sample multiple conformations and show the best-aligned ones with the reference ones.
posed method, In most cases, EVEN with larger normalization constants can still outperform other baselines significantly.
A.3.2 MOLECULAR CONFORMATION GENERATION
Dataset For each dataset, 40, 000 molecules randomly drawn and 5 most likely conforma- tions (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 formations in the test set for GEOM-QM9 and GEOM-Drugs datasets, respectively. The distances distributions task are evaluated the ISO17 dataset, where we follow the setup in (Simm & Hemández- 2019).
Implementation Details Besides the geometric input, we feed the node type, edge type and rel- ative distances 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 {ar)=_ be positive geometric progression scheme with common ratio, we split the noise range into levels. For the reverse process, we find the performance difference between the ALD sampler and the PC sampler is marginal,
do compare this point with quantitative results. We choose the PC sampler as our evolving block and respectively set the itera- tion steps of predictor and corrector to 10 and 100. The sample step size 7] chosen according to (Song & Ermon, 2020). We keep all hyper-parameters mentioned in the forward and reverse process the same (Shi al., 2021). The results reported in Table 2 are copied from (Shi al., 2021) considering that we rigorously evaluate EVEN on the same benchmark and data split setting.
Conformation Generation Here we introduce the calculation equation of RMSD:
RMSD(R, Rt) min
||
(A.23)
where n denotes the number of heavy atoms.
We visualize several conformations in the Drugs dataset in Figure 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 Distances
To evaluate the distribution of the generated conformations, uti- lize 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
21

---

Under review as a conference paper at ICLR 2022
Table 6: Accuracy of the distributions over distances generated by different approaches compared to the ground-truth.
Single
Pair
All
Method
Mean
Median
Mean
Median
Mean
Median
RDKit
3.4513
3.1602
3.8452
3.6287
4.0866
3.7519
CGCF
0.4490
0.1786
0.5509
0.2734
0.8703
0.4447
ConfGF
0,3684
0.2358
0.4582
0.3206
0.6091
0,4240
EVEN 0.1317
0.0420
0.1787
0.0695
0.3185
0.1142
outperforms the previous SOTA (ConfGF), demonstrating the strong capacity of the proposed model modeling molecular dynamics data. particular, EVFN reduces the MMD by magnitude in both Single-median and Pair-Median metrics.
22


Reviews:
[{"review_id": "e2IgawSKqoZ", "reviewer": "Reviewer_gmFP", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Transforming particle interaction into a rotationally invariant reference frames is a common approach in molecular mechanics, usually achieved with inner coordinates such as distances, angles, dihedral angles etc. These have the disadvantage that some coordinate information is lost, however they fulfill all equivariance / invariance properties w.r.t. rotation, reflection, translation and permutation in local reference frames. \n\nThe authors claim that \"EVFN is able to represent tensor information losslessly\". This is also true for previous equivariant approaches such as TFN, EGNN, PaiNN or NequIP, as long as no distance cutoff is used (see for example Dym and Maron 2020). However, in practice the distance cutoff is essential not only for computational efficiency but also to facilitate generalization across local environments. This is not possible in the proposed approach:\nThe EVFN transform is rotationally equivariant and translational invariant only with respect to global transforms. This has serious consequences for its capability to generalize, in particular for larger systems. As an extreme example, take two molecules (or clouds of Newtonian particles) at a large enough distance such that the inter-molecular interactions are negligible. Since the proposed coordinate is not translational invariant, the pairwise representations are not invariant to rotating the two molecules separately. Therefore, the proposed representation is not suitable to model general many-particle systems.\n\nDespite these flaws, the method shows good performance on the toy systems and conformer search. A reason might be that for small systems generalization is not the deciding factor, and the baseline methods do not retain the necessary information. To better judge this trade-off, it would be helpful to apply the proposed approach to the related task of predicting molecular force fields (e.g. on the MD17 dataset where a lot of baselines are available for comparison).\n\nFurther comments:\n- Centering the system prevents application of the method to a box with periodic boundary conditions. Are there any ideas how to overcome this?\n- The background section could be made more accessible for readers not familiar with group theory and tensor algebra, since this part is not really necessary to understand the method and proofs in Eq. 3.1-3.3.\n- Sec 4.2: the ISO17 dataset was introduced by Schütt et al (2017)\n- Why are the equivariance errors in Table 1 slightly higher for EVFN than for the other equivariant methods?", "summary_of_the_paper": "The paper proposes an rotationally equivariant neural network based on a transform of particle pairs into a rotationally invariant reference frame. This is applied to trajectory prediction of many-body particle systems and molecular conformer search.", "summary_of_the_review": "The proposed transformation has the flaws outlined above which might prevent generalization across local environments. More experiments are required to judge how severe these issues are in practice. Therefore, I can not recommend acceptance at this stage.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Dia_fiq_qF8", "reviewer": "Reviewer_YiG5", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n1. Well written and easy to understand.\n2. The construction of the equivariant basis seems interesting.\n\n## Weakness\n1. The connection between `EVFN` and `EGNN` is not well discussed. Essentially, `EGNN`  takes $||x_i – x_j||^2$ as the geometric scalar and apply $\\phi(\\cdot)$ to obtain the edgewise message  ( see Eq(3)  in [1]). After obtaining edgewise message, `EGNN`  uses $\\phi_x$ to project $m_{ij}$ to a scalar and then vectorizes it by $x_i – x_j$ in Eq(4) in [1]. The key differences of `EVFN` are 1). Extending the equivariant basis from one dimension to three dimensions. 2). Replacing the `MLP` model $\\phi_e$ as well as the aggregation operation in Eq(5) and Eq(6) in [1] with an existing `GraphTransformer`  block. In this vein, `EVFN`  can be viewed as the extension of `EGNN` with the new equivariant bases.\n\n2. The motivation behind the equivariant basis construction is not well elaborated.  In the multi-body system modeling, the equivariant basis $\\frac{x_i -x_j}{||x_i-x_j||}$ has its physical meaning which indicates the force directions. However,  the benefit of using the outer product of two coordinates to construct the equivariant basis is not clear, since there is no explicit physical meaning of the outer product of two coordinates. It’s better to give more in-depth discussions about the choice of such equivariant bases here. \n\n3. There is no ablation study on the new equivariant basis and the ` GraphTransformer ` on multi-body system modeling.  As mentioned before,  the new-added equivariant bases have no explicit physical meaning of modeling the multi-body system, the main impact factor of the performance improvement of `EVFN` looks vague. \nOn the other hand, the ablation study on molecular conformation generation even implies that the contribution of ` GraphTransformer ` is slightly larger than that of the new equivariant basis (`EVFN w/o GT` gains the larger performance drop on 5 out of 8 metrics in Table 3. ).  If the performance improvement comes from `GraphTransformer`, the overall contribution of this paper is limited. Therefore, it is crucial to conduct the ablation study of the different components of `EVFN` to show the contributions of the new equivariant basis on multi-body system modeling. For example, apply the new equivariant bases on the `EGNN`  framework. \n\n4. The number of training samples used in the multi-body modeling experiment is not reported. Meanwhile, there is no analysis of the impact of the number of training samples for the different baselines.\n\n5. From Table 1, the result of `GCN`  surpasses many sophisticated baselines in the interpolation and extrapolation tasks, such as `EGNN`.   This is quite weird. Why?  \n\n\n[1] E(n) Equivariant Graph Neural Networks\n", "summary_of_the_paper": "This paper proposes a new model: Equivariant Vector Field Network (`EVFN`), which aims to solve the multi-body system modeling problem. Specifically, `EVFN`\n introduces a new kind of equivariant basis which considers more interactions between two particles and employs a graph transformer architecture to learn the edge-wise embeddings.\n", "summary_of_the_review": "Overall, I think this paper is not ready to be accepted at this time. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BeURY1MgazR", "reviewer": "Reviewer_V2nd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#### Strengths:\n- The paper is well organized and well written (besides some non-standard use of English). \n- The mathematical presentation is fairly clear.\n\n#### Weaknesses:\nIt doesn't seem like the model produces equivariant representations. I agree that the input to the graph transformer model is equivariant to SE(3) transformations, but without some kind of constraint, I don't see how the output could be. Moreover, the empirical evidence does not support this claim (Table 1, col 7).\n\n#### Questions:\n- “So far, we have achieved permutation and SE(3) equivariance” (p.5) isn’t it SE(3) invariance? -- the initial centering giving translation invariance, and the scalarization block giving SO(3) invariance\n- Did the authors experiment with other kinds of SO(3)-equivariant frames in the scalarization block? I wonder if computing the frame from the neighborhood (e.g., Darboux frames) would improve computational efficiency and/or performance.\n\n#### Possible typos:\n- “without concerning about breaking” → without concern about breaking (p. 9)\n- “verify the effectiveness of the propose method” → verify the effectiveness of the proposed method (p. 9)\n- “The trajectory points are uniformly sampleed” →  The trajectory points are uniformly sampled (p. 6)\n- “We sample 1, 30 and 30 trajectories” → We sample 1, ? and 30 trajectories (p. 6)\n- “(GNN), which is superior to model” (p. 1)\n- “‘Vector filed’” → ‘vector field’ (p.2)\n- “We use x(t) to denote” → We use x_i(t) (p.2)\n- “Then each point of” → Then each point (p.2)\n- “Let {Y^i} be the dual frame of” → … be the dual frame on(?)\n- “‘gradient filed’” → gradient field\n- “As shown in Table 1, with the original input, EVFN outperforms all other equivariant methods in the interpolation and extrapolation tasks” (p. 7) -- that’s not what table 1 shows.\n- “Fij ((a(t), b(t), c(t)))” (p. 4) -- extra parentheses\n- Should there be a time index in A5-A8", "summary_of_the_paper": "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. ", "summary_of_the_review": "The main claim of the paper is that there is no information lost in the representations learned. As far as I can tell, this claim is not supported by the theory or experiments. Moreover, it does not seem that the learned representations are equivariant.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "2KF8c8FJjqs", "reviewer": "Reviewer_BR6H", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Initial Recommendation: Reject\n\nReason: In my view, in current format, the weaknesses outweigh the strengths of the paper. Please see details below.\n\nStrengths\n1. The idea using a scalarization block, a graph transformer and a vectorization block is interesting\n2. The results on the synthetic task and the molecular conformer generation task are promising\n\nWeaknesses:\n1. For me, the math appears to be imprecise/ wrong at many places which makes it a difficult read - for example, starting in  Def 2.1 shouldn't it  be r vectors and s covectors? Also in the line, the group actions from SO(3) acting should involve the dual representations?\n2. For me, the authors make claims without sufficient proofs - for e.g. on page 5 - the bolded line doesn't have a proof - can you please show the same, especially given that the basis associated with ij pair is different. Also the proof for Proposition A.2 is not easy to follow or check.\n3. Scalability Concerns - In the tensor example above eq 3.4, the authors provide an example for the s=0 case (2,0) - and this in $\\mathbb{R}^3$ alone is of size 9 (3x3) - and this needs to be saved/ computed for each i,j pair. The authors also note that in their case, they only use the (1,0) type tensors. Please add scalability results especially against E(N) GNN which the authors compare against.\n4. Appears to be missing some relevant work - for e.g. [1], etc. Moreover the authors only appear to compare their model with the computationally efficient models like E(N) GNN - while ignoring models like LieConv[2] , LieTransformer [3], etc which can be used for the same task (Newtonian Many-Body System task). \n\nReferences:\n1. Anderson, Brandon, Truong-Son Hy, and Risi Kondor. \"Cormorant: Covariant molecular neural networks.\" arXiv preprint arXiv:1906.04015 (2019).\n2. Finzi, Marc, et al. \"Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data.\" International Conference on Machine Learning. PMLR, 2020.\n3. Hutchinson, Michael J., et al. \"Lietransformer: equivariant self-attention for lie groups.\" International Conference on Machine Learning. PMLR, 2021.", "summary_of_the_paper": "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.\n", "summary_of_the_review": "Overall, there are certain merits for the proposed architectures, but the paper is hard to read due to some imprecise math, lack of clarity, proofs, etc. I initially suggest rejection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a symmetry-informed neural network for modelling many-body systems. The network is empirically evaluated in the tasks of predicting Newtonian trajectories and molecular conformations.

All four reviewers are critical of the paper and recommend rejection (one weak, three strong). The reviews have flagged weaknesses and quality issues with several aspects of the submission, including the proposed methodology, the novelty of the contribution, and the clarity of the presentation. Although detailed clarifications were provided by the authors, most of the reviewers' concerns remain, and the consensus among reviewers remains to reject the paper.

Consequently, the current version of the paper does not appear to meet the quality standards for acceptance to ICLR.
Decision: Reject