Geometric Algebra Transformer

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.

Strengths

(1) The research topic is interesting and useful, because the geometric data is widely existed in the application. (2) The experimental results seem good.

Weaknesses

(1) Although the authors give some background in this paper, it is still not friendly for readers with little knowledge in geometric algebra. (2) In this paper, there is not any ablation study.

Questions

(1) In the section 4.1, it should be Fig.3 instead of Fig.1. The authors can check this. (2) In Fig.3, it seems the number of training samples does not affect the performance of GATr, why?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

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.

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.

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. 2. The proposed method is the first equivariant networks with respect to $E(3)$, in contrast with the previous models dealing with $O(3)$ equivariance. 3. 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. 4. 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. 2. 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. 3. 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.

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. 2. Would there any future directions to measure the expressiveness of GATrs? 3. 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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The authors addressed the limitations in the paper.

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. The main improvement over previous work is the fully equivariant formulation.

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). The 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. The 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. The 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. Due to the limited practical results and some conceptual limitations I would give only a weakly positive recommendation for this (otherwise) well-rounded submission.

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? Does 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?)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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. It might also be worth considering a bit more in depth in how far expressiveness is constrained by the choice of equivariant operations allowed.

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. This 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. The 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. The specific technical contributions are: - block-by-block review and adaptation to geometric algebra and E(3) equivariance of the transformer architecture, - an overall E(3) equivariant architecture readily applicable, - two applications with demonstrated better performance thanks to this achieved equivariance.

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. Geometric 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. The 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. The 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.

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. Switching 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). The 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: - Highly accurate protein structure prediction with AlphaFold - https://www.nature.com/articles/s41586-021-03819-2 See Methods / IPA section + Supplementary Materials / Invariant point attention (IPA) for SE(3) invariance (a special case of equivariance) - Cormorant: Covariant Molecular Neural Networks - https://arxiv.org/abs/1906.04015 - Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds - https://arxiv.org/abs/1802.08219 - Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs - https://arxiv.org/abs/2206.11990 - Equivariant Transformer Networks - https://arxiv.org/abs/1901.11399 It 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.

Questions

See weaknesses above, could the authors: - elaborate on the specific choice of equivariance and target applications, - discuss and illustrate the practical consequences of the efficiency cost on the proposed applications and whether some specialization was used in the applications.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A