On the Expressive Power of Geometric Graph Neural Networks

TL;DR

A geometric version of the Weisfeiler-Leman graph isomorphism test (GWL) is proposed for discriminating geometric graphs while respecting the underlying physical symmetries: permutations, rotation, reflection, and translation.

Abstract

The expressive power of Graph Neural Networks (GNNs) has been studied extensively through the Weisfeiler-Leman (WL) graph isomorphism test. However, standard GNNs and the WL framework are inapplicable for geometric graphs embedded in Euclidean space, such as biomolecules, materials, and other physical systems. In this work, we propose a geometric version of the WL test (GWL) for discriminating geometric graphs while respecting the underlying physical symmetries: permutations, rotation, reflection, and translation. We use GWL to characterise the expressive power of geometric GNNs that are invariant or equivariant to physical symmetries in terms of distinguishing geometric graphs. GWL unpacks how key design choices influence geometric GNN expressivity: (1) Invariant layers have limited expressivity as they cannot distinguish one-hop identical geometric graphs; (2) Equivariant layers distinguish a larger class of graphs by propagating geometric information beyond local neighbourhoods; (3) Higher order tensors and scalarisation enable maximally powerful geometric GNNs; and (4) GWL's discrimination-based perspective is equivalent to universal approximation. Synthetic experiments supplementing our results are available at \url{https://github.com/chaitjo/geometric-gnn-dojo}

Figures (7)

• Figure 1: Axes of geometric GNN expressivity: (1) Body order: increasing scalarisation body order builds expressive local descriptors; (2) Tensor order: higher order tensors determine the relative orientation of neighbourhoods; and (3) Depth: deep equivariant layers propagate geometric information beyond local neighbourhoods.
• Figure 2: Geometric Weisfeiler-Leman Test. GWL distinguishes non-isomorphic geometric graphs $\mathcal{G}_1$ and $\mathcal{G}_2$ by injectively assigning colours to distinct neighbourhood patterns, up to global symmetries. Each iteration expands the neighbourhood from which geometric information can be gathered (shaded for node $i$). Example inspired by schutt2021equivariantmp, ${\mathfrak{G}} = O(d)$.
• Figure 3: Invariant GWL Test. IGWL cannot distinguish $\mathcal{G}_1$ and $\mathcal{G}_2$ as they are $1$-hop identical: The ${\mathfrak{G}}$-orbit of the $1$-hop neighbourhood around each node is the same, and IGWL cannot propagate geometric orientation information beyond $1$-hop.
• Figure 4: Invariant and equivariant functions on geometric graphs. (a) Geometric graphs are attributed graphs embedded in Euclidean space. The geometric attributes transform along with Euclidean transformations of the system, such as global rotations ${\mathfrak{G}}$. (b) The output of ${\mathfrak{G}}$-invariant functions remains unchanged under transformations of the input. (c) For ${\mathfrak{G}}$-equivariant functions, transformations of the input must result in the output transforming equivalently.
• Figure 5: Geometric GNN message passing. ${\mathfrak{G}}$-invariant layers only propagate local scalar quantities such as distances (SchNet, equation \ref{['eq:schnet']}) or distances and angles (DimeNet, equation \ref{['eq:dimenet']}). In contrast, ${\mathfrak{G}}$-equivariant layers propagated geometric quantities such as vectors and relative positions (PaiNN, equation \ref{['eq:painn-v']}) or higher order tensors (Tensor Field Network, equation \ref{['eq:e3nn-1']}).

Theorems & Definitions (80)

• Theorem 1
• Proposition 1
• Proposition 1
• Proposition 1
• Proposition 1
• Proposition 1
• Proposition 1
• Theorem 2
• Proposition 2
• Proposition 2