SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ilyes Batatia
Engineering Laboratory,
Cambridge, CB2 1PZ UK
ENS Paris-Saclay, Université Paris-Saclay
91190Gif-sur-Yvette,France
ilyes.batatia@ens-paris-saclay.fr
Mario Geiger
Department of Electrical Engineering
and Computer Science,
Cambridge, MA, USA
EIA University, FTA Group
Antioquia, Colombia
Tess Smidt
Department of Electrical Engineering and Computer Science,
Lior Silberman
Vancouver, BC, Canada V6T 1Z2
Christoph Ortner
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).
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, 2016a; Esteves et al., 2017), machine learning force fields with  $\mathrm{O}(3)$  symmetry (Anderson et al., 2019; Bartók et al., 2013; Batzner et al., 2022; Batatia 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.
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.
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 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-nm, 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 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.
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  $\operatorname {Sp}(2n,\mathbb{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  $\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  $\operatorname {GL}(n,F)$  such that every finite-dimensional representation of  $G$  on an  $F$ -vectorspace can be decomposed as a sum of irreducible representations.
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  $\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 = (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, 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 : \operatorname {m s e t s} (\Omega) \rightarrow 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 methods 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 cases 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 , \quad 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 : \text {m s e t s} (\Omega) \underset {\text {e m b e d d i n g}} {\longrightarrow} V \underset {\text {p a r a m e t e r i z a t i o n}} {\longrightarrow} V \underset {\text {r e a d o u t}} {\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 = \mathrm{O}(3)$  and  $x_{i}$  are vectorial interatomic distances. Aspects of the ACE framework were incorporated into  $\mathrm{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  $\Phi$  in terms of body orders, and truncate the expansion at a finite order  $N$ :
$$
\Phi^ {(N)} (X) = \varphi_ {0} + \sum_ {i} \varphi_ {1} \left(x _ {i}\right) + \sum_ {i _ {1}, i _ {2}} \varphi_ {2} \left(x _ {i _ {1}}, x _ {i _ {2}}\right) + \dots + \sum_ {i _ {1}, \dots , i _ {N}} \varphi_ {N} \left(x _ {i _ {1}}, \dots , x _ {i _ {N}}\right), \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)_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. 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.,
$$
(\phi_ {k} (x)) _ {k} \in V \quad \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} \left(x _ {1}, \dots , x _ {n}\right) = \sum_ {k _ {1}, \dots , k _ {n}} w _ {k _ {1} \dots k _ {n}} \prod_ {s = 1} ^ {n} \phi_ {k _ {s}} \left(x _ {s}\right). \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), \quad \boldsymbol {A} _ {\boldsymbol {k}} = \prod_ {s = 1} ^ {n} A _ {k}, \quad \Phi^ {(N)} = \sum_ {\boldsymbol {k} \in \mathcal {K}} w _ {\boldsymbol {k}} \boldsymbol {A} _ {\boldsymbol {k}}, \tag {7}
$$
where  $\mathcal{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  $\mathrm{msets}(\Omega)$  into the space  $V$ . The tensorial product features (basis functions)  $A_{k}$  form a complete linear basis of multi-set functions on  $\Omega$  and the weights  $w_{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 § 4.4.
While the standard tensor product embeds  $(\otimes_{s=1}^{n}\phi_{k_s})_k \colon \Omega^n \to V^n$ , the  $n$ -correlations  $\mathbf{A}_k$  are symmetric tensors and embed  $(\mathbf{A}_k)_k \colon \mathrm{msets}(\Omega) \to \mathrm{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  $\rho = (\rho_{kk'}) \colon G \to \mathrm{GL}(V)$  such that
$$
\phi_ {k} \circ g = \sum_ {k ^ {\prime}} \rho_ {k k ^ {\prime}} (g) \phi_ {k ^ {\prime}}, \tag {8}
$$
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  $(\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_{k_1 k_2}^{\alpha, K}$  to the equation
$$
\sum_ {k _ {1} ^ {\prime} k _ {2} ^ {\prime}} C _ {k _ {1} ^ {\prime} k _ {2} ^ {\prime}} ^ {\boldsymbol {\alpha}, K} \rho_ {k _ {1} ^ {\prime} k _ {1}} (g) \rho_ {k _ {2} ^ {\prime} k _ {2}} (g) = \sum_ {K ^ {\prime}} \rho_ {K K ^ {\prime}} (g) C _ {k _ {1} k _ {2}} ^ {\boldsymbol {\alpha}, K ^ {\prime}}; \tag {9}
$$
or, written in operator notation,  $C^\alpha \rho \otimes \rho = \rho C^\alpha$ . We will call the  $C_k^{\alpha ,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  $\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^\alpha$  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  $\operatorname {Sym}^n (V)$  becomes (cf. § A.9)
$$
\sum_ {\boldsymbol {k} ^ {\prime}} \mathcal {C} _ {\boldsymbol {k} ^ {\prime}} ^ {\boldsymbol {\alpha}, K} \bar {\rho} _ {\boldsymbol {k} ^ {\prime} \boldsymbol {k}} = \sum_ {K ^ {\prime}} \rho_ {K K ^ {\prime}} \mathcal {C} _ {\boldsymbol {k}} ^ {\boldsymbol {\alpha}, K ^ {\prime}} \quad \forall K, \quad \boldsymbol {k} = (k _ {1}, \dots , k _ {N}), \quad g \in G,
$$
$$
\text {w h e r e} \quad \overline {{\rho}} _ {\boldsymbol {k} ^ {\prime} \boldsymbol {k}} = \sum_ {\substack {\boldsymbol {k} ^ {\prime \prime} = \pi \boldsymbol {k} ^ {\prime} \\ \pi \in S _ {n}}} \rho_ {\boldsymbol {k} ^ {\prime \prime} \boldsymbol {k}} \quad \text {a n d} \quad \rho_ {\boldsymbol {k} ^ {\prime} \boldsymbol {k}} = \prod_ {t = 1} ^ {n} \rho_ {\boldsymbol {k} _ {t} ^ {\prime} k _ {t}}. \tag{10}
$$
Due to the symmetry of the  $(A_{k})_{k}$  tensors  $\mathcal{C}_k^{\alpha ,K}$  need only be computed for ordered  $k$  tuples and the sum  $\sum_{k'}$  also runs only over ordered  $k$  tuples. Again, the index  $\alpha$  enumerates a basis of the space of solutions. Equivalently, (10) can be written in compact notation as  $\mathcal{C}^{\alpha}\overline{\rho} = \rho \mathcal{C}^{\alpha}$ . 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
$$
\boldsymbol {B} _ {\alpha} ^ {K} = \sum_ {\boldsymbol {k} ^ {\prime}} \mathcal {C} _ {\boldsymbol {k} ^ {\prime}} ^ {\alpha , K} \boldsymbol {A} _ {\boldsymbol {k} ^ {\prime}}. \tag {11}
$$
The equivariance of (11) is easily verified by applying a transformation  $g \in G$  to the input (cf § A.6).
Universality: In the limit as the correlation order  $N\to \infty$  , the features  $(B_{\alpha}^{K})_{K,\alpha}$  form a complete basis of smooth equivariant multi-set functions, in a sense that we make precise in Appendix A.7. Any equivariant property  $\Phi_V:\Omega \rightarrow V$  can be approximated by a linear model
$$
\Phi_ {V} ^ {K} = \sum_ {\alpha} c _ {\alpha} ^ {K} B _ {\alpha} ^ {K}, \tag {12}
$$
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  $(\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 = (r_i,Z_i)$ ,
$$
\phi_ {k} (x _ {i}) = \phi_ {z n l m} (x _ {i}) = \delta_ {z Z _ {i}} R _ {n l} (r _ {i}) Y _ {l} ^ {m} (\hat {\boldsymbol {r}} _ {i}).
$$
Only the  $(lm)$  channels are involved in the representation of  $\mathrm{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  $(\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  $\rho_{k,k'}$ , that is, (8) becomes  $\phi_{ck} \circ g = \sum_{k'} \rho_{kk'}(g) \phi_{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
$$
B _ {c \alpha} ^ {K} = \sum_ {\mathbf {k} ^ {\prime}} C _ {\mathbf {k} ^ {\prime}} ^ {\alpha , K} \tilde {A} _ {c \mathbf {k} ^ {\prime}}, \quad \text {w h e r e} \tag {13}
$$
$$
\tilde {A} _ {c \boldsymbol {k}} = \prod_ {t = 1} ^ {n} \left(\sum_ {c ^ {\prime}} w _ {c c ^ {\prime}} \sum_ {x \in X} \phi_ {c ^ {\prime} k _ {t}} (x)\right). \tag {14}
$$
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 \to \infty$ . A smooth equivariant property  $\Phi$  may be approximated to within arbitrary accuracy by an expansion  $\Phi^K(X) \approx \sum_{c,\alpha} c_\alpha^K B_{c,\alpha}^K(X)$ . Since the embedding  $\tilde{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 re-expanding previous features in a new cluster expansion (Batatia et al., 2022b). The multi-set  $X$  is endowed with extra features,  $\pmb{h}_i^t = (h_{i,cK}^t)_{c,K}$ , that are updated for  $t \in \{1,\dots,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} = \left(x _ {i}, \boldsymbol {h} _ {i} ^ {t}\right) \tag {15}
$$
$$
\phi_ {c k} ^ {t} \left(x _ {i}, \boldsymbol {h} _ {i} ^ {t}\right) = \sum_ {\boldsymbol {\alpha}} w _ {\boldsymbol {\alpha}} ^ {t, c k} \sum_ {k ^ {\prime}, k ^ {\prime \prime}} C _ {k ^ {\prime} k ^ {\prime \prime}} ^ {\boldsymbol {\alpha}, k} h _ {i, c k ^ {\prime}} ^ {t} \phi_ {c k ^ {\prime \prime}} \left(x _ {i}\right) \tag {16}
$$
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, c K} ^ {t} = \sum_ {\boldsymbol {\alpha}} W _ {\boldsymbol {\alpha}} ^ {t, c K} \boldsymbol {B} _ {c \boldsymbol {\alpha}} ^ {t, K}. \tag {17}
$$
The gathered message  $\pmb{m}_i^t = (m_{i,cK}^t)_{c,k}$  is then used to update the particle states,
$$
x _ {i} ^ {t + 1} = \left(x _ {i}, \boldsymbol {h} _ {i} ^ {t + 1}\right), \quad \boldsymbol {h} _ {i} ^ {t + 1} = U _ {t} \left(\boldsymbol {m} _ {i} ^ {t}\right), \tag {18}
$$
where  $U_{t}$  can be an arbitrary 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} ^ {\text {l o c}} \left(x _ {i} ^ {t}\right), \quad \text {r e s p e c t i v e l y}, \quad y = \sum_ {t = 1} ^ {T} \mathcal {R} _ {t} ^ {\text {g l o b}} \left(\left\{x _ {i} ^ {t} \right\} _ {i}\right). \tag {19}
$$
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: