Equivariant Vector Field Network for Many-body System Modeling

1 INTRODUCTION . Modeling many-body systems has been a long-standing challenge in scientific fields from classical and quantum physics ( Carleo & Troyer , 2017 ; Zhang et al. , 2018 ; Satorras et al. , 2021c ) to structural biology ( Senior et al. , 2020 ; Shi et al. , 2021 ) , due to its high numerical complexity and complicated evolving mechanism . Graph neural network ( GNN ) , which is superior to model the high-dimensional structured data with permutation equivariance , brings a new opportunity to model the many-body systems in an end-to-end manner . Since many-body physical systems follow many physical constraints like SE ( 3 ) symmetry , pure black-box GNN models show limitations on generalization in this scenario and symmetry-preserving GNN models have become a hot research direction . The core question to be solved for developing general equivariant GNN models is how to conduct nonlinear operations on tensors in a reference-free way . To represent and manipulate equivariant tensor of arbitrary orders , some approaches resort to equivariant function spaces such as spherical harmonics ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ; Bogatskiy et al. , 2020 ; Fuchs et al. , 2021 ) or lifting the spatial space to high-dimensional spaces such as Lie group space ( Cohen & Welling , 2016 ; Cohen et al. , 2018 ; 2019b ; Finzi et al. , 2020 ; Hutchinson et al. , 2021 ) . Since no restriction on the order of tensors is imposed on these methods , sufficient expressive power of these models is guaranteed . Unfortunately , transforming a many-body system into those high-dimensional spaces or calculating equivariant functions usually brings excessive computational cost and great optimization difficulty , which is unacceptable in some real-world scenarios . To remedy this issue , Satorras et al . ( 2021c ) proposed EGNN to directly implement equivariant operations in the original space , providing an efficient way to preserve equivariance without performing complex space transformations . Detailed experiments in ( Satorras et al. , 2021c ) have shown that preserving equivariance without transforming the space is theoretically possible and computationally efficient in practice . However , one trade-off of EGNN is abandoning a certain amount of tensor information 1 , which causes the equivariance function class EGNN can approximate to be restricted . This drawback may become serious when modeling complex dynamical scenarios such as molecular simulation , where geometric information ( e.g. , angular potentials and rotatable bonds ) plays an important role in inducing conformation changes ( Klicpera et al . ( 2020 ) ; Xu et al . ( 2021 ) ) . To mitigate this issue , we propose a new model called Equivariant Vector Field Network ( EVFN ) to fit the gradient vector fields ( Song & Ermon , 2019 ; Shi et al. , 2021 ) of many-body systems . With a scalarization block and a vectorization block , EVFN is able to represent tensor information losslessly in the original space and outputs equivariant vector fields with no restriction on the direction . Inspired by the scalarization technique from differential geometry ( Kobayashi , 1963 ; Hsu , 2002 ) , EVFN first introduces a tuple of complete basis associated with each particle pair that preserves permutation and SE ( 3 ) symmetry under global rotation and translation transformation on the manybody system . Based on this basis , the scalarization block losslessly transforms the geometric information into SO ( 3 ) -invariant scalar representations . In principle , the scalar representations can be fed into any permutation-equivariant network to implement complex nonlinear transformations . Moreover , the vectorization block could reverse the scalars back to the vector field without sacrificing geometric information with the complete basis . Once obtained the estimated gradient field , we can predict a certain state or the whole dynamical trajectory of a 3D many-body system via an integration procedure . We evaluate the proposed framework on two many-body scenarios that require equivariance : ( 1 ) the simulated Newtonian many-body dynamics trajectory prediction and ( 2 ) the real-world molecular conformation generation . Our model achieves best or competitive results in various types of datasets . 2 BACKGROUND . In this section , we first introduce some basic concepts on the notion of equivariance and tensor field and then describe the scalarization technique from differential geometry . Finally , we define the ‘ vector field ’ as the differential of a many-body system . Let X = ( x1 , . . . , xN ) ∈ RN×3 be a many-body system living in R3 , where N is the number of particles . We use xi ( t ) to denote the position of the particle i at time t. SE ( 3 ) group and Equivariance In the Euclidean space R3 we can consider affine transformations that preserve the distance between any two points , i.e. , the isometric group SE ( 3 ) . We call it the symmetry group w.r.t . the Euclidean metric , and it turns out that SE ( 3 ) can be generated by the translation group and the rotation group SO ( 3 ) . Once we have a symmetry group , it ’ s valid to define quantities that are “ equivariant ” under the symmetry group . Given a function f : Rm → Rn , assuming the symmetry group G acts on Rm and Rn , then f is G-equivariant if f ( gx ) = gf ( x ) , ∀x ∈ Rm and g ∈ G. For SO ( 3 ) group , if n = 1 , i.e. , the output of f is a scalar , then the group action on R1 is the identity map , in this case f should be SO ( 3 ) -invariant ( Thomas et al. , 2018 ) : f ( gx ) = f ( x ) . The notion of tensor field can be defined for general Riemannian manifold ( see Definition 2.1.10 of ( Jost & Jost , 2008 ) ) . Let { ∂ ∂xi } 3i=1 and { dxi } 3i=1 be the tangent vectors and dual vectors in R3 respectively , and ⊗ denotes the tensor product . Then recall the definition of tensor field for R3 w.r.t . the SO ( 3 ) group : Definition 2.1 . A ( r , s ) - tensor field θ is a multi-linear map from a collection of r vectors and s dual vectors in R3 to R : θ ( x ) = θi1···irj1···js ∂ ∂xi1 ⊗ · · · ⊗ ∂ ∂xir ⊗ dxj1 ⊗ · · · ⊗ dxjs . It implies that under SO ( 3 ) coordinate transformation g : = { gij } 1≤i , j≤n , the tensor field θ transforms equivariantly : θ i′1···i ′ r j′1···j′s = gi′1i1 · · · gi′rirg T j1j ′ 1 · · · gTjsj′sθ i1···ir j1···js , where g T is the inverse of g. 1For example , dihedral angle is a function of position vectors rather than the position ’ s norm , which are the input of EGNN . Detailed definition is given in appendix A.2.4 Equivariant Vector field To model X ( t ) , a natural way is to estimate its differential dX ( t ) dt and apply an ODE solver to integrate the differential to obtain the dynamic trajectory or a state at a given time . Due to the SO ( 3 ) symmetry , we define such differential as an equivariant vector field . Most 3D real-world scenarios adopt first-order and second-order equivariant vector fields to depict their dynamic evolving mechanisms , which are also the modeling targets in this paper . A typical secondorder vector field is the acceleration field of Newtonian systems . Gradient field is a widely used terminology meaning the first-order derivative w.r.t . a scalar function ( Jost & Jost , 2008 ; Song & Ermon , 2019 ) . To generate molecular conformations ( i.e . equilibrium state ) with a single stage , ( Shi et al. , 2021 ) define “ gradient field ” to serve as pseudo-forces acting on each particle . By evolving the particles following the direction of the gradient field , the nonequilibrium system will finally converge to an equilibrium state . Gradient field is a special case of first-order vector field . 3 METHODOLOGY . 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.1 SCALARIZATION BLOCK . The scalarization block is designed to transform geometric tensors into edgewise SO ( 3 ) -invariant scalar features by introducing a novel tuple of complete basis . Given a particle xi ( t ) , define neighbor N ( xi ( t ) ) as the particles that react with xi ( t ) . Then we can consider a particle pair ( xi ( t ) , xj ( t ) ) , where xj ( t ) ∈ N ( xi ( t ) ) . Suppose we take the positions of the two particles as the relevant geometric information of edge ⟨i , j⟩ , then the edgewise SO ( 3 ) -invariant scalars could be defined as tij : = Scalarize ( xi ( t ) , xj ( t ) , Fij ) , where Scalarize is the scalarization operation under an edgewise dynamical basis Fij defined below . Equivariant basis construction For the particle pair ( xi ( t ) , xj ( t ) ) , let a ( t ) = xi ( t ) −xj ( t ) ∥xi ( t ) −xj ( t ) ∥ and × denotes the cross product of two vectors , define b ( t ) = xi ( t ) × xj ( t ) ∥xi ( t ) × xj ( t ) ∥ and c ( t ) = a ( t ) × b ( t ) , ( 3.1 ) Then we build a SO ( 3 ) -equivariant basis Fij : = ( a ( t ) , b ( t ) , c ( t ) ) . In practice we add a small constant ϵ to the normalization factor in case that xi and xj collapse . Under the condition that the matrix ( a ( t ) , b ( t ) , c ( t ) ) is non-degenerate , Fij formulates a complete orthonormal basis ( frame ) of the tangent space at xi ( t ) . Note that this is a dynamical basis w.r.t . t and the construction process of such basis is permutation-equivariant . Since the Euclidean metric is flat , the dual basis ( living in the cotangent space of xi ( Hsu , 2002 ) ) of Fij is just its transpose : ( aT ( t ) , bT ( t ) , cT ( t ) ) . The ‘ bad ’ event that Fij is degenerate for all neighbors happens only when all particles are restricted to a straight line , which is a measure zero set in R3 . Therefore , we assume Fij is non-degenerate from now on . Proof for SO ( 3 ) -equivariance of Fij is provided in proposition A.1 . Invariant scalarization of geometric tensors With the complete equivariant basis , we can realize the scalarization operation ( Kobayashi , 1963 ) in an elementary way . First of all , notice that under the basis Fij ( ( a ( t ) , b ( t ) , c ( t ) ) ) , the position vector of xk naturally owns a ‘ coefficient ’ or ‘ scalar ’ representation : ( xk · a ( t ) , xk · b ( t ) , xk · c ( t ) ) . ( 3.2 ) We define the process obtaining such scalars as Scalarize operation . Here we demonstrate that the set of obtained coefficients ( 3.2 ) is actually a SO ( 3 ) -invariant scalar tuple . Let g ∈ SO ( 3 ) be an arbitrary orthogonal transformation , then xk → gxk and ( a ( t ) , b ( t ) , c ( t ) ) → ( g · a ( t ) , g · b ( t ) , g · c ( t ) ) . Therefore ( 3.2 ) undergoes ( xk · a ( t ) , xk · b ( t ) , xk · c ( t ) ) → ( gxk · ga ( t ) , gxk · gb ( t ) , gxk · gc ( t ) ) = ( xk · a ( t ) , xk · b ( t ) , xk · c ( t ) ) , ( 3.3 ) where we use the fact that gT g = I to get the last line . It is easy to prove that the Scalarize operation could transform arbitrary geometric tensors into SO ( 3 ) -invariant scalars . Taking ( 2,0 ) -type tensors as an example , by extending the complete basis ( a , b , c ) through tensor product , it ’ s easy to check that { a⊗ a , b⊗ b , c⊗ c , a⊗ b , b⊗ a , a⊗ c , c⊗ a , b⊗ c , c⊗ b } forms an equivariant basis of the ( 2 , 0 ) -type tensor space . Then the scalarization of a tensor is just the linear combination coefficients under this basis . In the same way as ( 3.3 ) , we can prove that the coefficients are also SO ( 3 ) -invariant scalars . Given a ( 2 , 0 ) -symmetric tensor θ ( e.g. , energymomentum tensor ) , under the complete basis Fij = ( a , b , c ) , θ can be expressed as : θ =θaaa⊗ a+ θbbb⊗ b+ θccc⊗ c+ θab ( a⊗ b+ b⊗ a ) + θac ( a⊗ c+ c⊗ a ) + θbc ( b⊗ c+ c⊗ b ) . ( 3.4 ) The scalars tuple tij : = { θaa , θab , . . . } are the scalarization of θ under the equivariant basis Fij , which are SO ( 3 ) -invariant . Since any nonlinear transformations acting on SO ( 3 ) -invariant scalars are still SO ( 3 ) -invariant , therefore the scalars tuple can be fed into any neural network architectures without concerns about breaking the equivariance symmetry . Our scalarization is inspired by the scalarization technique on the frame bundle in differential geometry ( Hsu , 2002 ) , the equivalence of the scalarization technique and ( 3.4 ) is given in proposition A.2 . In practice , we focus on the scalarization of ( 1 , 0 ) -type tensors ( i.e. , vectors ) . Since the most common geometric information of input is vectors in real-world scenarios . We define the Scalarize operation as : Scalarize ( xi , xj , Fij ) = ( xi · a ( t ) , xi · b ( t ) , xi · c ( t ) , xj · a ( t ) , xj · b ( t ) , xj · c ( t ) ) ( 3.5 )

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