Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing

We thank the reviewer for the constructive feedback.

Q1/W1 (Supporting the main claim experimentally)
Our main claim is the following: tensorial message passing enhances performance in local canonicalization-based message passing frameworks. In the paper, we introduce tensorial message passing (Sec. 4.2) and explain that our formalism is a strict generalization, an extension, to standard local canonicalization based MPNNs (Fig. 1). In practice, we verify in several experiments on several datasets that indeed tensorial messages (in which we combine tensorial with scalar features) outperform MPNNs that use purely scalar messages (see Tab. 1, Tab. 2 and Tab. 5). Beyond that, Tab. 3 shows that our proposed framework substantially outperforms MPNNs with scalar messages, even when the local frames are chosen completely randomly. This highlights the key advantage of our approach: in many geometric tasks, it is indeed beneficial to enable direct communication of tensorial properties directly during message passing. Therefore, we firmly believe that we provide substantial experimental evidence for our main claim.

We agree that there are indeed more challenging tasks that could further highlight the strengths of our method. To address this, we have conducted a number of experiments on the ScanObjectNN dataset [3], please take a look at the respective overall comment.

W2 (Related equivariant approaches)
Thanks for pointing out that related works are missing in the comparison tables. Since we mention TFN [1] and VNN [2] in the related work, we will add them to the comparison. If the reviewer has more suggestions, we are open to incorporating more models into our comparison.

W3 (Relation to TFN [1] and VNN [2])
We cite the papers TFN [1] and VNN [2] in the paragraph "Equivariance using tensorial internal representations" in the related works section. For convenience, we repeat the differences to our approach here:

• Tensor Field Networks: TFN maintains equivariance throughout the network by employing equivariant layers, such as convolutions based on the tensor product. These layers rely on specialized components, including specific non-linearities or norm layers, to ensure equivariance is preserved. In contrast, our approach, like other methods based on canonicalization (see "Equivariance by canonicalization"), estimates a canonical orientation of the input features, in doing so avoiding architectural constraints imposed by specialized layers. This allows to integrate our framework more easily with existing architectures, in order to make non-equivariant architectures equivariant.

• Vector Neuron Networks: VNN uses vectors within the model, ensuring that these vectors are processed in a way that preserves equivariance. Unlike TFN, VNN does not rely on the tensor product but still requires specialized layers to maintain equivariance. While VNNs are restricted to internal vector representations, our architecture imposes no such limitations, allowing for greater flexibility in the choice of representations.

[1] Thomas et al. "Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds", 2018

[2] Deng et al. "Vector Neurons: A General Framework for SO(3)-Equivariant Networks", Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021

[3] Uy et al. "Revisiting Point Cloud Classification: A New Benchmark Dataset and Classification Model on Real-World Data", IEEE/CVF International Conference on Computer Vision (ICCV), 2019

Thanks to the authors for their efforts in responding to all the reviewers.

Some more experiments are added to the revised version. From a theoretical perspective, I think this paper provides enough content to the community. But from a practical perspective, the problem this paper is dealing with is still a little bit toy and fake, but I believe much progress in real-world problem solving is inspired by theoretical progress. So I raised my score from 5 to 6.

We thank the reviewer for the constructive feedback.

Q1.1 (Local frame estimation with predicted vectors $v_1,v_2$ (close to) zero, cf. Eq. 4)

In practice, the prediction of vectors $v_1,v_2$ close or equal to zero is extremely rare and does not cause practical instabilities. However, we identified two primary cases in which this situation could arise:

1. Nodes with no neighbors: However, if a node has no neighbors, it does not participate in message passing (given the same radius cutoff), rendering the local frame for such a node irrelevant.
2. Highly symmetric neighborhoods: For nodes within a highly symmetric neighborhood, establishing a meaningful local frame becomes infeasible due to the lack of distinguished directions.

To address these situations, we considered two options: either one sets the local frame to zero, excluding the node from tensorial message passing or one replaces the (almost) zero vector with a random vector. We chose the latter option, as it forces the model to learn to disregard the direction of degenerate coordinate axes when no distinguished direction can be inferred.

Q1.2 (Robustness of the local frame prediction)

We are currently investigating the robustness of the estimated local frames and of the full models to noise by adding jitter to the positions of the input point clouds. We will present the results within the next days.

Q2 (Relation to gauge equivariant networks)

Gauge equivariant neural networks, such as the method described in [1], share a similar philosophy with our approach, including parallel transport ("frame-to-frame transitions" in our framework, cf. Fig. 1). However, key differences are:

1. Dimensionality: The method in [1] is specifically tailored for 2D meshes, whereas our architecture generalizes to arbitrary-dimensional Euclidean spaces.
2. Treatment of local frames: In [1], random local frames are chosen for each node, and the model is designed to be equivariant to these frames, ensuring that outputs are invariant to the specific choice of frame. In contrast, our method explicitly leverages meaningful local frames, making the model dependent on the chosen frames. This dependency allows our architecture to utilize the geometric information embedded in these frames, enhancing its expressive power in certain applications.
3. Constrained model architecture: The equivariance to local frames in [1] is achieved through specialized kernel constraints and specialized non-linearities, while our method imposes no restriction on the used architecture. Our approach thus offers access to a much larger design space.

Furthermore, one may draw a connection to one of our ablation studies: in Tab. 3 we present results for our formalism paired with purely random local frames. In this case, the network is forced to predict (approximately) the same outputs irrespective of the choice of local coordinate frames, which can be interpreted as a form of learned gauge invariance.

[1] Pim De Haan, Maurice Weiler, Taco Cohen, Max Welling, "Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs" ICLR, 2021

We thank the reviewer for the constructive feedback.

W1 + W4 (Local frames estimation does not require point normals + PCA-based local frames)
We suspect a misunderstanding. To clarify, the construction of local reference frames in our approach does not depend on point normals. Instead, the vectors used to construct the orthonormal frames are predicted as a learned weighted sum of the relative neighbor positions (see Eq. 4).

If the reviewer finds that an additional explanation in a certain section of the paper could help to prevent a potential misunderstanding, we would gladly refine the text accordingly.

We thank the reviewer for the instructive suggestion to compare against local frames given by local PCA. We have re-trained our best-performing model on classification and on normal vector regression, and have replaced the learned local frames with local frames computed through local PCA (same cutoff radius as for the learned frames).

We find that for the classification task, PCA-based local frames even yield a slight improvement over the learned frames (accuracy: 88.9 vs 88.7) while for normal regression the learned and iteratively refined local frames perform better (cosine similarity: 0.87 vs 0.88).

Generally, we would like to highlight the flexibility in the choice of local frames as a strength of our formalism. Depending on the task and the structure of the input, different equivariant local frames can be superior and our framework can be readily applied to all equivariant estimators of rule-based or learned local frames.

W1 (Robustness to noise and real-world datasets)
We are currently investigating the robustness of the estimated local frames and the full models to noise in the form of adding jitter to the positions of the input point clouds. We will present the results within the next days. Furthermore, we conducted a number of experiments on the real-world dataset ScanObjectNN [1], for the results please see the overall comment.

W2 (Application to feature-based point cloud registration)
We agree with the reviewer that our framework provides an excellent formalism for building models for equivariant prediction of point features to form point correspondences for point cloud registration. We are currently working on several applications of the framework presented in the paper. We see point cloud registration as a promising avenue for future research, though setting up a new pipeline and conducting thorough experiments lies beyond the scope of this rebuttal.

W3 (Tensorial messages outperform scalar messages even with random frames)
Tab. 3 shows that our proposed framework (tensorial messages) outperforms traditional equivariant MPNNs based on local canonicalization (scalar messages), even when the local frames are chosen completely randomly. This is not a contradiction, but really highlights the key advantage of our framework: in many geometric tasks, it is indeed beneficial to enable direct communication of tensorial properties directly during message passing between local frames. In general, we believe that informative local frames improve model performance, but even with random local frames the preservation of tensorial features is beneficial (cf. Fig 1, the vectorial feature is preserved, irrespective of the local frames). The results in Tab. 3 show that the tensorial messages provide a larger gain compared to informative local frames. When the local frames are chosen fully at random, the network must learn to produce the same output irrespective of the local coordinate frame, which draws a connection to works on gauge equivariance (see also Answer to Q2 of reviewer NqWe).

[1] Uy et al. "Revisiting Point Cloud Classification: A New Benchmark Dataset and Classification Model on Real-World Data", IEEE/CVF International Conference on Computer Vision (ICCV), 2019