Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing

TL;DR

This work introduces an $\,\mathrm{O}(d)$-equivariant message passing framework built on local canonicalization, enabling exact symmetry preservation across arbitrary dimensions. By predicting equivariant local frames and transforming features into these local frames, the method achieves invariant node representations and uses tensorial messages to exchange directional information between frames, updating frames iteratively. The approach is demonstrated by adapting PointNet++ to the $\,\mathrm{O}(3)$ group, obtaining state-of-the-art results on normal vector regression and competitive performance on segmentation and classification, while enabling direct comparison with data augmentation baselines. The framework is architecture-agnostic and provides a scalable, principled path to incorporate higher-order geometric information in graph-based learning for 3D and higher-dimensional data.

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.

Figures (5)

• Figure 1: Limitation of scalar message passing. (a) The upper node $j$ is sending a characteristic direction in its neighborhood (encoded in a vector) to an adjacent node $i$. If the local neighborhood of $j$ is rotated, the vector and the equivariant local frame rotate along. The coordinates of the vector are invariant and thus, with scalar message passing, node $i$ will receive the same message, despite the two geometries being different. (b) Tensorial message passing overcomes this limitation so that directional information can be sent consistently.
• Figure 2: Expressive $\boldsymbol{\mathrm O(d)}$-equivariant message passing based on local canonicalization with tensorial messages. Based on the input geometry, one predicts an equivariant local frame $R_i$ at each node $i$. The geometric input node features $F_i$ are transformed from the global reference frame into the respective local frames, yielding coordinates $f_i$ invariant to the choice of global frame. In order to communicate geometric information consistently during message passing, node features are treated as vectors and tensors which are transformed from one local frame into another. After each message passing layer, the local frames are refined to incorporate newly aggregated geometric information. Finally, the geometric node features are transformed back from the local into the global frame to produce an equivariant output.
• Figure 3: Illustration of tensorial message passing between local frames. Node $j$ sends a vectorial feature $f_j$ from its local frame $R_j$ to node $i$ with local frame $R_i$. Through the change of basis, the vectorial information can be received in the other coordinate frame without loss of information. Due to the equivariance of the local frames, the local frame coordinates of geometric objects are invariant under global transformation.
• Figure 4: Data efficiency of built-in equivariance vs. data augmentation.
• Figure 5: Robustness analysis on model performance and local frame estimation. Best-performing models using our framework (learned frames + refining frames + tensor messages) trained with and without input jitter and evaluated on data of varying noise scale. Unsurprisingly, the models trained with jitter are more robust, showing that, in our framework, robustness to noise can be learned from noisy data. The coordinate axes of the local frames are compared against predictions without noise. The first coordinate axis (preserved by Gram-Schmidt) carries most geometric information and is most robust. The second and third coordinate axis may be degenerate in symmetric cases and are consequently less robust. Furthermore, the latter axes also inherit perturbations from the first (and second) axis through the Gram-Schmidt orthogonalization.