Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing
This work addresses the need for flexible equivariant models in geometric deep learning, though it is incremental as it builds on existing local canonicalization methods.
The authors tackled the problem of enforcing rotational and reflectional symmetries in geometric deep learning by enhancing equivariant message passing with tensorial messages, achieving state-of-the-art results on normal vector regression and competitive performance on other 3D point cloud tasks.
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.