Weisfeiler Leman for Euclidean Equivariant Machine Learning
This work addresses the problem of designing expressive and efficient equivariant machine learning models for applications like molecular dynamics and point cloud processing, representing a significant but incremental advance over prior invariant methods.
The paper extends the Weisfeiler-Leman hierarchy to prove that a modified graph neural network architecture achieves universal approximation for continuous equivariant functions on point clouds, setting new state-of-the-art results on N-Body dynamics and GEOM-QM9 tasks.
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.