LieTransformer: Equivariant self-attention for Lie Groups

Introduction

Group equivariant neural networks are useful architectures for problems with symmetries that can be described in terms of a group (in the mathematical sense). Convolutional neural networks (CNNs) are a special case that deal with translational symmetry, in that when the input to a convolutional layer is translated, the output is also translated. This property is known as translation equivariance, and offers a useful inductive bias for perception tasks which usually have translational symmetry. Constraining a linear layer to obey this symmetry, resulting in a covolutional layer, greatly reduces the number of parameters and computational cost. This has led to the success of CNNs in multiple domains such as computer vision (Krizhevsky et al., 2012) and audio (Graves & Jaitly, 2014). Following on from this success, there has been a growing literature on the study of group equivariant CNNs (G-CNNs) that generalise CNNs to deal with other types of symmetries beyond translations, such as rotations and reflections.

Most works on group equivariant NNs deal with CNNs i.e. linear maps with shared weights composed with pointwise non-linearities, building on the result that group equivariant linear maps (with mild assumptions) are necessarily convolutions (Kondor & Trivedi, 2018; Cohen et al., 2019; Bekkers, 2020). However there has been little work on non-linear group equivariant building blocks. In this paper we extend group equivariance to self-attention (Vaswani et al., 2017), a non-trivial non-linear map, that has become a prominent building block of deep learning models in various data modalities, such as natural-language processing (Vaswani et al., 2017; Brown et al., 2020), computer vision (Zhang et al., 2019; Parmar et al., 2019b), reinforcement learning (Parisotto et al., 2020), and audio generation (Huang et al., 2019).

We thus propose LieTransformer, a group invariant Transformer built from group equivariant LieSelfAttention layers. It uses a lifting based approach, that relaxes constraints on the attention module compared to approaches without lifting. Our method is applicable to Lie groups and their discrete subgroups (e.g. cyclic groups $C_{n}$ and dihedral groups $D_{n}$) acting on homogeneous spaces. Our work is very much in the spirit of Finzi et al. (2020), our main baseline, but for group equivariant self-attention instead of convolutions. Among works that deal with equivariant self-attention, we are the first to propose a methodology for general groups and domains (unspecified to 2D images (Romero et al., 2020; Romero & Cordonnier, 2021) or 3D point clouds (Fuchs et al., 2020)). We demonstrate the generality of our approach through strong performance on a wide variety of tasks, namely shape counting on point clouds, molecular property regression and modelling particle trajectories under Hamiltonian dynamics.

Background

This section lays down some of the necessary definitions and notations in group theory and representation theory in an informal and intuitive manner. For a more formal presentation of definitions, see Appendix B.

Loosely speaking, a group $G$ is a set of symmetries, with each group element $g$ corresponding to a symmetry transformation. These group elements ($g,g^{\prime}\in G$) can be composed ($gg^{\prime}$) or inverted ($g^{-1}$), just like transformations. An example of a discrete group is $C_{n}$, the set of rotational symmetries of a regular $n$-gon. The group consists of $n$ such rotations, including the identity. An example of a continuous (infinite) group is $SO(2)$, the set of all 2D rotations about the origin. $C_{n}$ is a subset of $SO(2)$, hence we call $C_{n}$ a subgroup of $SO(2)$. Note that $SO(2)=\{g_{\theta}:\theta\in[0,2\pi)\}$ can be parameterised by the angle of rotation $\theta$. Such groups that can be continuously parameterised by real values are called Lie groups.

Here the action $\rho$ is replaced by the action of the group on itself. If we wish to handle multiple channels of data, e.g. RGB images, we can stack these feature maps together, transforming in a similar manner.

Now let us define the notion of $G$-equivariance.

We say that a map $\Phi:V_{1}\rightarrow V_{2}$ is $G$-equivariant with respect to actions $\rho_{1},\rho_{2}$ of $G$ acting on $V_{1},V_{2}$ respectively if: $\Phi[\rho_{1}(g)f]=\rho_{2}(g)\Phi[f]$ for any $g\in G,f\in V_{1}$.

In the above example of rotating RGB images, we have $G=SO(2)$ and $\rho_{1}=\rho_{2}=\pi$. Hence the equivariance of $\Phi$ with respect to $SO(2)$ means that rotating an input image and then applying $\Phi$ yields the same result as first applying $\Phi$ to the original input image and then rotating the output, i.e. $\Phi$ commutes with the representation $\pi$.

The end goal for group equivariant neural networks is to design a neural network that obeys certain symmetries in the data. For example, we may want an image classifier to output the same classification when the input image is rotated. So in fact we want a $G$-invariant neural network, where the output is invariant to group actions on the input space. Note that $G$-invariance is a special case of $G$-equivariance, where $\rho_{2}$ is the trivial representation i.e. $\rho_{2}(g)$ is the identity map for any $g\in G$. Invariant maps are easy to design, by discarding information, e.g. pooling over spatial dimensions is invariant to rotations and translations. However, such maps are not expressive as they fail to extract high-level semantic features from the data. This is where equivariant neural networks become relevant; the standard recipe for constructing an expressive invariant neural network is to compose multiple equivariant layers with a final invariant layer. It is a standard result that such maps are invariant (e.g. Bloem-Reddy & Teh (2020)) and a proof is given in Appendix C for completeness.

2 Equivariant Maps on Homogeneous Input Spaces

Using this isomorphism, we can map each point in $\mathcal{X}$ to a set of group elements in $G$, i.e. mapping each data pair $(x_{i},\mathtt{f}_{i})$ to (possibly multiple) pairs $\{(g,\mathtt{f}_{i})|g\in s(x_{i})H\}$. This can be thought of as lifting the feature map $f_{\mathcal{X}}:x_{i}\mapsto\mathtt{f}_{i}$ defined on $\mathcal{X}$ to a feature map $\mathcal{L}[f_{\mathcal{X}}]:g\mapsto\mathtt{f}_{i}$ defined on $G$ (Kondor & Trivedi, 2018). Let $\mathcal{I}_{U}$ denote the space of such feature maps from $G$ to $\mathcal{F}$. Subsequently, we may define group equivariant maps as functions from $\mathcal{I}_{U}$ to itself, which turns out to be a simpler task than defining equivariant maps directly on $\mathcal{X}$.

The group equivariant convolution (Cohen & Welling, 2016; Cohen et al., 2018; Finzi et al., 2020; Romero et al., 2020) is an example of such a group equivariant map that has been studied extensively. Specifically, the group equivariant convolution $\Psi:\mathcal{I}_{U}\rightarrow\mathcal{I}_{U}$ is defined as:

LieTransformer

LieTransformer is composed of a lifting layer followed by residual blocks of LieSelfAttention layers, LayerNorm and pointwise MLPs, all of which are equivariant with respect to the regular representation, followed by a final invariant G-pooling layer (c.f. Appendix H for more details on these layers). We summarise the architecture in Figure 1 and describe its key components below.

Recall from Section 2.2 that the lifting $\mathcal{L}$ maps $f_{\mathcal{X}}$ (supported on $\bigcup_{i=1}^{n}\{x_{i}\}\subset\mathcal{X}$) to $\mathcal{L}[f_{\mathcal{X}}]$ (supported on $\bigcup_{i=1}^{n}s(x_{i})H\subset G$) such that:

This can be thought of as extending the domain of $f_{\mathcal{X}}$ from $\mathcal{X}$ to $G$ while preserving the feature values $\mathtt{f_{i}}$, mapping $(x_{i},\mathtt{f}_{i})\mapsto(g,\mathtt{f}_{i})$ for $g\in s(x_{i})H$ (c.f. Figure 2). Subsequently we may design $G$-equivariant maps on the space of functions on $G$, which is a simpler task than desgining $G$-equivariant maps directly on $\mathcal{X}$ (e.g. Cohen et al. (2018)).

As in Equations 2 and 3, we define the representation $\pi$ on $f_{\mathcal{X}}$ and $\mathcal{L}[f_{\mathcal{X}}]$ as:

where $u\in G$. Note that the actions correspond to mappings $(x_{i},\mathtt{f}_{i})\mapsto(ux_{i},\mathtt{f}_{i})$ and $(g,\mathtt{f}_{i})\mapsto(ug,\mathtt{f}_{i})$ respectively.

We need to ensure that lifting preserves equivariance, which is why we need the space to be homogeneous with respect to the action of $G$ on $\mathcal{X}$.

The lifting layer $\mathcal{L}$ is equivariant with respect to the representation $\pi$.

Intuition for proof. When $x_{i}$ is shifted by $u\in G$, the lifted coset $s(x_{i})H$ is also shifted by $u$, i.e. $x_{i}\mapsto ux_{i}\Rightarrow s(x_{i})H\mapsto us(x_{i})H$. See Appendix C for full proof.

2 LieSelfAttention

Let $f\triangleq\mathcal{L}[f_{\mathcal{X}}]$, hence $f$ is defined on the set $G_{f}=\cup_{i=1}^{n}s(x_{i})H$. We define the LieSelfAttention layer in Algorithm 1, where self-attention (see Appendix D for the original formulation) is defined across the elements of $G_{f}$. There are various choices for functions content-based attention $k_{c}$ , location-based attention $k_{l}$ , $F$ that determines how to combine the two to form unnormalised weights, and the choice of normalisation of weights. See Appendix E for a non-exhaustive list of choices of the above, and also for the details of the multi-head generalisation of LieSelfAttention.

LieSelfAttention is equivariant with respect to the regular representation $\pi$.

Intuition for proof. LieSelfAttention can be thought of as a map $\Phi:(g,f(g))_{g\in G_{f}}\mapsto(g,f_{out}(g))_{g\in G_{f}}$, and equivariance holds if $\forall u\in G$, $\Phi$ maps $(ug,f(g))_{g\in G_{f}}$ to $(ug,f_{out}(g))_{g\in G_{f}}$. Now note that $\Phi$ is a function of $g\in G_{f}$ only via $g^{-1}g^{\prime}$ for $g^{\prime}\in G_{f}$, and $g^{-1}g^{\prime}$ is invariant to the group action $g\mapsto ug,g^{\prime}\mapsto ug^{\prime}$. This is enough to show that $\Phi$ satisfies the above condition for equivariance. See Appendix C for full proof.

Replace $G_{f}\triangleq\cup_{i=1}^{n}s(x_{i})H$ with a finite subset $\hat{G}_{f}\triangleq\cup_{i=1}^{n}s(x_{i})\hat{H}$ where $\hat{H}$ is a finite subset of $H$ sampled uniformly. We refer to $|\hat{H}|$ as the number of lift samples.

(Optional, for computational efficiency) Further replace $\hat{G}_{f}$ with uniform samples from the neighbourhood $\mathtt{nbhd}_{\eta}(g)\triangleq\{g^{\prime}\in\hat{G}_{f}:d(g,g^{\prime})\leq\eta\}$ for some threshold $\eta$ where distance is measured by the log map $d(g,g^{\prime})=||\nu[\log(g^{-1}g^{\prime})]||$ (c.f. Appendix F).

See Figure 2 for a visualisation. Due to MC estimation we now have equivariance in expectation as Finzi et al. (2020). For sampling within the neighbourhood, we can show that the resulting LieSelfAttention is still equivariant in expectation given that the distance is a function of $g^{-1}g^{\prime}$ (c.f. Appendix C).

Related Work

Equivariant maps with/without lifting Equivariant neural networks can be broadly categorised by whether the input spatial data is lifted onto the space of functions on group $G$ or not. Without lifting, the equivariant map is defined between the space of functions/features on the homogeneous input space $X$, with equivariance imposing a constraint on the parameterisation of the convolutional kernel or attention module (Cohen & Welling, 2017; Worrall et al., 2017; Thomas et al., 2018; Kondor et al., 2018; Weiler et al., 2018b, a; Weiler & Cesa, 2019; Esteves et al., 2020; Fuchs et al., 2020). In the case of convolutions, the kernel is expressed using a basis of equivariant functions such as circular or spherical harmonics. However with lifting, the equivariant map is defined between the space of functions/features on $G$, and aforementioned constraints on the convolutional kernel or attention module are relaxed at the cost of an increased dimensionality of the input to the neural network (Cohen & Welling, 2016; Cohen et al., 2018; Esteves et al., 2018; Finzi et al., 2020; Bekkers, 2020; Romero & Hoogendoorn, 2020; Romero et al., 2020; Hoogeboom et al., 2018). Our method also uses lifting to define equivariant self-attention.

Equivariant self-attention Most of the above works use equivariant convolutions as the core building block of their equivariant module, drawing from the result that bounded linear operators are group equivariant if and only if they are convolutions (Kondor & Trivedi, 2018; Cohen et al., 2019; Bekkers, 2020). Such convolutions are used with pointwise non-linearities (applied independently to the features at each spatial location/group element) to form expressive equivariant maps. Exceptions to this are Romero et al. (2020) and Fuchs et al. (2020) that explore equivariant attentive convolutions, reweighing convolutional kernels with attention weights. This gives non-linear equivariant maps with non-linear interactions across spatial locations/group elements. Instead, our work removes convolutions and investigates the use of equivariant self-attention only, inspired by works that use stand-alone self-attention on images to achieve competitive performance to convolutions (Parmar et al., 2019a; Dosovitskiy et al., 2020). Furthermore, Romero et al. (2020) focus on image applications (hence scalability) and discrete groups (p4, p4m), and Fuchs et al. (2020) focus on 3D point cloud applications and the $SE(3)$ group with irreducible representations acting on functions on $\mathcal{X}$. Instead we use regular representations actingon functions on $G$, and give a general method for Lie groups acting on homogeneous spaces, with a wide range of applications from dealing with point cloud data to modelling Hamiltonian dynamics of particles. This is very much in the spirit of Finzi et al. (2020), except for self-attention instead of convolutions. In concurrent work, Romero & Cordonnier (2021) describe group equivariant self-attention also using lifting and regular representations. Their analogue of location-based attention are group invariant positional encodings. The main difference between the two works is that Romero & Cordonnier (2021) specify methodology for discrete groups applied to image classification only and it is not clear how to extend their approach to Lie groups. In contrast, our method provides a general formula for (unimodular) Lie groups and their discrete subgroups for the aforementioned applications.

Experiments

We consider three different tasks that have certain symmetries, highlighting the benefits of the LieTransformer: (1) Counting shapes in 2D point cloud of constellations (2) Molecular property regression and (3) Modelling particle trajectories under Hamiltonian dynamics.The code for our experiments is available at: https://github.com/oxcsml/lie-transformer

We first consider the toy, synthetic task of counting shapes in a 2D point cloud $\{x_{1},x_{2},...,x_{K}\}$ of constellations (Kosiorek et al., 2019), mainly to check that LieTransformer has the correct invariance properties. We use $\mathtt{f}_{i}=1$ for all points. Each example consists of points in the plane that form the vertices of a pattern. There are four types of patterns: triangles, squares, pentagons and the ‘L’ shape, with varying sizes, orientation, and number of instances per pattern (see Figure 3 (right)). The task is to classify the number of instances of each pattern, hence is invariant to 2D roto-translations $SE(2)$.

We first create a fixed training set $D_{\text{train}}$ and test set $D_{\text{test}}$ of size 10,000 and 1,000 respectively. We then create augmented test sets $D^{T2}_{\text{test}}$ and $D^{SE2}_{\text{test}}$ that are copies of $D_{\text{test}}$ with arbitrary transformations in $T(2)$ and $SE(2)$ respectively. In Table 1, we evaluate the test accuracy of LieTransformer at convergence with and without data augmentation during training time – $D^{T2}_{\text{train}}$ and $D^{SE2}_{\text{train}}$ indicate random $T(2)$ and $SE(2)$ augmentations respectively to each batch of $D_{\text{train}}$ at every training iteration. We evaluate the test performance of LieTransformer-T2 and LieTransformer-SE2 that are invariant to $T(2)$ and $SE(2)$ respectively, against the baseline SetTransformer (Lee et al., 2019), a Transformer-based model that is permutation invariant, but not invariant to rotations nor translations. We use a similar number of parameters for each model. See Appendix I.1 for further details on the setup.

Note that the test accuracy of LieTransformer-T2 and LieTransformer-SE2 remains unchanged when the train/test set is augmented with $T(2)$. For LieTransformer-SE2, this is not quite true for $SE(2)$ augmentations because the model is only $SE(2)$ equivariant in expectation and not exactly equivariant given a finite number of lifts samples ($|\hat{H}|$). However the changes in accuracy for $SE(2)$ augmentation are much smaller compared to LieTransformer-T2. The test accuracy of SetTransformer, the non-invariant baseline, is always lower than LieTransformer. Note that LieTransformer-T2 does slightly better than LieTransformer-SE2 on $D_{\text{test}}$ and $D^{T2}_{\text{test}}$. We suspect that the variance in the sampling of the lifting layer for LieTransformer-SE2 is making optimisation more difficult, and will continue to explore these results.

In Figure 3 (left), we report the equivariance error of LieTransformer-SE2 when increasing the number of lift samples ($|\hat{H}|$) used in the Monte Carlo approximation of LieSelfAttention. As expected the invariance error decreases monotonically with the number of lift samples, and already with 3 lift samples, the error is small ($\approx 10^{-6}$).

2 QM9: Molecular Property Regression

We apply the LieTransformer to the QM9 molecule property prediction task (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014). This dataset consists of 133,885 small inorganic molecules described by the location and charge of each atom in the molecule, along with the bonding structure of the molecule. The dataset includes 19 properties of each molecule, such as various rotational constants, energies and enthalpies, and 12 of these are used as regression tasks. We expect these molecular properties to be invariant to 3D roto-translations $SE(3)$. We follow the customary practice of performing hyperparameter search on the $\epsilon_{\text{HOMO}}$ task and use the same hyperparameters for training on the other 11 tasks. Further details of the exact experimental setup can be found in Appendix I.2.

We trained four variants of both LieTransformer and LieConv, namely the $T(3)$ and $SE(3)$ invariant models with and without $SO(3)$ (rotation) data augmentation. We set $x_{i}$ to be the atomic position and $\mathtt{f}_{i}$ to be the charge. Table 2 shows the test error of all models and baselines on the 12 tasks. The table is divided into 3 sections. Upper: non-invariant models specifically designed for the QM9 task. Middle: invariant models specifically designed for the QM9 task. Lower: invariant models that are general-purpose. We show very competitive results, and perform best of general-purpose models on 8/12 tasks. In particular when comparing against LieConv, we see better performance on the majority of tasks, suggesting that the attention framework is better suited to these tasks than convolutions.

As expected for both LieTransformer and LieConv, the $SE(3)$ models tend to outperform the $T(3)$ models without $SO(3)$ data augmentation (on 10/12 tasks and 7/12 tasks respectively), showing that being invariant to rotations improves generalisation. Moreover the $SE(3)$ models perform similarly with and without augmentation, whereas the $T(3)$ models greatly benefit from augmentation, showing evidence that the $SE(3)$ models are indeed invariant to rotations. However the $T(3)$ models with augmentation outperform the $SE(3)$ counterparts on most tasks for both LieTransformer-SE3 and LieConv-SE3. As for the experiments in Section 5.1, we suspect that the variance in the sampling of the lifting layer of $SE(3)$ models, along with the $SE(3)$ log-map (Appendix F) in the location attention is making optimisation more difficult, and plan to continue investigating the source of this discrepancy in performance. Note however that LieTransformer-SE3 and LieConv-SE3 tend to outperform the irreducible representation (irrep) based $SE(3)$-Transformer and TFN. This can be seen as further evidence that regular representation approaches tend to outperform irrep approaches, in line with the empirical observations of Weiler & Cesa (2019).

3 Modelling Particle Trajectories with Hamiltonian Dynamics

We also apply the LieTransformer to a physics simulation task in the context of Hamiltonian dynamics, a formalism for describing the evolution of a physical system using a single scalar function $H(q,p)$, called the Hamiltonian.

However we know a-priori that such physical systems have symmetries, namely conserved quantities such as linear and angular momentum. A notable result is Noether’s theorem (Noether, 1918), which states that the system has a conserved quantity if and only if the Hamiltonian is group-invariant. For example, translation invariance of the Hamiltonian implies conservation of momentum and rotation invariance implies conservation of angular momentum. Hence in our experiments, we parameterise the Hamiltonian $H_{\theta}$ by a LieTransformer and endow it with the symmetries corresponding to the conservation laws of the physical system we are modelling. We test our model on the spring dynamics task proposed in Sanchez-Gonzalez et al. (2019) – we consider a system of $6$ particles with randomly sampled massses in 2D, where each particle connected to all others by springs. This system conserves both linear and angular momentum, so the ground truth Hamiltonian will be both translationally and rotationally invariant, that is, $SE(2)$-invariant. We simulate this system for 500 timesteps from random initial conditions and use random subsets of length 5 from these roll-outs to train the model (see Appendix I.3 for full experimental details).

We compare our method to different parameterisations of $H_{\theta}$, namely Fully-connected network (Chen et al., 2018), Graph Network (Sanchez-Gonzalez et al., 2019) and LieConv. Only LieTransformer and LieConv incorporate invariance. In Figures 4, 5, and 6, we use LieTransformer-T2 and LieConv-T2 since Finzi et al. (2020) report that there are numerical instabilities for LieConv-SE2 on this task, due to which LieConv-T2 is their default model and performs the best. However in Figure 7, we also consider $SE(2)$-invariant versions of both models with modifications to the lifting procedure, which fixed the instabilities as outlined in Appendix F.

Figure 4 compares the performance of all methods as a function of the number of training examples. LieTransformer is highly data-efficient: the inductive bias from the symmetries of the Hamiltonian allow us to accurately learn the dynamics even from a small training set. Our method consistently outperforms non-invariant methods (fully-connected and graph networks), typically by 1-3 orders of magnitude. Furthermore, our method outperforms LieConv for most data sizes except the largest sizes where the errors are similar, suggesting that the attention framework more suited for this task.

Figure 5 shows the test error as function of the roll-out time step for a training data size of 10,000 (corresponding plots for other training data sizes are included in Appendix J.1). Here we show that the LieTransformer shows better generalisation than LieConv across all roll-out lengths, the error being low ($<10^{-3}$) for 100 step-roll-outs even though we only train on 5-step roll-outs. We also include example trajectories of our model in Figure 6 (more examples can be found in the appendix, including ones where LieConv performs better than LieTransformer) illustrating the accuracy of our model on this task.

Lastly, Figure 7 compares LieConv and LieTransformer for different model sizes (number of parameters) and equivariance groups. We first note LieTransformer outperforms LieConv given a fixed model size and group. For $T(2)$-invariant models, our method benefits from a larger model, whereas LieConv deteriorates (LieConv-T(2) (895K) is their default architecture on this task). However, for both methods, the $SE(2)$-invariant models perform at par with or better than their $T(2)$-invariant counterparts despite having smaller model sizes. In particular, LieTransformer-SE(2) (139K) outperforms all other models in this comparison despite having the smallest number of parameters, which highlights the advantage of incorporating the correct task symmetries into the architecture and the attention framework. Overall, we have shown that our model is suitable for use in a neural ODE setting that requires equivariant drift functions.

Limitations and Future Work

From the algorithmic perspective, LieTransformer shares the weakness of LieConv in being memory-expensive ($O(|\hat{G}_{f}||\mathtt{nbhd}_{\eta}|)$ memory cost (Appendix G) due to: 1. The lifting procedure that increases the number of inputs by $|\hat{H}|$, and 2. Quadratic complexity in the number of inputs from having to compute the kernel value at each pair of inputs. Although the first is a weakness shared by all lifting-based equivariant neural networks, the second can be addressed by incorporating works that study efficient variants of self-attention (Wang et al., 2020; Kitaev et al., 2020; Zaheer et al., 2020; Katharopoulos et al., 2020). An alternative is to incorporate information about pairs of inputs (such as bonding information for the QM9 task) as masking in self-attention (c.f. Appendix I.2).

From the methodological perspective, a key weakness of the LieTransformer that is also shared with LieConv is its approximate equivariance due to MC estimation of the integral in LieSelfAttention for the case where $H$ is infinite. The aforementioned directions for memory-efficiency can help to reduce the approximation error by allowing to use more lift samples ($|\hat{H}|$). Other directions include incorporating the notion of steerability (Cohen & Welling, 2017) to deal with vector fields in an equivariant manner (given inputs $(x_{i},\mathtt{f}_{i})$, the group acts non-trivially on $\mathtt{f}_{i}$ as well as $x_{i}$), and extending to non-homogeneous input spaces as outlined in Finzi et al. (2020).

The authors would like to thank Adam R.Kosiorek for setting up the initial codebase at the beginning of the project, and David W. Romero & Jean-Baptiste Cordonnier for useful discussions. Michael is supported by the EPSRC Centre for Doctoral Training in Modern Statistics and Statistical Machine Learning (EP/S023151/1). Charline acknowledges funding from the EPSRC grant agreement no. EP/N509711/1. Sheheryar wishes to acknowledge support from Aker Scholarship. Emilien acknowledges support of his PhD funding from Google DeepMind. Yee Whye Teh’s research leading to these results has received funding from the European Research Council under the European Union’s Seventh Framework Programme (FP7/2007-2013) ERC grant agreement no. 617071.

We would also like to thank the Python community (Van Rossum & Drake Jr, 1995; Oliphant, 2007) for developing the tools that enabled this work, including Pytorch (Paszke et al., 2017), NumPy (Oliphant, 2006; Walt et al., 2011; Harris et al., 2020), SciPy (Jones et al., 2001), and Matplotlib (Hunter, 2007).

References

Appendix A Contributions

Charline and Yee Whye conceived the project and Yee Whye initially came up with an equivariant form of self-attention.

Through discussions between Michael, Charline, Hyunjik and Yee Whye, this was modified to the current LieSelfAttention layer, and Michael derived the equivariance of the LieSelfAttention layer.

Michael, Sheheryar and Hyunjik simplified the proof of equivariance and further developed the methodology for the LieTransformer in its current state, and created links between LieTransformer and other related work.

Michael wrote the initial framework of the LieTransformer codebase. Charline and Sheheryar wrote the code for the shape counting experiments, Michael wrote the code for the QM9 experiments, Sheheryar wrote the code for the Hamiltonian dynamics experiments, after helpful discussions with Emilien.

Charline carried out the experiments for Table 1, Michael carried out most of the experiments for Table 2 with some help from Hyunjik, Sheheryar carried out the experiments for Figure 3b and all the Hamiltonian dynamics experiments.

Hyunjik wrote all sections of the paper except the experiment sections: the shape counting section was written by Charline, the QM9 section by Michael and the Hamiltonian dynamics section was written by Emilien and Sheheryar.

Appendix B Formal definitions for Groups and Representation Theory

A group ${G}$ is a set endowed with a single operator $\cdot:{G}\times{G}\mapsto{G}$ such that

Associativity: $\forall g,g^{\prime},g^{\prime\prime}\in{G},(g\cdot g^{\prime})\cdot g^{\prime\prime}=g\cdot(g^{\prime}\cdot g^{\prime\prime})$

Identity: $\exists e\in{G},\,\forall g\in{G}\,\,g\cdot e=e\cdot g=g$

Invertibility: $\forall g\in{G},\exists g^{-1}\in{G},g\cdot g^{-1}=g^{-1}\cdot g=e$

A Lie group is a finite-dimensional real smooth manifold, in which group multiplication and inversion are both smooth maps.

Let $S$ be a set, and let $\operatorname{Sym}(S)$ denote the set of invertible functions from $S$ to itself. We say that a group ${G}$ acts on $S$ via an action $\rho:G\rightarrow\operatorname{Sym}(S)$ when $\rho$ is a group homomorphism: $\rho(g_{1}g_{2})(s)=(\rho(g_{1})\circ\rho(g_{2}))(s)~{}\forall s\in S$.

If $S$ is a vector space $V$ and this action is, in addition, a linear function, i.e. $\rho:G\rightarrow GL(V)$, where $GL(V)$ is the set of linear invertible functions from $V$ to itself, then we say that $\rho$ is a representation of ${G}$.

Appendix C Proofs

The function composition $f\circ f_{K}\circ...\circ f_{1}$ of several equivariant functions $f_{k}$, $k\in\{1,2,...,K\}$ followed by an invariant function $f$, is an invariant function.

Consider group representations $\pi_{1},\ldots,\pi_{K}$ that act on $f_{1},\ldots f_{K}$ respectively, and representation $\pi_{0}$ that acts on the input space of $f_{1}$. If each $f_{k}$ is equivariant with respect to $\pi_{k},\pi_{k-1}$ such that $f_{k}\circ\pi_{k-1}=\pi_{k}\circ f_{k}$, and $f$ is invariant such that $f\circ\pi_{k}=f$, then we have

hence $f\circ f_{K}\circ...\circ f_{1}$ is invariant. ∎

The group equivariant convolution $\Psi:\mathcal{I}_{U}\rightarrow\mathcal{I}_{U}$ defined as: $[\Psi f](g)\triangleq\int_{G}\psi(g^{\prime-1}g)f(g^{\prime})dg^{\prime}$ is equivariant with resepct to the regular representation $\pi$ of $G$ acting on $\mathcal{I}_{U}$ as $[\pi(u)f](g)\triangleq f(u^{-1}g)$.

The second equality holds by invariance of the left Haar measure. ∎

Note $\mathcal{L}[\pi(u)f_{\mathcal{X}}](g)=\mathtt{f}_{i}$ for $g\in s(ux_{i})H$ and $[\pi(u)\mathcal{L}[f_{\mathcal{X}}]](g)=\mathcal{L}[f_{\mathcal{X}}](u^{-1}g)=\mathtt{f}_{i}$ for $g\in us(x_{i})H$. Hence $\mathcal{L}[\pi(u)f_{\mathcal{X}}]=\pi(u)\mathcal{L}[f_{\mathcal{X}}]$ because the two cosets are equal: $s(ux_{i})H=us(x_{i})H~{}\forall u\in G$. Note that this holds because:

If $g\in s(x_{i})H=\{g\in G|gx_{0}=x_{i}\}$, then $g$ maps $x_{0}$ to $x_{i}$, hence $ug$ maps $x_{0}$ to $ux_{i}$. So if $g\in s(x_{i})H$ then $ug\in s(ux_{i})H=\{g\in G|gx_{0}=ux_{i}\}$, the set of all $g$ that map $x_{0}$ to $ux_{i}$. In summary, $us(x_{i})H\subset s(ux_{i})H$

Conversely, if $g\in s(ux_{i})H$, then we know that $u^{-1}g$ maps $x_{0}$ to $x_{i}$, so $u^{-1}g\in s(x_{i})H$, hence $g\in us(x_{i})H$. In summary, $s(ux_{i})H\subset us(x_{i})H$

We have shown $us(x_{i})H\subset s(ux_{i})H$ and $s(ux_{i})H\subset us(x_{i})H$, thus $s(ux_{i})H=us(x_{i})H$.

$f$ is defined on the set $G_{f}=\cup_{i=1}^{n}s(x_{i})H$ (i.e. union of cosets corresponding to each $x_{i}$). Note $G_{\pi(u)f}=uG_{f}$, and $G_{f}$ does not depend on the choice of section $s$.

Note that for all provided choices of $k_{c}$ and $k_{l}$, we have:

Hence for all choices of $F$, we have that

We thus prove equivariance for the below choice of LieSelfAttention $\Phi:\mathcal{I}_{U}\rightarrow\mathcal{I}_{U}$ that uses softmax normalisation, but a similar proof holds for constant normalisation. Let $A_{f}(g,g^{\prime})\triangleq\exp(\alpha_{f}(g,g^{\prime}))$, hence Equation (10) also holds for $A_{f}$.

Then we can show that $\Phi$ is equivariant with respect to the representation $\pi$ as follows:

Appendix D Introduction to Self-Attention

Note $XW^{Q}(XW^{K})^{\top}$ is the Gram matrix for the dot-product kernel, and softmax normalisation is a particular choice of normalisation. Hence MSA can be generalised to other choices of kernels and normalisation that are equally valid (Wang et al., 2018; Tsai et al., 2019).

Appendix E LieSelfAttention: Details

We explore the following non-exhaustive list of choices for content-based attention, location-based attention, combining content and location attention and normalisation of weights:

Content-based attention $k_{c}(f(g),f(g^{\prime}))$:

Location-based attention $k_{l}(g^{-1}g^{\prime})$ for Lie groups $G$:

MLP: $\mathtt{MLP}(\nu[\log(g^{-1}g^{\prime})])$

Combining content and location attention $\alpha_{f}(g,g^{\prime})$:

Additive: $k_{c}(f(g),f(g^{\prime}))+k_{l}(g^{-1}g^{\prime})$

MLP: $\mathtt{MLP}[\mathtt{Concat}[k_{c}(f(g),f(g^{\prime})),k_{l}(g^{-1}g^{\prime})]]$

Multiplicative: $k_{c}(f(g),f(g^{\prime}))\cdot k_{l}(g^{-1}g^{\prime})$

Note that the MLP case is a strict generalisation of the additive combination, and for this option $k_{c}$ and $k_{l}$ need not be scalars.

Normalisation of weights $\{w_{f}(g,g^{\prime})\}_{g^{\prime}\in G_{f}}$:

Softmax: $\operatorname{softmax}\left(\{\alpha_{f}(g,g^{\prime})\}_{g^{\prime}\in G_{f}}\right)$

Constant: $\{\frac{1}{|G_{f}|}\alpha_{f}(g,g^{\prime})\}_{g^{\prime}\in G_{f}}$

We also outline how to extend the single-head LieSelfAttention described in Algorithm 1 extends to Multihead equivariant self-attention. Let $M$ be the number of heads, assuming it divides $d_{v}$, with $m$ indexing the head. Then the output of each head is:

Appendix F Lie Algebras and Log maps

In this section we briefly introduce Lie algebras and log maps, mainly summarising relevant sections of Hall (2015). See the reference for a formal and thorough treatment of Lie groups and Lie algebras.

where $t^{\prime}=V^{-1}t$, $V=\begin{bmatrix}a&-b\\ b&a\end{bmatrix}$, $a\triangleq\frac{\sin\theta}{\theta}$, $b\triangleq\frac{1-\cos\theta}{\theta}$

where $\cos\theta=\frac{\operatorname{Tr}(R)-1}{2}$. Note that the Taylor expansion of $\theta/\sin\theta$ should be used when $\theta$ is small.

where $t^{\prime}=V^{-1}t,r^{\prime}=\nu[\log(R)]$, $V=I+\frac{1-\cos\theta}{\theta^{2}}(R-R^{\top})+\frac{\theta-\sin\theta}{\theta^{3}}(R-R^{\top})^{2}$.

Appendix G Memory and Time Complexity Comparison with LieConv

Outputs: $\{g,\frac{1}{|\mathtt{nbhd}(g)|}\sum_{g^{\prime}\in\mathtt{nbhd}(g)}k_{L}(g^{-1}g^{\prime})f(g^{\prime})\}_{g\in G_{f}}$ where

$\mathtt{nbhd}(g)=\{g^{\prime}\in G_{f}:\nu[\log(g)]<r\}$. Let us assume that $|\mathtt{nbhd}(g)|\approx n~{}\forall g$.

There are (at least) two ways of computing LieConv: 1. Naive and 2. PointConv Trick.

Time: Compute $k_{L}(g^{-1}g^{\prime})f(g^{\prime})\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g)$. This requires $O(|G_{f}|nd_{out}d_{v})$ flops.

One-line summary: instead of applying a shared linear map then summing across $\mathtt{nbhd}$, first sum across $\mathtt{nbhd}$ then apply the linear map.

Details: $k_{L}(g^{-1}g^{\prime})=\mathtt{MLP}_{\theta}(\nu[\log(g^{-1}g^{\prime})])=\mathtt{reshape}(HM(g^{-1}g^{\prime}),[d_{out},d_{v}])$ where

The trick assumes $d_{mid}\ll d_{out}d_{v}$, and reorders the computation as:

Memory: Store $M(g^{-1}g^{\prime})~{}\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g)$, and store $H$. This requires $O(|G_{f}|nd_{mid}+d_{out}d_{v}d_{mid})$ memory.

Time: Compute $\sum_{g^{\prime}\in\mathtt{nbhd}(g)}M(g^{-1}g^{\prime})\otimes f(g^{\prime})$ via matrix multiplication: $\begin{bmatrix}|&&|\\ M(g^{-1}g^{\prime}_{1})&\ldots&M(g^{-1}g^{\prime}_{n})\\ |&&|\end{bmatrix}\begin{bmatrix}\text{---}&f(g^{\prime}_{1})&\text{---}\\ &\vdots&\\ \text{---}&f(g^{\prime}_{n})&\text{---}\end{bmatrix}$. This requires $O(d_{v}nd_{mid})$ flops.

Then multiply by $H$, requiring $O(d_{v}d_{out}d_{mid})$ flops.

This is done for each $g\in G_{f}$, so the total number of flops is $O(|G_{f}|d_{v}d_{mid}(n+d_{out}))$.

G.2 Equivariant Self-Attention

Outputs: $\{g,f(g)+\sum_{g^{\prime}\in\mathtt{nbhd}(g)}w_{f}(g,g^{\prime})W^{V}f(g^{\prime})\}_{g\in G_{f}}$ where

$\mathtt{nbhd}(g)=\{g^{\prime}\in G_{f}:\nu[\log(g)]<r\}$. Let us assume that $|\mathtt{nbhd}(g)|\approx n~{}\forall g$.

$\{w_{f}(g,g^{\prime})\}_{g^{\prime}\in G_{f}}=\operatorname{softmax}\left(\{\alpha_{f}(g,g^{\prime})\}_{g^{\prime}\in G_{f}}\right)$

$\alpha_{f}(g,g^{\prime})=k_{f}(f(g),f(g^{\prime}))+k_{x}(g^{-1}g^{\prime})$

Memory: Store $\alpha_{f}(g,g^{\prime})$ and $W^{V}f(g^{\prime})~{}\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g)$. This requires $O(|G_{f}|nd_{v})$ memory.

Time: Compute $k_{f}(f(g),f(g^{\prime}))$ and $w_{f}(g,g^{\prime})~{}\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g)$. This requires $O(|G_{f}|nd_{v}^{2})$ flops.

With multihead self-attention ($M$ heads), the output is:

Memory: Store $\alpha_{f}^{m}(g,g^{\prime})$ and $W^{V,m}f(g^{\prime})~{}\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g),m\in\{1,\ldots,M\}$. This requires $O(M|G_{f}|n+M|G_{f}|nd_{v}/M)=O(|G_{f}|n(M+d_{v}))$ memory.

Time: Compute $k_{f}^{m}(f(g),f(g^{\prime}))$ and $w_{f}^{m}(g,g^{\prime})~{}\forall g\in G_{f},g^{\prime}\in\mathtt{nbhd}(g),m\in\{1,\ldots,M\}$. This requires $O(M|G_{f}|nd_{v}d_{v}/M)=O(|G_{f}|nd_{v}^{2})$ flops.

Appendix H Other Equivariant/Invariant building blocks

G-Pooling is simply averaging over the features across the group: Inputs: $\{f(g)\}_{g\in G_{f}}$ Output: $\bar{f}(g)\triangleq\frac{1}{|G_{f}|}\sum_{g\in G_{f}}f(g)$ Note that G-pooling is invariant with respect to the regular representation.

Pointwise MLPs are MLPs applied independently to each $f(g)$ for $g\in G_{f}$. It is easy to show that any such pointwise operations are equivariant with respect to the regular representation.

LayerNorm (Ba et al., 2016) is defined as follows:

Outputs: $\{g,\beta\odot\frac{f(g)-m(g)}{\sqrt{v(g)+\epsilon}}+\gamma\}_{g\in G_{f}}$ where

BatchNorm We also describe BatchNorm (Ioffe & Szegedy, 2015) that is used in (Finzi et al., 2020) for completeness:

Inputs: $\smash{\{g,f^{b}(g)\}_{g\in G_{f},b\in\mathcal{B}}}$ where

$G_{f}$ defined as in Section 3.1, $\mathcal{B}$ is the batch of examples.

Outputs: $\smash{\{g,\beta\odot\frac{f^{b}(g)-\mathbf{m}(g)}{\sqrt{\mathbf{v}(g)+\epsilon}}+\gamma\}_{g\in G_{f},b\in\mathcal{B}}}$ where

$\mathtt{nbhd}(g)=\{g^{\prime}\in G_{f}:\nu[\log(g)]<r\}$.

A moving average of $\mathbf{m}(g)$ and $\mathbf{v}(g)$ are tracked during training time for use at test time. It is easy to check that both BatchNorm and LayerNorm are equivariant wrt the action of the regular representation $\pi$ on $f$ (for BatchNorm, note $g^{\prime}\in\mathtt{nbhd}(g)$ iff $u^{-1}g^{\prime}\in\mathtt{nbhd}(u^{-1}g)$).

Appendix I Experimental details

Each training / test example consists of up to two instances of each of the following shapes: triangles, squares, pentagons and the ”L” shape. The $x_{i}$ are the 2D coordinates of each point and $\mathtt{f}_{i}=1$ for all points.

We performed an architecture search on the LieTransformer first and then set the architecture of the SetTransformer such that the models have a similar number of parameters (1075k for the SetTransformer and 1048k for both LieTransformer-T2 and LieTransformer-SE2) and depth.

Model architecture. The architecture used for the SetTransformer (Lee et al., 2019) consists of 8 layers in the encoder, 8 layers in the decoder and 4 attention heads. No inducing points were used.

The architecture used for both LieTransformer-T2 and LieTransformer-SE2 is made of 10 layers, 8 heads and feature dimension $d_{v}=128$. The dimension of the kernel used is 12. One lift sample was used for each point.

Training procedure. We use Adam (Kingma & Ba, 2015) with parameters $\beta_{1}=0.5$ and $\beta_{2}=0.9$ and a learning rate of $1e-4$. Models are trained with mini-batches of size 32 until convergence.

I.2 QM9

For the QM9 experiment setup we follow the approach of Anderson et al. (2019) for parameterising the inputs and for the train/validation/test split. The $\mathtt{f}_{i}$ is a learnable linear embeddings of the vector $[1,c_{i},c_{i}^{2}]$ for charge $c_{i}$, with different linear maps for each atom type. We split the available data as follows: 100k samples for training, 10$\%$ for a test set and the rest used for validation.

In applying our model to this task, we ignore the bonding structure of the molecule. As noted in (Klicpera et al., 2019) this should not be needed to learn the task, although it may be helpful as auxiliary information. Given most methods compared against do not use such information, we follows this for a fair comparison (an exception is the $SE(3)$-Transformer (Fuchs et al., 2020) that uses the bonding information). It would be possible to utilise the bonding structure both in the neighbourhood selection step and as model features by treating only atoms that are connected via a bond to another atom as in the neighbourhood of that atom.

We performed architecture and hyperparameter optimisation on the $\epsilon_{HOMO}$ task and then trained with the resulting hyperparameters on the other 11 tasks. LieTransformer-T3 uses 13 layers of attention blocks (performance saturated at 13 layers), using 8 heads ($M$) in each layer and feature dimension $d_{v}=848$. The attention kernel uses the $linear-concat-linear$ feature embedding, identity embedding of the Lie algebra elements, and an MLP to combine these embeddings into the final attention coefficients. The final part of the model used had minor differences to the one in diagram 1. Instead of a global pooling layer followed by a 3 layer MLP, a single linear layer followed by global pooling was used. A single lift sample was used (since $H=\{e\}$ for $T(3)$-invariant models), with the radius of the neighbourhood chosen such that $|\mathtt{nbhd}_{\eta}(g)|=50~{}\forall g\in G$ and we uniformly sample 25 points from this neighbourhood. LieTransformer-SE3 used a similar hyperparameter setting, except using 30 layers (performance saturated at 30 layers) and 2 lift samples were used for each input point ($|\hat{H}|=2$), with the radius of the neighbourhood $\eta$ chosen such that the $|\mathtt{nbhd}_{\eta}(g)|=25~{}\forall g\in G$ and we uniformly sample 20 points from this neighbourhood. All models were trained using Adam, with a learning rate of $3e-4$ and a batch size of 75 for 500 epochs. For the LieConv models we used the hyperparameter setting that was used in Finzi et al. (2020).

Training these models with $T(3)$ and $SE(3)$ equivariance took approximately 3 and 6 days respectively on a single Nvidia Tesla-V100.

I.3 Hamiltonian dynamics

Spring dynamics simulation. We exactly follow the setup described in Appendix C.4 of Finzi et al. (2020) for generating the trajectories used in the train and test data.

Model architecture. In all results shown except Figures 7 and 8, we used a LieTransformer-T2 with 5 layers, 8 heads and feature dimension $d_{v}=160$. The attention kernel uses dot-product for the content component, a 3-layer MLP with hidden layer width 16 for the location component and addition to combine the content and location attention components. (Also, see end of Appendix F for a relevant discussion about the use of the $\log$ map in location-attention $k_{l}$ for this task.) We use constant normalisation of the weights. We observed a significant drop in performance when, instead of constant normalisation, we used softmax normalisation (which caused small gradients at initialization leading to optimization difficulties). The architecture had 842k parameters. Our small models (with 139k parameters) in Figures 7 and 8 use 3 layers and feature dimension $d_{v}=80$, keeping all else fixed. LieTransformer-SE(2) and LieConv-SE(2) in Figures 7 and 8 used 2 lift samples, which were deterministically chosen to be $\hat{H}=C_{2}<SO(2)$, where $C_{2}$ is the group of $180^{\circ}$ rotations. In fact, this yields exact equivariance to $T(2)\rtimes C_{2}$. Note that the true Hamiltonian $H(\mathbf{q},\mathbf{p})$ for the spring system separates as $H(\mathbf{q},\mathbf{p})=K(\mathbf{p})+V(\mathbf{q})$ where $K$ and $V$ are the kinetic and potential energies of the system respectively. Following Finzi et al. (2020), our model parameterises the potential term $V$. In particular, $x_{i}$ is particle $i$’s position and $\mathtt{f}_{i}=(m_{i},k_{i})$ where $m_{i}$ is its mass and $k_{i}$ is used to define the spring constants: $k_{i}k_{j}$ is spring constant for the spring connecting particles $i$ and $j$ (see Appendix C.4 of Finzi et al. (2020) for details).

Training details. To train the LieTransformer, we used Adam with a learning rate of 0.001 with cosine annealing and a batch size of 100. For a training dataset of size $n$, we trained the model for $400\sqrt{3000/n}$ epochs (although we found model training usually converged with fewer epochs). When $n\leq 100$, we used the full dataset in each batch. For training the LieConv baseline, we used their default architecture (with 895k parameters) and hyperparameter settings for this task, except for the number of epochs which was $400\sqrt{3000/n}$ to match the setting used for training LieTransformer.This yields better results for LieConv compared to those reported by Finzi et al. (2020), where they use fewer total epochs. The small LieConv models (with 173k parameters) in Figures 7 and 8 use 3 layers and 192 channels (instead of the default 4 layers and 384 channels). Lastly, only for the data efficiency results in Figure 4, we used early stopping by validation loss and generated nested training datasets as the training size varies, keeping the test dataset fixed.

Loss computation. One small difference between our setup and that of Finzi et al. (2020) is in the way we compute the test loss. Since we compare models’ losses not only over 5-step roll-outs but also longer 100-step roll-outs, we average the individual time step losses using a geometric mean rather than an arithmetic mean as in Finzi et al. (2020). Since the losses for later time steps are typically orders of magnitude higher than for earlier time steps (see e.g. Figure 5), a geometric mean prevents the losses for later time steps from dominating over the losses for the earlier time steps. During training, we use an arithmetic mean across time steps to compute the loss for optimization, exactly as in Finzi et al. (2020). This applies for both LieTransformer and LieConv.

Appendix J Additional experimental results