Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing

## Abstract

Abstract In many applications of geometric deep learning, the choice of global coordinate frame is arbitrary,
and predictions should be independent of the reference frame. In other words, the network should be equivariant with respect to rotations and reflections of the input, i.e. the transformations of O ⁢ ( d ) O 𝑑 \mathrm{O}(d) roman_O ( italic_d ) . We present a novel framework for building equivariant message passing architectures and modifying existing non-equivariant architectures to be equivariant. Our approach is based on local coordinate frames, between which geometric information is communicated consistently by including tensorial objects in the messages. Our framework can be applied to message passing on geometric data in arbitrary dimensional Euclidean space. While many other approaches for equivariant message passing require specialized building blocks, such as non-standard normalization layers or non-linearities, our approach can be adapted straightforwardly to any existing architecture without such modifications. We explicitly demonstrate the benefit of O ⁢ ( 3 ) O 3 \mathrm{O}(3) roman_O ( 3 ) -equivariance for a popular point cloud architecture and produce state-of-the-art results on normal vector regression on point clouds.

## 1 Introduction

Point cloud data is ubiquitous. Scans of 3D scenes, molecules, astrophysical simulations and earth science data can all be represented as nodes with features and positions in Euclidean space. Message passing neural networks are often used to extract and combine these node features and are used in various tasks.
In application domains in which inputs and outputs are governed by known symmetries, it may be desirable or required to enforce these.
One such approach is to build equivariant architectures in order to learn a function that behaves in a well-defined manner under transformations of the input, e.g. rotations and reflections which together form the group of $\mathrm{O}(d)$. One special case of equivariance is invariance, in which the output does not change if a transformation is applied to the input. Clearly, for a consistent prediction of molecular energies or normal vectors of 3D surfaces, the outputs must transform according to a well-defined geometric rule. For instance, a normal vector should rotate along as the 3D object is rotated. In other words, the choice of the coordinate system (which is often arbitrary) should not affect the predictions. While the idea of equivariance is conceptually appealing, it is also practically relevant: Built-in equivariance is known to enhance the performance and data efficiency of neural networks in several settings (Weiler et al., ; Batzner et al., ). For instance, in deep learning based molecular dynamics simulations, exact equivariance can be crucial for the simulation stability (Fu et al., ).

Much work has gone into the development of specific architectures, designed to ensure equivariance under rotations and reflections. Existing approaches rely on specialized non-linearities, norm layers and tensor products to achieve equivariance. Although $\mathrm{O}(d)$-equivariance can be an extremely helpful prior in many deep learning pipelines, this may so far have hindered practitioners from using equivariant architectures. With the framework presented in this paper, we aim to address this issue and close the gap by providing a novel, practical perspective on building $\mathrm{O}(d)$-equivariant architectures. Concretely, we make the following contributions:

- •
We present a novel framework for building $\mathrm{O}(d)$-equivariant message passing networks based on local frames, which enables the communication of geometric information in the form of tensorial messages.
- •
We show explicitly how our framework can be used to modify any existing message passing architecture to make it $\mathrm{O}(d)$-equivariant, without restricting the architecture to specialized building blocks.
- •
As a concrete example, we present an $\mathrm{O}(3)$-equivariant version of the widely-used PointNet++ architecture (Qi et al., ), which produces competitive and sometimes even state-of-the-art results.
- •
Our framework allows for a direct comparison between equivariant architectures and non-equivariant ones. In the comparison, we find that exact, built-in equivariance via Tensor Frames is more data-efficient and outperforms our strong baseline models trained via data augmentation.

Figure: Figure 1: Tensor Frames formalism for exact $\mathrm{O}(d)$-equivariant message passing. Based on the input geometry, one predicts a local frame at each node in an equivariant fashion. The geometric node features $F_{i}$ are transformed from the global frame of reference into the local frames $R_{i}$, yielding coordinates $f_{i}$ invariant to the choice of global frame. During message passing, geometric information from the neighborhood of each node is combined into updated node features. The key mechanism is that by transforming the geometric features from one local frame into the other, the messages are able to communicate geometric objects in the form of vectors and tensors. Finally, the geometric node features are transformed back from the local into the global frame to produce an equivariant output.
Refer to caption: x1.png

## 2 Related work

#### Approximate equivariance.

Many popular architectures for learning on point cloud data rely on message passing (Qi et al., ; Liu et al., ; Wang et al., ).
In general, the simplest way to achieve approximate equivariance with respect to a set of transformations is to augment the training data, i.e. present the model with randomly transformed data samples during training. Data augmentation is of course completely independent of the model and does not constrain the architecture. However, the equivariance must be learned and is thus not guaranteed, meaning that it may not generalize to out-of-distribution samples; and the learned equivariance is not exact.

An alternative approach to achieve invariance is to learn a canonical global orientation of the point cloud in the first layers and then transform the input accordingly before using it in the main part of the network. This factors out the orientation of the input and the output will be invariant. Several methods have been developed to predict the global orientation, based on subnetworks (Qi et al., ), principal component analysis (Li et al., ) or as a combination of local orientations (Zhao et al., ; Zhang et al., ; Zhao et al., ).
However, most such approaches yield only approximate invariance, notably when the predicted orientations are not consistent for different orientations of the input. In contrast, our framework guarantees exact equivariance.

#### Invariance using scalar internal representations.

A simple way of achieving exact equivariance is to extract invariant features from the input geometry, such as distances or angles, and only include these in the message passing (Schütt et al., ; Li et al., ). While this allows using typical deep learning building blocks (linear layers, activations, norm layers, etc.), these approaches are not able to communicate non-scalar geometric information (such as directions) during message passing and produce only invariant but not equivariant outputs.

#### Equivariance using tensorial internal representations.

Going one step further, several works have included vectors into the message passing formalism (Deng et al., ; Satorras et al., ). This enables the prediction of vectorial and tensorial properties (Schütt et al., ). Furthermore, it has been shown that including non-scalar geometric features in the message passing can enhance the performance, even when the model outputs are invariant quantities (Fuchs et al., ; Brandstetter et al., ).
However, including geometric vectors in neural networks comes at a cost: one can no longer treat every internal activation like an individual number, but the coordinates of vectors have to be processed jointly to maintain equivariance. In contrast to our framework, these approaches typically rely on carefully designed non-linearities, norm layers and special operators to communicate between scalar and vectorial features.

Beyond vectorial representations, recently, many works have explored including higher-order tensorial representations (Frank et al., ; Batatia et al., ; Liao et al., ; Musaelian et al., ; Remme et al., ; Simeon and De Fabritiis, ).
Most existing approaches are based on the irreducible representations of $\mathrm{SO}(3)$ (the representations under which spherical harmonics transform). Message passing is then often implemented in the form of convolutions with kernels whose angular part is fixed to spherical harmonics. During the convolution, tensorial objects of different orders are mixed through the tensor product, see e.g. (Thomas et al., ) for details.
Regarding the irreducible representations, several works have reported that including higher-order tensor fields is particularly helpful in tasks where angular information matters, e.g. for predicting forces in molecules (Zitnick et al., ; Batzner et al., ).
In our framework, one is not restricted to using irreducible representations and the message passing kernels are not restricted to the spherical harmonics times a learned radial function.
Below, we define the tensor representation (Eq.(2)), which can be implemented efficiently and is used as feature representation in our framework.

#### Equivariance based on local frames.

Similar to equivariance by global orientation estimation, one can achieve equivariance using local reference frames (Wang and Zhang, ; Lou et al., ).
One advantage of using local reference frames over a single global one is that substructures in the data, which are geometrically akin, obtain similar local features. This is a desirable property regarding generalization.
Most closely related to our work, Luo et al. () equivariantly predict a local coordinate frame for each node, into which the geometric input features are transformed. Thereby, the local coordinates of the node features become independent of the global orientation of the input and are thus invariant.
However, during message passing, Luo et al. () no longer leverage the local coordinate frames to transform features between the frames, which substantially limits the communication between nodes that have different local frames (cf. Fig. 3). Our framework overcomes this limitation, by enabling the communication of tensorial objects in the messages, resulting in a strictly more general and expressive formalism for equivariant message passing.

#### Constructing local reference frames.

Geometrically informative local frames should be constructed robustly, i.e. small changes in the local geometry should not change them drastically. Secondly, they should be predicted equivariantly, i.e. if the input is transformed the local frames must transform accordingly.
How the local frames are constructed from the local geometry strongly depends on the input. For molecules, many works rely on handcrafted frames based on nearest neighbor directions (Wu et al., ; Kramer et al., ; Zhang et al., ). For point clouds describing 3D shapes, local frames are typically based on the surface normal and other characteristic local directions, inferred from handcrafted rules or in a learned fashion (Tombari et al., ; Petrelli and Di Stefano, ; Yang et al., ; Melzi et al., ; Zhu et al., ).
Our proposed framework works for any method for constructing the local frames so that it can be chosen problem-specific.

## 3 Preliminaries

The set of transformations that can be applied to a data sample typically forms a group in the mathematical sense. Therefore, we will formally define a group representation, which characterizes the well-defined transformation behavior of node features in our message passing framework.

#### Group representation.

Given a group $G$ and vector space $V$, a group representation $\rho$ is a mapping from $G$ to the invertible matrices $GL(V)$ that fulfills

$$ $\displaystyle\rho(g_{1}g_{2})=\rho(g_{1})\rho(g_{2}),\quad\forall g_{1},g_{2} \in G,$ (1) $$

where $g_{1}g_{2}$ is the group product and $\rho(g_{1})\rho(g_{2})$ is a matrix product. The representation specifies how elements of the group act on vectors $v\in V$, i.e. in components $(\rho(g)v)_{i}=\rho(g)_{ij}v_{j}$. The condition (1) implies that $\rho(g^{-1})=(\rho(g))^{-1}$.

#### Equivariance.

Let $G$ be a group and $V,W$ two vector spaces equipped with group representations $\rho_{i}$ and $\rho_{o}$ respectively. A function $f:V\to W$ is said to be equivariant under the group $G$ if the following holds:

$$ $\displaystyle\rho_{o}(g)f(x)=f(\rho_{i}(g)x),\quad\forall g\in G,$ $$

where the input to $f$ transforms under the representation $\rho_{i}:V\to GL(V)$ and its output under the representation $\rho_{o}:W\to GL(W)$. If $\rho_{o}(g)=id$ for all $g\in G$, the function $f$ is said to be invariant.

#### Tensor representation.

Let us consider the group of rotations and reflections $O(d)$. The most common representation of $O(d)$ is in terms of $d\times d$ orthogonal matrices. Given an abstract group element $g\in G$, let us denote the corresponding $d\times d$ matrix by $R(g)$. Since the abstract group element $g$ is often directly associated with its corresponding $d\times d$ matrix $R$, we sometimes write $\rho(R(g))$ or even only $\rho(R)$ instead of the formally correct $\rho(g)$.

While a $d$-dimensional vector transforms by contraction of its only index, i.e. $(R(g)\mathbf{v})_{i}=R(g)_{ij}v_{j}$, one defines higher-order tensor representations based on the following transformation rule:

$$ $\displaystyle T^{\prime}_{i_{1}...i_{n}}=R(g)_{i_{1}j_{1}}\ ...\ R(g)_{i_{n}j_ {n}}T_{j_{1}...j_{n}}.$ (2) $$

A tensor $T_{j_{1}...j_{n}}$ with $n$ indices is said to have order $n$. All indices run over $d$ dimensions.
One may easily check that the transformation behavior in Eq. 2 fulfills condition (1) and defines a representation (see App. A).
The vector space on which this representation acts is $d^{n}$-dimensional.
The fact that the orthogonal matrices also include reflections can be used to distinguish the transformation behavior of geometric objects with respect to orientation-reversing transformations (with determinant $-1$).
Since the determinant is multiplicative, the following transformation behavior defines another representation, namely the one for pseudotensors:

$$ $\displaystyle P^{\prime}_{i_{1}...i_{n}}=\det(R(g))R(g)_{i_{1},j_{1}}\ ...\ R( g)_{i_{n},j_{n}}P_{j_{1}...j_{n}}.$ (3) $$

For instance, a pseudovector does not change sign under reflections, e.g. $\rho(r)\mathbf{v}=\mathbf{v}$ for a three-dimensional pseudovector and a pure reflection $r\in\mathrm{O}(3)$ (Jeevanjee, ).

## 4 Methods

The central idea behind the Tensor Frames approach (Fig. 1) can be summarized as follows: for every node, one predicts an orthonormal local frame in an equivariant fashion, meaning that it transforms consistently as the input point cloud is flipped or rotated. Then, one expresses the geometric node features in the respective local frame by change of basis from the global reference frame to the local one. Crucially, since the local frames are chosen equivariantly, the node features expressed in the local frames are invariant under $\mathrm{O}(3)$-transformations of the input (see Fig. 2).
Below, we provide a proof that this indeed holds for all geometric objects, irrespective of their representation.
The invariant coordinates can then be processed by arbitrary functions without breaking the invariance.
The key ingredient in the Tensor Frames approach is the following: During message passing, node features are transformed from one local frame into the other, which enables direct communication of geometric features, like vectors and tensors, between nodes. At the final layer, one transforms the invariant numbers back to a geometric object in the global frame, using the desired output representation. One thereby obtains an equivariant prediction.

Figure: Figure 2: 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}$. By the change of basis, the vectorial information can be received in the other coordinate frame without loss of information. Through the equivariance of the local frames, the local frame coordinates of geometric objects are invariant under global transformation $\hat{R}$.
Refer to caption: x2.png

### 4.1 Equivariance by local frames

Let us first describe how to predict the local frames in an equivariant manner, and illustrate how they are used to construct an equivariant pipeline. For concreteness, we will describe the procedure in three-dimensional Euclidean space, though all statements generalize straightforwardly to higher dimensions.

We adapt the simple approach by (Wang and Zhang, ) to learn the prediction of equivariant local frames during training.
Given an input graph of nodes with coordinates $\mathbf{x}_{i}$ and input node features $F_{i}$, we equivariantly predict two vectors $\mathbf{v}_{i,1},\mathbf{v}_{i,2}$ for each node:

$$ $\displaystyle\mathbf{v}_{i,k}=\sum_{j\in\mathcal{N}(i)}\omega(\left\|\mathbf{x }_{i}-\mathbf{x}_{j}\right\|)\ \phi(F_{i}^{(s)},F_{j}^{(s)},e_{ij}^{(s)},\left \|\mathbf{x}_{i}-\mathbf{x}_{j}\right\|)_{k}\ \frac{\mathbf{x}_{i}-\mathbf{x}_ {j}}{\left\|\mathbf{x}_{i}-\mathbf{x}_{j}\right\|},\quad k\in\{1,2\},$ (4) $$

where $\left\|.\right\|$ denotes the Euclidean norm and $\mathcal{N}(i)$ the neighborhood of node $i$. For point cloud data without any edges, the neighborhood is obtained from a radius graph. $\phi$ is a standard MLP that receives the scalar (invariant) part of the node features $F_{i}^{(s)},F_{j}^{(s)}$ and edge attributes $e_{ij}^{(s)}$ (if available) as inputs. The vectors $\mathbf{v}_{i,1}$ and $\mathbf{v}_{i,2}$ are computed as weighted sums of the normalized edge vectors, with the two outputs of $\phi$ being the respective weights. $\omega$ is an envelope function adapted from (Gasteiger et al., ), which goes to zero smoothly at the cutoff radius $r_{c}$ (see App. B for details).
Using the Gram-Schmidt procedure, one equivariantly constructs two orthonormal vectors from $\mathbf{v}_{i,1},\mathbf{v}_{i,2}$ (see App. B). A third vector is obtained from the vector product between these two, yielding an orthonormal basis $\mathbf{n}_{i,k}\in\mathbb{R}^{3},\ k=1,2,3$. However, the vector product $\mathbf{n}_{i,3}=\mathbf{n}_{i,1}\times\mathbf{n}_{i,2}$ always results in a right-handed local frame so that the handedness would not change under reflection of the input. Hence, such local frames would be $\mathrm{SO}(3)$- but not $\mathrm{O}(3)$-equivariant. A simple solution is to flip the third vector based on the relative position of the local center of mass:

$$ $\displaystyle\mathbf{n}_{i,3}=\begin{cases}\mathbf{n}_{i,1}\times\mathbf{n}_{i ,2}\quad\text{ if }\quad(\mathbf{n}_{i,1}\times\mathbf{n}_{i,2})\cdot\mathbf{ \bar{r}}\geq 0\\ -\mathbf{n}_{i,1}\times\mathbf{n}_{i,2}\hskip 56.9055pt\text{ else }\end{cases },\text{ with }\mathbf{\bar{r}}:=\sum_{j\in\mathcal{N}(i)}\omega(\left\| \mathbf{x}_{j}-\mathbf{x}_{i}\right\|)(\mathbf{x}_{j}-\mathbf{x}_{i}),$ (5) $$

where $\cdot$ denotes the standard dot-product. The computation of the direction $\mathbf{\bar{r}}$ is smoothed using the same envelope $\omega$ as in Eq. (4).

#### Invariance of node features expressed in local frames.

The $3\times 3$ orthogonal matrix which transforms a vector from the global frame of reference into the local frame at node $i$ is given by

$$ $\displaystyle R_{i}=(\mathbf{n}_{i,1},\mathbf{n}_{i,2},\mathbf{n}_{i,3})^{ \mathrm{T}}=\begin{pmatrix}\rule[2.15277pt]{14.22636pt}{0.55pt}\ \mathbf{n}_{i ,1}\ \rule[2.15277pt]{14.22636pt}{0.55pt}\\ \rule[2.15277pt]{14.22636pt}{0.55pt}\ \mathbf{n}_{i,2}\ \rule[2.15277pt]{14.22 636pt}{0.55pt}\\ \rule[2.15277pt]{14.22636pt}{0.55pt}\ \mathbf{n}_{i,3}\ \rule[2.15277pt]{14.22 636pt}{0.55pt}\\ \end{pmatrix}$ (6) $$

This can be easily seen from $R_{i}\mathbf{n}_{i,1}=(1,0,0)^{\mathrm{T}}$ etc., due to the orthonormality.
Under any global transformation $R(\hat{g})=:\hat{R}\in\mathrm{O}(3)$ the node positions, local frames and input features transform according to

$$ $\displaystyle\mathbf{x}^{\prime}_{i}=\hat{R}\,\mathbf{x}_{i},\quad\ \mathbf{n} ^{\prime}_{i,k}=\hat{R}\,\mathbf{n}_{i,k}\ \forall i,\ k=1,2,3\quad\text{ and }\quad F^{\prime}_{i}=\rho(\hat{g})F_{i},$ (7) $$

which implies the following transformation rule for the matrix $R_{i}$:

$$ $\displaystyle R^{\prime}_{i}$ $\displaystyle=R_{i}\hat{R}^{\mathrm{T}}=R_{i}\hat{R}^{-1}\text{ , since}$ (8) $\displaystyle[R_{i}^{\prime}]_{mn}$ $\displaystyle=[(\mathbf{n}^{\prime}_{i,1},\mathbf{n}^{\prime}_{i,2},\mathbf{n} ^{\prime}_{i,3})^{\mathrm{T}}]_{mn}=(\mathbf{n}^{\prime}_{i,m})_{n}=[\hat{R}]_ {nl}(\mathbf{n}_{i,m})_{l}=(\mathbf{n}_{i,m})_{l}[\hat{R}^{\mathrm{T}}]_{ln}=[ R_{i}\hat{R}^{\mathrm{T}}]_{mn}.$ $$

In components $(\mathbf{n}_{i,m})_{l}$ denotes the $l$-th component of the $m$-th basis vector at node $i$.
Using
Eq. (7) and (8), we can now show that node features expressed in local frames are invariant w.r.t. transformations of the inputs. Indeed, let $g_{i}$ denote the abstract group element associated with the local frame transformation $R_{i}$, then the features transformed into the local frames are given by $\rho(g_{i})F_{i}=:f_{i}$ and the invariance follows as

$$ $\displaystyle f^{\prime}_{i}=\rho(g_{i}\hat{g}^{-1})F^{\prime}_{i}=\rho(g_{i}) \rho(\hat{g}^{-1})\rho(\hat{g})F_{i}=\rho(g_{i})\rho(\hat{g})^{-1}\rho(\hat{g} )F_{i}=\rho(g_{i})F_{i}=f_{i}.$ (9) $$

#### Equivariance of the output.

After transforming the node features into the local frames, the node features can be processed and combined using arbitrary functions for message passing, such as standard MLPs, non-linearities and norm layers. Afterwards, the invariant features are transformed back into the global frame to obtain an equivariant prediction $\rho_{Y}(g_{i}^{-1})f_{i}=:Y_{i}$. Indeed, the prediction transforms equivariantly under any global transformation $\hat{g}$:

$$ $\displaystyle Y^{\prime}_{i}=\rho_{Y}(\hat{g}g_{i}^{-1})f^{\prime}_{i}=\rho_{Y }(\hat{g})\rho_{Y}(g_{i}^{-1})f_{i}=\rho_{Y}(\hat{g})Y_{i}.$ (10) $$

Crucially, Eq. (10) holds for any representation of the output $\rho_{Y}$. This means that, after applying multiple message passing layers to the invariant node features, it is merely a matter of interpretation to decide which numbers shall be combined into a common geometric object and which object it should be. Thereby, our pipeline allows for an equivariant prediction of any geometric object as required by the given problem.

### 4.2 Tensor Frames message passing

So far, we have seen how to achieve equivariance by transforming into equivariant local frames and then performing message passing on the invariant node features $f_{i}$. In a general form, the invariant message passing without tensorial messages for node feature $f_{i}$ in the $k$-th layer can be written as

$$ $\displaystyle f_{i}^{(k)}=\psi^{(k)}\bigg{(}f_{i}^{(k-1)},\bigoplus_{j\in \mathcal{N}}\phi^{(k)}\left(f_{i}^{(k-1)},f_{j}^{(k-1)},\rho_{e}(g_{i})e_{ij}, R_{i}(\mathbf{x}_{i}-\mathbf{x}_{j})\right)\bigg{)},$ (11) $$

where $\psi^{(k)}$ and $\phi^{(k)}$ are arbitrary non-linear functions and $\bigoplus_{j\in\mathcal{N}}$ denotes the aggregation over neighbors.

Figure: Figure 3: Limitation of scalar message passing. The upper node is sending a non-transforming (scalar) message that encodes a characteristic direction in its neighborhood. Since the direction is encoded relative to the local frame of the sending node, both messages (left and right) are the same to the receiving node.
Refer to caption: x3.png

The message passing defined by Eq. (11) differs from vanilla message passing only by the transformation of the edge vectors $\mathbf{x}_{i}-\mathbf{x}_{j}$ and the input edge features $e_{ij}$, which transform by $\rho_{e}$. Both expressions $\rho_{e}(g_{i})e_{ji}$ and $R_{i}(\mathbf{x}_{i}-\mathbf{x}_{j})$ are invariant under global transformations as instances of Eq. (9). Together with the invariance of the node features, this guarantees the invariance of Eq. (11).

The message passing of previous works like (Luo et al., ) can be expressed in the form of Eq. (11).
However, it has important implications that the invariant node features are expressed in different local frames: message passing in the form of Eq. (11) does not allow for the direct communication of geometric information: For example, a node may send a message encoding a characteristic direction in its neighborhood only relative to its own local frame. If, however, the receiving node has no or incomplete knowledge about the local frame orientation of the sending node, it cannot process this directional information (cf. Fig. 3). Our framework remedies exactly this weakness by incorporating proper tensorial messages in the message passing between local frames. As in the final transformation from the local frames back to the global one, the invariant features $f_{i}$ can be interpreted as coordinates of a geometric object already during message passing.
As a hyperparameter of the model, one chooses the transformation behavior of $f_{i}$ as the direct sum of multiple tensor and pseudotensor representations. The representation defines how $f_{i}$ is transformed from one local frame to the other, based on Eq. (2) and (3).
During training, the network learns to communicate vectorial and tensorial information through the respective feature channels, simply because they transform accordingly. That is, if a node would like to send a direction in the form of a vector to its neighbor, it will store the three coordinates in three channels of $f_{i}$ which by design transform like a vector (as illustrated in Fig. 2).
As our main result, this yields the following strictly more general form of invariant message passing with tensorial messages between local frames:

$$ $\boxed{\begin{aligned} f_{i}^{(k)}=\psi^{(k)}\bigg{(}f_{i}^{(k-1)},\bigoplus_{ j\in\mathcal{N}}^{\vphantom{j\in\mathcal{N}}}\phi^{(k)}\left(f_{i}^{(k-1)}, \rho(g_{i}g_{j}^{-1})f_{j}^{(k-1)},\rho_{e}(g_{i})e_{ji},R_{i}(\mathbf{x}_{i}- \mathbf{x}_{j})\right)\bigg{)}.\end{aligned}}$ (12) $$

Indeed, the transformation of an invariant node feature $f_{j}$ from the local frame of node $j$ into the one of node $i$ is also invariant:

$$ $\displaystyle\rho(g_{i}g_{j}^{-1})f_{j}\ \xrightarrow{\ \ \text{global }\hat{g }\ \ }\ \rho((g_{i}\hat{g}^{-1})(g_{j}\hat{g}^{-1})^{-1})f^{\prime}_{j}=\rho(g _{i}\hat{g}^{-1}\hat{g}g_{j})f_{j}=\rho(g_{i}g_{j})f_{j}.$ (13) $$

The formalism in Eq. (12) allows modifying all existing message passing approaches of this form to be $\mathrm{O}(d)$-equivariant and communicate tensorial messages of arbitrary representations.
In practice, we opt for the tensor representation as feature representation in our networks since it can be implemented efficiently directly using the transformation matrices of the local frames, e.g. by utilizing highly optimized Einstein summation algorithms (cf. Eq. (2) and (3)).

### 4.3 Relation to data augmentation

As a direct consequence of condition (1), any group representation maps the identity element to the identity matrix. Thus, if all local frames were chosen to be the identity, Eq. (12) would simply turn into the usual non-invariant message passing. Similarly, choosing all local frames to be the same group element, i.e. $R_{i}=R(\tilde{g}),\ \tilde{g}\in O(d)\ \forall i$, is equivalent to a global transformation of the input data. In this case, tensorial messages do not change when sent from one local frame to the other since the change of basis $g_{i}g_{j}^{-1}$ is trivial. Therefore, choosing $\tilde{g}\in O(d)$ randomly for every training sample precisely amounts to data augmentation with random rotations and reflections.

One clear advantage of our framework for equivariance is that it allows for a direct comparison between equivariant message passing and data augmentation. Essentially all other works that present an equivariant pipeline do not compare against a non-equivariant baseline trained with data augmentation. Presumably, this is due to the fact that these approaches use specialized equivariant building blocks in their architecture, which do not have a straightforward non-equivariant equivalent. In our case, the architecture can be trained in both ways for a fair comparison using the same hyperparameters in the exact same architecture and optimizer.

## 5 Experiments

**Table 1: Normal vector regression on ModelNet40. We report cosine similarities (higher is better) between predicted and target normal vectors for three scenarios: 1. trained and evaluated with augmentations around the gravitational axis, 2. trained only with rotations around $z$ but evaluated using all transforms of $\mathrm{O}(3)$ or $\mathrm{SO}(3)$ and 3. trained and evaluated using all transforms. Our equivariant adaptation of PointNet++ (Qi et al., ) produces state-of-the-art results.**
| Method | $z/z$ | $z/\mathrm{SO}(3)$ | $\mathrm{SO}(3)/\mathrm{SO}(3)$ | equivariant |
| --- | --- | --- | --- | --- |
| Luo et al. (Luo et al., ) | 0.80 | 0.80 | 0.80 | ✓ |
|  | $z/z$ | $z/\mathrm{O}(3)$ | $\mathrm{O}(3)/\mathrm{O}(3)$ | equivariant |
| PointNet++ with data augmentation (ours) | 0.88 | 0.76 | 0.85 | ✗ |
| PointNet++ with scalar messages (ours) | 0.81 | 0.81 | 0.81 | ✓ |
| PointNet++ with tensor messages (ours) | 0.86 | 0.86 | 0.86 | ✓ |

**Table 2: Classification accuracies on ModelNet40. Our equivariant adaptation of PointNet++ (Qi et al., ) produces superior results over the vanilla PointNet++ (with training and evaluation setups as in Tab. 1).**
| Method | $z/z$ | $z/\mathrm{SO}(3)$ | $\mathrm{SO}(3)/\mathrm{SO}(3)$ | invariant |
| --- | --- | --- | --- | --- |
| CRIN (Lou et al., ) | 91.8 | 91.8 | 91.8 | ✓ |
|  | $z/z$ | $z/\mathrm{O}(3)$ | $\mathrm{O}(3)/\mathrm{O}(3)$ | invariant |
| PointNet++ with data augmentation (ours) | 89.4 | 17.1 | 86.6 | ✗ |
| PointNet++ with scalar messages (ours) | 87.6 | 87.6 | 87.6 | ✓ |
| PointNet++ with tensor messages (ours) | 88.3 | 88.3 | 88.3 | ✓ |

**Table 3: Learning informative local frames and tensorial messages are beneficial. Normal vector regression on ModelNet40 using adaptations of PointNet++ with learned vs. random local frames and tensorial vs. scalar messages. The model with random frames and tensor messages still significantly outperforms the model with learned frames but scalar messages, highlighting the benefit of the Tensor Frames approach.**
| Cosine similarity $\uparrow$ | tensorial messages | scalar messages |
| --- | --- | --- |
| learned local frames | 0.86 | 0.81 |
| random local frames | 0.84 | 0.79 |

Below, we present results for two point cloud experiments using the popular PointNet++ architecture (Qi et al., ), adapted to our equivariant framework. Similar to a U-Net for images (Ronneberger et al., ), the PointNet++ architecture combines an encoder that iteratively subsamples the point cloud with a decoder that vice versa upsamples the point cloud. The message passing in the encoder and decoder both follow the form of Eq. (12).
Compared to the original PointNet++, we make minor architectural changes, e.g. by introducing radial and angular embeddings (see App. C and D for details of the architecture and training setup).
The decoder part is needed for tasks that require per point predictions.
For tasks that only require a single global output, we prune the architecture by utilizing only the encoder.

We have trained different variants of our PointNet++ adaptation on normal vector regression and classification on the ModelNet40 dataset (Wu et al., ). It consists of 12,311 3D shapes of 40 different categories. We use the resampled version of the dataset for which normal vectors at all points are available and use the default train/test split.
A model, using the Tensor Frames approach (Eq.(12)), is compared against an equivariant model which uses the less general scalar message passing of Eq.(11). In both models, the local frames are learned via Eq.(4). Further, we compare against a PointNet++ variant trained with data augmentation. As described in Sec. 4.3, our framework allows for a direct comparison between equivariant message passing, based on the Tensor Frames formalism, and data augmentation. For all three models, we use the same hyperparameters for the architecture and optimizer. Results compared to the current state-of-the-art equivariant architecture can be found in Tab. 1 and 2. The model trained with data augmentation has slightly fewer learnable parameters since the local frames are not learned but chosen randomly for data augmentation (cf. Sec. 4.3). For the normal regression model, this difference amounts to $0.4\%$ of the total parameter count and for the classification model to $1.1\%$.

## 6 Discussion

Figure: Figure 4: Data efficiency of built-in equivariance vs. data augmentation.
Refer to caption: x4.png

On normal vector regression, our Tensor Frames adaptation of the PointNet++ architecture achieves state-of-the-art results. Normal regression is a task in which equivariance is certainly desirable and in which geometric information must be propagated precisely. Thus, tensorial message passing proves to be superior over scalar messages due to the limitations mentioned in Sec. 4.2 and Fig. 3. For shape classification the gain is less significant, indicating that geometric information, e.g. in the form of characteristic directions, may be less important. In both experiments, the networks that model exact equivariance yield slight improvements over data augmentation. Moreover, we demonstrate that informative local frames are indeed beneficial through an ablation study with randomly chosen local frames (see Tab. 3). We find that the model with random frames and tensor messages still significantly outperforms the model with learned frames but scalar messages. This highlights once more that the Tensor Frames approach enables to communicate geometric information more reliably.

We would like to stress that our PointNet++ trained with data augmentation already outperforms the published state-of-the-art of equivariant methods on normal regression.
In general, we advocate that equivariant architectures are also compared against a strong baseline trained with data augmentation.
Since equivariant methods do neither “waste” data nor network capacity to perform well on different input orientations, they are often said to be more data efficient (Batzner et al., ), meaning that they improve faster as more data becomes available (Hestness et al., ). We have trained a series of classification networks (as in Sec. 5) with the same architecture but different fractions of the training data. The misclassification rate is expected to fall exponentially as the training set increases. Indeed, in the log-log plot (test error versus fraction of training data, Fig. 4) the equivariant Tensor Frames approach shows a steeper slope, indicating better data efficiency than the same model trained with data augmentation. However, perhaps surprisingly, the error rate is not necessarily smaller for all dataset sizes, meaning that in some cases, purely in terms of accuracy, data augmentation may be favorable over built-in equivariance.
We see this as a big advantage of our framework, which allows for the parallel development of an exact equivariant model and an equally well-engineered non-equivariant baseline, that can be trained via data augmentation.

## 7 Conclusion

This work introduces Tensor Frames, a novel framework for building $\mathrm{O}(d)$-equivariant message passing architectures. We provide a well-motivated formalism through which existing non-equivariant networks can be adapted to be equivariant.
The presented approach provides a new perspective compared to existing approaches for exact equivariance, which do not use local frames but specialized tensorial operations.
Our method does not require the implementation of tensor products or specialized linear layers, norms or non-linearities but can be integrated straightforwardly with existing architectures.
Our method offers a strict generalization of existing approaches that achieve equivariance based on local frames. We demonstrate that equivariance via Tensor Frames can significantly improve the performance compared to methods that use message passing without tensorial messages.
Our framework can be used as a drop-in replacement for data augmentation to achieve exact, built-in equivariance and allows for a direct and fair comparison between the two approaches.
Tensor Frames open up a new possibility for evaluating the efficacy of equivariance as a model prior on numerous geometric machine learning tasks. Through this work, we hope to inspire researchers and practitioners alike to apply the Tensor Frames approach to different architectures in various domains.

## Acknowledgments and Disclosure of Funding

This work is supported by Deutsche Forschungsgemeinschaft (DFG) under Germany’s Excellence Strategy EXC-2181/1 - 390900948 (the Heidelberg STRUCTURES Excellence Cluster) and under project number 240245660 - SFB 1129, as well as by the Klaus Tschira Stiftung gGmbH (SIMPLAIX project P4).

## Appendix

## Appendix A Tensor representation

The tensor representation is introduced in (2). The tensor representation is a group representation of $\mathrm{O}(3)$ and is used in this paper to define the transformation behavior of the tensorial objects in the Tensor Frames message passing.

Let $g_{1},g_{2}\in\mathrm{O}(3)$ and $R(g_{1}),R(g_{2})$ the corresponding orthogonal $3\times 3$ matrices. Essentially, the representation property follows from the fact that the 3x3 matrices are a representation. With this fact, the proof that the tensor representation is a representation is given as follows:

$$ $\displaystyle R(g_{1}g_{2})_{i_{1}j_{1}}\ ...\ R(g_{1}g_{2})_{i_{n}j_{n}}T_{j_ {1}...j_{n}}$ $\displaystyle=R(g_{1})_{i_{1}k_{1}}\ ...\ R(g_{1})_{i_{n}k_{n}}R(g_{2})_{k_{1} ,j_{1}}\ ...\ R(g_{2})_{k_{n}j_{n}}T_{j_{1}...j_{n}}$ $\displaystyle=(R(g_{1})R(g_{2}))_{i_{1}j_{1}}\ ...\ (R(g_{1})R(g_{2}))_{i_{n}j _{n}}T_{j_{1}...j_{n}}.$ (14) $$

The pseudo-vector representation is also a representation, which can be shown similarly and by using that the determinant is multiplicative: $\det(AB)=\det A\det B$.

## Appendix B Learning local frames

Learning the local frames is an essential part of the proposed architecture, therefore we present some further considerations for predicting local frames in an $\mathrm{O}(3)$ equivariant way.

As described in Sec. 4.1, for each node $i$ one predicts two vectors $\mathbf{v}_{i,1}$ and $\mathbf{v}_{i,2}$ by summing over the relative positions of adjacent nodes, weighted by the output of an MLP and an envelope function $\omega$. The envelope function, adapted from [Gasteiger et al., ], is given by

$$ $\displaystyle w(r_{ij})=\left\{\begin{matrix}1-\frac{(p+1)(p+2)}{2}\left(\frac {r_{ij}}{r_{c}}\right)^{p}+p(p+2)\left(\frac{r_{ij}}{r_{c}}\right)^{p+1}-\frac {p(p+1)}{2}\left(\frac{r_{ij}}{r_{c}}\right)^{p+2}&r_{ij}<r_{c}\\ 0&r_{ij}\geq r_{c}\end{matrix}\right.$ (15) $$

and ensures a smooth transition at the cutoff radius $r_{c}$. Here, $r_{ij}=\left\|\mathbf{x}_{j}-\mathbf{x}_{i}\right\|$ is the relative distance between the nodes. The parameter $p$ is chosen to be $5$ for all experiments presented in this paper.
Afterwards, the two vectors $\mathbf{v}_{i,1},\mathbf{v}_{i,2}$ are use to construct two orthonormal vectors $\mathbf{n}_{i,1},\mathbf{n}_{i,2}$ using the Gram-Schmidt procedure:

$$ $\displaystyle\mathbf{n}_{i,1}$ $\displaystyle=\frac{\mathbf{v}_{i,1}}{\left\|\mathbf{v}_{i,1}\right\|}$ (16) $\displaystyle\mathbf{n}_{i,2}^{\prime}$ $\displaystyle=\mathbf{v}_{i,2}-(\mathbf{n}_{i,1}\cdot\mathbf{v}_{i,2})\ \mathbf{n}_{i,1}$ (17) $\displaystyle\mathbf{n}_{i,2}$ $\displaystyle=\frac{\mathbf{n}_{i,2}^{\prime}}{\left\|\mathbf{n}_{i,2}^{\prime }\right\|}$ (18) $$

The third vector is chosen to point in the same half-space as the local center of mass to ensure an $\mathrm{O}(3)$-equivariant construction. The estimate of the direction to the local center of mass is smoothed using the same envelope function $\omega$.

$$ $\displaystyle\mathbf{n}_{i,3}=\begin{cases}\mathbf{n}_{i,1}\times\mathbf{n}_{i ,2}\quad\text{ if }\quad(\mathbf{n}_{i,1}\times\mathbf{n}_{i,2})\cdot\mathbf{ \bar{r}}\geq 0\\ -\mathbf{n}_{i,1}\times\mathbf{n}_{i,2}\hskip 56.9055pt\text{ else }\end{cases },\text{ with }\mathbf{\bar{r}}:=\sum_{j\in\mathcal{N}(i)}\omega(r_{ij})( \mathbf{x}_{j}-\mathbf{x}_{i}),$ (19) $$

The third vector can not be learned. If one constrains the local coordinate frames to be orthonormal, the third vector is defined up to its sign; and predicting this sign is a non-differentiable operation. Our experiments have shown that orthonormal frames, which are associated with $\mathrm{O}(3)$ transformations are favorable and that a relaxation of the normalization or orthogonality of the basis vectors decreases the performance of the models.

## Appendix C Equivariant PointNet++ using our framework

PointNet++ is a widely-used architecture for point cloud tasks. It combines an encoder that iteratively down samples the point cloud with a decoder with upsampling [Qi et al., ]. Each layer in the encoder consists of the following steps:

- 1.
Use furthest point sampling to sample a subset of equally spaced nodes $N^{(k)}$.
- 2.
For each node $i$ in $N^{(k)}$ generate its neighborhood $\mathcal{N}(i)$ by finding all nodes within a specified radius.
- 3.
Send and aggregate messages from all neighbors, according to:
$\displaystyle f_{i}^{(k,enc)}=\max_{j\in\mathcal{N}(i)}\ \phi\left(f_{j}^{(k-1
,enc)},\mathbf{x}_{j}-\mathbf{x}_{i}\right),$
(20)
with the channel-wise maximum as an aggregation function.
- 4.
Continue with the next layer, but keep only the nodes $N^{(k)}$.

Since Eq. (20) follows precisely the form of Eq. (12), the message passing formula can easily be adapted to the Tensor Frames formalism by

$$ $\displaystyle f_{i}^{(k,enc)}=\max_{j\in\mathcal{N}(i)}\ \phi\left(\rho(g_{i}g _{j}^{-1})f_{j}^{(k-1,enc)},R_{i}(\mathbf{x}_{j}-\mathbf{x}_{i})\right),$ (21) $$

where $R_{i}$ is the local frame of node $i$ and $\rho(g_{i}g_{j}^{-1})$ the representation under which the node features are transformed from the local frames of node $j$ into the one at node $i$. We have further refined the messages by splitting the edge vectors $\mathbf{x}_{j}-\mathbf{x}_{i}$ into a radial and an angular embedding. For the angular embedding, we simply use the normalized direction. For the norm of the edge vector, we use a Gaussian embedding, with $k$ Gaussians-like functions, which have learnable mean and standard deviation $\mu_{k}$ and $\sigma_{k}$:

$$ $\displaystyle(\tilde{r}_{ij})_{k}=\exp\left(-\frac{(\left\|\mathbf{x}_{j}- \mathbf{x}_{i}\right\|-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)$ (22) $$

The complete, invariant message passing step reads

$$ $\displaystyle f_{i}^{(k,enc)}=\max_{j\in\mathcal{N}(i)}\ \phi\left(\rho(g_{i}g _{j}^{-1})f_{j}^{(k-1,enc)},\tilde{r}_{ij},R_{i}\frac{\mathbf{x}_{j}-\mathbf{x }_{i}}{\left\|\mathbf{x}_{j}-\mathbf{x}_{i}\right\|}\right).$ (23) $$

If the task requires one output for the entire point cloud, the node features at the nodes, remaining after the encoder $N^{(k^{*})}$, are pooled into one global feature. The node and the local frame to which one sends all these messages is the one that is closest to the center of mass of the point cloud, i.e.:

$$ $\displaystyle\mathbf{\hat{\imath}}=\operatorname*{arg\,max}_{i\in N^{(k^{*})}} \left\|\mathbf{x}_{i}-\mathbf{\bar{x}}\right\|\text{ with }\mathbf{\bar{x}}:= \sum_{i\in N^{(k^{*})}}\mathbf{x}_{i}\ \Big{/}\!\sum_{i\in N^{(k^{*})}}1$ (24) $\displaystyle f_{global}=\max_{j\in\mathcal{N}_{\mathbf{\hat{\imath}}}}\ \phi \left(\rho(g_{\mathbf{\hat{\imath}}}g_{j}^{-1})f_{j}^{(k_{max},enc)},\tilde{r} _{\mathbf{\hat{\imath}}j},R_{\mathbf{\hat{\imath}}}\frac{\mathbf{x}_{j}- \mathbf{x}_{\mathbf{\hat{\imath}}}}{\left\|\mathbf{x}_{j}-\mathbf{x}_{\mathbf{ \hat{\imath}}}\right\|}\right)$ (25) $$

Finally, the global features may be passed through an MLP to generate the output of the invariant message passing part of our architecture.

If the task requires one output per node in the input point cloud, one must upsample the nodes again after the encoder. To do this, one caches the positions and features of the nodes of the encoder layers and iteratively applies the following steps in the decoder:

- 1.
Let the input nodes to that layer be $N^{(k)}$. The features at these nodes are interpolated to the node features of the larger subset $N^{(k-1)}$ (reversing the subsampling of the encoder layers): For that, one finds for each node $i\in N^{(k-1)}$ its three closest neighbors in $N^{(k)}$ (let us denote their set by $\mathrm{NN}_{3}(i)$) and interpolates their features by inverse distance weighting:
$\displaystyle h_{i}=\frac{\sum_{j\in\mathrm{NN}_{3}(i)}\left\|\mathbf{x}_{j}-
\mathbf{x}_{i}\right\|^{-1}\ \rho(g_{i}g_{j}^{-1})\ f_{j}^{(k,dec)}}{\sum_{j
\in\mathrm{NN}_{3}(i)}\left\|\mathbf{x}_{j}-\mathbf{x}_{i}\right\|^{-1}}.$
(26)
- 2.
The interpolated features $h_{i}$ are then concatenated with the node features at the $k-1$-th layer of the encoder and embedded in an MLP to obtain the updated and upsampled node features:
$\displaystyle f_{i}^{(k-1,dec)}=\mathrm{MLP}\left(h_{i},f_{i}^{(k-1,enc)}
\right),$
(27)
which practically implements a skip connection between the activations in the encoder and the decoder.
Note that the decoder features are counted backward in $k$.
- 3.
Continue with the next layer with the set of nodes given by $N^{(k-1)}$.

Finally, the node features, back at the level of the input nodes, are brought into the desired output dimension by one last MLP layer.

## Appendix D Additional experiments and experimental details

All experiments are conducted on the ModelNet40 dataset [Wu et al., ]. In particular, we use the resampled version available at [https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip](https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip), which includes normal vectors for each point in the point cloud. We use the first 1024 points based on the ordering provided in this version of the dataset and normalize the point clouds to fit in the unit sphere. The ordering is based on furthest point sampling to evenly cover the surface of the 3D shapes.

#### Hyperparameter choices.

The hyperparameters chosen for our two main experiments are listed in Tab. 4.

**Table 4: Hyperparameter choices. The main hyperparameter choices for our models in the classification and normal vector regression task. Label smoothing only applies to the classification model.**
|  | normal vector regression | classification |
| --- | --- | --- |
| optimizer | AdamW | AdamW |
| weight decay | 5e-4 | 0.05 |
| learning rate | 5e-3 | 1e-3 |
| scheduler | Cosine-LR | Cosine-LR |
| epochs | 800 | 800 |
| warm up epochs | 10 | 10 |
| gradient clip | 0.5 | 0.5 |
| label smoothing | N.A. | 0.3 |
| loss | L1-loss | Cross-Entropy |

#### Architectural design.

The architectures used in our experiment can be summarized using the following short-hand:

- •
Encoding layer: $E(\texttt{in rep.},[\text{hidden layers}],\text{neighborhood radius},\text{
subsampling fraction})$ see Eq. (23)
- •
Decoding layer: $D(\texttt{in rep.},[\text{hidden layers}])$ see Eq. (27)
- •
MLP: $\texttt{MLP}(\texttt{in rep.},[\text{hidden layers}],\texttt{out rep.})$
- •
Output layer: $O(\texttt{in rep.},[\text{hidden layers}],\texttt{out rep.},\text{dropout})$ see Eq. (24)

Furthermore, we define the following notation to specify the feature representation used during message passing: The feature representations are a direct sum of tensor and pseudotensor representations. A representation is characterized by its order (i.e. the number of indices, cf. 3) and its behavior under parity (n for tensors and p for pseudotensors). Furthermore, we specify the multiplicities, that is, how often each representation appears in a direct sum representation. To give an example, the representation denoted as 8x0p+4x1n is the direct sum of 8 pseudoscalars and 4 vectors.

The architecture used for normal vector regression is described in Tab. 5. The number of Gaussian-like functions in the radial embedding is set to 64. The architecture used for classification is described in Tab. 6. Here, the number of Gaussian-like functions in the radial embedding is set to 16. For both experiments, the MLP used in the prediction of the local frames (according to Eq. (4)) has two hidden layers of dimension 128 each.

All fully connected linear layers are followed by batch normalization except the MLP in the output layer, where we do not use any normalization. As activation function, we use the SiLU function.

**Table 5: Architecture of the normal vector regression model.**
| Layer number | Layer |
| --- | --- |
| 1 | $E(\texttt{0x0n},[64],0.2,1.0)$ |
| 2 | $E(\texttt{64x0n+16x0p+16x1n+4x1p+4x2n+1x2p},[64],0.2,1.0)$ |
| 3 | $E(\texttt{64x0n+16x0p+16x1n+4x1p+4x2n+1x2p},[128],0.2,0.2)$ |
| 4 | $E(\texttt{128x0n+32x0p+32x1n+8x1p+8x2n+2x2p},[256],0.5,0.2)$ |
| 5 | $E(\texttt{256x0n+64x0p+64x1n+16x1p+16x2n+4x2p},[512],0.8,0.2)$ |
| 6 | $E(\texttt{512x0n+128x0p+128x1n+32x1p+32x2n+8x2p},[512],1.4,0.2)$ |
| 7 | $D(\texttt{512x0n+128x0p+128x1n+32x1p+32x2n+8x2p},[512])$ |
| 8 | $D(\texttt{512x0n+128x0p+128x1n+32x1p+32x2n+8x2p},[256])$ |
| 9 | $D(\texttt{256x0n+64x0p+64x1n+16x1p+16x2n+4x2p},[128])$ |
| 10 | $D(\texttt{128x0n+32x0p+32x1n+8x1p+8x2n+2x2p},[128])$ |
| 11 | $D(\texttt{64x0n+16x0p+16x1n+4x1p+4x2n+1x2p},[64])$ |
| 12 | $D(\texttt{64x0n+16x0p+16x1n+4x1p+4x2n+1x2p},[64])$ |
| 13 | $\texttt{MLP}(\texttt{64x0n+16x0p+16x1n+4x1p+4x2n+1x2p},[128,64,32],\texttt{1x1<br>n})$ |

**Table 6: Architecture of the classification model.**
| Layer number | Layer |
| --- | --- |
| 1 | $E(\texttt{0x0},[64,128],0.2,0.3)$ |
| 2 | $E(\texttt{80x0n+20x1n+6x2n},[128,256],0.8,0.4)$ |
| 3 | $E(\texttt{160x0n+40x1n+12x2n},[256,512],1.4,0.5)$ |
| 4 | $E(\texttt{0x0},[64],0.2,1.0)$ |
| 5 | $O(\texttt{320x0n+80x1n+24x2n},[512,256,128],\texttt{40x0n},0.5)$ |

#### Hardware and runtimes.

The training of the equivariant learned frames + tensor messages model for the normal vector regression task took 28h on a single NVIDIA H100 GPU (CPU: 2 x 32-Core Epyc 7452 + 8 x A100, RAM: 1024 GB). The training of the data augmented model took 19h on a single NVIDIA A100 GPU (CPU: 2 x 48-Core Xeon 8468V + 4 x H100, RAM: 1024 GB). Training the equivariant model on the same hardware as the model trained with data augmentation would take approximately 94h. Accordingly, the training of the equivariant model is about five times slower than the data-augmented version.
The equivariant classification model (learned frames + tensor messages) was trained for 19h on a single NVIDIA A100 (as above) and the data augmented version for 10h.

We also compare the interference times of the trained models on the normal vector regression tasks. The results of this ablation can be found in Tab. 7. While the train time is substantially more expensive if tensor messages are included in the model, the difference in the inference time is far less severe.

**Table 7: Evaluation runtimes. Average runtime for a single sample on normal vector regression (executed on an NVIDIA A100). Standard deviations are based on 10 loops over the test set. Tensor frames achieve exact equivariance, but the runtime of data augmentation is 35% faster.**
| Method | evaluation runtime |
| --- | --- |
| Learned frames + tensor messages | $(0.23\pm 0.06)s$ |
| Random frames + tensor messages | $(0.21\pm 0.06)s$ |
| Learned frames + scalar messages | $(0.18\pm 0.05)s$ |
| Random frames + scalar messages | $(0.16\pm 0.07)s$ |
| Data augmentation | $(0.15\pm 0.05)s$ |

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: Please address the problems in the weaknesses.
soundness: 3
strengths: 1. The paper is well-written and technically sound.
2. The paper provides comprehensive theoretical analysis.
3. The experiment results are promising.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 4
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).
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)