Equivariant Graph Hierarchy-Based Neural Networks
Abstract
Equivariant Graph neural Networks (EGNs) are powerful in characterizing the dynamics of multi-body physical systems.
Existing EGNs conduct flat message passing, which, yet, is unable to capture the spatial/dynamical hierarchy for complex systems particularly, limiting substructure discovery and global information fusion.
In this paper, we propose Equivariant Hierarchy-based Graph Networks (EGHNs) which consist of the three key components: generalized Equivariant
Matrix Message Passing (EMMP) , E-Pool and E-UnPool.
In particular, EMMP is able to improve the expressivity of conventional equivariant message passing,
E-Pool assigns the quantities of the low-level nodes into high-level clusters, while
E-UnPool leverages the high-level information to update the dynamics of the lowlevel nodes.
As their names imply, both E-Pool and E-UnPool are guaranteed to be E( n )-equivariant to meet the physical symmetry.
Considerable experimental evaluations verify the effectiveness of our EGHN on several applications including multi-object dynamics simulation, motion capture, and protein dynamics modeling.
Introduction
Understanding the multi-body physical systems is vital to numerous scientific problems, from microscopically how a protein with thousands of atoms acts and folds in the human body to macroscopically how celestial bodies influence each other’s movement.
While this is exactly an important form of expert intelligence, researchers have paid attention to teaching a machine to discover the physical rules from the observational systems through end-to-end trainable neural networks.
Specifically, it is natural to use Graph Neural Networks (GNNs), which is able to model the relations between different bodies into a graph and the inter-body interaction as the message passing thereon [1, 15, 24, 25, 20].
More recently, Equivariant GNNs (EGNs) [28, 8, 7, 26] have become a crucial kind of tool for representing multi-body systems.
One desirable property is that their outputs are equivariant with respect to any translation/orientation/reflection of the inputs.
With this inductive bias encapsulated,
EGN permits the symmetry that the physical rules keep unchanged regardless of the reference coordinate system, enabling more enhanced generalization ability.
Nevertheless, current EGNs only conduct flat message passing in the sense that each layer of message passing in EGN is formulated in the same graph space, where the spatial and dynamical information can only be propagated node-wisely and locally.
By this design, it is difficult to discover the hierarchy of the patterns within complex systems.
Hierarchy is common in various domains.
Imagine a complex mechanical system, where the particles are distributed on different rigid objects.
In this case, for the particles on the same object, their states can be explained as the relative states to the object (probably the center) plus the dynamics of the object itself.
We can easily track the behavior of the system if these “implicit” objects are detected automatically by the model we use.
Another example, as illustrated in Figure 1, is the dynamics of a protein.
Most proteins fold and change in the form of regularly repeating local structures, such as α -helix, β -sheet and turns.
By applying a hierarchical network, we are more capable of not only characterizing the conformation of a protein, but also facilitating the propagation between thousands of atoms in a protein by a more efficient means.
There are earlier works proposed for hierarchical graph modeling [12, 5, 32, 3, 17], but these studies focus mainly on generic graph classification, and more importantly, they are not equivariant.
In this paper, we propose Equivariant Graph Hierarchy-based Network (EGHN), an end-to-end trainable model to discover local substructures of the input systems, while still maintaining the
Euclidean equivariance.
In a nutshell, EGHN is composed of an encoder and a decoder.
The encoder processes the input system from fine-scale to coarse-scale, where an Equivariant-Pooling (E-Pool) layer is developed to group the low-level particles into each of a certain number of clusters that are considered as the particles of the next layer.
By contrast, the decoder recovers the information from the coarse-scale system to the fine-scale one, by using the proposed Equivariant-Up-Pooling (EUnPool) layer.
Both E-Pool and E-UnPool are equivariant with regard to Euclidean transformations via our specific design.
EGHN is built upon a generalized equivariant layer, which passes directional matrices over edges other than passing vectors in EGNN [26].
To verify the effectiveness of EGHN, we have simulated a new task extended from the N-body system [15], dubbed M -complex system, where each of the M complexes is a rigid object comprised of a set of particles, and the dynamics of all complexes are driven by the electromagnetic force between particles.
In addition to M-complex, we also carry out evaluations on two real applications: human motion caption [4] and the Molecular Dynamics (MD) of proteins [27].
For all tasks, our
EGHN outperforms state-of-the-art EGN methods, indicating the efficacy and necessity of the proposed hierarchical modeling idea.
Related Work
GNNs for modeling physical interaction.
Graph Neural Networks (GNNs) have been widely investigated for modeling physical systems with multiple interacting objects.
As pioneer attempts,
Interaction Networks [1], NRI [15], and HRN [19] have been introduced to reason about the physical interactions.
With the development of neural networks enforced by physical priors, many works resort to injecting physical knowledge into the design of GNNs.
As an example, inspired by
HNN [11], HOGN [24] models the evolution of interacting systems by Hamiltonian equations to obtain energy conservation.
Another interesting feature of physical systems lies in Euclidean equivariance, i.e. , translation, rotation, and reflection.
Several works first approach translation equivariance [29, 25, 20, 30].
Yet, dealing with rotation equivariance is non-trivial.
TFN [28] and
SE(3)-Transformer [8] leverages the irreducible representation of the SO(3) group, while LieConv [7] and LieTransformer [14] extend the realization of equivariance to Lie group.
Apart from these works that resort to group representation theory, a succinct equivariant message passing scheme on E( n ) group is depicted in EGNN [26].
GMN [13] further involves equivariant forward kinematics modeling particularly for constrained systems.
[2] generalizes EGNN to involve covariant information with steerable vectors.
[21] leverages frame averaging for general equivariance.
[18] mainly studies sign and basis invariance.
Despite the rich literature, these models either violate the equivariance, or inspect the system at a single granularity, both of which are vital aspects when tackling highly complicated systems like proteins.
E - U p P oo l
E - P oo l …
E - U n P oo l
E - P oo l 𝒢 (#$%)
𝑺 (#$%)
E MM
P …
Initial System: 𝒢 (/)
T arget System: 𝒢 (3) 𝒢 (#) 𝒢 (#4%)
Encoder
Decoder
E MM
P
E MM
P 𝒢 (#$%)
Hierarchical GNNs.
There are also works that explore the representation learning of GNNs in hierarchies.
Several GNNs [12, 5, 31] adopt graph coarsening algorithms to view the graph in multiple granularities.
[9] leverages a U-net architecture with top- k pooling.
Another line of work injects learnable pooling modules into the model.
A differentiable pooling scheme DiffPool [32] has been introduced to learn a permutation-invariant pooling in an end-to-end manner.
[3] replaces the aggregation in DiffPool by node dropping for saving the computational cost.
[17] further incorporates self-attention mechanism into the pooling network.
[6] leverages junction tree to model molecular graph in multiple hierarchies.
Nevertheless, these techniques, although permutation equivariant, lack the guarantee of geometric equivariance, limiting their generalization on real-world 3D physical data.
The Proposed EGHN
This section first introduces the notations of our task, and then presents the design of the EMMP layer, which is the basic function in EGHN.
Upon EMMP, the details of how the proposed E-Pool and E-UnPool work are provided.
Finally, the instantiation of the entire architecture is described.
Notations.
Each input multi-body system is modeled as a graph G consisting of N particles (nodes)
V and the interactions (edges) E among them.
For each node i , it is assigned with a feature tuple ( Z (0) i , h (0) i ) , where the directional matrix Z (0) i ∈ R n × m is composed of m n -dimension vectors, such as the concatenation of position x i ∈ R 3 and velocity v i ∈ R 3 , leading to Z (0)
i =
[ x i , v i ] ∈ R 3 × 2 ; h i ∈ R
c is the non-directional feature, such as the category of the atom in molecules.
The edges are represented by an adjacency matrix A ∈ R N × N , which can either be constructed according to the geometric distance or physical connectivity.
We henceforth abbreviate the entire information of a system, i.e. , ( { Z (0) i , h (0) i } Ni =1 , A ) as the notation G in if necessary.
We are mainly interested in investigating the dynamics of the input system G in .
To be formal, given the initial state ( Z (0) i , h (0) i ) of each particle, our task is to find out a function ϕ to predict its future state Z ( T ) i given the interactions between particles.
As explored before [28, 8, 7, 26], ϕ is implemented as a GNN to encode the inter-particle relation.
In addition, it should be equivariant to any translation/reflection/rotation of the input states, so as to obey the physics symmetry about the coordinates.
It means, ∀ g ∈ E ( n ) that defines the Euclidean group [26],
[Equation]
where g · Z (0) i conducts the orthogonal transformation as RZ (0) i for both the position and velocity vectors and is additionally implemented as the translation
x i + b for the position vector; the ellipsis denotes the input variables uninfluenced by g , including h (0) i and A .
As discussed in Introduction, existing equivariant models [28, 8, 7, 26] are unable to mine the hierarchy within the dynamics of the input system by flat message passing.
To address this pitfall,
EGHN is formulated in the encoder-decoder form:
[Equation]
Here, as illustrated in Figure 2, the encoder aims at clustering the particles of G in with similar dynamics into a group that is treated as the particle in the high-level graph G high (the number of the nodes in G high is smaller than G in ).
We have developed a novel component, E-Pool to fulfill this goal.
As for the decoder, it recovers the information of all particles in the original graph space under the guidance of the high-level system G high , which is accomplished by the proposed E-UnPool.
It is worth mentioning that both E-Pool and E-UnPool, as their names imply, are equivariant, and they are mainly built upon an expressive and generalized equivariant message passing layer, EMMP.
To facilitate the understanding of our model, we first introduce the details of this layer in what follows.
3.1 Equivariant Matrix Message Passing
Given input features { ( Z i , h i ) } Ni =1 , EMMP performs information aggregation on the same graph to obtain the new features { ( Z ′ i , h ′ i ) } Ni =1 .
The dimension of the output features could be different from
[Equation]
6) the input, unless the row dimension of Z ′
i should keep the same as Z i ( i.e. equal to n ).
In detail, one EMMP layer is updated by Eq.
3-6, where
MLP ( · ) is a Multi-Layer Perceptron, N ( i ) collects the neighbors of i , and ˆ Z ij ∈ R n × 2 m =
( Z i − ¯
Z
, Z j − ¯ Z ) is a concatenation of the translated matrices on the edge ij .
¯ Z is the mean of all nodes for the position vectors and zero for other vectors.
With the subtraction of ¯ Z , ˆ Z ij is ensured to be translation invariant, and then Z ′ i is translation equivariant after the addition of Z i in
Eq. 6.
Specifically, the MLP in Eq.
3 takes as input the concatenation of the E( n )-invariant
ˆ Z ⊤ ij ˆ Z ij , h i , and h j , mapping from R 2 m × 2 m +2 c to R 2 m × m + c , and the output is split into H ij ∈ R 2 m × m and h ij ∈ R c .
The formal proof for the E( n )-equivariance of EMMP is deferred to Appendix.
Distinct from EGNN [26], the messages to pass in EMMP are directional matrices other than vectors.
Although GMN [13] has also explored the matrix form, it is just a specific case of our EMMP by simplifying ˆ Z ij =
Z i − Z j .
Indeed, we have the following theorem for the comparison of expressivity between EMMP, EGNN, and GMN, with the proof in Appendix.
Theorem 1 EMMP can reduce to EGNN and GMN by specific choices of MLP in Eq.
Besides, since taking the inner product might induce a larger variance in the scale of input, in our implementation we also enforce a normalization ˆ Z ⊤ ij ˆ Z ij / ∥ ˆ Z ⊤ ij ˆ Z ij ∥ F before feeding the invariant ˆ Z ⊤ ij ˆ Z ij into the MLP in Eq.
3, following the suggestion by GMN for better numerical stability.
3.2 Equivariant Pooling
Inspired by DiffPool, we propose E-Pool, an equivariant pooling module.
Formally, E-Pool coarsens the low-level system G low =
( { ( Z low i , h low i ) } Ni =1 , A low ) into an abstract and high-level system
G high =
( { ( Z high i , h high i ) } K i =1 , A high ) with fewer particles, K < N .
For this purpose, we first perform EMMP (Eq. 3-6) over the input system G to capture the local topology of each node.
Then we apply the updated features of each node to predict which cluster it belongs to.
This can be realized by a Softmax layer to output a soft score for each of the K clusters.
The cluster is deemed as a node of the high-level system, and its features are computed as a weighted combination of the low-level nodes with the scores it just derives.
In summary, we proceed:
[Equation]
where Eq.
8 maps the invariant feature
h ′ i into the score
s i ∈ R K of cluster assignment with Softmax performed long the feature dimension, and the score matrix is given by S
=
[ s ij ] N × K with s i being its i -th row.
By this design, it is tractable to verify that E-Pool is guaranteed to be E( n ) equivariant (also permutation equivariant).
Specifically, the division by the row-wise sum (cid:80) Ni =1 s ij in Eq.
9 is essential, as it permits the translation equivariance, that is,
[Equation]
This particular property distinguishes our pooling from traditional non-equivariant graph pooling [32, 17].
Notice that the normalization in Eq.
10 is unnecessary since h i is a non-directional vector, but it is still adopted in line with Eq. 9.
In practice, it is difficult to attain desirable clusters by using the SoftMax layer solely; instead, the pooling results are enhanced if we regulate the training process with an extra reconstruction loss related to the score matrix, whose formulation will be given in § 3.4.
3.3 Equivariant UnPooling
E-UnPool maps the information of the high-level system G high back to the original system space G low , leading to an output system G out .
We project the features back to the space of the original low-level
Z
[Equation]
i
[Equation]
system by using the transposed scores derived in EPool.
Then, the projected features along with the low-level features are integrated by an E( n ) equivariant function to give the final output.
The procedure of E-UnPool is given by Eq.
12-15, where ˆ
Z i =
[ Z low
i
[Equation]
; Z
[Equation]
i − ¯ Z agg ] is the column-wise concatenation of the mean-translated low-level matrix Z low i and the high-level matrix Z agg i , analogous to
Eq.
One interesting point is that Eq.
12 is naturally equivariant in terms of translations, even without the normalization term used in Eq. 9.
This is because the score matrix is summed to 1 for each row, indicating that (cid:80) Kj =1 s ij ( Z high j + b ) =
(cid:80) Kj =1 s ij Z high j + b .
We have the following theorem guaranteeing the equivariance, with all proofs deferred to Appendix.
Theorem 2 EMMP, E-Pool, and E-UnPool are all E( n )-equivariant.
3.4 Instantiation of the Architecture
The overall architecture constitutes an equivariant U-Net [23] with skip-connections.
We design the overall architecture as a sequence of EMMP, E-Pool, and E-UnPool in an encoder-decoder fashion, as depicted in Figure 2.
The encoder is equipped with a certain number of E-Pools and EMMPs, while the decoder is realized with E-UnPools and EMMPs.
For each E-UnPool in the decoder, as already defined in § 3.3, it is fed with the output of the previous layer, the score matrix S from
E-Pool, and the low-level system G from EMMP in the corresponding layers of the encoder.
Here, the so-called corresponding layers in E-Pool and E-UnPool are referred to the ones arranged in an inverse order; for example, in Figure 2, the final E-Pool corresponds to the first E-UnPool.
With such design, it is straightforward, by the conclusion of Theorem 1, that the resulting EGHN still satisfies
E( n )-equivariance.
There is always one EMMP layer prior to each E-Pool or E-UnPool.
This external EMMP plays a different role from the internal EMMP used in E-Pool (Eq. 7).
One crucial difference is that they leverage different adjacency matrices.
As we have introduced before, the adjacency matrix
A can either be specified by geometric distance, i.e. , distance-based, or physical connectivity, i.e. connectivity-based.
The external EMMP exploits a distance-based A global whose element is valued if the distance between two particles is less than a threshold; by such means, we are able to characterize the force interaction between any two particles even they are physically disconnected.
In higher-layer external EMMP, its A global is created as a re-scored form (akin to Eq. 11) of A global in lower layer, where the score matrix is obtained by its front E-Pool.
For the internal EMMP in
E-Pool, it applies a connectivity-based A local that exactly reflects the physical connection between particles, for example, it is valued 1 if there is a bond between two atoms.
In this way, E-Pool pays more attention to locally-connected particles when conducting clustering.
Another minor point is that the external EMMP is relaxed as EGNN for only modeling the radial interaction, whereas the internal
EMMP uses the generalized form in § 3.1.
As we will show in our experiments, such design yields more favorable results compared with using any one of A global and A local only.
The training objective of EGHN is given by:
[Equation]
where ∥ · ∥ F computes the Frobenius norm, L is the number of E-Pools in the encoder, and λ is the trade-off weight.
The first term is to minimize the mean-square-error between the output state Z out
i and the ground truth Z gt i .
The second term is the connectivity loss that encourages more connects within the pooling nodes and less cuts among pooling clusters [33].
For training stability, we first perform row-wise normalization of ( S ( l ) )
⊤ A ( l − 1) S ( l ) before substituting it into Eq.
Experiments
We contrast the performance of the proposed EGHN against a variety of baselines including the equivariant and non-equivariant GNNs, on one simulation task: the M -complex system, and the two real-world applications: human motion capture and molecular dynamics on proteins.
We also carry out a complete set of ablation studies to verify the optimal design of our model.
4.1 Simulation Dataset: M -complex System
Data generation.
We extend the N -body simulation system from [15] and generate the M-complex simulation dataset, in order to introduce hierarchical structures in the data.
Specifically, we initialize a system with N charged particles { x i , v i , c i } Ni =1 distributed on M disjoint complex objects {S j } Mj =1 , where x i , v i , c
i are separately the position, velocity, and charge for each particle.
Within each complex S j , the particles are connected by rigid sticks, yielding geometric objects like sicks, triangles, tetrahedrons, etc.
The dynamics of all M complexes are driven by the electromagnetic force between every pair of particles.
The task here is to predict the final positions { x Ti } Ni of all particles when
T = 1500 given their initial positions and velocities.
Without knowing which complex each particle belongs to, we will also test if our EGHN can group the particles correctly just based on the distribution of the trajectories.
We independently sample J different systems, each of which has
M complexes with the number of particles sampled from a uniform distribution with mean N/M .
A dataset consists J systems with M complexes, N/M average size of complex is abbreviated as ( M, N/M, J ) .
We adopt Mean Squared Error (MSE) as the evaluation metric for the experiments.
Implementation details.
We assign the node feature as the norm of the velocity ∥ v i ∥ 2 , and the edge attribute as c i c j for the edge connecting node i and j , following the setting in [26].
We also concatenate an indicator, which is set as 1 if a stick presents and 0 otherwise, to the edge feature, similar to [13].
We use a fully connected graph (without self-loops) as A global , since the interaction force spans across each pair of particles in the system.
The adjacency matrix A reflects the connectivity of the particles formed by the complexes.
We set the number of clusters the same as the number of complexes in the dataset.
The comparison models include: Linear Prediction (Linear) [26],
SE(3)-Transformer (SE(3)-Tr.)
[8], Radial-Field (RF) [16], MPNN [10] and EGNN [26].
For all these models, we employ the codes and architectures implemented by [26].
Detailed hyper-parameter settings are in Appendix.
Results.
Table 1 reports the overall performance of the comparison models on eight simulation datasets with different configurations.
From Table 1, we have the following observations: 1.
Clearly,
EGHN surpasses all other approaches in all cases, demonstrating the general superiority of its design.
Increasing the number of complexes ( M ) or the number of particles ( N ) always increases the complexity of the input system, but this does not necessarily hinder the performance of EGHN.
For example, in both the single-system and multiple-system cases, EGHN even performs better when the system is changed from (5 , 5) to (5 , 10) and (10 , 10) .
We conjecture that, with more particles/complexes, larger systems also provide more data samples to enhance the training of EGHN.
When increasing the diversity of systems ( J ) by switching from the single-system mode to multisystem mode, the performance of EGHN only drops slightly, indicating its adaptability to various scenarios.
Visualization.
we visualize in Figure 3 the predictions of EGNN and our EGHN on the (3 , 3 , 1) scenario.
We find that EGHN predicts the movements of the rigid objects more accurately than EGNN, especially for the large objects.
In the right sub-figure, we also display the pooling results of EGHN, outputted by the score matrix of the final E-Pool layer.
It is observed that EGHN is
[Table]
� � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � � able to detect the correct cluster for each particle.
This is interesting and it can justify the worth of designing hierarchical architecture for multi-body system modeling.
4.2 Motion Capture
We further evaluate our model on CMU Motion Capture Databse [4].
We primarily focus on two activities, namely walking (Subject #35) [15] and running (Subject #9).
With regard to walking, we leverage the random split adopted by [13], which includes 200 frame pairs for training, 600 for validation, and another 600 for testing.
As for running, we follow a similar strategy and obtain a split with 200/240/240 frame pairs.
The interval between each pair is 30 frames in both scenarios.
In this task the joints are edges and their intersections are the nodes.
[Table]
MPNN
[Table]
[Table] [Table]
RF [16] 188.0 ±
[Table] [Table]
TFN [28] 32.0 ± 1.
56.6 ± 1.
SE(3)-Tr.
[8] 31.5 ± 2.
61.2 ± 2.
EGNN [26] 28.7 ± 1.
50.9 ± 0.
GMN [13] 21.6 ± 1.
44.1 ± 2.
EGHN 8.5 ± 2.
25.9 ± 0.
Implementation details.
As discussed in [8], many realworld tasks, including our motion capture task here, break the Euclidean symmetry along the gravity axis ( z − axis), and it is beneficial to make the equivariant models aware of where the top is.
To this end, we augment the node feature by the coordinate of the z − axis, resulting in models that are height-aware while still equivariant in the horizontal directions.
This operation is also applied to all baselines.
Since the interaction of human body works along the joints, we propose to involve the edge in A global if it connects the nodes within two hops in G .
For the number of clusters K , we empirically find that K = 5 yields promising results for both walking and running.
Results.
Table 2 summarizes the whole results of all models on two subjects.
Here, we supplement an additional baseline GMN [13] for its promising performance on this task.
Excitingly, EGHN outperforms all compared baselines by a large margin on both activities.
Particularly, on Subject #35,
[Figure]
�� �� �� �� � � � �� � �� �� �� �� �� �� �� �� �� �� �� �� � � � �� � �� �� �� �� �� �� �� �� �� �� �� �� � � � �� � �� �� �� �� ��
[Figure]
the prediction error of EGHN is 8 .
5 × 10 − 2 , which is much lower than that of the best baseline, i.e. ,
GMN ( 21 . 6 × 10 − 2 ).
Visualization.
To investigate why EGHN works, we depict the skeletons estimated by both EGNN and EGHN on Subject #9 in Figure 4.
It shows that EGHN is able to capture more fine-grained details on certain parts ( e.g. the junction between the legs and the body) than EGNN.
When we additionally visualize the pooling outcome in the right sub-figure, we interestingly find that EGHN is capable of classifying the two right-left hands into the same cluster even they are spatially disconnected.
A similar result is observed for the arms and feet.
This is reasonable as EGHN checks not only if two particles are spatially close to each other but also if they share the similar dynamics.
4.3 Molecular Dynamics on Proteins
We adopt AdK equilibrium trajectory dataset [27] via MDAnalysis toolkit [22] to evaluate our hierarchical model.
The AdK equilibrium trajectory dataset involves the MD trajectory of apo adenylate kinase simulated with explicit water and ions in NPT at 300 K and 1 bar.
The atoms’ positions of the protein are saved every 240 ps for a total of 1.004 µ s as frames.
Implementation details.
We split the dataset into train/validation/test sets along the timeline that contain 2481/827/878 frame pairs respectively.
We choose T = 15 as the span between the input and prediction frames.
We ignore the hydrogen atoms to focus on the prediction of large atoms.
We further establish the global adjacency matrix as the neighboring atoms within a distance of 10Å.
The atoms’ velocities of the protein at each frame are computed by subtracting the positions to the previous frame’s positions.
We further leverage MDAnalysis to extract the protein backbone in order to reduce the data scale.
Even so, TFN and SE( 3 )-Transformer still run out of memory, and thus we compare our model with the rest of baselines.
Detailed hyper-parameters are in Appendix.
[Table]
Results.
The prediction MSE is depicted in Table 3.
Our EGHN yields significantly lower error on protein
MD compared with the baselines, achieving 1.843 MSE, while the second best model MPNN has an MSE of 2.322.
However, MPNN is non-equivariant, and we find that its MSE will dramatically increase to 605.7 if we apply a random rotation of the protein during testing.
Compared with EGNN, our EGHN exhibits its superiority thanks to the hierarchical modeling, particularly favorable on large and complex systems like proteins.
Qualitative comparisons.
We visualize the protein structure of top-1 predictions generated by different models in cartoon format in Fig.
5 (a), with more visualization examples provided in
Appendix.
In Fig.
5 (a), the structure in red indicates the ground truth, while the other colors indicate the prediction.
We can observe that EGHN tracks the folding and dynamics of the protein more precisely than the baselines.
For example, in the the bottom region, EGHN gives a close-fitting result of the alpha helix structure while the predictions from MPNN and EGNN have an obvious shift compared with the ground truth.
To validate the power of the E-Pool, we further visualize the pooling clusters in Fig.
5 (b).
Interestingly, the pooling assignment exhibits certain clusters in some structures of the protein.
It suggests that EGHN discovers local repetitive sub-structures of the protein; for instance, it detects the alpha helix structure in the middle of the protein.
4.4 Ablation Studies
[Table]
EGHN ( K = 5 )
[Table]
EGHN ( K = 3 )
[Table] [Table]
[Table]
EGHN ( K = 8 ) 14.
[Table]
w/o Equivariance 19.
w/o Hierarchy 21.
Replace by EGNN 22.
w/o Connectivity loss 10.
A global only 17.
A local only 16.
We investigate the necessity of our proposed components on motion capture dataset in Table 4.
We study the following questions:
Q1.
How will the performance of EGHN change, if we vary the number of clusters ( K )?
We modify the number of clusters K from 5 toand 8, both of which yield worse performance.
Specifically, we find that decreasing K on “Run” results in a larger degradation of performance, possibly because the activity “Run” is with complicated kinematics and it will be more difficult to learn if the joints are shared across a too small number of clusters.
We provide potential guidance on choosing K in Appendix D.1.
Q2.
How do our proposed two components EMMP and hierarchical modeling contribute?
We replace all EMMP layers in our model by typical non-equivariant
MPNN, and the performance drops from 8.5 to 19.7 on Walk, supporting that maintaining equivariance is vital.
We further set s i
= 1 i in all E-Pool and E-UnPool and observe that removing hierarchy is detrimental to accurate prediction.
Moreover, by replacing all EMMPs with EGNNs, the performance also drops, which aligns with our analysis on the stronger expressivity of EMMP over EGNN.
Complete studies are deferred to Appendix D.2 and D.3.
Q3.
How does the connectivity loss (the second term in Eq.
16) help?
By dropping the connectivity loss, we observe a larger prediction error.
This justifies the necessity of using the connectivity loss to focus more on intra-cluster connections against the inter-cluster edges.
Q4.
How about using the same adjacency matrix for all EMMP instead of distinguishing them as A global in the external EMMPs and A local in internal
EMMPs as discussed in § 3.4?
When we apply A global or A local for all EMMPs, the performance drops dramatically, implying that the external and internal EMMPs play different roles in EGHN, and should be equipped with different adjacency matrices to model the interactions of different scopes.
Discussion
Limitation.
In the current form the number of clusters K is fixed in EGHN as an empirical hyperparameter.
Future works include extending E-Pool to dynamically adjust K for systems with different scales for enhancing the flexibility of the hierarchical model.
Conclusion.
In this paper, we propose Equivariant Graph Hierarchy-based Network (EGHN) to model and represent the dynamics of multi-body systems.
EGHN leverages E-Pool to group the low-level nodes into clusters, and E-UnPool to restore the low-level information from the high-level systems with the aid of the corresponding E-Pool layer.
The fundamental layer of EGHN lies in
Equivariant Matrix Message Passing (EMMP) to characterize the topology and dynamics expressively.
Experimental evaluations on M-complex systems, Motion-Capture, and protein MD, show that EGHN consistently outperforms other non-hierarchical EGNs as well as non-equivariant GNNs.
Acknowledgments and Disclosure of Funding
This work was jointly supported by the following projects: the Scientific Innovation 2030 Major Project for New Generation of AI under Grant No. 2020AAA0107300, Ministry of Science and Technology of the People’s Republic of China; the National Natural Science Foundation of
China (No.62006137); Guoqiang Research Institute General Project, Tsinghua University (No.
2021GQG1012);
Beijing Outstanding Young Scientist Program (No. BJJWZYJH012019100020098);
Tencent AI Lab Rhino-Bird Visiting Scholars Program (VS2022TEG001).
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] In Appendix.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Appendix.
(b) Did you include complete proofs of all theoretical results?
[Yes] In Appendix. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] Yes.
See Sec. 4 and Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Table 1, 2.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Sec. 4.
(b) Did you mention the license of the assets?
[Yes] In Appendix.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]