Equivariant Graph Hierarchy-Based Neural Networks
Authors: Jiaqi Han
Wenbing Huang
Tingyang Xu
Yu Rong
Venue: NeurIPS 2022 Accept (2022)
Keywords: equivariant graph neural network

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 low-level nodes. As their names imply, both E-Pool and E-UnPool are guaranteed to be equivariant to meet physic symmetry. Considerable experimental evaluations verify the effectiveness of our EGHN on several applications including multi-object dynamics simulation, motion capture, and protein dynamics modeling.

Equivariant Graph Hierarchy-Based Neural Networks
Jiaqi Han'; Wenbing Tingyag
Tingyang Xu', Yu Rong
Department of Computer Science and Technology, Tsinghua University Gaoling School of Artificial Intelligence, Renmin University of China Beijing Key Laboratory of Big Data Management and Analysis Methods, Beijing, China Tencent AI Lab
alexhan99max@gmail com, hwenbing@126.ccm com xuty0076gmail..om com, yu. ru..oonghhtmmil.... com
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 com- plex systems particularly, limiting substructure discovery and global information fusion. In this paper. propose Equivariant Hierarchy-base Graph Networks (EGHNs) which consist of the three key components: generalized Equivariant Matrix Message Passing (EMMP), E-Pool and R-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 low- level nodes. As their names imply, both E-Pool and E-UnPool are guaranteed be =quu--quuaa to meet the physical symmetry. Considerable experimental evaluations verify the effectiveness of our BGHN 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 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, re- searchers have paid attention to teaching ma- chine to discover the physical rules from the observational systems through end-to-end train- able neural networks. Specifically, it is natural to use Graph Neural Networks (GNNs), which able to model the relations between different bodies into graph and the inter-body interac- tion as the message passing thereon [1, 16, 25, 26, 21].
Initall Syseen Expett System
Figure I: The folding dynamics of proteins in the cartoon format.
More recently, Equivariant GNNs (EGNs) [29, 8, 27, 12] have become crucial kind of tool for representing multi-body systems. One desirable property is that their outputs are equivariant with
"This work is done when Jiaqi Han works as an intern in Tencent AT Lab.
CCoresponding authors: Wenbing Huang and Yu Rong.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

respect to any translaionooreenttioorreeeeiiii 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.
Hierareky 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 Figure 1, is the dynamics of a protein. Most proteins fold and change in the form of regularly repeating local structures, such as o-helix, 3-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 [13, 5, 33, 3, 18], but these studies focus mainly on generic graph classification, and more importantly, they are not equivariant.
In this paper, we propose Equivariant Graph Hierarchy-basee 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 decoder. The encoder processes the input system from fine-scale to coarse-scale, where an Equívariant-Poolle Pooling (F-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-Per UnPool) layer. Both E-Pool and E-UnPool are equivariant with regard to Euclidean transformations via our specific design. EGHN is built upon generalized equivariant layer, which passes directional matrices over edges other than passing vectors in EGNN [27].
(E-
To verify the effectiveness of EGHN, we have simulated new task extended from the N-body system [16], dubbed M M-copplex system, where each of the M complexes is 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 earry out evaluations on two real applications: human motion caption [4] and the Molecular Dynamics (MD) of proteins [28]. For all tasks, our EGHN outperforms state-of-the-art EGN methods, indicating the efficacy and necessity of the proposed hierarchical modeling idea.²
2
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 [16], and HRN [20] have been introduced reason about the physical interactions. With the development of neural networks enforced by physical priors, many works resort injecting physical knowledge into the design of GNNs. As an example, inspired by HNN [11], HOGN [25] 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 cquivariance [30, 26, 21. 31]. Yet, dealing with rotation equivariance is non-trivial. TFN [29] and SE(3)-Transforme) [8] leverages the irreducible representation of the SO(3) group, while LieConv [7] and LieTransformerr [15] extend the realization of equivariance to Lie group. Apart from these works that resort to group representation theory, succinet equivariant message passing scheme on E(n) group is depicted in EGNN [27]. GMN [14] further involves equivariant forward kinematics modeling particularly for constrained systems. [2] generalizes EGNN to involve covariant information with steerable vectors. [22] leverages frame averaging for general equivariance. [19] mainly studies sign and basis invariance. Despite the rich literature, these models either violate the equivariance,
CCode is available at https://githube www.///waag/////aaaaaa

---

Initial System: GoD
S1L-1)
Target System: g
Encoder
Decoder
E-Pool
EMMP
g(L-2)
E-Pool
EMMP
Pael
EMMP
G(L-1)
Figure 2: Illustration of the proposed EGHN. It consists an encoder and decoder, which are equipped with E-Pool and E-UnPool, respectively. E-UnPool takes as the input the previous output and the score matrix S from E-Pool and output the low-level system
or inspect the system at a single granularity, both of which are vital aspects when tackling highly complicated systems like proteins.
Hierarchical GNNs. There also works that explore the representation learning of GNNs in hierarchies. Several GNNs [13, 5, 32] adopt graph coarsening algorithms to view the graph in multplle granularities. [9] leverages U-net architecture with top-k pooling. Another line of work injects learnable pooling modules into the model. A differentiable pooling scheme DiftPool [33] has been introduced learn permutation-invarian aggregation in DiffPool by node dropping for saving the computational cost. [18] further incorporates self-attention mechanism into the pooling network. [6] leverages junction tree model molecular graph in multiple hierarchies. Nevertheless, these techniques, although permutation equivariant, lack
pooling an end-to-end manner. [3] replaces 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 graph consisting N particles (nodes) V and the interactions (edges) & among them. For each node i, it is assigned with feature tuple (Z) h (o)). where the directional matrix Z R is composed of m n-dimension vectors, such as the concatenation of position æ; E R and velocity E R& leading to z [x;, v; E R3x2, h; E Re is the non-directional feature, such as the category the atom in molecules. The edges are represented by an adjacency matrix A RRNN which can either be constructed according to the geometric distance or physical connectivity. We henceforth abbreviate the entire information of system, i.e., (f (z h Timls A) as the notation Gin if necessary.
(0)
We are mainly interested investigating the dynamics of the input system gin. To formal, given the initial state (Z h (0)) of each particle, our task is to find out a function & to predict its future state z(T) given the interactions between particles. As explored before [29, 8, 7, 27], d is implemented GNN to encode the inter-particl. relation. In addition, it should be equivariant to any teanslatio/refeectooooroaaooo of the input states, so as to obey the physics symmetry about the coordinates. It means, Vg E E(n) that defines the Euclidean group [27].
ollg z, 10) 12 - o( (0) N........
(1)
where g z conducts the orthogonal transformation as RZ for both the position and velocity vectors and is additionally implemented as the translation æ; + b for the position vector; the ellipsis
denotes the input variables uninfluenced by g. including h (0) and A.
As discussed in Introduction, existing equivariant models [29, 8, 7, 27] are unable to mine the hierarchy within the dynamics the input system by flat message passing. To address this pitfall, EGHN is formulated in the encoder-decoder form:
hhigh Encode gom Decode gi
(2)

---

Here, as illustrated in Figure 2, the encoder aims at clustering the particles of G with similar dynamics into group that is treated as the particle in the high-level graph Ghigh (the number of the nodes in Ghigh is smaller than Gd). We have developed novel component, F-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 Ghigh, 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 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 h.) 1i=1> obtain the new features {(Z), h The dimension of the output features could be different from the input, unless the row dimension of Z should (3) keep the same as Z: (i.e. equal to n1). detail, one EMMP layer is updated by Eq. 3-6, where MLP(-) is a Multi-Layer Perceptron, (i)) col- lects the neighbors of i, and Zij E R
EMMP performs information aggregation on the same graph to
H hes - MLP 2Jj (ishhhhhh)) hs))
Mij- HijHij,
(5) (Z.- 2, Z; Z) is a concatenation of the trans- lated matrices on the edge ij. 2 is the mean of (6) all nodes for the position vectors and zero for other vectors. With the subtraction of Z, Zi, is ensured to be translation invariant, and then Z is translation equivariant after the addition of Z; in Eq. 6. Specifically, the MLP in Eq. 3 takes as input the concatenation of the E(n)-invariant 2,J_ Zij. hi, and h2, mapping from R2mx 2m to R 2mxx++ and the output is split into H, E R2mxm and h.j E R The formal proof for the E(u)--qquiaariance of EMMP is deferred to Appendix.
h/- MLP his EjeNo) his
Z = Z + Σ JEN()) My.
Distinct from EGNN [27], the messages to pass in EMMP are directional matrices other than vectors. Although GMN [14] has also explored the matrix form, it is just a specific case of our EMMP by simplifying Zij = Z, - Zj- 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 Eq. 3.
Besides, since taking the inner product might induce a larger variance in the scale of input, in our implementation we also enforce a normalization 2.j 2.g//l 2,j Ziyll || r before feeding the invariant Z.J Zi 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 glow = (f(Z)owv 2low) " 1omh) into an abstract and high-level system ghigh = ( Ih ) K Alghh 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 clus- ter is deemed as node of the high-level system, and its features are computed as weighted combination of the low-level nodes with the scores it just derives. In summary, we proceed: where Eq. 8 maps the invariant fea- ture h² into the score s; E K of (8) cluster assignment with Softmax per- formed long the feature dimension, and the score matrix is given by S = (9) (stg Nx with 8; being its i-th row. By this design, it is tractable to verify that E-Pool is guaranteed to be E(n)
hZ,, h EMMP ( (Z! h N.N. Alow),
Softmax (MLP(h()),
N
hiigh
Sij Z/1
N Σ 8ij
1=
h
high
1
Sig h low
(10) equivariant (also permutation equiv- ariant). Specifically, the division by (11) the riw- sum Σ Siy in Eq. 9
N Σj=1 1 Sij im1
A high
S low S,

---

is essential, as it permits the translation equivariance, that is, Eis ty Σ Siy ( + b)
E Em Sig Z.) + b. This particular property distinguishes our pooling from traditional non-equivariant graph pooling [33, 18]. Notice that the normalization in Eq. 10 is unnecessary since h; is 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- E-UnPool maps the information of the high-level system Ghigh back to the original system space Glow, leading to an output system gam.. We project the features back to the space of the original low-level system by using the transposed scores derived in E- (12) Pool. Then, the projected features along with the low-level features are integrated by an E(n) equiv- ariant function to give the final output. The pro-
Z,
AER
Sij Zj Zhiih
K
h;
Sij 12 high
(13) cedure of E-UnPool is given by Eq. 12-15, where
Zi = 1Z ZZlww Z _ is the column-wiss low-level ma-
==1
concatenation of the mucan-rranslate (14) trix !0w and the high-level matrix 2 analogous to Eq. 3. One interesting point is that Eq. 12 is naturally (15) equivariant terms of translations, even without the normalization term used in Eq. 9. This is because the score matrix is summed to I for each row, indicating that Σø= $tg (Z)ighh + b) Σÿ= Stj Zihish + b. We have the following theorem guaranteeing the equivariance of the designed components, with all proofs deferred to Appendix.
h
MLP
Z Zig low
ace
2, qui + Z,
Theorem 2. EMMP, E-Pool, and E-UnPool are all El)ivvarrraiaan..
3.4 Instantiation of the Architecture
The overall architecture constitutes an equivariant U-Net [24] with skip-connections. We design the overall architecture sequence of EMMP, E-Pool, and E-UnPool in 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 B-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, 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 equivariivr
There is always one EMMP layer prior to each E-Pool or E-UnPool. This external EMMP plays 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-basee
1. The external EMMP exploits a distance-based Aglobal 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 Aglohal is created in lower layer, where the score matrix is obtained by its front E-Pool. 2. For the internal EMMP in E-Pool, it applies connectiviry- based A that exactly reflects the physical connection between particles, for example, it is valued 1 if there is 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 FMMP uses the generalized form in $ 3.1. As we will show in our experiments, such design yields more favorable results compared with using any of A and Apeal only.
re-scored form (akin to Eq. 11) of Agiobul
5

---

The training objective of EGHN is given by:
C
((S g(11))(0) 11
(16)
where || || 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 meen-square-error between the output state Z and the ground truth - The second term is the connectivity loss that encourages more connects within the pooling nodes and less cuts among pooling clusters [34]. For training stability, we first perform row-wise normalization of (S A (())) before substituting into Eq. 16.
4 Experiments
We contrast the performance of the proposed EGHN against variety of baselines including the equivariant and non-equivariant GNNs, on one simulation task: the M 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 --boddy simulation system from [16] and generate the M-complex simulation dataset, in order to introduce hierarchical structures in the data. Specifically, initialize system with N charged particles {æ,, Vt C ) distributed on M disjoint complex objects ر where 2;, Cg are separately the position, velocity, and charge for each particle. Within each complex Sy. 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 { J:" 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. 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 the norm of the velocity |/2:|12, and the edge attribute as c:ty for the edge connecting node i and j. following the setting in [27]. We also concatenate an indicator, which set 1 if a stick presents and 0 otherwise, feature, similar to [14]. We use a fully connected graph (without self-loops) as Agiohal, 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 (Lincar) [27], SE(3)- Trenssoraaf -Sr(3)-Tr.) [8], Radial-Fiell (RF) [17], MPNN [10] and EGNN [27]. For all these models, we employ the codes and architectures implemented by [27]. Detailed hyper- syper-parameter settings are in Appendix.
the edge
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. 2. 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) (5, 10) and (10, 10). We conjecture that, with more particles/compleess, 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 multi- system mode, the performance of EGHN only drops slightly, indicating its adaptability to various scenarios. Visualization. we visualize 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
Results. Table

---

Table 1: Prediction error (x) on various types of simulated datasets. The "Multiple System"
contains J = 5 different systems. For each column, (M,N//M N/M) indicates that each system contains
M complexes of average size N/M. Results averaged across runs. "OOM" denotes out of memory.
Single System
Multiple Systems
(3, 3)
5)
(5, 10)
(10, 10)
(3, 3)
(5)
(5, 10)
(10, 10)
Linear
35.15
2002
30. 100000
31 nooi
35.91 35.29
30.88
22049
TEN [29]
25. 1: 1015
26.01 29.35
20.2
0OM
27.33
29.01 25.57=0.44
OOM
-E(3)-Tr. [8]
27.12 H026
28.87+am
24.48
+0.35
00M
28.144am 28.6644.10
25,00=n28
OOM
MPNN [10]
16.00=an
17.55
16. 1520.s
10
16.762013 17.5820.1
16.55 16.05
RF [17]
g4.
18.37
17.08
15 18.57
18.55
17.24_a.11
19.34
EGNN [27]
12.69 s019
15.37
15.12
14.64
13.33 15,48
15.29
15.02 10L18
EGHN
11.58
14.42
14.29
13.0920mc
12.80
14.85
14.50=ace
13.1 13.1 2002
Figure 3: Visualization on M-complex systems. Left: the prediction of EGNN. Middle: the prediction of EGHN. Right: the pooling results of EGHN with each color indicating a cluster. In the left and middle figure, ground truth is in red, and prediction is in blue. Best viewed by colour printing.
able to detect the correct cluster for each particle. This 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) [16] and running (Subject #9). With regard to walking, we leverage the random split adopted by [14], which includes 200 frame pairs for training, 600 for validation, and another 600 for testing. As for running, with 200/240/240 frame pairs. The interval between each pair is 30 frames in both scenarios. this task the joints are edges and their intersections are the nodes.
follow similar strategy and obtain split
Implementation details. As discussed in [8], many real- world tasks, including our motion capture task here, break the Euclidean symmetry along the gravity axis (2-axit 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 axis, resulting in models that are height-aware while still equivariant in the horizontal directions. This operation is also applied to all basclines. Since the interaction of human body works along the joints, propose to involve the edge in Aglobal if it connects the
Table 2: MSE (x10-2) the motion capture dataset averaged across 3 runs
Subject #35 Subject #9 Walk
Run
MPNN [10]
36.1 +
66.4 +2.2
RF [17]
188.0 +1.9
521.3+2.3
TEN [29]
32.0 =1.8
56.6 11.7
SB(3)-Tr. [8]
31.5 =2.1
61.2 +2.3
EGNN [27]
28.7
50.9 +0.9
nodes within two hops in G. For the number of clusters K, empirically find that K = 5 yields promising results for both walking and running.
GMN [14]
21.6
44.1 12.3
EGHN
8.5 +2.2
25.9 +0.3
Results. Table 2 summarizes the whole results of all models on two subjects. Here, we supplement an additional baseline GMN [14] for promising performance on this task. Excitingly, EGHN outperforms all compared basclines by a large margin on both activities. Particularly, Subject #35,
7

---

23
23
20
15
10
10
-10
-2
Figure 4: Visualization on Motion Capture. Left: the prediction of EGNN. Middle: the prediction of EGHN. Right: the pooling results of EGHN with each color indicating a cluster. the left and middle figure, ground truth in red, and prediction in blue. Best viewed by zooming in.
MSE: 0.925
MSE: 1.920
MSE: 0,314
-10
MPNN
Alpha helix
EGNN (a)
Alpha helix
EGHN
Alpha helix
(b)
Figure 5: Visualization on the MDAnalysis dataset. (a) The predictions of MPNN, EGNN, and EGHN. Ground truth is in red. (b) The pooling assignment of EGHN.
the prediction error of EGHN is 8.5 10-2, whiic is much lower than that of the best baseline, i.e., GMN (21.6 10-2)
Visualization. To investigate why EGHN works, depict the skeletons estimated by both EGNN and EGHN on Subject #9 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 each other but also if they share the similar dynamics.
4.3 Molecular Dynamics on Proteins
We adopt AdK equilibrium trajectory dataset [28] via MDAnalysis toolkit [23] to evaluate our hierarchical model. The AdK equilibrium trajectory dataset involves the MD trajectory of apo adenylate kinase simulated with explicit water and ions NPT 300 K and bar. The atoms' positions of the protein are saved every 240 for total of 1.004 μus as frames.
Implementation details. We split the dataset into train/vallddationnest contain 2481 2481/827/878 frame pairs respectively. We choose 7 = 15 as the span between the input and prediction frames. We ignore the hydrogen atoms focus on the prediction large atoms. We further establish the global adjacency matrix as the neighboring atoms within distance of 10Å. The atoms' velocities of the protein at each frame are computed by subtracting the positions to the previous frame's 's positions. We further leverage MDAnalysis extract the protein backbone in order to reduce the data scale. Even so, TFN and (B633>TTranformed still run out of memory, and thus we compare our model with the rest of baselines. Detailed perem-parsmeters are in Appendix.
along the timeline that

---

Results. The prediction MSE is de- picted Table 3. Our EGHN yields significantly lower error on protein MD compared with the baselines, achieving 1.843 MSE, while the sec- ond 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 apply random rotation of the protein during testing. Compared with EGNN, our EGHN exhibits its superiority thanks to the hierarchical modeling, particularly favorable large and complex systems like proteins.
Table 3: Prediction enror (MSE) on protein MD.
Linear
RF [17]
MPNN [10]
EGNN [27]
EGHN
2.890
2.846
2.322
2.735
1.843
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. (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, the the bottom region, EGHN gives a fitte-nitting result of the alpha helix structure while the predictions from MPNN and EGNN have 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, detects the alpha helix structure in the middle of the protein.
4.4 Ablation Studies
We investigate the necessity of our proposed components on motion capture dataset in Ta- ble 4. We study the following questions:
Table 4: Ablation studies on the motion capture dataset. Numbers are MSE (x10-3).
Subject #35 Walk
Subject #9 Run
Q1. How will the performance of EGHN change, if vary the number of clusters (K)? We modify the number of clusters from 5 to and 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 com- plicated kinematics and will be more difficult to learn if the joints are shared across a too small number of clusters. We provide potential guid- ance on choosing K in Appendix D.I. Q2. How do our proposed two components EMMP and hi- erarchical modeling contribute? We replace all EMMP layers in our model by typical non -quivarrii MPNN, and the performance drops from 8.5 19.7 on Walk, supporting that maintaining equiv- ariance is vital. We further set 8; = 1, 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 larger predic- tion error. This justifies the necessity using the connectivity loss to focus more on intra-cluster connections against the inter-cluster edges. Q4. How about using the same adjucency matrix for all EMMP instead of distinguishing them as Aglobal in the external EMMPs and Asseal in internal EMMPs discussed in 3.4? When we apply Aglotal or Atocal for all EMMPs, the performance drops dramatically, implying that the external and internal EMMPs play different roles EGHN, and should be equipped with different adjacency matrices to model the interactions of different scopes.
EGHN (K 5)
8.5
25.9
EGHN (K = 3)
10.1
41.4
EGHN (K 8)
14.9
26.8
w/o Equivariance w/o Hierarchy Replace by EGNN
19.7
40.9
21.9
42.1
22.3
42.5
w/o Connectivity loss
10.5
28.8
globbal only
17.4
31.5
Alocal only
16.8
33.5
5 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- Network (EGHN) to model and represent the dynamics of multi-body systems. EGHN leverages E-Pool 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 well as non-equivariant GNNs.
Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers for their helpful suggestions. This work was jointly supported by the following projects: the Scientific Innovation 2030 Major Project for New Generation of AT under Grant No. 2020AAA01073000 Ministry of Science and Technology of the People's Republic of China: the National Natural Science Foundation of China
Guoqiang Research Institute General Project, Tsinghua University (No. 202 202IGQG1012) Beijing Outstanding Young Scientist Program (No. BUWZYYY002099100000000 Tencent Al Lab Rhino-Bird Visiting Scholars Program (V$2022TEG001).
References
[1] Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. arXiv preprint 1rXiv:16/2.00222,, 2016. 1,2
[2] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022. 2
[3] Cätälina Cangea, Petar Veličković, Nikola Jovanović, Thomas Kipf, and Pietro Liò. Towards sparse hierarchical graph classifiers, 2018. 2,
[4] CMU. Carnegie-melloe motion capture database. 2003. 2,7
[5] Chenhui Deng, Zhiqiang Zhao, Yongyu Wang, Zhiru Zhang, and Zhuo Feng. Graphzoom:
A multi-level spectral approach for accurate and scalable graph embedding. In International
Conference on Learning Representations, 2020. 2,
[6] M. Fey, J. G. Yuen, and F. Weichert. Hierarchical inter-message passing for learning on molecular graphs. In ICML Graph Representation Learning and Beyond (GRL+) Workhop,
2020. 3
[7] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference Machine Learning, pages 3165-3176 PMLR, 2020. 1,
[8] Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. Se (3)-transformmr:: 3d roto-translation equivariant attention networks. arXiv preprint arXiv:2006. 10500 2020. 1, 2, 3,
6.
[9] Hongyang Gao and Shuiwang Graph u-nets. In Kamalika Chaudhuri and Ruslan Salakhutdi- nov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 2092-2092. PMLR, 09-15 Jun 2019. 3
[10] Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pages 1263-1272. PMLR, 2017. 6, 7.9
[11] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue -uch E. Fox, and R. Garnett, editors. Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
2
[12] Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey, arXiv preprint arXiv:2202.077530 2022. 1
10

---

[13] Fenyu Hu, Yanqiao Zhu, Shu Wu, Liang Wang, and Tieniu Tan. Hierarchical graph convolu- tional networks for semi- supervised node classification. In Proceedings of the Twenty-Eight International Joint Conference Artificial Intelligence, (IJCAI), 2019. 2,3
[14) Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations, 2022. 2, 4, 6,
[15] Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: equivariant self-attention for lie groups. In International Conference Machine Learning, pages 4533-4543. PMLR, 2021. 2
[16] Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems, arXiv preprint 182.:0682..4687,, 2018. 1, 2, 6,
[17] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi-body systems with symmetric energies. arXiv preprint arXiv:
2019. 6, 7,9
[18] Junhyun Lee. Inyeop Lee, and Jaewoo Kang. Self-attention graph pooling. In Proceedings of the 36th International Conference on Machine Learning, 09-15 Jun 2019. 2, 3.5
[19] Derek Lim, Joshua David Robinson, Lingxiao Zhao, Tess Smidt, Suvrit Sra, Haggai Maron, and Stefanie Jegelka. Sign and basis invariant networks for spectral graph representation learning. In ICLR 2022 Workshop on Geometrical and Topological Representation Learning, 2022. 2
[20] Damian Mrowca, Chengxu Zhuang, Elias Wang. Nick Haber, Li Fei-Fei, Joshua B Tenenbaum, and Daniel LK Yamins. Flexible neural representation for physics prediction. arXiv preprint 180i 08047, 2018.
[21] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez- Sanchez-Gonzalez and Peter W Battaglia. Learning mesh-based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020. 1,2 2
[22] Omri Puny, Matan Atzmon, Heli Ben-Hamu, Edward Smith, Ishan Misra, Aditya Grover, and Yaron Lipman. Frame averaging for invariant and equivariant network design. arXiv preprins urXiv:2/10.03336 2021. 2
[23] Richard J. Gowers, Max Linke, Jonathan Barnoud, Tyler J. E. Reddy, Manuel N. Melo, Sean L. Seyler, Jan Domański, David L. Dotson, Sébastien Buchoux, Ian M. Kenney, and Oliver Beckstein. MDAnalysis: Python Package for the Rapid Analysis of Molecular Dynamics Simulations. In Sebastian Benthall and Scott Rostrup, editors, Proceedings of the 15th Python in Science Conference, pages 98 105, 2016. 8
[24] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net:
Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer- intervention, pages 234-241. Springer, 2015.
[25] Alvaro Sanchez- Gonnzlzz, Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint arXiv:1909. 120..222222 12790, 2019. 1,2
[26] Alvaro Sanchez-Gonzalez. Jonathan Godwin, Tobias Pfaff, Rex Ying. Jure Leskovec, and Peter Battaglia. Learning simulate complex physics with graph networks. In International Conference on Machine Learning, pages 8459 8468. PMLR, 2020. 1,2
[27] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E(n) equivariant graph neural networks. arXiv preprint arXiv:2 2021. 1, 2, 3, 4, 7,9
[28] Sean Seyler and Oliver Beckstein. Molecular dynamics trajectory for benchmarking mdanalysis, 6 2017. URL: https://figshare com/artides/Moolcuar__ddoo tory_for pory_for_bencheee www.beehwnnnknnn...bbaiissssssoe
trajec-
doi, 10:m9. 2,8
[29] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Cotation-and translation-equivaiinn neural networks for 3d point clouds. arXiv preprint arXiv:1802.00 022.0822/7 19, 2018 2,7
[30] Benjamin Ummenhofer. Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid sim- ulation with continuous convolutions. In International Conference on Learning Representations, 2019.
[31] Wenxuan WU, Zhongang Qi, and Li Fuxin. Pointconv: Deep convolutional networks on 3d point clouds, 2018. 2
11

---

[32] Yifan Xing, Tong He, Tianjun Xiao, Yongxin Wang, Yuanjun Xiong, Wei Xia, David Wipf, Zheng Zhang, and Stefano Soatto. Learning hierarchical graph neural networks for image clustering. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV), pages 3467-3477, October 2021.
[33] Zhitao Ying, Jiaxuan You, Christopher Morris, Xiang Ren, Will Hamilton, and Jure Leskovec. Hierarchical graph representation learning with differentiable pooling. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. 2, 3,5
[34] Junchi Yu, Tingyang Xu, Yu Rong, Yatao Bian, Junzhou Huang, and Ran He. Graph information bottleneck for subgraph recognition. In International Conference Learning Representations, 2020.
Checklist
1. 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 descrihe 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 experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Yes. See Sec. and Appendix.
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)? [Yes] See Table 1, 2.
(d) Did you include the total amount of compute and the type of resources used (c.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 vou'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, applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent participant compensation? [N/A]
12


Reviews:
[{"review_id": "Y2KoWYs2Njp", "reviewer": "Reviewer_LcX5", "summary": "The authors propose a architecture for E(n) equivariant graph neural networks with two improvements:\n1) Compared to the EGNN, the method includes as node features - besides a set of scalars - not a single E(n)-equivariant vector as feature, but a matrix of M such vectors. The message passing done with an unconstrained MLP that takes as input the invariant 2Mx2M matrix created by taking the pairwise norms of those vectors of both source and target nodes.\n2) The authors propose to use attention to pool to and from finer and coarser hierarchies and carefully normalize the attention weights to ensure equivariance.\n\nThe authors perform a thorough empirical evaluation on datasets with synthethic, motion capture and molecular data and demonstrate state-of-the-art performance.", "questions": "no questions", "limitations": "Contrary to what's stated in the checklist, the authors do not give any limitations of their work. This should happen in the main paper, not the appendix.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Originality\nThe method is an incremental improvement over prior work, so has limited originality.\n\n# Quality\nThe paper is of very high quality. The proposed method is sound and thoroughly evaluated on a diverse set of tasks.\nTwo points for improvement:\n- I would have liked to see a comparison to [1], as they also generalize over the EGNN and claim they improve its performance.\n- I'd think that the hierarchical architecture constitutes a U-Net with skip-connections. Could the authors add this to the text?\n\n[1] Brandstetter, Johannes, et al. \"Geometric and physical quantities improve e (3) equivariant message passing.\" (2021).\n\n# Clarity\nThe paper is very clearly written. Some minor potential points of improvement:\n- The experimental parameter $J$ \"different systems with varying combinations of ...\" is not clearly explained.\n- I'd like a more extensive explanation of the second loss term in (16)\n\n# Significance\nAs the method is easy to implement and improves performance compared to prior work on a domain - equivariant pointcloud prediction - that is widely applicable, this method can have significant impact in the community.\n\n# Conclusion\nThe proposed method is well-motivated, simple and achieves significant performance benefits over prior work. Therefore, I recommend acceptance of this paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VnQIiCQJjY", "reviewer": "Reviewer_2KPJ", "summary": "The paper proposes an architecture to model the dynamics of multi-body systems.\nThe main characteristics are equivariance with respect to Euclidean transformations $g \\in \\textbf{E}(n)$ of the inputs and ability to model the hierarchy present in such systems. This is achieved through EMMP and E-Pool/E-UnPool layers.\nThe main novelty claimed is that this is the first architecture to be _both_ equivariant and hierarchical.", "questions": "- What is the performance of the architecture if the external EMMPs are not relaxed as EGNNs?\n- Can you provide a better evaluation and discussion on the impact of the choice of number of clusters? \"Ablation Studies/Q1\" briefly describes the degradation with $K = 3$. Perhaps it's possible to show a plot for multiple $K$s?\n", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper is well written and describes in detail the architecture.\n- The authors empirically show that this architecture is capable of modeling complex hierarchical systems as their original goal.\n- Clustering of the nodes might help investigate and interpret the results of the architecture.\n\n### Weaknesses\n- The EMMP layer is only a slight variation of EGNN to directional matrices. Additionally EGNNs are still used in the proposed architecture as the external EMMPs. It's unclear whether using the new EMMP layer is beneficial at all.\n- There are many parameters such as number of clusters $k$ and adjacency matrix construction which have to be manually provided. The authors describe how $k$ might affect the performances of the model in the ablation studies but provide no guidance regarding their choice.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "HNMxB28B12L", "reviewer": "Reviewer_4LWy", "summary": "This paper introduces a hierarchical message passing neural network that is equivariant to Euclidean transformations. The building blocks of these networks are i) an equivariant message passing layer that generalises EGNN, ii) a pooling layer that can roughly be seen as an E(3)-equivariant adaptation of Diffpool, and iii) an unpooling layer that uses the clustering coefficients learned by the pooling layers. \nThe resulting architecture (called EGHN) is benchmarked on future-state prediction tasks in several settings. \n", "questions": "The originality of the work is not easy to assess. The pooling and unpooling layers (as well as the reconstruction loss) are in my opinion moderately original and have a strong similarity to Diffpool and Graph-U nets. I however find equations 3) and 4) very interesting, as they generalize EGNNs without requiring to compute spherical harmonics. Are these equations novel or are they borrowed from existing work?\n\n  - If they are novel, this should be clearly stated, as they are in my opinion one of the biggest contributions of the paper. \n     - If possible, you should also give intuition on i) what can be computed with these equations that cannot be computed easily with a EGNN ii) why they are compatible with equivariance. \n    - In the ablation study, it would be useful to assess by how much these equations improve upon EGNN. \n    - These equations should be connected to “Frame averaging for invariant and equivariant network design” (Puny et al. 2022) -- they also have some similarity to “Sign and basis invariant networks for spectral graph representation learning” (Lim, Robinson et al. 2022).\n\n  - If these equations are already used in existing work, this should be clearly stated as well. \n\n\nCould you also provide a comparison of the training time for the different architectures? It seems that your method is more costly than flat MPNNs, but it is not clear by how much.", "limitations": "There is no clear potential negative societal impact to this work. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Summary**: this is a overall a good work, but the ablation study is not properly performed. I am willing to raise my rating if the authors successfully demonstrate the need for the hierarchical architecture.\n\n\n**Strenghts**:\n  - The paper is well written and easy to understand.  \n  - To my knowledge, it is the first work that introduces E(3) equivariant hierarchical message passing networks. \n  - The work is well motivated, and looks particularly relevant to learning tasks on proteins.\n  - The presented experimental results are good, and the presented method seems to consistently outperform previous work. However, it is not clear that this is due to the hierarchical structure, and not to the fact that equation 3) and 4) provide a better representation power than EGNN. \n\n**Weaknesses**: \n\n  - The ablation study is not satisfactory. Since the pooling and unpooling layers contain learnable parameters and update the representation of each node, the hierarchy ablation study should not remove these layers, but rather set the cluster assignments to the identity  ($s_i = \\mathbb 1_i$).\nFurthermore, since the hierarchical aspect is key to your contribution, this ablation should be performed for all datasets, and not only for the motion capture. At first sight, it is for example not so clear that a hierarchical structure will be useful for the n-body system. \n\nThe ablation is all the more important as hierarchical graph neural networks are known to often have disappointing performance (cf. “Rethinking pooling in graph neural networks”, Mesquita et al. 2020). You therefore need to demonstrate that it is key to good performance, at least for some of the tasks. \n  - While most relevant works are cited in the related work section, credit is not really paid where it should be inside Section 3. For example, there is a strong similarity between Diffpool and the equivariant pooling in 3.3., and the two should be compared in details (cf. also the question on Equations 3 and 4 below).\n\n**Some advice**:\n  - It should be clear from the abstract that you are referring to E(3) equivariance. There are other works such as “Multiresolution equivariant graph variational autoencoder” from Hy and Kondor (2021) that deal with permutation equivariance, which can introduce some confusion.\n  - “Hierarchical inter message-passing for learning on molecular graphs” (Fey et al. 2020) is a relevant related work.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept