SlowGuess/NeurIPS_2023_OCR huggingface dataset:
showing the  $\operatorname{E}(n)$  equivariance of  $\phi_{\theta, \omega}$ .
Proposition 7. The proposed symmetrization  $\phi_{\theta, \omega}$  for special Euclidean group  $\mathrm{SE}(n)$  is equivariant.
Proof. We prove  $\phi_{\theta, \omega}(g' \cdot \mathbf{x}) = g' \cdot \phi_{\theta, \omega}(\mathbf{x})$  for all  $\mathbf{x} \in \mathcal{X}$  and  $g' \in \mathrm{SE}(n)$ , in an analogous manner to the proof of Proposition 6. From Eq. (29), we have:
$$
\phi_ {\theta , \omega} \left(g ^ {\prime} \cdot \mathbf {x}\right) = \mathbb {E} _ {p _ {\omega} \left(g \mid g ^ {\prime} \cdot \mathbf {x} - \overline {{g ^ {\prime} \cdot \mathbf {x}}} \mathbf {1} ^ {\top}\right)} \left[ \overline {{g ^ {\prime} \cdot \mathbf {x}}} \mathbf {1} ^ {\top} + g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \left(g ^ {\prime} \cdot \mathbf {x} - \overline {{g ^ {\prime} \cdot \mathbf {x}}} \mathbf {1} ^ {\top}\right)\right) \right]. \tag {35}
$$
In general, an element of special Euclidean group  $g' \in \mathrm{SE}(n)$  acts on data  $\mathbf{x} \in \mathbb{R}^{n \times d}$  via  $g' \cdot \mathbf{x} = \mathbf{Q}_{g'}^+ \mathbf{x} + \mathbf{t}_{g'} \mathbf{1}^\top$  where  $\mathbf{Q}_{g'}^+ \in \mathrm{SO}(n)$  is rotation component and  $\mathbf{t}_{g'} \in \mathbb{R}^n$  is translation [74, 41]. With this, the centroid of the transformed data  $g' \cdot \mathbf{x}$  is given as follows:
$$
\overline {{g ^ {\prime} \cdot \mathbf {x}}} = \overline {{\mathbf {Q} _ {g ^ {\prime}} ^ {+} \mathbf {x} + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top}}} = \overline {{\mathbf {Q} _ {g ^ {\prime}} ^ {+} \mathbf {x}}} + \mathbf {t} _ {g ^ {\prime}} = \mathbf {Q} _ {g ^ {\prime}} ^ {+} \overline {{\mathbf {x}}} + \mathbf {t} _ {g ^ {\prime}}, \tag {36}
$$
$$
\begin{array}{l} g ^ {\prime} \cdot \mathbf {x} - \overline {{g ^ {\prime} \cdot \mathbf {x}}} \mathbf {1} ^ {\top} = \mathbf {Q} _ {g ^ {\prime}} ^ {+} \mathbf {x} + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} - \mathbf {Q} _ {g ^ {\prime}} ^ {+} \overline {{\mathbf {x}}} \mathbf {1} ^ {\top} - \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} \\ = \mathbf {Q} _ {g ^ {\prime}} ^ {+} (\mathbf {x} - \bar {\mathbf {x}} \mathbf {1} ^ {\top}). \tag {37} \\ \end{array}
$$
Similar as in Proposition 6, subtracting centroid only leaves  $\mathrm{SO}(n)$  component. We then have:
$$
\begin{array}{l} \phi_ {\theta , \omega} (g ^ {\prime} \cdot \mathbf {x}) = \mathbb {E} _ {p _ {\omega} (g | \mathbf {Q} _ {g ^ {\prime}} ^ {+} (\mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top}))} \left[ \mathbf {Q} _ {g ^ {\prime}} ^ {+} \bar {\mathbf {x}} \mathbf {1} ^ {\top} + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} + g \cdot f _ {\theta} (g ^ {- 1} \cdot (\mathbf {Q} _ {g ^ {\prime}} ^ {+} (\mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top}))) \right] \\ = \mathbb {E} _ {p _ {\omega} (g | g ^ {\prime} \cdot (\mathbf {x} - \bar {\mathbf {x}} \mathbf {1} ^ {\top}))} \left[ g ^ {\prime} \cdot \bar {\mathbf {x}} \mathbf {1} ^ {\top} + g \cdot f _ {\theta} \left(g ^ {- 1} g ^ {\prime} \cdot (\mathbf {x} - \bar {\mathbf {x}} \mathbf {1} ^ {\top})\right) \right] + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top}, \tag {38} \\ \end{array}
$$
where, inside the expectation, we interpret the rotation component of  $g'$  as an element of the special orthogonal group  $\mathrm{SO}(n)$ . Similar as in Theorem 1, we introduce  $h = g'^{-1}g \in \mathrm{SO}(n)$  that  $g = g'h$ . As the distribution  $p_{\omega}$  is  $\mathrm{SO}(n)$  equivariant, we have  $p_{\omega}(g|g' \cdot (\mathbf{x} - \overline{\mathbf{x}}\mathbf{1}^{\top})) = p_{\omega}(g'^{-1}g|g'^{-1}g' \cdot (\mathbf{x} - \overline{\mathbf{x}}\mathbf{1}^{\top})) = p_{\omega}(h|\mathbf{x} - \overline{\mathbf{x}}\mathbf{1}^{\top})$ . We then rewrite the expectation with respect to  $h$ :
$$
\begin{array}{l} \phi_ {\theta , \omega} (g ^ {\prime} \cdot \mathbf {x}) = \mathbb {E} _ {p _ {\omega} (h | \mathbf {x} - \bar {\mathbf {x}} \mathbf {1} ^ {\top})} [ g ^ {\prime} \cdot \bar {\mathbf {x}} \mathbf {1} ^ {\top} + g ^ {\prime} h \cdot f _ {\theta} ((g ^ {\prime} h) ^ {- 1} g ^ {\prime} \cdot (\mathbf {x} - \bar {\mathbf {x}} \mathbf {1} ^ {\top})) ] + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} \\ = \mathbb {E} _ {p _ {\omega} (h | \mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top})} \left[ g ^ {\prime} \cdot \bar {\mathbf {x}} \mathbf {1} ^ {\top} + g ^ {\prime} h \cdot f _ {\theta} \left(h ^ {- 1} \cdot (\mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top})\right) \right] + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} \\ = \mathbf {Q} _ {g ^ {\prime}} ^ {+} \mathbb {E} _ {p _ {\omega} (h | \mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top})} [ \bar {\mathbf {x}} \mathbf {1} ^ {\top} + h \cdot f _ {\theta} (h ^ {- 1} \cdot (\mathbf {x} - \overline {{\mathbf {x}}} \mathbf {1} ^ {\top})) ] + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} \\ = \mathbf {Q} _ {g ^ {\prime}} ^ {+} \phi_ {\theta , \omega} (\mathbf {x}) + \mathbf {t} _ {g ^ {\prime}} \mathbf {1} ^ {\top} \\ = g ^ {\prime} \cdot \phi_ {\theta , \omega} (\mathbf {x}), \tag {39} \\ \end{array}
$$
showing the  $\mathrm{SE}(n)$  equivariance of  $\phi_{\theta ,\omega}$
Product Group  $H \times K$  For the product group  $H \times K$ , we assume that the base representation for each element  $g = (h, k)$  is given as a pair of representations  $\rho(g) = (\rho(h), \rho(k))$ . Without loss of generality, we further assume that the representation  $\rho(g)$  can be expressed as the Kronecker product  $\rho(g) = \rho(h) \otimes \rho(k)$  that acts on flattened data  $\operatorname{vec}(\mathbf{x})$  as  $\mathbf{x} \mapsto \operatorname{vec}^{-1}(\rho(g)\operatorname{vec}(\mathbf{x}))$ . This follows the standard approach in equivariant deep learning [30, 57] that deals with composite representations using direct sum and tensor products of base group representations.
Above approach applies to many practical product groups, including sets and graphs with Euclidean attributes  $(\mathbf{S}_n\times \mathrm{O}(d) / \mathrm{SO}(d)^6)$  and sets of symmetric elements  $(\mathbf{S}_n\times H)$  in general [59]. For
This is after handling the translation component of the Euclidean group  $\mathrm{E}(d) / \mathrm{SE}(d)$  as in Eq. (29).
Group  $G$
Orbits  $G_{\mathbf{x}}g$
Distribution  $p_{\omega}(g|\mathbf{x})$
Figure 3: Visual illustration of the symmetrization methods based on probabilities assigned upon the partitioning of the group  $G$  into orbits  $G_{\mathbf{x}}g$ . Note that, while we use concentric circles of different perimeters to illustrate each orbit, all orbits actually have an identical cardinality  $|G_{\mathbf{x}}g| = |G_{\mathbf{x}}|$ .
Frame  $F(\mathbf{x})$
Canonicalizer  $C_{\omega}(\mathbf{x})$
example, for the group  $\mathbf{S}_n\times \mathbf{O}(d)$  on data  $\mathbf{x}\in \mathbb{R}^{n\times d}$ , an element  $g = (h,k)$  has representation  $\rho (g) = \rho (h)\otimes \rho (k)\in \mathbb{R}^{nd\times nd}$  combined from permutation  $\rho (h)\in \mathbb{R}^{n\times n}$  and rotation  $\rho (k)\in \mathbb{R}^{d\times d}$ , which acts by  $\mathbf{x}\mapsto \mathrm{vec}^{-1}(\rho (g)\mathrm{vec}(\mathbf{x}))$  or more simply  $\mathbf{x}\mapsto \rho (h)\mathbf{x}\rho (k)^{\top}$ .
Now we recall that the  $p_{\omega}(g|\mathbf{x})$  for the product group  $H\times K$  is implemented as follows:
1. Sample noise  $\epsilon \in \mathcal{E}$  from i.i.d. normal  $\mathcal{N}(0,\eta^2)$  such that  $p(\epsilon)$  is invariant under representations of  $H$  and  $K$  that satisfy  $|\operatorname*{det}\rho '(h)| = 1$  and  $|\operatorname*{det}\rho '(k)| = 1$ , respectively. For example, for  $\mathrm{S}_n\times \mathrm{O}(d)$ , the noise  $\epsilon \in \mathbb{R}^{n\times d}$  that follows i.i.d. normal  $\mathcal{N}(0,\eta^2)$  is invariant under base representations of both  $\mathrm{S}_n$  and  $\mathrm{O}(d)$  which are orthogonal.
2. Use a  $H \times K$  equivariant neural network to obtain features  $(\mathbf{x}, \epsilon) \mapsto (\mathbf{Z}_H, \mathbf{Z}_K)$  where  $\mathbf{Z}_H$  is  $K$  invariant and  $\mathbf{Z}_K$  is  $H$  invariant. For example, for  $\mathbf{S}_n \times \mathrm{O}(d)$ , we expect node-level scalar features  $\mathbf{Z}_{\mathbf{S}_n} \in \mathbb{R}^n$  to be  $\mathrm{O}(d)$  invariant and  $d$  global rotary features  $\mathbf{Z}_{\mathrm{O}(d)} \in \mathbb{R}^{d \times d}$  to be  $\mathbf{S}_n$  invariant.
3. Apply postprocessing for  $H$  and  $K$  groups onto  $\mathbf{Z}_H$  and  $\mathbf{Z}_K$  respectively to obtain representations  $\mathbf{Z}_H\mapsto \rho (h)$  and  $\mathbf{Z}_K\mapsto \rho (k)$  of  $H$  and  $K$  groups respectively. For example, for  $\mathbf{S}_n\times \mathbf{O}(d)$ , we use argsort in Eq. (23) to obtain  $\mathbf{Z}_{\mathrm{S}_n}\mapsto \rho (h)$  and Gram-Schmidt process to obtain  $\mathbf{Z}_{\mathrm{O}(d)}\mapsto \rho (k)$ .
4. Combine the representations  $\rho(g) = (\rho(h), \rho(k))$  to obtain a representation for the  $H \times K$  group.
Proposition 8. The proposed distribution  $p_{\omega}(g|\mathbf{x})$  for the product group  $H\times K$  is equivariant.
Proof. By assumption,  $p(\epsilon)$  is invariant under representations of  $H$  and  $K$  that satisfy  $|\operatorname{det}\rho'(h)| = 1$  and  $|\operatorname{det}\rho'(k)| = 1$ , respectively. This implies  $H \times K$  invariance as well, since  $p(\epsilon) = p(h \cdot \epsilon) = p(k \cdot \epsilon)$  for all  $\epsilon \in \mathcal{E}, h \in H, k \in K$  gives  $p(k \cdot h \cdot \epsilon) = p(k \cdot (h \cdot \epsilon)) = p(h \cdot \epsilon) = p(\epsilon)$ , and Kronecker product of matrices of determinant 1 gives a matrix of determinant 1. Furthermore, the map  $(\mathbf{x}, \epsilon) \mapsto (\rho(h), \rho(k)) = \rho(g)$  is overall  $H \times K$  equivariant, since an input transformed with  $g' = (h', k')$  is first mapped by the equivariant neural network as  $(g' \cdot \mathbf{x}, g' \cdot \epsilon) \mapsto (h' \cdot \mathbf{Z}_H, k' \cdot \mathbf{Z}_K)$ , then postprocessed as  $(h' \cdot \mathbf{Z}_H, k' \cdot \mathbf{Z}_K) \mapsto (\rho(h') \rho(h), \rho(k') \rho(k)) = (\rho(h'), \rho(k')) \cdot (\rho(h), \rho(k)) = \rho(g') \rho(g)$ . Combining the above, by Theorem 3, the distribution  $p_{\omega}(g|\mathbf{x})$  is  $H \times K$  equivariant.
A.1.5 Proof of Proposition 1 and Proposition 2 (Section 2.4)
Before proceeding to proofs, we recall that the stabilizer subgroup  $G_{\mathbf{x}}$  of a group  $G$  for  $\mathbf{x}$  is defined as  $\{g' \in G : g' \cdot \mathbf{x} = \mathbf{x}\}$  and acts on a given group element  $g \in G$  through left multiplication  $g \mapsto g'g$ . For some  $g \in G$ , by  $G_{\mathbf{x}}g$  we denote its orbit under the action by  $G_{\mathbf{x}}$ , i.e., the set of elements in  $G$  to which  $g$  can be moved by the action of elements  $g' \in G_{\mathbf{x}}$ . Importantly, we can show the following:
Property 1. Any group  $G$  is a union of disjoint orbits  $G_{\mathbf{x}}g$  of equal cardinality.
Proof. Let us consider the equivalence relation  $\sim$  on  $G$  induced by the action of the stabilizer  $G_{\mathbf{x}}$ , defined as  $g \sim h \iff h \in G_{\mathbf{x}}g$ . The orbits  $G_{\mathbf{x}}g$  are the equivalence classes under this relation, and the set of all orbits of  $G$  under the action of  $G_{\mathbf{x}}$  forms a partition of  $G$  (i.e., the quotient  $G / G_{\mathbf{x}}$ ). Furthermore, since  $G_{\mathbf{x}} \leq G$  and right multiplication by some  $g \in G$  is a faithful action of  $G$  on itself, we have  $|G_{\mathbf{x}}g| = |G_{\mathbf{x}}|$  for all  $g \in G$ , which shows that all orbits  $G_{\mathbf{x}}g$  have equal cardinality.
The partition of group  $G$  into disjoint orbits  $G_{\mathbf{x}}g$  is illustrated in the first and second panel of Figure 3. We now show the following:
Property 2.  $G$  equivariant  $p_{\omega}(g|\mathbf{x})$  assigns identical probability to all elements on each orbit  $G_{\mathbf{x}}g$ .
Proof. With equivariance, we have  $p_{\omega}(g|\mathbf{x}) = p_{\omega}(g'g|g'\cdot \mathbf{x})$ . Since  $g^{\prime}\cdot \mathbf{x} = \mathbf{x}$  for all  $g^{\prime}\in G_{\mathbf{x}}$ , we have  $p_{\omega}(g|\mathbf{x}) = p_{\omega}(g'g|\mathbf{x})$  for all  $g^{\prime}\in G_{\mathbf{x}}$ ; all elements on orbit  $G_{\mathbf{x}}g$  have an identical probability.
Property 2 characterizes probability distributions over  $G$  that can be expressed with  $p_{\omega}(g|\mathbf{x})$ , which we illustrate in the third panel of Figure 3. Intuitively,  $p_{\omega}(g|\mathbf{x})$  assigns constant probability densities over each of the orbit  $G_{\mathbf{x}}g$  that partitions  $G$  as shown in Property 1. We now prove Proposition 1 and Proposition 2 by showing that  $p_{\omega}(g|\mathbf{x})$  can become frame and canonicalizer as special cases:
Proof. A frame is defined as a set-valued function  $F: \mathcal{X} \to 2^G \setminus \emptyset$  that satisfies  $G$  equivariance  $F(g \cdot \mathbf{x}) = gF(\mathbf{x})$  [74]. For some frame  $F$ , frame averaging is defined as follows:
$$
\frac {1}{| F (\mathbf {x}) |} \sum_ {g \in F (\mathbf {x})} \left[ g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right) \right], \tag {40}
$$
which can be equivalently written as the below expectation:
$$
\mathbb {E} _ {g \sim \operatorname {U n i f} (F (\mathbf {x}))} \left[ g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right) \right]. \tag {41}
$$
From Theorem 3 of [74], we have that  $F(\mathbf{x})$  is a disjoint union of equal size orbits  $G_{\mathbf{x}}g$ . Therefore,  $\mathrm{Unif}(F(\mathbf{x}))$  is a uniform probability distribution over the union of the orbits. This can be expressed by a  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  by assigning identical probability over all orbits in the frame  $F$  and zero probability to all orbits not in the frame (illustrated in the fourth panel of Figure 3). Therefore, probabilistic symmetrization can become frame averaging.
Proof. A canonicalizer is defined as a (possibly stochastic) parameterized map  $C_{\omega} : \mathcal{X} \to G$  that satisfies relaxed  $G$  equivariance  $C_{\omega}(g \cdot \mathbf{x}) = gg' C_{\omega}(\mathbf{x})$  for some  $g' \in G_{\mathbf{x}}$  [41]. For some canonicalizer  $C_{\omega}$ , canonicalization is defined as follows:
$$
g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right), \quad g = C _ {\omega} (\mathbf {x}). \tag {42}
$$
From relaxed  $G$  equivariance, we have  $C_{\omega}(\mathbf{x}) = g' C_{\omega}(\mathbf{x})$  for some  $g' \in G_{\mathbf{x}}$ . A valid choice for the canonicalizer  $C_{\omega}$  is a stochastic map that samples from the uniform distribution over a frame  $C_{\omega}(\mathbf{x}) \sim \mathrm{Unif}(F_{\omega}(\mathbf{x}))$  where the frame is assumed to always provide a single orbit  $F_{\omega}(\mathbf{x}) = G_{\mathbf{x}}g$ . In this case, canonicalization is equivalent to a 1-sample estimation of the below expectation:
$$
\mathbb {E} _ {g \sim \operatorname {U n i f} \left(F _ {\omega} (\mathbf {x})\right)} \left[ g \cdot f _ {\theta} \left(g ^ {- 1} \cdot \mathbf {x}\right) \right]. \tag {43}
$$
Furthermore, uniform distribution over the single-orbit frame  $\mathrm{Unif}(F_{\omega}(\mathbf{x}))$  can be expressed by a  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  by assigning nonzero probability to the single orbit  $G_{\mathbf{x}}g$  and assigning zero probability to the rest (illustrated in the last panel of Figure 3). Therefore, probabilistic symmetrization can become canonicalization.
A.2 Extended Related Work (Continued from Section 2.4)
Our work draws inspiration from an extensive array of prior research, ranging from equivariant architectures and symmetrization to general-purpose deep learning with transformers. This section outlines a comprehensive review of these fields, spotlighting ideas specifically relevant to our work.
**Equivariant Architectures** Equivariant architectures, defined by the group equivariance of their building blocks, have been a prominent approach for equivariant deep learning [12, 10]. These architectures have been primarily developed for data types associated with permutation and Euclidean group symmetries, including images [18, 19], sets, graphs, and hypergraphs [5, 57, 8], and geometric graphs [23, 84, 91]. Additionally, they have been extended to more general data types under arbitrary finite group [78] and matrix group symmetries [30]. However, they face challenges such as limited expressive power [101, 56, 64, 105, 40] and architectural issues like over-smoothing [70, 14, 69] and over-squashing [92] in graph neural networks. Our work aims to develop an equivariant deep learning approach that relies less on equivariant architectures, to circumvent these limitations and enhance parameter sharing and transfer across varying group symmetries.
Symmetrization Our approach is an instance of symmetrization for equivariant deep learning which aims to achieve group equivariance using base models with unconstrained architectures. This is in general accomplished by averaging over specific group transformations of the input and output such that the averaged output exhibits equivariance. This allows us to leverage the expressive power of the base model e.g., achieve universal approximation using an MLP [35, 20] or a transformer [104], and potentially share or transfer parameters across different group symmetries. Existing literature has explored the choices of group transformations and base models for symmetrization. A straightforward approach is to average over the entire group [102], which is suitable for small, finite groups [4, 65, 42, 94] and requires sampling-based estimation for large groups such as permutations [67, 68, 88, 21]. Recent studies have attempted to identify smaller, input-dependent subsets of the group for averaging. Frame averaging [74] employs manually discovered set-values functions called frames, which still demand sampling-based estimation for certain worst-case inputs. Canonicalization [41] utilizes a single group transformation predicted by a neural network, but sacrifices strict equivariance. Our approach jointly achieves equivariance and end-to-end learning by utilizing parameterized, input-conditional equivariant distributions. Furthermore, our approach is one of the first demonstrations of symmetrization for the permutation group in real-world graph recognition task. Concerning the base model, previous work mostly examined small base models like an MLP or partial symmetrization of already equivariant models like GNNs. Few studies have explored symmetrizing pre-trained models for small finite groups [4, 3], and to our knowledge, we are the first to investigate symmetrization of a pre-trained standard transformer for permutation groups or any large group generally.
Transformer Architectures A significant motivation of our work is to combine the powerful scaling and transfer capabilities of the standard transformer architecture [96] with equivariant deep learning. The transformer architecture has driven major breakthroughs in language and vision domains [96, 24, 13, 75], and proven its ability to learn diverse modalities [39, 38] or transfer knowledge across them [85, 54, 82, 25, 71, 52]. Although transformer-style architectures have been developed for symmetric data modalities like sets [51], graphs [103, 45, 43, 50, 66, 63], hypergraphs [17, 44], and geometric graphs [31, 55], they often require specific architectural modifications to achieve equivariance to the given symmetry group, compromising full compatibility with transformer architectures used in language and vision domains. Apart from a few studies on linguistic graph encoding with language models [79], we believe we are the first to propose a general framework that facilitates full compatibility of the standard transformer architecture for learning symmetric data. For example, we have shown that a pre-trained vision transformer could be repurposed to encode graphs.
Learning Distribution of Data Augmentations Since our approach parameterizes a distribution  $p_{\omega}(g|\mathbf{x})$  on a group for symmetrization of form  $\phi_{\theta,\omega}(\mathbf{x}) = \mathbb{E}_g[g \cdot f_{\theta}(g^{-1} \cdot \mathbf{x})]$  and learns it from data, one may find similarity to Augerino [7] and related approaches [77, 81, 95, 93, 80, 37] that learn distributions over data augmentations (e.g.,  $p_{\omega}(g)$ ) for a similar symmetrization. The key difference is that, while these approaches aim to discover underlying (approximate) symmetry constraint from data and searches over a space of different group symmetries, our objective aims to obtain an exact  $G$  equivariant symmetrization  $\phi_{\theta,\omega}(\mathbf{x})$  given the known symmetry group  $G$  of data (e.g.,  $G = S_n$  for graphs). Because of this, the symmetrizing distribution has to be designed differently. In our case, we parameterize the distribution  $p_{\omega}(g|\mathbf{x})$  itself to be equivariant to a specific given group  $G$ , while for augmentation learning approaches, the distribution  $p_{\omega}(g)$  is parameterized for a different purpose of covering a range of different group symmetry constraints and their approximations (e.g., a set of 2D affine transformations [7]). This leads to advantages of our approach if the symmetry group  $G$  is known, as (1) our approach can learn non-trivial and useful distribution  $p_{\omega}(g|\mathbf{x})$  per input data  $\mathbf{x}$  while keeping the symmetrized function  $\phi_{\theta,\omega}(\mathbf{x})$  exactly  $G$  equivariant, while augmentation
Table 5: Overview of the datasets.
Dataset	Symmetry	Domain	Task	Feat. (dim)
GRAPH8c EXP EXP-classify	SnInvariant	Graph Isomorphism	Graph Separation Graph Classification	Adj. (1)
n-body	Sn × E(3) Equivariant	Physics	Position Regression	Pos. (3) + Vel. (3) + Charge (1)
PATTERN	Sn Equivariant	Mathematical Modeling	Node Classification	Rand. Node Attr. (3) + Adj. (1)
Peptides-func Peptides-struct	Sn Invariant	Chemistry	Graph Classification Graph Regression	Atom (9) + Bond (3) + Adj. (1)
PCQM-Contact	Sn Equivariant	Quantum Chemistry	Link Prediction	Atom (9) + Bond (3) + Adj. (1)
Table 6: Statistics of the datasets.
Dataset	Size	Max # Nodes	Average # Nodes	Average # Edges
GRAPH8c	11,117	8	8	28.82
EXP	1,200	64	44.44	110.21
EXP-classify
n-body	7,000	5	5	Fully Connected
PATTERN	14,000	188	117.47	4749.15
Peptides-func	15,535	444	150.94	307.30
Peptides-struct
PCQM-Contact	529,434	53	30.14	61.09
learning does not guarantee equivariance for a given group in general and often has to reduce to trivial group averaging  $p_{\omega} = \mathrm{Unif}(G)$  to be exactly  $G$  equivariant, and (2) while augmentation learning has to employ regularization [7] or model selection [37] to prevent collapse to trivial symmetry that is the least constrained and would fit the training data most easily [37], our approach fixes and enforces equivariance for the given symmetry group  $G$  by construction, which allows us to use regular maximum likelihood objective for training without the need to address symmetry collapse.
A.3 Experimental Details (Section 3)
We provide details of the datasets and models used in our experiments in Section 3. The details of the datasets from the original papers [2, 1, 27, 28, 31, 84] can be found in Table 5 and Table 6.
A.3.1 Implementation Details of  $p_{\omega}$  for Symmetric Group  $S_{n}$  (Section 3.1, 3.3, 3.4)
In all experiments regarding the symmetric group  $\mathbf{S}_n$ , we implement the  $\mathbf{S}_n$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$ , i.e.,  $q_{\omega}:(\mathbf{x},\epsilon)\mapsto \mathbf{P}_{g}$  as a 3-layer GIN with 64 hidden dimensions [101] that has around 25k parameters. Specifically, given a graph  $\mathbf{x}$  with node features  $\mathbf{X}\in \mathbb{R}^{n\times d_{\mathrm{in}}}$  and adjacency matrix  $\mathbf{A}\in \mathbb{R}^{n\times n}$ , we first augment a virtual node [32] which is connected to all nodes to facilitate global interaction while retaining  $\mathbf{S}_n$  equivariance, as follows: