A Geometric Insight into Equivariant Message Passing Neural Networks on Riemannian Manifolds

In this work, we propose a geometric insight into equivariant message passing on Riemannian manifolds. As previously proposed, numerical features on Riemannian manifolds are represented as coordinate independent feature fields. To any coordinate-independent feature field on a manifold comes canonically attached an equivariant embedding of the principal bundle to the space of numerical features. The minimization of a twisted form of the Polyakov action with respect to the graph of this embedding yields an equivariant diffusion process on the associated vector bundle. By discretizing the diffusion equation flow for a fixed time step, we obtain a message passing scheme on the manifold. We propose a higher-order equivariant diffusion process equivalent to diffusion on the cartesian product of the base manifold. The discretization of the higher-order diffusion process on a graph yields a new general class of equivariant GNN generalizing the ACE and MACE formalism to Riemannian manifolds