SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\sum_ {m _ {1} ^ {\prime} m _ {2} ^ {\prime}} C _ {l _ {1} m _ {1} ^ {\prime} l _ {2} m _ {2} ^ {\prime}} ^ {L M} \rho_ {m _ {1} ^ {\prime} m _ {1}} ^ {l _ {1}} (g) \rho_ {m _ {2} ^ {\prime} m _ {2}} ^ {l _ {2}} (g) = \sum_ {M ^ {\prime}} \rho_ {M M ^ {\prime}} ^ {L} (g) C _ {l _ {1} m _ {1} l _ {2} m _ {2}} ^ {L M ^ {\prime}}. \tag {40}
$$
In this setting, our index  $\alpha$  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
$$
C _ {l _ {1} m _ {1} l _ {2} m _ {2}} ^ {\boldsymbol {\alpha}, L M} = \left\{ \begin{array}{l l} C _ {l _ {1} m _ {1} l _ {2} m _ {2}} ^ {L M}, & \text {i f} \boldsymbol {\alpha} = (l _ {1}, l _ {2}), \\ 0, & \text {o t h e r w i s e}, \end{array} \right. \tag {41}
$$
where  $C_{l_1m_1l_2m_2}^{LM}$  are the Clebsch-Gordan coefficients in the classical notation. Thus, the additional index  $\alpha$  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,
$$
\begin{array}{l} \boldsymbol {B} _ {\boldsymbol {\alpha}} ^ {K} \circ g = \sum_ {\boldsymbol {k}} \mathcal {C} _ {\boldsymbol {k}} ^ {\boldsymbol {\alpha}, K} \boldsymbol {A} _ {\boldsymbol {k}} \circ g \\ = \sum_ {\boldsymbol {k}} \mathcal {C} _ {\boldsymbol {k}} ^ {\boldsymbol {\alpha}, K} \left(\sum_ {\boldsymbol {k} ^ {\prime}} \prod_ {t} \rho_ {k _ {t}, k _ {t} ^ {\prime}} (g) \boldsymbol {A} _ {\boldsymbol {k} ^ {\prime}}\right) \\ = \sum_ {\boldsymbol {k} ^ {\prime}} \left(\sum_ {\boldsymbol {k}} \mathcal {C} _ {\boldsymbol {k}} ^ {\alpha , K} \prod_ {t} \rho_ {k _ {t}, k _ {t} ^ {\prime}} (g)\right) \boldsymbol {A} _ {\boldsymbol {k} ^ {\prime}} \tag {42} \\ = \sum_ {\boldsymbol {k} ^ {\prime}} \left(\sum_ {K ^ {\prime}} \rho_ {K K ^ {\prime}} (g) \mathcal {C} _ {\boldsymbol {k} ^ {\prime}} ^ {\boldsymbol {\alpha}, K ^ {\prime}}\right) \boldsymbol {A} _ {\boldsymbol {k} ^ {\prime}} \\ = \sum_ {K ^ {\prime}} \rho_ {K K ^ {\prime}} (g) \boldsymbol {B} _ {\boldsymbol {\alpha}} ^ {K ^ {\prime}}. \\ \end{array}
$$
A.7 Completeness of the basis and Universality of G-MACE
We explain in which sense the basis  $B_{\alpha}^{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  $\Phi^V$  embedded in  $V$  where we have a representation, i.e. the actual target property is  $\Phi$  is then given as a linear mapping from  $V$  to  $Z$ . For technical reasons, we require that only finitely many entries  $\Phi_K^V$  may be non-zero, but this is consistent with common usage. For example, if  $G = O(3)$  and if  $\Phi$  is a scalar, then  $\Phi_0^V = \Phi$ , while all other  $\Phi_{LM}^V \equiv 0$ . If  $\Phi$  is a covariant vector, then  $\Phi_{LM}^V$  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  $\Phi^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  $\Omega$  is compact.
Throughout the remainder of this section we may therefore assume that an  $N$  can be chosen as well as smooth components  $\varphi^{(n)}$  such that the resulting model  $\Phi^{V,N}$  approximates  $\Phi^V$  to within a target accuracy  $\epsilon$ ,
$$
\left| \Phi_ {K} ^ {V, N} (\pmb {x}) - \Phi_ {K} ^ {V} (\pmb {x}) \right| \leq \epsilon \qquad \forall \pmb {x} \in \mathrm {m s e t s} (\Omega).
$$
2. The density of the embedding. As already stated in the main text, if the components  $\varphi_K^{(n)}$  are smooth, and the embedding  $\{\phi_k\}_k$  is dense in the space of one-particle functions (5) then it follows that the  $\varphi_K^{(n)}$  can be expanded in terms of the tensor product basis  $\phi_k \coloneqq \bigotimes_{s=1}^n \phi_{k_s}$  to within arbitrary accuracy. The precise statement is the following standard result of approximation theory: if  $\operatorname{span}\{\phi_k\}_k$  are dense in  $C(\Omega)$ , then  $\operatorname{span}\{\phi_k\}_k$  are dense in  $C(\Omega^n)$ . That is, for any  $\epsilon > 0$ , there exist approximants  $p_K^{(n)}$  such that
$$
\left\| \varphi_ {K} ^ {(n)} - p _ {K} ^ {(n)} \right\| _ {\infty} \leq \epsilon .
$$
3. The density of the symmetrized basis. The next and crucial step is to show that, if the  $\varphi_K^{(n)}$  are equivariant, then the  $p_K^{(n)}$  may be chosen equivariant as well without loss of accuracy. If the group  $G$  is compact then the representations  $\rho$  can be chosen unitary (Broecker, 1985). In that case, the argument from (Dusson et al., 2022) can be used almost verbatim: let
$$
\bar {p} ^ {(n)} (\boldsymbol {x}) := \int_ {G} \rho (g) ^ {- 1} p ^ {(n)} (g \boldsymbol {x}) H (d g),
$$
where  $H$  is the normalized Haar measure then  $\bar{p}^{(n)}$  is equivariant by construction and
$$
\begin{array}{l} \left| \varphi^ {(n)} (\boldsymbol {x}) - \bar {p} ^ {(n)} (\boldsymbol {x}) \right| \\ = \left| \int_ {G} \rho (g) ^ {- 1} \left(\varphi^ {(n)} (g \boldsymbol {x}) - p ^ {(n)} (g \boldsymbol {x})\right) H (d g) \right| \\ \leq \int_ {G} \left| \varphi^ {(n)} (g \boldsymbol {x}) - p ^ {(n)} (g \boldsymbol {x}) \right| H (d g) \\ \leq \int_ {G} \| \varphi^ {(n)} - p ^ {(n)} \| _ {\infty} H (d g) \leq \epsilon . \\ \end{array}
$$
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_{\mathbb{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_{\mathbb{C}}$  results in an approximant that is not only equivariant w.r.t.  $K_{\mathbb{C}}$  but also equivariant w.r.t.  $G$ .
4. The density of the basis  $B_{\alpha}^{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  $\epsilon \rightarrow 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 \to \infty$ .
5. Pointwise convergence. To obtain density in the sense of pointwise convergence we first introduce the canonical cluster expansion without self-interacting terms
$$
\Phi_ {K} (\boldsymbol {x}) = \sum_ {n = 0} ^ {\infty} \sum_ {j _ {1} <   \dots <   j _ {n}} v _ {K} ^ {(n)} (x _ {j _ {1}}, \ldots x _ {j _ {n}}).
$$
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  $\pmb{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 \leq N$  and then expanding the potentials  $v_{K}^{(n)}$  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\rho_{1},\rho_{1})$  and  $(d\rho_{2},\rho_{2})$  their continuous and discrete generators. One can define a representation of the direct product group  $G_{1}\times G_{2}$  as
$$
\rho := \left(A _ {1} \mid A _ {2}, n _ {1} n _ {2}, \left\{d \rho_ {1} \left(X _ {i}\right) \otimes I _ {2} + I _ {1} \otimes d \rho_ {2} \left(\tilde {X} _ {j}\right) \right\} _ {i, j}, \left\{\rho \left(h _ {1}\right) \otimes I _ {2} + I _ {1} \otimes \rho \left(h _ {2}\right) \right\} _ {h _ {1}, h _ {2} \in \mathbf {H} _ {1}, \mathbf {H} _ {2}}\right) \tag {43}
$$
The following essential property holds: if  $\rho_{1}$  and  $\rho_{2}$  are irreducible representations of  $G_{1}$  and  $G_{2}$  then  $\rho_{1} \otimes \rho_{2}$  is an irreducible representation of  $G_{1} \times G_{2}$ . Moreover, for reducible Lie groups, all the irreps of  $G_{1} \times G_{2}$  are of this form. Therefore, one can construct all the irreps of  $G_{1} \times 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 \times S_{n}$ .
Let us give a non-trivial application of lie-nn for computing invariants for the product group of  $O(3) \times S_3$ . We consider a set of three vectors,
$$
a = \left[ \begin{array}{c} a _ {x} \\ a _ {y} \\ a _ {z} \end{array} \right] b = \left[ \begin{array}{c} a _ {x} \\ a _ {y} \\ a _ {z} \end{array} \right] c = \left[ \begin{array}{c} a _ {x} \\ a _ {y} \\ a _ {z} \end{array} \right]
$$
A given vector belongs 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,
$$
\operatorname {r e p} = \text {l i e . g r o u p _ {-}} \text {p r o d u c t (l i e . f i n i t e . S n _ {-} n a t u r a l (3) , l i e . i r r e p s . O 3 (1 = 1 , p = - 1))}
$$
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,
$$
\begin{array}{l} \text {q s = l i e . i n f e r c h a n g e o f b a s i s (} \\ \quad \text {l i e . g r o u p _ {p r o d u c t (l i e . f i n i t e . S n _ {t r i v i a l} (3) , l i e . i r r e p s . O 3 (l = 1 , p = - 1)) ,} \\ \quad \text {r e p ,} \\ \quad \text {)} \end{array}
$$
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,
qs = lie .infer_change_of BASIS(
lie . group _ product( lie . finite . Sn _ trivial (3), lie . irreps . O3(1 = 0, p = 1)),
lie.tensor_product(rep，rep），
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,
$$
e _ {1} = a _ {x} ^ {2} + a _ {y} ^ {2} + a _ {z} ^ {2} + b _ {x} ^ {2} + b _ {y} ^ {2} + b _ {z} ^ {2} + c _ {x} ^ {2} + c _ {y} ^ {2} + c _ {z} ^ {2}
$$
$$
e _ {2} = a _ {x} b _ {x} + a _ {x} c _ {x} + a _ {y} b _ {y} + a _ {y} c _ {y} + a _ {z} b _ {z} + a _ {z} c _ {z} + b _ {x} c _ {x} + b _ {y} c _ {y} + b _ {z} c _ {z}
$$
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.
qs = lie.infer_change_of BASIS(
lie . group _ product ( lie . finite . Sn _ trivial ( 3 ) , lie . irreps . O3(1 = 2 , p = 1)) ,
The computation tells us that there exists two such vectors,
$$
v _ {1} = \left[ \begin{array}{c} - \sqrt {2} a _ {x} a _ {z} - \sqrt {2} b _ {x} b _ {z} - \sqrt {2} c _ {x} c _ {z} \\ - \sqrt {2} a _ {x} a _ {y} - \sqrt {2} b _ {x} b _ {y} - \sqrt {2} c _ {x} c _ {y} \\ \frac {\sqrt {6} a _ {x} ^ {2}}{6} - \frac {\sqrt {6} a _ {y} ^ {2}}{3} + \frac {\sqrt {6} a _ {z} ^ {2}}{6} + \frac {\sqrt {6} b _ {x} ^ {2}}{6} - \frac {\sqrt {6} b _ {y} ^ {2}}{3} + \frac {\sqrt {6} c _ {x} ^ {2}}{6} + \frac {\sqrt {6} c _ {y} ^ {2}}{6} - \frac {\sqrt {6} c _ {z} ^ {2}}{3} + \frac {\sqrt {6} c _ {z} ^ {2}}{6} \\ - \sqrt {2} a _ {y} a _ {z} - \sqrt {2} b _ {y} b _ {z} - \sqrt {2} c _ {y} c _ {z} \\ \frac {\sqrt {2} a _ {x} ^ {2}}{2} - \frac {\sqrt {2} a _ {z} ^ {2}}{2} + \frac {\sqrt {2} b _ {x} ^ {2}}{2} - \frac {\sqrt {2} b _ {z} ^ {2}}{2} + \frac {\sqrt {2} c _ {x} ^ {2}}{2} - \frac {\sqrt {2} c _ {z} ^ {2}}{2} \end{array} \right]
$$
$$
v _ {2} = \left[ \begin{array}{c} - a _ {x} b _ {z} - a _ {x} c _ {z} - a _ {z} b _ {x} - a _ {z} c _ {x} - b _ {x} c _ {z} - b _ {z} c _ {x} \\ - a _ {x} b _ {y} - a _ {x} c _ {y} - a _ {y} b _ {x} - a _ {y} c _ {x} - b _ {x} c _ {y} - b _ {y} c _ {x} \\ \frac {\sqrt {3} a _ {x} b _ {x}}{3} + \frac {\sqrt {3} a _ {x} c _ {x}}{3} - \frac {2 \sqrt {3} a _ {y} b _ {y}}{3} - \frac {2 \sqrt {3} a _ {y} c _ {y}}{3} + \frac {\sqrt {3} a _ {z} b _ {z}}{3} + \frac {\sqrt {3} a _ {z} c _ {z}}{3} + \frac {\sqrt {3} b _ {x} c _ {x}}{3} - \frac {2 \sqrt {3} b _ {y} c _ {y}}{3} + \frac {\sqrt {3} b _ {z} c _ {z}}{3} \\ - a _ {y} b _ {z} - a _ {y} c _ {z} - a _ {z} b _ {y} - a _ {z} c _ {y} - b _ {y} c _ {z} - b _ {z} c _ {y} \\ a _ {x} b _ {x} + a _ {x} c _ {x} - a _ {z} b _ {z} - a _ {z} c _ {z} + b _ {x} c _ {x} - b _ {z} c _ {z} \end{array} \right]
$$
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,
$$
S _ {\nu} = B _ {\nu ; i j k l} x _ {i} y _ {j} z _ {k} w _ {l} \tag {44}
$$
where  $B$  is the change of basis that satisfies:
$$
B _ {\nu ; i j k l} = B _ {\nu ; \sigma (i j k l)} \forall \sigma \in S _ {4} \tag {45}
$$
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  $\Delta \eta = 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 architectures 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  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.999$ , and  $\epsilon = 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,
$$
h ^ {(t + 1)} = \operatorname {P o i n t N e t} \left(x y z ^ {(t)}, h ^ {(t)}\right) + \operatorname {M A C E} \left(x y z ^ {(t)}, h ^ {(t)}\right) \tag {46}
$$
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  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.999$ , and  $\epsilon = 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. Therefore, the results can be stored for subsequent use, ensuring that this phase does not affect the model's overall performance.
Figure 3: Generation time in seconds of Clebsch-Gordan coefficients as a function of the size of the representations for different groups. The representation size is computed as the product of the three representation sizes involved in the Clebsch-Gordan.
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.
Figure 4: Sparsity in the percentage of non-zero entries of Clebsch-Gordan coefficients as a function of the size of the representations for different groups. The representations size is computed as the product of the three representation sizes involved in the Clebsch-Gordan.
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.
Figure 5: Inference time for a single jet for a two layer LorentzMACE model as a function of the correlation order in the product basis and the maximal angular resolution. Timing made on a Nvidia A100 and averaged over a 100 runs.
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 implementation 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.