SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pim de Haan*
Sönke Behrends
{jbrehmer, pim, sbehrend, tacos}@qti.QUALCOMM.com
Problems involving geometric data arise in physics, chemistry, robotics, computer vision, and many other fields. Such data can take numerous forms, for instance points, direction vectors, translations, or rotations, but to date there is no single architecture that can be applied to such a wide variety of geometric types while respecting their symmetries. In this paper we introduce the Geometric Algebra Transformer (GATr), a general-purpose architecture for geometric data. GATr represents inputs, outputs, and hidden states in the projective geometric (or Clifford) algebra, which offers an efficient 16-dimensional vector-space representation of common geometric objects as well as operators acting on them. GATr is equivariant with respect to  $\mathrm{E}(3)$ , the symmetry group of 3D Euclidean space. As a Transformer, GATr is versatile, efficient, and scalable. We demonstrate GATr in problems from  $n$ -body modeling to wall-shear-stress estimation on large arterial meshes to robotic motion planning. GATr consistently outperforms both non-geometric and equivariant baselines in terms of error, data efficiency, and scalability.
From molecular dynamics to astrophysics, from material design to robotics, fields across science and engineering deal with geometric data such as positions, shapes, orientations, or directions. The geometric nature of data provides a rich structure: a notion of common operations between geometric types (computing distances between points, applying rotations to orientations, etc.), a well-defined behaviour of data under transformations of a system, and the independence of certain properties of coordinate system choices. When learning from geometric data, incorporating this rich structure into the architecture has the potential to improve the performance.
With this goal, we introduce the Geometric Algebra Transformer (GATr), a general-purpose network architecture for geometric data. GATr brings together three key ideas.
Geometric algebra: To naturally describe both geometric objects as well as their transformations in three-dimensional space, GATr represents data as multivectors of the projective geometric algebra  $\mathbb{G}_{3,0,1}$ . Geometric (or Clifford) algebra is an principled yet practical mathematical framework for geometrical computations. The particular algebra  $\mathbb{G}_{3,0,1}$  extends the vector space  $\mathbb{R}^3$  to 16-dimensional multivectors, which can natively represent various geometric types and  $\mathrm{E}(3)$  poses. Unlike the O(3) representations popular in geometric deep learning, this algebra can represent data that is not invariant to translations, such as absolute positions.
Equivalence: GATr is equivariant with respect to  $\mathrm{E}(3)$ , the symmetry group of three-dimensional space. To this end, we develop several new  $\mathrm{E}(3)$ -equivariant primitives mapping between multivectors, including equivariant linear maps, an attention mechanism, nonlinearities, and normalization layers.
$^{\dagger}$ Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc.
Transformer: Due to its favorable scaling properties, expressiveness, trainability, and versatility, the Transformer architecture [50] has become the de-facto standard for a wide range of problems. GATr is based on the Transformer architecture, in particular on dot-product attention, and inherits these benefits.
GATr thus combines two lines of research: the representation of geometric objects with geometric algebra [17, 18, 38], popular in computer graphics and physics and recently gaining traction in deep learning [5, 41, 46], and the encoding of symmetries through equivariant deep learning [12]. The result—to the best of our knowledge the first E(3)-equivariant architecture with internal geometric algebra representations<sup>3</sup—is a versatile network for problems involving geometric data.
We demonstrate GATr in three problems from entirely different fields. In an  $n$ -body modelling task, we compare GATr to various baselines. We turn towards the task of predicting the wall shear stress in human arteries, demonstrating that GATr scales to realistic problems with meshes of thousands of nodes. Finally, experiment with robotic motion planning, using GATr as the backbone of an E(3)-invariant diffusion model. In all cases, GATr substantially outperforms both non-geometric and equivariant baselines.
Our implementation of GATr is available at https://github.com/qualcomm-ai-research/geometric-algebra-transformer.
Geometric algebra We begin with a brief overview of geometric algebra (GA); for more detail, see e.g. Refs. [17, 18, 38, 41]. Whereas a plain vector space like  $\mathbb{R}^3$  allows us to take linear combinations of elements  $x$  and  $y$  (vectors), a geometric algebra additionally has a bilinear associative operation: the geometric product, denoted simply by  $xy$ . By multiplying vectors, one obtains so-called multivectors, which can represent both geometrical objects and operators. Like vectors, multivectors have a notion of direction as well as magnitude and orientation (sign), and can be linearly combined.
Multivectors can be expanded on a multivector basis, consisting of products of basis vectors. For example, in a 3D GA with orthogonal basis  $e_1, e_2, e_3$ , a general multivector takes the form
$$
x = x _ {s} + x _ {1} e _ {1} + x _ {2} e _ {2} + x _ {3} e _ {3} + x _ {1 2} e _ {1} e _ {2} + x _ {1 3} e _ {1} e _ {3} + x _ {2 3} e _ {2} e _ {3} + x _ {1 2 3} e _ {1} e _ {2} e _ {3}, \tag {1}
$$
with real coefficients  $(x_{s},x_{1},\ldots ,x_{123})\in \mathbb{R}^{8}$ . Thus, similar to how a complex number  $a + bi$  is a sum of a real scalar and an imaginary number, a general multivector is a sum of different kinds of elements. These are characterized by their dimensionality (grade), such as scalars (grade 0), vectors  $e_i$  (grade 1), bivectors  $e_i e_j$  (grade 2), all the way up to the pseudoscalar  $e_1\dots e_d$  (grade  $d$ ).
The geometric product is characterized by the fundamental equation  $vv = \langle v, v \rangle$ , where  $\langle \cdot, \cdot \rangle$  is an inner product. In other words, we require that the square of a vector is its squared norm. In an orthogonal basis, where  $\langle e_i, e_j \rangle \propto \delta_{ij}$ , one can deduce that the geometric product of two different basis vectors is antisymmetric:  $e_i e_j = -e_j e_i^5$ . Since reordering only produces a sign flip, we only get one basis multivector per unordered subset of basis vectors, and so the total dimensionality of a GA is  $\sum_{i=0}^{d} \binom{d}{k} = 2^d$ . Moreover, using bilinearity and the fundamental equation one can compute the geometric product of arbitrary multivectors.
The symmetric and antisymmetric parts of the geometric product are called the interior and exterior (wedge) product. For vectors  $x$  and  $y$ , these are defined as  $\langle x, y \rangle = (xy + yx)/2$  and  $x \wedge y \equiv (xy - yx)/2$ . The former is indeed equal to the inner product used to define the GA, whereas the latter is new notation. Whereas the inner product computes the similarity, the exterior product constructs a multivector (called a blade) representing the weighted and oriented subspace spanned by the vectors. Both operations can be extended to general multivectors [18].
The final primitive of the GA that we will require is the dualization operator  $x \mapsto x^*$ . It acts on basis elements by swapping "empty" and "full" dimensions, e.g. sending  $e_1 \mapsto e_{23}$ .
3Concurrently to our work, a similar network architecture was studied by Ruhe et al. [40]; we comment on similarities and differences in Sec. 5.
Indeed the imaginary unit  $i$  can be thought of as the bivector  $e_1e_2$  in a 2D GA.
The antisymmetry can be derived by using  $v^2 = \langle v, v \rangle$  to show that  $e_i e_j + e_j e_i = (e_i + e_j)^2 - e_i^2 - e_j^2 = 0$ .
Object / operator	Scalar 1	Vector e0 ei	Bivector e0i ej	Trivector e0ij e123	PS e0123
Scalar λ ∈ R	λ	0	0	0	0	0	0	0
Plane w/ normal n ∈ R3, origin shift d ∈ R	0	d	n	0	0	0	0	0
Line w/ direction n ∈ R3, orthogonal shift s ∈ R3	0	0	0	s	n	0	0	0
Point p ∈ R3	0	0	0	0	0	p	1	0
Pseudoscalar μ ∈ R	0	0	0	0	0	0	0	μ
Reflection through plane w/ normal n ∈ R3, origin shift d ∈ R	0	d	n	0	0	0	0	0
Translation t ∈ R3	1	0	0	1/2t	0	0	0	0
Rotation expressed as quaternion q ∈ R4	q0	0	0	0	qi	0	0	0
Point reflection through p ∈ R3	0	0	0	0	0	p	1	0
Table 1: Embeddings of common geometric objects and transformations into the projective geometric algebra  $\mathbb{G}_{3,0,1}$ . The columns show different components of the multivectors with the corresponding basis elements, with  $i,j\in \{1,2,3\}, j\neq i$ , i.e.  $ij\in \{12,13,23\}$ . For simplicity, we fix gauge ambiguities (the weight of the multivectors) and leave out signs (which depend on the ordering of indices in the basis elements).
Projective geometric algebra In order to represent three-dimensional objects as well as arbitrary rotations and translations acting on them, the 3D GA is not enough: as it turns out, its multivectors can only represent linear subspaces passing through the origin as well as rotations around it. A common trick to expand the range of objects and operators is to embed the space of interest (e.g.  $\mathbb{R}^3$ ) into a higher dimensional space whose multivectors represent more general objects and operators in the original space.
In this paper we work with the projective geometric algebra  $\mathbb{G}_{3,0,1}$  [17, 38, 41]. Here one adds a fourth homogeneous coordinate  $x_0e_0$  to the vector space, yielding a  $2^4 = 16$ -dimensional geometric algebra. The metric of  $\mathbb{G}_{3,0,1}$  is such that  $e_0^2 = 0$  and  $e_i^2 = 1$  for  $i = 1,2,3$ . As we will explain in the following, in this setup the 16-dimensional multivectors can represent 3D points, lines, and planes, which need not pass through the origin, and arbitrary rotations, reflections, and translations in  $\mathbb{R}^3$ .
Representing transformations In geometric algebra, a vector  $u$  can act as an operator, reflecting other elements in the hyperplane orthogonal to  $u$ . Since any orthogonal transformation is equal to a sequence of reflections, this allows us to express any such transformation as a geometric product of (unit) vectors, called a (unit) versor  $u = u_{1} \cdots u_{k}$ . Furthermore, since the product of unit versors is a unit versor, and unit vectors are their own inverse ( $u^{2} = 1$ ), these form a group called the Pin group associated with the metric. Similarly, products of an even number of reflections form the Spin group. In the projective geometric algebra  $\mathbb{G}_{3,0,1}$ , these are the double cover<sup>6</sup> of  $\mathrm{E}(3)$  and  $\mathrm{SE}(3)$ , respectively. We can thus represent any rotation, translation, and mirroring—the symmetries of three-dimensional space—as  $\mathbb{G}_{3,0,1}$  multivectors.
In order to apply a versor  $u$  to an arbitrary element  $x$ , one uses the sandwich product:
$$
\rho_ {u} (x) = \left\{ \begin{array}{l l} u x u ^ {- 1} & \text {i f} u \text {i s e v e n} \\ u \hat {x} u ^ {- 1} & \text {i f} u \text {i s o d d} \end{array} \right. \tag {2}
$$
Here  $\hat{x}$  is the grade involution, which flips the sign of odd-grade elements such as vectors and trivectors, while leaving even-grade elements unchanged. Equation 2 thus gives us a linear action (i.e. group representation) of the Pin and Spin groups on the  $2^{d}$ -dimensional space of multivectors. The sandwich product is grade-preserving, so this representation splits into a direct sum of representations on each grade.
Representing 3D objects Following Refs. [17, 38, 41], we represent planes with vectors, and require that the intersection of two geometric objects is given by the wedge product of their representations. Lines (the intersection of two planes) are thus represented as bivectors, points (the intersection of three planes) as trivectors. This leads to a duality between objects and operators, where objects are represented like transformations that leave them invariant. Table 1 provides a dictionary of these embeddings. It is easy to check that this representation is consistent with using the sandwich product for transformations.
This means that for each element of  $\mathrm{E}(3)$  there are two elements of  $\operatorname{Pin}(3,0,1)$ , e.g. both the vector  $v$  and  $-v$  represent the same reflection.
Equivalence We construct network layers that are equivariant with respect to  $\operatorname{E}(3)$ , or equivalently its double cover  $\operatorname{Pin}(3,0,1)$ . A function  $f: \mathbb{G}_{3,0,1} \to \mathbb{G}_{3,0,1}$  is  $\operatorname{Pin}(3,0,1)$ -equivariant with respect to the representation  $\rho$  (or  $\operatorname{Pin}(3,0,1)$ -equivariant for short) if
$$
f \left(\rho_ {u} (x)\right) = \rho_ {u} (f (x)) \tag {3}
$$
for any  $u \in \operatorname{Pin}(3,0,1)$  and  $x \in \mathbb{G}_{3,0,1}$ , where  $\rho_u(x)$  is the sandwich product defined in Eq. (2).
3 The Geometric Algebra Transformer
3.1 Design principles and architecture overview
Geometric inductive bias through geometric algebra representations GATr is designed to provide a strong inductive bias for geometric data. It should be able to represent different geometric objects and their transformations, for instance points, lines, planes, translations, rotations, and so on. In addition, it should be able to represent common interactions between these types with few layers, and be able to identify them from little data (while maintaining the low bias of large transformer models). Examples of such common patterns include computing the relative distances between points, applying geometric transformations to objects, or computing the intersections of planes and lines.
Following a body of research in computer graphics, we propose that geometric algebra gives us a language that is well-suited to this task. We use the projective geometric algebra  $\mathbb{G}_{3,0,1}$  and use the plane-based representation of geometric structure outlined in the previous section.
Symmetry awareness through E(3) equivariance Our architecture should respect the symmetries of 3D space. Therefore we design GATr to be equivariant with respect to the symmetry group E(3) of translations, rotations, and reflections.
Note that the projective geometric algebra naturally offers a faithful representation of  $\mathrm{E}(3)$ , including translations. We can thus represent objects that transform arbitrarily under  $\mathrm{E}(3)$ , including with respect to translations of the inputs. This is in stark contrast with most  $\mathrm{E}(3)$ -equivariant architectures, which only use  $\mathrm{O}(3)$  representations and whose features only transform under rotation and are invariant under translations—and hence can not represent absolute positions like GATr can. Those architectures must handle points in hand-crafted ways, like by canonicalizing w.r.t. the center of mass or by treating the difference between points as a translation-invariant vector.
Many systems will not exhibit the full E(3) symmetry group. The direction of gravity, for instance, often breaks it down to the smaller E(2) group. To maximize the versatility of GATr, we choose to develop a E(3)-equivariant architecture and to include symmetry-breaking as part of the network inputs, similar to how position embeddings break permutation equivariance in transformers.
Scalability and flexibility through dot-product attention Finally, GATr should be expressive, easy to train, efficient, and scalable to large systems. It should also be as flexible as possible, supporting variable geometric inputs and both static scenes and time series.
These desiderata motivate us to implement GATr as a transformer [50], based on attention over multiple objects (similar to tokens in MLP or image patches in computer vision). This choice makes GATr equivariant also with respect to permutations along the object dimension. As in standard transformers, we can break this equivariance when desired (in particular, along time dimensions) through positional embedding.
Like a vanilla transformer, GATr is based on a dot-product attention mechanism, for which heavily optimized implementations exist [14, 32, 37]. We will demonstrate later that this allows us to scale GATr to problems with many thousands of tokens, much further than equivariant architectures based on graph neural networks and message-passing algorithms.
Architecture overview We sketch GATr in Fig. 1. In the top row, we sketch the overall workflow. If necessary, raw inputs are first preprocessed into geometric types. The geometric objects are then embedded into multivectors of the geometric algebra  $\mathbb{G}_{3,0,1}$ , following the recipe described in Tbl. 1.
The multivector-valued data are processed with a GATr network. We show this architecture in more detail in the bottom row of Fig. 1. GATr consists of  $N$  transformer blocks, each consisting
Figure 1: Overview over the GATr architecture. Boxes with solid lines are learnable components, those with dashed lines are fixed.
of an equivariant multivector LayerNorm, an equivariant multivector self-attention mechanism, a residual connection, another equivariant LayerNorm, an equivariant multivector MLP with geometric bilinear interactions, and another residual connection. The architecture is thus similar to a typical transformer [50] with pre-layer normalization [2, 54], but adapted to correctly handle multivector data and be E(3) equivariant. We describe the individual layers below.
Finally, from the outputs of the GATr network we extract the target variables, again following the mapping given in Tbl. 1.
3.2 GATr primitives
Linear layers We begin with linear layers between multivectors. In Appendix A, we show that the equivariance condition of Eq. (3) severely constrains them:
Proposition 1. Any linear map  $\phi : \mathbb{G}_{d,0,1} \to \mathbb{G}_{d,0,1}$  that is equivariant to  $\operatorname{Pin}(d,0,1)$  is of the form
$$
\phi (x) = \sum_ {k = 0} ^ {d + 1} w _ {k} \langle x \rangle_ {k} + \sum_ {k = 0} ^ {d} v _ {k} e _ {0} \langle x \rangle_ {k} \tag {4}
$$
for parameters  $w \in \mathbb{R}^{d + 2}, v \in \mathbb{R}^{d + 1}$ . Here  $\langle x \rangle_k$  is the blade projection of a multivector, which sets all non-grade- $k$  elements to zero.
Thus,  $\operatorname{E}(3)$ -equivariant linear maps between  $\mathbb{G}_{3,0,1}$  multivectors can be parameterized with nine coefficients, five of which are the grade projections and four include a multiplication with the homogeneous basis vector  $e_0$ . We thus parameterize affine layers between multivector-valued arrays with Eq. (4), with learnable coefficients  $w_k$  and  $v_k$  for each combination of input channel and output channel. In addition, there is a learnable bias term for the scalar components of the outputs (biases for the other components are not equivariant).
Geometric bilinears Equivariant linear maps are not sufficient to build expressive networks. The reason is that these operations allow for only very limited grade mixing, as shown in Prop. 1. For the network to be able to construct new geometric features from existing ones, such as the translation vector between two points, two additional primitives are essential.
The first is the geometric product  $x, y \mapsto xy$ , the fundamental bilinear operation of geometric algebra. It allows for substantial mixing between grades: for instance, the geometric product of vectors consists of scalars and bivector components. The geometric product is equivariant (Appendix A).
The second geometric primitive we use is derived from the so-called join $^7$ $x, y \mapsto (x^* \wedge y^*)^*$ . This map may appear complicated, but it plays a simple role in our architecture: an equivariant map that involves the dual  $x \mapsto x^*$ . Including the dual in an architecture is essential for expressivity: in  $\mathbb{G}_{3,0,1}$ , without any dualization it is impossible to represent even simple functions such as the Euclidean
Technically, the join has an anti-dual, not the dual, in the output. We leave this detail out for notational simplicity. In our network architecture, it makes no difference for equivariance nor for expressivity whether the anti-dual or dual is used, as any equivariant linear layer can transform between the two.
distance between two points [17]; we show this in Appendix A. While the dual itself is not  $\operatorname{Pin}(3,0,1)$ -equivariant (w.r.t.  $\rho$ ), the join operation is equivariant to even (non-mirror) transformations. To make the join equivariant to mirrorings as well, we multiply its output with a pseudoscalar derived from the network inputs:  $x,y,z \mapsto \operatorname{EquiJoin}(x,y;z) = z_{0123}(x^{*} \wedge y^{*})^{*}$ , where  $z_{0123} \in \mathbb{R}$  is the pseudoscalar component of a reference multivector  $z$  (see Appendix B).
We define a geometric bilinear layer that combines the geometric product and the join of the two inputs as  $\operatorname{Geometric}(x, y; z) = \operatorname{Concatenate}_{\mathrm{channels}}(xy, \operatorname{EquiJoin}(x, y; z))$ . In GATr, this layer is included in the MLP.
Nonlinearities and normalization We use scalar-gated GELU nonlinearities [23] Gated  $\mathrm{GELU}(x) = \mathrm{GELU}(x_1)x$ , where  $x_{1}$  is the scalar component of the multivector  $x$ . Moreover, we define an  $\operatorname{E}(3)$ -equivariant LayerNorm operation for multivectors as  $\mathrm{LayerNorm}(x) = x / \sqrt{\mathbb{E}_c\langle x,x\rangle}$ , where the expectation goes over channels and we use the invariant inner product  $\langle \cdot ,\cdot \rangle$  of  $\mathbb{G}_{3,0,1}$ .
Attention Given multivector-valued query, key, and value tensors, each consisting of  $n_i$  items (or tokens) and  $n_c$  channels (key length), we define the E(3)-equivariant multivector attention as
$$
\operatorname {A t t e n t i o n} (q, k, v) _ {i ^ {\prime} c ^ {\prime}} = \sum_ {i} \operatorname {S o f t m a x} _ {i} \left(\frac {\sum_ {c} \langle q _ {i ^ {\prime} c} , k _ {i c} \rangle}{\sqrt {8 n _ {c}}}\right) v _ {i c ^ {\prime}}. \tag {5}
$$
Here the indices  $i, i'$  label items,  $c, c'$  label channels, and  $\langle \cdot, \cdot \rangle$  is the invariant inner product of the geometric algebra. Just as in the original transformer [50], we thus compute scalar attention weights with a scaled dot product; the difference is that we use the inner product of  $\mathbb{G}_{3,0,1}$ , which is the regular  $\mathbb{R}^8$  dot product on 8 of the 16 dimensions, ignoring the dimensions containing  $e_0$ . We compute the attention using highly efficient implementations of conventional dot-product attention [14, 32, 37]. As we will demonstrate later, this allows us to scale GATr to systems with many thousands of tokens. We extend this attention mechanism to multi-head self-attention in the usual way.
3.3 Extensions
Auxiliary scalar representations While multivectors are well-suited to model geometric data, many problems contain non-geometric information as well. Such scalar information may be high-dimensional, for instance in sinusoidal positional encoding schemes. Rather than embedding into the scalar components of the multivectors, we add an auxiliary scalar representation to the hidden states of GATr. Each layer thus has both scalar and multivector inputs and outputs. They have the same batch dimension and item dimension, but may have different number of channels.