A General Framework for Equivariant Neural Networks on Reductive Lie
  Groups
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 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).
Ilyes Batatia, Mario Geiger, Jose Munoz, Tess Smidt, Lior Silberman, Christoph Ortner
# 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 finite-dimensional 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, \(\mathrm{SO}(1,3)\) in High-Energy Physics, \(\mathrm{SU}(2^{N})\) in quantum mechanics, and \(\mathrm{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 \(\mathrm{O}(2)\) symmetry (Cohen and Welling, 2016; Esteves _et al._, 2017), machine learning force fields with \(\mathrm{O}(3)\) symmetry (Anderson _et al._, 2019; Bartok _et al._, 2013; Batzner _et al._, 2022; Batzia _et al._, 2022a) or jet tagging with \(\mathrm{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 \(\mathrm{O}(3)\)-invariant features. By combining the concepts of ACE and \(\mathrm{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 new framework generalizes the ACE (Drautz, 2019) and MACE (Batatia _et al._, 2022b) architectures to parameterize properties of point clouds that are equivariant under the action of a 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, \(\mathrm{SU}(N)\), \(\mathrm{SL}_{2}(\mathbb{C})\) and product groups.
* We demonstrate 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 \(\mathrm{GL}(n,\mathbb{C})\) the group of invertible \(n\times n\) matrices with complex entries. This includes well-known groups such as \(\mathrm{Sp}(2n,\mathbb{R})\) consisting of matrices of determinant

Figure 1: Examples of natural science problems and associated reductive Lie groups. For high energy physics, the Lorentz group \(\mathrm{SO}(1,3)\); for chemistry, the Euclidean group \(\mathrm{E}(3)\); for quantum-chromodynamics, the \(\mathrm{SU}(3)\) group.
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 \(\rho\colon G\to\mathrm{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 \(\mathrm{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 convolutionsConvolutional neural networks (CNNs), which are translation equivariant, have also been generalized to other symmetries. For example, G-convolutions [12] generalized CNNs to discrete groups. Steerable CNNs [12] generalized CNNs to \(O(2)\) equivariance and Spherical CNNs [12]\(O(3)\) equivariance. A general theory of convolution on any compact group and symmetric space was given by kondor2018deep. This work was further extended to equivariant convolutions on Riemannian manifolds by weiler2021deep.

ACEThe Atomic Cluster Expansion (ACE) [1] 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 [20].

e3nn + Equivariant MLPsThe e3nn library [12] provides a complete solution to build \(E(3)-\)equivariant neural networks based on irreducible representations. The Equivariant MLPs [12] 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 MACEEquivariant MPNNs [12, 13, 14, 15, 16, 17] 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 [1] generalized the \(O(3)\)-Equivariant MPNNs to build messages of arbitrary body order, outperforming other approaches on molecular tasks. [1] 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 \(\Omega\). We denote the set of all such multi-sets by \(\mathrm{msets}(\Omega)\). For example, in molecular modeling, \(x_{i}\) might describe the position and species of an atom and therefore \(x_{i}=(\mathbf{r}_{i},Z_{i})\in\mathbb{R}^{3}\times\mathbb{Z}\), while in high energy physics, one commonly uses the four-momentum \(x_{i}=(E_{i},\mathbf{p}_{i})\in\mathbb{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

\[\Phi\colon\mathrm{msets}(\Omega)\to Z \tag{1}\]

i.e., \(X\mapsto\Phi(X)\in 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 developed in this article are also applicable to fixed-length multi-sets, in which case \(\Phi\) is simply a permutation-invariant function defined on some \(\Omega^{N}\). Mappings that are not permutation-invariant are special case with several simplifications.

In many applications, especially in the natural sciences, particle properties satisfy additional symmetries. When a group \(G\) acts on \(\Omega\) as well as on \(Z\) we say that \(\Phi\) is \(G\)**-equivariant** if

\[\Phi\circ g=\rho_{Z}(g)\Phi,\qquad g\in G \tag{2}\]

where \(\rho_{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 \(\Phi\), we consider model architectures of the form

\[\Phi\colon\mathrm{msets}(\Omega)\underset{\text{embedding}}{ \longrightarrow}V\underset{\text{parameterization}}{\longrightarrow}V \underset{\text{readout}}{\longrightarrow}Z, \tag{3}\]
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 \(\Phi\) for the specific case when \(G=\operatorname{O}(3)\) and \(x_{i}\) are vectorial interatomic distances. Aspects of the ACE framework were incorporated into \(\operatorname{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.

### Efficient many-body expansion

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

\[\Phi^{(N)}(X)=\varphi_{0}+\sum_{i}\varphi_{1}(x_{i})+\sum_{i_{1},i_{2}} \varphi_{2}(x_{i_{1}},x_{i_{2}})+\dots+\sum_{i_{1},\dots,i_{N}}\varphi_{N}(x_ {i_{1}},\dots,x_{i_{N}}), \tag{4}\]

where \(\varphi_{n}\) defines the \(n\)-body interaction. Formally, the expansion becomes systematic in the limit as \(N\to\infty\). The second step is the expansion of the \(n\)-particle functions \(\varphi_{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 \((\phi_{k})_{k}\) is a 1-particle basis if they are linearly independent on \(\Omega\) and any smooth 1-particle function \(\varphi_{1}\) (not necessarily equivariant) can be expanded in terms of \((\phi_{k})_{\mathbf{k}}\), i.e,

\[\varphi_{1}(x)=\sum_{k}w_{k}\phi_{k}(x). \tag{5}\]

For the sake of concreteness, we assume that \(\phi_{k}:\Omega\to\mathbb{C}\), but the range can in principle be any field. Let a complex vector space \(V\) be given, into which the particle embedding maps, i.e.,

\[(\phi_{k}(x))_{k}\in V\qquad\forall x\in\Omega.\]

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

\[\varphi_{n}(x_{1},\dots,x_{n})=\sum_{k_{1},\dots,k_{n}}w_{k_{1}\dots k_{n}} \prod_{s=1}^{n}\phi_{k_{s}}(x_{s}). \tag{6}\]

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

\[A_{k}=\sum_{x\in X}\phi_{k}(x),\qquad\mathbf{A_{k}}=\prod_{s=1}^{n}A_{k},\qquad \Phi^{(N)}=\sum_{\mathbf{k}\in\mathcal{K}}w_{\mathbf{k}}\mathbf{A_{k}}, \tag{7}\]

where \(\mathcal{K}\) is the set of all \(\mathbf{k}\) tuples indexing the features \(\mathbf{A_{k}}\). Since \(\mathbf{A_{k}}\) is invariant under permuting \(\mathbf{k}\), only ordered \(\mathbf{k}\) tuples are retained. The features \(A_{k}\) are an embedding of \(\operatorname{msets}(\Omega)\) into the space \(V\). The tensorial product features (basis functions) \(\mathbf{A_{k}}\) form a complete linear basis of multi-set functions on \(\Omega\) and the weights \(w_{\mathbf{k}}\) can be understood as a symmetric tensor. We will extend this linear cluster expansion model \(\Phi^{(N)}\) to a message-passing type neural network model in SS 4.4.

We remark that, while the standard tensor product embeds \((\otimes_{s=1}^{n}\phi_{k_{s}})_{\mathbf{k}}\colon\Omega^{n}\to V^{n}\), the \(n\)-correlations \(\mathbf{A_{k}}\) are _symmetric tensors_ and embed \((\mathbf{A_{k}})_{\mathbf{k}}\colon\operatorname{msets}(\Omega)\to\operatorname{Sym }^{n}V\).

### 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 \(\rho=(\rho_{kk^{\prime}})\colon G\to\operatorname{GL}(V)\) such that

\[\phi_{k}\circ g=\sum_{k^{\prime}}\rho_{kk^{\prime}}(g)\phi_{k^{\prime}}, \tag{8}\]

where for each \(k\) the sum over \(k^{\prime}\) 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 \((\phi_{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 \(\Omega\); 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\otimes V\to V\). Concretely, we seek all solutions \(C^{\mathbf{\alpha},K}_{k_{1}k_{2}}\) to the equation

\[\sum_{k^{\prime}_{1}k^{\prime}_{2}}C^{\mathbf{\alpha},K}_{k^{\prime}_{1}k^{\prime}_ {2}}\rho_{k^{\prime}_{1}k_{1}}(g)\rho_{k^{\prime}_{2}k_{2}}(g)=\sum_{K^{\prime }}\rho_{KK^{\prime}}(g)C^{\mathbf{\alpha},K^{\prime}}_{k_{1}k_{2}}; \tag{9}\]

or, written in operator notation,

\[C^{\mathbf{\alpha}}\rho\otimes\rho=\rho C^{\mathbf{\alpha}}. \tag{10}\]

We will call the \(C^{\mathbf{\alpha},K}_{\mathbf{k}}\)_generalized Clebsch-Gordan coefficients_ since in the case \(G=\mathrm{SO}(3)\) acting on the spherical harmonics embedding \(\phi_{lm}=Y_{l}^{m}\) those coefficients are exactly the classical Clebsch-Gordan coefficients. The index \(\mathbf{\alpha}\) enumerates a basis of the space of all solutions to this equation. For the most common groups, one normally identifies a canonical basis \(C^{\mathbf{\alpha}}\) and assigns a natural meaning to this index (cf. SS A.2). 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 \(\mathrm{Sym}^{n}(V)\) becomes (cf. SS A.6)

\[\begin{split}&\sum_{\mathbf{k}^{\prime}}C^{\mathbf{\alpha},K}_{\mathbf{k}^{ \prime}}\overline{\mathbf{\rho}}_{\mathbf{k}^{\prime}\mathbf{k}}=\sum_{K^{\prime}}\rho_{KK^ {\prime}}C^{\mathbf{\alpha},K^{\prime}}_{\mathbf{k}}\qquad\forall K,\quad\mathbf{k}=(k_{1 },\dots,k_{N}),\quad g\in G,\\ &\text{where}\qquad\overline{\mathbf{\rho}}_{\mathbf{k}^{\prime}\mathbf{k}}= \sum_{\begin{subarray}{c}\mathbf{k}^{\prime}=\mathbf{k}^{\prime}\\ \mathbf{\pi}\in S_{n}\end{subarray}}\mathbf{\rho}_{\mathbf{k}^{\prime\prime}\mathbf{k}} \qquad\text{and}\qquad\mathbf{\rho}_{\mathbf{k}^{\prime}\mathbf{k}}=\prod_{t=1}^{n}\rho_{k ^{\prime}_{t}k_{t}}.\end{split} \tag{11}\]

Due to the symmetry of the \((\mathbf{A}_{\mathbf{k}})_{\mathbf{k}}\) tensors \(\mathcal{C}^{\mathbf{\alpha},K}_{\mathbf{k}}\) need only be computed for ordered \(\mathbf{k}\) tuples and the sum \(\sum_{\mathbf{k}^{\prime}}\) also runs only over ordered \(\mathbf{k}\) tuples. Again, the index \(\mathbf{\alpha}\) enumerates a basis of the space of solutions. Equivalently, (11) can be written in compact notation as

\[\mathcal{C}^{\mathbf{\alpha}}\overline{\mathbf{\rho}}=\rho\mathcal{C}^{\mathbf{\alpha}}. \tag{12}\]

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

\[\mathbf{B}^{K}_{\mathbf{\alpha}}=\sum_{\mathbf{k}^{\prime}}C^{\mathbf{\alpha},K}_{\mathbf{k}^{ \prime}}\mathbf{A}_{\mathbf{k}^{\prime}}. \tag{13}\]

The equivariance of (13) is easily verified by applying a transformation \(g\in G\) to the input (cf SS A.3).

**Universality:** In the limit as the correlation order \(N\to\infty\), the features \((\mathbf{B}^{K}_{\mathbf{\alpha}})_{K,\mathbf{\alpha}}\) form a complete basis of smooth equivariant multi-set functions, in a sense that we make precise in Appendix A.4. Any equivariant property \(\Phi_{V}:\Omega\to V\) can be approximated by a linear model

\[\Phi_{V}^{K}=\sum_{\mathbf{\alpha}}c^{K}_{\mathbf{\alpha}}B^{K}_{\mathbf{\alpha}}, \tag{14}\]

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

### 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 \((\phi_{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}=(\mathbf{r}_{i},Z_{i})\),

\[\phi_{k}(x_{i})=\phi_{znlm}(x_{i})=\delta_{z_{Z_{i}}}R_{nl}(r_{i})Y_{l}^{m}( \hat{r}_{i}).\]

Only the \((lm)\) channels are involved in the representation of \(\mathrm{O}(3)\) hence there is considerable freedom in "compressing" the \((zn)\) 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 \((\phi_{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 the \(\rho_{k,k^{\prime}}\), that is, (8) becomes

\[\phi_{ck}\circ g=\sum_{k^{\prime}}\rho_{kk^{\prime}}(g)\phi_{ck^{\prime}}. \tag{15}\]

Therefore, manipulating only the \(c\) channel does not change any symmetry properties of the architecture. We can use this fact to admit a learnable embedding,

\[\tilde{A}_{ck}=\sum_{c^{\prime}}w_{cc^{\prime}}\phi_{c^{\prime}k},\]

This mechanism is employed in numerous architectures to reduce the dimensionality of the embedding, but the approach taken in by Darby _et al._ (2022) and Batatia _et al._ (2022b) and followed here is the exact opposite: we allow many more learnable \(c\) channels but then decouple them resulting in a much lower-dimensional basis of \(n\)-correlations, defined by

\[\tilde{\mathbf{A}}_{ck}=\prod_{t=1}^{n}\Bigg{(}\sum_{c^{\prime}}w_{cc^{\prime}} \sum_{x\in X}\phi_{c^{\prime}k_{t}}(x)\Bigg{)}. \tag{16}\]

The resulting symmetrized basis is then obtained by

\[\mathbf{B}_{c\mathbf{\alpha}}^{K}=\sum_{\mathbf{k^{\prime}}}C_{\mathbf{k^{\prime}}}^{\mathbf{ \alpha},K}\tilde{\mathbf{A}}_{ck^{\prime}}. \tag{17}\]

Following the terminology of Darby _et al._ (2022) we call this architecture the tensor-reduced ACE, or, \(G\)-TRACE. There are numerous natural variations on its construction, 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\rightarrow\infty\). A smooth equivariant property \(\Phi\) may be approximated to within arbitrary accuracy by an expansion \(\Phi^{K}(X)\approx\sum_{c,\mathbf{\alpha}}c^{K}_{\mathbf{\alpha}}\mathbf{B}_{c,\mathbf{\alpha} }^{K}(X)\). Since the embedding \(\tilde{A}_{ck}\) is learnable, this is a _nonlinear model_. We refer to SS A.4 for the details.

### G-MACE, Multi-layer cluster expansion

The \(G\)-equivariant cluster expansion is readily generalized to a multi-layer architecture by re-expanding previous features in a new cluster expansion (Batatia _et al._, 2022b). The multi-set \(X\) is endowed with extra features, \(\mathbf{h}_{i}^{t}=(h_{i,ck}^{t})_{c,k}\), that are updated for \(t\in\{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

\[x_{i}^{t} =(x_{i},\mathbf{h}_{i}^{t}) \tag{18}\] \[\phi_{ck}^{t}(x_{i},\mathbf{h}_{i}^{t}) =\sum_{\mathbf{\alpha}}w_{\mathbf{\alpha}}^{t,ck}\sum_{k^{\prime},k^{ \prime\prime}}C_{k^{\prime}k^{\prime\prime}}^{\mathbf{\alpha},k}h_{i,ck^{\prime}} ^{t}\phi_{ck^{\prime\prime}}(x_{i}) \tag{19}\]

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

\[m_{i,cK}^{t}=\sum_{\mathbf{\alpha}}W_{\mathbf{\alpha}}^{t,cK}\mathbf{B}_{c\mathbf{\alpha}}^{t,K}. \tag{20}\]

The gathered message \(\mathbf{m}_{i}^{t}=(m_{i,cK}^{t})_{c,k}\) is then used to update the particle states,

\[x_{i}^{t+1}=(x_{i},\mathbf{h}_{i}^{t+1}),\qquad\mathbf{h}_{i}^{t+1}=U_{t}\big{(}\mathbf{m}_ {i}^{t}), \tag{21}\]

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\),

\[y_{i}=\sum_{t=1}^{T}\mathcal{R}_{t}^{\mathrm{loc}}(x_{i}^{t}),\qquad\text{ respectively,}\qquad y=\sum_{t=1}^{T}\mathcal{R}_{t}^{\mathrm{glob}}(\{x_{i}^{t}\}_{i}). \tag{22}\]
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 \((\phi_{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\in\mathrm{msets}(\Omega)\), 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 (11) for a given tuple of representations (see 13). To facilitate this task, we have implemented an open source software library, lie-nn1. In this section we review the key techniques employed in this library.

Footnote 1: [https://github.com/lie-nn/lie-nn](https://github.com/lie-nn/lie-nn)

### 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 algebraThe Lie groups we study can be realized as closed subgroups \(G\subset\mathrm{GL}_{n}(\mathbb{R})\) of the general linear group. In that case their Lie algebras can be concretely realized as \(\mathfrak{g}=\mathrm{Lie}(G)=\{X\in M_{n}(\mathbb{R})\mid\forall t\in\mathbb{R }:\exp(tX)\in G\}\) where \(\exp(X)=1+X+\frac{1}{2}X^{2}...\) is the standard matrix exponential. It turns out that \(\mathfrak{g}\subset M_{n}(\mathbb{R})\) is a linear subspace closed under the commutator bracket \([X,Y]=XY-YX\).

Structure theoryWe fix a linear basis \(\{X_{i}\}\subset\mathfrak{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}]=\sum_{k}A_{ijk}X_{k}\), in that if \(X=\sum_{i}a_{i}X_{i}\) and \(Y=\sum_{j}b_{j}X_{j}\) then \([X,Y]=\sum_{k}\left(\sum_{i,j}A_{ijk}a_{i}b_{j}\right)X_{k}\). The classification of reductive groups provides convenient generating sets for their Lie algebras (or their complexifications). One identifies a large commutative subalgebra \(\mathfrak{h}\subset\mathfrak{g}\) (sometimes of \(\mathfrak{g}_{C}=\mathfrak{g}\otimes_{\mathbb{R}}\mathbb{C}\)) with basis \(\{H_{i}\}\) so that most (or all) of the other generators \(E_{\alpha}\) can be chosen so that \([H_{i},E_{\alpha}]=\alpha(H_{i})E_{\alpha}\) for a linear function \(\alpha\) on \(\mathfrak{h}\). These functions are the so-called _roots_ of \(\mathfrak{g}\). Structural information about \(\mathfrak{g}\) is commonly encoded pictorially via the _Dynkin diagram_ of \(\mathfrak{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}=\mathfrak{su}(n+1),B_{n}=\mathfrak{so}(2n+1),C_{n}=\mathfrak{sp}(2n),D_{ n}=\mathfrak{so}(2n)\) 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 \(\mathbf{H}\subset G\) containing a representative from each connected component.

Representation theoryThe 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 \(\mathfrak{h}\) ("highest

Figure 2: Examples of Dynkin diagrams and their associated group class.
weight"). Further given a highest weight \(\lambda\) 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.

### Numerical Computations in lie-nn

The most basic class of the lie-nn library encodes a group \(G\) and infinitesimal representation \(d\rho\) of \(\mathfrak{g}\) using the tuple

\[\rho:=(A,n,\{d\rho(X_{i})\}_{i},\{\rho(h)\}_{h\in\mathbf{H}})\, \tag{23}\]

with \(A\) the structure constants of the group, \(n\) the dimension of the representation, and \(d\rho(X_{i})\) and \(\rho(h)\) being \(n\times 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, \(\forall X\in\mathfrak{g},\rho(e^{X})=e^{d\rho(X)}\).

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. A.6 for details).

To construct an irrep explicitly as in (23) 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\rho_{1},d\rho_{2}\) be two Lie algebra representations of interest. The tensor product on the Lie algebra \(d\rho_{1}\otimes d\rho_{2}(X)\) can be computed as,

\[d\rho_{1}\otimes d\rho_{2}\ (X)=d\rho_{1}(X)\otimes 1+1\otimes d\rho_{2}(X) \tag{24}\]

Therefore, given sets of generators of three representations \(d\rho_{1},d\rho_{2},d\rho_{3}\), the Clebsch-Gordan coefficients are the change of basis between \((d\rho_{1}(X)\otimes 1+1\otimes d\rho_{2}(X))\) and \(d\rho_{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.7).

## 6 Applications

### 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.
### 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},\mathbf{p}_{i})]_{i=1}^{N}\), where \(E_{i}\in\mathbb{R}^{+}\) and \(\mathbf{p}_{i}\in\mathbb{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. \(G\)-MACE achieves excellent accuracy, being the only arbitrary equivariant model to reach similar accuracy as PELICAN. We refer to Appendix A.8.1 for the details of the architecture.

### 3D Shape recognition

3D shape recognition from point clouds is of central importance for computer vision. We use the ModellNet10 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.8.2 for more details on the architecture.

\begin{table}
\begin{tabular}{l r r} \hline \hline Group & Application & Reference \\ \hline \(\mathrm{SU}(1)\) & Electromagnetism & (Lagrave _et al._, 2021) \\ \(\mathrm{SU}(3)\) & Quantum Chromodynamics & (Favoni _et al._, 2022) \\ \(\mathrm{SO}(3)\) & 3D point clouds & (Batatia _et al._, 2022a) \\ \(\mathrm{SO}^{+}(1,3)\) & Particle Physics & (Bogatskiy _et al._, 2020b) \\ \(\mathrm{SL}(3,\mathbb{R})\) & Point cloud classification & - \\ \(\mathrm{SU}(2^{N})\) & Entangled QP & - \\ \hline \(\mathrm{Sp}(N)\) & Hamiltonian dynamics & - \\ \(\mathrm{SO}(2N+1)\) & Projective geometry & - \\ \hline \hline \end{tabular}
\end{table}
Table 1: Lie groups of interests covered by the present methods and their potential applications to equivariant neural networks. The groups above the horizontal line are already available in lie-nn. The ones below the line fall within our framework and can be added.

\begin{table}
\begin{tabular}{l r r r r} \hline \hline Architecture & \#Params & Accuracy & AUC & \(\mathrm{Rej}_{30\%}\) \\ \hline \hline
**PELICAN** & 45k & **0.942** & **0.987** & \(\mathbf{2289\pm 204}\) \\
**partT** & 2.14M & 0.940 & 0.986 & \(1602\pm 81\) \\
**ParticleNet** & 498k & 0.938 & 0.985 & \(1298\pm 46\) \\
**LorentzNet** & 224k & **0.942** & **0.987** & \(2195\pm 173\) \\
**BIP** & 4k & 0.931 & 0.981 & \(853\pm 68\) \\
**LGN** & 4.5k & 0.929 & 0.964 & \(435\pm 95\) \\
**EFN** & 82k & 0.927 & 0.979 & \(888\pm 17\) \\
**TopoDNN** & 59k & 0.916 & 0.972 & \(295\pm 5\) \\
**LorentzMACE** & 228k & **0.942** & **0.987** & \(1935\pm 85\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Comparisons between state-of-the-art metrics on the Top-Tagging dataset. Scores were taken from (Bogatskiy _et al._, 2022; Qu _et al._, 2022; Qu and Gouskos, 2020; Munoz _et al._, 2022; Bogatskiy _et al._, 2020a; Komiske _et al._, 2019; Pearkes _et al._, 2017).

\begin{table}
\begin{tabular}{l r} \hline \hline Architecture & Accuracy \\ \hline PointNet (Qi _et al._, 2016) & 94.2 \\ PointNet ++ (Qi _et al._, 2017) & 95.0 \\ PointMACE (ours) & **96.1** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Accuracy in shape recognition.
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 Gabor Csanyi for his support.
