Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Universally Invariant Learning in Equivariant GNNs

Jiacheng Cen, Anyi Li, Ning Lin, Tingyang Xu, Yu Rong, Deli Zhao, Zihe Wang, Wenbing Huang

TL;DR

This work proves that a complete equivariant GNN can be achieved through two key components: a complete scalar function, referred to as the canonical form of the geometric graph; and a full-rank steerable basis set, and proposes an efficient algorithm based on two common models: EGNN and TFN.

Abstract

Equivariant Graph Neural Networks (GNNs) have demonstrated significant success across various applications. To achieve completeness -- that is, the universal approximation property over the space of equivariant functions -- the network must effectively capture the intricate multi-body interactions among different nodes. Prior methods attain this via deeper architectures, augmented body orders, or increased degrees of steerable features, often at high computational cost and without polynomial-time solutions. In this work, we present a theoretically grounded framework for constructing complete equivariant GNNs that is both efficient and practical. We prove that a complete equivariant GNN can be achieved through two key components: 1) a complete scalar function, referred to as the canonical form of the geometric graph; and 2) a full-rank steerable basis set. Leveraging this finding, we propose an efficient algorithm for constructing complete equivariant GNNs based on two common models: EGNN and TFN. Empirical results demonstrate that our model demonstrates superior completeness and excellent performance with only a few layers, thereby significantly reducing computational overhead while maintaining strong practical efficacy.

Universally Invariant Learning in Equivariant GNNs

TL;DR

This work proves that a complete equivariant GNN can be achieved through two key components: a complete scalar function, referred to as the canonical form of the geometric graph; and a full-rank steerable basis set, and proposes an efficient algorithm based on two common models: EGNN and TFN.

Abstract

Equivariant Graph Neural Networks (GNNs) have demonstrated significant success across various applications. To achieve completeness -- that is, the universal approximation property over the space of equivariant functions -- the network must effectively capture the intricate multi-body interactions among different nodes. Prior methods attain this via deeper architectures, augmented body orders, or increased degrees of steerable features, often at high computational cost and without polynomial-time solutions. In this work, we present a theoretically grounded framework for constructing complete equivariant GNNs that is both efficient and practical. We prove that a complete equivariant GNN can be achieved through two key components: 1) a complete scalar function, referred to as the canonical form of the geometric graph; and 2) a full-rank steerable basis set. Leveraging this finding, we propose an efficient algorithm for constructing complete equivariant GNNs based on two common models: EGNN and TFN. Empirical results demonstrate that our model demonstrates superior completeness and excellent performance with only a few layers, thereby significantly reducing computational overhead while maintaining strong practical efficacy.
Paper Structure (38 sections, 18 theorems, 38 equations, 4 figures, 15 tables, 4 algorithms)

This paper contains 38 sections, 18 theorems, 38 equations, 4 figures, 15 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Preliminaries
  5. Rethinking equivariant GNNs: a perspective based on basis construction
  6. Reformulating equivariant GNNs: a perspective based on output space
  7. Reach completeness: canonical form and full-rank basis set
  8. Implementation of Complete Models
  9. Experiment
  10. Expressivity
  11. Performance on Physical Systems
  12. Conclusion
  13. Limitations
  14. Theoretical Analysis and Proofs
  15. Detailed Preliminaries
  16. ...and 23 more sections

Key Result

Theorem 3.2

Given a geometric graph ${\mathcal{G}}$, suppose there is a matrix $\tilde{{\bm{V}}}^{(l)}({\mathcal{G}})$ with $C$ channels of $l$th-degree steerable features denoted as $\tilde{{\bm{v}}}_c^{(l)}({\mathcal{G}})$ satisfying $\mathop{\mathrm{span}}\nolimits(\tilde{{\bm{V}}}^{(l)}({\mathcal{G}}))={\ma

Figures (4)

  • Figure 1: Difference between current methods and our Method. (a) Current methods can be understood as an expansion of multi-body high-degree bases $\tilde{{\bm{b}}}_i^{(l)}$ coupled with the Clebsch–Gordan (CG) tensor product (\ref{['eq:basis_set_sh']}). The completeness is achieved with sufficiently large values of the body order and feature degree. (b) In contrast, our method proposes a novel expansion which does not necessarily require CG tensor product (\ref{['eq:new-expansion']}): the sum over a finite basis set with dynamic weights dependent on the input. The completeness is ensured if the weights (the scalar function $\Gamma({\mathcal{G}})$) are complete and the basis set $\tilde{{\bm{V}}}^{(l)}$ is full-rank.
  • Figure 2: Different vectors are decoupled by coloring, thereby expanding the space of virtual nodes. The yellow area in the figure represents the out space of 1st-degree functions. \ref{['fig:chiral1original']}: All nodes have the same features, and the output can only appear in a straight line in one-dimensional space. \ref{['fig:chiral1center']}: Calculate the distance to the center $\|\vec{{\bm{x}}}_{ic}\|$ to update each node feature so that each node has a different color (feature). \ref{['fig:chiral1colored']}: After coloring, all vectors are decoupled, so the output can be generated in the entire 3D space.
  • Figure 3: Visualization of MSE loss.
  • Figure 4: Examples to test whether the model can recognize chirality. \ref{['fig:inversion', 'fig:mirror']} are modified from GWL-test's implementation joshi2023expressive of IASR-test pozdnyakov2020incompleteness, while \ref{['fig:chiralcounterexample']} is designed by ourselves. \ref{['fig:inversion']}: Test whether models could distinguish this geometric graph with itself after inversion. \ref{['fig:mirror']}: Test whether models could distinguish this geometric graph with itself after mirror about the $xOz$-plane. \ref{['fig:chiralcounterexample']}: The task is same as \ref{['fig:inversion']}, while all nodes are on the unit sphere, and the center point coincides with the center of the sphere. Such counterexample cannot produce separable virtual nodes by simple center distance coloring.

Theorems & Definitions (40)

  • Example 3.1: Basis Set of Common Equivariant GNNs
  • Theorem 3.2: Dynamic Method
  • Definition 3.2: Geometric Isomorphism
  • Definition 3.2: Canonical Form of Geometric Graph
  • Theorem 3.3: General Canonical Form
  • Theorem 3.4: Faster Canonical Form
  • Theorem 3.5: Coloring on Asymmetric Graph
  • Theorem 3.6: Existence of Full-Rank Basis Set
  • Example A.1: Spherical Harmonics
  • Example A.2: Common Physical Quantities
  • ...and 30 more