Geometric Algebra Transformer
Authors: Johann Brehmer
Pim De Haan
Sönke Behrends
Taco Cohen
Venue: NeurIPS 2023 poster (2023)
Keywords: Geometry
geometric algebra
equivariance
transformer
TL;DR: GATr is a new general-purpose network for geometric data, based on geometric algebra representations + E(3) equivariance + a transformer architecture

Abstract: 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 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.

Geometric Algebra Transformer
Johann Brehmer'
Pim de Haan"
Sönke Behrends
Taco Cohen
Qualcomm Al Research
(jbrehmer, pim, sbehrend, facos)@@ti.quall.
Abstract
Problems involving geometric data arise 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 be applied such wide variety of geometric types while respecting their symmetries. In this paper we introduce the Geometric Algebra Transformer (GATr),
penrrol-purpose architecture for geometric data. GATr represents inputs, outputs, and hidden states in the projective geometric (or Clifford) algebra, which offers efficient 16-dimensional vector-space representation of common geometric objects as well as operators acting on them. GATr is equivariant with respect to E(3). the symmetry group of 3D Euclidean space. As Transformer, GATr is versatile, efficient, and scalable. We demonstrate GATr in problems from n-body modeling wall-shear-stres estimation large arterial meshes robotic motion planning. GATr consistently outperforms both non non-geometric and equivariant baselines in terms of error, data efficiency, and scalability.
Introduction
From molecular dynamics to astrophysics, from material design to robotics, fields across science and engineering deal with geometric data such positions, shapes, orientations, directions. The geometric nature of data provides a rich structure: a notion of common operations between geometric types (computing distances between points, applying rotations orientations, etc.), well-defined behaviour of data under transformations of 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, introduce the Geometric Algebra Transformer (GATr), general-purper 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-dimensioonl space, GATr represents data as multivectors of the projective geometric algebra G3.0.11 Geometric (or Clifford) algebra is an principled yet practical mathematical framework for geometrical computations. The particular algebra G3,0,1 extends the vector space R to 16-dimensional multivectors, which can natively represent various geometric types and 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. Equivariance: GATr is equivariant with respect to E(3), the symmetry group of direediimnnsioal space. To this end, we develop several new E(3) =33)-equuvariaat primitives mapping between multivectors, including equivariant linear maps, an attention mechanism, nonlinearities, and normalization layers.
"qqual contribution.
"Qullcom Al Research is an initiative of Qualcomm Technologies, Inc
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Transformer:
Due to its favorable scaling properties, expressiveness, trainability, and versatility, the Transformer architecture [50] has become the de-facto standard for wide range of problems. GATr is based on the Transformer architecture, in particular on dot-produe 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 to the best of our knowledge the first E(3) equivariant architecture with internal geometric algebra representations" -is versatile network for problems involving geometric data.
We demonstrate GATr in three problems from entirely different fields. In 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 GAT' scales to realistic problems with meshes of thousands of nodes. Finally, experiment with robotic motion planning, using GATr as the backbone an F(3)- invariant diffusion model. In all cases, GATr substantially outperforms both non-geometric and cquivariant baselines.
Our implementation of GATr is available at attps://githhb. www/qualcoooocccoacccccccac aite--aarrrreeeeeeeee geometric-algebra-eeeee algebra-transformm
2 Background
Geometric algebra We begin with brief overview of geometric algebra (GA); for more detail, see e. e.g. Refs. [17, 18, 38, 41]. Whereas plain vector space like R33 allows us to take linear combinations of elements and y (vectors), a geometric algebra additionally has a bilinear associative operation: the geometric product, denoted simply by zy. By multiplying vectors, one obtains so-called multivectors, which can represent both geometrical objects and operators. Like vectors, multivectors have a notion of direction well as magnitude and orientation (sign), and can be linearly combined.
Multivectors be expanded multivector basis, consisting of products of basis vectors. For example, in a 3D GA with orthogonal basis e1, e3, a general multivector takes the form
x x. 61 2262 + T363 + 1126162 31ge1 + 3236263 + 2123616263,
(1)
with real coefficients (x,, 21 T123) E R Thus, similar to how a complex number a + bi is a sum of a real scalar and imaginary number, general multivector is sum of different kinds of elements. These are characterized by their dimensionality (grade), such as scalars (grade 0), vectors e (grade 1), bivectors ejéy (grade 2), all the way up the pseudoscalar e1 ...ed
ea (grade d).
The geometric product is characterized by the fundamental equation ev = (v, v), where (:: is an inner product. In other words, we require that the square of vector is its squared norm. In an orthogonal basis, where (e;, x ōız, one can deduce that the geometric product of two different basis vectors is antisymmetric: e,t; = -e,e.". Since reordering only produces sign flip, we only get one basis multivector per unordered subset of basis vectors, and so the total dimensionality of GA is Eg = ( 2d. Moreover, using bilinearity and the fundamental equation 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 (ar, y) (zy + ya)/2 and x A y = (ry- /22 /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 H x². It acts on basis elements by swapping "empty" and "full" dimensions, e. g. sending 61 + 623-
Concurrently to our work. similar network architecture was studied by Ruhe et al. [40]: comment on similarities and differences in Sec. 5.
4 Indeed the imaginary unit i can he thought of as the bivector 61 in 2D GA.
'The antisymmetry can be derived by using ?² (v, u) to show that esey teje; -- -ef

---

Scalar Vector
Bivector
Trivector
PS C0123
Object/ / operator
co
coi
Ciy
coij C123
Scalar À E R
0
0
Plane w/ normal n. E IR', origin shift d E R RR. orthogonal shift I23
0
d
12
0
0
0
0
0
Line w/ direction
0
0
0
0
Point p E R
0
0
p
1
0
Pscudoscalar la E R
0
0
0
0
0
J4
Reflection through plane w/ normal n E IR3, origin shift df IR
0
d
0
0
0
0
0
Translation E
2
0
0
Rotation expressed quaternion E ue*
0
0
0
0
0
Point reflection through p E TR
Table 1: Embeddings common geometric objects and transformations into projective geometric algebra G3.0,1- The columns show different components the multivectors with the corresponding hasis elements, with j,j E 21, 3).J & i, i.c. ij E {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- diree-dimessonall objects as well as arbitrary rotations and translations acting them, the 3D GA is not enough: 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. IR$) 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 33.0,1 38, 41]. Here one adds a fourth homogeneous coordinate TgCg to the vector space, yielding a 24 16-dimensional geometric algebra. The metric of G is such that ej = 0 and ef = 1 for 1, 2, 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 R3.
Representing transformations In geometric algebra, a vector u can act as an operator, reflecting other elements in the hyperplane orthogonal to Since any orthogonal transformation is equal to sequence reflections, this allows us to express any such transformation as a geometric product of (unit) vectors, called a (unit) versor u U1 24g Furthermore, since the product of unit versors is a unit versor, and unit vectors are their own inverse (u² = 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 G3,0, these are the double cover of E(3) and SE(3), respectively. We can thus represent any rotation, translation, and mirroring the symmetries of three- dimensional space-ar as G3,0,1 multivectors.
In order to apply versor to an arbitrary element x, uses the sandwich product:
Pu(x) (uruu- if u is even
(2)
(uru-1 u is odd
Here & is the grade involution, which flips the sign odd-grade elements such as vectors and trivectors, while leaving even-grade elements unchanged. Equation thus gives linear action ei.e. group representation) of the Pin and Spin groups on the 2d -i-mensiiioal space of multivectors. The sandwich product is grade-preservvin tations on each grade.
this representation splits into direct sum of represen-
Representing 3D objects Following Refs. [17, 38, 41], we represent planes with vectors, and re- quire that the intersection of two geometric objects is given by the wedge product of their representa- tions. Lines (the intersection of two planes) thus represented as bivectors, points (the intersection of three planes) as trivectors. This leads to duality between objects and operators, where objects represented like transformations that leave them invariant. Table I 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 (33) there are two elements of Pin(3, 0, 1), e. both the vector and represent the same reflection.
3

---

Equivariance We construct network layers that are equivariant with respect to E(3), or equivalently its double cover Pin 0, 1). A function f G3.0. - G3.0, 1 is Pin(3, 0, eqquivarian with respect to the representation (or Pin(3, 0, -qquivaaria for short) if
f(p((z)) pu(f(x))
(3)
for any a E Pin(3,0,1) and x E G3.0.1, where Pa(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 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 objects, 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 G3,0,1 1 and use the plane-based representation of geometric structure outlined 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 E(3), including translations. We can thus represent objects that transform arbitrarily under E(3). including with respect to translations of the inputs. This is in stark contrast with most E(3)--quivvaiaat architectures, which only O(3) representations and whose features only transform under rotation and 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. t. the center of vector.
mass or by treating the difference between points as a translation-invane
Many systems will not exhibit the full B(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 E(3) ((3--equivariant architecture and to include symmetry breakknn 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,
break this equivariance when desired (in particular, along time dimensions) through positional embedding.
Like vanilla transformer, GATr is based 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 pesssage-passing algorithms.
Architecture overview We sketch GATr in Fig. 1. In the top row, sketch the overall workflow. necessary, raw inputs are first preprocessed into geometric types. The geometric objects are then embedded into multivectors of the geometric algebra Ga,o, 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

---

Extrout from geometiic dggebra
Emieed geometric cigebra
Preprocess
GATr
nputs
Geometric types
Multivector scular impuss
Muttivecto & scolar outouts
Row oulpuls
Eoui layer norm
Equi oyer nor
GATI
Equi
Equi
Geo
Equi
Equi
Geo
Gated GELU
Equii
Eaui
Overview over the GATr architecture. Boxes with solid lines are Icarnable components, those with dashed lines are fixed.
Figure I:
of equivariant multivector LayerNorm, an equivariant multivector self-attention mechanism, residual connection, another equivariant LayerNorm. an equivariant multivector MLP with geometric bilinear interactions, and another residual connection. The architecture is thus similar to 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 GATr network we extract target variables, again following the mapping given in Tbl. 1.
3.2 GATr primitives
We begin with linear layers between multivectors. Appendix A, we show that equivariance condition of Eq. (3) severely constrains them:
Linear layers
Proposition 1. Any linear map Ø G..u,1 + Ga,001 that is equivariant to Pin(d, 0, I) is of the form
d+1
d
((x)
Σ w + Σ
vxco(o)
k=
for parameters IRd+2 Rd+1, Here (x)x is the blade projection of multivector, which sets all non graderk elements to zero.
Thus, -e3uiiaquv linear maps between G3,0.1 multivectors be parameterized with nine coefficients, five of which are the grade projections and four include a multiplication with the homogeneous basis vector eg. We thus parameterize affine layers between multivector-valued arrays with Eq. (4), with learnable coefficients channel. In addition, there learnable bias for the scalar components of the outputs (biases for the other components are not equivariant).
and e, for each combination of input channel and output
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 the geometric product z, y->+> + zy, 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 equivariant (Appendix A),
The second geometric primitive we use is derived from the so-called join? x. y + (x" A "" This map may appear complicated, but it plays simple role in our architecture: equivariant map that involves the dual x + x². Including the dual an architecture is essential for expressivity: in G3,0.1 without any dualization is impossible represent even simple functions such as the Euclidean
Techhnically, 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, any equivariant linear layer transform between the two.
5

---

distance between two points [17]; we show this in Appendix A. While the dual itself not Pin 0, 1)- equivariant (w. t. p). the join operation is equivariant to even make the join equivariant to mirrorings well, we multiply its output with pseudoscalar derived from the network inputs: z; y, -> EquiJoin(z; y; 2) = 20123(x* y²)", where 20123 E IR the pseudoscalar component of a reference multivector = (see Appendix B).
transformations. To
We define a geometric bilinear layer that combines the geometric product and the join of the two inputs as Geometric(z, y;; = Concatenatedhannell EquiJoin(a (z, yr z)). In GATr, this layer is included the MLP.
Nonlinearities and normalization We use scalar-gated GELU nonlincarities [23] GatedGELU(z)) = GELU(Z)) where ø the scalar component of the multivector æ. More- over, we define an -equivvaval LayerNorm operation for multivectors as LayerNorm(r)) = x/ (1,, where the expectation goes over channels and we use the invariant inner product (s, of G3,0.1.
Attention Given multivector-valued query, key, and value tensors, each consisting of items (or tokens) and n. channels (key length), we define the (33)-qquivvriint multivector attention as
Ln kie) v8n,
Attention(q, k v)re
Σ Softmax,
Vie!
(5)
Here the indices i, i' label items, c, e' label channels, and (:-)) is the invariant inner product of the geometric algebra. Just as in the original transformer [50], we thus compute scalar attention weights with scaled dot product; the difference is that we use the inner product of G3,0,1 which is the regular R dot product on of the 16 dimensions, ignoring the dimensions containing eq. 8 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 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.
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 Eq. (5), and from the auxiliary scalars, using a regular scaled dot- not-product attention. The two attention maps are summed before computing the softmax, and the normalizing factor adapted. In all other layers, the sealar 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-awan dot-produce attention The dot-product attention in Eq. (5) ignores the 8 di- mensions involving the basis element ég- These dimensions vary under translations, and thus their straightforward Euclidean inner product violates equivariance. We can, however, extend the atten- tion 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, non- linear query features ((q) and key features (kk) and extend the attention weights in Eq. (5) as (gre, kiz) (gve: kic) + ((gre) (k()).), adapting the normalization appropriately. We define these nonlinear features in Appendix B.
'We also experimented with attention mechanisms that use the geometric product rather than the dot product. but found a worse performance in practice.

---

Token number generalization
E(3) generallzation
101
MLP
GCA-GNN
10a
Transformer
Mean squared etror
SE(3)-TTT.
SEGNN
10-1
GATT (ours)
-0-2
10
102
10
104
105
102
16
10c
103 10
103
104
105
Training samples
Training samples
Training samples
Figure 2: r--body dynamics experiments. We show the error in predicting future positions of planets as a function of the training datasct size. Out of five independent training runs, the mean and standard error are shown. Left: Evaluating without distribution shift. GATr (-)i is more sample efficient than the equivariant SEGNN [6] (--) and (BE(33)rannformm)) [20] (-)) and clearly outperforms non- con-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 straightforward geometric interpretation. When the trivector components of queries and keys represent 3D points (see Tbl. 1), sb()) ((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 sequence, encode their position using "rotary positional" " 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.)
Experiments
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 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 (3333-raanffomme [20], as well as the geometric-algete based but not equivariant- GCA-GNN [41]. The experiment is described in detail in Appendix C.
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 hoi-equivariant baselines, including the geometric- all--allgeaaaaaaee GCA-GNN. Compared to the equivariant SEGNN and SE(3) Transformer, GATr is more sample-
"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 a rotation of 3D space, but rather how the position in sequence is embedded via sinusoids in the scalar channels of keys and queries. Using positional encoding thus does not affect E(3) equivariance.
SSnne auxiliary scalar representations and multivectors mix in the attention mechanism, the positional embeddings also affect the multivector processing.
7

---

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 the middle and right panels of Fig. 1. When evaluating
larger number of bodies than trained on, methods that attention weights (GATr, Transformer, S-(3)- Transfforrr)) generalize best. Finally, the performance of the Ee3)-qquivariant GATr, SEGNN, and SE(3)- Transsorrrerrmer does not drop when evaluated on spatially translated data, while the non-equivariant baselines fail in this setting.
softmax
4.2 Wall-shear-steess estimation on large meshes of human arteries
Next, we turn towards realistic experiment in-
100
volving more complex geometric objects. We study the prediction of the wall shear stress ex- erted by the blood flow on the arterial wall, an
important predictor of aneurysms. While the
Mean approximmatoon error 1%)
wall shear stress can be computed with computa- tional fluid dynamics, simulating a single artery take many hours, and efficient neural sur- rogates 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).
GATT (ours)
O-
GATT canonical
Transformer
Transformer canonical
10:
PointNet++
We train GATr dataset of arterial meshes and simulated wall shear stress published by Suk et al. [48]. They are compared to Trans- former and to the results reported by Suk et al. [48], including the equivariant GEM-CNN [15] and the non-equivariant PointNet++ [36]." !! See Appendix C.2 for experiment details.
PointNet+ canonical
GEM-CNN
GEM-CNN canonical
160
Number training samples
Figure 3: Arterial wall-shear-stres: estimation [48]. 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 mark- ers) and on a version of the dataset in which all artery meshes are canonically oriented (hollow markers), With- out canonicalization, GATr (-) predicts wall shear stress more precisely and is more sample-efficiet than the baselines.
The results are shown in Fig. 3. On non- canonicalized data, with randomly rotated meshes, GATr improves upon all previous meth- ods and sets new state of the art.
We also experiment with canonicalization: rotat- ing the arteries such that blood always flows in the same direction. This helps the Transformer to be almost competitive with GATr. However, canonicalization is only feasible for relatively straight arteries as in this dataset, not in more complex scenarios with branching and turning arteries. We find likely that GATr will be more robust in such scenarios.
4.3 Robotic planning through invariant diffusion
In our third experiment, show how GATr defines an E(3)- diffusion model, that can be used for model-based reinforcement learning and planning, and that this combination is well-suited to solve robotics problems. We follow Janner al. [27], who propose treat learning world model and planning within that model as a unified generative modeling problem. After training a diffusion model [45] on offline trajectories, one use it in planning loop, sampling from it conditional on the current state, desired future states, or to maximize a given reward, as needed.
We use GATT 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 E(3) S,-in -niinrariant diffusion model [29].
GATr-Diffuser demonstrated on the problem of Kuka robotic gripper stacking blocks using the "unconditional" environment introduced by Janner et al. [27]. We train GATr-Difluser on the offline
Il We also tried to run SEGNN [5] and SE(3)-Transfon Transformer [20] as additional baselines, but were not able to fit them into memory 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 origi- nal Diffuser model and a new Transformer back- bone for the Diffuser model. In addition, we show the published results of Diffuser [27], the equivariant EDGI [7], and the offline RL algo-
GATi-Diffuser (ours)
Transfomer-
Diffuser (reproduced)
EDGI
Diffuser
CQL
BCQ
Reward
rithms CQL [30] and BCQ [21] 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. is also clearly more sample-efficient matching the performance of a Diffuser model or Transformer trained on the full dataset even when training only 1% of the trajectories.
101
102
10a
10f
Training trajectories
Figure 4: Diffusion-basee robotic planning. We show normalized rewards (higher is better) function of training dataset size. GATr (-) is more successful block stacking and more sample-efficient than the baselines, including the original Diffuser [27] (-)) and modification of it based a Transformer (--). In grey, show results reported by Brehmer et al. [7] and Janner et al. [27].
4.4 Scaling
Finally, we study GATr's computational require- ments and scaling. We measure the memory usage and compute time of forward and back- ward passes on synthetic data a function of the number of items. GATr is compared to Transformer, SEGNN [6] in the official implementation, and the 3E(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 dif- ferences between the architectures, is impossible 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. For few tokens, GATr slower than 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
when pre-compiling the computation graph.
GPU memory m usage [GB]
For larger problems, compute and memory are dominated by the pairwise interactions the at-
10._
tention mechanism. Here GATr and the base- line Transformer perform on par, as they use the same efficient implementation of dot dot-product at- tention. 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 guadratically, but the equivari- ant baselines have 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.
10%
10-1
10a
10a-
per
GATr (ours)
Transformer SE(3)-THf.
10-1_
SEGNN
10a
10a
10a
5 Related work
Number of items
Figure 5: Computational cost and scaling. We measure peak GPU memory usage (top) and wall time (bottom) per combined forward and backward pass as func- tion of the number of items in synthetic data. Despite some compute overhead, GATr (-) scales just like the Transformer (--) and orders of magnitude more favor- ably than the equívariant baselines (-.--),
Geometric algebra Geometric (or Clifford) algebra was first conceived in the 19th century [10, 22] and has been used widely in quantum physies [16, 34]. Recently, has found new pop- ularity in computer graphics [18]. In particular,

---

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 machine learning in various ways. Spellings [46] use G3 geometric products compute notation-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 al. [41]. These works also use multivector features (the latter even of the same C3,0,1), and process them with operations such as the geometrie/ / 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 experiments.
Concurrently to this publication, Ruhe et [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 Gy.g.r. while we focus on the projective algebra G3,0,1 with its faithful E33 representations, with our theory also applicable to the group E'n) and the algebras C,000 and C,.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-passiingg graph network, while our GATr is Transformer.
Equivariant deep learning Equivariance to symmetries [11] is the primary design principle in modern geometrie deep learning [8]. Equivariant networks have been applied successfully in areas such as medical imaging [3 53] and robotics [7, 26, 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 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 S0(3) equivariance.
Discussion
We introduced the Geometric Algebra Transformer (GATr), general-purpose architecture for geo- metric data, and implemented at https: //gi thub. com/ cww/qualccccccccccccccccccccccca comm- ai rew.archeeeeeeeeeeee research/geometricc geometricc al algebra- sebbr-ttransforeee 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. Empiri- cally, GATr outperforms non-gev non-gcometric, equivariant, and geometric algebra- non-equivariant baselines across three experiments.
At the same time, GATr scales much better than most geometric networks. This 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 understanding the details of the method. However, given 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 universal approximator, which is an interesting direction for future work.
Given the promising results presented in this work, look forward further study the potential of Geometric Algebra Transformers in problems from molecular dynamics to robotics.
10

---

Acknowledgements We would like to thank Joey Bose, Johannes Brandstetter, Gabriele Cesa, Steven De Keninck, Daniel Dijkman, Leo Dorst, Mario Geiger, Jonas Köhler, Ekdeep Singh Lubana, Evgeny Mironov, and David Ruhe for generous guidance regarding geometry, enthusiastic encourage- ment on equivariance, and careful counsel concerning computing complications.
References
[1] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in Neural Information Processing Systems, 32, 2019. (Cited on page 10)
[2] Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv:1809.108: 10853, 2018. (Cited on pages and 22)
[3] lyes Batatia, David Peter Kovacs, Gregor N Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Advances in Neural Information Processing Systems, 2022. (Cited on page 10)
[4] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nat. Commun., 13(1):2453, May 2022. (Cited on page 10)
[5] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for PDE modeling. NKKY::22090093344 2022. (Cited on pages 2, 8, and 10)
[6] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve E(3) equivariant message passing. In International Conference on Learning Representations, 2022. (Cited on pages 7, 9, 10, 22, and 25)
[7] Johann Brehmer, Joey Bose, Pim De Haan, and Taco Cohen. EDGI: Equivariant diffusion for planning with embodied agents. In Advances in Neural Information Processing Systems, volume 37, 2023. (Cited on pages 9, 10, and 24)
[8] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. 2021. (Cited on page 10)
[9] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. /604.0660 74, 2016. (Cited on page 23)
[10] William Kingdon Clifford. Applications of Grassmann's Extensive Algebra. Amer. J. Math., (4):350-358, 1878. (Cited page 9)
[11] Taco Cohen. Equivariant Convolutional Networks. PhD thesis, University of Amsterdam, 2021. (Cited on page 10)
[12] Taco Cohen and Max Welling, Group equivariant convolutional networks. In International Conference on Machine Learning, pages 2990-2999, PMLR, 2016. (Cited on page 2)
[13] Erwin Coumans and Yunfei Bai. PyBullet, a Python module for physics simulation for games, robotics and machine learning. http: httt::/ppbullet..om org, 2016-2019. (Cited on page 23)
[14] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-effiiient exact attention with I0 IO-awareness. Advances in Neural Information Processing Systems, 35:16344-16335 2022. (Cited on pages and 6)
[15] Pim De Haan, Maurice Weiler, Taco Cohen, and Max Welling. Gauge equivariant mesh CNNs: Anisotropic convolutions on geometric graphs. International Conference on Learning Representutions, 2021. (Cited pages and 23)
[16] Doran and A Lasenby. Geometric algebra for physicists. Cambridge University Press, 2003.
(Cited on page 9)
[17] Leo Dorst. A guided tour to the plane-based geometric algebra pga. 2020. URL https: // http://arr..rrr//ooomaaa////AAA pdf. (Cited pages 2, 3, 6, 10, and 19)
11

---

[18] Leo Dorst, Daniel Fontijne, and Stephen Mann. Geometric Algebra for Computer Science: An Object-oriented Approach to Geometry. Morgan Kaufmann Series Computer Graphics. Morgan Kaufmann, Amsterdam, 2007. ISBN 978-0-12-369465-2 (Cited on pages 2, 9, 10, and 15)
[19] Thorben Frank, Oliver Thorsten Unke, and Klaus Robert Muller. So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems. October 2022. (Cited on page 10)
[20] Fabian B Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SB(3)- 3D Rots-Trannlttion equivariant attention networks. In Advances in Neural Information Processing Systems, 2020. (Cited on pages 8, 9, 10, and 22)
[21] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062. PMLR, 2019. (Cited on pages 9 and 24)
[22] Hermann Grassmann. Die lineale Ausdehnungslehre. Otto Wigand, Leipzig, 1844. (Cited on page 9)
[23] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). 100.v:1006.0881155 2016. (Cited on pages 6 and 22)
[24] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv: /2/80 [cs], December 2019. (Cited on page 7)
[25] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion prohabilistic models. In Neural Information Processing Systems, 2020. (Cited on page 24)
[26] Haojie Huang, Dian Wang, Robin Walters, and Robert Platt. Equivariant transporter network. In Proceedings of Robotics: Science and Systems, 2022. (Cited on page 10)
[27] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning. 2022. (Cited on pages 8, 9, 23, and 24)
[28] John Jumper, Richard Eyans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588 2021. (Cited on page 10)
[29] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR, 2020. (Cited page 8)
[30] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine, Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179- 1 1191, 2020. (Cited on pages and 24)
[31] Maxime W Lafarge, Erik J Bekkers, Josien P W Pluim, Remco Duits, and Mitko Veta. Roto- translation equivariant convolutional networks: Application to histopathology image analysis. Med. Image Anal., 68, 2021. (Cited page 10)
[32] Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xFormers: A modular and hackable Transformer modelling library. https://githh com/ facebokrresearch/ffooo
2022. (Cited on pages 23, and 24)
[33] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXi::200.111990 2022. (Cited page 10)
[34] Pertti Lounesto. Clifford Algebras and Spinors. London Mathematical Society Lecture Note. Cambridge University Press, 2001. (Cited on page
12

---

[35] Alexandre Milesi.
Accelerating open- open-source model implementation.
training using an NVIDIA https://developer: avidii.....iiiia............iai. nwidia. www.ia.coo/oooooo.o.ooooooo. accelerating- -ew--eeee-tttt-e--ee----eeee-ee ormers- www.eoootrrrraeeeeeeeeeeeeeeeeee saraining-usinee
vodel-impleeennatiion 2021. (Cited on pages and 25)
[36] Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++: Deep hierarchical feature learning point sets in metric space. Advances in Neural Information Processing Systems, 30, 2017. (Cited on pages 8 and 23)
[37] Markus N Rabe and Charles Staats. Self-attention does not need (n²²) memory. arXiv:2112.05682, 2021. (Cited on pages and 6)
[38] Martin Roelfs and Steven De Keninck. Graded symmetry groups: plane and simple. arXiv:2107.03771, 2021. (Cited on pages 2, 3, 10, and 15)
[39] David W Romero, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogendoom. Attentive group equivariant convolutional networks. In Proceedings of the 37th International Conference on Machine Learning, volume 119, pages 81 3188-8199, JMLR.org, July 2020. (Cited on page 10)
[40] David Ruhe, Johannes Brandstetter, and Patrick Forré. Clifford group equivariant neural networks. Advances in Neural Information Processing Systems, volume 37, 2023. (Cited on pages and 10)
[41] David Ruhe, Jayesh K Gupta, Steven de Keninck, Max Welling, and Johannes Brandstetter. Geometric clifford algebra networks. In International Conference on Machine Learning, 2023. (Cited on pages 2, 3,7, 7, 10, and 22)
[42] Víctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9323-9332, PMLR, 2021. (Cited on page 10)
[43] Noam Shazeer. Fast transformer decoding: One write-head is all you need. 2019. (Cited on page 21)
[44] Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B Tenenbaum, Alberto Rodriguez, Pulkit Agrawal, and Vincent Sitzmann. Neural descriptor fields: SE(3)-Equivariant object representations for manipulation. In ICRA, 2022. (Cited on page 10)
[45] Jascha Sohl-Dicke Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2266622 2265. PMLR, 2015. (Cited on page 8)
[46] Matthew Spellings.
Geometric algebra attention networks for small point clouds. 10.02391 2021. (Cited on pages and 10)
[47] Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, arXiv:2104. 09864, 2021. (Cited on page 7)
[48] Julian Suk, Pim de Haan, Phillip Lippe, Christoph Brune, and Jelmer M Wolterink. Mesh neural networks for (3)-equivariant hemodynamics estimation the artery wall. arXiv:2212.05002, 2022. (Cited on pages 8, 22, and 23)
[49] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and raanslatoo-cquiivaiian neural networks for 3d point clouds. arXiv: arXiv:/802.088. /802.08299 2018. (Cited on page 10)
[50] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- fion Processing Systems, volume 30, 2017. (Cited on pages 4, 5, and 6)
[51] Dian Wang, Robin Walters, Xupeng Zhu, and Robert Platt. Equivariant Q learning in spatial action spaces, 2021. (Cited on page 10)
13

---

[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 scans with equivariant CNNs. MIA, 2019. (Cited 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 the transformer architecture. In International Conference on Machine Learning, pages pages and 22)
PMLR, 2020. (Cited on
[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)
14

---

A Theoretical results
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 ã (-1)xz. Note that this an algebra automorphism zy zy, and also an A-algebra automorphism. The reversal in a linear involutive bijection : Gn,orr R0,,0 which sends a k-blade 21 A../A...AA.r to the reverse;
linear involutive bijection : Gx,o,- : Gn,o.r which sends k-blade
rgA...AAAA21 + with +- if l: E (0, 1,4, 5,,.,,800. ....
and = otherwise. Note that this is an anti- anti-automorphism (contravariant functor): zy yz.
Here we denote the sandwich action of u Pin(n, 0, multivector x not as (r.( but as u(zr]. For odd u, u|æ]] uāu- while for even u, u[x] uou-1 The sandwich action is linear by linearity of the and bilinearity of the geometric product. Furthermore, note that for any particular u E Pin(n, 0, r). the action is geometric algebra homomorphism: u(ab) uabu wău uubu = u[a] v (6]. By linearity and a symmetrization argument [18, Sec 7.1J. one can show that it also A-algebra homomorphism (outermorphism): u(a b) usa] A u[b].
Let/ 2 k. Given a h-vector a and (-vector b, define the left contraction as a == {ab) which is a I - k-vector. For k = and b a blade b = b1 A A bt. Geometrically, a b is the projection of a to the space spanned by the vectors b,. Thus we have that a-b= 0 > Vi, (a, b,) = 0 [18, Sec 3.2.3], in which their inner product is zero. Futhermore, we define a vector a to be tangential to blade if a A = O.
define and ò to be orthogonal. In particular, two vectors a, b are orthogonal if
In the projective algebra, blade is defined to be ideal if it can be written blade y.
eg A y for another
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 Gn,0,0 and then also in projective geometric algebra G,,0.1.
Proposition The grade projection (-)x is equivariant [18, Sec 13.2.3).
Proof. Choose an (-blade x A a2 A A Let u be a 1-versor. As the action u is an outermorphism, u[æ] u[a1] ... u[a]] is an I-blade. Now if & then (r)x = 0 and thus u[(x)))] = (u(==)) If = k, then (I)x and thus u[(x)x] = (u[=]) As the grade projection is linear, cquivariance extends to any multivector.
Proposition The following map is equivariant: & G3.0.1 + G3.0,1 : + eg..
Proof. Let u be 1-versor, then u acts on a multivector as + u[x] uău where è is the grade involution. Note that eo is invariant: u/eul = -ueuu = -guuu
where ueg - Cout because -eg u. Then is equivariant, as the action is an algebra homomorphism: u[o(x)] = u/eoz] = uegzu uēpu "uuu (eo) [xxr] = eoufat (u((x)). It follows that @ also equivariant to any product of vectors, i.e. any versor
and eg are orthogonal: uey (u, co) + eg = -co
Euclidean geometric algebra Before constructing the most general equivariant linear map between multivectors projective geometric algebra, we begin with the Euclidean case Cn,0,0-
Theorem (Cartan-Dieuodonn Dieuodonné). Every orthogonal transformation of an n-dimensional space can be decomposed into at most n reflections in hyperplanes.
Proof. This theorem is proven Roelfs and De Keninck [38].
Lemma 1. In the n-dimensional Euclidean geometric algebra Gn,o,o. the group Pin(n, 0) acts transitively on the space of k-blades of norm X R>0
15

---

Proof. As the Pin group preserves norm, choose À = 1 without loss of generality. Any k- of unit norm can be written by Gram-Schmin factorization as the wedge product of orthogonal vectors of unit norm x = 01 A A ... A vy- Consider another k-blade y = w1 A W2 A A wy with uy orthonormal. We'll construct u E Pin(n, 0, 0) such that u [xr] yy-
Choose n k additional orthonormal vectors Uk+1,...., Un and wk+1,... Wy to form orthonormal bases. Then, there exists a unique orthogonal transformation E IK" that maps v; into w, for all i E {I,..., n.. By the Cartan-Dieuodonné theorem 1, this orthogonal transformation can be expressed as the product of reflections, thus there exists u Pin(n, 0, 0) such that u(0t] = w4- As the action is A --algebra homomorphism (u(a A b] u(al A u(b], for any multivectors b), we have that u [r] y.
Lemma 2. In the Euclidean fr = 0) or projective (r = 1) geometric algebra Gn,u.r let x be a k- u x 0 and u|r)
blade. Let u be I-versor. Then u |
+
u uA
Proof. Let x be a k-blade and u a vector of unit norm. We can decompose u into u = + v with tAz 0 (the part tangential to the subspace of r)) and v u|z = 0 (the normal part). This decomposition is unique unless r is ideal in the projective GA, in which case the Eo component of u is both normal and tangential, and we choose Euclidean.
In either case, note the following equalities: at (-1)4- xv = (-)) of = to and note #d * O such that vfr = Ax, which can be shown e.g. by picking a basis. Then:
u[x] (-1)a(t (t + 0) v)t v) (t + c- u)z 1=11 ||012) 2utæ
We have u[xx] x x + vl. 0. If x is not ideal, this implies that either v 0 (thus u A 0 and u(x] = -x) or f = 0 (thus u/ z = 0 and [x]] = r)). If x is ideal, this implies that either oc eo (thus u A r = 0 and u u[x]] r) or t O (thus u| x = O and u|æ]] = 2).
Lemma 3. Let r E (0, 13. Any linear Pin(n, 0, r--quivaaria map é Ga,o, + Gn.o.r can be decomposed into a sum of equivariant maps & = Likm Olkm with Ulm equivariantly mapping k- blades to l-blades. If = 0 (Euclidean algebra) or k < n + 1, such map oue is defined by the image of any one non-ideal k-blade, like e Instead, if = 1 (projective algebra) and k = n 1. then such map is defined by the image of pseudoscalar, like coi....
Proof. The (in 0, r) group action maps fe-vccoor to k-vectors. Therefore, can be decomposed
into equivariant maps from grade & to grade 1: ((x)) = Eix ()(((=))) with ΦI having I-vectors as image, and all k--vectors in the kernel, for k/' * k. Let be non-ideal k-blade (or pseudoscalar if
= n + 1). By lemmas and 4, in both Euclidean and projective GA, the span of the k-vectors in the orbit of contains any v-vector. So diR, is defined by the I-vector y = @u.(x) Any l-vector can
be decomposed as a finite sum of (-blades: = 311 + J.. We can define @leem (I) = you extended to all l-vectors by equivariance, and note that Φn Em Ólim.
Proposition 4, For an n-dimensional Euclidean geometric algebra Gn,0.0- any linear endomorphism 6: : Gn,0,0 + Gn,0,0 that is equivariant to the Pin(n, 0,0,,,000) group (equivalently O(n) is of the n
type 6(x) = =E k= Wx wx(z)), k. for parameters
go+1
Proof. By decomposition of Lemma 3, let ó map from k-blades to I-blades. Let x be k-blade. Let u be a 1-versor. By Lemma 2, if u is orthogonal to x, u[o(z)] = d(u(z)) = d(z) and u is also orthogonal to ((r). If ufr 0. then u[o(x)) = ((u(z))
((-z)
((x) and u//(2z) = 0. Thus
any vector in x is in d(x) and any vector orthogonal to x is orthogonal to o(x), this implies o(z)
WK for some wi E RR. By Lemma 3, we can extend Ø to d(y) = way for any kevector y.
Projective geometric algebra How about equivariant linear maps in projective geometric algebra? The degenerate metric makes the derivation more involved, but the end we will arrive at a result that is only slightly more general.
Lemma 4. The Pin group of the projective geometric algebra, Pin(n, O, 1), acts transitively on the space of k-blades with positive norm ||>||| = A O. Additionally, the group acts transitively on the space of zero-norm k-blades of the form x eo A y (called ideal blades), with ||y|| K.
16

---

Proof. Let 2 21 A... A my be k-blade with positive norm X. All vectors 2, can be written as I; = 2/ + 6,e0. for a nonzero Euclidean vector i; (meaning with no ep component) and 6; R. because if 142 = 0. the norm of æ would have been O. Orthogonalize them as 4r2 = 32 (1), 11 etc., resulting in x = z A ... A x with r/ = o/ + Eee0 with orthogonal r/-
Define the translation t + 2 E; 8:cu A vy. which makes x' Euclidean: t|x²|| = e/ A A.. v/6- By Lemma 1, the Euclidean Pin group Pin(n, 0, 0), which is a subgroup of (in 0, 1), acts transitively on Euclidean k-blades of given norm. Thus, in the projective geometric algebra Pin(n, 0, 1), any two k--lades of equal positive norm À are related by a translation to the origin and then a Pin(n, 0, transformation.
For the ideal blades, let x = en A y. with |lll| = k. We take y to be Euclidean without loss of generality, For any g E Pin(n, O, 1), gleo] co. g|a]] = eo A gly]. Consider another x' Eo A y with |Iv'll = k and taking y' Euclidean. As Pin(n, 0,0) 0) acts transitively on Euclidean (A- - 1)-blades with norm K, let g (in 0,0) 0) such that glvl y'. Then g|ø]
We can now construct the most general equivariant linear map between projective geometrie algebras, key ingredient for GATr:
Proposition 5. For the projective geometric algebra G.,o.,, any linear endomorphism ø G..o. 7 Gn.0.1 that is equivariant to the group Pin(n, 0, r) (equivalently to E(n)) is of the type ((x) I K= wp ( + Lk=o U)ea(T) k. for parameters w E RR++2 E R=+11
Proof. Following Lemma 3, decompose & into linear equivariant map from k-blades I-blades. For k < n + 1, let = e12... Then following Lemma 2, for any 1 S. i < k. e, Ax = 0, e = =x, and ex
-0(x) and thus e; x(x) 0. Therefore, we can write <(x) = x A /1 A A Wk for / k vectors yj orthogonal to x.
(e([=) Φ(-2)
Also, again using Lemma 2, for ls < i is n. e 0 = c o(x) = e c 0 = Vi, (e, Uj) = 0. Thus, yj is orthogonal to all e, with 1 i n. Hence, I= k or = =k +1
and y x E0.
For = n + 1, x = k. By a similar argument, all invertible vectors u tangent to x must be tangent to (xx), thus we find that e(x)) = x A y for some blade y. For any non-zero o(x), x 1, and thus (x) x x. By Lemma 3, by equivariance and linearity, this fully defines Φ.
A.2 Bilinear maps
Next, we turn towards bilinear operations. In particular, show that the geometric product and the join are equivariant.
For the geometric product, equivariance is straightforward: Any transformation gives a homomorphism of the geometric algebra, as for any multivectors x. y. u u[zyy] = uzzu uzyu" uZu -uju- = ufx]u[y). The geometric product is thus equivariant.
Pin r),
Dual and join in Euclidean algebra For the join and the closely related dual, we again begin with the Euclidean geometric algebra, before turning to the projective case later.
The role of the dual is to have a bijection :
Gn,0,0 + Gn,0.. that maps veveccors (n -e-vectors. For the Euclidean algebra, with a choice of pseudoscalar I, we can define a dual as:
x²' -=xI--=a .t
(6)
This dual is bijective, and involutive up to a sign: (1, 5, 8, ...} and yy for n E (2, 6, 7,.... We choose I instead of I in the definition of the dual that given vectors F1, Tn. the dual of the multivector scalar of the oriented volume spanned by the vector. We denote the inverse of the dual as x- = xI. Expressed in basis, the dual yields the complementary indices and sign. For example, for 3 and I = 1223, we have (e1)* = -223, (e12)"
= Hy, with +y = for n e
21 A ....n. is given by the
Via the dual, we can define the bilinear join operation, for multivectors x, y:
x V y (x" ^g") ((II) A (yī))I.
17

---

Lemma's 5. In Euclidean algebra G,0000 the join is Spin(n, 0, 0) equivariant. Furthermore, it is Pin(n, O, equivariant if and only if is even.
Proof. The join is equivariant to the transformations from the group Spin(n, O, 0), which consists of the product of an even amount of unit vectors, because such transformations leave the pseudoscalar invariant, and the operation consists otherwise of equivariant geometric and wedge products.
However, let 612... I E Pin(n, 0, 0) be the point reflection, which negates vectors of odd grades by the grade involution: T[z] 2. Let x be a k-vector and an (-vector. Then x V is a vector of (and zero if k + < n). Given that the join is bilinear, the inputs transform (-1)*+/ under the point reflection, while the transformed output gets sign (-1)kti-n, Thus for odd the join is not Pin(n, 0, 0) equivariant.
grade ((n k) + (n - ()) =
To address this, given a pseudoscalar XI, we can create an equivariant Euclidean join via: EquiJoin(z,, y, XI) V v) X(x Ay")
Proposition 6. In Euclidean algebra Gn,0,0, the equivariant join EquiJoin is Pin(n, 0, 0) equivari- ant.
Proof. The EquiJoin is multilinear operation, so for k-vector and I-vector y. under point reflection, the input gets a sign (-1)&+4ffn while the output is still a k + I - n vector and gets sign (-1) These signs differ by even (-1)2n = 1 and thus EquiJoin is Pin(n, 0, --equivvariant.
We prove two equalities the Euclidean join which we use later.
Lemma 6. In the algebra Gn,a,,, let be a vector and x, y be multivectors. Then
w y) (u]x) Vy x ( (-1)"
(8)
and
Proof.
For the first statement, let a be k-vector and an I-vector. Then note the following two identities:
a V b = (a bī) 2n (a (2- (a" b)) = =a'b,
(uja)* (va) 1 (vaT) -vva")=~++ )- =a())
Combining these and the associativity of gives:
For the second statement, swapping k-vector and l-vector b incurs V (a" -b)
(-1)(m=-)(n-0( A -a*)= (bV a). Then get:
a a (ujb) (- (m- (-)))))=--000000---------- (0
(- (-1)
(-
= 1))=
(-I) Vb
(-I)--
-(1-)" vb.
This generalizes to multivectors x, by linearity.
18

---

Dual and join in projective algebra For the projective algebra Gn,00 with its degenerate inner product, the dual definition of Eq. 6 unfortunately does not yield a bijective dual. For example. epco12... O. For bijective dual that yields the complementary indices on basis elements, different definition is needed. Following Dorst [17], we use the right complement. This involves choosing an orthogonal basis and then for basis k-vector define the dual x² be the basis + - k-vector such that A x² = I, for pseudoscalar I = c012...n For example, this gives dual 601 623, so that 6g1 A 623 C0123.
This dual is still easy to compute numerically, but it can no longer be constructed solely from operations available to us in the geometric algebra. This makes it more difficult reason about equivariance. Proposition the algebra G,,,1 the join V (a* A b is equivariamt to Spin(n, (, 0, 1).
Proof. Even though the dual is not 6n,0. operation, we express the join in the algebra as follows. We decompose k-vector as æ = to + eop> into a Euclidean k-vector to and a Euclidean (k- 1)-vector Ps. Then Dorst [17, Eq (35)] computes the following expression
(lx Cops) V (1y + eoPy) = ((tz + eapx)) A (ly + )oPy)))
= 1, V Py + (-1)" P> VFace by + ea VFac Py),
(10)
where the Euclidean join of vectors a, in the projective algebra is defined to equal the join of the corresponding vectors in the Euclidean algebra:
VFa b = ((a (b (be12...)))12....
The operation a VBae 6 is Spin(n, 0, 0) equivariant, as discussed in Lemma 5. For any rotation E Spin(n, 0, 1) (which Euclidean), we thus have r|a V'ue r|a] Vue r(6]. This makes the PGA dual in Eq. (10) equivariant to the rotational subgroup Spin 0, 0) c Spin(n, 0)
We also need to show equivariance to translations. Let be Euclidean vector and translation. Translations act by shifting with e, times a contraction: T[æ] = x cu(u) x). This acts on the decomposed æ in the following way: 7[tz coPa] r[t,] + coP= t, eo co v vstz).
cbv/2 2
We thus get:
[x0] V rly]
= [ (
(lz + colp= - w(()) V (1, + co(D, - v((p))
V yy-d - to V (u]by) - (-1)"oJto Veuc by eo (px Vau (ujty) + (u]tz) VEue Py) x y eo (px V (w]4y) (w]4x) V Pn)
Used (10) & linearity
Used (9)
x y eo (-(-1)"u Pr VEuc by + (u]tz) Vuuc Ps) y co ((-1)"(v (v]P²) V by (w)(k) VEae Py) = I y - e (v) ((1))"p. Veuc ty + to VEuc pyl)
Used (9)
Used (8)
T[x vl
The join is thus equivariant 12 translations and rotations and is therefore Spin(n,, 0, equivariant.
Similar the Euclidean case, obtain full Pin(n, 0, 1) equivariance via multiplication with pseudoscalar. We thus also use the EquiJoin from Eq. (7) in the projective case.
A.3 Expressivity
As also noted in Ref. [17], in the projective algebra, the geometric product itself is unable to compute many quantities. It is thus insufficient build expressive networks. This follows from the fact that the geometric product preserves norms.
The authors agree with the reader that there must be an easier way prove this.
19

---

Lemma 7. For the algebra Gn, for multivectors x, y. we have ||ayil
||a|| [vll-
Proof. | = Tyzy = zyÿē = x -1ly||2 x: |||| ||
Hence, any null vector in the algebra can never be mapped to non-null vector, including scalars. The projective algebra can have substantial information encoded as null vector, such as the position of points. This information can never influence scalars or null vectors. For example, there is no way to compute the distance (a scalar) between points just using the projective algebra. In the GATr architecture, the input to the MIPs that operate on the scalars, or the attention weights, thus could not be affected by the null information, had we only used the geometric product multivectors.
To address this limitation,
besides the geometric product also the join. The join able to compute such quantities, For example, given the Euclidean vector 1 we can map a null vector
Cu12...k to non-null vector 612... X 6
B
Architecture
Equivariant join One of the primitives GATr is the equivariant join EquiJoin(z; y; x), which hidden states of the neural network after the previous layer. The nature of = is different: it is a reference multivector and only necessary to ensure that the function correctly changes sign under mirrorings of the inputs. We find it beneficial to choose this reference multivector z based on the input data rather than the hidden representations, and choose it as the mean of all inputs the network.
define in Eq. (7). For and y,
Auxiliary scalars In addition to multivector representations, GATr supports auxiliary scalar repre- sentations, for instance describe non-geometric side information such as positional encodings or diffusion time embeddings, In most layers, these scalar variables are processed like in standard Transformer, with two exceptions. linear layers, we allow for the scalar components of multivectors and the auxiliary scalars to freely mix. In the attention operation, we compute attention weights
Eolape MV LMV + E. Po
Softmax;
(11)
VS +
where qth MV and MV are query and key multivector representations, 9d and k are query and key scalar representations, nar is the number of multivector channels, and n, is the number of sealar channels.
Distance-aware dot-product attention we argue Sec. 3.3, it can be beneficial to extend queries and keys with nonlinear features. We use the following choice:
Pa Ei
Σ, K 22
(()) = w(no)
9i 1u1
and ((k) = w(k)o)
2h h
with w(x) = x22
x
Qi 4/2
26,0 6
9,00 9/3/
20, K33
Here the index \i denotes the trivector component with all indices but i. Then
(()) (k:)
7 0,0 2
where we use the shorthand ū (I) 322 " 3)) When the trivector components of queries and keys represent 3D points, with q:u = ky = 1, it becomes proportional to the pairwise negative squared Euclidean distance between the points.
In this notation, trivector (g.o. g) transforms under rotation R E O(3) and translation i'e R
20

---

as (9:0: Ø) + 19:0. RQT q,ot). is easy to see that the inner product is invariant:
b(a) - (k)) + -(no)w (Ri (.od) 9.o (RK + ko
9:0 k) + grok -
= (k(g)))((k))))))))) o 9. k
= =(g) (k:).
The normalization function w could have been chosen as ((t = 1, but this would make the attention logit fourth-order polynomial of the keys and queries, which tends to lead to instabilities when taking the exponential in the softmax. The choice above makes the attention logit a quadratic function for larger values of the keys and queries, just like regular dot- pro--roouct attention.
For the interested reader, we' dd like to note that the above inner product is closely related to the inner product in another geometric algebra, the conformal geometric algebra. We plan exploring this connection in future work.
With this final piece, GATr's attention mechanism computes attention weights from three sources: the 1 i inner product the multivector queries and keys (q, kj. the distance-aware inner product of the nonlinear features ((4) (k),, and the Euclidean inner product of the auxiliary scalars ds k.
We find it beneficial to add learnable weights prefactors each of these three terms. The attention weights are then given by
(A En
MV LMV, Nic
E. dlg (KLM + Ee Fich
Softmax;
V13nNV + n2
with learnable, head-specific ,,, 7,
This may look complicated, but all terms can be summarized in a single Euclidean dot product between query features and key features. We can therefore use efficient implementations of dot- product attention compute GATr's attention.
Multi-query attention To reduce memory use, we also consider version of GATr that uses multi- query attention [43] instead of multi-head attention, sharing the keys and values among attention heads.
C Experiments
C.1
n-body dynamics
Dataset sample:
We generate a synthetic n-body dataset for objects by following these steps for each
1. The masses n objects are sampled from log-uniform distributions. For one object (the star), use mo [1, 10]; for the remaining objects (the planets), (Following common practice in theoretical physics, we use dimensionless quantities such that the gravitational constant is 1.)
m, [0.01, 0.1].
2. The initial positions of all bodies are sampled. We first use a heliocentric reference frame. Here the initial positions of all bodies are sampled. The star is set to the origin, while the planets are sampled uniformly on a plane within a distance r E [0.1, 1.0) from the star.
3. The initial velocities are sampled. In the heliocentric reference frame, the star is at rest. The planet velocities are determined by computing the velocity of a stable circular orbit corresponding the initial positions and masses, and then adding isotropic Gaussian noise (with standard deviation 0.01) to it.
4. We transform the positions and velocities from the heliocentric reference frame to a global reference frame by applying a random translation and rotation to it. The translation is sampled from a multivariate Gaussian with standard deviation 20 and zero mean (except for the domain generalization evaluation set, where we use mean of (200, 0, 0)²). The rotation sampled from the Haar measure on S0(3). addition, we apply random permutation of the bodies.
21

---

5. We compute the final state of the system by evolving it under Newton's equations of motion, using Euler's method and 100 time steps with time interval of 10 each.
6. Finally, samples in which any bodies have traveled more than distance of 2 (the diamater of the solar system) are rejected. (Otherwise, rare gravitational slingshot effects dominate the regression loss and all methods become unreliable.)
We generate training datasets with with n. 4 and 5000 samples; a regular evaluation set with n = 4 and 5000 samples; a number- generalization evaluation set with an additional translation (see step 4 above), and 5000 samples.
and between 100 and 10 samples; validation dataset
and 5000 samples; and E(3) generalization set with n
All models are tasked with predicting the final object positions given the initial positions, initial velocities, and masses.
Models For the GATr model, we embed object masses as scalars, positions as trivectors, and velocities (like translation vectors) as bivectors. We use 10 attention blocks, 16 multivector and 128 scalar channels, and 8 attention heads, resulting in 1.9 million parameters. We use multi-head attention and do not use the distance- attention mechanism (as it led to similar results in short test).
For the Transformer baseline, we follow pre-layer normalization [2, 54] architecture with GELU activations [23] in the MLP block. We use 10 attention blocks, 384 channels, and 8 attention heads, resulting in 1 1.8 million parameters.
We also compare to a simple MLP with GELU actívations. We test models with 2, 5, and 10 layers, each with 384 channels. We report results for the best-performing 2-layer MLP, which has 0.2 million parameters.
Next, compare GATr to the crometri-allgebra-buscee but not equivariant, GCA-GNN [41]. We follow the hyperparameters reported by Ruhe et al. [41] for their Tetris experiment, arriving at an architecture with 1.5 million parameters. We also experiment with the GCA-MLP architecture [41], but find worse results.
Finally, we compare to two equivariant baselines: SEGNN [6] and the (rrn)-Tranformerr [20]. In both cases we use the code released by the authors and the hyperparameters they recommend for (slightly different) n--oody experiments. This leads to models with 0.1 (SEGNN) and 0.3 (SE(3)- Transformer) million parameters. For SEGNN, we vary the number of nearest neighbours between and the number of objects in the scene (corresponding a fully connected graph) and show the best result.
Training All models are trained by minimizing L2 loss on the final position of objects. We train for 50 000 steps with the Adam optimizer, using batch size of 64 and exponentially decaying the learning rate from 3 10 to 3 10-6
C.2 Arterial wall-shear-stress estimation
Dataset We use the single-artery wall-shear-stres dataset published by Suk al. [48]. consists of 2000 meshes of human arteries, of which 1600 are used for training the remaining for validation and evaluation. Each mesh has around 7000 nodes.
We experiment both with randomly rotated meshes as well as canonicalized version, in which all meshes point in the same direction.
Models For the GATr model, we embed object the positions of the mesh nodes as trivectors, the mesh surface normals as vectors, and the geodesic distance to the inlet as scalars. We use multi-query attention, the distance-aware attention mechanism, and learnable prefactors for the attention weights. Our model has 10 attention blocks, 8 multivector channels and 32 scalar channels, and 4 attention heads, resulting in 0.2 million parameters.
For the Transformer baseline, we follow a pre-layer normalization [2, 54] architecture with GELU activations [23] in the MIP block. We use 10 attention blocks, 160 channels, and 4 attention heads, resulting in 1.7 million parameters.
In addition, we compare to the results reported by Suk et al. [48].
22

---

Mean approximation enror [%]
Method
Random orientations
Canonical orientations
GATr
5.6
5.5
Transformer
10.5
6.0
PointNet++ [36]
12.3
8.6
GEM-CNN [15]
7.7
7.8
Table 2: Arterial wall-shear-stress estimation [48]. We show the approximation (lower better), reporting results both on randomly oriented training and test samples (left) on a version of dataset in which all artery meshes are canonically oriented (right). We compare GATr to our own Transformer baseline as well as to the results reported by Suk et al. [48].
Method
Mean approx. error [%]
GATr (default)
6.1
Without multivector representations (= Transformer)
10.5
Without aux. scalar representations
10.4
Without equivariance constraints on linear maps
9.7
Less wide (4 MV channels 8 scalar channels)
12.0
Less deep (5 blocks)
7.8
Table Ablation experiments on the arterial wall well-sshar-sttee stress dataset. We report the mean approximation error (lower better) on the validation set after training for 100 000 steps.
Training We train the models on a Lg loss on wall shear stress. We train for 200 000 steps with the Adam optimizer, using batch size of 8 and exponentially decaying the learning rate from 3 10-1 to 3 10-6
To fit the models into memory, we use mixed 16-bit / 32-bit precision, gradient checkpointing [9], and multi-query attention. Most importantly, we compute dot-product attention with the implementation by [32].
Results
The results of our experiments are shown in Fig. and in Tbl.
Ablation studies We also perform ablation experiments to study the impact of multiple design de- cisions on the performance. All major design choices appear to play an important role: without multi- vector representations, without auxiliary scalar representations, without equivariance constraints, GATr performs much worse. Further reducing the model's 's width also severaly harms the performance.
C.3 Robotic planning through invariant diffuusin
Environment We use the block stacking environment from Janner al. [27]. consists of a Kuka robotic arm interacting with four blocks on table, simulated with PyBullet [13]. 'he state consists of seven robotic joint angles as well as the positions and orientations of the four blocks. We consider the task of stacking four blocks top of each other in any order. The reward is the stacking success probability and is normalized such that means that no blocks are ever successfully stacked, while 100 denotes perfect block stacking.
Dataset and data parameterization We train models on the offline trajectory dataset published by Janner et al. [27]. It consists of 11 000 expert demonstrations.
To facilitate a geometric interpretation, we re-parameterize the environment state into the positions and orientations of the robotic endeffector as well the four blocks. The orientations of all objects are given by two direction vectors. In addition, there are attachment variables that characterize whether the endeffector is in contact with either of the four blocks. In this parameterization, the environment state is dimeessii.....
We train models in this geometric parameterization of the problem. To map back to the original parameterization in terms of joint angles, we use a simple inverse kinematics model that solves for the joint angles consistent with given endeffector pose.
23

---

Method
Reward
GATr-Diffuser (ours)
74.8 + 1.7
Senatormer-Diffusee
69.8 + 1.9
Diffuser [27] (reproduced)
57.7 + 1.8
Diffuser [27]
58.7 2.5
EDGI [7]
62.0 + 2.1
CQL [30] BCQ [21]
24.4
0.0
Table 4: Diffusion- biffusion-based robotic planning. We show the normalized cumulative rewards achieved on a robotic block stacking task [27]. where 100 is optimal and means that each block stacking task is completed successfully, while O corresponds to a failure to stack any blocks. We show the mean and standard error over 200 evaluation episodes. The top three results were computed in the GATr code hase, the hottom four taken from the literature [7, 271.
Models For the GATr model we embed object positions as trivectors, object orientations as oriented planes, gripper attachment variables as scalars, and the diffusion time as scalars. We use axial attention, alternating between attending over time steps and over objects, with positional embeddings along the time axis. Neither multi-query attention nor the distance-aware attention mechanism are used. We train three models, with 10, 20, and 30 Transformer blocks. In each case we use 16 multivector plus 32 scalar channels and 8 attention heads. We report results for the best-performing model with 20 Transformer blocks and 4.0 million parameters.
For the Transformer baseline, also use axial attention, pre-layer normalization, GELU activations, and rotary positional embeddings. We experiment with 6 models, with 10, 20, or 30 Transformer blocks and 144 or 384 channels. All use 8 attention heads. We report results for the best model, which has 20 Transformer blocks, 144 channels, and 3.5 million parameters.
For the Diffuser baseline, we follow the architecture and hyperparameters described by Janner et al. [27]. The model has 65.1 million parameters.
All models are embedded in diffusion pipeline described by Ho et al. [25], using the diffusion time embedding and hyperparameter choices Ref. [27]. In particular, we use univariate Gaussian base densities and 1000 diffusion steps.
Training We train all models by minimizing the simplified diffusion loss proposed by Ho et al. [25]. For our GATr models and the Diffuser baselines with the Adam optimizer, exponentially decaying the learning rate from 3 10 10-4 to 3 10 66, This setup did not work well for the Diffuser model, where (following Janner et al. [27]) we use L, loss and a low constant learning rate instead.
L2 loss and train for 200 000 steps
Evaluation All models are evaluated by rolling out least 200 episodes in a block stacking environment and reporting the mean task and the standard error. We use the planning algorithm and parameter choices of Janner et al. [27] (we do not optimize these, as our focus in this work is on architectural improvements). It consists of sampling trajectories of length 128 from the model, conditional on the current state; then executing these in the environment using PyBullet's PID controller. Each rollout consists of three such phases.
Results
The results of our experiments are shown in Fig. and in Tbl. 4.
C.4
Scaling
Setup To study computational requirements and scaling behaviour, we measure the wall time and peak GPU memory usage during combined forward and backward pass for various models. We use random Gaussian input data with batch size 4, a variable number of items (tokens), and four input channels. After a number of warmup steps, measure the resource use over ten successive steps and report the average time and peak memory usage. For measurements, we use mixed 16-bit/32-bii precision, gradient checkpointing, and efficient attention implementation by Lefaudeux et al. [32].
24

---

Models The GATr model has 10 blocks, 8 multivector and 16 scalar channels, and 4 attention heads. We use queti attention, the distance-aware attention mechanism, and learnable prefactors for the attention weights.
For the baselines, choose hyperparameters match GATr in terms of depth (number of blocks), width (total size of hidden channels), and token connectivity (allowing all tokens to interact with each other in each layer). For the Transformer, this leads to the same choices for GATr, except that 144 hidden channels are used.
For SEGNN and the 3E(3)--raasformen
fully connected graphs, 0(3) representations
< 1, a total of 144 hidden features distributed between scalar and vector representations, and 10 layers. We use the official implementation by Brandstetter et al. [6] for SEGNN and the optimized implementation by Milesi [35] for the Transsrrofformmrr
25


Reviews:
[{"review_id": "X9hZ5NLOKx", "reviewer": "Reviewer_8oHp", "summary": "This paper presents an adaptation of the transformer architecture to enforce equivariance to E(3), which can be helpful to a variety of practical applications necessitating such geometric inductive bias.\n\nThis is achieved partly by leveraging geometric algebra rather than vector algebra, which allows for a simple and elegant representation of all 3D sub-spaces from points to volumes and operations within the same mathematical framework via a single product, the geometric product.\n\nThe authors review the different operations making up a standard transformer and demonstrate how to adapt them to geometric algebra and bake such equivariance building block by building block (from linear layers, to activation to attention). To illustrate the benefit of this architecture, they take two applications, a n-body simulation and a robotic manipulation where they demonstrate superior performance of the GATr architecture compared to vanilla transformer or MLP (first application) or against a number of baselines include an equivariant planning method.\n\nThe specific technical contributions are:\n- block-by-block review and adaptation to geometric algebra and E(3) equivariance of the transformer architecture,\n- an overall E(3) equivariant architecture readily applicable,\n- two applications with demonstrated better performance thanks to this achieved equivariance.\n", "questions": "See weaknesses above, could the authors:\n- elaborate on the specific choice of equivariance and target applications,\n- discuss and illustrate the practical consequences of the efficiency cost on the proposed applications and whether some specialization was used in the applications.\n\n", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The presentation is solid and easy to follow even for someone not necessarily familiar with geometric algebra. That written, a reader unfamiliar with it, will  likely be required to ramp up on the topic via the references to fully grasp the contributions and be able to leverage them.\n\nGeometric algebra is typically a niche mathematical toolkit, which has recently enjoyed a bit more visibility in the graphics community. Shedding more light on this and the practical applications it can have when combined with successful machine learning tools is valuable for the community.\n\nThe authors gradually demonstrate and achieve E(3) equivariance with their GATr architecture very methodologically. This architecture can indeed benefit to a number of practical applications that require strong geometric inductive bias.\n\nThe proposed evaluations of the architecture, while limited to only two applications, show convincing results with clearly superior performance with the proposed architecture, including (and perhaps more convincingly) against another equivariant baseline in the second application.\n\n", "weaknesses": "The proposed transformer architecture focuses on achieving E(3) equivariance, enforcing inductive bias that make this architecture particularly suitable to applications with such constraints. However a number of practical applications would require equivariance to SE(3) only rather than E(3), these include fairly important ones in physics, chemistry and biology due to chirality. In the same register, equivariance in 2D would also be rather useful but is not covered in this submission.\n\nSwitching to geometric algebra representation carries a significant cost both in terms of storage and computation: each element of the n-D algebra requires 2^n cofficients and the geometric product translates to a 2^n x 2^n matrix. Libraries implementing geometric algebra have to rely on specialization, static or dynamic code generation or caching to achieve computational efficiency closer to standard vector algebra. The authors are touching on this important limitations in the corresponding section of the paper, it would have been relevant to include an assessment of the computational cost in practice in the proposed evaluations (or strategies used to alleviate it).\n\nThe submission appears to lack of number of relevant references that achieve some form of equivariance, especially related to the applications evoked in the first weakness above:\n- Highly accurate protein structure prediction with AlphaFold - https://www.nature.com/articles/s41586-021-03819-2\n    See Methods / IPA section + Supplementary Materials / Invariant point attention (IPA) for SE(3) invariance (a special case of equivariance)\n- Cormorant: Covariant Molecular Neural Networks - https://arxiv.org/abs/1906.04015\n- Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds - https://arxiv.org/abs/1802.08219\n- Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs - https://arxiv.org/abs/2206.11990\n- Equivariant Transformer Networks - https://arxiv.org/abs/1901.11399\n\nIt was surprising to not find a mention of equivariance to translation and CNNs. While the second application includes an equivariant baseline, the first one does not. A simple baseline (possibly from the references) with some form of translation (and / or rotation) invariance would strengthen the results of the evaluation, given that the current comparison does not have any geometric equivariance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review, helpful suggestions, and kind words.\n\n\n> A number of practical applications would require equivariance to SE(3) only. [...] Equivariance in 2D would also be rather useful.\n\nGood point, we agree. There are two strategies to address this: either designing separate architectures for each relevant symmetry subgroup of E(3), or using a universal E(3)-equivariant architecture and breaking the symmetry down to a smaller subgroup through **symmetry-breaking inputs**. In this paper we choose the latter strategy, as we find the universality of the architecture appealing.\n\nConcretely, consider the robotic manipulation experiment. Here the direction of gravity clearly breaks the E(3) symmetry to E(2). We use an E(3)-equivariant architecture, but address the smaller symmetry group by providing a vector pointing in the direction of gravity as part of the network inputs. The results show that that strategy works, but it would indeed be interesting to compare to an E(2)-equivariant architecture.\n\nAs another example, relevant e.g. in molecular settings, chirality symmetry can be broken by adding a pseudoscalar input feature, whose sign encodes a choice of orientation.\n\n\n> Each element of the n-D algebra requires 2^n cofficients and the geometric product translates to a 2^n x 2^n matrix\n\nFair point, and it's even worse: the geometric product is a $2^n \\times 2^n \\times 2^n$ tensor, though it is quite sparse and can be stored efficiently. Luckily GPUs are very efficient in matrix multiplication, making the overhead manageable.\n\n\n> Include an assessment of the computational cost in practice in the proposed evaluations (or strategies used to alleviate it).\n\nThanks for the suggestion. We added an experiment measuring the compute time and GPU memory GATr requires as a function of the problem size. We show and discuss the results in the overall response to all reviewers. In short, GATr indeed has some overhead in terms of compute time over a baseline transformer, but scales well to larger problems thanks to the dot-product attention. In particular, GATr **scales much better than the equivariant baselines** SEGNN and SE(3)-Transformer.\n\nStrategies we used to alleviate compute and memory requirements include using efficient implementations of dot-product attention, mixed precision training, gradient checkpointing, and optimized contraction paths in various tensor multiplications (think `einsum`). There is still room for further optimization though, for instance by compilation of the computational graph. We will document all of these tricks and discuss their benefits in the final version of the paper.\n\n\n> The submission appears to lack of number of relevant references that achieve some form of equivariance.\n\nThank you for the pointers. We will add and discuss these references.\n\n\nWe hope that we could address your concerns and look forward to discussing further. Please also have a quick look at our overall response, where we present a few new results. In addition to the scaling experiment, we compare GATr to more equivariant baselines. Finally, we test GATr on a realistic arterial shear stress estimation problem, a benchmark with large geometric structures (meshes of 7000 nodes) and small training sets (1600 training samples). Here GATr sets a new state of the art."}, {"review_id": "aHivU89em7", "reviewer": "Reviewer_wqKZ", "summary": "The paper proposes a (transformer-based) network architecture that operates on a projective geometric algebra that can encode simple 3D transformations and 3D objects (such as points, planes, etc) in a unified description. The network performs operations on these mathematical objects that are (also in composition) equivariant, leading to a frame-invariant computation that is desired for many geometric applications of deep learning, while the projective approach extends the set of transformations that can be employed to cover all of E(3). The paper shows two prototypical examples with geometric learning in gravitational particle simulation and planning of the motion of a robot arm.\nThe main improvement over previous work is the fully equivariant formulation.", "questions": "In how far could permutation invariance (or the lack thereof) become an issue, for example in a complex molecular dynamics simulation with heterogeneous atomistic components and interactions or more complex tasks (e.g. predicting forces or electron densities)? In general, could the presented design perform the same tasks as a (steerable) group-convolutional architecture?\n\nDoes the design come with potential limitations in expressiveness over non-geometric networks? (Appendix A.3 and Section 3.2 show that there are some improvements over prior approaches, but are there some known negative results for the actual/novel approach?)", "limitations": "Networks that obtain equivariance by employing an algebra of covariant and invariant operations are an established idea for simpler sets of geometric objects, such as collections of tensor fields of different type (e.g., scalars and vectors). The fundamental limitation of this approach is the lack of permutation invariance: As every object receives a specific, learned computation, it becomes non-trivial to achieve permutational and geometric invariance at the same time. Solutions often require substantial architectural restrictions (such as restrictions to commutative uniform operations such as pooling) that might reduce expressivity. In the case of the transformer architecture, the transformer is by design permutationally invariant under its input tokens, but problems arise once multiple geometric objects are embedded into the same token (these cannot be easily modeled as indistinguishable). I am lacking experience with transformer architectures to see in how far this might be an issue for typical applications, but it is probably worth discussing.\n\nIt might also be worth considering a bit more in depth in how far expressiveness is constrained by the choice of equivariant operations allowed.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written, easy to read and self-contained despite building on a larger body of previous work with mathematical concepts that are not \"mainstream\" (yet).\n\nThe basic idea appears to be sound and useful: The ability to perform computations on a somewhat richer set of geometric objects and transformations is probably a useful building block for geometric learning algorithms. The specific architecture ensures frame-invariant computations, which is required for many applications.", "weaknesses": "One weakness of the proposed approach is the limited invariance wrt to permutations, see \"limitations\" below. Similarly, it is not clear in how far expressiveness is restricted by the restriction to the specific operations employed.\n\nThe paper is also somewhat incremental in nature, as it adds a selection of specific operations that preserve equivariance, but some of the techniques required appear to be non-obvious.\n\nThe examples are rather basic; they provide a nice proof-of-concept but do not yet show whether the approach is competitive in more realistic problem scenarios or competitive benchmarks.\n\nDue to the limited practical results and some conceptual limitations I would give only a weakly positive recommendation for this (otherwise) well-rounded submission.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive review as well as the kind words.\n\n\n> In how far could permutation invariance (or the lack thereof) become an issue?\n\nOur architecture is fully **permutation equivariant**: a permutation of the input tokens leads to the same permutation of the output tokens. In this aspect it is exactly like a standard Transformer. We will emphasize this property more in the final version of the paper.\n\nIf permutation equivariance is *not* desired in a problem, this equivariance can be broken on the level of the inputs, for instance through positional embeddings. We make use of this option in the robotic manipulation experiments, which have two token dimensions: one between the different objects (along which we maintain permutation equivariance), and one over time steps (here we use positional embeddings to break the permutation symmetry). \n\n\n> Problems arise once multiple geometric objects are embedded into the same token.\n\nIndeed. This is why, in our experiments, we always embed different objects into different tokens. Different properties of the same object (like position, orientation, and velocity of the same object) are embedded into different channels of the same token. We will clarify the embeddings in the appendix.\n\n\n> Does the design come with potential limitations in expressiveness over non-geometric networks?\n\nThis is a great question. We address it in our overall response. In short, while we do not have a formal universality theorem, we have some arguments and empirical evidence that GATr is fully expressive.\n\n\n> In general, could the presented design perform the same tasks as a (steerable) group-convolutional architecture?\n\nWe believe it can. Based on the theoretical and empirical arguments we give in our overall response, we believe that GATr is expressive enough to represent any equivariant map, just like group-convolutional architectures.\n\nIn some sense, our empirical test of expressiveness explicitly tests this: it confirms that GATr can represent any randomly initialized SEGNN network [6], which can be written as group convolutions. Please have a look at our overall response for details.\n\n\n> The paper is also somewhat incremental in nature.\n\nWe politely disagree. Not only is our paper the first to propose equivariant networks for projective geometric algebra representations, it also has two properties that are missing from almost all geometric deep learning works.\n\nFirst, it has **faithful translational E(3) representations**: we can represent points, which transform under translation, not only relative vectors, which are invariant to translation.\n\nSecond, it computes pairwise interactions through dot-product attention, which turns out to be much more scalable than the message-passing operations in most equivariant architectures.\n\nWe study the scaling behaviour in detail in a new experiment and find that **we can scale GATr to 100 times as many tokens** as SEGNN or the SE(3)-Transformer; see the overall response for a discussion and the rebuttal result page for results.\n\n\n> The examples are rather basic.\n\nWe agree, and added a **new experiment on a more realistic problem**: wall shear stress estimation on artery meshes. This is a challenging problem because the meshes are large (**7000 nodes**) and the datasets small (1600 training samples). GATr has no problem scaling to this problem and outperforms all baselines on this task. Please have a look at the overall response for a discussion and the rebuttal result page for the results.\n\n\nThank you again for the review. We hope that we could address your concerns and look forward to discussing further. Please also have a quick look at our overall response. In addition to the scaling and artery experiments, we compare GATr to more equivariant baselines."}, {"review_id": "wDZq7UkOSL", "reviewer": "Reviewer_yyFc", "summary": "This paper introduces the Geometric Algebra Transformer, a novel architecture that leverages the projective geometric algebra to embed inputs, outputs, and hidden features. The transformer incorporates specific linear layers, nonlinear layers, normalization layers, and attention mechanisms designed for the projective geometric algebra. As a general equivariant architecture, the proposed transformer demonstrates significant improvements in experiments, highlighting its effectiveness compared to other models that do not consider geometric information in their features.", "questions": "1. Is there any theoretical relation or comparison between the GATrs and the previous $E(3)$ equivariant networks with steerable features, since they all consider the geometric information in the network.\n\n2. Would there any future directions to measure the expressiveness of GATrs?\n\n3. For the nonlinear layers, would the nonliearity acting on scalar components or the norm impact the expressivity of the network, and does there exist other nonlinearities like the previous approaches in $E(3)$ equivariant network?\n", "limitations": "The authors addressed the limitations in the paper.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper stands out as the first to integrate projective geometric algebra into an equivariant network based on $E(3)$. Unlike previous $E(n)$ networks, the utilization of projective algebra enables the representation of a wider range of geometric elements such as data points, rays, planes, and transformations. The design of the required transformer layers is both straightforward to implement and versatile, facilitating scalability and offering numerous applications in machine learning, computer vision, and robotics. \n\n2. The proposed method is the first equivariant networks with respect to $E(3)$, in contrast with the previous models dealing with $O(3)$ equivariance.\n\n3. The experimental results are impressive, showcasing significant performance enhancements and efficient training data utilization. These outcomes underscore the efficacy and potential of the proposed model.\n\n4. Despite the complexity associated with projective geometric algebra, the paper effectively conveys its ideas through clear explanations and illustrative figures. The authors ensure that readers without prior knowledge of projective algebra can grasp the general concepts and insights presented in the work.", "weaknesses": "1. The proposed model's requirement of 16-dimensional vectors for feature representation may present scalability challenges compared to traditional transformers. Similar to previous equivariant networks utilizing different types of features, this can lead to memory overhead and increased computational complexity. Further experimental analysis in this regard would be valuable to better understand the implications and potential trade-offs.\n\n2. Correct if I'm wrong. When the model adopts the rotary positional embeddings, I think it cannot contain the permutation and rotation equivariance since the transformation will change the corresponding positional embeddings.\n\n3. The experiments might lack one baseline, that is the $SE(3)$ equivariant transformer. I think it would be beneficial to provide the insights of the priority of projective geometric algebra representation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the in-depth review, kind praise, and on-point questions.\n\n\n> The proposed model's requirement of 16-dimensional vectors for feature representation may present scalability challenges compared to traditional transformers.\n\nGood point. To address it, we added a study of the computational requirements and scalability of GATr. Please have a look at the overall response for the results and a discussion.\n\nIn a nutshell, we find that GATr indeed has some overhead in terms of compute time compared to a regular Transformer when working with problems with a small or moderate number of tokens. However, the scaling behaviour to many tokens is the same as that of a Transformer, as both architectures are bottlenecked by the same dot-product attention mechanism. When using an efficient implementation of dot-product attention, we are able to **scale GATr to tens of thousands of tokens**, while equivariant baselines like SEGNN [6] or SE(3)-Transformer [19] run out of memory much sooner.\n\n\n> When the model adopts the rotary positional embeddings, I think it cannot contain the permutation and rotation equivariance.\n\nWe need to apologize for this confusing terminology, which we carelessly adapted from non-geometric transformers. \"Position\" here means position in a sequence, not geometric position in 3D. \"Rotary\" does not refer 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 rotational equivariance. We will use better terminology in the final version.\n\nPositional embeddings do indeed break permutation equivariance. We only use them to indicate the position in a sequence when this is desired. In our experiments, this is only the case for the time dimension of the robotic manipulation problem.\n\n\n> The experiments might lack one baseline, that is the $SE(3)$ equivariant transformer.\n\nThank you for the suggestion. We added the **SE(3)-Transformer** to our n-body experiment. For good measure, we also added another popular equivariant method, the **SEGNN** [6]. Please have a look at our overall response and the rebuttal result page.\n\nBoth the SE(3)-Transformer and SEGNN perform on par with GATr when trained on enough samples, but GATr is more sample-efficient, outperforming both baselines when training on little data. This provides further evidence that not only equivariance, but also the geometric algebra representations provide a useful inductive bias for geometric problems.\n\n\n> Is there any theoretical relation or comparison between the GATrs and the previous $E(3)$ equivariant networks with steerable features?\n\nThere are two key differences between GATr and prior work, like SEGNN [6]. Firstly, we use a geometric algebra as a representation, rather than E(3) irreducible representations. This is on the one hand a richer representation, as it can represent the E(3) group faithfully having features that transform (non-trivially) under translations. On the other hand, when decomposing the algebra into O(3) irreps, one sees that it does not contain irreducible representation of frequency higher than 1, which are often used in other works.\n\nA second difference is that we compute pair-wise interaction via attention, which in effect just takes linear combinations of representations weighted by an invariant scalar. This is different from SEGNN, which uses a neural network to compute pairwise interaction. It remains to be seen theoretically how limiting these choices are, but as we write in our overall response, we have good reasons to believe this does not harm GATr's expressivity.\n\n> Would there be any future directions to measure the expressiveness of GATrs?\n\nThank you for the suggestion. We address it in our overall response. In short, while we do not have a formal universality theorem, we have some arguments and empirical evidence that GATr is fully expressive. We aim to develop a formal result in future work.\n\n> Would the nonlinearity acting on scalar components or the norm impact the expressivity of the network? [...] Are there other nonlinearities like the previous approaches in $E(3)$ equivariant network?\n\nEquivariant neural networks with infinite groups (like E(3)) typically use three kinds of non-linearities: (1) polynomials via tensor products [52], (2) multipling a reprensentation by a scalar (derived from its norm or a scalar gate) [53], (3) approximate equivariance by sampling to a permutation representation, then using a point-wise non-linearity [54]. GATr uses the first two: the geometric product is a polynomial akin to tensor products and we gate the multivectors by multiplication with a scalar (on which we can equivariantly apply point-wise non-linearities like GELU). Therefore, we believe our choice of non-linearities is very similar to mainstream equivariant networks. We thus see no reason why this should hinder our expressivity.\n\n\nThank you again for the review. We hope that we could address your questions and look forward to discussing further. Please also have a look at our overall response. In addition to the scaling analysis and the comparison to more baselines, we also test GATr on a arterial shear stress estimation problem, a realistic benchmark with large geometric structures (meshes of 7000 nodes) and small training sets (1600 training samples). Here GATr sets a new state of the art.\n\n**References**\n\nSee paper, plus:\n\n[52] R. Kondor et al., \"Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network\", NeurIPS 2018\n\n[53] M. Weiler et al., \"3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data\",  NeurIPS 2018\n\n[54] T. Cohen el al., \"Spherical CNNs\", ICLR 2018"}, {"review_id": "SQ55irGQzm", "reviewer": "Reviewer_qqQC", "summary": "This paper is about geometric algebra, which is interesting but not widely known. In this paper, the authors propose a transformer-based model for geometric data, which can represent general geometric data and transformation. The authors also apply this method in two cases, which shows the effectiveness of the proposed method.", "questions": "(1)\tIn the section 4.1, it should be Fig.3 instead of Fig.1. The authors can check this.\n\n(2)\tIn Fig.3, it seems the number of training samples does not affect the performance of GATr, why?\n", "limitations": "From the Figure 1, it seems we need to obtain the geometric types from a processing block. I wonder how the accuracy of this block. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "(1)\tThe research topic is interesting and useful, because the geometric data is widely existed in the application. \n\n(2)\tThe experimental results seem good.\n", "weaknesses": "(1)\tAlthough the authors give some background in this paper, it is still not friendly for readers with little knowledge in geometric algebra. \n\n(2)\tIn this paper, there is not any ablation study.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive review and helpful questions.\n\n\n> There is no ablation study.\n\nThank you for the suggestion. We ran an ablation study on the main architectural decisions in GATr.\n\nWe performed this experiment on a new real-world problem: estimating the wall shear stress from artery meshes. We discuss this dataset in the overall response and show results on it on the rebuttal result page.\n\nHere we provide some additional ablation results. We report the normalized mean approximation error on the validation set after training for 100k steps:\n\n| Method | approximation error [%] |\n| --- | --- |\n| GATr (standard) | **6.1** |\n| Without multivector reps | 10.5 |\n| Without scalar reps | 10.4 |\n| Without equivariance constraints | 9.7 |\n| Less wide | 12.0 |\n| Less deep | 7.8 |\n\nWe conclude that all the main design choices in GATr are important for the strong performance on this dataset, and we cannot reduce the capacity much from the run. We have not yet systematically tested increasing the model capacity. We will add an extended version of this ablation study to the final version of our paper.\n\n\n> In the section 4.1, it should be Fig. 3 instead of Fig. 1.\n\nThanks for catching this, we will fix it.\n\n\n> In Fig. 3, it seems the number of training samples does not affect the performance of GATr.\n\nIt is hard to see in this log-scale plot, but GATr's error does actually decrease with training set size, from 6e-4 to 4e-4 mean absolute error. But we were also surprised to see how strong GATr performs even with very little training data. This is testament to the strong geometric inductive bias from equivariance and geometric algebra representations.\n\n\n> From the Figure 1, it seems we need to obtain the geometric types from a processing block. I wonder how the accuracy of this block.\n\nThis is a great question, and the answer will depend on the application. In our experiments, we assume knowledge of positions and/or momenta of the objects, so this is no source of errors.\n\nIn contrast, real-world robotics applications will require estimating the positions and poses of objects from sensor inputs, which will lead to some approximation error. We believe that such errors can be quite low when e.g. depth cameras are used. We are planning to study integrating this preprocessing step into the pipeline in future work.\n\n\nThank you again for the review. We hope that we could address your questions and look forward to discussing further. Please also have a look at our overall response, where we present some new results. In addition to the artery experiment, we also compare GATr to more equivariant baselines. Finally, we study its computational requirements and scaling behaviour, showing that GATr can be scaled to much larger systems than equivariant baselines.\n"}]

Author rebuttal:
Thanks to the reviewers for their insightful feedback. We are delighted that they find the key ideas behind GATr - geometric algebra representations, equivariance, and a transformer architecture - "sound and useful" (reviewer **wqKZ**) and "very methodological" (**8oHP**). They appreciated both the benefits of true translational E(3) representations and the scalability of the transformer implementation (**yyFc**) and commented positively on the "good" (**qqQC**) experimental results. We are particularly happy that the reviewers found our presentation "well written, easy to read and self-contained" (**wqKZ**).

## Comparison to equivariant baselines

At the suggestion of reviewers **8oHP** and **yyFc**, we added a comparison of GATr to other equivariant architectures on the n-body experiment. We chose the **SE(3)-Transformer** [19], suggested by reviewer **yyFc**, and **SEGNN** [6].

In Fig. 1 of the attached PDF, we show that all three equivariant architectures perform similarly given enough data. However, **GATr is clearly more sample-efficient**.

We also added a comparison to the GCA-GNN architecture [46], also based on projective geometric algebra representations, but not equivariant. GATr clearly outperforms GCA-GNN.

## Expressiveness

Reviewers **wqKZ** and **yyFc** asked about the expressiveness of our architecture. We are interested in theoretically analyzing the expressiveness of GATr as future work, but until such analysis has taken place, we would like to make two observations.

First, Villar et al. [47] show that universal expressiveness of E(3)-equivariant networks can be achieved using merely scalars and vectors, without the need fo higher-order irreps. Scalars and vectors are expressible in the geometric algebra, so we don't think our choice of representation limits expressivity. In addition, we don't think that the transformer architecture is limiting, as they are known to be universal approximators in the non-geometric setting [48].

Second, we studied expressivity empirically. We randomly initialized SEGNN networks [6], which are similar to achitectures known to be universal equivariant approximators [49], and trained GATr networks to imitate them. Our models were always able to reach close to 0 loss. A thorough version of this experiment will be included in the final version.

## Computational cost and scalability

Reviewers **8oHP** and **yyFc** suggested to study the computational costs and scaling behaviour of our architecture. We measured the time and memory requirements of GATr as a function of the number of tokens. We compare to multiple baselines, including a heavily optimized Nvidia implementation of the SE(3)-Transformer [50].

We show the results in Fig. 2 of the rebuttal result page. For few tokens, GATr is slower than a Transformer. However, this difference can be lowered with larger batch sizes or pre-compilation of 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.

## Arterial wall shear stress experiment

Reviewer **wqKZ** asked about the performance on more realistic scenarios. We address this with a new experiment on a dataset of synthetic meshes of realistic human arteries [51]. The task is to predict 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. As the meshes are large (around **7000 nodes**) and the datasets small (1600 training meshes), this is a challenging problem.

We show results in Fig. 3 of the attached PDF. Among the methods that do not canonicalize the data, GATr improves upon all previous methods and sets a **new state of the art** on this problem. It outperforms both a Transformer baseline as well as the (equivariant and non-equivariant) results reported by Suk et al. [51].

Canonicalizing the data (rotating arteries such that blood always flows in the same direction) helps the Transformer to be competitive with GATr. However, canonicalization is only feasible for relatively straight arteries as in this dataset, not in more complex scenarios with branching and turning arteries. We find it likely that GATr will be more robust in such scenarios.

## Conclusions

The results we add in this rebuttal highlight the two strengths of GATr: First, its geometric inductive bias allows GATr to perform on par or better than strong equivariant baselines. Second, thanks to its dot-product attention, it scales much more favorably to large problems with thousands of tokens.

## References

See paper, plus:

[46] D. Ruhe et al., "Geometric Clifford Algebra Networks", ICML 2023

[47] S. Villar et al., "Scalars are universal: Equivariant machine learning, structured like classical physics", NeurIPS 2021

[48] C. Yun et al.,"Are Transformers universal approximators of sequence-to-sequence functions?", ICLR 2020

[49] N. Dym & H. Maron, "On the Universality of Rotation Equivariant Point Cloud Networks", arXiv:2010.02449

[50] A. Milesi, "Accelerating SE(3)-Transformers Training Using an NVIDIA Open-Source Model Implementation", 2021

[51] J. Suk et al., "Mesh convolutional neural networks for wall shear stress estimation in 3D artery models", MICCAI 2021


Meta-review: Thanks for clarifying a few of the loose ends in the rebuttal. All the reviewers liked the submission. Please consider including the relevant bits from the rebuttal in the final version of the paper. Congratulations on a nice piece of work.
Decision: Accept (poster)