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

Reviews and author responses:

 Dear Reviewer 2KPJ,

Thank you for the support and comments which help improve the paper!

Best,

Authors

 I thank the authors for their reply. I believe they have fairly address my concerns. Thus, I have increased my score.

 Dear Reviewer LcX5,

Thank you for the support and helpful advice on the paper!

Best,

Authors

 Dear Reviewer 4LWy,

Thank you for the supportive comments and helpful suggestions on improving the paper!

Best,

Authors

 I thank the authors for their reply. My score remains the same.

 The authors have addressed my main concern, which was to demonstrate that the hierarchical scheme can be beneficial. I believe that this is a good paper, and will raise my grade accordingly to 7 (Accept). 



 
> **Q1: I however find equations 3) and 4) very interesting, as they generalize EGNNs without requiring computing spherical harmonics. Are these equations novel or are they borrowed from existing work? If they are novel, this should be clearly stated, as they are in my opinion one of the biggest contributions of the paper.**

Thank you for recognizing our contribution! Yes, these equations are novel.
As the vital part of EMMP, equations (3) and (4) construct the messages as directional matrices other than vectors. This construction provides a generalized form of EGNN, and is theoretically guaranteed to enhance the model expressivity (c.f. Theorem 1 and its proof in Appendix A.1). The extensive ablation studies below verify the effectiveness of such design.

> **If possible, you should also give intuition on i) what can be computed with these equations that cannot be computed easily with an EGNN ii) why they are compatible with equivariance.**

Thanks for raising this valuable point. The intuition lies in: 
i) The term $\hat{\mathbf{Z}}\_{ij}$ in Eq. (3,4) will degenerate to a relative coordinate $\mathbf{r}\_{ij}=\mathbf{x}\_i-\mathbf{x}\_j$ for EGNN, which means the message $\mathbf{M}\_{ij}$ in Eq. (4) can only be colinear with the direction of $\mathbf{r}\_{ij}$. This will make $\mathbf{M}\_{ij}$ unable to simulate the case when the interaction between particles is vertical to $\mathbf{r}\_{ij}$, such as the torque incurred by angular rotations. In contrast, by passing the message related to $\hat{\mathbf{Z}}\_{ij}$, EMMP is able to simulate all directions of interaction, exhibiting stronger expressivity. 
ii) EMMP still maintains equivariance since the inner-product $\hat{\mathbf{Z}}\_{ij}^\top\hat{\mathbf{Z}}\_{ij}$ is invariant and then $\mathbf{H}\_{ij}$ is invariant, thus $\mathbf{M}\_{ij}$ is equivariant. More details are provided in Appendix A2. 



> **In the ablation study, it would be useful to assess by how much these equations improve upon EGNN.**

Thanks for the advice. We add extra ablation studies to the experiments. The results are illustrated below.


|                        | MocapWalk | MocapRun | Complex (3, 3) | Complex (5, 5) |
| --------------------   | --------- | -------- | ---------- | ---------- |
| EGHN                   | **8.5**       | **21.9**     | **11.58**      | **14.42**      |
| + replace EMMP by EGNN | 22.3      | 42.5     | 12.51      | 15.77      |

The results indicate that EMMP, as a building block of EGHN, offers considerable improvements over EGNN in different scenarios.

> **These equations should be connected to “Frame averaging for invariant and equivariant network design” (Puny et al. 2022) -- they also have some similarities to “Sign and basis invariant networks for spectral graph representation learning” (Lim, Robinson et al. 2022).**

Thanks, the raised papers are related and have been discussed in the revision. While Puny et al. (2022) discuss equivariance based on group average in a general sense and Lim, Robinson et al. (2022) focus mainly on sign and basis invariance, our work aims at deriving E(n)-equivariant graph networks based on typical computations such as inner-products. Moreover, both Puny et al. (2022) and Robinson et al. (2022) require computing the bases via spectral decomposition. We have added discussions of these works in the paper. Please kindly refer to the related work in the revised version. 


> **Q2: Could you also provide a comparison of the training time for the different architectures?**

We evaluate the training time on simulation and motion capture datasets for the baselines and EGHN. The table below depicts the average training time per epoch (in seconds). All models are trained on a NVIDIA V100 GPU.


|                | MPNN | TFN  | SE3-Tr. | EGNN | GMN  | EGHN |
| ----------     | ---- | ---- | ------- | ---- | ---- | ---- |
| Complex (3, 3) | 1.21 | 7.81 | 23.25   | 1.45 | 1.58 | 1.69 |
| MocapWalk      | 0.92 | 6.85 | 18.96   | 1.21 | 1.49 | 1.41 |

EGHN is almost as efficient as EGNN and GMN, while only adding marginal computational overhead compared to MPNN, since the computations related to equivariance and pooling are efficient. The irreps-based methods TFN and SE(3)-Transformer yield significantly longer training time.

 We thank the reviewer for the detailed feedback and helpful advice!

> **Q1 in Quality: I would have liked to see a comparison to [A], as they also generalize over the EGNN and claim they improve its performance.**

We extra evaluate SEGNN [1] on $M$-complex dataset in both single-system and multiple-system scenarios. The results are presented below.

Single-system:
| Complex | (3, 3) | (5, 5) | (5, 10) | (10, 10) |
| ----- | ----- | ----- | ------- | -------- |
| EGNN  | 12.69 | 15.37 | 15.12   | 14.64    |
| SEGNN | 14.04 | 15.62 | 15.01   | 14.31    |
| EGHN  | **11.58** | **14.42** | **14.29**   | **13.09**    |


Multi-system:
| Complex | (3, 3) | (5, 5) | (5, 10) | (10, 10) |
| ----- | ------- | ------- | ---------- | ----------- |
| EGNN  | 13.33   | 15.48   | 15.29      | 15.02       |
| SEGNN | 13.88   | 16.01   | 15.41      | 14.78       |
| EGHN  | **12.8**    | **14.85**   | **14.5**       | **13.11**       |

SEGNN performs better than EGNN particularly when the system is large (e.g., on (5, 10) or (10, 10). Still, EGHN consistently outperforms these baselines by a significant margin.


[1] Brandstetter, Johannes, et al. "Geometric and physical quantities improve e (3) equivariant message passing." (2021).

> **Q2 in Quality: I'd think that the hierarchical architecture constitutes a U-Net with skip-connections. Could the authors add this to the text?**

Thank you for the advice! Yes, our architecture constitutes an equivariant U-Net with skip connections. We have revised the manuscript to reflect this point. 

> **Q1 in Clarity: The experimental parameter $J$  "different systems with varying combinations of ..." is not clearly explained.**

Thanks for pointing it out! We have made this point clearer by adding this explanation to the manuscript: Specifically, 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$.

> **Q2 in Clarity: I'd like a more extensive explanation of the second loss term in (16)**

We are glad to provide more explanations for the connectivity loss. Intuitively, the connectivity loss encourages pooling assignments with more edges within the pooled clusters and fewer in between. In particular, the loss reaches its minimum, i.e., 0, if and only if node $i$ and $j$ belong to the same cluster for each edge $(i, j)\in\mathcal{E}$.  We have added the above explanations including case studies in Appendix for better illustration.



 We thank the reviewer for the detailed reviews and helpful suggestions!

> **Weakness 1: The hierarchy ablation study should not remove these layers, but rather set the cluster assignments to the identity. Furthermore, since the hierarchical aspect is key to your contribution, this ablation should be performed for all datasets, and not only for the motion capture.**

Thank you for pointing it out! We agree with the reviewer on the setup of this important ablation study. Below we summarize the results of more ablation studies on all datasets (simulation, mocap, and protein), where `EGHN w/o hier` is implemented by setting the cluster assignment to identity, i.e., $\mathbf{s}_i=\mathbb {1}_i$, per the reviewer's advice.

|               | Complex (3,3) | Complex (5,5) | Mocap Walk | Mocap Run | Protein MD |
| ------------- | --------- | --------- | ---------- | --------- | ------- |
| EGHN          | **11.58**     | **14.42**     | **8.5**        | **25.9**      | **1.84**   |
| EGHN w/o hier | 12.24     | 15.18     | 21.9       | 42.1      | 2.00   |

As illustrated in the table above, the hierarchical structure is consistently beneficial to the model performance across $M$-complex simulation, Motion Capture, and Protein MD. This supports the validity and efficacy of our designed equivariant hierarchy module.

> **At first sight, it is for example not so clear that a hierarchical structure will be useful for the n-body system.**

For $N$-body systems previously studied in [1] where particles are separated and only interact by Coulomb or gravitational force, hierarchy may not have much impact. 
However, in the $M$-complex dataset discussed in this paper, the particles are connected with sticks to form multiple objects, i.e., polyhedrons. In physics, the dynamics of these objects can be decomposed into two parts: (1) the movement of the center of mass (CoM), and (2) the movements of the particles with respect to the CoM. This indicates that each object forms a certain hierarchy and there is crucial geometrical information shared within the same object. In our experiments, we also verify this intuition by showing that EGHN achieves much lower MSE than EGHN without hierarchy (the table above in the response to Weakness 1) and its pooling assignment exactly reflects the belonged object of each node (Figure 3).

[1] Satorras et al. E(n) equivariant graph neural networks.

> **Weakness 2: 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.**

We agree that our equivariant pooling shares a similar idea to DiffPool: both models learn the pooling assignment by themselves in a differentiable, end-to-end manner. We have revised the manuscript, in Sec. 3.3, to pay sufficient credit to DiffPool, which states: Inspired by DiffPool that learns the clustering assignment by the model in an end-to-end fashion, we propose E-Pool, an equivariant pooling module that learns to coarsen the input purely from data.

Meanwhile, it's worth noticing that the differences between our E-Pool/E-UnPool and DiffPool are clear.
* Both our E-Pool and E-UnPool are equivariant. The equivariant constraint is non-trivial to inject into a differentiable pooling module. For example, we require to derive equivariant message passing in Eq. (7,14,15)and consider specific normalization to ensure translation equivariant in Eq. (10).
* We leverage E-Pool and E-UnPool in an encoder-decoder style, whereas DiffPool focuses on graph-level prediction tasks with no decoders applied in the original work. 

On the experimental side (Table 4), we have also compared with "DiffPool" in the ablation study, where we replace the equivariant layers with plain MPNNs. The resulting model (dubbed `w/o Equivariance`) can be viewed as an instantiation of DiffPool on our 3D graph data. The results indicate that our EGHN yields significant enhancement over this variant.

> **Advice 1: It should be clear from the abstract that you are referring to E(3) equivariance.**

Thanks. We have revised the abstract by stating our model is E($n$)-equivariant.

> **Advice 2: Fey et al. 2020 is a relevant related work.**

Thanks. We add discussions of this paper in Related Work.


 We appreciate the reviewer for the insightful comments! We have conducted extra experiments to address the concerns.

> **Weakness 1 & Q1： It's unclear whether using the new EMMP layer is beneficial at all. What is the performance of the architecture if the external EMMPs are not relaxed as EGNNs?**

We add this ablation study that compares the performance of different choices between internal/external EMMP/EGNN. The experimental results are exhibited below.

| Internal | External | MocapWalk | MocapRun | Complex (3, 3) | Complex (5, 5) |
| -------- | -------- | --------- | -------- | ---------- | ---------- |
| EGNN     | EGNN     | 22.3      | 42.5     | 12.51      | 15.77      |
| EMMP     | EGNN     | 8.5       | 21.9     | **11.58**      | 14.42      |
| EMMP     | EMMP     | **8.1**       | **21.1**     | 11.82      | **14.36**      |

We have the following observations for the experiments here:
* When applying EMMP in either internal or external message passing, the performance consistently improves against EGNN. This verifies that the proposed EMMP is potentially more advantageous for modeling interactions, which aligns with our theoretical analysis that EMMP is more expressive than EGNN (c.f. Theorem 1).
* Compared with external EMMP, more significant improvements are obtained when applying EMMP as the internal message passing layers (e.g., 22.3 $\rightarrow$ 8.5 on MocapWalk). Note that the internal message passing layers are those right before our pooling layer, which are responsible for passing and aggregating messages towards the high-level cluster nodes. Therefore, we speculate the reason might be that compared with the flat message passing layers (the external EMMPs), the internal EMMPs require much higher expressivity and capacity since they need to fuse the message of all nodes towards their corresponding cluster nodes.
* In the Complex (3, 3) scenario, changing from EGNN to EMMP in external message passing slightly affects the performance, probably because the interactions between nodes in $M$-complex are Coulomb forces which can be well covered by EGNN. Nevertheless, on the mocap dataset where interactions are much more complicated, leveraging EMMP is consistently more advantageous over EGNN.



> **Weakness 2 & Q2: Can you provide a better evaluation and discussion on the impact of the choice of the 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?**

Nice advice! We thoroughly investigate how the number of clusters influences the model performance on all datasets. 
 - For $M$-complex System, we sweep over 1 to 5 in the Complex (3, 3) single system.
 - For Mocap dataset, we sweep over 1 to 8.
 - For Protein MD, we vary $K$ from 1, 5, 10, 15, 20, 25. The results are depicted below.

Here are the results:

On $M$-complex (3,3) simulation (9 nodes):

| $K$  | 1     | 2     | 3         | 4     | 5     |
| ---- | ----- | ----- | -----     | ----- | ----- |
| MSE  | 14.86 | 13.21 | **11.58** | 12.05 | 12.92 |


On Mocap Walk (31 nodes):
| $K$  | 1    | 2    | 3    | 4        | 5    | 6    | 7    | 8    |
| ---- | ---- | ---- | ---- | ----     | ---- | ---- | ---- | ---- |
| MSE  | 19.8 | 16.8 | 10.1 | **8.1**  | 8.5  | 10.5 | 11.2 | 14.9 |


On Protein MD (855 nodes):

| $K$  | 1     | 5     | 10    | 15        | 20    | 25     | 
| ---- | ----- | ----- | ----- | -----     | ----- | -----  |
| MSE  | 2.132 | 2.234 | 2.127 | **1.843** | 2.189 | 2.2452 |

We also plot the corresponding figures and add them to Figure 8 in Appendix D.1.
From the above tables, we observe that: 
 - On all datasets, the performance degenerates when $K=1$, since all nodes in the system are pooled into one cluster and therefore there are no learnable cluster assignments. It verifies the necessity of modeling hierarchies in multi-body systems.
 - The systems with a larger scale enjoy larger $K$ in practice. It indicates that for the systems with a larger number of nodes, it is beneficial to choose larger $K$ to better model their complex hierarchies.
 - For the Complex (3,3) system, it is interesting that the best performance is obtained when $K=3$, since it contains 3 disjoint complexes. This implies that it is also possible to choose $K$ by some prior knowledge assessed from data. 

We are glad to offer more explanations if there is still any question!

 The paper proposes an architecture to model the dynamics of multi-body systems.
The 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.
The main novelty claimed is that this is the first architecture to be _both_ equivariant and hierarchical. ### Strengths
- The paper is well written and describes in detail the architecture.
- The authors empirically show that this architecture is capable of modeling complex hierarchical systems as their original goal.
- Clustering of the nodes might help investigate and interpret the results of the architecture.

### Weaknesses
- 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.
- 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. - What is the performance of the architecture if the external EMMPs are not relaxed as EGNNs?
- 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/A

 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. 
The resulting architecture (called EGHN) is benchmarked on future-state prediction tasks in several settings. 
 **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.


**Strenghts**:
  - The paper is well written and easy to understand.  
  - To my knowledge, it is the first work that introduces E(3) equivariant hierarchical message passing networks. 
  - The work is well motivated, and looks particularly relevant to learning tasks on proteins.
  - 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. 

**Weaknesses**: 

  - 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$).
Furthermore, 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. 

The 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. 
  - 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).

**Some advice**:
  - 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.
  - “Hierarchical inter message-passing for learning on molecular graphs” (Fey et al. 2020) is a relevant related work.


 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?

  - If they are novel, this should be clearly stated, as they are in my opinion one of the biggest contributions of the paper. 
     - 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. 
    - In the ablation study, it would be useful to assess by how much these equations improve upon EGNN. 
    - 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).

  - If these equations are already used in existing work, this should be clearly stated as well. 


Could 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. There is no clear potential negative societal impact to this work. 

 The authors propose a architecture for E(n) equivariant graph neural networks with two improvements:
1) 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.
2) The authors propose to use attention to pool to and from finer and coarser hierarchies and carefully normalize the attention weights to ensure equivariance.

The authors perform a thorough empirical evaluation on datasets with synthethic, motion capture and molecular data and demonstrate state-of-the-art performance. # Originality
The method is an incremental improvement over prior work, so has limited originality.

# Quality
The paper is of very high quality. The proposed method is sound and thoroughly evaluated on a diverse set of tasks.
Two points for improvement:
- I would have liked to see a comparison to [1], as they also generalize over the EGNN and claim they improve its performance.
- I'd think that the hierarchical architecture constitutes a U-Net with skip-connections. Could the authors add this to the text?

[1] Brandstetter, Johannes, et al. "Geometric and physical quantities improve e (3) equivariant message passing." (2021).

# Clarity
The paper is very clearly written. Some minor potential points of improvement:
- The experimental parameter $J$ "different systems with varying combinations of ..." is not clearly explained.
- I'd like a more extensive explanation of the second loss term in (16)

# Significance
As 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.

# Conclusion
The proposed method is well-motivated, simple and achieves significant performance benefits over prior work. Therefore, I recommend acceptance of this paper. no questions 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.

Meta-review: All reviewers agree that this work makes progress on the front of E3-equivariant message-passing neural networks. Though not surprising in light of previous work (e.g., on graph coarsening and pooling for GNNs), the introduced modifications appear to be novel in the context of EGNN. The proposed ideas are well motivated and explained while also yielding numerical benefits.

Decision: Accept