Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing

Authors: Peter Lippmann, Gerrit Gerhartz, Roman Remme, Fred A. Hamprecht
Primary area: learning on graphs and other geometries & topologies
Keywords: equivariance, message passing, tensor representation, local frames, geometric deep learning

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.

Decision: accept (poster)
Reviewer scores: 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730668819579

thread_type:
official_review

<thread_object>
creator: reviewer_au99

<content>
summary:
The main claim of the paper is that maintaining local equivariant tensor geometric feature in the graph network is better than first canonicalize the feature in local frame and do invariant message passing.  It also proposes a way to pass tensorial features in graph network. The proposed message passing is experimented on toy shape dataset.

strengths:
- The paper has a smooth and easy-to-flow presentation into equivariance.
- Maintaining local geometric features may be useful in more general settings (but not the examples shown in the paper, see weakness below). Which in the long run may benefit the equivariance community.

weaknesses:
- The main concern is that all the experiments are conducted on rigid objects. However, the reviewer believes that the main advantage of the local geometric features preserving throughout the network is to deal with some non-rigid, multi-body, or deformable objects. Indeed there is no strict equivariance in deformation but it is where the local feature should make a difference. Just as shown in Fig.1 in the paper, the geometric feature should help recognize the pattern of the sub-part when it deforms or move. However, the main experiment is conducted on the modelnet rigid object, which we know the performance is quite saturated, and the reviewer believes that a robust global PCA plus any modern large point network will outperform an equivariant network in such an easy setting.
- Again, the comparison does not capture the full equivariant network baselines. We know that there are many more equivariant point networks compared on the same benchmark but they are not included in the table.
- Some more clear discussion of the difference between the proposed message passing with previous ones like TFN or VNN should be highlighted in the paper.

questions:
The main concern is that the experiments are not convincing for the main claim, some more challenging cases (e.g. multi body objects)  should be included to show the effectiveness of the local geo features
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701267350

thread_type:
official_review

<thread_object>
creator: reviewer_nqwe

<content>
summary:
This work proposes an extension of unconstrained message-passing architectures that makes them equivariant by canonicalizing the messages received by each node to its local frame.  It enables local canonicalization of arbitrary types of tensorial messages, which extends previous works that restrict the allowed tensor type of messages (e.g. only allowing for scalars or vectors).  In addition to the local canonicalization methodology, the authors introduce a mechanism for learning the local frame of each node, which is then refined in the later layers of the network.  In the experimental section, the authors evaluate their proposed method on various rotational equivariant point-cloud tasks and provide ablation studies that showcase how the individual parts of the proposed framework affect its performance and generalization.

strengths:
- The authors describe the proposed framework in detail, providing clear intuition about the specific problems each part of the framework addresses.
- The simplicity of the proposed framework allows it to be easily applied to widely used non-equivariant message-passing architectures with minimal modifications.
- The experimental results demonstrate how the proposed local canonicalization benefits the performance of the baseline model when it is used in tasks with various types of tensorial outputs (e.g.  normal regression or point cloud segmentation).

weaknesses:
- The proposed canonicalization procedure assumes that the inferred output vectors $v_{1},v_{2}$  are non-zero. While the authors describe how they resolve ambiguities when $v_{1},v_{2}$ are parallel, they do not explain how they handle the case where the vectors are close to zero, which makes frame selection highly sensitive to small perturbations due to noise.
-  While in Section 2 the authors mention previous works on local-canonicalization during message passing, they do not discuss work on gauge equivariant neural networks, such as the work:  
 [1] Pim De Haan, Maurice Weiler, Taco Cohen, Max Welling, "Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs" ICLR (2021)  
which also transforms geometric features from one local frame to another during the message passing, performed in their case during the mesh convolution.

questions:
- How does the proposed method handle cases where the predicted vectors $v_{1},v_{2}$ are zero or close to zero? Additionally, how sensitive is the frame selection mechanism when different levels of noise are added to the input point clouds? Does this sensitivity change in cases of more symmetric objects?
- An addition of a discussion of gauge equivariant neural networks will benefit the completeness of the related section of this work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730994765653

thread_type:
official_review

<thread_object>
creator: reviewer_xpaf

<content>
summary:
This paper focuses on the equivariant message passing and proposes a formalism which together with local canonicalization enables consistent communication of geometric features between different nodes. This method solves the problem of communicating geometric information between local patches with different coordinate frames and can be combined with other point cloud methods, achieving state-of-the-art results in the experiments.

strengths:
1. The paper is well-written and technically sound.
2. The paper provides comprehensive theoretical analysis.
3. The experiment results are promising.

weaknesses:
1. The proposed method relies on point normals to establish local reference frames. However, estimating accurate normals is difficult for *real-world* point clouds due to severe noise. So I expect to see results on real-world tasks rather than only synthetic datasets.
2. An important application of invariance and equivariance is point cloud registration. I expect to see the effectiveness of the proposed method on real-world point cloud registration tasks, such as 3DMatch and 3DLoMatch.
3. In Tab.3, the tensor messages surprisingly outperforms the model with scalar messages under random local frames. The random local frames affect the performance of the model in the form of noise, but these noises help tensor messages perform better. I am very curious about its reason.
4. In Tab.2, I notice refining frames brings marginal improvements, which may indicate that this step fails to obtain better normals. For comparison, I expect to see the results with (1) PCA-based normals and (2) ground-truth normals.

questions:
Please address the problems in the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732119670070

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We cordially thank the reviewers for the time invested in reviewing our paper and appreciate their helpful comments.

&nbsp;

**Real-world dataset (W1 xPaF, W1 aU99)**
Reviewer xPaF expressed interest in the applicability to real-world datasets and reviewer aU99 suspects that for (more challenging) real-world datasets our formalism could even demonstrate larger performance improvements compared to the baseline models. To provide experimental evidence that our framework applies favorably also to real-world data, we have now conducted a series of experiments on the ScanObjectNN dataset [1] (model variants and evaluation setup as in the paper, cf. Tab 1). The dataset contains scanned indoor scene data subject to realistic measurement noise and includes deformable objects (e.g. the bag class) and multi-body objects (e.g. a shelf with objects in it). We find our method is consistently better in the $\mathrm{O}(3)/\mathrm{O}(3)$ setting than data augmentation. Furthermore, tensorial messages outperform the scalar messages by a margin, showing that our main claim is also true for noisy real-world datasets.

&nbsp;

ScanObjectNN classification accuracies:

| Method                                                    | z/z  | z/O(3) | O(3)/O(3) |
|-----------------------------------------------------------|:----:|:------:|:---------:|
| Data augmentation                                         | 87.0 | 13.3   | 70.3      |
| Learned frames + scalar messages (ours)                   | 71.9 | 71.9   | 71.9      |
| Learned frames + refining frames + scalar messages (ours) | 73.7 | 73.7   | 73.7      |
| Learned frames + tensor messages (ours)                   | 79.8 | 79.8   | 79.8      |
| Learned frames + refining frames + tensor messages (ours) | 81.0 | 81.0   | 81.0      |

$z/z$ - trained and evaluated with augmentations around the gravitational axis  
$z/\mathrm{O}(3)$ - trained only with rotations around z but evaluated using all O(3) transformations  
$\mathrm{O}(3)/\mathrm{O}(3)$ - trained and evaluated with augmentations all O(3) transformations  

&nbsp;

[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
</content>
</thread_object>

</thread>
<thread>
creation date:
1732291320723

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Once again, we would like to thank all reviewers for their time and thoughtful feedback. We greatly appreciate their constructive comments, which have helped improve the quality of our paper.

In response to all of your helpful comments, we have conducted the following experiments and added them to the paper:
- A series of experiments on the real-world dataset ScanObjectNN. We could validate our main claim that tensorial message passing significantly enhances local canonicalization also on noisy, deformable and many-body data (W1 xPaF, W1 aU99).  
[paragraph from l. 899 + Tab. 6]
- Experiments using PCA-based local frames. The models with PCA-based local frames performed well on all tasks, indicating that local PCA-based frames are a very reasonable choice. It further highlights that our framework can be readily applied with both rule-based or learned local frame estimators. This is crucial for the adaptation to new datasets and domains (W4 xPaF).  
[paragraph from l. 958]
- Investigation of the robustness of our model and our local frame prediction. We find that the local frames and model predictions are reasonably robust to input noise and that robustness can be further improved by training with noisy samples. Together with the experiments on ScanObjectNN this underscores the real-world applicability of our method (W1 xPaF, Q1.2 NqWe).  
[paragraph from l. 909 + Fig. 5]

Additionally, we made the following changes to the paper:
- we discuss the relation to gauge equivariant networks (NqWe Q2). [l. 142]
- we included TFN and VNN as competitive related works for comparison in the results tables (aU99 W2). [Tab. 2 + Tab. 4]

We have highlighted new additional paragraphs and figures in green and refer to their locations in the paper in [...].  

We hope that you find these changes satisfactory. We would be happy to answer any further questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper suggests a way to adapt any message passing architecture to be globally $O(d)$ equivariant by choosing (in an equivariant manner) a local frame at each node, encoding the local $R^d$ feature in this local frame, performing message passing between nodes considering these local representations and frames, and finally moving information back to global frames. 

The paper is well written and suggests a simple algorithm to implement on top of existing message passing architectures. Some critique about this paper was based on: 1) the fact that the method might not be robust to noise (in frame selection process); 2) comparison to relevant previous works that also provide $O(d)$ equivariant networks is missing; and 3) the tasks/dataset considered in the experimental section are a bit too synthetic to exemplify the full merit of the method. In their rebuttal, the authors responded to these points and supported their response with adequate experiments, to the satisfaction of most reviewers.