Weisfeiler Leman for Euclidean Equivariant Machine Learning (2402.02484v3)

Abstract: The $k$-Weisfeiler-Leman ($k$-WL) graph isomorphism test hierarchy is a common method for assessing the expressive power of graph neural networks (GNNs). Recently, GNNs whose expressive power is equivalent to the $2$-WL test were proven to be universal on weighted graphs which encode $3\mathrm{D}$ point cloud data, yet this result is limited to invariant continuous functions on point clouds. In this paper, we extend this result in three ways: Firstly, we show that PPGN can simulate $2$-WL uniformly on all point clouds with low complexity. Secondly, we show that $2$-WL tests can be extended to point clouds which include both positions and velocities, a scenario often encountered in applications. Finally, we provide a general framework for proving equivariant universality and leverage it to prove that a simple modification of this invariant PPGN architecture can be used to obtain a universal equivariant architecture that can approximate all continuous equivariant functions uniformly. Building on our results, we develop our WeLNet architecture, which sets new state-of-the-art results on the N-Body dynamics task and the GEOM-QM9 molecular conformation generation task.

• The paper extends the 2-WL test to integrate both position and velocity data within 3D point cloud processing.
• It demonstrates that a Point Pair Graph Network can simulate the 2-WL test uniformly, reducing computational complexity.
• The paper introduces WeLNet, a universally equivariant architecture that achieves state-of-the-art results on key benchmark tasks.

Extension of the 2-Weisfeiler-Leman Test for Euclidean Equivariant Machine Learning

This paper presents substantial advancements in the understanding and application of graph neural networks (GNNs) and their capacity to handle Euclidean data through the lens of the Weisfeiler-Leman (WL) graph isomorphism test hierarchy. Traditionally, the k-Weisfeiler-Leman (k-WL) test is employed to evaluate the expressiveness of GNNs. The research focuses particularly on the 2-WL test, exploring its capability to act as a universal measure when applied to continuous functions on three-dimensional (3D) point cloud data.

The authors extend the applicability of the 2-WL test in three critical dimensions:

1. Integration of Position and Velocity: The work demonstrates that the 2-WL test can be effectively extended to handle point clouds that encompass both particle positions and velocities, a configuration prevalent in many practical applications.
2. Simulation with PPGN: It is proven that the Point Pair Graph Network (PPGN) can simulate the 2-WL test across all point clouds with optimized complexity.
3. Equivariant Architecture Development: With minor modifications to the PPGN architecture, the researchers develop an architecture that is universally equivariant, capable of approximating any continuous equivariant function.

Based on these theoretical extensions, the authors introduce the Weisfeiler-Leman Network (WeLNet), a novel architecture that efficiently processes position-velocity pairs while maintaining complete Euclidean equivariance under permutations, rotations, and translations. WeLNet is not only proven complete in theory but also achieves state-of-the-art results empirically on benchmark tasks, notably the N-body dynamics simulation and the GEOM-QM9 molecular conformation generation task.

Key Contributions and Empirical Results

1. Enhanced 2-WL Applicability: By extending 2-WL to encompass velocity in addition to positional data, the paper addresses a significant gap in the theoretical modeling of dynamic particle systems. This contributes to a more comprehensive interpretation and manipulation of such systems using GNNs, underpinned by a structure that reflects their intrinsic symmetries.
2. Cardinality Reduction in PPGN Simulation: The research illustrates that the PPGN, with a fixed set of parameters, can deliver a uniform separation of a continuous family of 2-WL-separable graphs. This is particularly impactful for handling weighted graphs derived from point clouds, reducing memory and runtime complexity and making the approach computationally feasible for practical deployment.
3. Development of Universal Equivariant Architectures: The insights garnered from extending PPGN capabilities are harnessed to craft an architecture capable of universally approximating continuous equivariant functions. This innovation addresses a notable gap for functions that are both permutation and rotation equivariant, a challenge not previously solved.

WeLNet, through these contributions, signifies an advance in modeling dynamics in Euclidean spaces, offering a promising foundation for future applications in physics simulations and molecular modeling. The experiments conducted verify the theoretical predictions and demonstrate WeLNet's superiority compared to existing benchmarks.

Implications and Speculative Outlook

The implications of this research are vast, both theoretically and practically, in the ongoing development of AI applications handling complex geometric datasets. The presented methodologies offer a refined approach to embedding central Euclidean structures and symmetries, potentially inspiring further investigations and refinements in GNN frameworks.

Speculatively, the methods open pathways for heightened understanding and manipulation of large-scale molecular and physical dynamics within computational domains, aiding advancements in materials science, drug discovery, and theoretical physics. Future work could potentially explore simplification or expansion of these methods to broader dimensions and more complex data structures, cementing the foundational role of GNNs in diverse scientific applications.

In conclusion, this paper stands as a testament to the potential of extending traditional graph theoretical methods into rich domains like point cloud data processing, setting the stage for enriched applicability of machine learning models in scientific explorations involving symmetry-respecting structures.