Equivariant Graph Hierarchy-Based Neural Networks
Anonymous Author(s)
AffiliationAddress email
Abstract
Equivariant Graph neural Networks (EGNs) are powerful in characterizing thedynamics of multi-body physical systems.
Existing EGNs conduct flat messagepassing, which, yet, is unable to capture the spatial/dynamical hierarchy for complex systems particularly, limiting substructure discovery and global informationfusion.
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, EMMPis 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 beequivariant to meet the physical symmetry.
Considerable experimental evaluationsverify the effectiveness of our EGHN on several applications including multi-objectdynamics simulation, motion capture, and protein dynamics modeling.
Introduction
Understanding the multi-body physical systemsis vital to numerous scientific problems, frommicroscopically how a protein with thousandsof atoms acts and folds in the human body tomacroscopically how celestial bodies influenceeach other’s movement.
While this is exactlyan important form of expert intelligence, researchers have paid attention to teaching a machine to discover the physical rules from theobservational systems through end-to-end trainable neural networks.
Specifically, it is naturalto use Graph Neural Networks (GNNs), whichis able to model the relations between differentbodies into a graph and the inter-body interaction as the message passing thereon [1, 14, 23, 24, 19].
More recently, Equivariant GNNs (EGNs) [27, 7, 6, 25] have become a crucial kind of tool forrepresenting multi-body systems.
One desirable property is that their outputs are equivariant withrespect 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 referencecoordinate 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 formulatedin 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 withincomplex systems.
Hierarchy is common in various domains.
Imagine a complex mechanical system, where the particlesare distributed on different rigid objects.
In this case, for the particles on the same object, their statescan be explained as the relative states to the object (probably the center) plus the dynamics of theobject itself.
We can easily track the behavior of the system if these “implicit” objects are detectedautomatically by the model we use.
Another example, as illustrated in Figure 1, is the dynamics of aprotein.
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 onlycharacterizing the conformation of a protein, but also facilitating the propagation between thousandsof atoms in a protein by a more efficient means.
There are earlier works proposed for hierarchicalgraph modeling [11, 5, 31, 3, 16], but these studies focus mainly on generic graph classification, andmore importantly, they are not equivariant.
In this paper, we propose Equivariant Graph Hierarchy-based Network (EGHN), an end-to-endtrainable 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 encoderprocesses 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 thatare considered as the particles of the next layer.
By contrast, the decoder recovers the informationfrom 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 transformationsvia our specific design.
EGHN is built upon a generalized equivariant layer, which passes directionalmatrices over edges other than passing vectors in EGNN [25].
To verify the effectiveness of EGHN, we have simulated a new task extended from the N-bodysystem [14], dubbed M -complex system, where each of the M complexes is a rigid object comprisedof a set of particles, and the dynamics of all complexes are driven by the electromagnetic forcebetween 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 [26].
For all tasks, our
EGHN outperforms state-of-the-art EGN methods, indicating the efficacy and necessity of theproposed hierarchical modeling idea.
Related Work
GNNs for modeling physical interaction.
Graph Neural Networks (GNNs) have been widelyinvestigated for modeling physical systems with multiple interacting objects.
As pioneer attempts,
Interaction Networks [1] have been introduced to reason about the physical interactions, NRI [14]conducts relational inference for the system in a variational manner, and HRN [18] further facilitates the prediction of complex dynamical systems by manually grouping the particles.
With thedevelopment of neural networks enforced by physical priors, many works resort to injecting physicalknowledge into the design of GNNs.
As an example, inspired by HNN [10], HOGN [23] models theevolution of interacting systems by Hamiltonian equations to obtain energy conservation.
Anotherinteresting feature of physical systems lies in Euclidean equivariance, which contains translation,rotation, and reflection.
Several works first approach translation equivariance [28, 24, 19, 29].
Yet,dealing with rotation equivariance is non-trivial.
TFN [27] and SE(3)-Transformer [7] leverages theirreducible representation of the SO(3) group, while LieConv [6] and LieTransformer [13] extend therealization of equivariance to Lie group.
Apart from these works that resort to group representationtheory, a succinct equivariant message passing scheme on E( n ) group is depicted in EGNN [25].
GMN [12] further involves equivariant forward kinematics modeling particularly for constrained systems.
[2] generalizes EGNN to involve covariant information with steerable vectors.
[20] leveragesframe averaging for general equivariance.
[17] mainly studies sign and basis invariance.
Despitethe rich literature, these models either violate the equivariance, or inspect the system at a singlegranularity, both of which are vital aspects when tackling highly complicated systems like proteins.
Conventional Hierarchical GNNs.
There are also works that explore the representation learningof GNNs in a hierarchical fashion.
Several GNNs [11, 5, 30] adopt graph coarsening algorithmsto view the graph in multiple granularities.
[8] leverages a U-net architecture with top- k pooling.
Another line of work injects learnable pooling modules into the model.
A differentiable poolingscheme DiffPool [31] has been introduced to learn a permutation-invariant pooling in an end-to-end
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 𝒢 (#$%) manner.
[3] replaces the aggregation in DiffPool by node dropping for saving the computationalcost.
[16] further incorporates self-attention mechanism into the pooling network.
Nevertheless,these techniques, although permutation equivariant, lack the guarantee of geometric equivariance,which limits their generalization on real-world physical data with 3D information.
The Proposed EGHN
In this section, we first introduce the notations and formulation of our task, and then follow them upby presenting the design of the EMMP layer, which is the basic function in EGHN.
Upon EMMP,we provide the details of how the proposed E-Pool and E-UnPool work.
Finally, we describe theinstantiation of the entire architecture.
[Figure]
3.1 Notations and Formulation
[Figure]
Each input multi-body system is modeled as a graph G consisting of N particles (nodes) V and theinteractions (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 theconcatenation 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 arerepresented by an adjacency matrix A ∈ R N × N , which can either be constructed according to thegeometric distance or physical connectivity.
We henceforth abbreviate the entire information of asystem, 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 itsfuture state Z ( T ) i given the interactions between particles.
As explored before [27, 7, 6, 25], ϕ isimplemented as a GNN to encode the inter-particle relation.
In addition, it should be equivariant toany translation/reflection/rotation of the input states, so as to obey the physics symmetry about thecoordinates.
It means, ∀ g ∈ E ( n ) that defines the Euclidean group [25],
[Equation]
where g · Z (0) i conducts the orthogonal transformation as RZ (0) i for both the position and velocityvectors and is additionally implemented as the translation
x i + b for the position vector; the ellipsisdenotes the input variables uninfluenced by g , including h (0) i and A .
As discussed in Introduction, existing equivariant models [27, 7, 6, 25] are unable to mine thehierarchy 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 similardynamics into a group that is treated as the particle in the high-level graph G high (the number of thenodes 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 theguidance of the high-level system G high , which is accomplished by the proposed E-UnPool.
It isworth mentioning that both E-Pool and E-UnPool, as their names imply, are equivariant, and theyare mainly built upon an expressive and generalized equivariant message passing layer, EMMP.
Tofacilitate the understanding of our model, we first introduce the details of this layer in what follows.
3.2 Equivariant Matrix Message Passing
Given input features { ( Z i , h i ) } Ni =1 , EMMP performs information aggregation on the same graph toobtain the new features { ( Z ′ i , h ′ i ) }
Ni =1 .
The dimension of the output features could be different from
[Equation]
[Equation]
the input, unless the row dimension of Z ′
i shouldkeep 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 ofall nodes for the position vectors and zero forother vectors.
With the subtraction of ¯ Z , ˆ Z ij isensured 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 andh ij ∈ R c .
The formal proof for the E( n )-equivariance of EMMP is deferred to Appendix.
Distinct from EGNN [25], the messages to pass in EMMP are directional matrices other than vectors.
Although GMN [12] has also explored the matrix form, it is just a specific case of our EMMPby simplifying ˆ Z ij
=
Z i − Z j .
Indeed, we have the following theorem for the comparison ofexpressivity 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 ourimplementation 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.3 Equivariant Pooling
Inspired by DiffPool, we propose E-Pool, an equivariant pooling module.
Formally, E-Pool coarsensthe 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 ) } Ki =1 , A high ) with fewer particles, K < N .
For this purpose, we firstperform 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.
Thiscan 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 weightedcombination 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 ofcluster 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 verifythat E-Pool is guaranteed to be E( n )equivariant (also permutation equivariant).
Specifically, the division bythe row-wise sum (cid:80) Ni
=1 s ij in Eq.
9 isessential, as it permits the translationequivariance, that is,
[Equation]
b .
This particular property distinguishes our pooling from traditional non-equivariant graph pooling [31, 16].
Noticethat the normalization in Eq.
10 is unnecessary since h i is a non-directional vector, but it is stilladopted in line with Eq. 9.
In practice, it is difficult to attain desirable clusters by using the SoftMaxlayer solely; instead, the pooling results are enhanced if we regulate the training process with an extrareconstruction loss related to the score matrix, whose formulation will be given in § 3.5.
3.4 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
[Equation]
system by using the transposed scores derived in EPool.
Then, the projected features along with thelow-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 =
[Equation]
− ¯ Z low ; Z agg i − ¯ Z agg ] is the column-wiseconcatenation 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 naturallyequivariant in terms of translations, even without thenormalization term used in Eq.
This is because thescore matrix is summed to 1 for each row, indicating that (cid:80) Kj =1 s ij ( Z high j + b [Equation] =
[Equation]
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.5 Instantiation of the Architecture
The overall architecture constitutes an equivariant U-Net [22] with skip-connections.
We design theoverall 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, asalready defined in § 3.4, 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 aninverse order; for example, in Figure 2, the final E-Pool corresponds to the first E-UnPool.
With suchdesign, 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 playsa different role from the internal EMMP used in E-Pool (Eq. 7).
One crucial difference is thatthey 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 isvalued if the distance between two particles is less than a threshold; by such means, we are able tocharacterize 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 globalin 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 betweenparticles, for example, it is valued 1 if there is a bond between two atoms.
In this way, E-Pool paysmore attention to locally-connected particles when conducting clustering.
Another minor point is thatthe external EMMP is relaxed as EGNN for only modeling the radial interaction, whereas the internal
EMMP uses the generalized form in § 3.2.
As we will show in our experiments, such design yieldsmore favorable results compared with using any one of A global and A local only.
The training objective of EGHN is given by:
[Equation]
[Table]
where ∥ · ∥ F computes the Frobenius norm, L is the number of E-Pools in the encoder, and λ is thetrade-off weight.
The first term is to minimize the mean-square-error between the output state Z out iand the ground truth Z gt i .
The second term is the connectivity loss that encourages more connectswithin the pooling nodes and less cuts among pooling clusters [32].
For training stability, we firstperform 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 theequivariant and non-equivariant GNNs, on one simulation task: the M -complex system, and the tworeal-world applications: human motion capture and molecular dynamics on proteins.
We also carryout 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 [14] and generate the M-complexsimulation dataset, in order to introduce hierarchical structures in the data.
Specifically, we initialize asystem 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 eachcomplex 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 betweenevery 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 particlebelongs to, we will also test if our EGHN can group the particles correctly just based on thedistribution 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 theedge attribute as c i c j for the edge connecting node i and j , following the setting in [25].
Wealso concatenate an indicator, which is set as 1 if a stick presents and 0 otherwise, to the edgefeature, similar to [12].
We use a fully connected graph (without self-loops) as A global , since theinteraction force spans across each pair of particles in the system.
The adjacency matrix A reflects theconnectivity of the particles formed by the complexes.
We set the number of clusters the same as thenumber of complexes in the dataset.
The comparison models include: Linear Prediction (Linear) [25],
SE(3)-Transformer (SE(3)-Tr.)
[7], Radial-Field (RF) [15], GNN and EGNN [25].
For all thesemodels, we employ the codes and architectures implemented by [25].
Detailed hyper-parametersettings are in Appendix.
Results.
Table 1 reports the overall performance of the comparison models on eight simulationdatasets 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 thecomplexity 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 betterwhen the system is changed from (5 , 5) to (5 , 10) and (10 , 10) .
We conjecture that, with moreparticles/complexes, larger systems also provide more data samples to enhance the training of EGHN.3.
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 variousscenarios.
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 accuratelythan EGNN, especially for the large objects.
In the right sub-figure, we also display the poolingresults of EGHN, outputted by the score matrix of the final E-Pool layer.
It is observed that EGHN isable to detect the correct cluster for each particle.
This is interesting and it can justify the worth ofdesigning 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 twoactivities, namely walking (Subject #35) [14] and running (Subject #9).
With regard to walking, weleverage the random split adopted by [12], which includes 200 frame pairs for training, 600 forvalidation, and another 600 for testing.
As for running, we follow a similar strategy and obtain a splitwith 200/240/240 frame pairs.
The interval between each pair is 30 frames in both scenarios.
In thistask the joints are edges and their intersections are the nodes.
Implementation details.
As discussed in [7], many real-world tasks, including our motion capturetask here, break the Euclidean symmetry along the gravity axis ( z − axis), and it is beneficial to makethe equivariant models aware of where the top is.
To this end, we augment the node feature by thecoordinate of the z − axis, resulting in models that are height-aware while still equivariant in thehorizontal directions.
This operation is also applied to all baselines.
Since the interaction of human
[Figure]
body works along the joints, we propose to involve the edge in A global if it connects the nodes withintwo hops in G .
For the number of clusters K , we empirically find that K = 5 yields promising resultsfor both walking and running.
[Table]
Results.
Table 2 summarizes the whole results of all models on two subjects.
Here, we supplement an additionalbaseline GMN [12] for its promising performance on thistask.
Excitingly, EGHN outperforms all compared baselines by a large margin on both activities.
Particularly, on
Subject #35, the prediction error of EGHN is 8 .
5 × 10 −
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 showsthat EGHN is able to capture more fine-grained details oncertain parts ( e.g. the junction between the legs and thebody) than EGNN.
When we additionally visualize thepooling outcome in the right sub-figure, we interestingly find that EGHN is capable of classifyingthe two right-left hands into the same cluster even they are spatially disconnected.
A similar resultis observed for the arms and feet.
This is reasonable as EGHN checks not only if two particles arespatially close to each other but also if they share the similar dynamics.
4.3 Molecular Dynamics on Proteins
We adopt AdK equilibrium trajectory dataset [26] via MDAnalysis toolkit [21] to evaluate ourhierarchical model.
The AdK equilibrium trajectory dataset involves the MD trajectory of apoadenylate 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 thatcontain 2481/827/878 frame pairs respectively.
We choose T = 15 as the span between the inputand prediction frames.
We ignore the hydrogen atoms to focus on the prediction of large atoms.
Wefurther 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 theprevious frame’s positions.
We further leverage MDAnalysis to extract the protein backbone in orderto reduce the data scale.
Even so, TFN and SE( 3 )-Transformer still run out of memory, and thus wecompare our model with the rest of baselines.
Detailed hyper-parameters are in Appendix.
Linear
[Table]
[Table] [Table]
Results.
The prediction MSE is depicted in Table 3.
Our EGHN yieldssignificantly 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 duringtesting.
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 bydifferent 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 indicatethe prediction.
We can observe that EGHN tracks the folding and dynamics of the protein moreprecisely than the baselines.
For example, in the the bottom region, EGHN gives a close-fitting resultof the alpha helix structure while the predictions from MPNN and EGNN have an obvious shiftcompared with the ground truth.
To validate the power of the E-Pool, we further visualize the poolingclusters in Fig.
5 (b).
Interestingly, the pooling assignment exhibits certain clusters in some structuresof the protein.
It suggests that EGHN discovers local repetitive sub-structures of the protein; forinstance, it detects the alpha helix structure in the middle of the protein.
4.4 Ablation Studies
[Table]
EGHN ( K = 8 )
[Table]
[Table]
w/o Equivariance
[Table]
w/o Hierarchy 21.
Replace by EGNN 22.
w/o Connectivity loss 10.
A global only
[Table]
A local only 16.
We investigate the necessity of our proposedcomponents 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 )?
Wemodify the number of clusters K from 5 to 3and 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 difficultto learn if the joints are shared across a too smallnumber of clusters.
We provide potential guidance on choosing K in Appendix D.1.
Q2.
Howdo 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 removinghierarchy 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 EMMPover EGNN. Complete studies are deferred to Appendix D.2 and D.3.
Q3.
How does the connectivityloss (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-clusterconnections against the inter-cluster edges.
Q4.
How about using the same adjacency matrix forall EMMP instead of distinguishing them as A global in the external EMMPs and A local in internal
EMMPs as discussed in § 3.5?
When we apply A global or A local for all EMMPs, the performancedrops dramatically, implying that the external and internal EMMPs play different roles in EGHN, andshould 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 empiricalhyperparameter.
Future works include extending E-Pool to dynamically adjust K for systems withdifferent scales for enhancing the flexibility of the hierarchical model.
Conclusion.
In this paper, we propose Equivariant Graph Hierarchy-based Network (EGHN) tomodel and represent the dynamics of multi-body systems.
EGHN leverages E-Pool to group thelow-level nodes into clusters, and E-UnPool to restore the low-level information from the high-levelsystems 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 EGHNconsistently outperforms other non-hierarchical EGNs as well as non-equivariant GNNs.
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]
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.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental 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.
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
[Table]
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[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]