Published in Transactions on Machine Learning Research (01/2023)
VN-Transformer: Rotation-Equivariant Attention for Vector
Neurons
Serge Assaad ∗ serge.assaad@duke.edu
Duke University
Carlton Downey cmdowney@waymo.com
Waymo LLC
Rami Al-Rfou rmyeid@waymo.com
Waymo LLC
Nigamaa Nayakanti nigamaa@waymo.com
Waymo LLC
Ben Sapp bensapp@waymo.com
Waymo LLC
Reviewed on OpenReview: https://openreview.net/forum?id=EiX2L4sDPG
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
developedframeworkofferingasimpleyeteffectiveapproachforderivingrotation-equivariant
analogsofstandardmachinelearningoperationsbyextendingone-dimensionalscalarneurons
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-worlddatasets;(iii)wederivearotation-equivariantmechanismformulti-scalereduction
of point-cloud resolution, greatly speeding up inference and training; (iv) we show that
small tradeoffs in equivariance ((cid:15)-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.
∗WorkdoneduringaninternshipatWaymoLLC.
1
3202
naJ
42
]VC.sc[
3v67140.6022:viXraPublished in Transactions on Machine Learning Research (01/2023)
1 Introduction
(cid:296)(cid:1044)(cid:1109)(cid:939)(cid:937)
A chair – seen from the front, the back, the top, or the side – is still a
chair. When driving a 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
worldincontext(seefigureontheright). Unfortunately,typicalmachine
learning models struggle to preserve equivariance/invariance when
appropriate –it is indeed challenging to equipneural networks with the (cid:126)(cid:94)(cid:1074)(cid:940)(cid:1075) (cid:126)(cid:94)(cid:1074)(cid:940)(cid:1075)
right inductive biases to represent 3D objects in an equivariant manner.
Modeling spatial data is a core component in many domains such as
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 in the real world. Preserving rotation invariance/equivariance 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 a 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); Qi et al. (2017).
Recently, approaches jointly addressing rotation invariance and equivariance are gaining momentum. These
(cid:28)(cid:271)(cid:265)(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)(cid:341)
can roughly be categorized into modeling invariance or equivariance by: (i) data augmentation, (ii) canonical
pose estimation, and (iii) model construction. Approaches (i) and (ii) do not guarantee exact equivariance
since they rely on model parameters to learn the right inductive biases (Qi et al., 2017; Esteves et al., 2018;
Jaderberg et 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 (iii) do provide equivariance
guarantees, including the Tensor Field Networks of Thomas et al. (2018) and the SE(3)-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. (2021) proposed a simple and generalizable framework, dubbed Vector Neurons
(VNs), that can be used to 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.
WhileDengetal.(2021)developedaframeworkandbasiclayers,manyissuesrequiredforpracticaldeployment
on real-world applications remain unaddressed. A summary of our contributions and their motivations are as
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 et al.,
2017), and they have become an essential modeling component in most domains. Our primary contribution
of this paper is to develop a VN formulation of soft attention by generalizing scalar inner-product based
attention to matrix inner-products (i.e., the Frobenius inner product). Thanks to Transformers’ ability to
model functions on sets (since they are permutation-equivariant), they are a 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.
Direct point-set input. The original VN paper relied on edge convolution as a pre-processing step
to 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 run
2Published in Transactions on Machine Learning Research (01/2023)
(cid:296)(cid:1044)(cid:1109)(cid:944) (cid:296)(cid:1044)(cid:1109)(cid:944)
(cid:65)(cid:265)(cid:334)(cid:65)(cid:265)(cid:177)(cid:334)(cid:299)(cid:241)(cid:177)(cid:177)(cid:299)(cid:265)(cid:241)(cid:177)(cid:310)(cid:1109)(cid:265)(cid:203)(cid:310)(cid:259)(cid:177)(cid:1109)(cid:203)(cid:303)(cid:259)(cid:177)(cid:303)(cid:1109)(cid:303)(cid:296)(cid:303)(cid:299)(cid:1109)(cid:213)(cid:296)(cid:209)(cid:299)(cid:213)(cid:1044)(cid:209)(cid:1044) (cid:38)(cid:298)(cid:38)(cid:315)(cid:298)(cid:241)(cid:334)(cid:315)(cid:177)(cid:241)(cid:334)(cid:299)(cid:241)(cid:177)(cid:177)(cid:299)(cid:265)(cid:241)(cid:177)(cid:310)(cid:1109)(cid:265)(cid:310)(cid:299)(cid:310)(cid:177)(cid:1109)(cid:310)(cid:253)(cid:299)(cid:213)(cid:177)(cid:203)(cid:253)(cid:213)(cid:310)(cid:271)(cid:203)(cid:299)(cid:310)(cid:271)(cid:341)(cid:299)(cid:341)
on accelerated hardware. Our proposed rotation-
equivariant attention mechanism learns higher-level fea-
tures directly from single points for arbitrary point-clouds
(see Section 4).
(cid:119)(cid:271)(cid:119)(cid:271)(cid:271)(cid:259)(cid:1109)(cid:271)(cid:1219)(cid:259)(cid:1109)(cid:1109)(cid:87)(cid:1219)(cid:1109)(cid:81)(cid:87)(cid:119)(cid:81)(cid:119) (cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:87)(cid:1080)(cid:81)(cid:87)(cid:119)(cid:81)(cid:119)
Handling points augmented with non-spatial at-
tributes. Real-world point-cloud datasets have compli-
cated features sets – the [x,y,z] spatial dimensions are
(cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:65)(cid:265)(cid:1080)(cid:334)(cid:65)(cid:265)(cid:177)(cid:334)(cid:299)(cid:241)(cid:177)(cid:177)(cid:299)(cid:265)(cid:241)(cid:177)(cid:310)(cid:265)(cid:310) (cid:157)(cid:213)(cid:157)(cid:241)(cid:233)(cid:213)(cid:238)(cid:241)(cid:233)(cid:310)(cid:213)(cid:238)(cid:209)(cid:310)(cid:213)(cid:1109)(cid:119)(cid:209)(cid:271)(cid:1109)(cid:119)(cid:271)(cid:271)(cid:259)(cid:271)(cid:259)
typically augmented with crucial non-spatial attributes
(cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:132)(cid:299)(cid:1080)(cid:177)(cid:132)(cid:265)(cid:299)(cid:177)(cid:303)(cid:265)(cid:232)(cid:271)(cid:303)(cid:232)(cid:299)(cid:271)(cid:264)(cid:299)(cid:213)(cid:264)(cid:299)(cid:213)(cid:1109)(cid:299)(cid:1109) (cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:132)(cid:299)(cid:1080)(cid:177)(cid:132)(cid:265)(cid:299)(cid:177)(cid:303)(cid:265)(cid:232)(cid:271)(cid:303)(cid:232)(cid:299)(cid:271)(cid:264)(cid:299)(cid:213)(cid:264)(cid:299)(cid:213)(cid:1109)(cid:299)(cid:1109)
[[x,y,z];[a]] where a can be high-dimensional. For exam-
(cid:38)(cid:265)(cid:38)(cid:203)(cid:265)(cid:271)(cid:203)(cid:209)(cid:271)(cid:213)(cid:209)(cid:299)(cid:213)(cid:299) (cid:38)(cid:265)(cid:38)(cid:203)(cid:265)(cid:271)(cid:203)(cid:209)(cid:271)(cid:213)(cid:209)(cid:299)(cid:213)(cid:299)
ple, Lidar point-clouds have intensity & elongation 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
(cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:87)(cid:1080)(cid:81)(cid:87)(cid:119)(cid:81)(cid:119) (cid:156)(cid:88)(cid:156)(cid:1080)(cid:88)(cid:87)(cid:1080)(cid:81)(cid:87)(cid:119)(cid:81)(cid:119)
scope of their work to spatial point-cloud data, limiting
theapplicabilityoftheirmodelsforreal-worldpoint-clouds
with attributes. We investigate two mechanisms to inte- (cid:65)(cid:265)(cid:296)(cid:65)(cid:265)(cid:315)(cid:296)(cid:310)(cid:315)(cid:1109)(cid:296)(cid:310)(cid:271)(cid:1109)(cid:296)(cid:241)(cid:265)(cid:271)(cid:310)(cid:241)(cid:1080)(cid:265)(cid:203)(cid:310)(cid:1080)(cid:259)(cid:271)(cid:203)(cid:315)(cid:259)(cid:271)(cid:209)(cid:315)(cid:209) (cid:65)(cid:265)(cid:296)(cid:65)(cid:265)(cid:315)(cid:296)(cid:310)(cid:315)(cid:1109)(cid:310)(cid:310)(cid:299)(cid:1109)(cid:177)(cid:310)(cid:253)(cid:299)(cid:213)(cid:177)(cid:203)(cid:253)(cid:213)(cid:310)(cid:271)(cid:203)(cid:310)(cid:299)(cid:271)(cid:341)(cid:299)(cid:341)
grate attributes into equivariant models w(cid:28)(cid:271)(cid:265)(cid:367)h(cid:209)(cid:213)(cid:28)(cid:265)(cid:271)(cid:310)i(cid:241)(cid:265)(cid:177)(cid:367)(cid:259)l(cid:1109)(cid:1219)(cid:209)(cid:1109)(cid:213)e(cid:296)(cid:265)(cid:299)(cid:310)(cid:271)(cid:241)(cid:177)(cid:296)(cid:259)(cid:299)(cid:1109)(cid:1219)(cid:241)(cid:213)(cid:1109)p(cid:310)(cid:296)(cid:177)(cid:299)(cid:299)(cid:271)(cid:341)(cid:296)(cid:299)r(cid:241)(cid:213)(cid:310)(cid:177)e(cid:299)(cid:341)serving
(a)Rotation-invariantclas- (b) Rotation-equivariant
rotation equivariance (see Section 5). sificationmodel. trajectory forecasting
model.
Equivariantmulti-scalefeaturereduction. Practical Figure 1: VN-Transformer (“early fusion”) mod-
data structures such as Lidar point-clouds are extremely
els. Legend: ( ) SO(3)-equivariant features; ( )
large, consisting of hundreds of objects each with poten-
SO(3)-invariantfeatures;( )Non-spatialfeatures.
tially millions of points. To handle such computationally
challenging situations, we design a rotation-equivariant mechanism for multi-scale reduction of point-cloud
resolution. This mechanism learns how to pool the point set in a context-sensitive manner leading to a
significant reduction in training and inference latency (see Section 6).
(cid:15)-approximateequivariance. WhenattemptingtoscaleupVNmodelsfordistributedacceleratedhardware
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 (cid:15)-approximate equivariance, and use it to show that
small tradeoffs in equivariance can be controlled to obtain large improvements in numerical 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 7).
Empirical analysis. Finally, we evaluate our VN-Transformer on (i) the ModelNet40 shape classification
task, (ii) a modified ModelNet40 which includes per-point non-spatial attributes, 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
et al. (2021).
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
et 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).
Rotation-equivariant 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., 2017;
Chidester et al., 2018). Worrall & Brostow (2018) extended this work to 3D perception, and Veeling et al.
(2018) demonstrated the promise of rotation-equivariant models for medical images.
3Published in Transactions on Machine Learning Research (01/2023)
Equivariant point-cloud models. Thomas et al. (2018) 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)-Transformer” by adding an
attention mechanism for TFNs. One of the key ideas behind this body of work is to create highly restricted
weight matrices that commute with rotation operations by construction (i.e., WR = RW). In contrast,
we propose a simpler alternative: a “VN-Transformer” which guarantees equivariance for arbitrary weight
matrices, removing the need for the complex mathematical machinery of the SE(3)-Transformer. For a
detailed comparison with Fuchs et al. (2020), see Appendix A.
Controllable approximate equivariance. Finzietal.(2021)proposedequivariantpriorsonmodelweight
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
“(cid:15)-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 SE(3)-Transformer account for non-spatial attributes associated
with each point (e.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. (2017), 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 et al.,
2020; Yang et 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 a compelling alternative to convolutional
networks.
3 Background
3.1 Notation & preliminaries
Dataset. Suppose we have a dataset D (cid:44) {X ,Y }P , where p ∈ {1,...,P} is an index into a point-
p p p=1
cloud/label pair {X ,Y } – we omit the subscript p whenever it is unambiguous to do so. X ∈X ⊂RN×3
p p
is a single 3D point-cloud with N points. In a classification problem, Y ∈ Y ⊂ {1,...,κ}, where κ is the
number of classes. In a regression problem, we might have Y ∈Y ⊂RNout×Sout where N
out
is the number
of output points, and S is the dimension of each output point (with N = S = 1 corresponding to
out out out
univariate regression).
Index notation. We use “numpy-like” indexing of tensors. Assuming we have a tensor Z ∈ RA×B×C,
we present some examples of this indexing scheme: Z(a) ∈ RB×C, Z(:,:,c) ∈ RA×B, Z(alo:ahi) ∈
R(ahi−alo+1)×B×C.
Rotations & weights. Suppose we have a tensor V ∈RN×C×3 and a rotation matrix R∈SO(3), where
SO(3) is the three-dimensional rotation group. We denote the “rotation” of the tensor by VR∈RN×C×3,
defined as: (VR)(n) (cid:44) V(n)R, ∀n ∈ {1,...,N} – in other words, the rotated tensor VR is simply the
concatenation of the N individually rotated matrices V(n)R∈RC×3. Additionally, if we have a matrix of
weights W ∈RC0×C, we define the product WV ∈RN×C0×3 by (WV)(n) (cid:44)WV(n).
Invariance and equivariance.
Definition 1 (Rotation Invariance). f :X →Y is rotation-invariant if ∀R ∈SO(3), X ∈X, f(XR)=
f(X).
Definition 2 (RotationEquivariance). f :X →Y isrotation-equivariantif ∀R∈SO(3), X ∈X, f(XR)=
f(X)R.
For simplicity, we defined invariance/equivariance as above instead of the more general f(Xρ (g)) =
X
f(X)ρ (g), which requires background on group theory and representation theory.
Y
4Published in Transactions on Machine Learning Research (01/2023)
Proofs. We defer proofs to Appendix C. (cid:296)(cid:1044)(cid:1109)(cid:938)(cid:943)
3.2 The Vector Neuron (VN) framework
In the Vector Neuron framework (Deng et al., 2021), the au-
thors represent a single point (e.g., in a hidden layer of a neu-
ral network) as a matrix V(n) ∈ RC×3 (see inset figure), where
V ∈RN×C×3 can be thought of as a tensor representation of the
entire point-cloud. This representation allows for the design of
SO(3)-equivariant analogs of standard neural network operations.
(cid:28)(cid:259)(cid:177)(cid:303)(cid:303)(cid:241)(cid:203)(cid:177)(cid:259)(cid:1109)(cid:265)(cid:213)(cid:315)(cid:299)(cid:271)(cid:265)(cid:303) (cid:156)(cid:213)(cid:203)(cid:310)(cid:271)(cid:299)(cid:1109)(cid:88)(cid:213)(cid:315)(cid:299)(cid:271)(cid:265)(cid:303)
VN-Linear layer. Asanillustrativeexample, theVN-LinearlayerisafunctionVN-Linear(·;W):RC×3 →
RC0×3, defined by VN-Linear(V(n);W)(cid:44)WV(n), where W ∈RC0×C is a matrix of learnable weights. This
operation is rotation-equivariant: VN-Linear(V(n)R;W)=WV(n)R=(WV(n))R=VN-Linear(V(n);W)R.
Deng et al. (2021) also develop VN analogs of common deep network layers ReLU, MLP, BatchNorm, and
Pool. For further definitions and proofs of equivariance, see Appendix C. For further details, we point the
(cid:28)(cid:271)(cid:265)(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)(cid:341)
reader to Deng et al. (2021).
4 The VN-Transformer
In this section, we extend the ideas presented in Deng et al. (2021) to design a “VN-Transformer” that enjoys
the rotation equivariance property.
4.1 Rotation-invariant inner product
The notion of an inner product between tokens is 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(n),V(n0) ∈
RC×3 is defined by hV(n),V(n0)i (cid:44)PC P3 V(n,c,s)V(n0,c,s) =PC V(n,c)V(n0,c)(cid:124).
F c=1 s=1 c=1
This choice of inner product is convenient because of its rotation invariance property, stated below.
Proposition 1. The Frobenius inner product between Vector Neuron representations V(n),V(n0) ∈RC×3 is
rotation-invariant, i.e. hV(n)R,V(n0)Ri =hV(n),V(n0)i , ∀R∈SO(3).
F F
Proof.
C C C
hV(n)R,V(n0)Ri =X (V(n,c)R)(V(n0,c)R)(cid:124) =X V(n,c)RR(cid:124) V(n0,c)(cid:124) =X V(n,c)V(n0,c)(cid:124) =hV(n),V(n0)i
(1)
F F
c=1 c=1 c=1
Thisrotation-invariantinnerproductbetweenVNrepresentationsallowsustoconstructarotation-equivariant
attention operation, detailed in the next section.
4.2 Rotation-equivariant attention
ConsidertwotensorsQ∈RM×C×3 andK ∈RN×C×3,whichcanbethoughtofassetsofM (resp. N)tokens,
eachaC×3matrix. UsingtheFrobeniusinnerproduct, wecandefineanattentionmatrixA(Q,K)∈RM×N
between the two sets as follows:
A(Q,K)(m) (cid:44)softmax(cid:16) √1 (cid:2) hQ(m),K(n)i (cid:3)N (cid:17) , (2)
F
3C n=1
5Published in Transactions on Machine Learning Research (01/2023)
This attention operation is rotation-equivariant w.r.t. simultaneous (cid:1044)(cid:1109)(cid:1044)(cid:1109)(cid:1044)(cid:1109)
rotation of all inputs: Following Vaswani et al. (2017), we divide the
√
inner products by 3C since Q(m),K(n) ∈RC×3. From Proposition 1,
A(QR,KR)=A(Q,K) ∀R∈SO(3).
Finally, we define the operation VN-Attn : RM×C×3 × RN×C×3 ×
(cid:156)(cid:88)(cid:1080)(cid:81)(cid:177)(cid:341)(cid:213)(cid:299)(cid:88)(cid:271)(cid:299)(cid:264)
RN×C0×3 →RM×C0×3 as:
(cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)
N
X
VN-Attn(Q,K,Z)(m) (cid:44) A(Q,K)(m,n)Z(n). (3)
n=1
Proposition 2. VN-Attn(QR,KR,ZR)=VN-Attn(Q,K,Z)R. (cid:156)(cid:88)(cid:1080)(cid:81)(cid:177)(cid:341)(cid:213)(cid:299)(cid:88)(cid:271)(cid:299)(cid:264)
Proof. (cid:156)(cid:88)(cid:1080)(cid:87)(cid:315)(cid:259)(cid:310)(cid:241)(cid:238)(cid:213)(cid:177)(cid:209)(cid:1)(cid:376)(cid:265)
N
X (cid:1044)(cid:1109)(cid:1044)(cid:1109)(cid:1044)(cid:1109)
VN-Attn(QR,KR,ZR)(m) = A(QR,KR(cid:157)(cid:177)(cid:341)(cid:264))(cid:271)(cid:1109)((cid:1229)(cid:1109)(cid:28)m(cid:271)(cid:265)(cid:367)(cid:209),(cid:213)n(cid:265)(cid:310)(cid:241)(cid:177))(cid:259)(cid:1109)(cid:1219)Z(cid:1109)(cid:119)(cid:299)(cid:271)(cid:296)((cid:299)(cid:241)(cid:213)n(cid:310)(cid:177)(cid:299))(cid:341)R (4)
n=1
" # Figure 2: VN-Transformer
N
( =∗) X A(Q,K)(m,n)Z(n) R=VN-Attn(Q,K,Z)(m)R, (5) encoder block architecture.
VN-MultiheadAttn and VN-
n=1
LayerNorm are defined in
where (∗) holds since A(QR,KR)=A(Q,K) (which follows straight-
equation 6 and equation 7,
forwardly from Proposition 1 and equation 2).
respectively. VN-MLP is a
composition of VN-Linear,
This is extendable to multi-head attention with H heads,
VN-BatchNorm, and VN-ReLU
VN-MultiHeadAttn:RM×C×3×RN×C×3×RN×C0×3 →RM×C0×3:
layers from Deng et al. (2021).
VN-MultiHeadAttn(Q,K,Z)(cid:44)WO(cid:2) VN-Attn(WQQ,WKK,WZZ)(cid:3)H
, (6)
h h h h=1
where WQ,WK ∈ RP×C, WZ ∈ RP×C0 are feature, key, and value projection matrices of the h-th head
h h h
(respectively),andWO ∈RC0×HP isanoutputprojectionmatrix(Vaswanietal.,2017).1 VN-MultiHeadAttn
is also rotation-equivariant, and is the key building block of our rotation-equivariant VN-Transformer.
We note that a similar idea was proposed in Fuchs et al. (2020) – namely, they use inner products between (cid:296)(cid:1044)(cid:1109)(cid:938)(cid:940)
equivariant representations (obtained from the TFN framework of Thomas et al. (2018)) to create a rotation-
invariant attention matrix and a rotation-equivariant attention mechanism. Our attention mechanism can be
thought of as the same treatment applied to the VN framework. For a more detailed comparison between
this work and the proposal of Fuchs et al. (2020), see Appendix A.
4.3 Rotation-equivariant layer normalization
Deng et al. (2021) allude to a rotation-equivariant version of the
well-known layer normalization operation (Ba et al., 2016), but
do not explicitly provide it – we do so here for completeness (see
(cid:122)(cid:271)(cid:335)(cid:1080)(cid:335)(cid:241)(cid:303)(cid:213)(cid:1109) (cid:81)(cid:177)(cid:341)(cid:213)(cid:299)(cid:1109)
Figure 3):
(cid:265)(cid:271)(cid:299)(cid:264) (cid:88)(cid:271)(cid:299)(cid:264)
VN-LayerNorm(V(n))(cid:44) (7)
(cid:20) V(n,c) (cid:21)C (cid:12)LayerNorm(cid:16)(cid:2) ||V(n,c)|| (cid:3)C (cid:17) 1 , F eqig uu ivr ae ria3 n: t lV ayN e- rL na oy re mrN alo ir zm at: ionr .ot “a Lt ai yo en r-
||V(n,c)|| 2 c=1 2 c=1 1×3 Norm” is the standard layer normaliza-
where (cid:12) is an elementwise product, LayerNorm: RC →RC is the tion operation of Ba et al. (2016). (cid:12)
layer normalization operation of Ba et al. (2016), and 1 is a and (cid:11) are row-wise multiplication and
1×3
row-vector of ones. division.
(cid:28)(cid:271)(cid:265)(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)(cid:341)
1Inpractice,wesetH andP suchthatHP =C0.
6Published in Transactions on Machine Learning Research (01/2023)
4.4 Encoder architecture
Figure 2 details the architecture of our proposed rotation-equivariant VN-Transformer encoder. The encoder
is structurally identical to the original Transformer encoder of Vaswani et al. (2017), with each operation
replaced by its rotation-equivariant VN analog.
5 Non-spatial attributes
“Real-world” point-clouds are typically augmented w(cid:56)(cid:241)i(cid:233)t(cid:315)h(cid:299)(cid:213)(cid:56)(cid:1109)(cid:942)(cid:241)(cid:233)c(cid:202)(cid:315)r(cid:299)(cid:213)u(cid:1109)(cid:942)-(cid:202)
(cid:296)(cid:1044)(cid:1109)(cid:938)(cid:939) (cid:296)(cid:1044)(cid:1109)(cid:938)(cid:939)
cial meta-data such as intensity & elongation for Lidar
(cid:87)(cid:81)(cid:119) (cid:87)(cid:81)(cid:119) (cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)(cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)
point-cloudsand&sensortypeformulti-sensorpoint-clouds. (cid:119)(cid:271)(cid:271)(cid:259) (cid:119)(cid:271)(cid:271)(cid:259) (cid:157)(cid:213)(cid:241)(cid:233)(cid:238)(cid:310)(cid:213)(cid:157)(cid:209)(cid:213)(cid:1109)(cid:119)(cid:241)(cid:233)(cid:271)(cid:238)(cid:271)(cid:310)(cid:259)(cid:213)(cid:209)(cid:1109)(cid:119)(cid:271)(cid:271)(cid:259)
Handling such point-clouds while still satisfying equivari- (cid:87)(cid:81)(cid:119) (cid:87)(cid:81)(cid:119) (cid:156)(cid:88)(cid:1080)(cid:132)(cid:299)(cid:177)(cid:265)(cid:156)(cid:303)(cid:88)(cid:232)(cid:271)(cid:1080)(cid:132)(cid:299)(cid:264)(cid:299)(cid:177)(cid:213)(cid:265)(cid:299)(cid:303)(cid:1109)(cid:232)(cid:271)(cid:299)(cid:264)(cid:213)(cid:299)(cid:1109)
(cid:38)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:38)(cid:299)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:299)
ance/invariance w.r.t. spatial inputs would be useful for (cid:203)(cid:271)(cid:265)(cid:203)(cid:177)(cid:310)(cid:203)(cid:1109)(cid:271)(cid:265)(cid:203)(cid:177)(cid:310)(cid:1109)
(cid:203)(cid:271)(cid:265)(cid:203)(cid:177)(cid:310)(cid:203)(cid:1109)(cid:271)(cid:265)(cid:203)(cid:177)(cid:310)(cid:1109)
many applications.
We investigate two strategies (late fusion and early fusion)
(cid:56)(cid:259)(cid:177)(cid:376)(cid:213)(cid:265)(cid:56)(cid:259)(cid:177)(cid:376)(cid:213)(cid:265)
to handle non-spatial attributes while maintaining rotation (cid:156)(cid:88)(cid:1080)(cid:65)(cid:265)(cid:334)(cid:177)(cid:156)(cid:299)(cid:88)(cid:241)(cid:177)(cid:1080)(cid:265)(cid:65)(cid:310)(cid:265)(cid:334)(cid:177)(cid:299)(cid:241)(cid:177)(cid:265)(cid:310)
equivariance/invariance w.r.t. spatial dimensions: (cid:156)(cid:88)(cid:1080)(cid:132)(cid:299)(cid:177)(cid:265)(cid:156)(cid:303)(cid:88)(cid:232)(cid:271)(cid:1080)(cid:132)(cid:299)(cid:264)(cid:299)(cid:177)(cid:213)(cid:265)(cid:299)(cid:303)(cid:1109)(cid:232)(cid:271)(cid:299)(cid:264)(cid:213)(cid:299)(cid:1109) (cid:156)(cid:88)(cid:1080)(cid:132)(cid:299)(cid:177)(cid:265)(cid:156)(cid:303)(cid:88)(cid:232)(cid:271)(cid:1080)(cid:132)(cid:299)(cid:264)(cid:299)(cid:177)(cid:213)(cid:265)(cid:299)(cid:303)(cid:1109)(cid:232)(cid:271)(cid:299)(cid:264)(cid:213)(cid:299)(cid:1109)
(cid:38)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:38)(cid:299)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:299) (cid:38)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:38)(cid:299)(cid:265)(cid:203)(cid:271)(cid:209)(cid:213)(cid:299)
Late fusion. In this approach, we propose to incorporate (cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)(cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119) (cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)(cid:156)(cid:88)(cid:1080)(cid:87)(cid:81)(cid:119)
non-spatial attributes into the model at a later stage, where
we have already processed the spatial inputs in a rotation-
(cid:65)(cid:265)(cid:296)(cid:315)(cid:310)(cid:1109)(cid:296)(cid:271)(cid:65)(cid:265)(cid:241)(cid:265)(cid:296)(cid:310)(cid:315)(cid:1080)(cid:310)(cid:203)(cid:1109)(cid:296)(cid:259)(cid:271)(cid:271)(cid:315)(cid:241)(cid:209)(cid:265)(cid:310)(cid:1109)(cid:1080)(cid:203)(cid:259)(cid:271)(cid:315)(cid:209)(cid:1109) (cid:65)(cid:265)(cid:296)(cid:315)(cid:310)(cid:1109)(cid:310)(cid:299)(cid:177)(cid:65)(cid:265)(cid:253)(cid:213)(cid:296)(cid:203)(cid:315)(cid:310)(cid:310)(cid:271)(cid:1109)(cid:310)(cid:299)(cid:299)(cid:341)(cid:177)(cid:253)(cid:213)(cid:203)(cid:310)(cid:271)(cid:299)(cid:341)
equivariant fashion – our “late fusion” models for classi(cid:28)(cid:271)(cid:265)fi(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)c(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)a(cid:296)(cid:299)(cid:241)(cid:213)(cid:28)(cid:310)(cid:177)(cid:271)(cid:299)(cid:265)(cid:341)-(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)(cid:341)
tion and trajectory prediction are shown in Figure 4. (a) Rotation-invariant (b) Rotation-equivariant
classificationmodel. trajectory forecasting
Early fusion. Early fusion is a simple yet powerful way to model.
process non-spatial attributes (Jaegle et al., 2021). In this
approach, we do not treat non-spatial attributes differently Figure 4: VN-Transformer (“late fusion”)
(see Figure 1) – we simply concatenate spatial & non-spatial models.
inputs before feeding them into the VN-Transformer. The Legend: ( ) SO(3)-equivariant features; ( )
VN representations obtained are C × (3 + d ) matrices SO(3)-invariant features; ( ) Non-spatial
A
(instead of C×3). features.
(cid:87)(cid:315)(cid:259)(cid:310)(cid:241)(cid:62)(cid:213)(cid:177)(cid:209)(cid:1)(cid:376)(cid:265)
6 Rotation-equivariant multi-scale feature aggregation
Jaegle et al. (2021) recently proposed an attention-based architec-
ture, PerceiverIO, which reduces the computational complexity of
(cid:156)(cid:88)(cid:1080)
the vanilla Transformer by reducing the number of tokens (and (cid:87)(cid:315)(cid:259)(cid:310)(cid:241)(cid:62)(cid:213)(cid:177)(cid:209)(cid:1)(cid:376)(cid:265)
their dimension) in the intermediate representations of the net-
(cid:156)(cid:88)(cid:1080)(cid:87)(cid:213)(cid:177)(cid:265)(cid:119)(cid:299)(cid:271)(cid:253)(cid:213)(cid:203)(cid:310)
work. They achieve this reduction by learning a set Z ∈RM×C0
of “latent features,” which they use to perform QKV attention
with the original input tokens X ∈ RN×C (with M << N and
Figure 5: Rotation-equivariant latent
C0 <<C). Finally, they perform self-attention operations on the
resulting M ×C0 array, leading to a O(M2C0) (cid:157)r(cid:177)u(cid:341)(cid:264)n(cid:271)(cid:1109)(cid:1229)(cid:1109)(cid:28)t(cid:271)(cid:265)i(cid:367)m(cid:209)(cid:213)(cid:265)(cid:310)e(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:119)i(cid:299)(cid:271)n(cid:296)(cid:299)(cid:241)s(cid:213)(cid:310)(cid:177)t(cid:299)(cid:341)ead of features for Vector Neurons.
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
latentfeatureswouldviolateequivarianceinourcase,sincetheselearnablefeatureswouldhavenoinformation
about the original input’s orientation. To remedy this, we instead propose to a learn a transformation from
the inputs to the latent features (where the number of latent features is much smaller than the number
of original inputs). Specifically, we propose to use a mean projection function VN-MeanProject(V)(m) (cid:44)
h i
W(m) 1 PN V(n) , where W ∈ RM×C0×C is a learnable tensor. VN-MeanProject is both rotation-
N n=1
equivariant and permutation-invariant. We then perform VN-MultiHeadAttention between the resulting
latent features and the original inputs V to get a smaller set of VN representations. The architecture diagram
for our proposed rotation-equivariant “latent feature” mechanism is shown in Figure 5.
7Published in Transactions on Machine Learning Research (01/2023)
7 (cid:15)-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
to frequent underflow issues on distributed accelerated hardware. We found that introducing small and
controllable additive biases fixes these issues – we modify the VN-Linear layer by adding a bias with
controllable norm:
VN-LinearWithBias(V(n);W,U,(cid:15))(cid:44)WV(n)+(cid:15)U, U(c) (cid:44)B(c)/||B(c)|| , (8)
2
where (cid:15)≥0 is a hyperparameter controlling the bias norm, and B ∈RC0×3 is a learnable matrix. This leads
to significant improvements in training stability and model quality. In principle, VN-LinearWithBias is not
equivariant, but its violation of equivariance can be bounded.
Work on equivariance by construction typically treats rotation equivariance as a binary idea – a model is
either equivariant, or it is not. This can be relaxed by asking: how large is the violation of equivariance? We
quantify this with the equivariance violation metric, defined by:
∆(f,X,R)(cid:44)||f(XR)−f(X)R|| . (9)
F
If ∆(f,X,R)≤(cid:15), we say f is (cid:15)-approximately equivariant.
We bound the equivariance violation of VN-LinearWithBias(·;W,U,(cid:15)):RC×3 →RC0×3 below:
√
Proposition 3. VN-LinearWithBias is (2(cid:15) C0)-approximately equivariant (tight when R=−I).
A natural next question is: how do such equivariance violations propagate through a deep model?
Proposition 4. Suppose we have K functions f
k
: X
k
→ X
k+1
(with X
k
⊂ RCk×3,X
k+1
⊂ RCk+1×3,
k ∈{1,...,K}) s.t.
1. f is (cid:15) -approximately equivariant for all k ∈{1,...,K}
k k
2. f is L -Lipschitz (w.r.t. ||·|| ) for all k ∈{2,...,K}.
k k F
Then, the composition f ◦···◦f is (cid:15) -approximately equivariant, where
K 1 1...K
(cid:15) (cid:44)L (···(L (L (cid:15) +(cid:15) )+(cid:15) )+···)+(cid:15) . (10)
1...K K 3 2 1 2 3 K
Intuitively, each layer f “stretches” the equivariance violation error of the previous layers by its Lipschitz
k
constant L , and adds its own violation (cid:15) to the total error.
k k
8 Experiments
Table 1: ModelNet40 test accuracy. Top block shows
SO(3)-invariant baselines taken from Deng et al. (2021),
8.1 Rotation-invariant classification included here for convenience.
Figure 1a shows our proposed VN-Transformer Model Acc. # Params
architecture for classification. It consists of
TFN (Thomas et al., 2018) 88.5% –
rotation-equivariant operations (VN-MLP and
RI-Conv (Zhang et al., 2019) 86.5% –
VN-Transformer Encoder blocks), followed by GC-Conv (Zhang et al., 2020) 89.0% –
an invariant operation (VN-Invariant), and fi-
VN-PointNet (Deng et al., 2021) 77.2% 2.20M
nally standard Flatten/Pool/MLP operations
VN-DGCNN (Deng et al., 2021) 90.0% 2.00M
to get class predictions. The resulting log-
VN-Transformer (ours) 90.8% 0.04M
its/class predictions are rotation-invariant.
We evaluate our VN-Transformer classifier on the commonly used ModelNet40 dataset (Wu et al., 2015), a
40-class point-cloud classification problem. In Table 1, we compare our model with recent rotation-invariant
models. TheVN-TransformeroutperformsthebaselineVNmodelswithordersofmagnitudefewerparameters.
Furthermore, we dispense with the computationally expensive edge-convolution used as a preprocessing step
8(cid:296)(cid:1044)(cid:1109)(cid:938)(cid:944)
Published in Transactions on Machine Learning Research (01/2023)
0.94
0.92
0.90
0.88
0.86
0 5 10 15 20 25 30 35 40
# Nearest Neighbors
Figure 6: Example point-clouds from the ModelNet40
Polka-dot dataset. Cyan points correspond to a =0,
i
and pink points correspond to a =1. Note that the
i
“airplane” class has a narrower polka-dot radius than
the “table” class, since we make the polka-dot radius
dependent on the object class.
(cid:28)(cid:271)(cid:265)(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)(cid:341)
ycaruccA
1.2
1.0
0.8
0.6
0.4
0.2
0.0
pudeepS
evitaleR
Accuracy Speedup
Figure 7: VN-Transformer ModelNet40 accuracy and
relative training speed vs. number of nearest neigh-
bors used in edge-convolution preprocessing. Speed
is computed relative to zero neighbors (i.e., no 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-Transformer’s performance is relatively unaffected
(see Figure 7).
1.00
0.98
0.96
0.94
0.92
0.90
0.88
0.86
0.84
32 64 128 256 512
# Latent Features
ycaruccA
2.2
2.0
1.8
1.6
1.4
1.2
1.0
pudeepS
evitaleR
Accuracy Speedup
Table2: TestsetaccuracyonModelNet40“Polka-
Dot” dataset. Top block shows ModelNet40 re-
sults (i.e., only spatial inputs). Bottom block
showsModelNet40Polka-dotresults. (cid:15)isthebias
norm in VNLinearWithBias of eq. equation 8.
Model (cid:15) Fusion Features Acc.
VN-PointNet 0 – [x,y,z] 77.2%
VN-Transformer 0 – [x,y,z] 88.5%
VN-Transformer 10−6 – [x,y,z] 90.8%
Figure 8: Test set accuracy on ModelNet40 vs.
VN-PointNet 0 Early [x,y,z,a] 82.0%
number of latent features (described in Figure VN-Transformer 0 Early [x,y,z,a] 91.1%
5). Latent features provide a ∼2x speedup with VN-Transformer 10−6 Early [x,y,z,a] 95.4%
minimal acc. degradation (vs. row 3 of Table 2).
VN-Transformer 10−6 Late [x,y,z,a] 91.0%
8.2 Classification with non-spatial attributes
Toevaluateourmodel’sabilitytohandlenon-spatialattributes,wedesignamodifiedversionoftheModelNet40
dataset, called ModelNet40 “Polka-dot,” which we construct (from ModelNet40) as follows: to each point
[x ,y ,z ], we append a ∈ {0,1}. Within a radius r of a randomly chosen center, we randomly select
i i i i
30 points and set a = 1. We make the “polka-dot” radius r depend on the label y ∈ {1,...,40} via
i
r(y) (cid:44) r + 1 (y−1)(r −r ), where r = 0.3,r = 1. Figure 6 shows example point-clouds from the
lo 39 hi lo lo hi
ModelNet40 Polka-dot dataset. By generating ModelNet40 Polka-dot in 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 is no useful information in the non-spatial attributes alone
since all point-clouds have PN a =30). The results on ModelNet40 Polka-dot are shown in Table 2. Both
i=1 i
VN-PointNet and VN-Transformer benefit significantly from the binary polka-dots, suggesting they are able
to effectively combine spatial and non-spatial information.
9Published in Transactions on Machine Learning Research (01/2023)
8.3 Latent features
Figure 8 shows our results on ModelNet40 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 ∼2x latency improvement (in training steps/sec) with minimal (∼1.7%) accuracy
degradation (compared with row 3 of Table 2). This suggests a real benefit of the latent feature mechanism
in time-sensitive applications such as autonomous driving.
8.4 Rotation-equivariant motion forecasting
Table 3: Average (cid:8)Distance Error on WOMD. ↓
Figure 1b shows our proposed rotation-equivariant ar-
Lower is better. z = random z-axis rotations
chitectureformotionforecasting. Inmotionforecasting
used as data augmentation at training time. (cid:15) is
the goal is to predict the [x,y,z] locations of an agent
the bias in the VNLinearWithBias layers (equa-
for a sequence of future timesteps, given as input the
tionequation8). Thebottomblockusesthespeed
past locations of the agent. We evaluate the model
a as an added input feature (via early fusion).
on a simplified version of the Waymo Open Motion
Dataset (WOMD; Ettinger et al., 2021): Model (cid:15) Features ADE (↓)
• We select 4904 trajectories (3915 for training, 979 Transformer
(cid:8)
– [x,y,z] 5.01
for testing). Transformer+ z – [x,y,z] 4.51
• Each trajectory consists of 91 [x,y,z] points for a VN-Transformer 0 [x,y,z] 4.91
single vehicle sampled at 5 Hz. VN-Transformer 10−6 [x,y,z] 3.95
• We use the first 11 points (the past) as input and VN-Transformer 0 [x,y,z,a] 5.01 (cid:296)(cid:1044)(cid:1109)(cid:938)(cid:945)
we predict the remaining 80 points (the future). VN-Transformer 10−6 [x,y,z,a] 3.67
We evaluate the quality of our trajectory forecast-
(cid:1074)(cid:938)(cid:1075)(cid:1109)
ing models using the Average Distance Error (ADE):
ADE(Y ,Yˆ)(cid:44) 1 PT ||Y(t)−Yˆ(t)|| , where T is the
i i T t=1 i i 2
number of time-steps in the output trajectory and
Y ,Yˆ ∈RT×3 are the ground-truth trajectory and the
i i
predicted trajectory, respectively. Results are shown in (cid:1074)(cid:939)(cid:1075)(cid:1109)
Table 3. Adding training-time random rotations about
thez-axisyieldsimprovestheperformanceofthevanilla
Transformer, andtheVN-Transformeroutperformsthe (cid:1074)(cid:177)(cid:1075)(cid:1109) (cid:1074)(cid:202)(cid:1075)(cid:1109) (cid:1074)(cid:203)(cid:1075)(cid:1109)
vanillaTransformer(withouttheneedfortrain-timero- Figure 9: Example predictions on WOMD. Leg-
tation augmentations, thanks to equivariance). Figure end: Rectangle=currentcarposition,Redpoints
9 shows example predictions on WOMD. Equivariance = input trajectory, Colored streaks = predicted
violations of the vanilla Transformer models (columns (cid:28)(cid:271)(cid:265)(cid:367)(cid:209)(cid:213)(cid:265)(cid:310)t(cid:241)(cid:177)(cid:259)(cid:1109)(cid:1219)r(cid:1109)(cid:296)(cid:299)(cid:271)(cid:296)a(cid:299)(cid:241)(cid:213)(cid:310)(cid:177)(cid:299)j(cid:341)ectories. Columns are different trajectory
(a) and (b)) are clearly demonstrated here, in contrast mo (cid:8)dels, with (a) Transformer, (b) Transformer
with the equivariant VN-Transformer (column (c)). + z augmentations, and (c) VN-Transformer.
Row (1) is one example in the dataset. Row (2)
is a 45◦ rotation of the input points in Row (1).
9 Conclusion
In this paper, we introduced the VN-Transformer, a rotation-equivariant Transformer model based on the
Vector Neurons framework. VN-Transformer is 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 et 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
havenotaddressedothertypesofinvariance/equivariance(e.g.,scaleinvariance)inthiswork; (ii)Proposition
4 shows an error bound on the total equivariance violation of the network with L -Lipschitz layers. We
k
know the Lipschitz constants of VN-Linear and VN-LinearWithBias (see Appendix C), but we have not yet
determined them for other layers (e.g., VN-ReLU, VN-MultiHeadAttn). We will address these gaps in future
10Published in Transactions on Machine Learning Research (01/2023)
work, and we will also leverage VN-Transformers to obtain state-of-the-art performance on a 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/2104.13478, 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 arXiv:1512.03012, 2015. (Cited on Page 2, 3)
Benjamin Chidester, Minh N. Do, and Jian Ma. Rotation equivariance and invariance in convolutional neural
networks, 2018. (Cited on Page 3)
Taco Cohen and Max Welling. Group equivariant convolutional networks. In Maria Florina Balcan and
Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning,
volume 48 of Proceedings of Machine Learning Research, pp. 2990–2999, New York, New York, USA, 20–22
Jun 2016. PMLR. URL https://proceedings.mlr.press/v48/cohenc16.html. (Cited on Page 3)
Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas Guibas. Vector
neurons: A general framework for so(3)-equivariant networks, 2021. (Cited on Page 2, 5, 6, 8, 9, 10, 14, 15,
20, 21)
JacobDevlin, Ming-WeiChang, KentonLee, andKristinaToutanova. Bert: Pre-trainingofdeepbidirectional
transformers for language understanding, 2019. (Cited on Page 4)
AlexeyDosovitskiy,LucasBeyer,AlexanderKolesnikov,DirkWeissenborn,XiaohuaZhai,ThomasUnterthiner,
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 Allen-Blanchette, Xiaowei Zhou, and Kostas Daniilidis. Polar transformer networks.
In International Conference on Learning Representations, 2018. (Cited on Page 2, 3)
ScottEttinger, ShuyangCheng, BenjaminCaine, ChenxiLiu, HangZhao, SabeekPradhan,YuningChai, 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 IEEE/CVF International
Conference on Computer Vision (ICCV), pp. 9710–9719, October 2021. (Cited on Page 10, 15)
MarcAntonFinzi,GregoryBenton,andAndrewGordonWilson. Residualpathwaypriorsforsoftequivariance
constraints. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural
Information Processing Systems, 2021. URL https://openreview.net/forum?id=k505ekjMzww. (Cited
on Page 4)
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d roto-translation
equivariant attention networks, 2020. (Cited on Page 2, 4, 6, 14)
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. In 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, 28:2017–2025, 2015. (Cited on Page 2, 3)
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, 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 on Page 7)
11Published in Transactions on Machine Learning Research (01/2023)
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak
Shah. Transformers in vision: A survey. ACM Computing Surveys, Jan 2022. ISSN 1557-7341. doi:
10.1145/3505244. URL http://dx.doi.org/10.1145/3505244. (Cited on Page 2, 4)
Alex Krizhevsky, Ilya 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, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.
Albert: A lite bert for self-supervised learning of language representations, 2020. (Cited on Page 4)
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/CVF Conference on Computer
Vision and Pattern Recognition, pp. 12697–12705, 2019. (Cited on Page 3)
JuhoLee,YoonhoLee,JungtaekKim,AdamR.Kosiorek,SeungjinChoi,andYeeWhyeTeh. Settransformer:
A framework for attention-based permutation-invariant 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 4)
Ilya 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 networks.
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
detectionfromrgb-ddata. InProceedingsoftheIEEEconferenceoncomputervisionandpatternrecognition,
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-equivariant neural networks for 3d point clouds, 2018. (Cited on
Page 2, 4, 6, 8)
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, 5, 6, 7, 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 https://arxiv.org/abs/2201.11969. (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 on Page 3)
DanielE.Worrall,StephanJ.Garbin,DaniyarTurmukhambetov,andGabrielJ.Brostow. Harmonicnetworks:
Deep translation and rotation equivariance. In Proceedings of the IEEE Conference on Computer Vision
and Pattern Recognition (CVPR), July 2017. (Cited on 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 8, 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 on Page 3)
12Published in Transactions on Machine Learning Research (01/2023)
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. Xlnet:
Generalized autoregressive pretraining for language understanding, 2020. (Cited on Page 4)
ManzilZaheer, SatwikKottur, SiamakRavanbakhsh, BarnabasPoczos, RuslanSalakhutdinov, andAlexander
Smola. Deep sets, 2018. (Cited on 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 https://arxiv.org/abs/1908.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: 10.1109/3DV50981.2020.00031. (Cited on Page 8)
Yin Zhou and Oncel Tuzel. Voxelnet: End-to-end learning for point cloud based 3d object detection. In
Proceedings of the IEEE conference on computer vision and pattern recognition,pp.4490–4499,2018. (Cited
on Page 3)
13Published in Transactions on Machine Learning Research (01/2023)
VN-Transformer (ours) SE(3)-Transformer (with ‘=1)
(Fuchs et al., 2020)
Weight shape WQ ∈RC0×C (C,C0= # of channels) W11 ∈ R3×3 (∈ R(2‘0+1)×(2‘+1) in
Q
the general case)
Action on input Left-multiply: Right-multiply:
WQV(n)(cid:0)
V(n)
∈RC×3(cid:1) f1W11(cid:0)
f1
∈R1×3(cid:1)
i Q i
Requirement for equiv- None, WQ ∈RC0×C is arbitrary W11 =
Q
ariance P2 PJ ϕ11(||x ||)Y (x /||x ||)Q11 ,
J=0 m=−J J i Jm i i Jm
(ϕ= radial neural net., Y =
Jm
spherical harmonics, Q11 =
Jm
Clebsch-Gordan coefficients)
Table 4: Attention query computation in VN-Transformer vs. SE(3)-Transformer. WQ ∈ RC0×C is the
VN-Transformer weight matrix used to compute the query (C and C0 are the number of input and query
channels, respectively).
A Detailed comparison with SE(3)-Transformer (Fuchs et al., 2020)
A.1 Attention computation
There is a rich literature on equivariant models using steerable kernels, and the SE(3)-Transformer is the
closest development in this field to our work. Here, we make a detailed comparison between our work and
the SE(3)-Transformer (and related steerable kernel-based models). For simplicity, we will compare the
VN-Transformer with only spatial features vs. the SE(3)-Transformer with only type-1 features (i.e., spatial
features).
The key difference is in the way the weight matrices are defined and how they interact with
the input.
Specifically,Fuchsetal. use3×3weightmatricesW11,W11,W11 ∈R3×3thatactonthespatial/representation
Q K V
dimension of the input points (e.g., via f1W11 where f1 ∈ R1×3).2 As a result, in order to guarantee
i Q i
equivariance they need to design these matrices W11,W11,W11 such that they each commute with a rotation
Q K V
operation (in general (f1R)W11 6= (f1W11)R – this depends on the choice of W11), hence the need for
i Q i Q Q
the machinery of Clebsch-Gordan coefficients, spherical harmonics, and radial neural nets to construct the
weights.
In contrast, in our proposed attention mechanism, the matrices WQ,WK,WZ ∈RC0×C act on the channel
dimension of the input (e.g., via WQV(n), where V(n) ∈ RC×3) and not the spatial dimension. As a
result, the operations WQV(n),WKV(n),WZV(n) are equivariant no matter the choice of WQ,WK,WZ,
since W(V(n)R) = (WV(n))R. This results in a significantly simpler construction of rotation-equivariant
attention that is (i) accessible to a 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 4 above.
A.2 VN-Linear vs. SE(3)-Transformer “self-interaction”
ThereisarelationshipbetweentheVN-LinearoperationofDengetal.(2021),andthe“linearself-interaction”
layers of Fuchs et al. (2020). Comparing equation (12) of Fuchs et al. (2020), repeated here for convenience:
X
f‘ = w‘‘ f‘ , (11)
out,i,c0 c0c in,i,c
c
2W‘‘0,W‘‘0,W‘‘0 ∈R(2‘+1)×(2‘0+1) inthegeneralcase,where‘,‘0∈{0,1,2}arefeaturetypes
Q K V
14Published in Transactions on Machine Learning Research (01/2023)
with the VN-Linear operation of Deng et al. (2021):
V(n) =WV(n), W ∈RC0×C,V(n) ∈RC×3, (12)
out
we see that these operations are identical.
However, our proposed VN-MultiHeadAttention is different and significantly simpler than the attention
mechanismoftheSE(3)-Transformer(seeSectionA.1), asitreliesonlyon(i)therotation-invariantFrobenius
inner product and (ii) straightforward multiplication by an arbitary weight matrix to compute the keys,
queries, and values (equivalent to VN-Linear/linear self-interaction). In that sense, it is closer in spirit to the
originalTransformerofVaswanietal. (replacingvectorinnerproductwithFrobeniusinnerproduct). Further,
as 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 SE(3)-Transformer.
B Experimental details
B.1 Datasets
ModelNet40 The ModelNet40 dataset (Wu et al., 2015) is publicly available at
https://modelnet.cs.princeton.edu, with the following comment under “Copyright”:
“All CAD models are downloaded from the Internet and the original authors hold the copyright of the CAD
models. The label of the data was obtained by us via Amazon Mechanical Turk service and it is provided freely.
This dataset is provided for the convenience of academic research only.”
Waymo Open Motion Dataset The Waymo Open Motion Dataset (Ettinger et al., 2021) is publicly
available at https://waymo.com/open/data/motion/ under a non-commercial use license agreement. Full
license details can be found here: https://waymo.com/open/terms/.
B.2 Hyperparameter tuning
Table 5 shows the hyperparameters we swept over for all our experiments on ModelNet40, ModelNet40
Polka-dot, and the Waymo Open Motion Dataset.
Hyperparameter Value/Range
Feature dimension of VN-Transformer {32, 64, 128, 256, 512, 1024}
Number of attention heads {4, 8, 16, 32, 64, 128}
Hidden layer dimension in encoder’s VN-MLP {32, 64, 128, 256, 512}
Learning rate 10−3
Learning rate schedule Linear decay
Optimizer AdamW (Loshchilov & Hutter, 2019)
Epochs 4000
(cid:15) of VN-LinearWithBias {0, 10−6}
Table 5: Model hyperparameter ranges for ModelNet40, ModelNet40 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 3 hours on 32 TPU cores.
15Published in Transactions on Machine Learning Research (01/2023)
C Proofs
In this section, for convenience we will treat all 3D vectors as row-vectors: e.g., x∈R1×3. We also note that,
while all our proofs of invariance/equivariance use matrices V(n) ∈RC×3, they can all be trivially generalized
to V(n) ∈RC×S.
C.1 Partial invariance & equivariance
Assumingtheinputpoint-cloudconsistsofspatialinputsX ∈X ⊂RN×3 andassociatednon-spatialattributes
A∈A⊂RN×dA (d
A
is the number of non-spatial attributes associated with each point), we would like our
model f :X ×A→Y to satisfy the following property:
Definition 4 (Partial rotation invariance). A model f : X ×A → Y satisfies partial rotation invariance
if ∀ X ∈X, A∈A, R∈SO(3), f(XR,A)=f(X,A).
Definition 5 (Partial rotation equivariance). A model f :X ×A→Y (where Y ⊂RNout×(3+dA)) satisfies
partial rotation equivariance if ∀ X ∈X, A∈A, R∈SO(3), f(XR,A)(:,:3) =f(X,A)(:,:3)R.
We show here that the models in Figure 1 satisfy partial rotation invariance and equivariance (respectively).
Consider two rotation matrices R ∈SO(d ) and R ∈SO(d ).
d1×d1 1 d2×d2 2
(cid:20) (cid:21)
R 0
Lemma 1. The matrix R (cid:44) d1×d1 d1×d2 is a valid rotation matrix in SO(d +d ).
(d1+d2)×(d1+d2) 0
d2×d1
R
d2×d2
1 2
Proof. We begin by showing that R(cid:124) =R−1 . We first compute R(cid:124) :
(d1+d2)×(d1+d2) (d1+d2)×(d1+d2) (d1+d2)×(d1+d2)
R ((cid:124)
d1+d2)×(d1+d2)
=(cid:20) R
0
dd 21 ×× dd 11 R0 d d1 2× ×d d2 2(cid:21)(cid:124) =(cid:20) R 0(cid:124) dd(cid:124) 11 ×× dd 21 R0(cid:124) d d(cid:124)2 2× ×d d1 2(cid:21) (=∗)(cid:20) R
0
dd− 211 ×× dd 11 R0 d d−1 21× ×d d2 2(cid:21) , (13)
where (∗) holds since R ∈SO(d ) and R ∈SO(d ) by assumption.
d1×d1 1 d2×d2 2
Now, we compute R−1 :
(d1+d2)×(d1+d2)
(cid:20) (cid:21)−1
R 0
R−1 = d1×d1 d1×d2 (14)
(d1+d2)×(d1+d2) 0
d2×d1
R
d2×d2
"(cid:2) R −0 R−1 0 (cid:3)−1 0 #
= d1×d1 d1 0×d2 d2×d2 d2×d1
(cid:2) R −0
d1× Rd −2
1 0 (cid:3)−1
(15)
d2×d1 d2×d2 d2×d1 d1×d1 d1×d2
(cid:20) R−1 0 (cid:21)
= 0d1×d1 Rd −1 1×d2 . (16)
d2×d1 d2×d2
Hence, we have that R−1 =R(cid:124) . Finally, we show that det(R )=1:
(d1+d2)×(d1+d2) (d1+d2)×(d1+d2) (d1+d2)×(d1+d2)
(cid:18)(cid:20) (cid:21)(cid:19)
det(R )=det R d1×d1 0 d1×d2 =det(R )det(R )(=∗)1·1=1, (17)
(d1+d2)×(d1+d2) 0
d2×d1
R
d2×d2
d1×d1 d2×d2
where (∗) holds since R ∈SO(d ) and R ∈SO(d ) by assumption.
d1×d1 1 d2×d2 2
Proposition 5. The VN-Transformer model f :X ×A→Y (where Y ⊂Rκ) shown in Figure 1a satisfies
partial rotation invariance.
Proof.
• For convenience, we reparametrize the model f as f
concat
: RN×(3+dA) → Rκ (with κ the number
of object classes) where f ([X,A])=f(X,A). It then suffices to show that f ([XR,A])=
concat concat
f ([X,A])R.
concat
16Published in Transactions on Machine Learning Research (01/2023)
• First, note that f is SO(3+d )-invariant, since it is composed of SO(3+d )-equivariant
concat A A
operations followed by a SO(3+d )-invariant operation.
A
(cid:20) (cid:21)
R 0
• Consider the matrix R (cid:44) 3×dA , where R ∈ SO(3) is an arbitrary 3-
(3+dA)×(3+dA) 0
dA×3
I
dA×dA
dimensional rotation. From Lemma 1, R ∈SO(3+d ).
(3+dA)×(3+dA) A
f ([X,A]R )(=∗) f ([X,A]) (18)
concat (3+dA)×(3+dA) concat
⇒f ([XR+A0 ,X0 +AI ])=f ([X,A]) (19)
concat dA×3 3×dA dA×dA concat
⇒f ([XR,A])=f ([X,A]), (20)
concat concat
where (∗) holds from SO(3+d )-invariance of f .
A concat
Proposition 6. The VN-Transformer model f :X ×A→Y (where Y ⊂RNout×(3+dA)) shown in Figure 1b
satisfies partial rotation equivariance.
Proof.
• For convenience, we reparametrize the model f as f
concat
: RN×(3+dA) → RNout×(3+dA) where
f ([X,A])=f(X,A). It then suffices to show that f ([XR,A])(:,:3) =f ([X,A])(:,:3)R.
concat concat concat
• First, note that f is SO(3+d )-equivariant, since it is composed of SO(3+d )-equivariant
concat A A
operations.
(cid:20) (cid:21)
R 0
• Consider the matrix R (cid:44) 3×dA , where R ∈ SO(3) is an arbitrary 3-
(3+dA)×(3+dA) 0
dA×3
I
dA×dA
dimensional rotation. From Lemma 1, R ∈SO(3+d ).
(3+dA)×(3+dA) A
(∗)
f ([X,A]R ) = f ([X,A])R (21)
concat (3+dA)×(3+dA) concat (3+dA)×(3+dA)
⇒f ([XR+A0 ,X0 +AI ])
concat dA×3 3×dA dA×dA
h i
= f ([X,A])(:,:3)R+f ([X,A])(:,4:)0 ,f ([X,A])(:,:3)0 +f ([X,A])(:,4:)I
concat concat dA×3 concat 3×dA concat dA×dA
(22)
⇒f ([XR,A])=(cid:2) f ([X,A])(:,:3)R,f ([X,A])(:,4:)(cid:3) (23)
concat concat concat
⇒f ([XR,A])(:,:3) =f ([X,A])(:,:3)R, (24)
concat concat
where (∗) holds from SO(3+d )-equivariance of f .
A concat
C.2 (cid:15)-approximate equivariance
√
Proposition3(Restated). VN-LinearWithBias(·;W,U,(cid:15))is(2(cid:15) C0)-approximatelyequivariant. Thisbound
is tight when R=−I .
3×3
17Published in Transactions on Machine Learning Research (01/2023)
Proof. Set f (cid:44)VN-LinearWithBias(·;W,U,(cid:15)):
f(XR)−f(X)R=(WXR+(cid:15)U)−(WX+(cid:15)U)R (25)
=(cid:15)U −(cid:15)UR (26)
⇒∆(f,X,R)2 =||f(XR)−f(X)R||2 (27)
F
C0
X
= ||(cid:15)(U(c)−U(c)R)||2 (28)
2
c=1
C0
X
=(cid:15)2 ||U(c)−U(c)R||2 (29)
2
c=1
C0
=(cid:15)2X ||U(c)||2+||U(c)R||2−2U(c)R(cid:124) U(c)(cid:124) (30)
2 2
c=1
C0
≤(cid:15)2X ||U(c)||2+||U(c)R||2+2U(c)U(c)(cid:124) (31)
2 2
c=1
C0
X
=(cid:15)2 4||U(c)||2 (32)
2
c=1
=4(cid:15)2C0 (33)
√
⇒∆(f,X,R)≤2(cid:15) C0. (34)
Lemma 2. Suppose
1. f :X
1
→X
2
(with X
1
⊂RC1×3,X
2
⊂RC2×3) is (cid:15) f-approximately equivariant.
2. g : X
2
→ X
3
(with X
3
⊂ RC3×3) is (cid:15) g-approximately equivariant and L g-Lipschitz (w.r.t. the
Frobenius norm).
Then the composition g◦f :X →X is (L (cid:15) +(cid:15) )-approximately equivariant.
1 3 g f g
Proof.
∆(g◦f,X,R)=||g(f(XR))−g(f(X))R|| (35)
F
=||g(f(XR))−g(f(X)R)+g(f(X)R)−g(f(X))R|| (36)
F
≤||g(f(XR))−g(f(X)R)|| +||g(f(X)R)−g(f(X))R|| (37)
F F
(∗)
≤ ||g(f(XR))−g(f(X)R)|| +(cid:15) (38)
F g
(∗∗)
≤ L ||f(XR)−f(X)R|| +(cid:15) (39)
g F g
=L ∆(f,X,R)+(cid:15) (40)
g g
(∗∗∗)
≤ L (cid:15) +(cid:15) , (41)
g f g
where (∗) holds from (cid:15) -approximate equivariance of g, (∗∗) holds because g is L -Lipschitz, and (∗∗∗) holds
g g
from (cid:15) -approximate equivariance of f.
f
Proposition 4 (Restated). Suppose we have K functions f
k
: X
k
→ X
k+1
(with X
k
⊂ RCk×3,X
k+1
⊂
RCk+1×3) for k ∈{1,...,K}, satisfying the following:
1. f is (cid:15) -approximately equivariant for all k ∈{1,...,K}.
k k
2. f is L -Lipschitz (w.r.t. ||·|| ) for all k ∈{2,...,K}.
k k F
18Published in Transactions on Machine Learning Research (01/2023)
Figure 10: Violations of equivariance in a neural network with L -Lipschitz and (cid:15) -approximately equivariant
k k
layers.
Then, the composition f ◦···◦f is (cid:15) -approximately equivariant, where:
K 1 1...K
(cid:15) (cid:44)L (···(L (L (cid:15) +(cid:15) )+(cid:15) )+···)+(cid:15) (42)
1...K K 3 2 1 2 3 K
Proof.
∆(f ◦···◦f ,X,R)=∆(f ◦(f ◦···◦f ),X,R) (43)
K 1 K K−1 1
(K)
≤ L ∆(f ◦···◦f ,X,R)+(cid:15) (44)
K K−1 1 K
(K−1)
≤ L (L (∆(f ◦···◦f ,X,R)+(cid:15) )+(cid:15) (45)
K K−1 K−2 1 K−1 K
.
.
. (46)
(1)
≤ L (···(L (L ∆(f ,X,R)+(cid:15) )+(cid:15) )+···)+(cid:15) (47)
K 3 2 1 2 3 K
≤L (···(L (L (cid:15) +(cid:15) )+(cid:15) )+···)+(cid:15) (48)
K 3 2 1 2 3 K
where (K)−(1) hold from applying inequality equation 40 from Lemma 2 K times (setting g (cid:44) f and
k
f (cid:44)f ◦...f at each step).
k−1 1
Figure 10 illustrates the propagation of equivariance violations through a composition of 3 functions.
Proposition 7. The VN-Linear(·;W):RC×S →RC0×S layer is σ(W)-Lipschitz w.r.t. the Frobenius norm,
where σ(W) is the spectral norm of W. The same holds for VN-LinearWithBias(·;W,U,(cid:15)).
Proof. Consider X ,X ∈RC×S. We can write:
1 2
S
||WX −WX ||2 =X ||WX(:,s)−WX(:,s)||2 (49)
1 2 F 1 2 2
s=1
S
=X ||W(X(:,s)−X(:,s))||2 (50)
1 2 2
s=1
S
≤X σ2(W)||X(:,s)−X(:,s)||2 (51)
1 2 2
s=1
S
=σ(W)2X ||X(:,s)−X(:,s)||2 (52)
1 2 2
s=1
=σ(W)2||X −X ||2 (53)
1 2 F
⇒||WX −WX || ≤σ(W)||X −X || (54)
1 2 F 1 2 F
To see this for VN-LinearWithBias, note that ||(WX +(cid:15)U)−(WX +(cid:15)U)||2 =||WX −WX ||2 – we can
1 2 F 1 2 F
then show the same result using the above proof.
19Published 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:
(cid:20) V(n,c) (cid:21)C (cid:18)h iC (cid:19)
VN-LayerNorm(V(n))(cid:44) (cid:12)LayerNorm ||V(n,c)|| 1 (55)
||V(n,c)|| 2 c=1 2 c=1 1×3
Proposition 8. VN-LayerNorm:RC×3 →RC×3 is rotation-equivariant.
Proof.
V(n,c)R (cid:18)h iC (cid:19)(c)
VN-LayerNorm(V(n)R)(c) = LayerNorm ||V(n,c)R|| (56)
||V(n,c)R|| 2 2 c0=1
(=∗) V(n,c)R LayerNorm(cid:18)h ||V(n,c)|| iC (cid:19)(c) (57)
||V(n,c)|| 2 2 c0=1
" V(n,c) (cid:18)h iC (cid:19)(c)#
= LayerNorm ||V(n,c)|| R (58)
||V(n,c)|| 2 2 c0=1
=VN-LayerNorm(V(n))(c)R (59)
=[VN-LayerNorm(V(n))R](c) (60)
⇒VN-LayerNorm(V(n)R)=VN-LayerNorm(V(n))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(n) ∈RC×3,
we compute a feature set q ∈RC×3:
q (cid:44)WV(n), W ∈RC×C. (62)
Then, we compute a set of C “learnable directions” k ∈RC×3:
k (cid:44)UV(n), U ∈RC×C. (63)
Note that W,U are learnable square matrices. Finally, we compute the output of the VN-ReLU operation
VN-ReLU(· ;W,U):RC×3 →RC×3 as follows:
(
q(c) if hq(c),k(c)i≥0
VN-ReLU(V(n))(c) (cid:44) (64)
q(c)−hq(c), k(c) i k(c) o.w.
||k(c)|| ||k(c)||
Otherwise stated: if the inner product between the feature q(c) and the learnable direction k(c) is positive,
return q(c), else return the projection of q(c) onto the plane defined by the direction k(c). It can be readily
shown that VN-ReLU is rotation-equivariant (for a proof, see Appendix C.5).
VN-Invariant layer VN-Invariant(· ;W):RC×3 →RC×3 is defined as:
VN-Invariant(V(n);W)(cid:44)V(n)VN-MLP(V(n);W)(cid:124)
, (65)
where VN-MLP(· ;W):RC×3 →R3×3 is a composition of VN-Linear and VN-ReLU layers, and W is the
set of all learnable parameters in VN-MLP. It can be easily shown that VN-Invariant is rotation-invariant
(see Appendix C.5 for a proof).
VN-Batch Norm, VN-Pool For rotation-equivariant analogs of the standard batch norm and pooling
operations, we point the reader to Deng et al. (2021).
20Published 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(· ;W):RC×3 →RC0×3 is rotation-equivariant.
Proof.
VN-Linear(V(n)R;W)(cid:44)WV(n)R=(WV(n))R=VN-Linear(V(n);W)R (66)
Proposition 10. (Deng et al., 2021) VN-ReLU:RC×3 →RC×3 is rotation-equivariant.
Proof.
(
q(c)R if hq(c)R,k(c)Ri≥0
VN-ReLU(V(n)R)(c) (=∗) (67)
q(c)R−hq(c)R, k(c)R i k(c)R o.w.
||k(c)R||2 ||k(c)R||2
(
q(c)R if hq(c),k(c)i≥0
(∗=∗) (68)
q(c)R−hq(c), k(c) i k(c)R o.w.
||k(c)||2 ||k(c)||2
"( #
q(c) if hq(c),k(c)i≥0
= R (69)
q(c)−hq(c), k(c) i k(c) o.w.
||k(c)|| ||k(c)||2
=VN-ReLU(V(n))(c)R (70)
=[VN-ReLU(V(n))R](c) (71)
⇒VN-ReLU(V(n)R)=VN-ReLU(V(n))R, (72)
where (∗) holds because q and k are rotation-equivariant w.r.t. V(n) and (∗∗) holds because vector inner
products are rotation-invariant.
Proposition 11. (Deng et al., 2021) VN-Invariant:RC×3 →RC×3 is rotation-invariant.
Proof.
VN-Invariant(V(n)R;W)=(V(n)R)VN-MLP(V(n)R;W)(cid:124) (73)
(=∗) V(n)Rh VN-MLP(V(n);W)Ri(cid:124) (74)
=V(n)RR(cid:124)VN-MLP(V(n);W)(cid:124) (75)
=V(n)VN-MLP(V(n);W)(cid:124) (76)
=VN-Invariant(V(n);W), (77)
where (∗) holds by equivariance of VN-MLP.
21