Title: A General Framework for Equivariant Neural Networks on Reductive Lie Groups
Venue: NeurIPS 2023
Accepted: yes
Keywords: equivariance

point clouds

machine learning

particle physics
Abstract: Reductive Lie Groups, such as the orthogonal groups, the Lorentz group, or the unitary groups, play essential roles across scientific fields as diverse as high energy physics, quantum mechanics, quantum chromodynamics, molecular dynamics, computer vision, and imaging. In this paper, we present a general Equivariant Neural Network architecture capable of respecting the symmetries of the finite-dimensional representations of any reductive Lie Group. Our approach generalizes the successful ACE and MACE architectures for atomistic point clouds to any data equivariant to a reductive Lie group action. We also introduce the lie-nn software library, which provides all the necessary tools to develop and implement such general G-equivariant neural networks. It implements routines for the reduction of generic tensor products of representations into irreducible representations, making it easy to apply our architecture to a wide range of problems and groups. The generality and performance of our approach are demonstrated by applying it to the tasks of top quark decay tagging (Lorentz group) and shape recognition (orthogonal group).

A General Framework for Equivariant Neural Networks on Reductive Lie Groups

Abstract

Reductive Lie Groups, such as the orthogonal groups, the Lorentz group, or the unitary groups, play essential roles across scientific fields as diverse as high energy physics, quantum mechanics, quantum chromodynamics, molecular dynamics, computer vision, and imaging. In this paper, we present a general Equivariant Neural Network architecture capable of respecting the symmetries of the finitedimensional representations of any reductive Lie Group G . Our approach generalizes the successful ACE and MACE architectures for atomistic point clouds to any data equivariant to a reductive Lie group action. We also introduce the lie-nn software library, which provides all the necessary tools to develop and implement such general G -equivariant neural networks. It implements routines for the reduction of generic tensor products of representations into irreducible representations, making it easy to apply our architecture to a wide range of problems and groups. The generality and performance of our approach are demonstrated by applying it to the tasks of top quark decay tagging (Lorentz group) and shape recognition (orthogonal group).

1 Introduction

Convolutional Neural Networks (CNNs) (LeCun et al. , 1989) have become a widely used and powerful tool for computer vision tasks, in large part due to their ability to achieve translation equivariance. This property led to improved generalization and a significant reduction in the number of parameters. Translation equivariance is one of many possible symmetries occurring in machine learning tasks.

A wide range of symmetries described by reductive Lie Groups is present in physics, such as O (3) in molecular mechanics, SO(1 , 3) in High-Energy Physics, SU(2 N ) in quantum mechanics, and SU(3) in quantum chromodynamics. Machine learning architectures that respect these symmetries often lead to significantly improved predictions while requiring far less training data. This has been demonstrated in many applications including 2D imaging with O(2) symmetry (Cohen and Welling, 2016a; Esteves et al. , 2017), machine learning force fields with O(3) symmetry (Anderson et al. , 2019; Bartók et al. , 2013; Batzner et al. , 2022; Batatia et al. , 2022a) or jet tagging with SO + (1 , 3) symmetry (Bogatskiy et al. , 2022; Li et al. , 2022).

One way to extend CNNs to other groups (Finzi et al. , 2020; Kondor and Trivedi, 2018) is through harmonic analysis on homogeneous spaces, where the convolution becomes an integral over the group. Other architectures work directly with finite-dimensional representations. We follow the demonstration of Bogatskiy et al. (2020a) who constructed a universal approximation of any equivariant map with a feed-forward neural network with vector activations belonging to finite-dimensional representations of a wide class of Lie groups. In this way, one can avoid computational challenges created by infinite-dimensional representations.

Alternatively, our current work can be thought of as a generalization of the Atomic Cluster Expansion (ACE) formalism of Drautz (2019) to general Lie groups. The ACE formalism provides a complete body-ordered basis of O(3) -invariant features. By combining the concepts of ACE and E(3) -equivariant neural networks, Batatia et al. (2022a) proposed the MACE architecture, which achieves state-of-the-art performance on learning tasks in molecular modelling. The present work generalizes the ACE and MACE architectures to arbitrary Lie groups in order to propose a generic architecture for creating representations of geometric point clouds in interaction.

Concretely, our work makes the following contributions:

We develop the G -Equivariant Cluster Expansion. This framework generalizes the ACE (Drautz, 2019) and MACE (Batatia et al. , 2022b) architectures to parameterize properties of point clouds, equivariant under a general reductive Lie group G .

We prove that our architecture is universal, even for a single layer.

We introduce lie-nn , a new library providing all the essential tools to apply our framework to a variety of essential Lie Groups in physics and computer visions, including the Lorentz group, SU( N ) , SL 2 ( C ) and product groups.

We illustrate the generality and efficiency of our general-purpose approach by demonstrating excellent accuracy on two prototype applications, jet tagging, and 3D point cloud recognition.

2 Background

We briefly review a few important group-theoretic concepts: A real (complex) Lie group is a group that is also a finite-dimensional smooth (complex) manifold in which the product and inversion of the group are also smooth (holomorphic) maps. Among the most important Lie groups are Matrix Lie groups, which are closed subgroups of GL( n, C ) the group of invertible n × n matrices with complex entries. This includes well-known groups such as Sp(2 n, R ) consisting of matrices of determinant one, that is relevant in Hamiltonian dynamics A finite-dimensional representation of the Lie group

G is a finite-dimensional vector space V endowed with a smooth homomorphism ρ : G → GL( V ) . Features in the equivariant neural networks live in these vector spaces. An irreducible representation V is a representation that has no subspaces which are invariant under the action of the group (other than { 0 } and V itself). This means that V can not be decomposed non-trivially as the direct sum of representations. A reductive group over a field F is a (Zariski-) closed subgroup of the group of matrices GL( n, F ) such that every finite-dimensional representation of G on an F -vectorspace can be decomposed as a sum of irreducible representations.

3 Related Work

Lie group convolutions Convolutional neural networks (CNNs), which are translation equivariant, have also been generalized to other symmetries. For example, G-convolutions (Cohen and Welling, 2016b) generalized CNNs to discrete groups. Steerable CNNs (Cohen and Welling, 2016a) generalized CNNs to O (2) equivariance and Spherical CNNs (Cohen et al. , 2018) O (3) equivariance. A general theory of convolution on any compact group and symmetric space was given by Kondor and Trivedi (2018). This work was further extended to equivariant convolutions on Riemannian manifolds by Weiler et al. (2021).

ACE The Atomic Cluster Expansion (ACE) (Drautz, 2019) introduced a systematic framework for constructing complete O (3) -invariant high body order basis sets with constant cost per basis function, independent of body order (Dusson et al. , 2022).

e3nn + Equivariant MLPs The e3nn library (Geiger and Smidt, 2022) provides a complete solution to build E (3) -equivariant neural networks based on irreducible representations. The Equivariant MLPs (Finzi et al. , 2021) include more groups, such as SO (1 , 3) , Z n , but are restricted to reducible representations making them much less computationally efficient than irreducible representations.

Equivariant MPNNs and MACE Equivariant MPNNs (Kondor et al. , 2018; Anderson et al. , 2019; Bogatskiy et al. , 2020a; Satorras et al. , 2021; Brandstetter et al. , 2022; Batzner et al. , 2022) have emerged as a powerful architecture to learn on geometric point clouds. They construct permutation invariants and group equivariant representations of point clouds. Successful applications include simulations in chemistry, particle physics, and 3D vision. MACE (Batatia et al. , 2022a) generalized the O (3) -Equivariant MPNNs to build messages of arbitrary body order, outperforming other approaches on molecular tasks. (Batatia et al. , 2022b) showed that the MACE design space is large enough to include most of the previously published equivariant architectures.

4 The G -Equivariant Cluster Expansion

We are concerned with the representation of properties of point clouds. Point clouds are described as multi-sets (unordered tuples) X = [ x i ] i where each particle x i belongs to a configuration domain Ω . We denote the set of all such multi-sets by msets(Ω) . For example, in molecular modeling, x i might describe the position and species of an atom and therefore x i = ( r i , Z i ) ∈ R 3 × Z , while in high energy physics, one commonly uses the four-momentum x i = ( E i , p i ) ∈ R 4 , but one could also include additional features such as charge, spin, and so forth. A property of the point cloud is a map

i.e., X ↦→ Φ( X ) ∈ Z , usually a scalar or tensor. The range space Z is application dependent and left abstract throughout this paper. Expressing the input as a multi-set implicitly entails two important facts: (1) it can have varying lengths; (2) it is invariant under the permutations of the particles. The methods developed in this article are also applicable to fixed-length multi-sets, in which case Φ is simply a permutation-invariant function defined on some Ω N . Mappings that are not permutation-invariant are special cases with several simplifications.

In many applications, especially in the natural sciences, particle properties satisfy additional symmetries. When a group G acts on Ω as well as on Z we say that Φ is G -equivariant if

where ρ Z ( g ) is the action of the group element g on the range space Z . In order to effectively incorporate exact group symmetry into properties Φ , we consider model architectures of the form

where the space V into which we 'embed' the parameterization is a possibly infinite-dimensional vector space in which a convenient representation of the group is available. For simplicity we will sometimes assume that Z = V .

The Atomic Cluster Expansion (ACE) framework (Drautz, 2019; Dusson et al. , 2022; Drautz, 2020)) produces a complete linear basis for the space of all 'smooth' G -equivariant properties Φ for the specific case when G = O(3) and x i are vectorial interatomic distances. Aspects of the ACE framework were incorporated into E(3) -equivariant message passing architectures, with significant improvements in accuracy (Batatia et al. , 2022a). In the following paragraphs we demonstrate that these ideas readily generalize to arbitrary reductive Lie groups.

4.1 Efficient many-body expansion

The first step is to expand Φ in terms of body orders, and truncate the expansion at a finite order N :

where φ n defines the n -body interaction. Formally, the expansion becomes systematic in the limit as N →∞ . The second step is the expansion of the n -particle functions φ n in terms of a symmetrized tensor product basis. To define this we first need to specify the embedding of particles x : A countable family ( ϕ k ) k is a 1-particle basis if they are linearly independent on Ω and any smooth 1-particle function φ 1 (not necessarily equivariant) can be expanded in terms of ( ϕ k ) k , i.e,

For the sake of concreteness, we assume that ϕ k : Ω → C , but the range can in principle be any field.We provide concrete examples of 1-particle bases in Appendix A.2. Let a complex vector space V be given, into which the particle embedding maps, i.e.,

As a consequence of (5) any smooth scalar n -particle function φ n can be expanded in terms of the corresponding tensor product basis,

Inserting these expansions into (4) and interchanging summation (see appendix for the details) we arrive at a model for scalar permutation-symmetric properties,

where K is the set of all k tuples indexing the features A k . Since A k is invariant under permuting k , only ordered k tuples are retained. The features A k are an embedding of msets(Ω) into the space V . The tensorial product features (basis functions) A k form a complete linear basis of multi-set functions on Ω and the weights w k can be understood as a symmetric tensor. We will extend this linear cluster expansion model Φ ( N ) to a message-passing type neural network model in § 4.4.

While the standard tensor product embeds ( ⊗ n s =1 ϕ k s ) k : Ω n → V n , the n -correlations A k are symmetric tensors and embed ( A k ) k : msets(Ω) → Sym n V .

The evaluation of the symmetric tensor features A k is the computational bottleneck in most scenarios, but efficient recursive evaluation algorithms (Batatia et al. , 2022a; Kaliuzhnyi and Ortner, 2022) are available. See Appendix A.13.2 for further discussion of model computational costs.

4.2 Symmetrisation

With (7) we obtained a systematic linear model for (smooth) multi-set functions. It remains to incorporate G -equivariance. We assume that G is a reductive Lie group with a locally finite representation in V . In other words we choose a representation ρ = ( ρ kk ′ ): G → GL( V ) such that

where for each k the sum over k ′ is over a finite index-set depending only on k . Most Lie groups one encounters in physical applications belong to this class, the affine groups being notable exceptions. However, those can usually be treated in an ad hoc fashion, which is done in all E (3) -equivariant architectures we are aware of. In practice, these requirements restrict how we can choose the embedding ( ϕ k ) k . If the point clouds X = [ x i ] i are already given in terms of a representation of the group, then one may simply construct V to be iterative tensor products of Ω ; see e.g. the MTP (Shapeev, 2016) and PELICAN (Bogatskiy et al. , 2022) models. To construct an equivariant two-particle basis we need to first construct the set of all intertwining operators from V ⊗ V → V . Concretely, we seek all solutions C α ,K k 1 k 2 to the equation

or, written in operator notation, C α ρ ⊗ ρ = ρC α . Wewill call the C α ,K k generalized Clebsch-Gordan coefficients since in the case G = SO(3) acting on the spherical harmonics embedding ϕ lm = Y l m those coefficients are exactly the classical Clebsch-Gordan coefficients. The index α enumerates a basis of the space of all solutions to this equation. For the most common groups, one normally identifies a canonical basis C α and assigns a natural meaning to this index (cf. § A.5). Our abstract notation is chosen because of its generality and convenience for designing computational schemes. The generalization of the Clebsch-Gordan equation (9) to n products of representations acting on the symmetric tensor space Sym n ( V ) becomes (cf. § A.9)

Due to the symmetry of the ( A k ) k tensors C α ,K k need only be computed for ordered k tuples and the sum ∑ k ′ also runs only over ordered k tuples. Again, the index α enumerates a basis of the space of solutions. Equivalently, (10) can be written in compact notation as C α ρ = ρ C α . These coupling operators for N products can often (but not always) be constructed recursively from couplings of pairs (9). We can now define the symmetrized basis

The equivariance of (11) is easily verified by applying a transformation g ∈ G to the input (cf § A.6).

Universality: In the limit as the correlation order N →∞ , the features ( B K α ) K, α form a complete basis of smooth equivariant multi-set functions, in a sense that we make precise in Appendix A.7. Any equivariant property Φ V : Ω → V can be approximated by a linear model

to within arbitrary accuracy by taking the number of terms in the linear combination to infinity.

4.3 Dimension Reduction

The tensor product of the cluster expansion in (7) is taken on all the indices of the one-particle basis. Unless the embedding ( ϕ k ) k is very low-dimensional it is often preferable to 'sketch' this tensor product. For example, consider the canonical embedding of an atom x i = ( r i , Z i ) ,

Only the ( lm ) channels are involved in the representation of O(3) hence there is considerable freedom in 'compressing' the ( z, n ) channels.

Following Darby et al. (2022) we construct a sketched G -equivariant cluster expansion: We endow the one-particle basis with an additional index c , referred to as the sketched channel, replacing the index k with the index pair ( c, k ) , and renaming the embedding ( ϕ ck ) c,k . In the case of three-dimensional particles one may, for example, choose c = ( z, n ) . In general, it is crucial that the representation

remains in terms of ρ k,k ′ , that is, (8) becomes ϕ ck ◦ g = ∑ k ′ ρ kk ′ ( g ) ϕ ck ′ . Therefore, manipulating only the c channel does not change any symmetry properties of the architecture. Generalizing Darby et al. (2022), the G -TRACE (tensor-reduced ACE) basis then becomes

This construction is best understood as an equivariance-preserving canonical tensor decomposition (Darby et al. , 2022). There are numerous variations, but for the sake of simplicity, we restrict our presentation to this one case.

Universality: Following the proof of Darby et al. (2022) one can readily see that the G -TRACE architecture inherits the universality of the cluster expansion, in the limit of decoupled channels # c →∞ . A smooth equivariant property Φ may be approximated to within arbitrary accuracy by an expansion Φ K ( X ) ≈ ∑ c, α c K α B K c, α ( X ) . Since the embedding ˜ A ck is learnable, this is a nonlinear model . We refer to § A.7 for the details.

4.4 G-MACE, Multi-layer cluster expansion

The G -equivariant cluster expansion is readily generalized to a multi-layer architecture by reexpanding previous features in a new cluster expansion (Batatia et al. , 2022b). The multi-set X is endowed with extra features, h t i = ( h t i,cK ) c,K , that are updated for t ∈ { 1 , ..., T } iterations. These features themselves are chosen to be a field of representations such that they have a well-defined transformation under the action of the group. This results in

The recursive update of the features proceeds as in a standard message-passing framework but with the unique aspect that messages are formed via the G -TRACE and in particular can contain arbitrary high correlation order:.

The gathered message m t i = ( m t i,cK ) c,k is then used to update the particle states,

where U t can be an arbitary fixed or learnable transformation (even the identity). Lastly, a readout function maps the state of a particle to a target quantity of interest, which could be local to each particle or global to the mset X ,

This multi-layer architecture corresponds to a general message-passing neural network with arbitrary body order of the message at each layer. We will refer to this architecture as G -MACE. The G -MACE architecture directly inherits universality from the G -ACE and G -TRACE architectures:

Theorem 4.1 (Universality of G -MACE) . Assume that the one-particle embedding ( ϕ k ) k is a complete basis. Then, the set of G -MACE models, with a fixed finite number of layers T , is dense in the set of continuous and equivariant properties of point clouds X ∈ msets(Ω) , in the topology of pointwise convergence. It is dense in the uniform topology on compact and size-bounded subsets.

5 lie-nn : Generating Irreducible Representations for Reductive Lie Groups

In order to construct the G-cluster expansion for arbitrary Lie groups, one needs to compute the generalized Clebsch-Gordan coefficients (10) for a given tuple of representations (see 11). To facilitate this task, we have implemented an open source software library, lie-nn 1 . In this section we review the key techniques employed in this library.

5.1 Lie Algebras of Reductive Lie Groups

Formally, the Lie algebra of a Lie group is its tangent space at the origin and carries an additional structure, the Lie bracket. Informally the Lie algebra can be thought of as a linear approximation to the Lie group but, due to the group structure, this linear approximation carries (almost) full information about the group. In particular the representation theory of the Group is almost entirely determined by the Lie algebra, which is a simpler object to work with instead of the fully nonlinear Lie group.

Lie algebra The Lie groups we study can be realized as closed subgroups G ⊂ GL n ( R ) of the general linear group. In that case their Lie algebras can be concretely realized as g = Lie( G ) = { X ∈ M n ( R ) | ∀ t ∈ R : exp( tX ) ∈ G } where exp( X ) = 1 + X + 1 2 X 2 ... is the standard matrix exponential. It turns out that g ⊂ M n ( R ) is a linear subspace closed under the commutator bracket [ X,Y ] = XY -Y X .

Structure theory We fix a linear basis { X i } ⊂ g , called a set of generators for the group. The Lie algebra structure is determined by the structure constants A ijk defined by [ X i , X j ] = ∑ k A ijk X k , in that if X = ∑ i a i X i and Y = ∑ j b j X j then [ X,Y ] = ∑ k ( ∑ i,j A ijk a i b j ) X k . The classification of reductive groups provides convenient generating sets for their Lie algebras (or their complexifications). One identifies a large commutative subalgebra h ⊂ g (sometimes of g C = g ⊗ R C ) with basis { H i } so that most (or all) of the other generators E α can be chosen so that [ H i , E α ] = α ( H i ) E α for a linear function α on h . These functions are the so-called roots of g . Structural information about g is commonly encoded pictorially via the Dynkin diagram of g , a finite graph the nodes of which are a certain subset of the roots. There are four infinite families of simple complex Lie algebras A n = su ( n +1) , B n = so (2 n +1) , C n = sp (2 n ) , D n = so (2 n ) and further five exceptional simple complex Lie algebras (a general reductive Lie algebra is the direct sum of several simple ones and its centre). The Lie algebra only depends on the connected component of G . thus when the group G is disconnected in addition to the infinitesimal generators { X i } one also needs to fix so-called "discrete generators", a subset H ⊂ G containing a representative from each connected component.

Representation theory The representation theory of complex reductive Lie algebras is completely understood. Every finite-dimensional representation is (isomorphic to) the direct sum of irreducible representations ("irreps"), with the latter parametrized by appropriate linear functional on h ("highest weight"). Further given a highest weight λ there is a construction of the associated irrep with an explicit action of the infinitesimal generators chosen above. The Weyl Dimension Formula gives the dimension of an irrep in terms of its highest weight.

5.2 Numerical Computations in lie-nn

The most basic class of the lie-nn library encodes a group G and infinitesimal representation dρ of g using the tuple

with A the structure constants of the group, n the dimension of the representation, and dρ ( X i ) and ρ ( h ) being n × n matrices encoding the action of the infinitesimal and the discrete generators respectively. The action of infinitesimal generators is related to the action of group generators by the exponential, ∀ X ∈ g , ρ ( e X ) = e dρ ( X ) . For finite groups, we assume that dρ ( X ) = 0 as they have only discrete generators.

As the building blocks of the theory irreps are treated specially; the package implements functionality for the following operations for each supported Lie group:

Constructing the irrep with a given highest weight.

Determining the dimension of an irrep.

Decomposing the tensor product of several irreps into irreps up to isomorphism (the selection rule , giving the list of irreducible components and their multiplicities).

Decomposing the tensor product of several irreps into irreps explicitly via a change of basis ("generalized Clebsch-Gordan coefficients").

Computating the symmetrized tensor product of the group (see. 5.3 and A.9 for details).

To construct an irrep explicitly as in (20) one needs to choose a basis in the abstract representation space (including a labeling scheme for the basis) so that we can give matrix representations for the action of generators. For this purpose, we use in lie-nn the Gelfand-Tsetlin (GT) basis (Gelfand and Tsetlin, 1950) and associated labeling of the basis by GT patterns (this formalism was initially introduced for algebras of type A n but later generalized to all classical groups). Enumerating the GT patterns for a given algebra gives the dimension of a given irrep, the selection rules can be determined combinatorially, and it is also possible to give explicit algorithms to compute Clebsch-Gordan coefficients (the case of A n is treated by Alex et al. (2011)). For some specific groups, simplifications to this procedure are possible and GT patterns are not required.

In some cases, one wants to compute coefficients for reducible representations or for representations where the analytical computation with GT patterns is too complex. In these cases, a numerical algorithm to compute the coefficients is required. Let dρ 1 , dρ 2 be two Lie aglebra representations of interest. The tensor product on the Lie algebra dρ 1 ⊗ dρ 2 ( X ) can be computed as,

Therefore, given sets of generators of three representations dρ 1 , dρ 2 , dρ 3 , the Clebsch-Gordan coefficients are the change of basis between ( dρ 1 ( X ) ⊗ 1 + 1 ⊗ dρ 2 ( X )) and dρ 3 ( X ) . One can compute this change of basis numerically via a null space algorithm. For some groups, one can apply an iterative algorithm that generates all irreps starting with a single representation, using the above-mentioned procedure (see A.10).

5.3 Symmetric Powers

Let V be a vector space and { e i } be a basis of V .The symmetric power of V , Sym n V , can be regarded as the space of homogeneous polynomials of degree n in the variables e i . The product basis in Equation 10 spans exactly this space. A basis of Sym n V can be constructed as,

If V λ if an irreducible representation of a reductive Lie group G with highest weight λ , Sym n V λ admits a decomposition into irreducible representations,

The generalized Clebsch Gordan coefficients in (11) represent the change of basis between Sym n V λ and one of the V µ . The following steps are taken to obtain these coefficients:

Construct the symmetric power basis as in (22)

Compute the coefficients c λ,µ , using Freudenthal's Formula or GT patterns.

For any µ with c λ,µ non-zero, find a basis of V µ , and compute the change of basis between the basis of Sym n V λ and V µ

Alternatively, if one simply has the Clebsch Gordan coefficients, the change of basis from V λ to some V µ , a new algorithm outlined in Appendix A.9 and implemented in lie-nn can construct the change of basis from Sym n V λ to V µ .

6 Applications

6.1 Lie groups and their applications

In Table 6.1 we give a non-exhaustive overview of Lie groups and their typical application domains, to which our methodology naturally applies.

Benchmarking our method on all of these applications is beyond the scope of the present work, in particular, because most of these fields do not have standardized benchmarks and baselines to

compare against. The MACE architecture has proven to be state of the art for a large range of atomistic modeling benchmarks (Batatia et al. , 2022a). In the next section, we choose two new prototypical applications and their respective groups to further assess the performance of our general approach.

6.2 Particle physics with the SO (1 , 3)

Jet tagging consists in identifying the process that generated a collimated spray of particles called a jet after a high-energy collision occurs at particle colliders. Each jet can be defined as a multiset of four-momenta [( E i , p i )] N i =1 , where E i ∈ R + and p i ∈ R 3 .

Current state-of-the-art models incorporate the natural symmetry arising from relativistic objects, e.g, the Lorentz symmetry, as model invariance. To showcase the performance and generality of the G -MACE framework we use the Top-Tagging dataset (Butter et al. , 2019), where the task is to differentiate boosted top quarks from the background composed of gluons and light quark jets. In Table 6.2, we can see that G -MACE achieves excellent accuracy, being the only arbitrary equivariant model to reach similar accuracy as PELICAN, which is an invariant model. We refer to Appendix A.11.1 for the details of the architecture.

6.3 3D Shape recognition

3D shape recognition from point clouds is of central importance for computer vision. We use the ModelNet10 dataset (Wu et al. , 2015) to test our proposed architecture in this setting. As rotated objects need to map to the same class, we use a MACE model with O (3) symmetry. To create an encoder version of G -MACE, we augment a PointNet++ implementation (Yan, 2019) with G -MACE layers. See the appendix A.11.2 for more details on the architecture. We see in Table 3 that MACE outperforms the non-equivariant baseline.

7 Conclusion

We introduced the G -Equivariant Cluster Expansion, which generalizes the successful ACE and MACE architectures to symmetries under arbitrary reductive Lie groups. We provide an open-source Python library lie-nn that provides all the essential tools to construct such general Lie-group equivariant neural networks. We demonstrated that the general G -MACE architecture simultaneously achieves excellent accuracy in Chemistry, Particle Physics, and Computer Vision. Future development will implement additional groups and generalize to new application domains.

Acknowledgments and Disclosure of Funding

IB's work was supported by the ENS Paris Saclay. CO's work was supported by NSERC Discovery Grant IDGR019381 and NFRF Exploration Grant GR022937. This work was also performed using resources provided by the Cambridge Service for Data Driven Discovery (CSD3).IB would like to thank Gábor Csányi for his support.

References

LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel, Neural Computation 1 , 541 (1989).

S. Cohen and M. Welling, ICLR 2017 (2016a), 10.48550/ARXIV.1612.08498.

Esteves, C. Allen-Blanchette, X. Zhou, and K. Daniilidis, 'Polar transformer networks,' (2017).

Anderson, T. S. Hy, and R. Kondor, in Advances in Neural Information Processing Systems , Vol. 32, edited by H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (Curran Associates, Inc., 2019).

P. Bartók, R. Kondor, and G. Csányi, Physical Review B 87 , 184115 (2013).

Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky, Nature Communications 13 , 2453 (2022).

Batatia, D. P. Kovács, G. N. C. Simm, C. Ortner, and G. Csányi, 'Mace: Higher order equivariant message passing neural networks for fast and accurate force fields,' (2022a).

Bogatskiy, T. Hoffman, D. W. Miller, and J. T. Offermann, Machine Learning and the Physical Sciences workshop, NeurIPS 2022 (2022), arXiv:2211.00454 [hep-ph] .

Li, H. Qu, S. Qian, Q. Meng, S. Gong, J. Zhang, T.-Y. Liu, and Q. Li, (2022), arXiv:2208.07814 [hep-ph] .

Finzi, S. Stanton, P. Izmailov, and A. G. Wilson, in Proceedings of the 37th International Conference on Machine Learning , Proceedings of Machine Learning Research, Vol. 119, edited by H. D. III and A. Singh (PMLR, 2020) pp. 3165-3176.

Kondor and S. Trivedi, in Proceedings of the 35th International Conference on Machine Learning , Proceedings of Machine Learning Research, Vol. 80, edited by J. Dy and A. Krause (PMLR, 2018) pp. 2747-2755.

Bogatskiy, B. Anderson, J. T. Offermann, M. Roussi, D. W. Miller, and R. Kondor, 'Lorentz Group Equivariant Neural Network for Particle Physics,' (2020a), arXiv:2006.04780 [hep-ph] .

Drautz, Phys. Rev. B 99 , 014104 (2019).

Batatia, S. Batzner, D. P. Kovács, A. Musaelian, G. N. C. Simm, R. Drautz, C. Ortner, B. Kozinsky, and G. Csányi, 'The design space of e(3)-equivariant atom-centered interatomic potentials,' (2022b).

Cohen and M. Welling, in Proceedings of The 33rd International Conference on Machine Learning , Proceedings of Machine Learning Research, Vol. 48, edited by M. F. Balcan and K. Q. Weinberger (PMLR, New York, New York, USA, 2016) pp. 2990-2999.

S. Cohen, M. Geiger, J. Köhler, and M. Welling, in International Conference on Learning Representations (2018).

Weiler, P. Forré, E. Verlinde, and M. Welling, 'Coordinate independent convolutional networks isometry and gauge equivariant convolutions on riemannian manifolds,' (2021).

Dusson, M. Bachmayr, G. Csányi, R. Drautz, S. Etter, C. van der Oord, and C. Ortner, Journal of Computational Physics 454 , 110946 (2022).

Geiger and T. Smidt, 'e3nn: Euclidean neural networks,' (2022).

Finzi, M. Welling, and A. G. Wilson, 'A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups,' (2021).

Kondor, Z. Lin, and S. Trivedi, in Advances in Neural Information Processing Systems , Vol. 31, edited by S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (Curran Associates, Inc., 2018).

G. Satorras, E. Hoogeboom, and M. Welling, 'E(n) equivariant graph neural networks,' (2021).

Brandstetter, R. Hesselink, E. van der Pol, E. J. Bekkers, and M. Welling, 'Geometric and physical quantities improve e(3) equivariant message passing,' (2022), arXiv:2110.02905 [cs.LG] .

Drautz, Phys. Rev. B 102 , 024104 (2020).

Kaliuzhnyi and C. Ortner, ArXiv e-prints 2202.04140 (2022).

Shapeev, Multiscale Model. Simul. 14 , 1153 (2016).

P. Darby, D. P. Kovács, I. Batatia, M. A. Caro, G. L. W. Hart, C. Ortner, and G. Csányi, 'Tensorreduced atomic density representations,' (2022).

M. Gelfand and M. L. Tsetlin, 825828 , 71 (1950).

Alex, M. Kalus, A. Huckleberry, and J. von Delft, Journal of Mathematical Physics 52 , 023507 (2011).

P.-Y. Lagrave, Y. Cabanes, and F. Barbaresco, in Geometric Science of Information , edited by F. Nielsen and F. Barbaresco (Springer International Publishing, Cham, 2021) pp. 577-584.

Favoni, A. Ipp, D. I. Müller, and D. Schuh, Phys. Rev. Lett. 128 , 032003 (2022).

Bogatskiy, B. Anderson, J. Offermann, M. Roussi, D. Miller, and R. Kondor, in Proceedings of the 37th International Conference on Machine Learning , Proceedings of Machine Learning Research, Vol. 119, edited by H. D. III and A. Singh (PMLR, 2020) pp. 992-1002.

Butter et al. , SciPost Phys. 7 , 014 (2019), arXiv:1902.09914 [hep-ph] .

Qu, C. Li, and S. Qian, 'Particle transformer for jet tagging,' (2022).

Qu and L. Gouskos, Phys. Rev. D 101 , 056019 (2020), arXiv:1902.08570 [hep-ph] .

M. Munoz, I. Batatia, and C. Ortner, 'Bip: Boost invariant polynomials for efficient jet tagging,' (2022).

T. Komiske, E. M. Metodiev, and J. Thaler, JHEP 01 , 121 (2019), arXiv:1810.05165 [hep-ph] .

Pearkes, W. Fedorko, A. Lister, and C. Gay, (2017), arXiv:1704.02124 [hep-ex] .

Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao, in 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (IEEE Computer Society, Los Alamitos, CA, USA, 2015) pp. 1912-1920.

X. Yan, github.com/yanx27 (2019).

C. R. Qi, H. Su, K. Mo, and L. J. Guibas, arXiv preprint arXiv:1612.00593 (2016).

C. R. Qi, L. Yi, H. Su, and L. J. Guibas, arXiv preprint arXiv:1706.02413 (2017).

Y. Shen, C. Feng, Y. Yang, and D. Tian, in CVPR (2018).

J. Li, B. M. Chen, and G. H. Lee, CoRR (2018).

S. Kumawat and S. Raman, in CVPR (2019).

I. Molev, 'A weight basis for representations of even orthogonal lie algebras,' (1999), arXiv:math/9902060 [math.RT] .

R. Steinberg, Bull. Amer. Math. Soc. 67 , 406 (1961).

M. Bachmayr, G. Dusson, C. Ortner, and J. Thomas, ArXiv e-prints 2109.14771 (2021).

J. Thomas, H. Chen, and C. Ortner, Arch. Ration. Mech. Anal. 246 (2022).

T. Broecker, Representation of Compact Lie Groups (Springer Berlin Heidelberg, 1985).

N. Bourbaki, Lie Groups and Lie Algebras (Springer Berlin, Heidelberg, 1989).

A Appendix

A.1 Background on Lie groups and Tensor Products of Representations

A.1.1 Lie groups

Lie groups generalize the idea of continuous symmetry such as the rotation symmetry of the round sphere S n (given by the action of the orthogonal group O( n ) ), the symmetry of a vector space under changes of basis (given by the action of the general linear group GL n ( R ) ), or the invariance of Minkowski spacetime under boosts (given by the Lorentz group). The symmetry groups of specific geometries were studied in the 19 th century, leading to Sophus Lie's study of general continuous symmetry. Other major contributors to our understanding include Klein, Riemann, Hilbert, Poincarré, Noether, Cartan, among many others. In particular the solution of Hilbert's 5 th problem by Gleason, Montgomery, and Zippin shows that under very general conditions a continuous symmetry group must be a Lie group. Symmetry being central to mathematics Lie groups appearl in many areas of mathematics, from differential geometry to number theory.

Let us give the general formal definition, as well as a more practical concrete one:

Definition A.1. A real (complex) Lie group is a group that is also a finite-dimensional smooth (complex) manifold in which the product and inversion of the group are also smooth (holomorphic) maps.

A linear Lie group is a subgroup of the group GL n ( R ) which is also a closed subset (in the sense of real analysis).

It is a fact that every closed subgroup of a Lie group is a Lie group, so every group of the second type is indeed also a group of the first type. Conversely while in general a Lie group is only locally isomorphic to a linear group (by a surjective group homomorphism), every compact Lie group is isomorphic to a linear group, and most Lie groups that arise in practice are linear.

When the Lie group G acts on a space X (e.g. as a symmetry group) it naturally also acts on spaces of functions on X (by translation), and this action is linear . In general. an action of G on a (suually complex) vector space by linear map is called a linear representation (of G ). In the context of our paper the space of functions on X is the space of features, and we achieve efficient computation in it by decomposing this representation into simpler constituents - in other words by understanding the representation theory of G .

Describing the representation theory of arbitrary Lie groups is a very difficult problem and an active area of research in pure mathematics. However, for certain classes of groups it is possible to say quite a bit (for example classical harmonic analysis can be viewed from the lens of the representation theory of commutative groups). An important class of groups whose representation theory can be completely understood is the class of reductive Lie groups, on which we focus in this paper.

For now let us give a concrete definition; an abstract will follow later once we introduce the necessary language.

Definition A.2. A linear Lie group G ⊂ GL n ( R ) is reductive if it is also closed under transpose: for every g ∈ G we have that g T ∈ G .

Most of the groups that are of interest in natural sciences are reductive Lie groups. This includes well-known groups such as the symplectic group Sp 2 n ( R ) relevant in Hamiltonian dynamics, and orthogonal group O( n ) which parametrizes rotations in R n .

A.1.2 Linear representations

Fix a Lie group G . A linear representation of G is a pair ( π, V ) where V is a complex vecctor space and π is an action of G on V by linear maps. In other words for each g ∈ G we have a complex-linear map π ( g ): V → V such that π ( g 1 g 2 ) = π ( g 1 ) π ( g 2 ) and such that π (1 G ) = Id V (in other words, π : G → GL( V ) is a group homomorphism).

For an example let the familiar rotation group O(3) act on the usual round sphere S 2 . Then O(3) acts on the space of functions on S 2 . This space is infinite-dimensional, but it turns out we can approximate it by a sum of finite-dimensional pieces. Indeed for each degree d ≥ 0 let V d be the space of polynomials in the variables x, y, z which are homogeneous of degree d (for example

V 0 = span { 1 } , V 1 = span { x, y, z } , V 2 = span { x 2 , xy, y 2 , . . . } . Letting O(3) act on R 3 by linear change of variable, it acts on each V d since such change of variable does not change the degree of homogeneity. Now let W d ⊂ V d be the subspace of harmonic polynomials, those polynomials p that satisfy ∂ 2 p ∂x 2 + ∂ 2 p ∂y 2 + ∂ 2 p ∂z 2 = 0 . Then W d is itself O(3) invariant since the Laplace operator is rotation-invariant, making W d a subrepresentation of V d . It is a (not immediately obvious) fact that each W d is in fact an irreducible representation: it has no subrepresentations of its own other than itself and the zero subspace. Further, if we interpret the harmonic polynomials in W d as functions on the unit sphere S 2 ⊂ R 3 we can view W d as a space of functions on the sphere. It then turns out that the spaces W d are linearly independent of eaceh other, and that their sum ⊕ d ≥ 0 W d is dense in any resaonable space of functions on the sphere (the space of continuous functions, the space of smooth functions, or the space of square-integrable functions, each with its respective notion of convergence). In other words, we have morally decomposed the space of functions on S 2 as the sum of O(3) -invariant subrepresentations, each of which is irreducible.

Decomposing a representation into its irreducible constituents is a major goal of representation theory (as well as a computational goal of this paper). When this is possible we can understand a general representations V of a group G by first enumerating its irreducible representations and then counting how many times each irreducible occurs as summand in V , known as the (this is very much analogous to understanding general positive integers via their prime factorization - here the irreducible representations play the role of the prime numbers).

Reductive groups are particularly amenable to this approach because their representations indeed do decompose as direct sums in this fashion. It is fact that a linear Lie group G is reductive if and only if each of its finite-dimensional representations is isomorphic to a direct sum of irreducible representations.

Before we delve further into representation theory we need a brief detour into the structure theory of Lie groups.

A.1.3 Lie algebras

As Lie groups are manifolds, understanding their representations theory usually leads to advanced geometrical and topological results. We will now describe how these difficult questions can be transformed into simpler linear algebra questions using the Lie algebras associated with Lie groups.

A Lie algebra is a vector space g over a field F together with a bilinear map [ · , · ] : g × g → g called the Lie bracket, satisfying the following axioms:

The Lie bracket is bilinear, that is for all x, y, z ∈ g and α, β ∈ F , [ αx + βy, z ] = α [ x, z ] + β [ y, z ] and [ x, αy + βz ] = α [ x, y ] + β [ x, z ] .

The Lie bracket is antisymmetric, that is for all x, y ∈ g , [ x, y ] = -[ y, x ] .

The Lie bracket satisfies the Jacobi identity, that is for all x, y, z ∈ g , [ x, [ y, z ]] + [ y, [ z, x ]] + [ z, [ x, y ]] = 0 .

Lie algebras can alternatively be seen as the tangent space of a Lie group to the identity:

Theorem A.3. Let G be a Lie group viewed as a smooth manifold, the tangent space T I G is a Lie algebra called the Lie algebra of G noted g . Moreover in the case of connected complex Lie groups, the category of representation of G , Rep ( G ) , is equivalent to the category of representation of g , Rep ( g ) .

This theorem is at the heart of the representation theory of Lie groups. As Lie algebra are vector spaces, the functorial equivalence between the two categories of representation translates geometrical and topological questions into simpler linear algebra questions.

A classification of Lie algebras was carried out by Cartan and later Gelfand work. In the case of complex Lie algebras, this classification is complete. Two major classes of Lie algebras emerged from Cartan's work, semi-simple Lie algebras and nilpotent Lie algebras.

Definition A.4 (Central series) . A central series of a Lie algebra g is a sequence of ideals ( g i ) i ∈ N such that g 0 = g and g i +1 = [ g i , g ] .

Definition A.5. A Lie algebra is called semi-simple if it does not contain any non-trivial ideals. A Lie algebra is called nilpotent if its central series terminates to zero.

An important result of this classification is that any Lie algebras over an algebrically closed field can be decomposed into semisimple parts and a nilpotent part. We call this the Jordan decomposition. Therefore, understanding arbitrary representations of a Lie algebra boils down to understanding representation of both semi-simple and nilpotent algebras.

Most reductive Lie groups of interest have a semi-simple Lie algebra as associated Lie algebra. These algebras are also the one for which the representation theory is the best understood.

The study of Lie algebras is decomposed in two subfields. The structure theory studies the classification of Lie algebras in terms of their eigenspaces (as vector spaces). The representation theory builds on the structure theory to classify representations of semi-simple Lie algebras. We will give a very short introduction to both of them.

A.1.4 Structure theory of semi-simple Lie algebras

Let g be a semi-simple complex Lie algebras. Let h be one maximal set of mutually commutative element of g . Maximal means here that no other such set contains strictly more element than h . We call h the Cartan algebra of g . As elements of h are mutually commutative, they share the same eigenspaces.

Therefore one can decompose g in terms of eigenspaces of elements of h . Let α ∈ h ∗ be a function in the dual of the Cartan algebra such that ∀ h i ∈ h , α ( h i ) = α i with α i the eigenvalue associated to h i . We call this linear operator the roots of g . Let R be the space of all this functions, called the root space. One therefore has the the following decomposition of the Lie algebra;

The structure theory of semi-simple complex Lie algebras is the study of the root space R and the associated eigenspaces g α . One can clasify complex semi-simple Lie algebras based on their type of root spaces. For real semi-simple Lie algebras, the structure theory essential reduces down to the complex case with a few extra steps.

A.1.5 Representation theory of complex semi-simple Lie algebras and Lie groups

The representation theory of complex reductive Lie algebras is completely understood. Every finite-dimensional representation is the direct sum of irreducible representations. Every irreducible representation is uniquely determined by a highest weight λ , which is an element of the dual of the Cartan algebra h . The dimension of each irreducible representation can be computed using the Weyl dimension formula.

where R + is a subset of the root space R , called the positive roots, and ρ is the half sum of the positive roots. Obtaining explicitly a basis for each representation is much harder. Gelfand and Tselin proposed a basis for the general case of SU ( N ) groups. This basis was recently generalized to other complex semi-simple algebras Molev (1999).

A.1.6 Representations of real semi-simple Lie groups

The representation theory of real semi-simple Lie groups can be elucidated by examining the complex semi-simple case through the application of the so-called Weyl unitary trick. This technique establishes a connection between the representations of real Lie groups and those of the complexification of their universal covering groups.

Definition A.6. Given a Lie algebra g , there exists a unique simply connected group ˆ G possessing g as its Lie algebra. A group homomorphism ψ exists from this unique simply connected Lie group to any group G sharing the same Lie algebra. We refer to ˆ G as the universal covering Lie group, and it is unique up to isomorphism.

An essential observation is that the irreducible representations of G form a subset of the irreducible representations of ˆ G . Understanding the irreducible representations of ˆ G suffices. Take the case of

the su (2) Lie algebra, for instance. The universal covering group is SU(2) , corresponding to the special unitary two-by-two matrices. The group SO(3) of special orthogonal matrices also has su (2) as its Lie algebra. Consequently, the irreducible representations of SO(3) are a subset of those of su (2) . Representations of su (2) can be indexed by a half-integer, known as the quantum number j . The integer representations are also irreducible representations of SO(3) , with non-integer ones being referred to as the spin representations.

Since the representation theory of complex Lie groups is well understood, an optimal approach to understanding the representation theory of a real Lie group involves studying the complexification of its universal cover.

Definition A.7. For a universal covering Lie group G with Lie algebra g , the universal covering Lie group ˆ G C , having the associated Lie algebra g C = g ⊗ C , is termed the complexification of ˆ G .

The group of interest G , its universal cover ˆ G and the complexification of the universal cover ˆ G C , all share the similar irreducible representations. Therefore one obtains the following chain of inclusion of the respective representations of the groups,

Let K be the maximal compact subgroup of ˆ G C and ρ λ a finite dimensional irreducible representation of K , with λ as its highest weight. It is possible to analytically continue ρ λ into an irreducible representation of ˆ G C , denoted as ρ ˆ G C λ , and also into the analytical conjugate ¯ ρ ˆ G C λ . Every finitedimensional representation of ˆ G C can be constructed as the tensor product form ρ ˆ G C λ ⊗ ¯ ρ ˆ G C λ ′ . Due to the fact that finite-dimensional representations of the universal covers and complexification of a semi-simple Lie group are isomorphic, one can express:

As universal covering groups contain more representations than the original cover, this mapping is not an isomorphism for all lambda. To illustrate this, consider the example of the Lorentz group SO (1 , 3) .

The universal cover of SO(1 , 3) is SL(2 , C ) , and the maximal compact subgroup of SL(2 , C ) is SU(2) . The irreducible representations of SU(2) are indexed by l ∈ N / 2 and correspond to the Wigner D matrices D l ( g ) , g ∈ SU(2) . The SU(2) group elements are parameterized by Euler angles that can be analytically continued to SL(2 , C ) as follows:

This enables the construction of the fundamental representation of SL(2 , C ) , resulting in representations of the Lorentz group corresponding to the product of Wigner D matrices:

The irreducible representations of SO(1 , 3) are indexed by a pair of integers ( l, k ) corresponding to the associated SU (2) representations in the tensor product.

A.1.7 Gefland-Tselin Basis

The enumeration of finite-dimensional representations of complex Lie groups has been well understood. However, the derivation of explicit matrix representations in specific bases and the tensor product decomposition of representations within these bases remain challenging and subject to ongoing research. The Gelfand-Tselin basis provides an explicit structure for irreducible representations, initially designed for the SU ( N ) group Gelfand and Tsetlin (1950); Alex et al. (2011), and subsequently generalized to encompass other classical Lie groups Molev (1999).

Central to the concept of GT patterns is the application of branching rules, whereby representations of a larger group are constructed by enumerating induced representations on its subgroups. Consider, for instance, the SU ( N ) case where irreducible representations are characterized by the highest weight vector λ , a vector of length N terminated by a zero, λ = ( λ 1 , . . . , λ N -1 , 0) . This highest weight

vector may spawn several induced representations of the subgroup SU ( N -1) , each described by N -1 vectors. This process can be iteratively executed in a descending chain:

This chain commences with the group of interest and culminates with the smallest subgroup. From this branching chain, triangular arrays known as GT-patterns can be constructed, with each row corresponding to a vector of a representation ranging from SU ( N ) to U (1) :

Adjacent rows within these patterns must adhere to a so-called "snake rule" to form an admissible GT pattern:

The dimension of a specific irreducible representation corresponds to the count of such admissible arrays, with the top row representing the highest weight vector. Representations sharing the same highest weight vectors interrelate through ladder operators. One key attribute of GT patterns is that the action of ladder operators on a GT pattern is analytically known, enabling the construction of the matrix representation for all representations using only the highest weight representation. Furthermore, constructing the highest weight matrix representation equals solving a linear system when the ladder operators are known analytically. Extending this methodology to other classical Lie groups (such as ( n ) , Sp( n ) , SL( n ) ) presents challenges as the ladder operators' entries become increasingly complex and recursive.

A.2 Example of one particle basis for O(3) and SO(1,3) groups

In the paper, the one-particle basis ϕ is left abstract, as it depends on the specific type of input point cloud and the group under consideration. We now describe two concrete examples, one involving isometry the other Lorentz equivariance.

A.2.1 Isometry equivariance

In the first example we consider a point cloud comprised of atoms, for example representing a molecule. Here, the point states are denoted by x i = ( r i , Z i ) ∈ R 3 × Z , where r i corresponds to the position in 3D and Z i , its atomic number, is a categorical variable. The one-particle basis must form a complete basis for R 3 × Z , which is both rotationally equivariant and translationally invariant. The translation is typically handled by referring to the atom's state in relation to a central value or another point. For simplicity, we can assume r i the particle's position relative to the point cloud's centre of mass, ensuring translational invariance.

The space R 3 can be represented as the product R × S 2 , with S 2 being the sphere. As such, we choose a basis for R such as Bessel functions, Chebyshev polynomials or a Fourier basis, denoted as R n . For S 2 , we utilise the spherical harmonics Y lm , which form a basis of functions in the sphere. Consequently, the one-particle basis can be expanded as:

where r i corresponds to the norm of r i and ˆ r i = r i r i .

The case of the Lorentz group is more complex than the (3) case. Consider a point cloud made up of elementary particles, represented by their 4-momenta x i = ( E i , p i ) ∈ R 4 . Although the 4-momenta is a vector, the natural metric in this space is not the Euclidean metric but rather a pseudo metric known as the Minkowski metric. This pseudo-metric encapsulates the unique role that time plays in special relativity. In natural units, the pseudo norm η of a 4-vector u = ( u 0 , u 1 , u 2 , u 3 ) is defined as

To construct the one-particle basis, it is necessary to expand a four-vector in terms of the basis of representations of the Lorentz groups. As outlined in Appendix A.1.6, irreducible representations

of the Lorentz group are characterised by a tuple of integers ( l, k ) . The 4-vectors correspond to the fundamental, irreducible representation, namely the (1 , 1) representation. Analogous structures to the spherical harmonics can be obtained by examining the symmetric powers of the irreducible representations. These are analogous to harmonic polynomials on Minkowski space, much as spherical harmonics are harmonic polynomials on the sphere S 2 . We thus define the Minkowski spherical harmonics as the following symmetric tensor product,

These spherical harmonics form a harmonic polynomial basis on Minkowski space. Therefore, one can create the one-particle basis in a manner akin to the O (3) case, as follows:

Where p i = η ( p i , p i ) , and ˆ p i = p i p i + ϵ , with ϵ being a small constant that prevents divergence when p i is zero due to the non-positiveness of the Minkowski norm.

A.3 Extended Related Work

Convolutional Neural Networks (CNNs), which are translation equivariance, initiated the utilization of data symmetry in machine learning architectures. Throughout time, CNNs have been extended to include other symmetries as well. Central to all these generalizations is the group averaging operation,

Where x denotes the input signal or feature, f is the convolution kernel, G represents the group of interest, and dg is an invariant measure on G . This transformation is essential, as it converts any convolution into a group invariant convolution. The feasibility of this approach largely depends on the computational simplicity of the integral. The most straightforward instance occurs for finite groups, as explained by G-convolutions (Cohen and Welling, 2016b), where the integral simplifies to a sum,

In the context of G being a compact group, the invariant measure dg is unique, referred to as the Haar measure of the group, allowing the integral to be computed numerically, e.g., on a grid. Steerable CNNs (Cohen and Welling, 2016a), LieConv Finzi et al. (2020) and Spherical CNNs (Cohen et al. , 2018) extended CNNs to O (2) and O (3) equivariance, respectively. A general theory for convolution on any compact group and symmetric space was developed by Kondor and Trivedi (2018), and further extended to equivariant convolutions on Riemannian manifolds by Weiler et al. (2021). However, this approach has several limitations,

The direct computation of the integral can be numerically unstable and inefficient, even for relatively small groups like O (3) .

In the case of non-compact groups, a unique invariant measure is absent, and the integral diverges.

Across these methods, the convolution kernel f is usually constrained to a two-body operator

.

In the case of compact groups, the integral over the group may be calculated via an alternative means. There exists a linear operator, called the Clebsch Gordan operator C , such that,

Therefore, the complex integral over the group becomes a linear operation. The form of this operator depends on the basis in which f is expended. In the case of G = O (3) and if this basis is carefully chosen, the entries of this operator are known analytically. This approach is numerically stable and more efficient and was taken by numerous works, including ACE Dusson et al. (2022), Cormorant Anderson et al. (2019), NequIP Batzner et al. (2022), or MACE Batatia et al. (2022b). The central aim of our work is to show that this approach can also be generalized to all reductive Lie groups, even non-compact ones, and provide tools to do so.

A.4 Proof of (7)

This statement follows closely the arguments by Dusson et al. (2022); Drautz (2020) and others.

A.5 Custom notation and indexing

We briefly contrast our notation for Clebsch-Gordan coefficients (10) with the standard notation. By means of example, consider the group SO (3) in which case the Clebsch-Gordan equations are written as

In this setting, our index α simply enumerates all possible such coefficients. One can often assign a natural meaning to this index, e.g., for the group SO (3) it is given by the pair of angular quantum numbers ( l 1 , l 2 ) . Specifically, in this case, we obtain

where C LM l 1 m 1 l 2 m 2 are the Clebsch-Gordan coefficients in the classical notation. Thus, the additional index α is not really required in the case of SO (3) , nor our other main example, SO (1 , 3) . Our notation is still useful to organize the computations of equivariant models, especially when additional channels are present, which is usually the case. Moreover, it allows for easy generalization to other groups where such a simple identification is not possible (Steinberg, 1961).

A.6 Equivariance of G-cluster expansion

The equivariance of the G-cluster expansion is easily verified by applying a transformation g to the input,

A.7 Completeness of the basis and Universality of G-MACE

We explain in which sense the basis B K α is a complete basis, and briefly sketch how to prove this claim. The argument is contained almost entirely in (Dusson et al. , 2022) and only requires a single modification, namely Step 3 below, using a classical argument from representation theory. We will therefore give only a very brief summary and explain that necessary change.

We start with an arbitrary equivariant property Φ V embedded in V where we have a representation, i.e. the actual target property is Φ is then given as a linear mapping from V to Z . For technical reasons, we require that only finitely many entries Φ V K may be non-zero, but this is consistent with common usage. For example, if G = O (3) and if Φ is a scalar, then Φ V 0 = Φ , while all other Φ V LM ≡ 0 . If Φ is a covariant vector, then Φ V LM is non-zero if and only if L = 1 ; and so forth. For other groups, the labeling may differ but the principle remains the same.

1. Convergence of the cluster expansion. The first step in our parameterisation is to approximate Φ V in terms of a truncated many-body expansion (4). It is highly application-dependent on how fast this expansion converges. Rigorous results in this direction in the context of learning interatomic potentials can be found in (Bachmayr et al. , 2021; Thomas et al. , 2022). A generic statement can be made if the number of input particles is limited by an upper bound, in which case the expansion becomes exact for a finite N . This case leads to the uniform density result stated in Theorem 4.1. We adopt this setting for the time being and return to the pointwise convergence setting below.

In the uniform convergence setting we also require that the domain Ω is compact.

Throughout the remainder of this section we may therefore assume that an N can be chosen as well as smooth components φ ( n ) such that the resulting model Φ V,N approximates Φ V to within a target accuracy ϵ ,

The density of the embedding. As already stated in the main text, if the components φ ( n ) K are smooth, and the embedding { ϕ k } k is dense in the space of one-particle functions (5) then it follows that the φ ( n ) K can be expanded in terms of the tensor product basis ϕ k := ⊗ n s =1 ϕ k s to within arbitrary accuracy. The precise statement is the following standard result of approximation theory: if span { ϕ k } k are dense in C (Ω) , then span { ϕ k } k are dense in C (Ω n ) . That is, for any ϵ > 0 , there exist approximants p ( n ) K such that

3. The density of the symmetrized basis. The next and crucial step is to show that, if the φ ( n ) K are equivariant, then the p ( n ) K may be chosen equivariant as well without loss of accuracy. If the group G is compact then the representations ρ can be chosen unitary (Broecker, 1985). In that case, the argument from (Dusson et al. , 2022) can be used almost verbatim: let

where H is the normalized Haar measure then ¯ p ( n ) is equivariant by construction and

If the group is not compact, then one can apply 'Weyl's Unitary Trick' (see (Bourbaki, 1989), Ch. 3): first, one complexifies the group (if it is real) and then constructs a maximal compact subgroup K C of the complexification. This new group K will have the same representation as G and because it is compact, that representation may again be chosen as unitary. Therefore, symmetrizing p ( n ) with

respect to K C results in an approximant that is not only equivariant w.r.t. K C but also equivariant w.r.t. G .

The density of the basis B K α . As the last step, one can readily observe that the symmetrization and cluster expansion steps can be exchanged. I.e. first symmetrizing and then employing the steps (7) result in the same model. Letting ϵ → 0 in the foregoing argument while fixing the number of particles # x results in all errors vanishing. Note that this will in particular require taking N →∞ .

Pointwise convergence. To obtain density in the sense of pointwise convergence we first introduce the canonical cluster expansion without self-interacting terms

The difference here is that the summation is only over genuine sub-clusters. Because of this restriction, the series is finite for all multi-set inputs x . In other words, it converges in the pointwise sense.

One can easily see that v n can be chosen (explicitly) to make this expansion exact. After truncating the expansion at finite n ≤ N and then expanding the potentials v ( n ) K one can exactly transform the canonical cluster expansion into the self-interacting cluster expansion. This procedure is detailed in (Dusson et al. , 2022; Drautz, 2020).

The arguments up to this point establish the claimed universality for the linear ACE model. The corresponding universality of the TRACE model follows immediately from (Darby et al. , 2022). Since a single layer of the MACE model is a TRACE model, this completes the proof of Theorem 4.1.

A.8 Product of groups

Let G 1 and G 2 be two reductive Lie groups (or finite groups). Let A 1 and A 2 be the two structure constants of G 1 and G 2 , ( dρ 1 , ρ 1 ) and ( dρ 2 , ρ 2 ) their continuous and discrete generators. One can define a representation of the direct product group G 1 × G 2 as

The following essential property holds: if ρ 1 and ρ 2 are irreducible representations of G 1 and G 2 then ρ 1 ⊗ ρ 2 is an irreducible representation of G 1 × G 2 . Moreover, for reducible Lie groups, all the irreps of G 1 × G 2 are of this form. Therefore, one can construct all the irreps of G 1 × G 2 this way. It is of particular interest in the case of equivariant message passing networks on points clouds, where the group of interest is G × S n .

Let us give a non-trivial application of lie-nn for computing invariants for the product group of O (3) × S 3 . We consider a set of three vectors,

A given vector belong to the product representation of the l = 1 vector of the O (3) group and the natural representation of S 3 . In lie-nn, we can define the product representation as follows,

If one wants to know the unique permutation invariant vector that one can construct from the three vectors a , b and c , one can use the following code,

As expected, the output is the sum of the three vectors, a + b + c , corresponding to the only permutation invariant vector that can be constructed from the three vectors.

Consider now the case where we want to extract a permutation invariant scalar from the set of products of two vectors. In this case, we can use the following code,

In this code, we first construct the tensor product representation and then seek the permutation invariant scalar in it. The result is two distinct scalars,

Finally we consider a case that would be very hard to do by hand. It is the case of extracting a permutation invariant l = 2 vector for this tensor product of the original set of three vectors.

The computation tells us that there exists two such vectors,

A.9 Symmetric Tensor products

The permutation group is an important concept in the context of tensor products. It can be useful to focus on a subset of the full tensor product space that exhibits certain permutation equivariance. For example, the spherical harmonics are defined as the permutation-invariant part of a tensor product.

The symmetric tensor product can be thought of as a change of basis, or projector, from the tensor product to the symmetric part of the tensor product. In the case of a tensor product of correlation order four we have,

where B is the change of basis that satisfies:

We propose in lie-nn a new algorithm used to calculate B . The Symmetric Tensor Product is calculated using a tree structure, starting at the leaves and progressing towards the trunk. The leaves are the basis of the individual indices, and they are combined and constrained at each step to impose symmetry.

A.10 Computing the irreps from input representations

For some groups, the computation of the generators X can become a very involved task. However in most applications, the data itself is already given in a form of a representation. One approach proposed by (Finzi et al. , 2021) is to not work in the space of irreps but the space of polynomials of the input representation. This approach has the advantage of requiring little previous knowledge of

the group. However it is also much less efficient than using irreps. One alternative way is to consider polynomials of the input representation, that are reducible and then compute the block diagonalisation to project down to irreps subspace. One can then work directly as polynomials in this subspace and compute Clebsch-Gordan coefficients numerically. We provide routines in lie-nn to carry out these operations from any given input representation.

A.11 Details of numerical experiments

A.11.1 Jet Tagging

Dataset The dataset (Butter et al. , 2019) was generated using a Pythia , Delphes , and FastJet (using cuts for the jet's kinematics on ∆ η = 2 , R = 0 . 8 ) to simulate the response of the ATLAS detector at the Large Hadron Collider (LHC). The dataset is released under the "Creative Commons Attribution 4.0" license. The entire dataset contains 2 millions jets with a 60/20/20 for training, validation, and testing balanced splits.

Model The model uses 3 layers of the G -MACE architecture to generate the Lorentz group equivariant representation of each jet. For the 1 particle basis, we use a product of radial features on the Minkowski distances, and SO (1 , 3) spherical harmonics. The radial features are computing by passing a logarithmic radial basis as in (Bogatskiy et al. , 2022) into a [64 , 64 , 64 , 512] MLP using SiLU nonlinearities on the outputs of the hidden layers. The internal representations used are (0 , 0) and (1 , 1) . We use 72 channels for each representation. For the embedding, and readout out, we use similar achitectures to LorentzNet.

Training Models were trained on an NVIDIA A100 GPU in single GPU training. Typical training time for the dataset is up to 72 hours. Models were trained with AMSGrad variant of Adam, with default parameters of β 1 = 0 . 9 , β 2 = 0 . 999 , and ϵ = 10 -8 . We used a learning rate of 0 . 0035 and a batch size of 64. The model was trained for 80 epochs with 2 epochs of linear learning rate warmup and followed by a phase of cosine annealing LR scheduling.

A.11.2 3D shape recognition

Dataset ModelNet10 (Wu et al. , 2015) is a synthetic 3D object point clouds dataset containing 4,899 pre-aligned shapes from 10 categories. The dataset is split into 3,991 (80%) shapes for training and 908 (20%) shapes for testing. We were unable to find a license.

Model The model uses a three-layer encoder architecture following the PointNet++ one. We use an encoder of the full point cloud into sub-point clouds of sizes [1024 , 256 , 128] . Each PointNet layer maps a point cloud of size N t to one of size N t +1 . We compute the node features as the sum of the PointNet output and the MACE output,

Training Models were trained on an NVIDIA A100 GPU in single GPU training. The typical training time for the dataset is up to 12 hours. Models were trained with the AMSGrad variant of Adam, with default parameters of β 1 = 0 . 9 , β 2 = 0 . 999 , and ϵ = 10 -8 .

A.12 Limitations and Future Work

The spectrum of potential applications of the present method is very large. In this paper, we focus on a subset of applications that have known benchmarks and baselines. A broader range of groups is implemented in the lie-nn library. Future work should focus on applying this architecture to tasks with domain-specific knowledge.

A.13 Computational cost

A.13.1 Clebsch Gordan generation

Generation time The Clebsch Gordan generation time depends on the size of the representations. In Table 3, we plot the generation time as a function of the size of the representations for different groups. The generation time is found to range from a matter of milliseconds for smaller representations to a few seconds for larger ones. The generation of CGs constitutes a preprocessing step, separate from the model inference computations. Therfore, the results can be stored for subsequent use, ensuring that this phase does not affect the model's overall performance.

0

25

50

75

100

Representations size

0.000

0.025

0.050

0.075

0.100

0.125

0.150

0.175

0.200

Time (s)

SU(2)

0

1000

2000

3000

Representations size

0

5

10

15

20

25

SU(3)

0

500

1000

1500

2000

Representations size

0

5

10

15

20

25

SU(4)

Sparsity One distinguishing feature of the Clebsch-Gordan coefficient is their sparsity. This sparsity comes from the selection rule, which is unique for each group. These selection rules also give rise to a sparsity structure that can be exploited to achieve the best efficiency. In Figure 4 we compare the sparsity percentage as a function of the size of the representations for different groups.

0

20

40

60

80

100

Representations size

0.125

0.150

0.175

0.200

0.225

0.250

0.275

0.300

0.325

Sparsity %

SU(2)

0

1000

2000

3000

4000

5000

Representations size

0.025

0.050

0.075

0.100

0.125

0.150

0.175

0.200

0.225

SU(3)

0

1000

2000

3000

4000

Representations size

0.02

0.04

0.06

0.08

0.10

0.12

0.14

0.16

SU(4)

A.13.2 G-MACE computational cost

The computational bottleneck of traditional equivariant MPNNs is the equivariant tensor product on the edges of the graph used to form the message. In G-MACE the edge operation is the pooling operation to compute the features A k . Correlations A k are computed through the tensor product of the product basis, an operation that is carried out on nodes. In typical models, the correlations are the bottleneck. Since the number of nodes is orders of magnitudes smaller than the number of edges in most applications we are envisioning, it is a significant computational advantage to organize the computation in this way. We return to reviewing the cost of the product basis below.

2

3

4

0.0005

0.0010

0.0015

0.0020

0.0025

0.0030

0.0035

0.0040

Time / per jet (s)

lmax

=(1,1)

lmax

=(2,2)

lmax

=(3,3)

Correlation

Obtaining the equivariant basis B K involves simply linear operation. The Clebsch-Gordan coefficients are very sparse, with a well defined structural sparsity, which can be easily exploited for constructing highly efficient low level code on both CPUs and GPUs.

Thus, we return to the product basis A k , which is normally the computational bottleneck. It has structural sparsity due to its symmetry under permuting the k tuples and due to the structural sparsity in the Clebsch-Gordan coefficients. Despite that sparsity it was shown in Kaliuzhnyi and Ortner (2022) that in theory the product basis can be evaluated recursively with O (1) operations per feature, however, this requires effective use of sparse tensor formats which is challenging to optimize for GPU architectures since a naive implemention relies on random memory access patterns. Instead, the current MACE code employs a similar recursive evaluation algorithm described in Batatia et al. (2022a), which employs full tensor formats, but avoids the random memory access. For the most common case of 3-correlation models we estimate that this code reaches within a factor 3-5 of the optimal performance of hypothetical sparse tensor implementation.

Peer reviews:
comments: SUMMARY:
The paper presents a novel and highly general Equivariant Neural Network (ENN) architecture that is capable of respecting the symmetries of the finite-dimensional representations of any reductive Lie Group G. The proposed approach generalizes the successful ACE and MACE architectures for atomistic point clouds to any data equivariant to a reductive Lie group action. The authors demonstrate the generality and performance of their approach by applying it to two different tasks, namely top quark decay tagging (Lorentz group) and shape recognition (orthogonal group). The results presented in the paper are convincing and showcase the potential of the proposed architecture.

STRENGTHS:
* The paper is well-organized and well-written, providing a clear overview of the problem and proposed solution.
* The paper addresses a significant and challenging problem by proposing a highly general ENN architecture that can handle symmetries of the finite-dimensional representations of any reductive Lie Group G, which contributes to the field of equivariant neural networks.

WEAKNESSES:
Related concerns are discussed in the questions section.

QUESTIONS:
* A more comprehensive discussion comparing methods based on Lie Groups, such as LieConv, should be included in the paper.
* The results presented in the paper might not adequately convey the effectiveness of the proposed method. In order to reinforce the claims and emphasize the practical applicability of the approach, it would be valuable for the authors to incorporate comparisons of experimental results from methods like MACE and LieConv, using examples such as those from the QM9 dataset.

LIMITATIONS:
There are no potential negative societal impacts of this work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a class of G-equivariant neural networks for reductive lie groups that generalizes the previous ACE and MCAE models (which are designed to be equivariant with respect to the orthogonal group O(3)) to more general irreducible Lie groups.  A software library (Lie-NN) has also been developed and released for implementation of the proposed method. The experimental results on jet tagging and 3d point cloud classification support the authors' claims about the equivariant model improving model performance.\n\nSTRENGTHS:\n- The proposed approach seems novel; I am not aware of other works that design a G-equivariant network in this way.\n- The paper is generally well-written and technically correct (although I have not performed an in-depth check for the proof of the universality theorem). Occasional variable definitions are missing but I had no trouble following the general flow of the paper.\n- The experimental results appear to validate the claimed advantages of the proposed method (although more discussion on the computational aspects is needed to assess practicality; see below).\n\nWEAKNESSES:\n- Since this work is a generalized version of the previous ACE and MACE methods, a brief review of ACE and MACE would be helpful to understand and evaluate the novelty of the method. The current version is presented in a way that contains elements of previous existing methods, without a clear delineation of the new and original aspects of the proposed method. Exactly what is the precise nature of the extensions to ACE and MACE should be made clearer.\n- While the technical contents seem correct, some effort at providing intuitive explanation and justifications at key places would be helpful to understanding the paper.  Descriptive figures, intuitive examples (e.g., basis including 1-particle basis or Clebsch-Gordan coefficients) come to mind as examples. In particular, I am curious about how to construct a 1-particle basis for specific Lie groups, e.g., SO(3).\n- The formulation in the case of the product group in Section 4 seems to be missing, although there is a mention of the product group in the introduction and in Section 5. Is it trivial to design an equivariant model for the product group G_1 x G_2 with the formulation used in this method?  My initial impression is that there may be some subtleties involved, such as the order of the group actions.\n- A discussion of the computational aspects of this method is missing.  When using this model, the Clebsch-Gordan coefficients must be calculated numerically (as mentioned in Section 5), and it would be helful to mention calculation times and errors. Does the calculation time vary depending on the number of basis components? I am also curious about how much time it takes for iteration on backpropagation compared to classic MLP models.\n- Continuing with the above comment, since the Clebsch-Gordan coefficients are numerically calculated, some numerical calculation errors are inevitable. Is the equivariance of the model maintained in the presence of numerical errors? If not, some discussion, even qualitative, should be provided (e.g., to what extent the model remains equivariant) although quantitative results (e.g., experimental results, figures) would clearly be preferable.\n\nQUESTIONS:\nThese have been mentioned for the most part in the weaknesses section.\n\nLIMITATIONS:\nNo potential negative societal impact of this work as far as I can tell. As mentioned earlier, computational aspects of the method would be helpful. I don't want the authors to feel compelled to show that this method is immediately computationally practical, as that is not the only measure of the value and worth of any new idea, but it would be helpful to indicate to other researchers on what the computational limitations are, and to possibly spur interest in finding improvements.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors generalize MACE, a point cloud network that uses higher-order interactions via tensor products of basis expansions of the features, to being equivariant to arbitrary reductive Lie groups. The paper shows that this setup inherits universality properties from MACE. A generic method to compute the a basis for the equivariant linear maps between tensor products of the representations is implemented.\nWhile the paper is mostly theoretical, the authors show strong performance on a SO(1, 3) equivariant task and a point cloud task.\n\nSTRENGTHS:\n- The proposed method is an elegant generalization of a popular prior method\n- The paper is well-written and mostly easy to read.\n- The code is included\n- The proposed method could be a useful turn-key equivariance solution for researchers with a niche symmetry problem.\n\nWEAKNESSES:\n- Figure 1 of the paper appears to suggest that the method works on E(3), but this can't be a reductive group, as the 4D homogenous representation is not decomposable as a sum of irreducibles (contradicting lines 67-68). I suppose the authors mean to write O(3) and treat the translation by canonicalization, but they should clarity that.\n- The point could experiment should include more baselines.\n- As the key contribution of the paper mostly lies in implementing the necessary computations for generic reductive Lie groups (sec 5), it would make sense to allocate more space to that and less to the prior sections, which are mostly already covered in prior works. Alternatively / in addition, it would be good to give more background on the material in section 5 in the appendix.\n\nQUESTIONS:\n- It would be helpful to include/repeat the definition of B used in equation (20).\n\nLIMITATIONS:\nThe authors should clarify better that not all Lie groups of interest are reductive.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposes a framework for building equivariant neural networks on reductive lie groups. The proposed method first constructs a linear model for multi-set functions which is then symmetrised to generate a complete basis of equivariant multi-set functions. The model is also extended to a multi-layer architecture by a message-passing scheme.

STRENGTHS:
* The idea of generalizing ACE and MACE frameworks to arbitrary reductive Lie groups is new. 
* Proofs are given in the supplemental material to support theoretical claims in the paper. 
* A library is provided for developing G-equivariant neural networks.

WEAKNESSES:
* The paper has some minor errors in writting. For instance, Table 2 caption should be on top of the table. Tables (e.g., Tabs. (2) and (3)) that report experimental results should be mentioned in the text. 
* Lack of comparison against state-of-the-art methods for validating the effectiveness of the proposed method on 3D shape  recognition.

QUESTIONS:
* The computations in Eqs. (7) together with (16) seem to be expensive. How these can be done in practice ?
* It would be interesting to show the impact of higher order messages on the performance of G-MACE in terms of accuracy and computation time. I didn't find this study in the paper and supplemental material.

LIMITATIONS:
Limitations are discussed in the supplemental material.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07