SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This additional scalar information interacts with the multivector data in two ways. In linear layers, we allow the auxiliary scalars to mix with the scalar component of the multivectors. In the attention layer, we compute attention weights both from the multivectors, as given in Eq. (5), and from the auxiliary scalars, using a regular scaled dot-product attention. The two attention maps are summed before computing the softmax, and the normalizing factor is adapted. In all other layers, the scalar information is processed separately from the multivector information, using the unrestricted form of the multivector map. For instance, nonlinearities transform multivectors with equivariant gated GELUs and auxiliary scalars with regular GELU functions. We describe the scalar path of our architecture in more detail in Appendix B.
Distance-aware dot-product attention The dot-product attention in Eq. (5) ignores the 8 dimensions involving the basis element  $e_0$ . These dimensions vary under translations, and thus their straightforward Euclidean inner product violates equivariance. We can, however, extend the attention mechanism to incorporate more components, while still maintaining E(3) equivariance and the computational efficiency of dot-product attention. To this end, we define certain auxiliary, nonlinear query features  $\phi(q)$  and key features  $\psi(k)$  and extend the attention weights in Eq. (5) as  $\langle q_{i'c}, k_{ic} \rangle \rightarrow \langle q_{i'c}, k_{ic} \rangle + \phi(q_{i'c}) \cdot \psi(k_{ic})$ , adapting the normalization appropriately. We define these nonlinear features in Appendix B.
<sup>8</sup>We also experimented with attention mechanisms that use the geometric product rather than the dot product, but found a worse performance in practice.
Figure 2:  $n$ -body dynamics experiments. We show the error in predicting future positions of planets as a function of the training dataset size. Out of five independent training runs, the mean and standard error are shown. Left: Evaluating without distribution shift. GATr (--) is more sample efficient than the equivariant SEGNN [6] (--) and SE(3)-Transformer [20] (--) and clearly outperforms non-equivariant baselines (---, ---, ---). Middle: Evaluating on systems with more planets than trained on. Transformer architectures generalize well to different object counts. GCA-GNN has larger errors than the visible range. Right: Evaluating on translated data. Because GATr is E(3) equivariant, it generalizes under this domain shift.
Our choice of these nonlinear features not only maintains equivariance, but has a straightforward geometric interpretation. When the trivector components of queries and keys represent 3D points (see Tbl. 1),  $\psi(q) \cdot \phi(k)$  is proportional to the pairwise negative squared Euclidean distance. GATr's attention mechanism is therefore directly sensitive to Euclidean distance, while still respecting the highly efficient dot-product attention format.
Positional embeddings GATr assumes the data can be described as a set of items (or tokens). If these items are distinguishable and form a sequence, we encode their position using "rotary positional"[9] embeddings [47] in the auxiliary scalar variables.[10]
Axial attention The architecture is flexible about the structure of the data. In some use cases, there will be a single dimension along which objects are organized, for instance when describing a static scene or the time evolution of a single object. But GATr also supports the organization of a problem along multiple axes, for example with one dimension describing objects and another time steps. In this case, we follow an axial transformer layout [24], alternating between transformer blocks that attend over different dimensions. (The not-attended dimensions in each block are treated like a batch dimension.)
4.1  $n$ -body dynamics
We start our empirical demonstration of GATr with an  $n$ -body dynamics problem, on which we compare GATr to a wide range of baselines. Given the masses, initial positions, and velocities of a star and a few planets, the goal is to predict the final position after the system has evolved under Newtonian gravity for 1000 Euler integration time steps. We compare GATr to a Transformer and an MLP, the equivariant SEGNN [6] and SE(3)-Transformer [20], as well as the geometric-algebra-based—but not equivariant—GCA-GNN [41]. The experiment is described in detail in Appendix C.1.
In the left panel of Fig. 2 we show the prediction error as a function of the number of training samples used. GATr clearly outperforms all non-equivariant baselines, including the geometric-algebra-based GCA-GNN. Compared to the equivariant SEGNN and SE(3)-Transformer, GATr is more sample-
<sup>9</sup>This terminology, which stems from non-geometric transformers, can be confusing. "Position" here means position in a sequence, not geometric position in 3D. "Rotary" does not refer to a rotation of 3D space, but rather to how the position in a sequence is embedded via sinusoids in the scalar channels of keys and queries. Using positional encoding thus does not affect E(3) equivariance.
<sup>10</sup>Since auxiliary scalar representations and multivectors mix in the attention mechanism, the positional embeddings also affect the multivector processing.
efficient, while all three methods reach the same prediction error when trained on enough data. This provides evidence for the usefulness of geometric algebra representations as an inductive bias.
GATr also generalizes robustly out of domain, as we show in the middle and right panels of Fig. 1. When evaluating on a larger number of bodies than trained on, methods that use a softmax over attention weights (GATr, Transformer, SE(3)-Transformer) generalize best. Finally, the performance of the E(3)-equivariant GATr, SEGNN, and SE(3)-Transformer does not drop when evaluated on spatially translated data, while the non-equivariant baselines fail in this setting.
4.2 Wall-shear-stress estimation on large meshes of human arteries
Next, we turn towards a realistic experiment involving more complex geometric objects. We study the prediction of the wall shear stress exerted by the blood flow on the arterial wall, an important predictor of aneurysms. While the wall shear stress can be computed with computational fluid dynamics, simulating a single artery can take many hours, and efficient neural surrogates can have substantial impact. However, training such neural surrogates is challenging, as meshes are large (around 7000 nodes in our data) and datasets typically small (we work with 1600 training meshes).
We train GATr on a dataset of arterial meshes and simulated wall shear stress published by Suk et al. [48]. They are compared to a Transformer and to the results reported by Suk et al. [48], including the equivariant GEM-CNN [15] and the non-equivariant PointNet++ [36].<sup>11</sup> See Appendix C.2 for experiment details.
The results are shown in Fig. 3. On noncanonicalized data, with randomly rotated meshes, GATr improves upon all previous methods and sets a new state of the art.
We also experiment with canonicalization: rotating the arteries such that blood always flows in be almost competitive with GATr. However, car arteries as in this dataset, not in more complex so it likely that GATr will be more robust in such s
Figure 3: Arterial wall-shear-stress estimation [48]. We show the mean approximation error (lower is better) as a function of training dataset size, reporting results both on randomly oriented training and test samples (solid markers) and on a version of the dataset in which all artery meshes are canonically oriented (hollow markers). Without canonicalization, GATr  $(\longrightarrow)$  predicts wall shear stress more precisely and is more sample-efficient than the baselines.
4.3 Robotic planning through invariant diffusion
In our third experiment, we show how GATr defines an E(3)-invariant diffusion model, that it can be used for model-based reinforcement learning and planning, and that this combination is well-suited to solve robotics problems. We follow Janner et al. [27], who propose to treat learning a world model and planning within that model as a unified generative modeling problem. After training a diffusion model [45] on offline trajectories, one can use it in a planning loop, sampling from it conditional on the current state, desired future states, or to maximize a given reward, as needed.
We use a GATr model as the denoising network in a diffusion model and to use it for planning. We call this combination  $GATr$ -Diffuser. Combining the equivariant GATr with an invariant base density defines an  $\mathrm{E}(3) \times \mathrm{S}_n$ -invariant diffusion model [29].
GATr-Diffuser is demonstrated on the problem of a Kuka robotic gripper stacking blocks using the "unconditional" environment introduced by Janner et al. [27]. We train GATr-Diffuser on the offline
<sup>11</sup>We also tried to run SEGNN [5] and SE(3)-Transformer [20] as additional baselines, but were not able to fit them into memory on this task.
trajectory dataset published with that paper and then use it for planning, following the setup of Janner et al. [27]. We compare our GATr-Diffuser model to a reproduction of the original Diffuser model and a new Transformer backbone for the Diffuser model. In addition, we show the published results of Diffuser [27], the equivariant EDGI [7], and the offline RL algorithms CQL [30] and BCQ [21] as published in Ref. [27]. The problem and hyperparameters are described in detail in Appendix C.3.
As shown in Fig. 4, GATr-Diffuser solves the block-stacking task better than all baselines. It is also clearly more sample-efficient, matching the performance of a Diffuser model or Transformer trained on the full dataset even when training only on  $1\%$  of the trajectories.
4.4 Scaling
Finally, we study GATr's computational requirements and scaling. We measure the memory usage and compute time of forward and backward passes on synthetic data as a function of the number of items. GATr is compared to a
Transformer, to SEGNN [6] in the official implementation, and to the SE(3)-Transformer [20]; for the latter we use the highly optimized implementation by Milesi [35]. We choose hyperparameters for the four architectures such that they have the same depth and width and require that the methods allow all items to interact at each step (i.e. fully connected graphs). Because of the fundamental differences between the architectures, it is impossible to find fully equivalent settings; our results should thus be interpreted with care. The details of our experiment are described in Appendix C.4.
We show the results in Fig. 5. For few tokens, GATr is slower than a Transformer. However, this difference is partially due to the low utilization of the GPUs in this test; GATr is closer to the
Transformer when using larger batch sizes or when pre-compiling the computation graph.
For larger problems, compute and memory are dominated by the pairwise interactions in the attention mechanism. Here GATr and the baseline Transformer perform on par, as they use the same efficient implementation of dot-product attention. In terms of memory, both scale linearly in the number of tokens, while the equivariant baselines scale quadratically. In terms of time, all methods scale quadratically, but the equivariant baselines have a worse prefactor than GATr and the Transformer. All in all, we can easily scale GATr to tens of thousands of tokens, while the equivariant baselines run out of memory two orders of magnitude earlier.
Geometric algebra Geometric (or Clifford) algebra was first conceived in the 19th century [10, 22] and has been used widely in quantum physics [16, 34]. Recently, it has found new popularity in computer graphics [18]. In particular,
Figure 4: Diffusion-based robotic planning. We show normalized rewards (higher is better) as a function of training dataset size. GATr  $(\text{一})$  is more successful at block stacking and more sample-efficient than the baselines, including the original Diffuser [27]  $(\text{一} \cdot)$  and our modification of it based on a Transformer  $(\text{一} \cdot)$ . In grey, we show results reported by Brehmer et al. [7] and Janner et al. [27].
Figure 5: Computational cost and scaling. We measure peak GPU memory usage (top) and wall time (bottom) per combined forward and backward pass as a function of the number of items in synthetic data. Despite some compute overhead, GATr  $(\text{一})$  scales just like the Transformer  $(\text{一})$  and orders of magnitude more favorably than the equivariant baselines  $(\text{一}, \text{一})$ .
the projective geometric algebra used in this work [17, 38] and a conformal model [18] are suitable for 3D computations.
Geometric algebras have been used in machine learning in various ways. Spellings [46] use  $\mathbb{G}_{3,0,0}$  geometric products to compute rotation-invariant features from small point clouds. Unlike us, they do not learn internal geometric representations.
Our work was inspired by Brandstetter et al. [5] and Ruhe et al. [41]. These works also use multivector features (the latter even of the same  $\mathbb{G}_{3,0,1}$ ), and process them with operations such as the geometric/sandwich product, Clifford Fourier transforms and Clifford convolutions. The main difference to our work is that GATr is E(3) equivariant, while both of these works are not. We compare to the GCA-GNN network from Ref. [41] in our experiments.
Concurrently to this publication, Ruhe et al. [40] also study equivariant, geometric algebra-based architectures. While some of our and their findings overlap, there are several differences: They develop theory for generic geometric algebras of the form  $\mathbb{G}_{p,q,r}$ , while we focus on the projective algebra  $\mathbb{G}_{3,0,1}$  with its faithful  $\mathrm{E}(3)$  representations, with our theory also applicable to the group  $E(n)$  and the algebras  $\mathbb{G}_{n,0,0}$  and  $\mathbb{G}_{n,0,1}$ . Ruhe et al. [40] also finds the linear maps of Eq. (4), but does not discuss the join or distance-aware dot-product, which we found to be essential for performance in the projective algebra. Moreover, they propose an MLP-like architecture and use it in a message-passing graph network, while our GATr is a Transformer.
**Equivariant deep learning** Equivariance to symmetries [11] is the primary design principle in modern geometric deep learning [8]. Equivariant networks have been applied successfully in areas such as medical imaging [31, 53] and robotics [7, 26, 44, 51, 52, 55], and are ubiquitous in applications of machine learning to physics and chemistry [1, 3, 4, 19, 28, 33].
In recent years, a number of works have investigated equivariant Transformer and message-passing architectures [3, 4, 6, 19, 20, 39, 42, 49]. These works are generally more limited in terms of the types of geometric quantities they can process compared to our multivector features. Furthermore, our architecture is equivariant to the full group of Euclidean transformations, whereas previous works focus on SO(3) equivariance.
We introduced the Geometric Algebra Transformer (GATr), a general-purpose architecture for geometric data, and implemented it at https://github.com/qualcomm-ai-research/geometric-algebra-transformer. We argued and demonstrated that GATr effectively combines structure and scalability.
GATr incorporates geometric structure by representing data in projective geometric algebra, as well as through E(3) equivariance. Unlike most equivariant architectures, GATr features faithful E(3) representations, including absolute positions and equivariance with respect to translations. Empirically, GATr outperforms non-geometric, equivariant, and geometric algebra-based non-equivariant baselines across three experiments.
At the same time, GATr scales much better than most geometric networks. This is because GATr is a Transformer and computes pairwise interactions through dot-product attention. Using recent efficient attention implementations, we demonstrated that we can scale GATr to systems with many thousands of tokens and fully connected interactions. For few tokens and small batch sizes, GATr has some computational overhead, which we hope to address in future implementations.
One drawback of our approach is that since geometric algebra is not widely known yet, it may present an obstacle to understanding the details of the method. However, given a dictionary for embeddings of common objects and a library of primitives that act on them, using this framework is no more difficult than using typical neural network layers grounded in linear algebra. Another potential downside is that GATr is not yet shown to be a universal approximator, which is an interesting direction for future work.
Given the promising results presented in this work, we look forward to further study the potential of Geometric Algebra Transformers in problems from molecular dynamics to robotics.
Acknowledgements We would like to thank Joey Bose, Johannes Brandstetter, Gabriele Cesa, Steven De Keninck, Daniel Dijkman, Leo Dorst, Mario Geiger, Jonas Kohler, Ekdeep Singh Lubana, Evgeny Mironov, and David Ruhe for generous guidance regarding geometry, enthusiastic encouragement on equivariance, and careful counsel concerning computing complications.
[52] Dian Wang, Mingxi Jia, Xupeng Zhu, Robin Walters, and Robert Platt. On-Robot learning with equivariant models. In CoRL, 2022. (Cited on page 10)
[53] Marysia Winkels and Taco Cohen. Pulmonary nodule detection in CT scans with equivariant CNNs. MIA, 2019. (Cited on page 10)
[54] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020. (Cited on pages 5 and 22)
[55] Xupeng Zhu, Dian Wang, Ondrej Biza, Guanang Su, Robin Walters, and Robert Platt. Sample efficient grasp learning using equivariant models. In Robotics: Science and Systems XVIII. Robotics: Science and Systems Foundation, June 2022. (Cited on page 10)
In this section, we state or prove several properties of equivariant maps between geometric algebras that we use in the construction of GATr.
The grade involution is a linear involutive bijection  $\widehat{\cdot}:\mathbb{G}_{n,0,r}:\mathbb{G}_{n,0,r}$ , which sends a  $k$ -blade  $x$  to  $\widehat{x} = (-1)^k x$ . Note that this is an algebra automorphism  $\widehat{xy} = \widehat{xy}$ , and also an  $\wedge$ -algebra automorphism. The reversal in a linear involutive bijection  $\widetilde{\cdot}:\mathbb{G}_{n,0,r}:\mathbb{G}_{n,0,r}$  which sends a  $k$ -blade  $x = x_1\wedge x_2\wedge \ldots \wedge x_k$  to the reverse:  $\widetilde{x} = x_k\wedge \ldots \wedge x_2\wedge x_1 = \pm x$  with  $+x$  if  $k\in \{0,1,4,5,\ldots ,8,9,\ldots \}$  and  $-x$  otherwise. Note that this is an anti-automorphism (contravariant functor):  $\widetilde{xy} = \widetilde{y\bar{x}}$ .
Here we denote the sandwich action of  $u \in \operatorname{Pin}(n,0,r)$  on a multivector  $x$  not as  $\rho_u(x)$ , but as  $u[x]$ . For odd  $u$ ,  $u[x] = u\widehat{x}u^{-1}$ , while for even  $u$ ,  $u[x] = uxu^{-1}$ . The sandwich action is linear by linearity of the  $\widehat{\cdot}$  and bilinearity of the geometric product. Furthermore, note that for any particular  $u \in \operatorname{Pin}(n,0,r)$ , the action is a geometric algebra homomorphism:  $u[ab] = u\widehat{a}b u^{-1} = u\widehat{a}u^{-1}u\widehat{b}u^{-1} = u[a]u[b]$ . By linearity and a symmetrization argument [18, Sec 7.1], one can show that it also a  $\wedge$ -algebra homomorphism (outermorphism):  $u[a \wedge b] = u[a] \wedge u[b]$ .
Let  $l \geq k$ . Given a  $k$ -vector  $a$  and  $l$ -vector  $b$ , define the left contraction as  $a \rfloor b := \langle ab \rangle_{l-k}$ , which is a  $l-k$ -vector. For  $k = 1$ , and  $b$  a blade  $b = b_1 \wedge \ldots \wedge b_l$ . Geometrically,  $a \rfloor b$  is the projection of  $a$  to the space spanned by the vectors  $b_i$ . Thus we have that  $a \rfloor b = 0 \iff \forall i, \langle a, b_i \rangle = 0$  [18, Sec 3.2.3], in which case we define  $a$  and  $b$  to be orthogonal. In particular, two vectors  $a, b$  are orthogonal if their inner product is zero. Furthermore, we define a vector  $a$  to be tangential to blade  $b$  if  $a \wedge b = 0$ .
In the projective algebra, a blade  $x$  is defined to be ideal if it can be written as  $x = e_0 \wedge y$  for another blade  $y$ .
A.1 Linear maps
We begin with Pin-equivariant linear maps. After some technical lemmata, we prove the most general form of linear equivariant maps in the Euclidean geometric algebra  $\mathbb{G}_{n,0,0}$ , and then also in projective geometric algebra  $\mathbb{G}_{n,0,1}$ .
Proposition 2. The grade projection  $\langle \cdot \rangle_{k}$  is equivariant [18, Sec 13.2.3].
Proof. Choose an  $l$ -blade  $x = a_{1} \wedge a_{2} \wedge \ldots \wedge a_{l}$ . Let  $u$  be a 1-versor. As the action  $u$  is an outermorphism,  $u[x] = u[a_{1}] \wedge \ldots \wedge u[a_{l}]$  is an  $l$ -blade. Now if  $l \neq k$ , then  $\langle x \rangle_{k} = 0$  and thus  $u[\langle x \rangle_{k}] = \langle u[x] \rangle_{k}$ . If  $l = k$ , then  $\langle x \rangle_{k} = x$  and thus  $u[\langle x \rangle_{k}] = \langle u[x] \rangle_{k}$ . As the grade projection is linear, equivariance extends to any multivector.
Proposition 3. The following map is equivariant:  $\phi : \mathbb{G}_{3,0,1} \to \mathbb{G}_{3,0,1} : x \mapsto e_0x$ .
Proof. Let  $u$  be a 1- versor, then  $u$  acts on a multivector as  $x \mapsto u[x] = u\hat{x}u^{-1}$ , where  $\hat{x}$  is the grade involution. Note that  $e_0$  is invariant:  $u[e_0] = -ue_0u^{-1} = e_0uu^{-1} = e_0$ , where  $ue_0 = -e_0u$  because  $u$  and  $e_0$  are orthogonal:  $ue_0 = \langle u, e_0 \rangle + u \wedge e_0 = -e_0 \wedge u = -e_0 \wedge u$ . Then  $\phi$  is equivariant, as the action is an algebra homomorphism:  $u[\phi(x)] = u[e_0x] = \widehat{ue_0xu^{-1}} = u\hat{e}_0u^{-1}u\hat{x}u^{-1} = u[e_0]u[x] = e_0u[x] = \phi(u[x])$ . It follows that  $\phi$  is also equivariant to any product of vectors, i.e. any versor  $u$ .
Euclidean geometric algebra Before constructing the most general equivariant linear map between multivectors in projective geometric algebra, we begin with the Euclidean case  $\mathbb{G}_{n,0,0}$ .
Theorem 1 (Cartan-Dieuodonne). Every orthogonal transformation of an  $n$ -dimensional space can be decomposed into at most  $n$  reflections in hyperplanes.
Proof. This theorem is proven in Roelfs and De Keninck [38].
Lemma 1. In the  $n$ -dimensional Euclidean geometric algebra  $\mathbb{G}_{n,0,0}$ , the group  $\operatorname{Pin}(n,0,0)$  acts transitively on the space of  $k$ -blades of norm  $\lambda \in \mathbb{R}^{>0}$ .
Proof. As the Pin group preserves norm, choose  $\lambda = 1$  without loss of generality. Any  $k$ -blade  $x$  of unit norm can be written by Gram-Schmidt factorization as the wedge product of  $k$  orthogonal vectors of unit norm  $x = v_{1} \wedge v_{2} \wedge \ldots \wedge v_{k}$ . Consider another  $k$ -blade  $y = w_{1} \wedge w_{2} \wedge \ldots \wedge w_{k}$  with  $w_{i}$  orthonormal. We'll construct a  $u \in \operatorname{Pin}(n,0,0)$  such that  $u[x] = y$ .
Choose  $n - k$  additional orthonormal vectors  $v_{k + 1},\dots,v_n$  and  $w_{k + 1},\dots,w_n$  to form orthonormal bases. Then, there exists a unique orthogonal transformation  $\mathbb{R}^n\to \mathbb{R}^n$  that maps  $v_{i}$  into  $w_{i}$  for all  $i\in \{1,\ldots ,n\}$ . By the Cartan-Dieudonne theorem 1, this orthogonal transformation can be expressed as the product of reflections, thus there exists a  $u\in \mathrm{Pin}(n,0,0)$  such that  $u[v_i] = w_i$ . As the  $u$  action is a  $\wedge$ -algebra homomorphism  $(u[a\wedge b] = u[a]\wedge u[b]$ , for any multivectors  $a,b$ ), we have that  $u[x] = y$ .
Lemma 2. In the Euclidean  $(r = 0)$  or projective  $(r = 1)$  geometric algebra  $\mathbb{G}_{n,0,r}$ , let  $x$  be a  $k$ -blade. Let  $u$  be a 1- versor. Then  $u[x] = x \iff u\rfloor x = 0$  and  $u[x] = -x \iff u\wedge x = 0$ .