Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing

TL;DR

This paper introduces a flexible framework for equivariant message passing in geometric deep learning using local reference frames and tensorial messages, achieving state-of-the-art results on 3D point cloud tasks.

Contribution

It presents a novel tensorial message passing framework based on local canonicalization that can be integrated with any architecture to improve equivariance in geometric data processing.

Findings

Tensorial messages outperform previous methods in normal vector regression.

The framework achieves state-of-the-art results on certain 3D point cloud benchmarks.

It can be adapted to existing architectures to enhance equivariance.

Abstract

In numerous applications of geometric deep learning, the studied systems exhibit spatial symmetries and it is desirable to enforce these. For the symmetry of global rotations and reflections, this means that the model should be equivariant with respect to the transformations that form the group of $\mathrm O(d)$. While many approaches for equivariant message passing require specialized architectures, including non-standard normalization layers or non-linearities, we here present a framework based on local reference frames ("local canonicalization") which can be integrated with any architecture without restrictions. We enhance equivariant message passing based on local canonicalization by introducing tensorial messages to communicate geometric information consistently between different local coordinate frames. Our framework applies to message passing on geometric data in Euclidean spaces of arbitrary dimension. We explicitly show how our approach can be adapted to make a popular existing point cloud architecture equivariant. We demonstrate the superiority of tensorial messages and achieve state-of-the-art results on normal vector regression and competitive results on other standard 3D point cloud tasks.