SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Introduction

Self-attention mechanisms have enjoyed a sharp rise in popularity in recent years. Their relative implementational simplicity coupled with high efficacy on a wide range of tasks such as language modeling , image recognition , or graph-based problems , make them an attractive component to use. However, their generality of application means that for specific tasks, knowledge of existing underlying structure is unused. In this paper, we propose the SE(3)-Transformer shown in Fig. 1, a self-attention mechanism specifically for 3D point cloud and graph data, which adheres to equivariance constraints, improving robustness to nuisance transformations and general performance.

Point cloud data is ubiquitous across many fields, presenting itself in diverse forms such as 3D object scans , 3D molecular structures , or $N$-body particle simulations . Finding neural structures which can adapt to the varying number of points in an input, while respecting the irregular sampling of point positions, is challenging. Furthermore, an important property is that these structures should be invariant to global changes in overall input pose; that is, 3D translations and rotations of the input point cloud should not affect the output. In this paper, we find that the explicit imposition of equivariance constraints on the self-attention mechanism addresses these challenges. The SE(3)-Transformer uses the self-attention mechanism as a data-dependent filter particularly suited for sparse, non-voxelised point cloud data, while respecting and leveraging the symmetries of the task at hand.

Self-attention itself is a pseudo-linear map between sets of points. It can be seen to consist of two components: input-dependent attention weights and an embedding of the input, called a value embedding. In Fig. 1, we show an example of a molecular graph, where attached to every atom we see a value embedding vector and where the attention weights are represented as edges, with width corresponding to the attention weight magnitude. In the SE(3)-Transformer, we explicitly design the attention weights to be invariant to global pose. Furthermore, we design the value embedding to be equivariant to global pose. Equivariance generalises the translational weight-tying of convolutions. It ensures that transformations of a layer’s input manifest as equivalent transformations of the output. SE(3)-equivariance in particular is the generalisation of translational weight-tying in 2D known from conventional convolutions to roto-translations in 3D. This restricts the space of learnable functions to a subspace which adheres to the symmetries of the task and thus reduces the number of learnable parameters. Meanwhile, it provides us with a richer form of invariance, since relative positional information between features in the input is preserved.

The works closest related to ours are tensor field networks (TFN) and their voxelised equivalent, 3D steerable CNNs . These provide frameworks for building SE(3)-equivariant convolutional networks operating on point clouds. Employing self-attention instead of convolutions has several advantages. (1) It allows a natural handling of edge features extending TFNs to the graph setting. (2) This is one of the first examples of a nonlinear equivariant layer. In Section 3.2, we show our proposed approach relieves the strong angular constraints on the filter compared to TFNs, therefore adding representational capacity. This constraint has been pointed out in the equivariance literature to limit performance severely . Furthermore, we provide a more efficient implementation, mainly due to a GPU accelerated version of the spherical harmonics. The TFN baselines in our experiments leverage this and use significantly scaled up architectures compared to the ones used in .

We introduce a novel self-attention mechanism, guaranteeably invariant to global rotations and translations of its input. It is also equivariant to permutations of the input point labels.

We show that the SE(3)-Transformer resolves an issue with concurrent SE(3)-equivariant neural networks, which suffer from angularly constrained filters.

We introduce a Pytorch implementation of spherical harmonics, which is 10x faster than Scipy on CPU and $100-1000\times$ faster on GPU. This directly addresses a bottleneck of TFNs . E.g., for a ScanObjectNN model, we achieve $\approx 22\times$ speed up of the forward pass compared to a network built with SH from the lielearn library (see Appendix C).

Code available at https://github.com/FabianFuchsML/se3-transformer-public

Background And Related Work

where we used a softmax as a nonlinearity acting on the weights. In general, the number of query vectors does not have to equal the number of input points . In the case of self-attention the query, key, and value vectors are embeddings of the input features, so

where $\{h_{Q},h_{K},h_{V}\}$ are, in the most general case, neural networks . For us, query $\textbf{q}_{i}$ is associated with a point $i$ in the input, which has a geometric location $\textbf{x}_{i}$. Thus if we have $n$ points, we have $n$ possible queries. For query $\textbf{q}_{i}$, we say that node $i$ attends to all other nodes $j\neq i$.

Motivated by a successes across a wide range of tasks in deep learning such as language modeling , image recognition , graph-based problems , and relational reasoning , a recent stream of work has applied forms of self-attention algorithms to point cloud data . One such example is the Set Transformer . When applied to object classification on ModelNet40 , the input to the Set Transformer are the cartesian coordinates of the points. Each layer embeds this positional information further while dynamically querying information from other points. The final per-point embeddings are downsampled and used for object classification.

Permutation equivariance A key property of self-attention is permutation equivariance. Permutations of point labels $1,...,n$ lead to permutations of the self-attention output. This guarantees the attention output does not depend arbitrarily on input point ordering. Wagstaff et al. recently showed that this mechanism can theoretically approximate all permutation equivariant functions. The SE(3)-transformer is a special case of this attention mechanism, inheriting permutation equivariance. However, it limits the space of learnable functions to rotation and translation equivariant ones.

2 Graph Neural Networks

Attention scales quadratically with point cloud size, so it is useful to introduce neighbourhoods: instead of each point attending to all other points, it only attends to its nearest neighbours. Sets with neighbourhoods are naturally represented as graphs. Attention has previously been introduced on graphs under the names of intra-, self-, vertex-, or graph-attention . These methods were unified by Wang et al. with the non-local neural network. This has the simple form

where $w$ and $h$ are neural networks and $\mathcal{C}$ normalises the sum as a function of all features in the neighbourhood $\mathcal{N}_{i}$. This has a similar structure to attention, and indeed we can see it as performing attention per neighbourhood. While non-local modules do not explicitly incorporate edge-features, it is possible to add them, as done in Veličković et al. and Hoshen .

3 Equivariance

Given a set of transformations $T_{g}:\mathcal{V}\to\mathcal{V}$ for $g\in G$, where $G$ is an abstract group, a function $\phi:\mathcal{V}\to\mathcal{Y}$ is called equivariant if for every $g$ there exists a transformation $S_{g}:\mathcal{Y}\to\mathcal{Y}$ such that

The indices $g$ can be considered as parameters describing the transformation. Given a pair $(T_{g},S_{g})$, we can solve for the family of equivariant functions $\phi$ satisfying Equation 4. Furthermore, if $(T_{g},S_{g})$ are linear and the map $\phi$ is also linear, then a very rich and developed theory already exists for finding $\phi$ . In the equivariance literature, deep networks are built from interleaved linear maps $\phi$ and equivariant nonlinearities. In the case of 3D roto-translations it has already been shown that a suitable structure for $\phi$ is a tensor field network , explained below. Note that Romero et al. recently introduced a 2D roto-translationally equivariant attention module for pixel-based image data.

Group Representations In general, the transformations $(T_{g},S_{g})$ are called group representations. Formally, a group representation $\rho:G\to GL(N)$ is a map from a group $G$ to the set of $N\times N$ invertible matrices $GL(N)$. Critically $\rho$ is a group homomorphism; that is, it satisfies the following property $\rho(g_{1}g_{2})=\rho(g_{1})\rho(g_{2})$ for all $g_{1},g_{2}\in G$. Specifically for 3D rotations $G=SO(3)$, we have a few interesting properties: 1) its representations are orthogonal matrices, 2) all representations can be decomposed as

Eq. 7 and Eq. 9 present the convolution in message-passing form, where messages are aggregated from all nodes and feature types. They are also a form of nonlocal graph operation as in Eq. 3, where the weights are functions on edges and the features $\{\textbf{f}_{i}\}$ are node features. We will later see how our proposed attention layer unifies aspects of convolutions and graph neural networks.

Method

Here, we present the SE(3)-Transformer. The layer can be broken down into a procedure of steps as shown in Fig. 2, which we describe in the following section. These are the construction of a graph from a point cloud, the construction of equivariant edge functions on the graph, how to propagate SE(3)-equivariant messages on the graph, and how to aggregate them. We also introduce an alternative for the self-interaction layer, which we call attentive self-interaction.

Given a point cloud $\{(\textbf{x}_{i},\textbf{f}_{i})\}$, we first introduce a collection of neighbourhoods $\mathcal{N}_{i}\subseteq\{1,...,N\}$, one centered on each point $i$. These neighbourhoods are computed either via the nearest-neighbours methods or may already be defined. For instance, molecular structures have neighbourhoods defined by their bonding structure. Neighbourhoods reduce the computational complexity of the attention mechanism from quadratic in the number of points to linear. The introduction of neighbourhoods converts our point cloud into a graph. This step is shown as Step 1 of Fig. 2.

2 The SE(3)-Transformer

The SE(3)-Transformer itself consists of three components. These are 1) edge-wise attention weights $\alpha_{ij}$, constructed to be SE(3)-invariant on each edge $ij$, 2) edge-wise SE(3)-equivariant value messages, propagating information between nodes, as found in the TFN convolution of Eq. 7, and 3) a linear/attentive self-interaction layer. Attention is performed on a per-neighbourhood basis as follows:

These components are visualised in Fig. 2. If we remove the attention weights then we have a tensor field convolution, and if we instead remove the dependence of $\textbf{W}_{V}$ on $(\textbf{x}_{j}-\textbf{x}_{i})$, we have a conventional attention mechanism. Provided that the attention weights $\alpha_{ij}$ are invariant, Eq. 10 is equivariant to SE(3)-transformations. This is because it is just a linear combination of equivariant value messages. Invariant attention weights can be achieved with a dot-product attention structure shown in Eq. 11. This mechanism consists of a normalised inner product between a query vector $\textbf{q}_{i}$ at node $i$ and a set of key vectors $\{\textbf{k}_{ij}\}_{j\in\mathcal{N}_{i}}$ along each edge $ij$ in the neighbourhood $\mathcal{N}_{i}$ where

3 Node and Edge Features

Point cloud data often has information attached to points (node-features) and connections between points (edge-features), which we would both like to pass as inputs into the first layer of the network. Node information can directly be incorporated via the tensors $\textbf{f}_{j}$ in Eqs. 6, LABEL: and 10. For incorporating edge information, note that $\textbf{f}_{j}$ is part of multiple neighbourhoods. One can replace $\textbf{f}_{j}$ with $\textbf{f}_{ij}$ in Eq. 10. Now, $\textbf{f}_{ij}$ can carry different information depending on which neighbourhood $\mathcal{N}_{i}$ we are currently performing attention over. In other words, $\textbf{f}_{ij}$ can carry information both about node $j$ but also about edge $ij$. Alternatively, if the edge information is scalar, it can be incorporated into the weight matrices $\textbf{W}_{V}$ and $\textbf{W}_{K}$ as an input to the radial network (see step 2 in Fig. 2).

Experiments

We test the efficacy of the SE(3)-Transformer on three datasets, each testing different aspects of the model. The N-body problem is an equivariant task: rotation of the input should result in rotated predictions of locations and velocities of the particles. Next, we evaluate on a real-world object classification task. Here, the network is confronted with large point clouds of noisy data with symmetry only around the gravitational axis. Finally, we test the SE(3)-Transformer on a molecular property regression task, which shines light on its ability to incorporate rich graph structures. We compare to publicly available, state-of-the-art results as well as a set of our own baselines. Specifically, we compare to the Set-Transformer , a non-equivariant attention model, and Tensor Field Networks , which is similar to SE(3)-Transformer but does not leverage attention.

Similar to , we measure the exactness of equivariance by applying uniformly sampled SO(3)-transformations to input and output. The distance between the two, averaged over samples, yields the equivariance error. Note that, unlike in Sosnovik et al. , the error is not squared:

In this experiment, we use an adaptation of the dataset from Kipf et al. . Five particles each carry either a positive or a negative charge and exert repulsive or attractive forces on each other. The input to the network is the position of a particle in a specific time step, its velocity, and its charge. The task of the algorithm is then to predict the relative location and velocity 500 time steps into the future. We deliberately formulated this as a regression problem to avoid the need to predict multiple time steps iteratively. Even though it certainly is an interesting direction for future research to combine equivariant attention with, e.g., an LSTM, our goal here was to test our core contribution and compare it to related models. This task sets itself apart from the other two experiments by not being invariant but equivariant: When the input is rotated or translated, the output changes respectively (see Fig. 3).

We trained an SE(3)-Transformer with 4 equivariant layers, each followed by an attentive self-interaction layer (details are provided in the Appendix). Table 1 shows quantitative results. Our model outperforms both an attention-based, but not rotation-equivariant approach (Set Transformer) and a equivariant approach which does not levarage attention (Tensor Field). The equivariance error shows that our approach is indeed fully rotation equivariant up to the precision of the computations.

2 Real-World Object Classification on ScanObjectNN

Object classification from point clouds is a well-studied task. Interestingly, the vast majority of current neural network methods work on scalar coordinates without incorporating vector specific inductive biases. Some recent works explore rotation invariant point cloud classification . Our method sets itself apart by using roto-translation equivariant layers acting directly on the point cloud without prior projection onto a sphere . This allows for weight-tying and increased sample efficiency while transporting information about local and global orientation through the network - analogous to the translation equivariance of CNNs on 2D images. To test our method, we choose ScanObjectNN, a recently introduced dataset for real-world object classification. The benchmark provides point clouds of 2902 objects across 15 different categories. We only use the coordinates of the points as input and object categories as training labels. We train an SE(3)-Transformer with 4 equivariant layers with linear self-interaction followed by max-pooling and an MLP. Interestingly, the task is not fully rotation invariant, in a statistical sense, as the objects are aligned with respect to the gravity axis. This results in a performance loss when deploying a fully SO(3) invariant model (see Fig. 4(a)). In other words: when looking at a new object, it helps to know where ‘up’ is. We create an SO(2) invariant version of our algorithm by additionally feeding the $z$-component as an type-0 field and the $x$, $y$ position as an additional type-1 field (see Appendix). We dub this model SE(3)-Transformer +z. This way, the model can ‘learn’ which symmetries to adhere to by suppressing and promoting different inputs (compare Fig. 4(a) and Fig. 4(b)). In Table 2, we compare our model to the current state-of-the-art in object classificationAt time of submission, PointGLR was a recently published preprint . The performance of the following models was taken from the official benchmark of the dataset as of June 4th, 2020 (https://hkust-vgd.github.io/benchmark/): 3DmFV , PointNet , SpiderCNN , PointNet++ , DGCN .. Despite the dataset not playing to the strengths of our model (full SE(3)-invariance) and a much lower number of input points, the performance is competitive with models specifically designed for object classification - PointGLR , for instance, is pre-trained on the larger ModelNet40 dataset . For a discussion of performance vs. number of input points used, see Section D.1.2.

3 QM9

The QM9 regression dataset is a publicly available chemical property prediction task. There are 134k molecules with up to 29 atoms per molecule. Atoms are represented as a 5 dimensional one-hot node embeddings in a molecular graph connected by 4 different chemical bond types (more details in Appendix). ‘Positions’ of each atom are provided. We show results on the test set of Anderson et al. for 6 regression tasks in Table 3. Lower is better. The table is split into non-equivariant (top) and equivariant models (bottom). Our nearest models are Cormorant and TFN (own implementation). We see that while not state-of-the-art, we offer competitive performance, especially against Cormorant and TFN, which transform under irreducible representations of SE(3) (like us), unlike LieConv(T3), using a left-regular representation of SE(3), which may explain its success.

Conclusion

We have presented an attention-based neural architecture designed specifically for point cloud data. This architecture is guaranteed to be robust to rotations and translations of the input, obviating the need for training time data augmentation and ensuring stability to arbitrary choices of coordinate frame. The use of self-attention allows for anisotropic, data-adaptive filters, while the use of neighbourhoods enables scalability to large point clouds. We have also introduced the interpretation of the attention mechanism as a data-dependent nonlinearity, adding to the list of equivariant nonlinearties which we can use in equivariant networks. Furthermore, we provide code for a speed up of spherical harmonics computation of up to 3 orders of magnitudes. This speed-up allowed us to train significantly larger versions of both the SE(3)-Transformer and the Tensor Field network and to apply these models to real-world datasets.

Our experiments showed that adding attention to a roto-translation-equivariant model consistently led to higher accuracy and increased training stability. Specifically for large neighbourhoods, attention proved essential for model convergence. On the other hand, compared to convential attention, adding the equivariance constraints also increases performance in all of our experiments while at the same time providing a mathematical guarantee for robustness with respect to rotations of the input data.

Broader Impact

The main contribution of the paper is a mathematically motivated attention mechanism which can be used for deep learning on point cloud based problems. We do not see a direct potential of negative impact to the society. However, we would like to stress that this type of algorithm is inherently suited for classification and regression problems on molecules. The SE(3)-Transformer therefore lends itself to application in drug research. One concrete application we are currently investigating is to use the algorithm for early-stage suitability classification of molecules for inhibiting the reproductive cycle of the coronavirus. While research of this sort always requires intensive testing in wet labs, computer algorithms can be and are being used to filter out particularly promising compounds from large databases of millions of molecules.

Acknowledgements and Funding Disclosure

We would like to express our gratitude to the Bosch Center for Artificial Intelligence and Konincklijke Philips N.V. for funding our work and contributing to open research by publishing our paper. Fabian Fuchs worked on this project while on a research sabbatical at the Bosch Center for Artificial Intelligence. His PhD is funded by Kellogg College Oxford and the EPSRC AIMS Centre for Doctoral Training at Oxford University.

References

Appendix A Group Theory and Tensor Field Networks

A group is an abstract mathematical concept. Formally a group $(G,\circ)$ consists of a set $G$ and a binary composition operator $\circ:G\times G\to G$ (typically we just use the symbol $G$ to refer to the group). All groups must adhere to the following 4 axioms

Closure: $g\circ h\in G$ for all $g,h\in G$

Associativity: $f\circ(g\circ h)=(f\circ g)\circ h=f\circ g\circ h$ for all $f,g,h\in G$

Identity: There exists an element $e\in G$ such that $e\circ g=g\circ e=g$ for all $g\in G$

Inverses: For each $g\in G$ there exists a $g^{-1}\in G$ such that $g^{-1}\circ g=g\circ g^{-1}=e$

In practice, we omit writing the binary composition operator $\circ$, so would write $gh$ instead of $g\circ h$. Groups can be finite or infinite, countable or uncountable, compact or non-compact. Note that they are not necessarily commutative; that is, $gh\neq hg$ in general.

Groups are useful concepts, because they allow us to describe the structure of transformations, also sometimes called actions. A transformation (operator) $T_{g}:\mathcal{X}\to\mathcal{X}$ is an injective map from a space into itself. It is parameterised by an element $g$ of a group $G$. Transformations obey two laws:

Closure: $T_{g}\circ T_{h}$ is a valid transformation for all $g,h\in G$

Identity: There exists at least one element $e\in G$ such that $T_{e}[\textbf{x}]=\textbf{x}$ for all $\textbf{x}\in\mathcal{X}$,

where $\circ$ denotes composition of transformations. For the expression $T_{g}[\textbf{x}]$, we say that $T_{g}$ acts on x. It can also be shown that transformations are associative under composition. To codify the structure of a transformation, we note that due to closure we can always write

If for any $x,y\in\mathcal{X}$ we can always find a group element $g$, such that $T_{g}[x]=y$, then we call $\mathcal{X}$ a homogeneous space. Homogeneous spaces are important concepts, because to each pair of points $x,y$ we can always associate at least one group element.

As written in the main body of the text, equivariance is a property of functions $f:\mathcal{X}\to\mathcal{Y}$. Just to recap, given a set of transformations $T_{g}:\mathcal{X}\to\mathcal{X}$ for $g\in G$, where $G$ is an abstract group, a function $f:\mathcal{X}\to\mathcal{Y}$ is called equivariant if for every $g$ there exists a transformation $S_{g}:\mathcal{Y}\to\mathcal{Y}$ such that

If $f$ is linear and equivariant, then it is called an intertwiner. Two important questions arise: 1) How do we choose $S_{g}$? 2) once we have $(T_{g},S_{g})$, how do we solve for $f$? To answer these questions, we need to understand what kinds of $S_{g}$ are possible. For this, we review representations.

However, there are many more representations of $SO(3)$ than just the 3D rotation matrices. Among representations, two representations $\rho$ and $\rho^{\prime}$ of the same dimensionality are said to be equivalent if they can be connected by a similarity transformation

We also say that a representation is reducible if is can be written as

As it turns out, all linear representations of compact groupsOver a field of characteristic zero. (such as $SO(3)$) can be decomposed into a direct sum of irreps, as

where Q is an orthogonal, $N\times N$, change-of-basis matrix ; and each $\textbf{D}_{J}$ for $J=0,1,2,...$ is a $(2J+1)\times(2J+1)$ matrix known as a Wigner-D matrix. The Wigner-D matrices are the irreducible representations of $SO(3)$. We also mentioned that vectors transforming according to $\textbf{D}_{J}$ (i.e. we set $\textbf{Q}=\textbf{I}$), are called type-$J$ vectors. Type-0 vectors are invariant under rotations and type-1 vectors rotate according to 3D rotation matrices. Note, type-$J$ vectors have length $2J+1$. In the previous paragraph we mentioned that irreps act on orthogonal subspaces $\mathcal{X}_{0},\mathcal{X}_{1},...$. The orthogonal subspaces corresponding to the Wigner-D matrices are the space of spherical harmonics.

where $\textbf{D}_{J}$ is the $J$th Wigner-D matrix and $\textbf{D}_{J}^{*}$ is its complex conjugate. They form an orthonormal basis for (the Hilbert space of) square-integrable functions on the sphere $L^{2}(S^{2})$, with inner product given as

So $\left\langle Y_{Jm},Y_{J^{\prime}m^{\prime}}\right\rangle_{S^{2}}=\delta_{JJ^{\prime}}\delta_{mm^{\prime}}$, where $Y_{Jm}$ is the $m$th element of $\textbf{Y}_{J}$. We can express any function in $L^{2}(S^{2})$ as a linear combination of spherical harmonics, where

where each $\textbf{f}_{J}$ is a vector of coefficients of length $2J+1$. And in the opposite direction, we can retrieve the coefficients as

following from the orthonormality of the spherical harmonics. This is in fact a Fourier transform on the sphere and the the vectors $\textbf{f}_{J}$ can be considered Fourier coefficients. Critically, we can represent rotated functions as

In the main text we introduced the Clebsch-Gordan coefficients. These are used in the construction of the equivariant kernels. They arise in the situation where we have a tensor product of Wigner-D matrices, which as we will see is part of the equivariance constraint on the form of the equivariant kernels. In representation theory a tensor product of representations is also a representation, but since it is not an easy object to work with, we seek to decompose it into a direct sum of irreps, which are easier. This decomposition is of the form of Eq. 20, written

In Tensor Field Networks and 3D Steerable CNNs , the authors solve for the intertwiners between SO(3) equivariant point clouds. Here we run through the derivation again in our own notation.

Now let’s apply the equivariance condition to this expression, then

Now we notice that this expression should also be equal to Eq. 32, which is the convolution with an unrotated point cloud. Thus we end up at

which is sometimes refered to as the kernel constraint. To solve the kernel constraint, we notice that it is a linear equation and that we can rearrange it as

where we used the identity $\text{vec}(\textbf{AXB})=(\textbf{B}^{\top}\otimes\textbf{A})\text{vec}(\textbf{X})$ and the fact that the Wigner-D matrices are orthogonal. Using the Clebsch-Gordan decomposition we rewrite this as

Appendix B Recipe for Building an Equivariant Weight Matrix

One of the core operations in the SE(3)-Transformer is multiplying a feature vector f, which transforms according to $SO(3)$, with a matrix W while preserving equivariance:

In contrast to Weiler et al. , we do not voxelise space and therefore x will be different for each pair of points in each point cloud. However, the same $\textbf{Y}_{J}\left(\textbf{x}\right)$ will be used multiple times in the network and even multiple times in the same layer. Hence, precomputing them at the beginning of each forward pass for the entire network can significantly speed up the computation. The Clebsch-Gordan coefficients do not depend on the relative positions and can therefore be precomputed once and stored on disk. Multiple libraries exist which approximate those coefficients numerically.

Appendix C Accelerated Computation of Spherical Harmonics

The spherical harmonics (SH) typically have to be computed on the fly for point cloud methods based on irreducible computations, a bottleneck of TFNs . Thomas et al. ameliorate this by restricting the maximum type of feature to type-2, trading expressivity for speed. Weiler et al. circumvent this challenge by voxelising the input, allowing them to pre-compute spherical harmonics for fixed relative positions. This is at the cost of detail as well as exact rotation and translation equivariance.

The number of spherical harmonic lookups in a network based on irreducible representations can quickly become large (number of layers $\times$ number of points $\times$ number of neighbours $\times$ number of channels $\times$ number of degrees needed). This motivates parallelised computation on the GPU - a feature not supported by common libraries. To that end, we wrote our own spherical harmonics library in Pytorch, which can generate spherical harmonics on the GPU. We found this critical to being able to run the SE(3)-Transformer and Tensor Field network baselines in a reasonable time. This library is accurate to within machine precision against the scipy counterpart scipy.special.sph_harm and is significantly faster. E.g., for a ScanObjectNN model, we achieve $\sim 22\times$ speed up of the forward pass compared to a network built with SH from the lielearn library. A speed comparison isolating the computation of the spherical harmonics is shown in Fig. 5. Code is available at https://github.com/FabianFuchsML/se3-transformer-public. In the following, we outline our method to generate them.

The tesseral/real spherical harmonics are given as

where $P_{J}^{|m|}$ is the associated Legendre polynomial (ALP), $\theta\in[0,2\pi)$ is azimuth, and $\phi\in[0,\pi]$ is a polar angle. The term $P_{J}^{|m|}$ is by far the most expensive component to compute and can be computed recursively. To speed up the computation, we use dynamic programming storing intermediate results in a memoization.

We make use of the following recursion relations in the computation of the ALPs:

To understand how we recurse, we consider an example. Fig. 6 shows the space of $J$ and $m$. The black vertices represent a particular ALP, for instance, we have highlighted $P_{3}^{-1}(x)$. When $m<0$, we can use Eq. 49 to compute $P_{3}^{-1}(x)$ from $P_{3}^{1}(x)$. We can then use Eq. 50 to compute $P_{3}^{1}(x)$ from $P_{2}^{1}(x)$ and $P_{1}^{1}(x)$. $P_{2}^{1}(x)$ can also be computed from Eq. 50 and the boundary value $P_{1}^{1}(x)$ can be computed directly using Eq. 48. Crucially, all intermediate ALPs are stored for reuse. Say we wanted to compute $P_{4}^{-1}(x)$, then we could use Eq. 49 to find it from $P_{4}^{-1}(x)$, which can be recursed from the stored values $P_{3}^{1}(x)$ and $P_{2}^{1}(x)$, without needing to recurse down to the boundary.

Appendix D Experimental Details

A particularity of object classification from point clouds is the large number of points the algorithm needs to handle. We use up to 200 points out of the available 2024 points per sample and create neighbourhoods with up to 40 nearest neighbours. It is worth pointing out that especially in this setting, adding self-attention (i.e. when comparing the SE(3) Transformer to Tensor Field Networks) significantly increased the stability. As a result, whenever we swapped out the attention mechanism for a convolution to retrieve the Tensor Field network baseline, we had to decrease the model size to obtain stable training. However, we would like to stress that all the Tensor Field networks we trained were significantly bigger than in the original paper , mostly enabled by the faster computation of the spherical harmonics.

For the ablation study in Fig. 4, we trained networks with 4 hidden equivariant layers with 5 channels each, and up to representation degree 2. This results in a hidden feature size per point of

We used 200 points of the point cloud and neighbourhood size 40. For the Tensor Field network baseline, in order to achieve stable training, we used a smaller model with 3 instead of 5 channels, 100 input points and neighbourhood size 10, but with representation degrees up to 3.

We used 1 head per attention mechanism yielding one attention weight for each pair of points but across all channels and degrees (for an implementation of multi-head attention, see Section D.3). For the query embedding, we used the identity matrix. For the key embedding, we used a square equivariant matrix preserving the number of degrees and channels per degree.

For the quantitative comparison to the start-of-the-art in Table 2, we used 128 input points and neighbourhood size 10 for both the Tensor Field network baseline and the SE(3)-Transformer. We used farthest point sampling with a random starting point to retrieve the 128 points from the overall point cloud. We used degrees up to 3 and 5 channels per degree, which we again had to reduce to 3 channels for the Tensor Field network to obtain stable training. We used a norm based non-linearity for the Tensor Field network (as in ) and no extra non-linearity (beyond the softmax in the self-attention algorithm) for the SE(3) Transformer.

For all experiments, the final layer of the equivariant encoder maps to 64 channels of degree 0 representations. This yields a 64-dimensional SE(3) invariant representation per point. Next, we pool over the point dimension followed by an MLP with one hidden layer of dimension 64, a ReLU and a 15 dimensional output with a cross entropy loss. We trained for 60000 steps with batch size 10. We used the Adam optimizer with a start learning of 1e-2 and a reduction of the learning rate by 70% every 15000 steps. Training took up to 2 days on a system with 4 CPU cores, 30 GB of RAM, and an NVIDIA GeForce GTX 1080 Ti GPU.

The input to the Tensorfield network and the Se(3) Transformer are relative x-y-z positions of each point w.r.t. their neighbours. To guarantee equivariance, these inputs are provided as fields of degree 1. For the ‘+z‘ versions, however, we deliberately break the SE(3) equivariance by providing additional and relative z-position as two additional scalar fields (i.e. degree 0), as well as relative x-y positions as a degree 1 field (where the z-component is set to 0).

D.1.2 Number of Input Points

Limiting the input to 128/200 points in our experiments on ScanObjectNN was not primarily due to computational limitations: we conducted experiments with up to 2048 points, but without performance improvements. We suspect this is due to the global pooling. Examining cascaded pooling via attention is a future research direction. Interestingly, when limiting other methods to using 128 points, the SE(3)-Transformer outperforms the baselines (PointCNN: $80.3\pm 0.8\%$, PointGLR: $81.5\pm 1.0\%$, DGCNN: $82.2\pm 0.8\%$, ours: $85.0\pm 0.7\%$). It is worth noting that these methods were explicitly designed for the well-studied task of point classification whereas the SE(3)-Transformer was applied as is. Combining different elements from the current state-of-the-art methods with the geometrical inductive bias of the SE(3)-Transformer could potentially yield additional performance gains, especially with respect to leveraging inputs with more points. It is also worth noting that the SE(3)-Transformer was remarkably stable with respect to lowering the number of input points in an ablation study (16 points: $79.2\%$, 32 points: $81.4\%$, 64 points: $82.5\%$, 128 points: $85.0\%$, 256 points: $82.6\%$).

D.1.3 Sample Complexity

Equivariance is known to often lead to smaller sample complexity, meaning that less training data is needed (Fig. 10 in Worrall et al. , Fig. 4 in Winkels and Cohen , Fig. 4 in Weiler et al. ). We conducted experiments with different amounts of training samples $N_{\text{samples}}$ from the ScanObjectNN dataset. The results showed that for all $N_{\text{samples}}$, the SE(3)-Transformer outperformed the Set Transformer, a non-equivariant network based on attention. The performance delta was also slightly higher for the smallest $N_{\text{samples}}$ we tested ($3.1\%$ of the samples available in the training split of ScanObjectNN) than when using all the data indicating that the SE(3)-Transformer performs particularly well on small amounts of training data. However, performance differences can be due to multiple reasons. Especially for small datasets, such as ScanObjectNN, where the performance does not saturate with respect to amount of data available, it is difficult to draw conclusions about sample complexity of one model versus another. In summary, we found that our experimental results are in line with the claim that equivariance decreases sample complexity in this specific case but do not give definitive support.

D.1.4 Baselines

We originally replicated the implementation proposed in for their object classification experiment on ModelNet40 . However, most likely due to the relatively small number of objects in the ScanObjectNN dataset, we found that reducing the model size helped the performance significantly. The reported model had 128 units per hidden layer (instead of 256) and no dropout but the same number of layers and type of non-linearity as in .

We used the same architecture as in their object classification experiment on ModelNet40 with an ISAB (induced set attention block)-based encoder followed by PMA (pooling by multihead attention) and an MLP.

D.2 Relational Inference

Following Kipf et al. , we simulated trajectories for $5$ charged, interacting particles. Instead of a 2d simulation setup, we considered a 3d setup. Positive and negative charges were drawn as Bernoulli trials ($p=0.5$). We used the provided code base https://github.com/ethanfetaya/nri with the following modifications: While we randomly sampled initial positions inside a $^{3}$ box, we removed the bounding-boxes during the simulation. We generated $5\text{k}$ simulation samples for training and $1\text{k}$ for testing. Instead of phrasing it as a time-series task, we posed it as a regression task: The input data is positions and velocities at a random time step as well as the signs of the charges. The labels (which the algorithm is learning to predict) are the positions and velocities 500 simulation time steps into the future.

We trained each model for 100,000 steps with batch size 128 using an Adam optimizer . We used a fixed learning rate throughout training and conducted a separate hyper parameter search for each model to find a suitable learning rate.

We trained an SE(3)-Transformer with 4 equivariant layers, where the hidden layers had representation degrees $\{0,1,2,3\}$ and 3 channels per degree. The input is handled as two type-1 fields (for positions and velocities) and one type-0 field (for charges). The learning rate was set to 3e-3. Each layer included attentive self-interaction.

We used 1 head per attention mechanism yielding one attention weight for each pair of points but across all channels and degrees (for an implementation of multi-head attention, see Section D.3). For the query embedding, we used the identity matrix. For the key embedding, we used a square equivariant matrix preserving the number of degrees and channels per degree.

All our baselines fulfill permutation invariance (ordering of input points), but only the Tensor Field network and the linear baseline are SE(3) equivariant. For the Tensor Field Network baseline, we used the same hyper parameters as for the SE(3) Transformer but with a linear self-interaction and an additional norm-based nonlinearity in each layer as in Thomas et al. . For the DeepSet baseline, we used 3 fully connected layers, a pooling layer, and two more fully connected layers with 64 units each. All fully connected layers act pointwise. The pooling layer uses max pooling to aggregate information from all points, but concatenates this with a skip connection for each point. Each hidden layer was followed by a LeakyReLU. The learning rate was set to 1e-3. For the Set Transformer, we used 4 self-attention blocks with 64 hidden units and 4 heads each. For each point this was then followed by a fully connected layer (64 units), a LeakyReLU and another fully connected layer. The learning rate was set to 3e-4.

For the linear baseline, we simply propagated the particles linearly according to the simulation hyperparamaters. The linear baseline can be seen as removing the interactions between particles from the prediction. Any performance improvement beyond the linear baseline can therefore be interpreted as an indication for relational reasoning being performed.

D.3 QM9

The QM9 regression dataset is a publicly available chemical property prediction task consisting of 134k small drug-like organic molecules with up to 29 atoms per molecule. There are 5 atomic species (Hydrogen, Carbon, Oxygen, Nitrogen, and Flourine) in a molecular graph connected by chemical bonds of 4 types (single, double, triple, and aromatic bonds). ‘Positions’ of each atom, measured in ångtröms, are provided. We used the exact same train/validation/test splits as Anderson et al. of sizes 100k/18k/13k.

The architecture we used is shown in Table 4. It consists of 7 multihead attention layers interspersed with norm nonlinearities, followed by a TFN layer, max pooling, and two linear layers separated by a ReLU. For each attention layer, shown in Fig. 7, we embed the input to half the number of feature channels before applying multiheaded attention . Multiheaded attention is a variation of attention, where we partition the queries, keys, and values into $H$ attention heads. So if our embeddings have dimensionality $(4,16)$ (denoting 4 feature types with 16 channels each) and we use $H=8$ attention heads, then we partition the embeddings to shape $(4,2)$. We then combine each of the 8 sets of shape $(4,2)$ queries, keys, and values individually and then concatenate the results into a single vector of the original shape $(4,16)$. The keys and queries are edge embeddings, and thus the embedding matrices are of TFN type (c.f. Eq. 8). For TFN type layers, the radial functions are learnable maps. For these we used neural networks with architecture shown in Table 5.

where LN is layer norm applied across all features within a feature type. For the TFN baseline, we used the exact same architecture but we replaced each of the multiheaded attention blocks with a TFN layer with the same output shape.

The input to the network is a sparse molecular graph, with edges represented by the molecular bonds. The node embeedings are a 6 dimensional vector composed of a 5 dimensional one-hot embedding of the 5 atomic species and a 1 dimension integer node embedding for number of protons per atom. The edges embeddings are a 5 dimensional vector consisting of a 4 dimensional one-hot embedding of bond type and a positive scalar for the Euclidean distance between the two atoms at the ends of the bond. For each regression target, we normalised the values by mean and dividing by the standard deviation of the training set.

We trained for 50 epochs using Adam at initial learning rate 1e-3 and a single-cycle cosine rate-decay to learning rate 1e-4. The batch size was 32, but for the TFN baseline we used batch size 16, to fit the model in memory. We show results on the 6 regression tasks not requiring thermochemical energy subtraction in Table 3. As is common practice, we optimised architectures and hyperparameters on $\varepsilon_{\text{HOMO}}$ and retrained each network on the other tasks. Training took about 2.5 days on an NVIDIA GeForce GTX 1080 Ti GPU with 4 CPU cores and 15 GB of RAM.

D.4 General Remark

Across experiments on different datasets with the SE(3)-Transformer, we made the observation that the number of representation degrees have a significant but saturating impact on performance. A big improvement was observed when switching from degrees $\{0,1\}$ to $\{0,1,2\}$. Adding type-3 latent representations gave small improvements, further representation degrees did not change the performance of the model. However, higher representation degrees have a significant impact on memory usage and computation time. We therefore recommend representation degrees up to 2, when computation time and memory usage is a concern, and 3 otherwise.