SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jinwoo Kim Tien Dat Nguyen Ayhan Suleymanzade
Hyeokjun An Seunghoon Hong
We present a novel framework to overcome the limitations of equivariant architectures in learning functions with group symmetries. In contrary to equivariant architectures, we use an arbitrary base model such as an MLP or a transformer and symmetrize it to be equivariant to the given group by employing a small equivariant network that parameterizes the probabilistic distribution underlying the symmetrization. The distribution is end-to-end trained with the base model which can maximize performance while reducing sample complexity of symmetrization. We show that this approach ensures not only equivariance to given group but also universal approximation capability in expectation. We implement our method on various base models, including patch-based transformers that can be initialized from pretrained vision transformers, and test them for a wide range of symmetry groups including permutation and Euclidean groups and their combinations. Empirical tests show competitive results against tailored equivariant architectures, suggesting the potential for learning equivariant functions for diverse groups using a non-equivariant universal base architecture. We further show evidence of enhanced learning in symmetric modalities, like graphs, when pretrained from non-symmetric modalities, like vision. Code is available at https://github.com/jw9730/lps.
Many perception problems in machine learning involve functions that are invariant or equivariant to certain symmetry group of transformations of data. Examples include learning on sets and graphs, point clouds, molecules, proteins, and physical data, to name a few [10, 12]. Equivariant architecture design has emerged as a successful approach, where every building block of a model is carefully restricted to be equivariant to a symmetry group of interest [5, 12, 23]. However, equivariant architecture design faces fundamental limitations, as individual construction of models for each group can be laborious or computationally expensive [91, 58, 64, 45], the architectural restrictions often lead to limited expressive power [101, 56, 105, 40], and the knowledge learned from one problem cannot be easily transferred to others of different symmetries as the architecture would be incompatible.
This motivates us to seek a symmetrization solution that can achieve group invariance and equivariance with general-purpose, group-agnostic architectures such as an MLP or a transformer. As a basic form of symmetrization, any parameterized function  $f_{\theta} : \mathcal{X} \to \mathcal{Y}$  on vector spaces  $\mathcal{X}, \mathcal{Y}$  can be made invariant or equivariant by group averaging [102, 67], i.e., averaging over all possible transformations of inputs  $\mathbf{x} \in \mathcal{X}$  and outputs  $\mathbf{y} \in \mathcal{Y}$  by a symmetry group  $G = \{g\}$ :
$$
\phi_ {\theta} (\mathbf {x}) = \frac {1}{| G |} \sum_ {g \in G} g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right), \tag {1}
$$
where  $\phi_{\theta}$  is equivariant or invariant to the group  $G$ . An important advantage is that the symmetrized function  $\phi_{\theta}$  can leverage the expressive power of the base function  $f_{\theta}$ ; it has been shown that  $\phi_{\theta}$  is a universal approximator of invariant or equivariant functions if  $f_{\theta}$  is a universal approximator [102],
Probabilistic Symmetrization
Figure 1: Overview of probabilistic symmetrization. We symmetrize an unconstrained base function  $f_{\theta}$  into an equivariant function  $\phi_{\theta, \omega}$  for group  $G$  using a learned equivariant distribution  $p_{\omega}(g|\mathbf{x})$ .
which includes an MLP [35] or a transformer [104]. On the other hand, an immediate challenge is that for many practical groups involving permutation and rotations, the cardinality of the group  $|\bar{G}|$  is large or infinite, so the exact averaging is intractable. Due to this, existing symmetrization approaches often focus on small finite groups [4, 65, 94, 42], manually derive smaller subsets of the entire group e.g., a frame [74] to average over [67, 68], or implement a relaxed version of equivariance [41, 83].
An alternative method for tractability is to interpret Eq. (1) as an expectation with uniform distribution  $\mathrm{Unif}(G)$  over the compact group  $G$  [62], and use sampling-based average to estimate it [102, 67, 68]:
$$
\phi_ {\theta} (\mathbf {x}) = \mathbb {E} _ {g \sim \operatorname {U n i f} (G)} \left[ g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right) \right], \tag {2}
$$
where  $g \cdot f_{\theta}(g^{-1} \cdot \mathbf{x})$  serves as an unbiased estimator of  $\phi_{\theta}(\mathbf{x})$ . While simple and general, this approach has practical issues that the base function  $f_{\theta}$  is burdened to learn all equally possible group transformations, and the expectedly high variance of the estimator can lead to challenges in sampling-based training due to large variance of gradients as well as sample complexity of inference.
Our key idea is to replace the uniform distribution  $\mathrm{Unif}(G)$  for the expectation in Eq. (2) with a parameterized distribution  $p_{\omega}(g|\mathbf{x})$  in a way that equivariance and expressive power are always guaranteed, and train it end-to-end with the base function  $f_{\theta}$  to directly minimize task loss. We show that the distribution  $p_{\omega}(g|\mathbf{x})$  only needs to satisfy one simple condition to guarantee equivariance and expressive power: it has to be probabilistically equivariant [9]. This allows us to generally implement  $p_{\omega}(g|\mathbf{x})$  as a noise-outsourced map  $(\mathbf{x},\epsilon)\mapsto g$  with an invariant noise  $\epsilon$  and a small equivariant network  $q_{\omega}$ , which enables gradient-based training with reparameterization [47]. As  $p_{\omega}$  is trained, it can enhance the learning of  $f_{\theta}$  by producing group transformations of lower variance compared to  $\mathrm{Unif}(G)$  so that  $f_{\theta}$  is less burdened, and coordinating with  $f_{\theta}$  to maximize task performance. We refer to our approach as probabilistic symmetrization. An overview is provided in Figure 1.
We implement and test our method with two general-purpose architectures as the base function  $f_{\theta}$ : MLP and transformer. In particular, our transformer backbone is architecturally identical to patch-based vision transformers [26], which allows us to initialize most of its parameters from ImageNet-21k pretrained weights [89] and only replace the input and output projections to match task dimensions. We implement the conditional distribution  $p_{\omega}(g|\mathbf{x})$  for a wide range of practical symmetry groups including permutation  $(\mathbf{S}_n)$  and Euclidean groups  $(\mathrm{O} / \mathrm{SO}(d)$  and  $\mathrm{E} / \mathrm{SE}(d))$  and their product combinations (e.g.,  $\mathbf{S}_n \times \mathbf{O}(3)$ ), all of which are combinatorial or infinite groups. Empirical tests on a wide range of invariant and equivariant tasks involving graphs and motion data show competitive results against tailored equivariant architectures as well as existing symmetrization methods [74, 41]. This suggests the potential for learning invariant or equivariant functions for diverse groups with a group-agnostic general-purpose backbone. We further show evidence that pretraining from non-symmetric modality (vision) leads to enhanced learning in symmetric modality (graphs).
2 Probabilistic Symmetrization for Equivalence
We introduce and analyze our approach called probabilistic symmetrization, which involves an equivariant distribution  $p_{\omega}$  and group-agnostic base function  $f_{\theta}$ , in Section 2.1. We then describe implementation of  $p_{\omega}$  for practical groups including permutations and rotations in Section 2.2. Then, we describe our choice of base function  $f_{\theta}$  focusing on MLP and, transformers in particular, in Section 2.3. All proofs can be found in Appendix A.1.
Problem Setup In general, our goal is to construct a function  $\phi : \mathcal{X} \to \mathcal{Y}$  on finite vector spaces  $\mathcal{X}, \mathcal{Y}$  that is invariant or equivariant to symmetry specified by a group  $G = \{g\}^1$ . This is formally described by specifying how the group acts as transformations on the input and output. A group representation  $\rho : G \to \mathrm{GL}(\mathcal{X})$ , where  $\mathrm{GL}(\mathcal{X})$  is the set of all invertible matrices on  $\mathcal{X}$ , associates each group element  $g \in G$  to an invertible matrix  $\rho(g)$  that transforms a given vector  $\mathbf{x} \in \mathcal{X}$  through  $\mathbf{x} \mapsto g \cdot \mathbf{x} = \rho(g)\mathbf{x}$ . Given that, a function  $\phi : \mathcal{X} \to \mathcal{Y}$  is  $G$  equivariant if:
$$
\phi \left(\rho_ {1} (g) \mathbf {x}\right) = \rho_ {2} (g) \phi (\mathbf {x}), \quad \forall \mathbf {x} \in \mathcal {X}, g \in G, \tag {3}
$$
where the representations  $\rho_{1}$  and  $\rho_{2}$  are on the input and output, respectively.  $G$  invariance is a special case of equivariance when the output representation is trivial,  $\rho_{2}(g) = \mathbf{I}$ .
2.1 Probabilistic Symmetrization
To construct a  $G$  equivariant function  $\phi_{\theta}$ , group averaging symmetrizes an arbitrary base function  $f_{\theta}:\mathcal{X}\to \mathcal{Y}$  by taking expectation with uniform distribution over the group (Eq. (2)). Instead, we propose to use an input-conditional parameterized distribution  $p_{\omega}(g|\mathbf{x})$  and symmetrize  $f_{\theta}$  as follows:
$$
\phi_ {\theta , \omega} (\mathbf {x}) = \mathbb {E} _ {p _ {\omega} (g | \mathbf {x})} \left[ \rho_ {2} (g) f _ {\theta} \left(\rho_ {1} (g) ^ {- 1} \mathbf {x}\right) \right], \tag {4}
$$
where the distribution  $p_{\omega}(g|\mathbf{x})$  itself satisfies probabilistic  $G$  equivariance:
$$
p _ {\omega} (g | \mathbf {x}) = p _ {\omega} \left(g ^ {\prime} g \mid \rho_ {1} \left(g ^ {\prime}\right) \mathbf {x}\right), \quad \forall \mathbf {x} \in \mathcal {X}, g, g ^ {\prime} \in G. \tag {5}
$$
Importantly, we show that probabilistic symmetrization with equivariant  $p_{\omega}$  guarantees equivariance as well as expressive power of the symmetrized  $\phi_{\theta,\omega}$ .
Theorem 1. If  $p_{\omega}$  is  $G$ -equivariant, then  $\phi_{\theta, \omega}$  is  $G$ -equivariant for arbitrary  $f_{\theta}$ .
Proof. The proof can be found in Appendix A.1.1.
Theorem 2. If  $p_{\omega}$  is  $G$  equivariant and  $f_{\theta}$  is a universal approximator, then  $\phi_{\theta,\omega}$  is a universal approximator of  $G$  equivariant functions.
Proof. The proof can be found in Appendix A.1.2.
While the base function  $f_{\theta}$  that guarantees universal approximation can be chosen in a group-agnostic manner e.g., an MLP [35, 20] or a transformer on token sequences [104], the distribution  $p_{\omega}$  needs to be instantiated group specifically to satisfy  $G$  equivariance. A simplistic choice is using uniform distribution  $\mathrm{Unif}(G)$  for all inputs  $\mathbf{x}$  with no parameterization (reducing to group averaging), which is technically equivariant and therefore guarantees equivariance and universality. However, appropriately parameterizing and learning  $p_{\omega}$  can provide distinguished advantages compared to the uniform distribution, as it can (1) learn from data to collaborate with (pre-trained) base function  $f_{\theta}$  to maximize task performance and (2) learn to produce more consistent samples  $g \sim p_{\omega}(g|\mathbf{x})$  that can offer more stable gradients for the base function  $f_{\theta}$  during early training.
We now provide a generic blueprint of  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  for any compact group  $G$ . Our goal is to sample  $g \sim p_{\omega}(g|\mathbf{x})$  to obtain group representation  $\rho(g)$  for symmetrization (Eq. (4)) in a differentiable manner so that  $p_{\omega}$  can be trained end-to-end. Since we only need sampling and there is no need to evaluate likelihoods, we simply implement  $p_{\omega}(g|\mathbf{x})$  as a noise-outsourced, differentiable transformation  $q_{\omega}(\mathbf{x}, \epsilon)$  of a noise variable  $\epsilon \in \mathcal{E}$  that directly outputs a group representation  $\rho(g)$ :
$$
\rho (g) = q _ {\omega} (\mathbf {x}, \epsilon), \quad \epsilon \sim p (\epsilon), \tag {6}
$$
In this paper, we assume the group  $G$  to be compact.
where  $q_{\omega}$  is  $G$  equivariant and  $p(\epsilon)$  is  $G$  invariant under an appropriate representation  $\rho'$ :
$$
q _ {\omega} \left(\rho_ {1} (g) \mathbf {x}, \rho^ {\prime} (g) \boldsymbol {\epsilon}\right) = \rho (g) q _ {\omega} (\mathbf {x}, \boldsymbol {\epsilon}), \quad p (\boldsymbol {\epsilon}) = p \left(\rho^ {\prime} (g) \boldsymbol {\epsilon}\right), \quad \forall \mathbf {x} \in \mathcal {X}, \boldsymbol {\epsilon} \in \mathcal {E}, g \in G. \tag {7}
$$
Given above implementation, we can show the  $G$  equivariance of  $p_{\omega}$ :
Theorem 3. If  $q_{\omega}$  is  $G$  equivariant and  $p(\epsilon)$  is  $G$  invariant under representation  $\rho'$  that  $|\operatorname{det} \rho'(g)| = 1 \forall g \in G$ , the distribution  $p_{\omega}(g|\mathbf{x})$  characterized by  $q_{\omega} : (\mathbf{x}, \epsilon) \mapsto \rho(g)$  is  $G$  equivariant.
Proof. The proof can be found in Appendix A.1.3.
In practice, one can use any available  $G$  equivariant neural network to implement  $q_{\omega}$ , e.g., a graph neural network for the symmetric group  $\mathbf{S}_n$ , or an equivariant MLP which can be constructed for any matrix group [30]. Since we expect most of the reasoning to be done by the base function  $f_{\theta}$ , the equivariant network  $q_{\omega}$  can be small and relatively less expressive. This allows us to get less affected by their known issues in expressiveness and scaling [101, 70, 14, 40]. For the noise  $\epsilon \sim p(\epsilon)$ , simple choices often suffice for  $G$  invariance. For example, standard normal  $\epsilon \sim \mathcal{N}(0,\mathbf{I}_n)$  provides invariance for the symmetric group  $\mathbf{S}_n$  as well as the (special) orthogonal group  $\mathrm{O}(n)$  and  $\mathrm{SO}(n)$ .
One important detail in designing  $q_{\omega}$  is constraining its output to be a valid group representation  $\rho(g)$ . For this, we apply appropriate postprocessing to refine neural network features into group representations, e.g., Gram-Schmidt orthogonalization to obtain a representation  $\rho(g) \in \mathbb{R}^{n \times n}$  of the orthogonal group  $g \in \mathrm{O}(n)$ . Importantly, to not break the  $G$  equivariance of  $q_{\omega}$ , this postprocessing needs to be equivariant itself, e.g., Gram-Schmidt process is itself  $\mathrm{O}(n)$  equivariant [41]. Implementations of  $p_{\omega}$  for a range of practical symmetry groups are provided in detail in Section 2.2.
2.2 Equivariant Distribution  $p_{\omega}$
We present implementations of the  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  for a range of practical groups demonstrated in our experiments (Section 3). Formal proofs of correctness are in Appendix A.1.4.
Symmetric Group  $\mathbf{S}_n$  The symmetric group  $\mathbf{S}_n$  over a finite set of  $n$  elements contains all permutations of the set, which describes symmetry to ordering desired for learning set and graph data. The base representation is given by  $\rho(g) = \mathbf{P}_g$  where  $\mathbf{P}_g \in \{0,1\}^{n \times n}$  is a permutation matrix for  $g$ .
To implement  $S_{n}$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  that provides permutation matrices  $\mathbf{P}_g$  from graph data  $\mathbf{x}$ , we use the following design. We first sample invariant noise  $\epsilon \in \mathbb{R}^{n\times d}$  from i.i.d. uniform Unif[0,  $\eta$ ] with noise scale  $\eta$ . For the  $S_{n}$  equivariant map  $q_{\omega}:(\mathbf{x},\epsilon)\mapsto \rho (g)$ , we first use a graph neural network (GNN) as an equivariant map  $(\mathbf{x},\epsilon)\mapsto \mathbf{Z}$  that outputs nodewise scalar  $\mathbf{Z}\in \mathbb{R}^n$ . Then, assuming  $\mathbf{Z}$  is tie-free<sup>2</sup>, we use below argsort operator  $\mathbf{Z}\mapsto \mathbf{P}_g$  to obtain permutation matrix [73, 98]:
$$
\mathbf {P} _ {g} = \operatorname {e q} \left(\mathbf {Z} \mathbf {1} ^ {\top}, \mathbf {1} \operatorname {s o r t} (\mathbf {Z}) ^ {\top}\right), \tag {8}
$$
where eq denotes elementwise equality indicator. The argsort operator is  $S_{n}$  equivariant, i.e., it maps  $\mathbf{P}_{g'}\mathbf{Z} \mapsto \mathbf{P}_{g'}\mathbf{P}_g$  for all  $\mathbf{P}_{g'} \in S_n$ . To backpropagate through  $\mathbf{P}_g$  during training, we use straight-through gradient estimator [6] with an approximate permutation matrix  $\hat{\mathbf{P}}_g \approx \mathbf{P}_g$  obtained from a differentiable relaxation of argsort operator [61, 33, 98]; details can be found in Appendix A.3.1.
Orthogonal Group  $\mathrm{O}(n)$ ,  $\mathrm{SO}(n)$  The orthogonal group  $\mathrm{O}(n)$  contains all roto-reflections in  $\mathbb{R}^n$  around origin, and the special orthogonal group  $\mathrm{SO}(n)$  contains all rotations without reflections. These groups describe rotation symmetries desirable in learning geometric data. The base group representation for  $\mathrm{O}(n)$  is given by  $\rho(g) = \mathbf{Q}_g$  where  $\mathbf{Q}_g$  is the orthogonal matrix for  $g$ . For  $\mathrm{SO}(n)$ , the representation is  $\rho(g) = \mathbf{Q}_g^+$  where  $\mathbf{Q}_g^+$  is the orthogonal matrix of  $g$  with determinant  $+1$ .
To implement  $\mathrm{O}(n) / \mathrm{SO}(n)$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  that provides orthogonal matrices given input data  $\mathbf{x}$ , we use the following design. We first sample invariant noise  $\epsilon \in \mathbb{R}^{n\times d}$  from i.i.d. normal  $\mathcal{N}(0,\eta^2)$  with noise scale  $\eta$ . For the  $\mathrm{O}(n) / \mathrm{SO}(n)$  equivariant map  $p_{\omega}:(\mathbf{x},\epsilon)\mapsto \rho (g)$ , we first use an equivariant neural network as a map  $(\mathbf{x},\epsilon)\mapsto \mathbf{Z}$  that outputs  $n$  features  $\mathbf{Z}\in \mathbb{R}^{n\times n}$ . Then, assuming  $\mathbf{Z}$  is full-rank<sup>3</sup>, we use Gram-Schmidt process for orthogonalization, which is differentiable
2This can be assumed since the invariant noise  $\epsilon \in \mathbb{R}^{n\times d}$  serves as tiebreaker between the  $n$  nodes.
In practice, we add a random Gaussian matrix of a tiny magnitude to  $\mathbf{Z}$  to prevent rank collapse.
and  $\mathbf{O}(n)$  equivariant [41]. This completes the postprocessing  $\mathbf{Z} \mapsto \mathbf{Q}_g$  for  $\mathrm{O}(n)$ . For  $\mathrm{SO}(n)$ , we can further use a simple scale operator  $\mathbf{Q} \mapsto \mathbf{Q}_g^+$  to set the determinant of orthogonalized matrix to  $+1$ :
$$
\operatorname {s c a l e}: \left[ \begin{array}{c c c} \mathbf {Q} _ {1} & \dots & \mathbf {Q} _ {n} \end{array} \right] \mapsto \left[ \begin{array}{c c c} \det  (\mathbf {Q}) \cdot \mathbf {Q} _ {1} & \dots & \mathbf {Q} _ {n} \end{array} \right], \tag {9}
$$
The scale operator is differentiable and  $\mathrm{SO}(n)$  equivariant, i.e., it maps  $\mathbf{Q}_{g'}^+ \mathbf{Q} \mapsto \mathbf{Q}_{g'}^+ \mathbf{Q}_g^+$  for all  $\mathbf{Q}_{g'}^+ \in \mathrm{SO}(n)$ , thereby completing the postprocessing  $\mathbf{Z} \mapsto \mathbf{Q}_g^+$  for  $\mathrm{SO}(n)$ .
Euclidean Group  $\mathrm{E}(n)$ ,  $\mathrm{SE}(n)$  The Euclidean group  $\mathrm{E}(n)$  contains all roto-translations and reflections in  $\mathbb{R}^n$  and their combinations, and the special Euclidean group  $\mathrm{SE}(n)$  contains all roto-translations without reflections. These groups are desired in learning physical systems such as a particle in motion. Formally, the Euclidean group is given as a combination of orthogonal group and translation group  $\mathrm{E}(n) = \mathrm{O}(n) \ltimes \mathrm{T}(n)$ , and similarly  $\mathrm{SE}(n) = \mathrm{SO}(n) \ltimes \mathrm{T}(n)$ . As the translation group  $\mathrm{T}(n)$  is non-compact which violates our assumption for symmetrization, we handle it separately. Following prior work [74, 41], we subtract the centroid  $\bar{\mathbf{x}}$  from the input data  $\mathbf{x}$  as  $\mathbf{x} - \bar{\mathbf{x}}$ , and add it to the rotation symmetrized output as  $\bar{\mathbf{x}} + g \cdot f_{\theta}(g^{-1} \cdot (\mathbf{x} - \bar{\mathbf{x}})))$  where  $g$  is sampled from  $\mathrm{O}(n) / \mathrm{SO}(n)$  equivariant  $p_{\omega}(g|\mathbf{x} - \bar{\mathbf{x}})$ . This makes the overall symmetrized function  $\mathrm{E}(n) / \mathrm{SE}(n)$  equivariant.
Product Group  $H \times K$  While we have described several groups individually, in practice we often encounter product combinations of groups  $G = H \times K$  that describe joint symmetry to each group  $H$  and  $K$ . For example,  $S_{n} \times O(3)$  describes joint symmetry to permutations and rotations, which is desired in learning point clouds, molecules, and particle interactions. In general, an element of  $H \times K$  is given as  $g = (h,k)$  where  $h \in H$  and  $k \in K$ , and group operations are applied elementwise  $gg' = (h,k)(h',k') = (hh',kk')$ . The base group representation is accordingly given as pair of representations  $\rho(g) = (\rho(h),\rho(k))$ . While a common approach to handling  $H \times K$  is partially symmetrizing on  $H$  and imposing  $K$  equivariance on the base architecture (e.g., rotational symmetrization of a graph neural network for  $S_{n} \times O(3)$  equivariance [74, 41]), we extend to full symmetrization on  $H \times K$  since our goal is not imposing any constraint on the base function  $f_{\theta}$ .
To implement  $H \times K$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  that gives  $\rho(g) = (\rho(h), \rho(k))$  from data  $\mathbf{x}$ , we use the following design. We first sample invariant noise  $\epsilon$  from i.i.d. normal  $\mathcal{N}(0, \eta^2)$  with scale  $\eta$ . For the  $H \times K$  equivariant map  $q_{\omega} : (\mathbf{x}, \epsilon) \mapsto (\rho(h), \rho(k))$ , we employ a  $H \times K$  equivariant neural network as a map  $(\mathbf{x}, \epsilon) \mapsto (\mathbf{Z}_H, \mathbf{Z}_K)$  such that the postprocessing for each group  $H$  and  $K$  provides maps  $\mathbf{Z}_H \mapsto \rho(h)$  and  $\mathbf{Z}_K \mapsto \rho(k)$  respectively, leading to full representation  $\rho(g) = (\rho(h), \rho(k))$ . For this whole procedure to be  $H \times K$  equivariant, it is sufficient to have  $\mathbf{Z}_H$  be  $K$  invariant and  $\mathbf{Z}_K$  be  $H$  invariant. These are special cases of  $H \times K$  equivariance, and is supported by a range of equivariant neural networks especially regarding  $S_n \times O(3)$  or  $S_n \times SO(3)$  equivariances [23].
2.3 Base Function  $f_{\theta}$
We now describe the choice of group-agnostic base function  $f_{\theta} : \mathbf{x} \mapsto \mathbf{y}$ . As group symmetry is handled by the equivariant distribution  $p_{\omega}(g|\mathbf{x})$ , any symmetry concern is hidden from  $f_{\theta}$ , allowing the inputs  $\mathbf{x}$  and outputs  $\mathbf{y}$  to be treated as plain multidimensional arrays. This allows us to implement  $f_{\theta}$  with powerful general-purpose architectures, namely as an MLP that operates on flattened vectors of inputs and outputs, or a transformer as we describe below.
Let inputs  $\mathbf{x} \in \mathcal{X} = \mathbb{R}^{n_1 \times \ldots \times n_a \times c}$  and outputs  $\mathbf{y} \in \mathcal{Y} = \mathbb{R}^{n_1' \times \ldots \times n_b' \times c'}$  be multidimensional arrays with  $c$  and  $c'$  channels, respectively. Our transformer base function  $f_{\theta}: \mathbf{x} \mapsto \mathbf{y}$  is given as:
$$
f _ {\theta} = \text {d e t o k i n i z e} \circ \text {t r a n s f o r m e r} \circ \text {t o k i n i z e}, \tag {10}
$$