A new perspective on building efficient and expressive 3D equivariant graph neural networks

Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these networks through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding (LSE) and frame transition encoding (FTE). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out the design space for future developments of equivariant graph neural networks. Our codes are available at \url{https://github.com/yuanqidu/LeftNet}.

Paper

References (64)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer b3yr8/10 · confidence 3/52023-06-08

Summary

This paper presents the expressive power of equivariant graph neural networks in the context of 3D local isomorphism from a 2D local isomorphism perspective. Similar to the subgraph isomorphism of GNNs, the paper proposes three types of isomorphism in 3D space: Tree Isometric, Triangular Isometric, Subgraph Isometric. Based on the above definitions, the paper defines geometric Weisfeiler-Lehman (WL) tests. Based on the conclusions in paper [1], this paper enhances the expressive power of equivariant GNN by incorporating mutual 3D substructures. On the other hand, to further enhance the expressive power of the equivariant graph neural network, this paper introduces an SE(3)-invariant encoder to exploit mutual 3D structures. Through a theoretical analysis (theoreom 4.1), the paper demonstrates that simple aggregation of local messages cannot approximate global interactions, thus the authors further enhance the model's expressive power through the frame transition matrix. The effectiveness of the proposed model is demonstrated by experimental results on scalar and vector properties of molecules on QM9. [1] A New Perspective on "How Graph Neural Networks Go Beyond Weisfeiler-Lehman?"

Strengths

1. To the best of my knowledge, this is the first paper to analyze the expressive power of equivariant neural networks from the 3D isomorphism perspective. This could inspire future work to analyze the expressive power of current equivariant neural networks from this perspective. 2. Based on the analysis proposed in the paper on the expressivity of equivariant neural networks, the paper incorporates mutual 3D substructures into 3D GNNs (LSE in this paper) to further enhance their expressive power. The effectiveness of the proposed analysis and methods has been demonstrated experimentally. 3. To reduce the computational cost of neural sheaf diffusion, the paper proposes a novel frame transition matrix to express global geometric information. The computational complexity analysis in line 304 validates the effectiveness of the proposed method.

Weaknesses

1. I have to say that the notation in this paper is confusing, which imposes a burden on me. For instance, h is defined as a node feature on line 81, and f is a tensor-valued function on line 70. However, f is then used as a bijection on line 112. Does this mean the definition of h on line 112 is different from that on line 81? 2. Although the analysis of the expressivity of 3D GNNs is the first from the 3D isomorphism perspective, this paper seems to be based on the analysis in [1]. Moreover, the method proposed by the analysis in [1] can incorporate structural properties into many GNNs, such as GCN, GIN, and GAT. Can this paper similarly incorporate mutual 3D substructures into other 3D GNNs? 3. Recently, some work [2] has analyzed the expressive power of GNNs from the perspective of the attention mechanism, so how can we analyze the expressive power of 3D GNNs from the attention mechanism? [1] A New Perspective on "How Graph Neural Networks Go Beyond Weisfeiler-Lehman?" [2] Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism

Questions

See my comments on weakness.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

In my personal opinion, one of the limitations of this paper is that the notation is a bit confusing, and the paper's layout is somewhat disorganized, making it hard to follow. I would suggest that the authors provide a guideline at the end of the introduction to help readers understand the role of each section.

Reviewer gxFS7/10 · confidence 3/52023-07-06

Summary

This paper introduces local structure encoding and frame transition encoding for more expressive representation learning in 3D graph neural networks. The local structure encoding is inspired by observations in the proposed local hierarchy of 3D graph isomorphisms; specifically the observation that subgraph isomorphism has greater discriminative power than triangle or tree isomorphism. The frame transition encoding is inspired by the observation that not all types of invariant interactions between disjoint clusters can be expressed as functions of the union of those clusters. The authors show that the proposed 3D GNN exceeds the performance of or performs on par with a multitude of existing models.

Strengths

* Originality: The proposed work appears original. The authors introduce a classification of local and global structure, and use it to classify existing 3D GNNs. They also show how to incorporate the most expressive types of local and global structural elements into a new 3D GNN and show that it outperforms/performs on par with existing methods. * Quality/Clarity: The proposed work is of good quality. The paper is well written, preliminary ideas are presented in an accessible way; the local hierarchy of 3D graph isomorphisms is paired with a very nice figure; the work is well motivated; the proposed model outperforms/performs on par with existing methods. * Significance: The problem of building expressive representations in GNNs is challenging and of significance to the community.

Weaknesses

* Quality: see questions

Questions

* SpookyNet is not cited or discussed although it shows very strong performance in the analysis. This method appears to have local and global features, it would be interesting to see where it lands in the proposed classification.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors discuss some limitations of the method

Reviewer 7xD97/10 · confidence 4/52023-07-06

Summary

Authors analize the expressive power of 3D equivariant GNNs and introduce a new expressive equivariant GNN architecture based on local node and edge wise frames.

Strengths

The proposed way to achieve equivariance is well grounded and conceptually very interesting and novel. The construction is also based on solid theoretical motivation and achieved good results in practical experiments. The paper is well written

Weaknesses

The theory about the local hierarchy of 3D Graph Isomorphism (section 3) relatively straightforwardly follows from existing works on equivariant GNN expressive power and subgraph GNN expressive power. That being said, the overall package is still very solid.

Questions

The local frame construction reminded me a bit of Frame Averaging by Puny et al. (ICLR 2022). It would be interesting to see how they stack up.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Authors have addressed the limitations

Reviewer Wt9q5/10 · confidence 3/52023-07-14

Summary

This paper investigates expressive message passing architectures for processing 3D geometric graphs with permutation and Euclidean symmetries. The authors first provide an investigation of a local hierarchy of isomorphism separability. Then, the authors show cases where powerful local invariant models may fall short of encoding global structures due to information loss during propagation, and aim to reduce the information loss by turning to equivariant message passing incorporating frame transition matrices that are related to neural sheaf diffusion. Based on the theoretical framework, the authors propose a message passing architecture that is expressive in terms of the local isomorphism separability as well as incorporating frame transition. Experimental results on invariant molecular property prediction on QM9 and equivariant force prediction on MD17 shows that the proposed approach outperforms previous equivariant neural networks.

Strengths

S1. This paper addresses an important problem of understanding the theoretical expressiveness of geometric graph neural networks and developing architectures that are provably expressive. S2. The local isomorphism hierarchy proposed in Section 3 is novel as far as I can tell (in context of 3D geometric graphs) and correct. S3. The visual example shown in Figure 1 was helpful in understanding the proposed hierarchy of local isomorphism. S4. The connection of the proposed approach to sheaf diffusion mentioned in Appendix I was interesting. S5. The experimental results and ablation study on QM9 and MD17 seems to support the main arguments.

Weaknesses

W1. The equality in Eq. 3 seems not correct as it neglects permutation symmetry. If we consider positional features represented as matrices $\mathbf{X}\in\mathbb{R}^{n\times 3}$ (Line 77), the precise symmetry under consideration is a direct product of permutation group and Euclidean group $S_n \times SE(3)$ (Dym et al., 2020). I think it should be clarified that the equality in Eq. 3 is up to permutation of nodes. W2. In Line 111, I think {-tree, -triangular, -subgraph} should be {tree-, triangular-, subgraph-}. Same applies to Line 603-604 and Line 612. W3. In the definition of tree and triangular isometries (Line 114-118), is there any reason to use the term tree and triangular? I think the definitions are a straightforward extension of the subtree and overlap isomorphism proposed in Wijesinghe et al., 2022, respectively, and I see no reason to use different name. Especially, for the triangular isometry, I found the name misleading since the isometry condition is imposed on the entire set of triangles that share $e_{ui}$, rather than individual triangle. W4. In the definition of triangular isometry (Line 116), the condition is specified for each edge in the intersection of the local subgraphs $e_{iu}\in\mathbf{S}\_{i-j}$. Please correct me if I am wrong, but to make sense, I think the condition should be for each edge in one of the local subgraphs, $e_{iu}\in\mathbf{S}\_{i}$. W5. In Line 606, the term structure completeness is used without proper definition (although it is emphasized by boldface). Without clear definition, I was not able to properly understand and correctness of the argument in Line 608. W6. The argument in Line 162-164 is ambiguous; do you mean a composition of Atoministc Cluster Expansion and invariant map constructed using scalarization by edge-wise equivariant frames is a universal approximator of SE(3) invariant functions, so it can serve as $\phi$ in Theorem 3.1? Is the proof of universality available in literature? W7. For LSE, based on Appendix H, it seems that the hierarchy of local isomorphism proposed in this paper is a subcase of the geometric WL test proposed in Joshi et al., 2023. Since this paper provides 3-way classification while GWL provides a more sensitive separation of expressive power of geometric GNNs based on the orders of node tuple and feature tensors, one can argue that the local isomorphism hierarchy proposed in this paper is less of practical interest compared to GWL. W8. For FTE, in Section 4, frame transition is introduced to solve the limitation of SE(3) invariant messages as outlined in Theorem 4.1. But as far as I know, a majority of practical equivariant geometric GNNs in literature are already using equivariant (tensor) messages (Thomas et al., 2018; Fuchs et al., 2020; Satorras et al., 2022; Geiger et al., 2022; Kohler et al., 2019), and the expressive power of some of them (TFN (Thomas et al., 2018) and SE(3)-Transformers (Fuchs et al., 2020)) are shown to be universal i.e., maximally expressive (Dym et al., 2020). Is the motivation of frame transition also relevant for these architectures? Also, is there specific reasons to favor the proposed frame based propagation over these already equivariant message passing algorithms, despite the added complexity of incorporating node- and edge-wise frames? W9. In Line 214, does the frame here refer to frames in Riemannian manifold theory? Since there is no proper definition (Eq. 5 provides a description of key property but is not a definition), I was confused on what exactly are the frames mentioned here and was unable to verify whether Line 216-221 is correct. It seems the precise definition and descriptions of equivariant frame are completely deferred to Appendix B and F although it serves as a main component of the proposed architecture, which makes the paper hard to read and understand; I think key parts of the proposed algorithm should be described in the main text in a self-contained manner. W10. Overall, I find that the writing of the paper has room for improvement in terms of organization and readability. There are many cross-references back and forth across the paper, so I couldn't serially read through each section without going through a large portion of the entire text multiple text. For example, Line 236-237 (Section 4) refers to Figure 3 (Section 5), Line 555 (Appendix A) refers to Eq. 21 (Appendix F), Line 592 (Appendix C) refers to Figure 3 (Section 5), Algorithm 1 (Appendix C) refers to Eq. 4 (Section 3), and so on. W11. In Eq. 7, I don't see how the equation describes equivariance. I think it should be something like $\mathbf{m}(g\mathbf{x}\_u) = \sum\_{i=0}\^l\mathcal{M}\^i(g)\mathbf{m}\_i(\mathbf{x}\_u)$, is this a typo? W12. Typo in Line 1 of Algorithm 1, gragh -> graph W13. Specification of the output is missing in Algorithm 1, and the output in Algorithm 2 is not clear since normal and boldfaced h are mixed in the algorithm. Furthermore, In Line 7 of Algorithm 1, how do you scalarize the mutual 3D structure, which is a subgraph, based on edge frame $F_{ij}$? I don't see how this can be done in Eq. 8 which defines scalarization. Similarly, in Line 6 of Algorithm 2, how do you scalarize the mutual 3D structure without frame? Also, if the algorithm 2 is LSE only, why should one compute edge-wise frames in Line 5? W14. In the proof of Theorem 5.1, I can see that the proposed parameterization is universal, but does it correctly guarantee equivariance as well (as specified in Theorem 5.1)? The current proof seems to clearly show universality, but I am not sure about equivariance. Dym et al., On the universality of rotation equivariant point cloud networks (2020) Wijesinghe et al., A new perspective on "how graph neural networks go beyond weisfeiler-lehman" (2022) Joshi et al., On the Expressive Power of Geometric Graph Neural Networks (2023) Thomas et al., Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds (2018) Fuchs et al., SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks (2020) Satorras et al., E(n) Equivariant Graph Neural Networks (2022) Geiger et al., e3nn: Euclidean Neural Networks (2022) Kohler et al., Equivariant Flows: sampling configurations for multi-body systems with symmetric energies (2019)

Questions

The questions are included in the weakness section.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have addressed the limitations of their work in Section 7.

Reviewer b3yr2023-08-10

Thanks for your response. For the notation, I think $\varphi$ is ok. My concerns have been addressed. And I suggest to improve the writing and reorganize some parts.

Authorsrebuttal2023-08-10

Thanks

Thanks again for your suggestion! We will make sure to improve organization in the revised version.

Reviewer Wt9q2023-08-11

Some additional questions

I appreciate the response from the authors. I want to ask for some additional clarifications: Regarding optimization benefits in the response to W8, can you elaborate a bit more? In which ways did you observe enhanced stability and faster training for LSE + FTE compared to existing equivariant message passing algorithms? I might have missed it, but cannot find relevant results in the main text or supplementary material. Regarding ease of multi-modality pretraining in the response to W8, can you provide some example-based description? I looked up [1], but couldn't find evidences on how invariant scalar hidden layers can be easier to integrate across modalities.

Authorsrebuttal2023-08-11

Further response

Thanks for the quick response. We would like to make more clarifications on your further concerns. **optimization benefits**: Thanks for posing this question. Here, we empirically find that our LEFTNet is not sensitive with respect to Lr scheduler. We have tried ReduceLROnPlateau, StepLr, warm-up, and LEFTNet achieved quite similar results. While some baseline methods are relatively sensitive with respect to optimization tricks. Note that this empirical finding is informal and may also depend on datasets, and we cannot test all optimization tricks for all baseline methods, so we didn't include it in the paper. In fact, different algorithms may reveal different optimization benefits for different datasets. Therefore, we advocate exploring diverse FT realizations(including our realizations) in Table 4 to navigate optimization complexities. In conclusion, we will emphasize in section 7 that our theoretical analysis in this paper is restricted to the expressiveness power, while the optimization error for different algorithms remains a challenge that is worth future investigation. **integrate across modalities**: In the earlier response, we provided a general reference on multi-modality training. Here, we provide a multi-modality pre-training specific to 3D structures: [2]. In [2], the authors provide a general pretraining framework by maximalizing the mutual information between the 2D modality and 3D modality of molecules. Note that the contrastive learning between the latent representation of 2D and 3D in [2] must be invariant, otherwise, the EBM-NCE contrastive loss (the first part of Eq. 13 of [2]) itself no longer respects the symmetry. Since the 2D representation (usually comes from an ordinary GNN) is automatically invariant, the 3D latent representation paired with the 2D representation should also be invariant. [2] A Group Symmetric Stochastic Differential Equation Model for Molecule Multi-modal Pretraining, ICML 2023

Reviewer Wt9q2023-08-16

Response to the comment

Hello, thank you for the reminder. My major concerns were (in summary) (1) readability and clarity of the manuscript, and (2) the contributions of the proposed framework compared to GWL and existing SE(3)/E(3) equivariant message passing. With the response, I feel that the readability and clarity of the manuscript will be improved, while it might require a non-trivial amount of revisions. For comparison with message passing, the additional reports on practical optimization benefits seems interesting and necessary (since existing tensor message passing is theoretically already universal), and I highly recommend to organize and include such results in the manuscript. I am currently inclining towards acceptance of the paper, except that I expect the above revision to change the manuscript significantly and I am not able to review the revised version in this round. I have updated my score accordingly.

Authorsrebuttal2023-08-16

We will make the suggested revisions as promised in our updated manuscript. We sincerely appreciate your comments to help improve the clarity and quality of our paper!

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC