E QUIVARIANT C ONSTRAINED G RAPH M ECHANICS N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning to reason about relations and dynamics over multiple interacting objects is a challenging topic in machine learning.
The challenges mainly stem from that the interacting systems are exponentially-compositional, symmetrical, and commonly geometrically-constrained.
Current methods, particularly the ones based on equivariant Graph Neural Networks (GNNs), have targeted on the ﬁrst two challenges but remain immature for constrained systems.
In this paper, we propose
Graph Mechanics Network (GMN) which is combinatorially efﬁcient, equivariant and constraint-aware.
The core of GMN is that it represents, by generalized coordinates, the forward kinematics information (positions and velocities) of a structural object.
In this manner, the geometrical constraints are implicitly and naturally encoded in the forward kinematics.
Moreover, to allow equivariant message passing in GMN, we have developed a general form of orthogonality-equivariant functions, given that the dynamics of constrained systems are more complicated than the unconstrained counterparts.
Theoretically, the proposed equivariant formulation is proved to be universally expressive under certain conditions.
Extensive experiments support the advantages of GMN compared to the state-of-the-art GNNs in terms of prediction accuracy, constraint satisfaction and data efﬁciency on the simulated systems consisting of particles, sticks and hinges, as well as two real-world datasets for molecular dynamics prediction and human motion capture.
I NTRODUCTION
[Figure]
Representing and reasoning about the relations and dynamics of a group of interacting objects is among the core aspects of human intelligence (Tenenbaum et al., 2011).
As a motivating example, we consider the Nbody system (Kipf et al., 2018) where the movement of a single charged particle is attracted or repelled by other charged particles.
Physicists have revealed that this process can be modeled by Newton’s laws along with the Coulomb force.
One may wonder, however, if we can teach a machine to rediscover the underlying physics by solely observing the particles’ states.
This thinking has actually inspired the study of learning to model interacting systems, which now is a prevailing topic in machine learning (Battaglia et al., 2016;
Thomas et al., 2018; K¨ohler et al., 2019; SanchezGonzalez et al., 2019; Fuchs et al., 2020; Martinkus et al., 2021;
Satorras et al., 2021).
Learning to model interacting systems is challenging.
First, the systems are combinatorially complex, on account of that objects can be composed in combinatorially many possible arrangements (Battaglia et al., 2016).
This challenge, to some extent, can be addressed by making use of
Graph Neural Networks (GNNs) (Wu et al., 2020).
By regarding objects as nodes and interactions as edges, GNNs extract information via message passing, which is able to characterize arbitrarily ordered objects and combinatorial relations.
The second challenge is related to an important symmetry in physics: the model we use should be equivariant to any Euclidean transformation (transla- tion/reﬂection/rotation) of the input.
This complies with the fact that physics rules keep unchanged regardless of the reference coordinate system.
Several works (Fuchs et al., 2020; Satorras et al., 2021) have investigated equivariance upon GNNs and exhibited remarkable beneﬁts on N-body.
Another challenge, though less explored, is that the systems could be geometrically constrained.
Geometric constraints arise in common practical systems, for example, in robotics where the joints of mechanical arms should be linked one by one, or in biochemistry where the atoms of molecules are connected by chemical bonds.
When modeling these constrained systems, it is crucial to enforce the model to output legal predictions.
For instance, the lengths/angles of chemical bonds in a molecular closely determine its chemical property which will be changed dramatically if the structural constraints are broken.
As mentioned above, equivariant GNNs (Fuchs et al., 2020; Satorras et al., 2021) have achieved desirable performance on the N-body system—this system, nevertheless, lacks of constraint.
Considering constraints is not easy, as the dynamics of all elements within a constrained system evolve in a joint and complex manner.
Unfortunately, to the best of our knowledge, there is no research (particularly among equivariant GNN methods) that learns to model the dynamical systems of multiple interacting objects under geometrical constraints.
In this paper, we propose Graph Mechanics Network (GMN) that can tackle the above three challenges simultaneously: I.
To cope with the combinatorial complexity, GMN takes advantage of graph models to encode the states of and interactions between objects, analogous to previous approaches.
II.
For the constraint satisfaction, GMN resorts to generalized coordinates, a well known notion in conventional mechanics to model the dynamics of structural objects.
Here, a structural object (such as the stick and hinge in Fig. 1) is deﬁned as a set of multiple rigidly connected particles.
In GMN, the constraints are implicitly encoded in the forward kinematics that describes the
Cartesian states as the function of generalized coordinates.
This strategy can inherently maintain the constraints and requires no external regulation.
III.
GMN is equivariant.
In GMN, we need to compute the interaction forces in the Cartesian space (see Eq. 5), and infer the accelerations of the generalized coordinates based on the inverse dynamics (see Eq. 6), both of which are learned by a general form of equivariant functions (see Eq. 11).
Notably, the proposed equivariant formulation is more efﬁcient to compute compared to the previous versions based on spherical harmonics (Thomas et al., 2018; Fuchs et al., 2020), or more general than the one developed by EGNN (Satorras et al., 2021).
More importantly, we have theoretically discussed when and how our formulation can universally approximate any equivariant function.
To evaluate the advantages of GMN, we have constructed a simulated dataset composed of three types of objects: particles, sticks and hinges, which is a complex case of the N-body system (Kipf et al., 2018) and can be used as building blocks for common systems.
Under various scenarios in terms of different ratios of object types and different numbers of training data, we empirically verify the superiority of GMN compared to state-of-the-art models in prediction accuracy, constraint satisfaction and data efﬁciency. In addition, we also test the effectiveness of GMN on two real-world datasets: MD17 (Chmiela et al., 2017) and CMU Motion Capture (CMU, 2003).
R ELATED W ORK
Learning to simulate complex physical systems has been shown to greatly beneﬁt from using GNNs.
Interaction Network (IN) proposed by Battaglia et al. (2016) is the ﬁrst attempt for this purpose, and it can be deemed as a special kind of GNNs to learn how the system interacts and how the states of particles evolve.
Later researches have extended IN in different aspects: HRN (Mrowca et al., 2018) utilizes hierarchical graph convolution for tackling objects of various geometrical shapes and materials, NRI (Kipf et al., 2018) further explicitly infers the interactions with the help of a variational auto-encoder, and Hamiltonian graph networks (Sanchez-Gonzalez et al., 2019) equip GNNs with ordinary differential equations and Hamiltonian mechanics for energy conservation.
However, all above approaches have ignored the symmetry in physics and the GNN models they use are not Euclidean equivariant.
There is a subset of models (Ummenhofer et al., 2019; Sanchez-Gonzalez et al., 2020; Pfaff et al., 2020) that partially implement symmetries, but they only enforce translation equivariance but not rotation equivariance.
On the other hand, it is nontrivial to enforce rotation equivariance.
Tensor-Field networks (Thomas et al., 2018) uses ﬁlters built from spherical harmonics to allow 3D rotation equivariance, and this idea has been developed by SE(3) Transformer (Fuchs et al., 2020) that further takes the attention mechanism into account.
Anther class of works (Finzi et al., 2020a;
Hutchinson et al., 2021) resorts to the Lie convolution for the equivariance on any Lie group, based on lifting and sampling.
Recently, EGNN (Satorras et al., 2021) has proposed a simple yet effective form of equivariant message passing on graphs, which does not require computationally expensive higher-order representations while still achieving better performance on N-body.
As interpreted before, ensuring geometrical constraints is crucial for many practical systems, which, nevertheless, is seldom investigated in aforementioned works.
Although several attempts (Yang et al., 2020; Finzi et al., 2020b) have been proposed for learning to enforce constraints, they explicitly augment the training loss with soft Lagrangian regulation and thus are completely data-driven and have no guarantee of generalization for limited training data. DeLaN (Lutter et al., 2019) also employs generalized coordinates to describe the kinematics of the rigid object.
Nevertheless, it only target on the physical process of a single object, other than the complex systems with multiple rigid and structural objects that are the focus of this paper.
In DPI-Net (Li et al., 2018), the BoxBath task does share the similar strategy to us by ﬁrst predicting the canonical coordinates of the box and then using the forward kinematic model to obtain the Cartesian positions.
However, the passing messages in DPI-Net are scalars other than directional vectors (positions, velocities, and accelerations) used in our work. After all, both DeLaN and DPI-Net never study equivariant models.
G RAPH M ECHANICS N EURAL N ETWORK 3. N OTATIONS AND P RELIMINARIES
We begin with introducing the N-body system (Kipf et al., 2018), as illustrated in Fig.
1 (a).
This system consists of N interacting particles { P i } Ni =1 of the same mass, and the kinematics states of each particle are deﬁned as S i
=
( x i , v i ) , where x i , v i ∈ R 3 are the position and velocity vectors, respectively.
There could be certain non-vector information of each particle (such as charge), which is represented by a c -channel feature h i ∈
R c 1 .
Suppose the system we study is conservative and the dynamics is driven by the interaction force f ij ∈ R 3 between any pair of particles i and j .
According to Newton’s second law, the acceleration of particle i , a i ∈ R 3 is proportional to the aggregated force from other particles (cid:80) j (cid:54) =
i f ij .
All symbols will be speciﬁed with a superscript t for temporal denotations, e.g. , x ti indicating the position of particle i at time t .
In this paper, we are mainly concerned with the prediction task: we need to seek out a function φ
( { ( S 0 i , h 0 i ) } Ni =1 ) given the initial states { S 0 i } Ni =1 and features { h 0 i } Ni =1 to forecast the future states { S Ti } Ni
=1 at time T .
As presented in Introduction, two kinds of inductive biases have been explored previously.
The ﬁrst one is to apply the graph structure to capture the distribution of particle states and their interactions (Battaglia et al., 2016;
Kipf et al., 2018), where, particularly, the particle states S
i (along with h 0 i ) are as node features and the interaction forces f 0 ij as edge messages.
In this way, the transition function φ boils down to a GNN model.
The second inductive bias is that φ should be equivariant to any translation/reﬂection/rotation of the input states.
By saying equivariance, we imply
[Equation]
Here, O (3) deﬁnes the 3D orthogonal group (Fuchs et al., 2020) that consists of translation, reﬂection and rotation transformations; g · S 0 i denotes to perform transformation g on the states S 0 i , and it is instantiated as R ( x 0 i + b ) for the position and Rv
i for the velocity, where R ∈ R 3 × 3 is the orthogonal matrix and b ∈ R 3 is the translation vector.
Several works (Thomas et al., 2018; K¨ohler et al., 2019; Fuchs et al., 2020;
Satorras et al., 2021) have investigated both inductive biases, among which EGNN (Satorras et al., 2021) has achieved promising performance on N-body.
The typical process in EGNN iterates the following steps:
[Equation]
where the superscript l denotes the l -th layer; the acceleration a li returned by the message aggregation of ϕ egnn in Eq.
2 is adopted for the update of the velocity v li in Eq.
3 (multiplied by a scalar ψ ( h l − 1 i ) ∈ R ), followed by the renovation of the position x li in Eq.
The formulation of ϕ egnn is physically reasonable, since the interaction (actually the Coulomb force) between particles i and j truly depends on their relative position x
l − 1 ji = x
l − 1 i − x l − 1 j , node features h l − 1 i and h l − 1 j , and edge feature e ji .
To enable equivariant message passing, EGNN has developed a speciﬁc form of ϕ egnn to let a li be equivariant and h li be invariant in terms of the input x l − 1 ji , which will be presented in Eq. Notice that the computations for f li and h li are actually by two different functions, and we have abbreviated them into one in Eq. (and also Eq. 5 later), since they share the same inputs; besides, their parameters are shared following EGNN.
3. O UR G ENERAL A RCHITECTURE
Despite the desirable performance on N-body, existing methods are incapable of maintaining the geometric constraint.
In this section, we design a general architecture that intrinsically meets the requirement of geometry constraints by making use of the generalized coordinates.
The main idea is that the constraints are implicitly and exactly encoded in the forward kinematics that represents the Cartesian states of the system as the function of generalized coordinates.
We deﬁne O k
= { P i } n k i =
a structural object composed of n k rigidly connected particles.
Fig.
1 (b) illustrates two examples of the structural object, the stick with 2 connected particles and the hinge with 3 particles.
To preserve the distance, the dynamics of the two particles on a stick should be updated in a joint way, rather than fulﬁlling the independent process in EGNN (Eq. 3 and 4).
Besides, the force on each particle within O k will indirectly inﬂuence the dynamics of others through the physical connections.
This requires us to analyze the dynamics of the particles in O k as a whole, which is implemented by generalized coordinates.
There could be multiple generalized coordinates for each O k , some located in the Cartesian space but some in the angle space.
For instance, the states of a stick can be decoupled by two independent sets of generalized coordinates: the state of particleas the Cartesian coordinates and the relative rotation angles of particle 2 to 1 as the angle coordinates.
For conciseness, this section only focus on the Cartesian part which essentially determines the local coordinates in O k , with providing full examples in § 3.4.
We denote the position, velocity and acceleration of the generalized Cartesian coordinates as q k , ˙
q k and ¨ q k , respectively.
We now detail how to update the states of O k .
In Fig.
2, we ﬁrst compute the interaction force between each particle and others, and aggregate information of all particles within each structural object to infer the acceleration of the generalized coordinates (which is termed as the generalized acceleration henceforth) by inverse dynamics.
[Figure]
Then, the dynamical updates are carried out in the space of the generalized coordinates.
Finally, the updated generalized coordinates will be projected back to the particles’ states via the forward kinematics.
In form,
[Equation]
where, the elements of the system can be described in two views, {O k } Kk =1 as the objectlevel view and { P i } Ni =1 as the particle-level view; for distinction, we index the structural object with the subscript k and particles with i or j .
We explain each equation separately.
Interaction force (Eq. 5).
The interaction force f li is computed analogous to Eq. 2.
EGNN straightly regards the interaction force to be the acceleration of each particle in Eq.
But here, given the constraints between particles, we record the force as an intermediate variable that will contribute to the inference of the generalized acceleration in the next step.
Inverse dynamics (Eq. 6).
This step is the core of our methodology.
The generalized acceleration ¨ q lk is dependent to the forces f li on all particles within O k , and their relative positions x l − 1 ki
= x
l − 1 i − q l − 1 k and relative velocities v l − 1 ki
= v l − 1 i − ˙ q l − 1 k with regard to the generalized coordinates.
The formulation of Eq.
6 is physics-inspired.
In Appendix (Eq. 26), we have analytically derived the dynamics of hinges, where the acceleration of the hinge is indeed related to the forces, relative positions and relative velocities of all particles in each hinge.
This is reasonable in mechanics, since the cross product x l − 1 ki × f li yields the torque, and the relative velocity v l − 1 ki is related to the centrifugal force of particle
i around q k , both of which inﬂuence the acceleration ¨ q k .
Different from the analytical form which is always complex and hard to compute in practice, we will employ a learnable and equivariant function with universal expressivity.
The details are in § 3.3.
Generalized update (Eq. 7-8).
The updates of the generalized coordinates are akin to Eq.
3 andin EGNN, as the dimensions of the generalized coordinates have been made independent.
Notice that in Eq.
7 the scalar factor for ˙
q l − 1 k takes as input the summation of all hidden features, which is a generalized form of Eq. 3 for multiple particles in O k .
Forward kinematics (Eq. 9).
Once the generalized coordinates have been refreshed, we can derive the states of all particles in O k by proceeding the forward kinematics.
Different system of O k could have different type of the forward kinematics.
Here we denote it as the function FK ( · ) in general, while providing the speciﬁcations in § 3.4.
Although our method is initially targeted on structural objects, it is still applicable for free particles, if we set the generalized coordinates as the states of each particle and utilize the trivial forward kinematics (the identity mapping) in Eq. 9.
Alg. 1 has summarized the updates for all objects.
3. E QUIVARIANT M ESSAGE P ASSING
As shown in Eq. 1, the Euclidean equivariance on the estimation function φ is necessary for ensuring the physical symmetry.
When considering this property in our case, we demand the interaction force (Eq. 5) and generalized acceleration (Eq. 6) to be equivariant with respect to orthogonal transformations, while other equations are already equivariant 2 .
EGNN (Satorras et al., 2021) has developed a particular orthogonality-equivariant form for the acceleration output in Eq. 2:
[Equation]
where ρ w ( · ) is an arbitrary Multi-Layer Perceptron (MLP) with parameter w , and we have abbreviated other non-vector terms in Eq.
2 as h .
This formulation does satisfy the rotation equivariance, but it is unknown if it can be generalized to functions (such as Eq. 6) with multiple input vectors, and more importantly, its representation completeness is never explored rigorously.
In this section, we propose a general form of orthogonality-equivariant functions with necessary theoretical guarantees.
Without loss of generality, the target function we would like to enforce equivariance is denoted as ϕ
[Equation]
Z , h ) :
R d × m × R c → R d × m (cid:48) .
We deﬁne the below formulation,
[Equation]
It is easy to justify the function ϕ
( Z , h ) in Eq.
11 is equivariant to any orthogonal matrix i.e. , ϕ ( OZ , h ) = [Equation]
, h [Equation]
[Equation]
Apparently, Eq.
11 reduces to Eq.
10 by setting the number of vectors in Z as 1, namely, m = m (cid:48)
= 1 .
We immediately have the following theory.
Theorem 1.
If m ≥ d and the row rank of Z is full, i.e. rank ( Z )
=
d , then for any continuous orthogonality-equivariant function ˆ ϕ ( Z , h ) , there must exist an MLP ρ w satisfying (cid:107) ϕ
( Z , h ) − ˆ ϕ
( Z , h ) (cid:107) <
(cid:15) for arbitrarily small error (cid:15) .
The proof employs the universality of MLP (Cybenko, 1989; Hornik, 1991), with the entire details deferred in Appendix.
Theorem 1 is nutritive, as it characterizes the rich expressivity of formulating the equivariant function via Eq.
The condition holds in general for the function like Eq.
6 whose input Z = ( f i , x ki , v ki ) ∈ R 3 × 3 is of full rank when the force, position and velocity expand the whole space.
Indeed, this condition holds with probability 1, stated by the following corollary.
Corollary 1.
Assume m ≥ d , and also the entries of Z are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R .
Then, almost surely, the conclusion of Theorem 1 holds.
When the condition m ≥ d is invalid, the universal approximation still maintains if restricted in the linear subspace expanded by the columns of Z .
Corollary 2.
For any continuous orthogonality-equivariant function ˆ ϕ
( Z , h ) located in the linear subspace expanded by the columns of Z , the conclusion of Theorem 1 holds universally.
Corollary 2 tells that the message passing ϕ egnn in Eq.
10 is not universally expressive since m < d , and can only ﬁt the vectors parallel to x ( i.e. the relative position).
Yet, ϕ egnn is still physically complete for the Coulomb force that is oriented by the relative position between two particles.
In our experiments, we ﬁnd that more stable performance is delivered by further adding the normalization term speciﬁcally when m > 1 :
[Equation]
(12) where (cid:107) · (cid:107) F computes the Frobenius norm.
Notice that adding normalization does not change the conclusions in Theorem 1 and its two corollaries, since the norm is also a function of Z (cid:62) Z that can be approximated by MLP.
We implement ϕ 1 in Eq.
5 and ϕ 2 in Eq.
6 by using the general formulation Eq.
12, where
Z = x l − 1 ji and Z = ( f li , x l − 1 ki , v l − 1 ki ) , respectively.
Note that for the update of hidden feature h li in Eq.
5, we do not need equivariance but invariance, hence we set h li =
ρ w ( Z (cid:62) Z (cid:107) Z (cid:62) Z (cid:107) F , h l − 1 i ) by just keeping the invariant part in Eq.
3. K INEMATICS I MPLEMENTATIONS BY C ONSIDERING G ENERALIZED A NGLES
Implementation of hinges.
A hinge, as displayed in Fig.3 (a), consists of three particles 0, 1, 2, and two sticks 01 and 02.
The freedom degrees of this system can be explained in this way: particle 0 moves freely, and particles 1 and 2 can only rotate round particle
owing to the length constraint by the two sticks.
Hence, the generalized coordinates include the states of particledenoted as q 0 ∈ R 3 and the rotation Euler angles of stick 01 as θ 01 ∈ R 3 and 02 as θ 02 ∈ R 3 .
As q 0 are Cartesian, the dynamical update follows directly Eq.
With the forces { f li } 3 i = estimated in Eq.
5 and acceleration ¨ q l 0 by Eq.
6, the angle acceleration of θ 01 is calculated as
[Equation]
according to rigid mechanics, where the numerator is the relative torque, the denominator is the moment of inertia under unit mass, and × means the cross product.
Then the angle velocity is ˙ θ l 01 = ψ (cid:48) ( h l − 1+ [Equation]
+ [Equation] l
[Equation]
) ˙ θ l [Equation] + ¨ θ l 01 similar to Eq.
We do not need to record θ 01 , since the forward kinematics can be conducted without it.
In detail, the forward kinematics in Eq.
9 is:
[Equation]
(13) where rot ( ˙ θ l 01 ) indicates the rotation matrix around the direction of ˙
θ l 01 by absolute angle (cid:107) ˙
θ l 01 (cid:107) .
The dynamic updates for particle 2 is similar.
We put the whole details into Alg. 1 in Appendix.
Implementation of sticks.
For a stick with particles 1 and 2 in Fig 3 (b), we choose the center as the generalized Cartesian coordinate q 0 , and the rotation of particle 1 θ 01 and particle 2 θ( θ 01 = θ 02 ) as the generalized angle coordinates.
The dynamics propagation of q 0 is given by
Eq.
There are two choices to compute ¨ q l 0 , one using the general form in Eq.
6 and the other one leveraging a simpliﬁed version as ¨ q l 0 = (cid:80) i ∈O k ϕ ( f li ) by explicitly omitting the relative position and velocity.
The physical motivation of introducing the simpliﬁed version is that the acceleration of a stick center is only affected by the forces according to the theorem of the motion of the center of mass.
The angle accelerations are ¨ θ
[Equation]
01 =
[Equation]
[Table]
[Equation] [Table]
[Equation]
[Table] 2 , and the velocity becomes ˙ θ
[Equation]
01 =
ψ (cid:48) ( h l − 1+ h l − 1 2 ) ˙
θ l − 1 01 + ¨ θ l 01 .
The states of particles 1 and 2 are renewed as the same as Eq. 13.
[Table]
E XPERIMENTS 4. S IMULATION D ATASET : C ONSTRAINED N-B ODY
Datasets.
We inherit the 3D extension of Fuchs et al.
based on the N-body simulation introduced in Kipf et al.
For each trajectory, we provide the initial states of the system { S 0 i } Ni =1 , the particles’ charges { c i ∈ ( − 1 , 1) } Ni =1 and a conﬁguration indicating which particles are connected by sticks or hinges.
The task is to predict the ﬁnal positions { x Ti } Ni =1 of the particles when T = 1000 .
The validation and testing sets contain 2000 trajectories.
We evaluate the prediction error by the MSE metric.
Compared with the simulation conducted in Fuchs et al.
Satorras et al.
(2021), our dataset is more challenging in three senses: 1.
We consider systems with multiple scales, including 5, 10, and 20 particles in total, respectively.
We introduce to the system the dynamics of hinges and sticks (depicted in Appendix B), and construct various combinations between these objects.
We investigate the performance of each model across different scales of training set, e.g. , 500 and 1500, to see how the models perform with scarce or relatively abundant training data.
The system consisting of p isolated particles, s sticks and h hinges, is abbreviated as ( p, s, h ) .
Implementation details.
Following Satorras et al.
(2021), we use a linear mapping of the scale of initial velocity || v
i || 2 as the input node feature
h 0 i .
The edge feature is provided by a concatenation of the product of charges c i c j and an edge type indicator I ij , where I ij is valued as 0 if node i and j are disconnected, 1 if connected by a stick, and 2 if connected by a hinge.
Note that this edge type indicator is an augmentation over the original setting in EGNN, designed to enforce EGNN and other baselines the ability to distinguish different types of edges, namely, with or without constraints.
Other settings including the hyper-parameters are introduced in Appendix E.
Comparison with SOTAs.
Table 1 reports the performance of GMN and various compared models:
EGNN (Satorras et al., 2021) and its regulated version EGNNReg, SE3-Transformer (Fuchs et al., 2020), Radial-Field (RF) (K ¨ ohler et al., 2019), Tensor-Field-Network (TFN) (Thomas et al., 2018), and other two baselines, GNN and the Linear prediction (Satorras et al., 2021).
Regarding EGNN and EGNNReg, they share the same backbones ( i.e. ϕ 1 and ψ ) and training hyper-parameters (learning rates, layer number, etc) with our GMN for a fair comparison.
For EGNNReg, we explicitly involve a regularization term during training by enforcing the geometrical constraints, namely preserving the lengths between two particles on sticks and hinges; the regularization factor is ranged from 0.01 to 0.1, where the value giving the best performance is selected.
The default settings of
SE3-Transformer and TFN perform poorly on our experiments, hence we have tried our best to tune their hyper-parameters by validation.
From Table 1, we have these observations: 1.
GMN clearly outperforms all compared methods on all scenarios, suggesting the consistent superiority of our proposed methodology on constrained systems.
GMN is more robust when the complexity of the system increases.
On | Train | = 500 , for example, the performance of GMN degenerates slightly by increasing the number of sticks ( e.g. from (1,2,0) to (0,10,0)), or increasing the number of particles and hinges ( e.g. from (1,2,0) to (3,2,1)), while other methods (such as EGNN) drops signiﬁcantly in both cases.
Reducing the training size will hinder the performance of all compared methods remarkably.
On the contrary, GMN still performs promisingly in general.
For instance, on (3,2,1), EGNNReg becomes much worse from 2.62 to 7.01 when the volume of training data is decreased from 1500 to 500, whereas [Table] the change of GMN is smaller (2.10 v.s. 2.48).
This is reasonable as GMN has explicitly encoded the constraints as opposed to EGNN and EGNNReg that learn to remember constraints by training.
[Figure]
Data efﬁciency.
Fig.
4 depicts the prediction errors on (3,2,1) when the training size varies.
It is observed that GMN acts steadily, justifying its beneﬁt in data efﬁciency.
Once again, both EGNN and EGNNReg deliver much worse performance when the training size is small, and they approach GMN when the training dataset is enlarged sufﬁciently.
In physics, it is important to ensure the physics rules that are discovered by a relatively small number of experiments to be general enough for explaining universal phenomena.
Hence, data efﬁciency, as a key advantage of GMN, comes as an important requirement for learning to model physical systems.
Besides, by the comparison between GNN and other equivariant models, it is seen that equivariance is a crucial point for performance guarantee.
The generalization capability across different systems.
It is interesting to check how the models perform when trained on one system but tested on others.
Hence, Table 2 tests the generalization from (3,2,1) to (2,4,0) and (1,0,3).
Interestingly, the performances of all models on new systems are comparable with the original environment.
We conjecture that this ability could be attributed to the usage of GNN in capturing the combination diversity of the objects.
As before, GMN performs best. [Table]
Ablation studies.
The function ϕ 2 is designed to be equivariant in Eq.
To justify this necessity, we replace ϕ 2 with an MLP of the same size, and report the errors in Table 3, from which we conﬁrm that removing equivariance incurs detriment to the performance.
The default setting of GMN in
Eq. 6 is leveraging unshared acceleration inference ϕ 2 ( f li ) for sticks and ϕ 2 ( f li , x l − 1 ki , v l − 1 ki ) for hinges.
Here we investigate different cases by assigning the identical form of ϕ 2 for sticks and hinges with shared or unshared parameters.
Applying identical ϕ 2 ( f li , x l − 1 ki , v l − 1 ki ) mostly outperforms the case by using ϕ 2 ( f li ) , probably owing to the better expressivity of the former version.
Yet, both cases are worse than our design, implying that the dynamics of sticks and hinges should be modeled distinctly.
We have introduced a normalization term in Eq.
Tabledemonstrates that eliminating this term leads to divergence during training, possibly owing to he numerical instability in the forward/backward propagation. 4. A PPLICATIONS ON REAL - WORLD DATASETS
This subsection introduces how to apply GMN to practical applications including MD17 (Chmiela et al., 2017) and CMU Motion Capture (CMU, 2003).
It is not required to manually derive the entire kinematics for these complex systems; instead, each input system is decomposed as a set of particles and sticks ( e.g. , the circles in Fig. 5), which can be directly processed by our current formulation of
GMN without any modiﬁcation.
The core is modeling partial length-constraints of the input system with disjoint sticks. The full details of the kinematics decomposition trick are in Appendix C.
[Table]
MD17. We adopt MD17 (Chmiela et al., 2017) which involves the trajectories of eight molecules generated via molecular dynamics (MD) simulation.
Our goal here is to predict the future positions of the atoms given the current system state. We observe that the lengths of chemical bonds remain very stable during the simulation, making it reasonable to model the bonds as sticks.
We split the dataset into train/validation/test sets containing 500/2000/2000 frame pairs.
The complete implementation details are deferred to Appendix E. As presented in
Table 4, GMN outperforms other competitive equivariant models on 7 of the 8 molecules.
Particularly, on molecules with complex structures ( e.g. , Aspirin, Benzene, and Salicylic), the improvement of GMN is more signiﬁcant, showcasing the beneﬁt of constraint modeling on the bonds. Yet, we also observe that the constraint-aware models (GMN and EGNNReg) perform worse than others on Ethanol, possibly because Ethanol is a relatively small molecule with simple structure, where considering the bond constraints possibly makes less beneﬁt but instead hinders the learning. CMU Motion Capture. We use the motion data from the CMU Motion Capture Database (CMU, 2003), which contains the trajectory of human motion in various scenarios.
Different parts of the human body could be treated as hard rigid-body constraints.
We focus on the walking motion of single object (subject #35) containing 23 trials, similar to Kipf et al. We sample from the trials to get train/val/test sets containing 200/600/600 frame pairs with interval T = 30 .
As depicted in Table 5, GMN outperforms other models by a large margin, verifying the efﬁcacy of our equivariant constraint module on modeling complex rigid bodies. We further provide a visualization in Fig. 6, where GMN predicts the motion accurately while EGNN yields larger error.
[Table] �� � � � �� �� �� �� ���������������� �� �� �� � � � �� �� �� � � � �� �� �� �� ���������������� �� �� �� � � � �� �� �� � � � �� �� �� �� ���������������� �� �� �� � � � �� ��
C ONCLUSION
In this paper, we propose Graph Mechanics Networks (GMN) that are capable of characterising constrained systems of interacting objects.
The core is making use of generalized coordinates, by which the constraints are implicitly and exactly encapsulated in the forward kinematics.
To enable Euclidean symmetry, we have developed a general form of equivariant functions to simulate the interaction forces and backward dynamics, whose expressivity is theoretically justiﬁed.
For the simulated systems with particles, sticks and hinges, GMN outperforms existing methods regarding prediction error, constraint satisfaction and data efﬁciency.
Moreover, the evaluations on two realworld datasets support the generalization ability of GMN towards complex systems.
R EPRODUCIBILITY S TATEMENT
The complete proof of the theorems is provided in Appendix A.
The hyper-parameters and other experimental details are provided in Appendix E.
Our code is available at: https://github.com/GraphMechanicsNetwork/GraphMechanicsNetwork.
E THICS S TATEMENT
The research in this paper does NOT involve any human subject, and our dataset is not related to any issue of privacy and can be used publicly.
All authors of this paper follow the ICLR Code of Ethics (https://iclr.cc/public/CodeOfEthics)
A T HE PROOF OF T HEOREM
In the following we deﬁne the orthogonal group as O ( d ) = { O ∈ R d × d | O (cid:62) O = OO (cid:62) =
I d }
Prior to providing the proof for Theorem we ﬁrst list two necessary lemmas below.
Lemma 1.
For any two matrices Z 1 , Z 2 ∈ R
d
× m we have this equivalence: ∃ O ∈ O ( d ) , OZ 1 =
Z 2 ⇔ Z (cid:62) 1 Z 1 = Z (cid:62) 2 Z 2 .
Proof.
We only need to prove the “ ⇐ ” direction, as “ ⇒ ” holds clearly.
Suppose the SVD decomposition of Z 1 as Z
= U 1 S 1 V (cid:62) 1 with the left-singular matrix U 1 ∈ O ( d ) , the singular diagonal matrix S 1 ∈ R d × m , and the right-singular matrix V 1 ∈
O ( m ) .
Since Z (cid:62) 1 Z 1 =
Z (cid:62) 2 Z 2 , then there must exists a certain SVD decomposition of Z 2 that shares the same singular matrix and right-singular matrix with Z 1 , implying that Z 2 =
U 2 S 1 V (cid:62) 1 , where U
2 ∈ O ( d ) .
Hence,
Z
= U 2 S 1 V (cid:62)
= U 2 U (cid:62) 1 U 1 S 1 V (cid:62)
= U 2 U (cid:62) 1 Z 1 , which concludes the proof owing to the orthogonality of U 2 U (cid:62) 1 .
Lemma 2.
The function f : R d × m → R m (cid:48) is invariant on O ( d ) , namely, f ( OZ )
=
f ( Z ) , ∀ O ∈
O ( d ) , ∀ Z ∈ R d
× m , if and only if there exists function g : R m × m → R m (cid:48) satisfying f ( Z )
=
g ( Z (cid:62) Z ) .
Proof.
The sufﬁciency is obvious.
We now prove the necessity.
We deﬁne the equivalence class [ Z 0 ] =
{ Z | ∃ O ∈ O ( d ) , OZ =
Z 0 } .
Since f is invariant to the orthogonal transformation, it means f is actually a function on the equivalence class, i.e. , f ( Z )
= f
([ Z ]) .
On the other hand, according to Lemma 1, we have [ Z 1 ] =
[ Z 2 ] ⇔ Z (cid:62) 1 Z 1 =
Z (cid:62) 2 Z 2 , which implies the oneto-one correspondence between [ Z ] and Z (cid:62) Z ; hence, there must exist a function f (cid:48) leading to [ Z ] = f (cid:48)
( Z (cid:62) Z ) , and f (cid:48) is continuous in terms of any invariant metric such as the norm (cid:107) Z (cid:62) Z (cid:107) .
Overall, f ( Z ) = f
([ Z ]) =
f ( f (cid:48)
( Z (cid:62) Z )) :
= g ( Z (cid:62) Z ) .
We are now ready to prove Theorem 1 that is copied below for better readability.
Theorem 1.
If m ≥ d and the row rank of Z is full, i.e. rank ( Z )
=
d , then for any continuous orthogonality-equivariant function ˆ ϕ ( Z , h ) , there must exist an MLP ρ w satisfying (cid:107) ϕ
( Z , h ) − ˆ ϕ
( Z , h ) (cid:107) <
(cid:15) for arbitrarily small error (cid:15) .
Proof.
Without loss of generality, we will omit the non-vector term h , which does not change the story but let our proof more concise.
Because Z is of full row-rank, the columns of an arbitrary function ˆ ϕ
( Z ) can be represented as a linear combination of the columns of Z ; in other words, there must exist a function π : R d × m → R m × m (cid:48) giving rise to ˆ ϕ ( Z )
=
Z π ( Z ) .
Considering the orthogonality-equivariance, we derive the property of π ( Z ) as:
[Equation]
Since d ≤ m and the row-rank of Z is full, we perform the compact SVD decomposition on Z , namely, Z = U Z S Z V (cid:62) Z , where U Z ∈ O ( d ) , R d × d (cid:51) S Z > 0 , and V Z ∈ R m × d satisfying
V (cid:62)
Z V Z = I d .
Eq.
14 becomes:
[Equation]
Given that S Z = Eigen ( Z (cid:62) Z ) and V Z =
EigenVector ( Z (cid:62) Z ) are respectively the eigenvalues and eigenvectors of Z (cid:62) Z and are clearly invariant to the orthogonal transformation of Z .
Their values can be numerically approximated by iterative programs, such as the power method (Mises &
Pollaczek-Geiringer, 1929), thus can be treated as the continuous functions of Z (cid:62) Z .
Let us deﬁne g (cid:48)
( Z ) :=
S Z V (cid:62) Z
π
( Z ) .
Then, we keep deriving Eq.
15 by:
[Equation]
According to Lemma 2, the function g satisﬁes Eq.
if and only if it is written as g (cid:48)
( Z ) =
g ( Z (cid:62) Z ) for a certain function g .
By checking the formulation of ˆ ϕ ( Z ) as demonstrated before, we arrive at
[Equation]
Here the function η can be approximated by MLP whose universality has been justiﬁed by (Cybenko, 1989; Hornik, 1991).
The conclusion of Theorem 1 is proved.
Corollary 1.
Assume m ≥ d , and also the entries of Z are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R .
Then, almost surely, the conclusion of Theorem 1 holds.
Proof.
This is straightforward.
When rank ( Z ) <
d , the columns of Z are located in a subspace of R d (for example a line or a plane in the 3D space), whose measure is zero.
Therefore, the probability for making rank ( Z ) =
d is 1, and we almost surely have the same conclusion as Theorem 1.
Corollary 2.
For any continuous orthogonality-equivariant function ˆ ϕ
( Z , h ) whose output is located in the linear subspace expanded by the columns of Z , the conclusion of Theorem 1 holds universally.
Proof.
According to the deﬁnition of ˆ ϕ ( Z ) , we still obtain ˆ ϕ ( Z ) =
Z π ( Z ) .
Let us assume d > m (otherwise we can directly obtain Theorem 1), then the full SVD decomposition of Z is Z =
U Z S Z V (cid:62) Z , with U Z ∈ O ( d ) ,
S Z ∈ R d × m , and V Z ∈ O ( m ) .
But here, S Z is not strictly positive.
Suppose S Z
= (cid:18) S + (cid:19) , where S + > 0 .
We retain that S Z V (cid:62)
Z π ( Z ) =
g ( Z (cid:62) Z )
by imitating the proof in Eq.
Analogous to Eq.
17, we derive,
[Equation]
which concludes the proof.
B T HE DYNAMICS ANALYSES FOR STICKS AND HINGES
The analytical forms of the dynamics for sticks and hinges can be found from a mechanics book.
Here we derive the formulas on our own to make our paper more self-contained.
For simplicity, we consider all particles to be of the equal mass and the sticks of no mass.
Dynamics analysis of sticks.
In Fig.
7 (Left), we assume the forces acting on particles 1 and 2 are separately f 1 and f 2 .
By following the theorem of the motion of the center of mass, the acceleration of the center is given by
[Equation]
The rotation accelerations of particles 1 and 2 around the center are calculated by
[Equation]
20) where M deﬁnes the total torque and J is the moments of inertia.
Dynamics analysis of hinges.
The analysis for hinges is more complicated than sticks.
In Fig.(Right), the forces on particles 0, 1 and 2 are f 0 , f 1 and f 2 .
By using Newton’s second law,
[Equation]
(21) where the aggregated force is f = f 0 + f
+ f 2 .
In addition, the kinematics relations between the three particles show that
[Equation]
[Equation]
where ˙ θ 01 and ¨ θ 01 denote the speed and acceleration of the rotation angle of particle 1 around 0, and ν 01 is the corresponding linear velocity; the symbols ˙ θ 02 , ¨ θ
and ν 02 are deﬁned similarly.
Moreover, the relative rotation acceleration of particle 1 to 0 is caused by the external force f 1 and the inertia force − m a 0 acting on 1, which derives that
[Equation]
Analogously,
[Equation] [Equation]
After substituting Eq.
24 into Eq.
22, and Eq.
25 into Eq.
23, and then rearranging Eq.
21, we have
[Equation]
26) where a = f m − ˙ θ 01 × ν 01 − ˙ θ 02 × ν 02 − ( I − e 01 e (cid:62) 01 )
f 1 m − ( I − e 02 e (cid:62) 02 ) f 2 m with e 01 and erepresent the unit vectors along x 01 and x 02 , respectively.
We have applied a trick to simplify the derivation by observing
[Equation]
01 ) f 1 . Besides, I + e 01 e (cid:62) {+
+}
+ e 02 e (cid:62) 02 is invertible, hence Eq.
26 is always meaningful.
Substituting Eq.
26 back into Eq.
24 and Eq. 25 derives the values of ¨ θ 01 and ¨ θ 02 .
Note that in § 3.4, we employ the denotation of generalized coordinates by q 0 , ˙ q 0 and ¨ q 0 which indeed share the same meaning with the Cartesian coordinates x 0 , v 0 and a 0 of particle 0.
C K INEMATICS DECOMPOSITION
We manually build the forward kinematics (Eq. 9) of the stick (and hinge), which relies on the domain knowledge of the underlying physics. However, for complex systems, we no longer require to derive the kinematics of the entire system.
Instead, we propose kinematics decomposition, a simple yet effective trick that can decompose each input system (an arbitrary graph) into particles and sticks.
Taking the MD17 dataset as an example, for each molecule, we select certain bonds as sticks (the circles in Fig. 5) and the remaining atoms as isolated particles; in this way, we obtain a set of particles and sticks.
Note that different sticks are not allowed to intersect; otherwise, it will generate two values for the intersecting particle of two sticks and cause ambiguity if these two values are distinct.
Although this kind of kinematics decomposition will only maintain partial constraints, it greatly enlarges the application scope of our current formulation Eq. (5-9) without any revision. More importantly, our experiments verify that GMN by this formulation is sufﬁcient to surpass other methods on complex systems like molecules.
On CMU Motion Capture, since the motion graph contains no circle and is of the tree-like structure, it is tractable to derive the exact forward kinematics by recursive kinematics computation from the root node. Yet, we still encourage the usage of the above kinematics decomposition for its easy implementation and compatibility with our GMN.
Algorithm
Input: Initial states of all particles { S 0 i =
( x 0 i , v 0 i ) } Ni =1 and features { h 0 i } Ni ; Learnable equivariant functions ϕ 1 , ϕ 2 , ψ, ψ (cid:48) ; Layer number L .
Compute the generalized coordinates of all structural objects { ( q 0 k , ˙ q 0 k , { ˙ θ 0 ki } i ∈O k ) } Kk =1 .
for layer l = 1 to L do for particle i = 1 to N do
Calculate the interaction force f li and feature h li for each particle by:
[Equation]
(27) end for for object k = 1 to K do
Inference the generalized Cartesian acceleration ¨ q lk via:
[Equation] [Equation]
for hinges
[Equation]
Derive the generalized angle acceleration:
[Equation] [Equation]
for hinges
[Equation]
Update the positions and velocities as follows:
[Equation] [Equation]
[Equation]
Perform the forward kinematics for each particle in O k :
[Equation]
The predicted states of all particles { S Li } Ni =1 .
D F ULL ALGORITHMIC DETAILS
In the main body of the paper, for better readability, we ﬁrst introduce the general pipeline of our method in § 3.2 and then present the implementation details by taking the angels into account in § 3.4.
Here, we combine them into one singe algorithmic ﬂowchart in Alg. 1.
E M ORE EXPERIMENTAL DETAILS AND RESULTS
Hyper-parameters and baselines.
For GNN, RF, EGNN, EGNNReg, and GMN, we empirically ﬁnd that the following hyper-parameters generally work well, and use them across all experimental dim 64, and weight decay
SE3-Transformer and TFN do not perform well on our datasets, potentially due to the challenge of highly complex and constrained systems.
Consequently, we tune the hyper-parameters and adopt the following conﬁguration: batch size 100, learning rate 0.001, hidden dim 64, representation degrees 3 and weight decay 1 × 10 − 8 .
All models are evaluated with four layers.
Models are trained for 600 epochs on the simulation dataset, and 500 epochs on the real-world datasets.
EGNNReg is a variant of EGNN that explicitly adds the constraint error into its training loss by a regularization factor of λ .
In our experiments, we also treat λ as a hyper-parameter, and choose λ that yields the best performance within the range [0.01, 0.1].
Detailed experimental setup on MD17.
We randomly split the dataset into train/validation/test sets containing 500/2000/2000 frame pairs respectively.
We choose T = 5000 as the span between the input and prediction frames, and the difference in positions as the input velocity.
The hyperparameters of all models are kept the same as the synthetic dataset.
We masked out the hydrogen atoms, focusing on the prediction of large atoms.
We further augment the original molecular graph with 2-hop neighbors, and concatenate the hop index with atom numbers of the connected atoms as well as the edge type indicator as the edge feature, similar to Shi et al.
We use the norm of velocity concatenated with the atom number as the node feature.
We randomly select bonds without commonly-connected atoms as sticks, and the rest of atoms as isolated particles.
Although by this means not all the bond lengths are preserved, our experiment does illustrate that this is a simple but effective strategy of applying GMN on complex systems like molecules.
Detailed experimental setup on CMU Motion Capture.
We ﬁrst split the data into train/val/test sets containing 11/6/6 trials respectively.
We then sample from the trials to get 200/600/600 frame pairs with T = 30 .
We use the norm of velocity as node feature.
We also augment the edges with 2-hop neighbors, and use the edge type indicator as edge feature.
As for GMN, we sample 6 key bones of human body (as speciﬁed in Appendix C) as sticks, resulting in a system with 19 isolated particles and 6 sticks.
Empirically, we select the edges connecting nodes (0, 11), (2, 3), (7, 8), (12, 13), (17, 18), and (24, 25) as sticks on the motion capture dataset, which indeed are the key parts of human body like the arms and legs.
Visualization on the simulation dataset.
We visualize the prediction outcomes by GMN and
EGNN in Fig.
8 on the simulation dataset.
GMN is found to be able to track the ground-truth trajec-
[Figure]
tories accurately, whereas EGNN yields clear position errors and particularly breaks the constraints for sticks and hinges.
These results are consistent with the performance in Table 1.
More experimental results.
In Table 6, we provide a comprehensive performance comparison of different models under more scenarios with 500 training samples.
Clearly, GMN outperforms other models on all object combinations involved.
Besides, we provide the learning curve for GMN,
EGNN, and EGNNReg on the simulation dataset (3,2,1).
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
F M ORE A BLATIONS
Hinge treated as two sticks.
We treat 0-1, 0-2 as two sticks, and apply the stick FK respectively.
Afterwards, we translate the two sticks such that the 0s coincide at the midpoint of their predicted positions.
From Table 7, we ﬁnd that this strategy performs worse than the hinge-modeled GMN, since splitting the hinge into two sticks overlooks the kinematics at the connected point.
Yet and still, it performs better than EGNN, again verifying the beneﬁt of our proposed constraint modeling.
[Table]
GMN 2.
Charges as node or edge feature.
In the experiment we by default assign e ij =
c
i c j as the edge feature (denoted as “Edge + C”).
Here we instead concatenate c i to the node feature of particle i , and set all e ij = 0 (denoted as “Node + C”).
In Table 8 we observe that these alternatives on charges make very limited difference on performance, and indeed the models can learn the interaction of charges from node features, which is truly the case as depicted in Eq.
2 and 5.
[Table]
G C ONSTRAINT S ATISFACTION
Constraint error.
The constraint error is computed as the total change in the lengths of sticks and hinges between the input and output, averaged per trajectory.
Speciﬁcally, for hinges, the two edges are both considered.
Results.
One vital feature of GMN is that it maintains the geometrical constraints exactly and inherently.
To show this, Table 9 records the corresponding constraint errors of several typical models.
The results do verify our claim that GMN always outputs near-zero errors (all below 1e-4).
Although EGNNReg that augments EGNN with regulation helps in reducing the constraint errors, it is data-driven and limited by the number of training samples; further, the constraints are pursued softly, making it defective for the applications where hard constraints are indispensable.
[Table]
H M ORE DISCUSSIONS ON GENERALIZATION
In Table 2 we compare the generalization capability of GMN with other methods.
Here, we denote
GMN trained on (3,2,1) and tested across different scenarios as GMN-Transfer
[Table]
and GMN trained and tested on the same dataset as GMN-Original.
It is observed from Table 10 that GMN has strong generalization capability, since the transfer performance is very close to the original setting.
[Table]