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

Ilyes Batatia

Keywords: Geometry, equivariance, manifolds, message passing neural network

Abstract:

This work proposes 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 on the
manifold. To any coordinate-independent feature
field on a manifold comes attached an equivariant
embedding of the principal bundle to the space
of numerical features. We argue that the metric
this embedding induces on the numerical feature
space should optimally preserve the principal
bundle’s original metric. This optimality criterion
leads to the minimization of a twisted form
of the Polyakov action with respect to the
graph of this embedding, yielding an equivariant
diffusion process on the associated vector bundle.
We obtain a message passing scheme on the
manifold by discretizing the diffusion equation
flow for a fixed time step. 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 data on Riemannian
manifolds.

Type Of Submission: Proceedings Track (8 pages)

Submission Number: 12