E QUIVARIANT VECTOR FIELD NETWORK FOR MANYBODY SYSTEM MODELING
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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; 2019; 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, 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.
Xu et al.
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.
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.
B ACKGROUND
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 =
( x 1 , . .
, x N ) ∈ R N × 3 be a many-body system living in R 3 , where N is the number of particles.
We use x i ( t ) to denote the position of the particle i at time t .
SE(3) group and Equivariance
In the Euclidean space R 3 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 : R m → R n , assuming the symmetry group G acts on R m and R n , then f is G - equivariant if
[Equation]
For SO(3) group, if n = 1 , i.e., the output of f is a scalar, then the group action on R 1 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.
of (Jost & Jost, 2008)).
Let { ∂∂x i } 3 i =1 and { dx i } 3 i =1 be the tangent vectors and dual vectors in Rrespectively, and ⊗ denotes the tensor product.
Then recall the definition of tensor field for R 3 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 R 3 to R :
[Equation]
It implies that under SO(3) coordinate transformation g :=
{ g ij } 1 ≤ i,j ≤ n , the tensor field θ trans- forms equivariantly:
[Equation]
where g T is the inverse of g .
N o d e
F e a t u r e
E d g e
F e a t u r e
𝑿
∈
ℝ
𝑁× 𝑣 𝑖 𝑗 𝒂 𝒃
𝒄 ℱ
𝑖𝑗
Complete Basis Construction
S c a l a r i z a t
i o n
Graph Transformer 𝒙 𝑖 ,𝒙 𝑗
E m b e dd i n g 𝑡 𝑖𝑗 𝑒 𝑖𝑗 ℎ 𝑖
Vectorization 𝑚 𝑖𝑗 ⋮ 𝑚 𝑖 ⋮ 𝒂 𝒃 𝒄 𝜙 ∙ ∙
∙
A gg r e g a t i o n 𝑽 𝑖
C e n t r a li z a t i o n
Equivariant Vector field To model X ( t ) , a natural way is to estimate its differential d X ( 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.
M ETHODOLOGY
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.
3. S CALARIZATION B LOCK
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 x i ( t ) , define neighbor
N ( x i ( t )) as the particles that react with x i ( t ) .
Then we can consider a particle pair ( x i ( t ) ,
x j ( t )) , where x j ( t ) ∈
N ( x i ( 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 t ij :=
Scalarize( x i ( t ) , x j ( t ) , F ij ) , where Scalarize is the scalarization operation under an edgewise dynamical basis F ij defined below.
Equivariant basis construction For the particle pair ( x i ( t ) , x j ( t )) , let a ( t ) =
x i ( t ) − x j ( t ) ∥
[Equation]
i
[Equation]
t
[Equation]
− x j
[Equation]
t
[Equation]
∥ and × denotes the cross product of two vectors, define
[Equation]
Then we build a SO(3)-equivariant basis F ij :=
( a ( t ) , b ( t ) , c ( t )) .
In practice we add a small constant ϵ to the normalization factor in case that x i and x j collapse.
Under the condition that the matrix ( a ( t ) , b ( t ) , c ( t )) is non-degenerate, F ij formulates a complete orthonormal basis (frame) of the tangent space at x i ( 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 x i (Hsu, 2002)) of F ij is just its transpose: ( a T ( t ) ,
b T ( t ) , c T ( t )) .
The ‘bad’ event that F ij is degenerate for all neighbors happens only when all particles are restricted to a straight line, which is a measure zero set in R 3 .
Therefore, we assume F ij is non-degenerate from now on.
Proof for SO(3)-equivariance of F ij 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 F ij ( ( a ( t ) , b ( t ) , c ( t )) ), the position vector of x k naturally owns a ‘coefficient’ or ‘scalar’ representation:
[Equation]
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 x k → g x k and ( a ( t ) , b ( t ) , c ( t ))
→ ( g · a ( t ) , g · b ( t ) , g · c ( t )) .
Therefore (3.2) undergoes
[Equation]
where we use the fact that g T 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
[Equation]
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 F ij = ( a , b , c ) , θ can be expressed as:
[Equation]
The scalars tuple t ij :
=
{ θ aa , θ ab , . . .
} are the scalarization of θ under the equivariant basis F ij , 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:
[Equation]
3. G RAPH T RANSFORMER B LOCK
Input: X = ( x 1 , . . .
, x N ) ∈ R N × 3 , h i ∈
R h , e ij ∈ R e , G X //
Centralization
X ← Centralize ( X ) ; // Scalarization Block for x i ∈
X do for x j ∈ N ( x i ) do
F ij = EquiBasis ( x j , x j ) ; t ij = Scalarize ( F ij , x i , x j ) ; end end //
Graph Transformer Block m ij =
GraphTransformer ( t ij , h i , e ij ) ; //
Vectorization Block
V i = 1 N (cid:80) j ∈N ( x i ) Vectorize ( m ij , F ij ) ;
Output: V i
After encoding the geometric tensors into SO(3)- invariant scalars t ij , we first embed them alone with other pre-defined node/edge attributes ( h j , e ij ) 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 m ij by propagating and aggregating information on the graph G X .
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.
3. V ECTORIZATION BLOCK
Given the refined edgewise message m ij , the vectorization block is designed to transform these scalars back to equivariant vectors, which requires pairing m ij with the corresponding complete basis 1 F ij :=
( a , b , c ) .
More precisely, we first project m ij into a scalar triple { x 1 , x 2 , x 3 } , then define the vectorization process as:
[Equation]
We encapsulate the pairing process as V ij =
Vectorize( m ij , F ij ) .
It’s easy to check that the output follows the transformation rule of vectors.
Finally, we aggregate all vectors V ij associated with x i to estimate the ground-truth vector field V i .
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: Till now, we have introduced the basic blocks of EVFN.
We make some discussions on its property.
First, The SO(3)-equivariance of our model 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.
E XPERIMENTS
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)
. 4. N EWTONIAN 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 Newtonian forces under three settings.
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 ) ∈
R 4 × 3 and velocities V ( t ) ∈ R 4 × 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: f l ( v ( t ))
=
q v ( 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 x
i ( t ) , t ∈ { ∆ t, 2∆ t . . .
, T 1 } as the training labels and observations x i ( t ) , t ∈ { T 1 + ∆ t, T 1 + ∆ t, . .
, T 2 } as the validation set.
To evaluate the interpolation and extrapolation capacity of all methods, the observations x i ( t ) , t ∈ { 12 ∆ t, 32 ∆ t, . .
, T 1 + 12 ∆ t } and x i ( t ) , t ∈ { T 2 + ∆ t, T 2 + 2∆ t, . .
. , T 3 } 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 :
[Equation]
where L s 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.4) to integrate both the position and velocity trajectories.
Only the MSE between predicted position trajectory and ground truth is taken as the loss penalty:
[Equation]
where ( x t 0 , v t 0 ) 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] ⊤ , 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: T 1 = { 1 . 5 , 0 . 5 , 0 . 5 } , T 2 = { 2 , 0 .
55 } and T
= { 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¨ohler 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
[Table]
that EVFN exhibits stronger expressive power by representing geometric information losslessly 2 ; (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. M OLECULAR C ONFORMATION G ENERATION
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.4, A.3.2).
The optimization objective of EVFN Φ can be summarized as:
[Equation]
where λ ( t ) :
[0 , T ] → R + is a positive weighting function and ∇ H t i is the pre-computed gradient field of noisy atomic positions.
Once the score network is optimized, we can use an annealed
GCN
A.3.
[Table]
Langevin dynamics (ALD) sampler or an ODE-based PC sampler as the evolving block to generate conformations (Song et al., 2020).
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 (Sch ¨ utt 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 respectively.
[Equation]
where S g and S r denote generated and reference conformations, respectively
δ is a given RMSD threshold.
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 summarized 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.
In particular, compared with
ConfGF which employs a similar learning strategy with us, EVFN significantly increases 26 .
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
[Table]
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.
R ELATED W ORK
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 (Sch ¨ utt et al., 2018;
K ¨ ohler 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, EVFN can be viewed as the extension of
EGNN in the vein that EVFN extends the one basis (i.e., radical direction x i − x j ) 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.
; Shi et al.
; Wu et al.
(2021) model the equilibrium states by going through the process determined by the gradient of the density.
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 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.
P ROOF 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 ∈ SO
(3) , then under the action of g , the positions of the many-body system X ( t ) transform equivariantly:
[Equation]
Then from the definition of a ( t ) , we know that
[Equation]
For b ( t ) , since
[Equation]
where we have used g − 1 = g T for orthogonal matrix g to get the last line.
Therefore, b ( t ) g −→ g b ( t ) .
Applying (A.1) once again, we have c ( t ) g −→ g c ( t ) .
Frame bundle and scalarization technique A frame at x ∈ M is a linear isomorphism u from R d to the tangent space at x : T x M .
We use F ( M ) x to denote the space of all frames at x .
Then
GL ( d, R ) acts on F ( M ) x by
[Equation]
M.
Then the frame bundle
[Equation]
) 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, R d 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 { e i , 1 ≤ i ≤ d } be the canonical basis of R d , and { e i } the corresponding dual basis.
At each frame u , the vectors Y i :=
ue
i form a basis of T x M .
Let { Y i } be the dual frame of T ∗ x M , then a (r,s)-tensor θ can be expressed as
[Equation]
The scalarization of θ at u is
[Equation]
Through scalarization, a tensor field θ becomes an ordinary vector space valued function on F ( M ) :
[Equation]
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
( R 3 ) (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 R 3 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 ( R 3 ) .
A frame u ∈ O
( R 3 ) at π ( u ) ∈ R 3 can be transported to another point in R 3 by translation.
Therefore, we neglect the origin of the frame in the proof.
Let u e =
( e 1 , e 2 , e 3 ) be an equivariant basis of R 3 , then given a scalars tuple { θ i 1 ,...,i r } , we construct a vector-valued function on O ( R 3 ) by:
[Equation]
where g is the SO (3) transformation from u e to another frame u .
Moreover, ˜ θ i 1 ,...,i r ( u ) is SO(3)- equivariant, since
[Equation]
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 ˜
θ i 1 ,...,i r ( u ) induces a (r,0) tensor field on R 3 by following the definition of (2.1).
From scalarization ˜
θ
i 1 ,...,i r ( u ) on O ( R 3 ) to scalarization is obvious.
Note that any equivariant basis u e is also a point on O ( R 3 ) , therefore the scalars tuple is just the values of ˜ θ
i 1 ,...,i r at u e :
[Equation]
i
[Equation]
.,i
[Equation] [Equation]
...,i r
[Equation]
u e ) .
For general (r,s)-type tensors, the proof is the same by adding the dual basis of u e .
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 = x i × x j → b .
It implies that c → − c .
In conclusion, the orientation of the equivariant basis remains unchanged under reflection.
A. M ETHODOLOGY A.2. P ERMUTATION AND TRANSLATION EQUIVARIANCE
For a many-body system X ( t ) = ( x 1 ( t ) , . . .
, x n ( t ) , the centroid C ( t ) is defined by
[Equation]
Translating the reference by h , then
[Equation]
Therefore, recentering the reference’s origin to the centroid at the input’s time t = 0 , we have
[Equation]
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 ( x 1 , . . .
, x n ) is a permutation invariant multivariate continuous function, then
[Equation]
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 =
( v 1 , . . .
v n ) be its vector field, then v i ∈ R 3 corresponds the equivariant vector for x i ∈ { x 1 , . . .
, x n } .
Denote EVFN with parameters θ by Φ θ,t = { Φ θi } ni
=1 , then
[Equation]
= Φ θ i ( X ( t ) , t ) , for a fixed particle x i .
Suppose ( i 1 , . . .
, i k ) (neighbors of x i ) indicate indexes of particles which have interaction with x i , then obviously 1 ≤ k ≤ n − 1 .
By (A.4), Φ i ( X ( t ) , t ) is an aggregation of message from x i ’s neighbors, therefore we have:
[Equation]
) 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.
A.2. G RAPH T RANSFORMER BLOCK
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 m ij 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 t ij , node features h i and edge features m ij , GTB first embeds them into high-dimensional representations:
[Equation]
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 ϕ 1 m (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).
[Equation]
where ϕ 1 m , ϕ q , ϕ k , ϕ v , ϕ h and ϕ 2 m are fully connected networks.
α ij and M i denote the attention weights and the refined nodewise embeddings, respectively.
LayerNorm refers to the normalization layer adopted in (Vaswani et al., 2017).
A.2. E QUIVARIANCE OF EVFN
We need the following lemma:
Lemma A.4.
Suppose h is an invariant function, f 1 , . . .
, f k are arbitrary nonlinear functions.
Then, the composition f k ◦ · · · ◦ f 1 ◦ h is an invariant function.
Proof.
Consider the group action g acting on x ∈ R 3 , then f is invariant means that f ◦ g ( x ) :
= f ( gx ) =
f
( x ) .
We have
[Equation]
hence the composition f k ◦ · · · ◦ f 1 is invariant.
Let f 1 be the scalarization block.
Because the output of f 1 are scalars (from (3.3)) and SO(3)- scalars must be invariant under SO(3) group action, we conclude that f 1 is invariant.
Denote each graph-transformer by f i , i ∈ { 2 , . . .
, k } , then by Lemma A.4, it implies that the output of the graphtransformer block is invariant.
Finally, from the standard fact that scalars time 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.
A.2. E VOLVING BLOCK
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 ) =
{ x i ( t ) ∈ R 3 , i = 1 , 2 , · · · , n } whose underlying physical rules are unknown, let v i denote the velocity of particle x i .
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:
[Equation]
For the i -th particle,
[Equation]
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:
[Equation]
where W t is the Brownian motion and the initial state X (0) follows an unknown equilibrium distribution p 0 .
Denote the marginal distribution at time t by p t , then
[Equation]
so the Hamiltonian function H t at time t is entangled with p 0 .
According to the Liouville equation (the probability flow in Song et al.
(2020)), the reverse evolving process satisfies the following ODE:
[Equation]
The gradient field of the Hamiltonian function are also called the force field.
Therefore the NODE evolving block has the following form:
[Equation]
and for the i -th particle, following form:
[Equation]
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.
and our molecular experiment, we use the discretization of VP SDE (Song et al., 2020), where f ≡ 0 and
[Equation]
A.2. N EURAL 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 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
[Equation]
Taking the discrete layer index t to its continuous limit, we get
[Equation]
The discrete back-propagation method also has a continuous limit: the adjoint method Chen et al.
; Norcliffe et al.
A.2. S CALABILITY
Compared to EGNN, the extra cost of EVFN is to calculate the values of the scalars in t ij , ∀ i, j .
Compared to the computational cost of back-propagation in the neural networks, the cost of calculating the O (3 ∗ 3 ∗ N 2 ) (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.
A. E XPERIMENT A.3. N EWTONIAN MANY - BODY SYSTEM
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 ) ∈
R 4 × 3 and velocities V ( t ) ∈ R 4 × 3 .
The time evolution of the particles is given by
[Equation]
Gravity field This system consists of six particles under both the mutual newton’s gravitation and an external static force of the form: f g = (0 , 0 , η ) .
The time evolution of the particles are given by
[Equation]
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., f l ( v ) =
q v × 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:
[Equation]
[Table]
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 −for the other two datasets.
The number of training labels (i.e., trajectory points) for three tasks is 1 .
5 k, 100 k and 100 k, 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 4 , 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.
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
[Figure]
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 proposed method.
In most cases, EVFN with larger normalization constants can still outperform other baselines significantly.
A.3. M OLECULAR CONFORMATION GENERATION
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 [Figure] Setting Method
[Table] [Table] [Table] [Table]
[Table]
[Table]
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 &
Hern´andez-Lobato, 2019).
[Figure] 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 .
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:
[Equation]
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.
[Table]