SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We find it beneficial to add learnable weights as prefactors to each of these three terms. The attention weights are then given by
$$
\text {S o f t m a x} _ {i} \left(\frac {\alpha \sum_ {c} \langle q _ {i ^ {\prime} c} ^ {M V} , k _ {i c} ^ {M V} \rangle + \beta \sum_ {c} \phi (q _ {i ^ {\prime} c} ^ {M V}) \cdot \psi (k _ {i c} ^ {M V}) + \gamma \sum_ {c} q _ {i ^ {\prime} c} ^ {s} k _ {i c} ^ {s}}{\sqrt {1 3 n _ {M V} + n _ {s}}}\right)
$$
with learnable, head-specific  $\alpha ,\beta ,\gamma >0$
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 to compute GATr's attention.
Multi-query attention To reduce memory use, we also consider a version of GATr that uses multi-query attention [43] instead of multi-head attention, sharing the keys and values among attention heads.
C.1  $n$ -body dynamics
Dataset We generate a synthetic  $n$ -body dataset for  $n$  objects by following these steps for each sample:
1. The masses of  $n$  objects are sampled from log-uniform distributions. For one object (the star), we use  $m_0 \in [1, 10]$ ; for the remaining objects (the planets), we use  $m_i \in [0.01, 0.1]$ . (Following common practice in theoretical physics, we use dimensionless quantities such that the gravitational constant is 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_i \in [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 to 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 a mean of  $(200,0,0)^T$ ). The rotation is sampled from the Haar measure on  $SO(3)$ . In addition, we apply a random permutation of the bodies.
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 a time interval of  $10^{-4}$  each.
6. Finally, samples in which any bodies have traveled more than a distance of 2 (the diameter 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  $n = 4$  and between 100 and  $10^{5}$  samples; a validation dataset with  $n = 4$  and 5000 samples; a regular evaluation set with  $n = 4$  and 5000 samples; a number- generalization evaluation set with  $n = 6$  and 5000 samples; and a E(3) generalization set with  $n = 4$ , an additional translation (see step 4 above), and 5000 samples.
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-aware attention mechanism (as it led to similar results in a short test).
For the Transformer baseline, we follow a 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 11.8 million parameters.
We also compare to a simple MLP with GELU activations. 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, we compare GATr to the geometric-algebra-based, 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 SE(3)-Transformer [20]. In both cases we use the code released by the authors and the hyperparameters they recommend for (slightly different)  $n$ -body 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 3 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 a  $L_{2}$  loss on the final position of all objects. We train for 50 000 steps with the Adam optimizer, using a batch size of 64 and exponentially decaying the learning rate from  $3 \cdot 10^{-4}$  to  $3 \cdot 10^{-6}$ .
C.2 Arterial wall-shear-stress estimation
Dataset We use the single-artery wall-shear-stress dataset published by Suk et al. [48]. It 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 a 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 MLP 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].
Method	Mean approximation error [%]
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 mean approximation error (lower is better), reporting results both on randomly oriented training and test samples (left) and on a version of the 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 3: Ablation experiments on the arterial wall-shear-stress dataset. We report the mean approximation error (lower is better) on the validation set after training for 100000 steps.
Training We train the models on a  $L_{2}$  loss on wall shear stress. We train for 200 000 steps with the Adam optimizer, using a batch size of 8 and exponentially decaying the learning rate from  $3 \cdot 10^{-4}$  to  $3 \cdot 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. 3 and in Tbl. 2.
Ablation studies We also perform ablation experiments to study the impact of multiple design decisions on the performance. All major design choices appear to play an important role: without multi-vector representations, without auxiliary scalar representations, or without equivariance constraints, GATr performs much worse. Further reducing the model's width also severely harms the performance.
C.3 Robotic planning through invariant diffusion
Environment We use the block stacking environment from Janner et al. [27]. It consists of a Kuka robotic arm interacting with four blocks on a table, simulated with PyBullet [13]. The 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 on top of each other in any order. The reward is the stacking success probability and is normalized such that 0 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 as 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 49-dimensional.
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 a given end effector pose.
Method	Reward
GATr-Diffuser (ours)	74.8 ± 1.7
Transformer-Diffuser	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]	24.4
BCQ [21]	0.0
Table 4: Diffusion-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 0 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 base, the bottom four taken from the literature [7, 27].
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, we 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 a diffusion pipeline as described by Ho et al. [25], using the diffusion time embedding and hyperparameter choices of 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 we use an  $L_{2}$  loss and train for 200 000 steps with the Adam optimizer, exponentially decaying the learning rate from  $3 \cdot 10^{-4}$  to  $3 \cdot 10^{-6}$ . This setup did not work well for the Diffuser model, where (following Janner et al. [27]) we use a  $L_{1}$  loss and a low constant learning rate instead.
Evaluation All models are evaluated by rolling out at 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. 4 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 a combined forward and backward pass for various models. We use random Gaussian input data with a batch size of 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 all measurements, we use mixed 16-bit/32-bit precision, gradient checkpointing, and an efficient attention implementation by Lefaudeauux et al. [32].
Models The GATr model has 10 blocks, 8 multivector and 16 scalar channels, and 4 attention heads. We use multi-query attention, the distance-aware attention mechanism, and learnable prefactors for the attention weights.
For the baselines, we choose hyperparameters to 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 as for GATr, except that 144 hidden channels are used.
For SEGNN and the SE(3)-Transformer, we use fully connected graphs, O(3) representations  $\ell \leq 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 SE(3)-Transformer.