VN-Transformer: Rotation-Equivariant Attention for Vector Neurons
Authors: Serge Assaad
Carlton Downey
Rami Al-Rfou'
Nigamaa Nayakanti
Benjamin Sapp
Venue: Accepted by TMLR ()

Abstract: Rotation equivariance is a desirable property in many practical applications such as motion forecasting and 3D perception, where it can offer benefits like sample efficiency, better generalization, and robustness to input perturbations.
Vector Neurons (VN) is a recently developed framework offering a simple yet effective approach for deriving rotation-equivariant analogs of standard machine learning operations by extending one-dimensional scalar neurons to three-dimensional "vector neurons."
We introduce a novel "VN-Transformer" architecture to address several shortcomings of the current VN models. Our contributions are:
(i) we derive a rotation-equivariant attention mechanism which eliminates the need for the heavy feature preprocessing required by the original Vector Neurons models; (ii) we extend the VN framework to support non-spatial attributes, expanding the applicability of these models to real-world datasets; (iii) we derive a rotation-equivariant mechanism for multi-scale reduction of point-cloud resolution, greatly speeding up inference and training; (iv) we show that small tradeoffs in equivariance ($\epsilon$-approximate equivariance) can be used to obtain large improvements in numerical stability and training robustness on accelerated hardware, and we bound the propagation of equivariance violations in our models.
Finally, we apply our VN-Transformer to 3D shape classification and motion forecasting with compelling results.

Published in Transactions on Machine Learning Research (01/2023)
VN-Transformer Rotation-Equiivriant Attention for Vector Neurons
Serge Assaad Duke University
serge.ossondldubeeeel
Carlton Downey
cmdoumecykuurymo.ccmm
Waymo LLC
Rami Al-Rfou
rmyeid@uuaymo.com
Waymo LLC
Nigamaa Nayakanti
nigamaa@wwayo.com
Waymo LLC
Ben Sapp Waymo LLC
bensapp@raaymo.com com
Reviewed on OpenReview: https: // openreview. net/ forum? id= BXX2L45DPG
//
Abstract
Rotation equivariance is a desirable properly in many practical applications such as motion forecasting and 3D perception, where it can offer benefits like sample efficiency, better generalization, and robustness to input perturbations. Vector Neurons (VN) is a recently developed framework offering a simple yet effective approach for deriving rotation-equivariant analogs of standard machine learning operations by extending one-dimensional scalar neurons to three-dimensional "vector neurons." We introduce novel "VN-Transformer" architecture lo address several shortcomings of the current VN models. Our contributions are: (i) we derive ootation-equivatie attention mechanism which eliminates the need for the heavy feature preprocessing required by the original Vector Neurons models; (ii) we extend the VN framework to support non-spatial attributes, expanding the applicability of these models to real-world datasets; (iti) we derive a equuaiinn mechanism for multi-scale reduction of point-cloud resolution, greatly speeding up inference and training; (iv) we show that small tradcoffs in equivariance (c-apprxxmmaee oquivariance) can used obtain large improvements in numerical stability and training robustness accelerated hardware, and we bound the propagation of equivariance violations in models. Finally, apply our VN- Transformer to 3D shape classification and motion forccasting with compelling results.
Work done during internship Waymo I.C.C

---

Published in Transactions on Machine Learning Research (01/2023)
1 Introduction
A chair scen from the front, the back, the top, the side is is still chair. When driving car, our driving behavior is independent of our direction of travel. These simple examples demonstrate how humans excel at using rotation invariance and equivariance to understand the world in context (see figure on the right). Unfortunately, typical machine learning models struggle to preserve equivariance/ivvarian appropriate it is indeed challenging to equip neural networks with the right inductive biases to represent 3D objects in equivariant manner.
when
SO(3)
SO(3)
Modeling spatial data is core component in many domains such CAD, AR/VR, and medical imaging applications. In assistive robotics and autonomous vehicle applications, 3D object detection, tracking, and motion forecasting form the basis for how a robot interacts with humans the real world. Preserving rotation invariance/ invariacce/quuivariance /qquvvariance can improve training time, reduce model size, and provide crucial guarantees about model performance in the presence of noise. Spatial data is often represented as point-cloud data structure. These point-clouds require both permutation invariance and rotation equivariance to be modeled sufficiently well. Approaches addressing permutation invariance
include Zaheer et al. (2018): Lee et al. 2019 ct al. (0177)
Recently, approaches jointly addressing rolation invariance and equivariance are gaining momentum. These can roughly be categorized into modeling invariance or equivariance by: (i) data augmentation, (ii) canonical pose ostimation, and (tii) model construction. Approaches (i) and (ii) do not guarantee exact cquivariance since they rely on model parameters to learn the right. inductive biases (Qi et al. 2017: Esteves el al. 2018
Jaderberg al., 2015 Chang et al. 2015). Further, data augmentation makes training more costly and errors in pose estimation propagate to downstream tasks, degrading model performance. Moreover, pose estimation requires labeling objects with their observed poses. In contrast, proposals in (dit) do provide equivariance guarantees, including the Tensor Field Networks of Thomas et al. (2018) and the (Ba3)-Transformer of Fuchs et al. 2020) (see Section 2.. However, their formulations require complex mathematical machinery or are limited to specific network architectures.
Most recently, Deng et al. (021] proposed a simple and generalizable framework, dubbed Vector Neurons (VNs), that can be used replace traditional building blocks of neural networks with rotation-equivariant analogs. The basis of the framework is lifting scalar neurons to 3-dimensional vectors, which admit simple mappings of SO(3) actions to latent spaces.
While et al. (2021) developed a framework and basic layers, many issues required for practical deployment real-world applications remain unaddressed. A summary of our contributions and their motivations are follows:
VN-Transformer. The use of Transformers in deep learning has exploded in popularity in recent years as the de facto standard mechanism for learned soft attention over input and latent representations. They have enjoyed many successes in image and natural language understanding (Khan et al. 2022: Vaswani al., 2017), and they have become an essential modeling component most domains. Our primary contribution of this paper to develop VN formulation of soft attention by generalizing scalar inner-product based attention to matrix model functions on sets (since they are permutation-qquivaiaanl they are natural fit to model functions on point-clouds. Our VN-Transformer possesses all the appealing properties that have made the original Transformer so successful, as well as rotation equivariance as an added benefit.
(Ée., the Frobenius inner product). Thanks to Transformers' ability to
Direct point-sct input. The original VN paper relied on odge convolution as a pre-processing step capture local point-cloud structure. Such feature engineering is not data-driven and requires human involvement in designing and tuning. Moreover, the sparsity of these computations makes them slow to
2

---

Published in Transactions on Machine Learning Research (01/2023)
Invariant class pred.
Equivariant trajectory
accelerated hardware. Our proposed rotation- equivariant attention mechanism learns higher-level fea- tures directly from single points for arbitrary point-clouds (see Section
Pool MLP
VN-MLP
Handling points augmented with non-spatial at- tributes. Real-world point-cloud datasets have compli- cated features sets the [x, y, spatial dimensions are typically augmented with crucial non-spatial attributes [| y, 2]; [a]] where a can be high-dimensional. For exam- ple, Lidar point-clouds have intensity & clongation values associated with each point, multi-sensor point-clouds have modality types, and point-clouds with semantic type have semantic attributes. The VN framework restricted the scope of their work to spatial point-cloud data, limiting the applicability of their models for real-world point-clouds with attributes. We investigate two mechanisms inte- grate attributes into equivariant models while preserving rotation equivariance (see Section 5
VN-Invariant
Weighted Pool
VN-Transformer Encoder
VN-Transformer Encoder
VN-MLP
VN-MLP
Input point- clont
Input trajectory
(a) sification model.
clas-
(b) Rotation-equivariant forecasting
trajectory model.
Equivariant multi-scale feature reduction. Practical data structures such as Lidar point-clouds are extremely large, consisting of hundreds of objects each with poten- tially millions of points. To handle such computationally challenging situations, we design a rolation-equivariant mechanism for multi-scale reduction of point-cloud resolution. This mechanism learns how to pool the point set in context-sensitive manner leading to a significant reduction in training and inference latency (sce Scction 6
Figure 1: VN-Transformer ("early fusion") mod- els. Legend: S0(3)-equivariant features; (0(3)-invariant features; Non-spatial feat ures.
capproximate equivariance. When attempting to scale up VN models for distributed accelerated hardware we observed significant numerical stability issues. We determined that these stemmed from a fundamental limitation of the original VN framework, where bias values could not be included in linear layers while preserving equivariance. We introduce the notion of e-approximate equivariance, and use to show that small tradcoffss equivariance can be controlled obtain large improvements in numorical stability via the addition of small biases, improving robustness of training on accelerated hardware. Additionally, we theoretically bound the propagation of rotation equivariance violations in VN networks (see Section
Empirical analysis. Finally, we evaluate our VN-Transformer on (i) the ModelNet40 shape classification task, (id) & modified No which includes per-point non-spatial attribules, and (iii) a modified version of the Waymo Open Motion Dataset trajectory forecasting task (see Section 8).
2 Related work
The machine learning community has long been interested in building models that achieve equivariance to certain transformations, e.g., permutations, translations, and rotations. For a thorough review, see Bronstein p.21p.
et al.
Learned approximate transformation invariance. A very common approach is to learn robustness to input transforms via data augmentation (Zhou & Tuzel 2018 Qi et al. 2018 Krizhevsky et al. 2012 Lang
ct al.
2019 Yang et al. 2018) or by explicitly predicting transforms to canonicalize pose (Jaderberg et al.
2015
Hinton et al. 2018 Esteves et al. 2018 Chang et al. 2015
Rotaaioon-quivariant CNNs. Recently, there has been specific interest in designing rotation-equivariant image models for 2D perception tasks (Cohen & Welling 2016 Worrall et al. 2017 Marcos et al. 20177 Chidester et al. 20118). Worrall & Brostow (2018) extended this work to 3D perception, and Veeling et al. (2018) demonstrated the promise of rotation-equivaaiinn. models for medical images.
3

---

Published in Transactions on Machine Learning Research (01/2023)
(018 proposed Tensor Field Networks (TFNs), which use tensor representations of point-cloud data, Clebsch-Gordan coefficients, and spherical harmonic filters to build rotation-equivariant models. Fuchs et al. (2020) propose an "SE(3)-Transforme"" by adding an attention mechanism for TFNs. One of the key ideas behind this body of work to create highly restricted weight matrices that commute with rotation operations by construction (f.e., WR = RW). In contrast, we propose simpler alternative: a "VN-Transformer" which guarantees equivariance for arbitrary weight matrices, removing the need for the complex mathematical machinery of the Trans-Tra,,frrmer. For a detailed comparison with Fuchs et al. 2020 Appendix A
Equivariant point-cloud models. Thomas et al.
Controllable approximate equivariance. Finzi et al. (2021) proposed equivariant priors on model weight matrices to achieve approximate equivariance, and Wang et al. (2022) proposed a relaxed steerable 2D convolution along with soft equivariance regularization. In this work, we introduce the related notion of "c-approximate equivariance," " achieved by adding biases with small and controllable norms. We theoretically bound the equivariance violation introduced by this bias, and we also bound how such violations propagate through deep VN networks.
Non-spatial attributes. TFNs and the 3E(3)-Transforrm account for non-spatial attributes associated with each point (c.g., color, intensity), which they refer to as "type-0" features. In this work, we investigate two mechanisms (early & late fusion) to incorporate non-spatial data into the VN framework.
Attention-based architectures. Since the introduction of Transformers by Vaswani et al. (0117 self- attention and cross-attention mechanisms have provided powerful and versatile components which are propelling the field of natural language processing forward (Devlin et al. 2019 Liu et al.) 2019 Lan al. 2020: Yang al. 2020). Lately, so-called "Vision Transformers" (Dosovitskiy et al. 2021 Khan et al.) 2022 have had a similar impact on the field of computer vision, providing compelling alternative to convolutional
networks.
3 Background
3.1 Notation & preliminaries
Dataset. Suppose have dataset D Xp) where p {L...., P} is an index into point- cloud/label pair (Xp.Yp1 we omit the subscript p whenever it is unambiguous to do so. E X C Nx3 is a single 3D point-cloud with N points. In classification problem, Y [I,...,,k) [K,, where x is the number of classes. In regression problem, we might have Y ) C Nout x Sow where Nout is the number of output points, and Som is the dimension of cach output point (with Nout Sout = corresponding to univariate regression).
Index notation. We use "numpy-like" indexing of tensors. Assuming we have a tensor Z E AXxBxC we present some examples of this indexing scherne: Z(a) RBx Zlane) E RAXB Z ((ao:a E +0x=
Rotations & weights. Suppose we have tensor V E NxCx3 and rotation matrix R E SO(3), where SO(3) is the three-dimensionall rotation group. We denote the "rotation" of the tensor by VR NxCx3 defined as: (VR)()) & v(*) R, Vn E .4....,NY in other words, the rotated tensor VR is simply the concatenation of the N individually rotated matrices V(n R e RCx3 Additionally, we have a matrix of weights W RC'xC xC we define the product WV E NNC' x3 by (WV) A WV(").
Invariance and equivariance.
Definition 1 (Rotation Invariance). f: + y is Inlation-invoriinn if VR E SO(3), X XEX, f(XR) = f(X).
Definition 2 (Rotation Equivariance). f:X' + + is rotation-equivariant if VR SO(3), X X, f(XR) f(X)R.
For simplicity, we defined aaarrianceeqquivaianneee as above instead of the more general f(Xpx(() I ((X)pr which requires background group theory and representation theory.
4

---

Published in Transactions on Machine Learning Research (01/2023)
Proofs. We defer proofs to Appendix C
3.2 The Vector Neuron (VN) framework
In the Vector Neuron framework Deng al. 2021 the au- thors represent single point (e-g-, in a hidden layer of neu- ral network) as matrix V(n) E gxx3 (see inset figure), where V E VxCx3 can be thought of as a tensor representation of the entire point-cloud. This representation allows for the design of 50(3)-cquivariaa analogs of standard neural network operations.
Classical neurons
Vector Neurons
VN-Linear layer. As an illustrative example, the VN--inear layer function VN--inearr-;; ;W) RCx3 TR'x3 x3 defined by VN-L.neer(VVim) W) & WV(m), where W RC' xC is matrix of learnable weights. This operation is rofation-cquiiaarian:: VN-Lincar(V R; W) WV R (WV(") VN-Linear(VV"); W)R.
2021 also develop VN analogs of common deep network layers ReLU, MLP, BatchNorm, and Pool. For further definitions and proofs of equivariance, see Appendix [ For further details, we point the reader to Dong et al. [2021]
Deng et al.
4 The VN- VN-r
In this section, extend the ideas presented in Deng al. (021 design "VN-Transformer" that enjoys the rotation equivariance property.
4.1 Rotation-invariant inner product
The notion of an inner product between tokens central to the attention operation from the original Transformer (Vaswani et al. 2017). Consider the Frobenius inner product between two VN representations, defined below.
Definition 3 (Frobenius inner product). The Frobenius inner product between two matrices V(a), v(n)) e RC3 is defined by (V("), y(")) e E2 y(re.)) = - ((ac)y ,0),
This choice of inner product is convenient because of its rotation invariance property, stated below.
Proposition 1. The Probenius inner product between Vector Neuron representations V(n), v(n)) Cx3 is motation--mmariann,, tee. (V(m) R. v(n))
V())R)p R) (V("), v(n))p, VR E SO(3).
Proof.
(V(0) R,
R(n' vln')
BiF
Σ ylmol v RRTvinson
Σ pin (((n)) pin')s (1)
This rotation-invariant inner product between VN representations allows to construct rotation-equivariant allention operation, detailed in the next section.
4.2 Rotation-equivariant attention
Consider two lensors E M and K NxCxx3 which be thought of sets of M (resp. N) tokens, each C' x 3 matrix. Using the Frobenius inner product, we can define an attention matrix A(Q, K) MMx between the sets follows:
A(Q, K)(m) (m) A softmax
(Q)) k(())) |
(2)
V3C
5

---

Published in Transactions on Machine Learning Research (01/2023)
This attention operation is rotation-cquivariant w.r.t. simultancous rotation of all inputs: Following Vaswani et al. (2017), we divide the
(A
out
out
inner products by v3C since Q(m), K(a) RCx3 From Proposition 1
A(QR, K)) = K) VR e (3)
Finally, we define the operation VN-Attn MxCX3 x RNC33 x R RNx MxC'x as:
VN-LayerNorm
VN-MLP
VN-Attn(Q, K, 2)
Σ K) (m)))))m)))))))))) z(n)
(3)
WN-LayerNorm
Proposition 2. VN-Atta(QR, KR, ZR) VN-AAtn(Q,
Proof.
VN MN-MuutiheadAttn
VN-Attn(QR, KR, KR)
AQQR, (R))
(m..)
(1)
(N)
Figure 2:
VN-Transformer
A(Q, ((0, z(1) R= VN-Attn(Q. (,Z)(
(5) encoder block architecture.
VN-MultiheedAAtt and VN- defined in equalion G and equation 7 respectively. VN-MLP is a composition of VN-Linear, VN- TN-BatchNorm, and VN-ReLU layers from Deng al. 2021
LayerNorm
where (+) holds since A(QR, KR) = A(Q,K) (which follows straight-
forwardly from Proposition ] and equation 2
This is extendable to multi-head attention with H heads, VN-MuliiHeadAttu MxCx3 x RxCx3 x NNC'xx + Mx
VX-MultiHeedAAtnn( K, & w [VN-Atu(WWQQ
(6)
h=1*
where W W E Pxc W PxC' are feature, key, and value projection matrices of the h-th head (respectively), and po RC XHP is output projection matrix (Vaswani et al. 2017 VN-MulciHeaddAtn is also rotaationequuivarann, and is the key building block of our rotation-equiivariant VN- Transffrrmrr
We note that similar idea was proposed in Fuchs et al. 2020) namely, they use inner products between equivariant representations (obtained from the TFN framework of Thomas et al. (2018)) to create rotation- invariant attention matrix and a votation-equivariann attention mechanism. Our attention mechanism can be thought as the same treatment applied to the VN framework. For a more detailed comparison between this work and the proposal of Fuchs et al. (020 see Appendix A
4.3 Rotation-cqqivariann layer normalization
Deng al. 2021 allude to rotation-cquiivaiaan version of the
well-known layer normalization operation (Ba et al. 2016), but do not explicitly provide it we do here for completeness (see
Row-wise norm
Layer Norm
E Rx
Figure 3):
(V- & ((ne) C
(7)
R0x3
ERC
R
Figure 3: VN-LLyerNorm: rotation- equivariant layer normalization. "Layer Norm" is the standard layer normaliza- tion operation of Ba et al. (2016). O and are row-wise multiplication and division.
@ LayerNorm ([Iv(nic)))] 11x3.
[[]
where @ is an elementwise product, LayerNorm: BC + RR is the layer normalization operation of Ba et al. (016), and 1xx is a row-vector of ones.
In practice, we set H and such that HP= C'.

---

Published in Transactions on Machine Learning Research (01/2023)
4.4 Encoder architecture
Figure details the architecture of proposed rotation-oquivarian VN-Transformer encoder. The encoder is structurally identical to the original Transformer encoder of Vaswani et al. (0177), with each operation replaced by its Intatio-eqqivariaan VN analog.
5 Non-spatial attributes
"Real-world" point-clouds are typically augmented with cru- cial meta-data such as intensity & elongation for Lidar point-clouds and & sensor type for multi-sensor point-clouds. Handling such point-clouds while still satisfying equivari- ance/ /neeivaariance w.r.t. spatial inputs would be useful for many applications.
M..
WK-MPP
MLP
VN Transforms mir Encodar
concet
We investigate two strategies (late fusion and early fusion) to handle non-spatial attributes while maintaining rotation gqnivariance/ivaariin
Vie Invarint Encoder
Vk: Ri--Irrrrrmmrr VN-MPP
w.r.t. spatial dimensions:
Late fusion. In this approach, we propose to incorporate non-spatial ial attributes into the model a later stage, where we have already processed the spatial inputs in a rotation- equivariant fashion
VW-M ML-
Inputporrt-coo ood
Input Vnaectory
"late fusion" models for classifica- tion and trajectory prediction are shown in Figure 4 Early fusion. Early fusion is simple yet. powerful way to process non-spatial attributes (Jaegle et al. 2021). In this approach, we do not treat non-spatial attributes differently (see Figure D we simply concatenate spatial non-spatial inputs before feeding them into the VN-Transformer. The VN representations obtained are C x (3 + da) matrices (instead of C' x 3).
(a) Rotation-invariant classification model.
(b) Rotation--qquivariant trajectory forecasting model.
Figure 4: VN-Transformer (*late fusion") models.
Legend: (0(()-equivariant features; (3))-inva features; Non-spatial features.
6 Rotation-equivarian multi-scale feature aggregation
Jaegle et al. 2021) recently proposed an attention- architec- ture, PerceiverIO. which reduces the computational complexity of the vanilla Transformer by reducing the number of tokens (and their dimension) in the intermediate representations of the net- work. They achieve this reduction by learning a set Z E Mxx' of "latent features," which they use to perform QKV attention with the original input tokens X E NxC (with M << and C). Finally, they perform self-attention operations on the
K.V.
VN- MultiHeadAttn
VN-MeenProject
Figure 5: Rotation-equuivaiint latent features for Vector Neurons.
resulting x C' array, leading to a O(M²C") runtime instead of C' <<<
O(N2C) for each encoder self-attention operation, greatly improving time complexity during training and inference a boon for time-critical applications such as real-time motion forecasting. However, such learnable latent features would violate equivariance in our case, since these learnable features would have no information about the original input's orientation. To remedy this, we instead propose to a learn a transformation from the inputs to the latent fcatures (where the number of latent features is much smaller than the number of original inputs). Specifically, propose to use mean projection function PX-MeenPrroect(Vyeee W(m) ΣM v(m)), where W E BMxX'X xC is a learnable tensor. VN-McanProjeet is both rotation- oquivariant and permutation-invarian We then perform VN-MiuliHeaddAttottion between the resulting latent features and the original inputs V to get a smaller set of VN representations. The architecture diagram for our proposed rofation-equivariant "latent feature" mechanism is shown in Figure
7

---

Published in Transactions on Machine Learning Research (01/2023)
7 6-approximate equivariance
We noticed that distributed training on accelerated hardware is numerically unstable for points with small norms. This is unique to VN models - the VN-Linear layer does not include a bias vector, which leads frequent underflow issues on distributed accelerated hardware. We found that introducing small and controllable additíve biases fixes these issues - we modify the VN-Linear layer by adding a. bias with controllable norm:
WW-LinearWithBias((\W1 W,, c,U,,) wV( + cU, U & B ||
(8)
gC'xt where 2 0 is a hyperparameter controlling the bias norm, and B E R x3 is a learnable matrix. This leads to significant improvements in training stability and model quality. In principle, VN-Linear is not equivariant, but its violation equivariance can be bounded.
Work on equivariance by construction typically treats rotation equivariance as binary idea a model is either equivariant, it is not. This can be relaxed by asking: how large is the violation of equivariance? We quantify this with the equivariance violalion metric, defined by:
A(J.,, R) A (I ((X))))
(9)
If A(f,X,1) Se, say f is e-approzimately equivariant.
We bound the equivariance violation of (1-LiinarWiilBBies::;,,,,L0 RCx3 + R.Ox33 x3 below: Proposition 3. VN-Linear is
equivariant (tight when R = -I).
A natural next question is: how do such equivariance violations propagate through deep model?
Proposition 4. Suppose we have K functions fo Xss + Xx+1 (with Xs RCx3 Re+1 C RC411x & E (1....,k) K}) s.t.
fú is eu a-aaprorimately equivariant for all k E Kl,...,K}
2. fk is Lipischcit (w.r.t. || [[ for all lo E (2...., K}.
Then, the composition fK .... fi is E apprroiaaoa equivariant, where
cx... (oal...(aa(aa + e2) + es) + ...) eex.
(10)
Intuitively, each layer fa "stretches" the equivariance violation error of the previous layers by its Lipschitz constant Lg, and adds its violation eg to the total error.
8 Experiments
Table 1: ModelNet440 test accuracy. Top block shows
SO(3) -3vviinnar basclines taken from Deng et al. (2021) included here for convenience.
8.1 Rotation-invariant classification
Figure la shows our proposed VN- VN-Transformer architecture for classification. It consists of motation-equivariann operations (VN-MLP and VN-Transformer Encoder blocks), followed by an invariant operation (VN-Invariant), and fi- operations
Model
Acc.
# Params
TFN Thomas al.
2018
88.5%
RI-Conv (Zhang et al. 2019
86.5%
GC-Conv (Zhang et al.
2020
89.0%
VN-PointNeet Deng et al.
2021
77.2%
2.20M
nally standard get class predictions. The resulting log-
VN-DGCNN Deng et al.
2021
90.0%
2.00M
VN- Traasfoooreer (ours)
90.8%
0.04M
its/ als/clas predictions are rolation-invariant.
We evaluate our VN-Transformer classifier on the commonly used ModelNet400 dataset (Wu et al. 2015). a 40-class point-cloud classification problem. In Table [ we compare our model with recent rotation-invariin models. The VN-Tafoffrmeer outperforms the baseline VN models with orders of magnitude fewer parameters. Furthermore, we dispense with the computationally expensive edge-convolution used preprocessing step
8

---

Published in Transactions on Machine Learning Research (01/2023)
Accuracy
Speedup
0.94
1.0
Re Reletive Speedup
0.0
Accurapy
0.90
0.00
0.4
0.06
00.0
5
10
15
25 Neighbors
30
35
40
Figure 6: Example point-clouds from the ModelNet Polka-dot dataset. Cyan points correspond a, 0, and pink points correspond to a, 1. Note that the "airplanc" class has narrower polka-dot radius than the "table" class, since we make the polka-dot radius dependent on the object class.
Figure 7: VN-Transformer ModelNet40 accuracy and relative training speed vs. number of nearest neigh- bors used in edgo-convolution preprocessing. Speed is computed relative to zero neighbors (ie., edge- convolution). Edge-convolution slows training speed by ~5x and has little effect on model performance.
in the models of Deng et al. 2021 and find that the VN-Trrasformerr' performance relatively unaffected (see Figure
Accuracy
Speedup
1.00
-2.2
0.98
Table Test set accuracy on ModelNet40 "Polka- Dot" dataset. Top block shows ModelNet-40 re- sults (ie., only spatial inputs). Bottom block shows ModelNet40 Polka-dot results. is the bias norm in VNLinear WithBias of eq. equation
1..8
Rel ative Speciup
Accuracy
0.92
1.8
0.90
1.4
1.2
Model
Fusion
Features
Acc.
NN-PoontNet
0
at v.2]
77.2%
32
128 Latent Features
256
512
VN- TN-Trausfrmmer
0
88.5%
VN- TN-T
10-6
R.
90.8%
Figure 8: Test set accuracy on ModelNet40 vs. number of latent features (described in Figure 5). Latent features provide ~2x speedup with minimal acc. degradation (vs. row 3 of Table 2
VN-PointNet.
0
Early
R. P...2200
82.0%
VN- IN-Trausformer
0
Early
2.
91.1%
VN- Transfoorree
10-6
Early
sr 2.
95.4%
VN- Transformer
10
Late
2. 91.0%
8.2 Classification with non-spatial attributes
To evaluate model's ability handle non-spatial attributes, design modified version of the ModelNet40 dataset, called ModelNet40 "Polka-dot," which we construct (from ModelNet44)) as follows: to each point [r, Vi, 21]. we append E 13. Within radius r of a randomly chosen center, we randomly select 30 points and set a; = 1. We make the "polka-dot" radius r depend on the label y E (1,...,40) via r(y) è Tlo + 39 (y- 1) - -no), where nio = 0.3, Thi = 1. Figure 6 shows example point-clouds from the ModeINet40 Polka-dot dataset. By generating ModelNet40 Polka-dot this way, we directly embed class information into the non-spatial attributes. In order to perform well on this task, models need to effectively fuse spatial and non-spatial information (there
useful information in the non-spatial altributes alone since all point-clouds have Σ/" a, = 30). The results on ModelNet40 Polka-dot are shown in Table 2 Both VNN-oiit.Vet and Vn-Trassforrrr benefit significantly from the binary polka-dots, suggesting they are able to effectively combine spatial and non-spatial information.
9

---

Published in Transactions on Machine Learning Research (01/2023)
8.3 Latent features
Figure shows our results ModelNet440 when we reduce the number of tokens from 1024 (the number of points in the original point-cloud) to 32 using the latent feature mechanism presented in Figure 5 Using latent features provides a ~22 latency improvement (in training steps/sec) with minimal (~1.7%) accuracy degradation (compared with row of Table 2). This suggests a real benefit of the latent feature mechanism in time-sensitvve applications such as autonomous driving.
8.4 Rotation-equivariant motion forecasting
Table 3: Average Distance Error WOMD. + Lower better. used as data augmentation at training time, the bias in the VNLinearWithBias layers (equa- tion equation 8). The bottom block uses the speed an added input feature (via early fusion).
Figure 15 shows our proposed rotation-equivariian ar- chitecture for motion forecasting. In motion forecasting the goal is to predict the [ar, y; locations of an agent for sequence of future timesteps, given as input the past locations of the agent. We evaluate the model simplified version of the Waymo Open Motion
random z-axis rotations
Dataset (WOMD; Ettinger et al. 2021
Model
Features
ADE (L)
We select. 4904 trajectories (3915 for training, 979 for testing).
Transformer
[ar, y. =]
5.01
Transformer +
[at, y.
4.51
Each trajectory consists of 91 [2, points for a single vehicle sampled at 5 Hz.
VN-Thansformer
[xt [.
4.91
VN- Transforrre
[at M.
3.95
We use the first 11 points (the past) input and we predict the remaining 80 points (the future).
VN-Transformer
[at, =,
5.01
VN-Transformer
10
[ar, Uy 2,
3.67
We evaluate the quality of our trajectory forecast- ing models using the Average Distance Error (ADE): ADE(Y,, P:) + E [|x((0) P.19 l|2, where T is the number of time-steps in the output trajectory and Y.YY E TRx are the ground-truth trajectory and the predicted trajectory, respectively. Results are shown in Table 3 Adding training-time random rolations about the z-axis yields improves the performance of the vanilla Transformer, and the VN-Transformer outperforms the vanilla Transformer (without the need for train-time ro- tation augmentations, thanks to equivariance). Figure shows example predictions on WOMD. Equivariance violations of the vanilla Transformer models (columns (a) and (b)) are clearly remonstrated here, in contrast with the equivariant VN-Transformer (column (c)).
(
(2)
(a)
(b)
(c)
Figure 9: Example predictions on WOMD. Log- end: Rectangle current car position, Red points imput trajectory, Colored streaks predicted trajectories. Columns are different trajectory models, with (a) Transformer, (b) Transformer
+ augmentations, and (c) VN-Transformer. Row (1) is one example in the dataset. Row (2) is 45 rotation of the input points in Row (1).
9 Conclusion
In this paper, we introduced the VN-Transformer, a rotattinn-qqiivariant Transformer model based on the Vector Neurons framework. VN-Transformer a significant step towards building powerful, modular, and easy-to-use models that have appealing equivariance properties for point-cloud data.
Limitations of our work include: (i) similar to previous work (Deng al. 2021 Qi et al. 2017 we assume input data has been mean-centered. This is sensitive to outliers, and prevents us from making single-pass predictions for multi-object problems (we have to independently mean-center each agent first). Similarly, we invariance/eeuiivraann equivariance (e.g., scale invariance) in this work; (ii) Proposition shows an error bound on the total equivariance violation of the network with Ly-Lipschitz layers. We know the Lipschitz constants of VN-Linear and EN-LLnearWWthhB. (see Appendix Cj. but we have not yet determined them for other layers (e.g., VN-ReLU, VN-MultiTHeddAttn). We will address these gaps in future
have not addressed other types invariance/
10

---

Published in Transactions on Machine Learning Research (01/2023)
work, and we will also leverage Vr-araosforrrrs to obtain state-of-the-art performance number of key benchmarks such as the full Waymo Open Motion Dataset and ScanObjectNN.
References
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. (Cited on Page 6
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. ArXiv, abs/ 2bs/2104.13445 2021. (Cited on Page 3
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint ar Xiv: 1512. 03012, 2015. (Cited on Page 3
Benjamin Chidester, Minh N. Do, and Jian Ma. Rotation equivariance and invariance in convolutional neural networks, 2018. (Cited on Page
Taco Cohen and Max Welling. Group equivariant convolutional networks. In Maria Florina Balcan and Kilian Q Weinberger (eds.), Proceediiggs of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 2990-2999 New York, New York, USA, 20-22 (Cited on Page 3
Jun 2016. PMLR. URL
mmr. ptwps//4/////eeee wtt.//16..taaee
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector
neurons: A general framework for sq33)-equivariant networks, 2021. (Cited on Page 5
S
10
14
15
20
21
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019. (Cited on Page
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021. (Cited on Page 4
Carlos Esteves, Christine Alen-Blaancheter Xiaowei Zhou, and Kostas Daniilidis. Polar transformer networks In International Conference on Learning Representations, 2018. (Cited on Page 2 3
Scott Ettinger, Shuyang Cheng, Benjamin Caine, Chenxi Liu, Hang Zhao, Sabeek Pradhan, Yuning Chai, Ben Sapp, Charles R. Qi, Yin Zhou, Zoey Yang, Aurélien Chouard, Pei Sun, Jiquan Ngiam, Vijay Vasudevan, Alexander McCauley, Jonathon Shlens, and Dragomir Anguelov. Large scale interactive motion forecasting for autonomous driving: The waymo open motion dataset. In Proceedings of the CEEE/VVE International Conference on Computer Vision (ICCV), pp. 9710-9719, October 2021. (Citod on Page 10
15)
Marc Anton Finzi, Gregory Benton, and Andrew Gordon Wilson. Residual pathway priors for soft equivariance constraints. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances én Neural Information Processing Systems, 2021. URL http:///ppparrrieeeeeffeoororeee
(Cited
Page 4
Fabian B. Fuchs, Daniel F. Worrall, Volker Fischer, and Max Welling, Se(3)-transsformer:: 3d roto-translation equivariant attention networks, 2020. (Cited on Page 2
6
Geoffrey F. Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. Tn International
conference on learning representations, 2018. (Cited on Page 3
Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. Advances in neural information processing systems, 0122077-0025, 2015. (Cited on Page 2 33)
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin lonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver io: A general architecture for structured inputs & outputs, 2021. (Cited Page 7
11

---

Published in Transactions on Machine Learning Research (01/2023)
Salman Khan, Muzammal Nascer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak
Shah. Transformers in vision: A survey. ACM Compuling Surveys, Jan 2022. ISSN 1557-7341. doi:
10.1145/3505244 URL http://dxx doi org/10../1. 10111114/3500 1145/3505244 (Cited on Page 2 4
Alex Krizhevsky, Nya Sutskever, and Geoffrey E, Hinton. Imagenet classification with deep convolutional
neural networks. Advances in neural information processing systems, 25:1097-1105, 2012. (Cited on Page
3
Zhenzhong Lan, Mingda Chen, Scbastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations, 2020. (Cited on Page
Alex H Lang, Sourabh Vora, Holger Caesar, Lubing Zhou, Jiong Yang, and Oscar Beijbom. Pointpillars: Fast encoders for object detection from point clouds. In Proceedings of the IEEE/CWE Conference on Computer Vision and Pallern Recognition, pp. 12797, 2019. (Cited on Page 3
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-baser permutation-invante
neural networks, 2019. (Cited on Page 2
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019. (Cited
on Page
Nyy Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. (Cited on Page 15
Diego Marcos, Michele Volpi, Nikos Komodakis, and Devis Tuia. Rotation equivariant vector field nelworks. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017. (Cited on
Page 3
Charles R. Qi, Hao Su, Kaichun Mo, and Leonidas J. Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation, 2017. (Cited on Page 2 10
Charles R Qi, Wei Liu, Chenxia Wu, Hao Su, and Leonidas J Guibas. Frustum pointnets for 3d object detection from rgb-d data. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 918 927, 2018. (Cited on Page 3
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equiiariant neural networks for 3d point clouds, 2018. (Cited on Page
4 6
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017. (Cited on Page 2 4 6
15
Bastiaan S. Veeling, Jasper Linmans, Jim Winkens, Taco Cohen, and Max Welling. Rotation equivariant. cnns for digital pathology, 2018. (Cited on Page 3
Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics, 2022. URL http::::///aa....o///a//////////
(Cited on Page 4
Daniel Worrall and Gabriel Brostow, Cubenet: Equivariance to 3d rotation and translation. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018. (Cited Page
Daniel F. Worrall, Stephan J. Garbin, Daniyar Turmukhambetov, and Gabriel J. Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017. (Cited Page 3
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes, 2015. (Cited on Page 15
Bin Yang, Wenjie Luo, and Raquel Urtasun. Pixor: Real-time 3d object detection from point clouds. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, pp. 7652-7660, 2018.
(Cited Page 3
12

---

Published in Transactions on Machine Learning Research (01/2023)
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov,, and Quoc V. Le. Xinet: Generalized autoregressive pretraining for language understanding, 2020. (Cited on Page
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets, 2018. (Cited Page 2
Zhiyuan Zhang, Binh-Son Hua, David W. Rosen, and Sai-Kit Yeung. Rotation invariant convolutions for 3d point clouds deep learning, 2019. URL
06297 (Cited on Page 8
Zhiyuan Zhang, Binh-Son Hua, Wei Chen, Yibin Tian, and Sai-Kit Yeung. Global context aware convolutions for 3d point cloud understanding. In 2020 International Conference on 3D Vision (3DV), pp. 210 219, 2020. doi:
(Cited on Page
Yin Zhou and Oncel Tuzel. Voxelnet: End-to-end learning for point cloud based 3d object detection. In Proceedings of the IEEE conference on compuler vision and patlern recognilion, pp. 4490 4490-4499, 2018. (Cited on Page 3
13

---

Published in Transactions on Machine Learning Research (01/2023)
VN-Transformer (ours)
Trans-Transo (with & 1) (Fuchs et al., 2020)
Weight shape
W E RC'xO (C,C'= # of channels)
W E R (e p(a(41)x(2++11 in
the general case)
Action on input
Left- meft-multipl; wyyl") (y(r) E c
Right-multiply:
fW (ft ER
Requirement for equiv- None, w E RCx x arbitrary
W Σ m--j po
ariance
E ([|ad] ) (esl/////
(F= radial neural net., Yjm=
spherical harmonics, @ Clebsch-Gordan coefficients)
Table 4: Attention query computation in VN-Transformer vs. 3E(3)-Transformm WQ E CC'C is the VN-Transformer weight matrix used to compute the query (C and C" are the number of input and query channels, respectively).
A Detailed comparison with
(Fuchs et al. 2020)
A.1 Attention computation
There is a rich literature on equivariant models using steerable kernels, and the (e(())Transformet is the closest development this field our work. Here, we make a dctailod comparison between our work and the Tra3)-Traasformer (and related steerable kernel-based models). For simplicily, will compare the VN-Transformer with only spatial features vs. the features).
with only type-1 features (i.e., spatial
The key difference is in the way the weight matrices are defined and how they interact with the input.
Specifically, Fuchs al. use 3x3 weight matrices W Wil W R3x3 that act on the dimension of the input points (e.g., via f! where f E x1x3) As a result, in order to guarantee equivariance they need to design these matrices Will WT WH such that they each commute with rotation operation (in general (f, RRW + (f(W))) this depends on the choice of W hence the need for the machinery of Teebsch-Gordan coefficients, spherical harmonics, and radial neural nets construct the weights.
In contrast, in our proposed attention mechanism, the matrices WQ, WKK W E RC XC act on the channel dimension of the input (e.g., via wQy(m), where v(a) E RC33) and not the spatial dimension. As a result, the operations (aV(a), WK)(m,, (zy(") are equivariant no malter the choice of WQ, WK W since ((V(a) Rt) (WV This results in significantly simpler construction of rotation-equivariant attention that is (i) accessible to wider audience (i.e., it does not require an understanding of group theory, representation theory, spherical harmonics, Clebsch-Gordan coefficients, etc.) and (ii) much easier to implement.
For a side-by-side comparison of both attention query computations, see Table above.
A.2 VN-Linear vs. SE(3)- Transformer "self- interaction"
There is relationship between the VN-Linear operation of Deng et al. 2021) and the "linear relf-interaction"
layers of Fuchs et al. 2020). Comparing equation (12) of Fuchs et al. (2020) repeated here for convenience:
e Σ
el
(11)
Zwel' wit' wit (22++)))( in the general case, where E, E 1,2) are feature types
14

---

Published in Transactions on Machine Learning Research (01/2023)
with the VN-Linear operation of Deng et al. (2021
v(n)
WV W E CCx xc v(n) E RRX3
(12)
out
we see that these operations are identical.
However, our proposed VN-MuliHHeulAttention is different and significantly simpler than the attention mechanism of the (B(3)-Transformee (see Section A.1), as relies only on (t) the rotation-invarian Frobenius inner product and (di) straightforward multiplication by an arbitary weight matrix compute the keys, queries, and values (equivalent. to VN-Tinear /linear self-interaction). In that sense, it is closer in spirit to the original Transformer of Vaswani et al. (replacing vector inner product with Frobenius inner product). Further, explained previously, it does not require the special construction of weight matrices using Clebsch-Gordan coefficients, spherical harmonics, and radial neural networks as in the (rans--raaffrmmer.
B
Experimental details
B.1 Datasets
ModelNet40 The ModelNet40 dataset Wu et al. 2015 is publicly available at https://modelnet... cs. wwinceeenceoonoom edul with the following comment under "Copyright":
*All CAD models are dounloaded from the Internet and the original authors hold the copyright of the CAD models. The label of the dala was oblained by us via Amazon Mechanical Turk service and is provided freely. This dataset is provided for the convenience of academic research only."
Waymo Open Motion Dataset The Waymo Open Motion Dataset (Ettinger al. 2021) publicly available at httpp://wwaysee https////aayyoo..oo///a//ataaaaa com./open/dataa/ooiooll.// under non-commercial use license agreement. Full license details can be found here: https:// comp//e///e////////c
B.2 Hyperparameter tuning
Table 5 shows the hyperparameters we swept over for all our experiments on ModelNet400, ModelNet440 Polka-dot, and the Waymo Open Motion Datasct.
Hyperparameter
Value/Range
Feature dimension of VN-Transformer Number of attention heads
(32, 64, 128, 256, 512, 1024) {4, 8, 16, 32, 64, 128] (32, 64, 128, 256, 512)
Hidden layer dimension in encoder's VN-MLP Learning rate Learning rate schedule Optimizer
10
Lincar decay
AdamW (Loshchilov & Hutter 2019
Epochs
4000
E of NN-LLiearWWithBee
{0, 10-63
Table 5: Model hyperparameter ranges for ModelNet40, ModelNet400 Polka-dot, and Waymo Open Motion Dataset.
B.3 Compute infrastructure
We trained our models on TPU-v3 devices. which are accessible through Google Cloud. Our longest training jobs ran for less than hours 32 TPU cores.
15

---

Published in Transactions on Machine Learning Research (01/2023)
C Proofs
In this section, for convenience will treat all 3D vectors row-vectors: e.g., E RIx3, We also note that, while all our proofs of invariance/ hewinnceeiivviiinncee use matrices V(n) E RCx3, they can all be trivially generalized to v(n) E TRCxS
C.1 Partial invariance & equivariance
Assuming the input point-cloud consists of spatial inputs X E A' C N and associated non-spatial attributes A E A C IR Nxda (da is the number of non-spatial attributes associated with cach point), we would like our model f: X' x A + y to satisfy the following property:
Definition 4 (Partial rotation invariance). A model f : X' x A y satisfies portial rotation invariance if X E X, A A, RE SO(3), f(XR, A) = f(X, A).
Definition 5 (Partial rotation equivariance). A model f X x A + (where y C R (3N satisfies partial rolation equivariance if XEE, AE A, RE SO(3), J(XR, A) = [Xx
We show here that the models in Figure 1 Istiisly partial rotation invariance and equivariance (respectively). Consider two rotation matrices Rd, xd, E (SO(d,) and Rdgxda E (SO(d2).
Rd, xd, Odgxds
Od, xdz Royxds
Lemma 1. The metria Rids ((dd
is valid rotation matrir in SO(d, + d2).
Proof. We begin by showing that R] (ds))()d)
Rid (dd))(dd))))
We first compute R
R ((1+)))x((10))
Rds xdl Odexds
Odixds T Rayxda
R xd1 O xd2
da xdi R X2
0r
(+)
R Odgxds
xix
Odixda Raxxda
(13)
where (+) holds since Ra, xdi SO(d)) and Rasxda SO(d2) by assumption. Now, we compute Rid, (dxxxx)x
Rdy x1
Od, Xdyx Razxdz
Rdd x
(14)
Odzxds
[Rdxdds
Odg xdo Raixds 0daxdi]
Od xd2 Rdzxd= Od2 xd, Rt,xd xd, Od, xda] [1]
(15)
Odexdi
[
Od,
(16)
Odzxdi
Rd,xda
Hence, we have that Rid, (dd)
RR (( Finally, we show that det det +
det (Rt det ((Beg
[Rd, xdi (05yxd
Os, dz Rdgxds
) det det
1-1=1,
(17)
where (x) holds since Rd, (O(11) and Razxda (O(d2) by assumption.
Proposition 5. The VN-Transformer model f ::X A + y (where y C CR*) shown in Figure 1a
satisfies
partial rotation invariance.
Proof.
For convenience, we reparametrize the model f as fconeat Nx(+da)) R (with the number of object classes) where fooncut (X, f(X, It then suffices show that fconeat ([X R, Al) Sconcat (IX, Rl.
16

---

Published in Transactions on Machine Learning Research (01/2023)
First, note that feoncat is SO + da) -iviriiar,, since it is composed of SO da)-cvurrrr operations followed by SO(3 + da)-invariant operation.
R
03xda
Consider the matrix ((3td))((tda) A
where R E SO(3) is an arbitrary 3-
Idaxda] dimensional rotation. From Lemma R(3+da)x(3+da) 50(3++ +d)).
Odxx3
Sonneat (IX, A () (2) Seoneat (1X,
(18)
Sooncat (IX RR + A0 x3, X03xda Aldaxdal) S Sconcat ([X, A))
(19)
coon RX A]) = fconcat ((X, A)),
(20)
where (+) holds from SO(3 + d -iiiirrriinncee of feoncat.
Proposition 6. The VN- TN-Transformer model : X x A + (where c R (Noee((+dds) shown in Figure 1b solisfies partial rotation equiaariance.
Proof.
For convenience, we reparametrize the model f as Scoucat Nx++da) + IR None ((+ddx where Seoncat ([X, A)) = s(X, A). It then suffices to show that Sconcat (X R, A)) (..33 = Seoneat ( (X, A)) R. First, note that feoncat is SO(3 + -quieiuuva,,, since it composed of SO(3 + d eqquvvriiat operations.
R Oaxda
Consider the matrix ((++)))(3+da)) OdAx3
Idaxda] where R E SO(3) is an arbitrary 3- dimensional rotation. From Lemma ((stda)x(3tda) E S0(3 da).
Sconcaa ([X, A ((344)))((+44))) ( feoneal (IX, A)) (33 Sconcar (XR A0 X X03xda + Aldaxdal)
(21)
[feoncaa ()) R feomcar Al) (1:4:) 084x3, feomcas ([X, Al) Oaxds + fooncas Al) I
da xd
(22)
Scoucar (XR, [ecomcat (I) R, fcoment (IX, (I) Sconcar (XR, (1)6: Sconcat (X, (1) R.
(23)
(24)
where (+) holds from SO(3 + duuuquaivai of fconeat.
C.2 c-approximate equivariance
Proposition (Restated). VX-LinearWitthBassy W, U.e) is (2apCC")apprrommmaae) equivariont. This bound is tight when R -Tax3-
17

---

Published in Transactions on Machine Learning Research (01/2023)
Proof. Set A V W.U,c,
((XR) (1X=)-1XX)=(XXXRRG0 (WX eU)R
(25)
=cU-eUl E
(26)
AJ,, R (X -
(27)
Σ |E(U)6)) - -1())
(28)
2
10((] -U R13
(29)
10:00112 + [1 -0 RT
(30)
|0 + + + 20
(31)
11
62
Σ
411070011
(32)
4e20
(33)
R) S 22.
(34)
Lemma 2. Suppose
1. f: N1 + X2 (with X, c RCix Ny C R2x33) is C eaapproozimately equivariant.
2. 9: 22 + Xs (with X' C R0x33) is cgaapprozimately equivoriant and Lg Lipscchit (w.r.t. the
Frobenius norm).
Then the composition go f X1 + Xx is (Lies + =g))-approzimated equivariant.
Proof.
Algo f,X,R = (l((f(X))) -(((())
(35)
(l)(f(X))) ((((X)))) + -
(36)
(lo(F(X)))
-
(37)
[1o(((X)) +((()) +
(38)
(*))
( f(X)R +
(39)
= -I, X. R) 6g
(40)
" Lgeg + Cg:
(41)
where (+) holds from ey-approximate equivariance of g, (**) holds because g is Lipsshhhi, and (* (***) holds from approxiiee equivariance of
Proposition 4 (Restated). Suppose we have K functions Sk : Xs - Xe+i (with Xr C RC X C B3+11x3) for i e KK, satisfying the following:
1. fu is en-approrimately equivariant for all lo E (1,...,K] K}.
2. Jk. is Lapiiscchitz (w.r.t. || [|F) for all l E (2,..., K}.
18

---

Published in Transactions on Machine Learning Research (01/2023)
61
C3
f1
c1
f2
L261 + 6
f3
La(L2ei + + + Eg
Figure 10: Violations of equivariance in a neural network with L3- Li-Lipschitz and e--approximately equivariant layers.
Then, the composition fK .... fi is cI... -appprozimately equivariant, where:
eI..K t LL( (Ig(T1e1) + 62) + e3) + +...) eK
(42)
Proof.
AlfK o.... f,, =(fk (f..1.....0) RR),
(43)
(K) LKA(fK-1 0... 0,, 1,XX,, R) + +EK
(44)
(K-1)
< LR(LK-I(A(IK-200 Si,X, X, R) + +(-- EK
(45)
(46)
22 LK( (4s(,, R) + +2) + cs) + +...) + CK
(47)
1 (-x(----La(L25 + + +aa) + +
(48)
where (K) - (1) hold from applying inequality equation f 2 fw-1 o... fi each step).
40 from Lemma K times (setting g - fk. and
Figure 10 illustrates the propagation of equivariance violations through a composition of functions. Proposition 7. The VN-Linear(:; W) CRxS + xx's layer ((W)-Lipschile w.r.t. the Frobenius norm, where (W) is the spectral norm of W. The same holds for VN-LinearWithBBass WithBias(: W, U.c).
Proof. Consider X2 E RCxS We can write:
Wx. W X211
Σ [ W.
(49)
Σ (X -x,x
(50)
a=1
0² (W) -x!i>)
(51)
==1
oW)) Σ ||x((.)) x1
(52)
oWW)2 X2112
(53)
Wx. 5 a(W)| X2llF
(54)
To see this for VN-LinearWithBias note that I| X; cU) (W + eU) = I[X; then show the same result. using the above proof.
WX
we can
19

---

Published in Transactions on Machine Learning Research (01/2023)
C.3 Equivariance of VN-LayerNorm
We define the VN analog of the layer normalization operation as follows:
V(n.))
([Iv(aco111]
VN-LayerNorm(((())) A
C LayerNorm
1xx
(55)
||Vino0)|2 c=1
Proposition 8. VN-LayerNorm: :C -> Cx3 is rolation-equaaariant.
Proof.
V(n.e) R
(Inviiee) p|l = (6)
VN-LayerNNorm(Von (V))))
LayerNorm
(56)
[[((n,e) R|l2
V(n.o) R
(0)
LayerNorm
[
(57)
[V(n.))
r'/m
V(n.e)
(6)
[
LayerNorm
(58)
||V(n.))|22
o==1
NN-LayerNormmvvenjeer
(59)
[VN-LayerNorm(("0111
(60)
VN-LayeNoor(((*>) R) VN-LayerNormm (V()))))m))))))))) R
(61)
where (*) holds from invariance of vector norms to rotations.
C.4 Definitions of VN layers from Deng et al. (2021)
VN-ReLU layer The VN-ReLU layer is constructed as follows: from a given representation V E Cx3x
we compute a feature set g E Cxx,
q ف WV("), W RCxc
(62)
Then, we compute a set of "learnable directions" & E Rxx3
b UV(), U U
(63)
Note that W.U are learnable square matrices. Finally, compute the output of the VN-ReL. operation VN-ReLU(- WWU) Rx3 + g0x3 as follows:
(0 (0) ( 2 0
VN-
g/)) (1
gle) k(0) TROT [ o.W.
(64)
Otherwise stated: if the inner product between the feature g()) and the learnable direction k(() positive, return g(e), else return the projection of g(0) onto the plane defined by the direction k(c), can be readily shown that VN-ReLU is rotation-equivariann (for a proof, see Appendix C.5
VN-Invariant layer VN-Invariant(-- ;W): RCx3 + x is defined as:
VN-Invariant W) A (n) VN-MLP(V(") W)T.
(65)
where VN-MLP(-; W): Cx3 + 3x3 is composition of VN-Linear and VN-ReLU layers, and W is the set of all learnable parameters in VN-MLP. can be easily shown that VN-Invariant is rotation-invariant (see Appendix C.5 for a proof).
VN-Batch Norm, VN-Pool For votation-equivariin analogs of the standard batch norm and pooling operations, point the reader to Deng al. 22021.
20

---

Published in Transactions on Machine Learning Research (01/2023)
C.5
Invariance & equivariance of VN layers of Deng et al. (2021)
Proposition 9.
Deng et al. 2021 VN-Linear(: () RCx3 + RC is rolation-equivarian..
Proof.
VN--Lnear((V R: W)
WV R= (WV())R VN--Lnear(V(o)), W)R
(66)
Proposition 10. Deng al. 2021
VN-ReLU RRxx RC33 is rolation-equiuariant
Proof.
( if R, R3 20
1X (*)
(67)
g (g R, bO
12
O.W.
q if (g(0),((0) 20
(68)
gie) - (0 [mo011
M/
M
O.W.
g(e) if (q()), ((())
> 0
g(e) - (0))), |
k6
k
R
(69)
|
O.W.
(N- (c)
(70)
[VN-ReLU(( (1) (1)) (c)
(71)
VN-ReLU(y( R) = VN- )
(72)
where (+) holds because g and are rotation-equivariant w.r.t. (nn) and (**) holds because vector inner products are rotation-invariant.
Proposition 11.
Deng al. 2021) VN-Invariant: RRCx3 - Cx3 rolation-inwariant.
Proof.
VN-Invariant (V (v R) = B/WNBNNNLPPP(V
(73)
(+) (()) (N-MIP(V("):W)) W)R)
(74)
VVW.RRVVVMMMLVV W)T
(75)
-VWNN-MLP(((M)) W)T
(76)
VN-IIvarrintt (V)),
(77)
where (x) holds by cquivariance of VN-MLP.
21


Reviews:
[{"review_id": "AQ37hSXLvF", "reviewer": "Reviewer_x351", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This topic of this paper is significant in practice, and combination of vector neurons and transformers can be an interesting attempt. The following concerns need to be revised before publishing:\n1.\tThe authors have claimed a number of contributions in the paper, which addresses different problems for feature encoding on point clouds. The writing and organization of this paper can be less logical, and some technical proposal can be less relevant in a story of rotation-equivariant feature encoding such as non-spatial attributes.\n2.\tStatements in the introduction should be checked carefully, e.g. 1) it is suggested to rewrite the first paragraph to emphasize the challenge of rotation-invariance in the perspective of representation learning of 3D object shape; 2) how existing rotation-agnostic methods respectively can model rotation invariance and equivariance in the third paragraph and why? \n3.\tTopics investigated in the related works seem to lack of sufficient details, which makes this paper less convincing on comparison with recent works.\n4.\tProof of rotation equivariant attention is given in appendix, which is suggested to place it in the main text to support the proposed method, as rotation-equivariant attention module based on Proposition 1 could be straightforward.\n5.\tIn sec. 5, two fusion methods can be trivial, while partial rotation invariance and equivariance can be inspiring for other researchers. \n6.\tContent in Sec 6 and 7 seem to be implementation tricks of the proposed VN-Transformer rather than technical contributions, although they are probably important in the experiments.\n7.\tIn the experiments, multiple settings of rotation-equivariant feature learning should be compared, such as so(3)/so(3), z/so(3); Table 2, results of VN-DGCNN should be provided. \n  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel rotation equivariant network -- VN-Transformer for semantic classification and motion forecasting on arbitrarily posed point clouds, whose effectiveness is verified on the popular datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See strengths and weaknesses", "broader_impact_concerns": "More details of the investigated topic should be provided such as difference between rotation invariance and equivariance.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IWqtSeUFxW", "reviewer": "Reviewer_DJaz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed transformer architecture is simple and fits nicely with vector neurons.\nThe evaluation of the invariant task (modelnet40) supports the advantages of VN equivariant transformers versus vanilla VN. \nThe paper validates theoretically essential properties of the suggested model. The proofs seem to be solid.\n\nThe main observation of the paper regarding transformers (weighted sum of equivariant reps is equivariant)  has already been pointed out in [1] in more general settings. In a sense, the differences to [1] are the choice of the activation function (learnable in VN) and representation types (VN only allows type 1 and is thus simpler). Therefore, I think the authors should consider including such a discussion as part of the presentation of the method. Moreover, the evaluation should include a more detailed comparison to [1]. Currently, it only compares to a version of VN or compares to TFN and non-equivariant versions. Furthermore, the Modelnet benchmark also seems a bit odd, as the input data is aligned up to a 90 degrees rotations around the up direction. Why isn't the vnpointnet/vndgnn tested with bias?\n\n\nThe paper includes contributions that are independent of equivariant transformers: eps-approximation and non-spatial attributes. It is not clear why they are incorporated in an equivariant transformer paper. Moreover, the discussion on non-spatial features seems trivial, as they are merely seemed to be as invariant features. \n\n\n[1]: Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d roto-translation\nequivariant attention networks, 2020", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper extends the vector-neurons rotation equivariant point cloud network to support transformers. The suggested model is built on the facts that i) the dot product of equivariant representations can be used to model an invariant score in an attention layer, and ii) the weighted sum of equivariant representations is an equivariant representation.\nFurthermore, some additions related to the Vector-Neurons framework are proposed such as the addition of bias to the linear layers and the incorporation of non-spatial features. \nThe model is tested on one invariant task (classification) and one equivariant task (motion prediction).\n", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Please address the weakness stated above. ", "broader_impact_concerns": "The paper does not provide a broader impact statement. I do not have concerns. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZXmquvO4D4", "reviewer": "Reviewer_3fB3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n1. All the techniques introduced in this paper, including rotation-equivariant attention, dealing with non-spatial attributes, multi-scale feature aggregation, bias with bounded equivariance error, etc. are all very useful in my opinion, and directly address the drawbacks of the original vector neuron (VN) framework. I believe they are essential improvements to make the vector neuron framework attractive to users. The effects of these improvements are clearly demonstrated in the experiments section.\n2. This paper communicates their motivation very well, and it is easy for readers to follow their math derivation. For both related work and methodology, the paper is divided into several parts, corresponding to their main contributions. So it has a good writing structure. The diagrams are also helpful.\n\n**Weaknesses:**\n\nThis paper basically presents a batch of incremental techniques that can improve the vector neuron framework. So in terms of significance, this paper does not introduce ground-breaking results that can inspire lots of future research. (But I still think it presents many useful techniques.)\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper makes the following contributions, which are justified mathematically and backed by empirical results:\n1. They propose rotation-equivariant attention mechanisms for vector neurons, which eliminates the need for heavy feature preprocessing.\n2. They introduce the support of non-spatial attributes, which expand the applicability of the vector neuron framework.\n3. They propose multi-scale feature aggregation, which greatly speeds up training and inference.\n4. They show that a small bias (which introduces a small equivariance error) could accelerate hardware and improve numerical stability.\n5. They demonstrate that the proposed VN-transformer can successfully be applied to 3D shape classification and motion forecasting, and empirical results support the use of the techniques mentioned above.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Ideally, results presented in this paper should include standard error to account for the randomness of model initialization and optimization.\n2. It may be worthwhile to discuss whether this framework can be generalized beyond rotation equivariance, maybe in future work.\n", "broader_impact_concerns": "I have no ethical concern over this submission.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Three experts reviewed this paper, and asked for additional clarification. After rebuttal, all three reviewers agreed with the acceptance of this paper. The AE checked the reviews and rebuttal, and would recommend the acceptance of this paper. In particular, 
the authours proposed many practical techniques to address the drawbacks of the original vector neuron (VN) framework, including rotation-equivariant attention, dealing with non-spatial attributes, multi-scale feature aggregation, bias with bounded equivariance error, etc. These contributions are well supported by the experimental results. 
Decision: Accept (poster)