SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\text{tokenize} : \mathcal{X} \to \mathbb{R}^{m \times d}$  parses input array to a sequence of  $m$  tokens, transformer:  $\mathbb{R}^{m \times d} \to \mathbb{R}^{m \times d}$  is a standard transformer encoder on tokens used in language and vision [24, 26], and  $\text{detokenize} : \mathbb{R}^{m \times d} \to \mathcal{Y}$  decodes encoded tokens to output array. For the tokenizer and  $\text{detokenizer}$ , we can use linear projections on flattened chunks of the array, which directly extends flattened patch projections in vision transformers to higher dimensions [26, 85, 86]. This enables mapping between different dimensional inputs and outputs, e.g., for graph node classification with  $\mathbf{x} \in \mathbb{R}^{n \times n \times c}$  and  $\mathbf{y} \in \mathbb{R}^{n \times c'}$ , it is possible to use 2D patch projection for the input and 1D projection for the output.
Above choice of  $f_{\theta}$  offers important advantages including universal approximation [104] and ability to share and transfer the learned knowledge in  $\theta$  over different domains of different group symmetries. Remarkably, this allows us to directly leverage large-scale pre-trained parameters from data-abundant domains for learning on symmetric domains. In our experiments, we only replace the tokenizer and detokenizer of a vision transformer pre-trained on ImageNet-21k [99, 26] and fine-tune it to perform diverse  $S_{n}$  equivariant tasks such as graph classification, node classification, and link prediction.
2.4 Relation to Other Symmetrization Approaches
We discuss relation of our symmetrization method to prior ones, specifically group averaging [102, 4], frame averaging [74], and canonicalization [41]. An extended discussion on broader related work can be found in Appendix A.2. Since these symmetrization methods share common formalization  $\mathbf{y} = \mathbb{E}_g[g\cdot f_\theta (g^{-1}\cdot \mathbf{x})]$ , one can expect a close theoretical relationship between them. We observe that probabilistic symmetrization is quite general; based on particular choices of the  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$ , it can become most of the related symmetrization methods as special cases. This can be easily seen for group averaging [102], as the distribution  $p_{\omega}$  can reduce to the uniform distribution  $\mathrm{Unif}(G)$  over the group. Frame averaging [74] also takes an average, but over a subset of group given by a frame  $F:\mathcal{X}\to 2^{G}\setminus \emptyset$ ; importantly, it is required that the frame itself is  $G$  equivariant  $F(\rho (g)\mathbf{x}) = gF(\mathbf{x})$ . We can make the following connection between our method and frame averaging, by adopting the concept of stabilizer subgroup  $G_{\mathbf{x}} = \{g\in G:\rho (g)\mathbf{x} = \mathbf{x}\}$ :
Proposition 1. Probabilistic symmetrization with  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  can become frame averaging [74] by assigning uniform density to a set of orbits  $G_{\mathbf{x}}g$  for some group elements  $g$ .
Proof. The proof can be found in Appendix A.1.5.
Canonicalization [41] uses a single group element for symmetrization, produced by a trainable canonicalizer  $C_{\omega} : \mathcal{X} \to G$ . Here, it is required that the canonicalizer itself satisfies relaxed  $G$  equivariance  $C(\rho(g)\mathbf{x}) = gg'C(\mathbf{x})$  up to arbitrary action from the stabilizer  $g' \in G_{\mathbf{x}}$ . We now show:
Proposition 2. Probabilistic symmetrization with  $G$  equivariant distribution  $p_{\omega}(g|\mathbf{x})$  can become canonicalization [41] by assigning uniform density to a single orbit  $G_{\mathbf{x}}g$  of some group element  $g$ .
Assuming that stabilizer  $G_{\mathbf{x}}$  is trivial, this can be implemented with our method by removing random noise  $\epsilon$ , which reduces  $p_{\omega}$  to deterministic map  $\rho(g) = q_{\omega}(\mathbf{x})$ . We use this approach to implement canonicalizer for the  $S_{n}$  group, while [41] only provides canonicalizers for Euclidean groups.
We empirically demonstrate and analyze probabilistic symmetrization on a range of symmetry groups  $\mathbf{S}_n$ ,  $\mathrm{E}(3)$  ( $\mathrm{O}(3)$ ), and the product  $\mathbf{S}_n \times \mathrm{E}(3)$  ( $\mathbf{S}_n \times \mathrm{O}(3)$ ), on a variety of invariant and equivariant tasks, with general-purpose base functions  $f_{\theta}$  chosen as MLP and transformer optionally with pretraining from vision domain. Details of the experiments are in Appendix A.3, and supplementary experiments on other base functions and comparisons to other symmetrization approaches are in Appendix A.4.
3.1 Graph Isomorphism Learning with MLP
Building expressive neural networks for graphs  $(\mathbf{S}_n)$  has been considered important and challenging, as simple and efficient GNNs are often limited in expressive power to certain Weisfeiler-Lehman isomorphism tests like 1-WL [101, 56]. Since an MLP equipped with probabilistic symmetrization is in theory universal and  $\mathbf{S}_n$  equivariant, it has potential for graph learning that require high expressive power. To explicitly test this, we adopt the experimental setup of [74] and use two datasets on graph separation task  $(\mathbf{S}_n$  invariant). GRAPH8c [2] consists of all non-isomorphic connected graphs with 8 nodes, and EXP [1] consists of 3-WL distinguishable graphs that are not 2-WL distinguishable. We compare our method to standard GNNs as well as an MLP symmetrized with group averaging [102], frame averaging [74], and canonicalization [41]. Our method uses the same MLP architecture to symmetrization baselines, and its  $\mathbf{S}_n$  equivariant distribution  $p_{\omega}$  for symmetrization is implemented
Table 1: Results for  $S_{n}$  invariant graph separation. We use two tasks, one for counting pairs of graphs not separated by a model at random initialization (GRAPH8c and EXP), and one for learning to classify EXP to two classes (EXP-classify). For EXP-classify, we report the test accuracy at best validation accuracy. The columns arch. and sym. denote architectural and symmetrized equivariance, respectively. The results for baselines are from [74] except for MLP-Canonical. which is tested by us.
method	arch.	sym.	GRAPH8c ↓	EXP ↓	EXP-classify ↑
GCN [48]	Sn	-	4755	600	50%
GAT [97]	Sn	-	1828	600	50%
GIN [101]	Sn	-	386	600	50%
ChebNet [22]	Sn	-	44	71	82%
PPGN [56]	Sn	-	0	0	100%
GNNML3 [2]	Sn	-	0	0	100%
MLP-GA [102]	-	Sn	0	0	50%
MLP-FA [74]	-	Sn	0	0	100%
MLP-Canonical.	-	Sn	0	0	50%
MLP-PS (Ours), fixed ε	-	Sn	0	0	79.5%
MLP-PS (Ours)	-	Sn	0	0	100%
Epoch 1 Epoch 1000
Figure 2: Learned  $p_{\omega}(g|\mathbf{x})$  over time. The entropy of aggregated permutation matrices  $\bar{\mathbf{P}} = \sum \mathbf{P}_g / N$  from  $\mathbf{P}_g \sim p_{\omega}(g|\mathbf{x})$  for each input  $\mathbf{x}$  drops in early training, indicating that the distribution learns to produce lower-variance permutations as in below visualizations.
using a 3-layer GIN [101] which is 1-WL expressive. We use 10 samples for symmetrization during both training and testing. Further details can be found in Appendix A.3.3, and supplementary results on symmetrization of a different base model can be found in Appendix A.4.1.
The results are in Table 1. At random initialization, all symmetrization methods can provide perfect separation of all graphs, similar to PPGN [56] and GNNML3 [2] that are equivariant neural networks carefully designed to be 3-WL expressive. However, when trained with gradient descent to solve classification problem, naive symmetrization with group averaging fails, presumably because the MLP fails to adjust to equally possible 64! permutations of 64 nodes in maximum. On the other hand, our method is able to learn the task, achieving the same accuracy to frame averaging that utilizes costly eigendecomposition of graph Laplacian [74]. What makes our method work while group averaging fails? We conjecture this is since the distribution  $p_{\omega}(g|\mathbf{x})$  can learn to provide more consistent permutations during early training, as we illustrate in Figure 2. In the figure, we measured the consistency of samples from  $p_{\omega}(g|\mathbf{x})$  over training progress by sampling  $N = 50$  permutation matrices  $\mathbf{P}_g \sim p_{\omega}(g|\mathbf{x})$  and measuring the row-wise entropy of their average  $\bar{\mathbf{P}} = \sum \mathbf{P}_g / N$  for each input  $\mathbf{x}$ . The more consistent the sampled permutations, the sharper their average, and lesser the entropy. As training progresses,  $p_{\omega}$  learns to produce more consistent samples, which coincides with the initial increase in task performance. Given that, a natural question would be: if we enforce the samples  $g \sim p_{\omega}(g|\mathbf{x})$  to be consistent from the first place, would it work? To answer this, we also tested a non-probabilistic version of our model that uses a single permutation per input  $\rho(g) = q_{\omega}(\mathbf{x})$ , which is a canonicalizer under relaxed equivariance [41] as described in Section 2.4. As in Table 1, canonicalization fails, suggesting that probabilistic nature of  $p_{\omega}(g|\mathbf{x})$  can be beneficial for learning. For another deterministic version of our model made by fixing the noise  $\epsilon$  (Eq. (6)) at initialization, the performance drops to 79.5%, further implying that stochasticity of  $p_{\omega}(g|\mathbf{x})$  has a role.
Table 2: Results for  $S_{n} \times E(3)$  equivariant  $n$ -body problem. The columns arch. and sym. denote architectural and symmetrized equivariance, respectively. We report test MSE at best validation MSE, along with the standard deviation for GA and Ours where predictions are stochastic. The results for baselines are from [41] except symmetrized transformers which are tested by us.
method	arch.	sym.	Position MSE ↓
SE(3) Transformer [31]	Sn × SE(3)	-	0.0244
TFN [91]	Sn × SE(3)	-	0.0155
Radial Field [49]	Sn × E(3)	-	0.0104
EGNN [84]	Sn × E(3)	-	0.0071
GNN-FA [74]	Sn	E(3)	0.0057
GNN-Canonical. [41]	Sn	E(3)	0.0043
Transformer-Canonical.	-	Sn × E(3)	0.00508
Transformer-GA	-	Sn × E(3)	0.00414 ± 0.00001
Transformer-PS (Ours)	-	Sn × E(3)	0.00401 ± 0.00001
3.2 Particle Dynamics Learning with Transformer
Learning sets or graphs attributed with position and velocity in 3D  $(\mathbf{S}_n\times \mathrm{E}(3))$  is practically significant as they universally appear in physics, chemistry, and biology applications. While prior symmetrization methods employ an already  $\mathbf{S}_n$  equivariant base function and partially symmetrize the  $\mathrm{E}(3)$  part, we attempt to symmetrize the entire product group  $\mathbf{S}_n\times \mathbf{E}(3)$  and choose the base model  $f_{\theta}$  as a sequence transformer to leverage its expressive power. For empirical demonstration, we adopt the experimental setup of [41] and use the  $n$ -body dataset [84, 31] where the task is predicting the position of  $n = 5$  charged particles after certain time given their initial position and velocity in  $\mathbb{R}^3$  ( $\mathbf{S}_n\times \mathbf{E}(3)$  equivariant). We compare our method to  $\mathbf{S}_n\times \mathbf{E}(3)$  equivariant neural networks and partial symmetrization methods applying  $\mathrm{E}(3)$  symmetrization to GNNs. We also test prior symmetrization methods on the full group  $\mathbf{S}_n\times \mathbf{E}(3)$  along our method, but could not test for frame averaging since equivariant frames for the full group  $\mathbf{S}_n\times \mathbf{E}(3)$  was not available in current literature. Our method is implemented using a transformer with sequence positional encodings with around  $2.3\times$  parameters of the baselines, and the  $\mathbf{S}_n\times \mathbf{E}(3)$  equivariant distribution  $p_{\omega}$  for symmetrization is implemented using a 2-layer Vector Neurons [23] that has around  $0.03\times$  of parameters to the transformer. We use 20 samples for symmetrization during training, and use  $10\times$  sample size for testing since the task is regression where appropriate variance reduction is necessary to guarantee a reliable performance. Further details can be found in Appendix A.3.4, and supplementary results on  $\mathrm{E}(3)$  partial symmetrization of GNN base model can be found in Appendix A.4.2.
The results are in Table 2. We observe simple group averaging exhibits a surprisingly strong performance, as it achieves 0.00414 MSE and already outperforms previous state of the art 0.0043 MSE. This is because the permutation component of the symmetry is fairly small, with  $n = 5$  particles interacting with each other, such that combining it with an expressive base model  $f_{\theta}$  (a sequence transformer) can adjust to  $5! = 120$  equally possible permutations and their rotations in 3D. Nevertheless, our method outperforms group averaging and achieves a new state of the art 0.00401 MSE, presumably as the parameterized distribution  $p_{\omega}$  learns to further maximize task performance. On the other hand, the canonicalization approach, implemented by eliminating noise variable  $\epsilon$  from our method (Section 2.4), performs relatively poorly. We empirically observe that  $f_{\theta}$  memorizes the per-input canonical orientations provided by  $\rho(g) = q_{\omega}(\mathbf{x})$  that do not generalize to test inputs. This again shows that probabilistic nature of  $p_{\omega}(g|\mathbf{x})$  can be beneficial for performance.
3.3 Graph Pattern Recognition with Vision Transformer
One important goal of our approach, and symmetrization in general, is to decouple the symmetry of problem from the base function  $f_{\theta}$ , such that we can leverage knowledge learned from other symmetries by transferring the parameters  $\theta$ . We demonstrate an extreme case by transferring the parameters of a vision transformer [26] trained on large-scale image classification (translation invariant) to solve node classification on graphs ( $S_{n}$  equivariant) for the first time in literature. For this, we use the PATTERN dataset [27] that contains 14,000 purely topological random SBM graphs with 44-188 nodes, whose task is finding certain subgraph pattern by binary node classification.
Table 3: Results for  $S_{n}$  equivariant node classification on PATTERN. We report test accuracy at the best validation accuracy, along with the standard deviation for GA and Ours where predictions are stochastic. The results for GNN baselines are from [27].
method	pretrain.	Accuracy ↑
GCN [48], 16 layers	-	85.614
GAT [97], 16 layers	-	78.271
GatedGCN [11], 16 layers	-	85.568
GIN [101], 16 layers	-	85.387
RingGNN [16], 2 layers	-	86.245
RingGNN [16], 8 layers	-	diverged
PPGN [56], 3 layers	-	85.661
PPGN [56], 8 layers	-	diverged
ViT-GA, 1-sample	-	76.776 ± 0.137
ViT-GA, 10-sample	-	83.119 ± 0.048
ViT-GA, 1-sample	ImageNet-21k	81.407 ± 0.101
ViT-GA, 10-sample	ImageNet-21k	84.351 ± 0.053
ViT-FA	-	70.063
ViT-FA	ImageNet-21k	79.637
ViT-Canonical.	-	85.460
ViT-Canonical.	ImageNet-21k	86.166
ViT-PS (Ours), 1-sample	-	85.542 ± 0.012
ViT-PS (Ours), 10-sample	-	85.635 ± 0.021
ViT-PS (Ours), 1-sample	ImageNet-21k	86.226 ± 0.028
ViT-PS (Ours), 10-sample	ImageNet-21k	86.285 ± 0.015
Based on pre-trained ViT [26, 89], we construct the base model  $f_{\theta}$  by modifying only the input and output layers to take the flattened patches of 2D zero-padded adjacency matrices of size  $188 \times 188$  and produce output as 1D per-node classification logits of length 188 with 2 channels. In addition to standard GNNs<sup>4</sup>, we compare group averaging, frame averaging, and canonicalization to our method, and also test whether pre-trained representations from ImageNet-21k is beneficial for the task. For the  $S_{n}$  equivariant distribution  $p_{\omega}$  in our method and canonicalization, we use a 3-layer GIN [101] with only  $0.02\%$  of the base model parameters. Further details can be found in Appendix A.3.5.
The results are in Table 3. First, we observe that transferring pre-trained ViT parameters consistently improves node classification for all symmetrization methods. It indicates that some traits of the pre-trained visual representation can benefit learning graph tasks which vastly differ in both the underlying symmetry (translation invariance  $\rightarrow$ $S_{n}$  equivariance) and the data generating process (natural images  $\rightarrow$  random process of SBM). In particular, it is somewhat surprising that vision pretraining allows group averaging to achieve  $84.351\%$  accuracy, on par with GNN baselines, considering that memorizing all 188! equally possible permutations in this dataset is impossible. We conjecture that group averaged ViT can in some way learn meaningful graph representation internally to solve the task, and vision pretraining helps in acquiring the representation by providing a good initialization point or transferable computation motifs.
On the other hand, frame averaging shows a low performance,  $79.637\%$  accuracy with vision pretraining, which is also surprising considering that frames vastly reduce the sample space of symmetrization in general; in fact, the size of frame of each graph in PATTERN is exactly 1. We empirically observe that, unlike group averaging, ViT with frame averaging memorizes frames of training graphs rather than learning generalizable graph representations. In contrast, canonicalization that also uses a single sample per graph successfully learns the task with  $86.166\%$  accuracy. We conjecture that the learnable orderings provided by an equivariant neural network  $\rho(g) = q_{\omega}(\mathbf{x})$  is more flexible and generalizable to unseen graphs compared to frames computed from fixed graph Laplacian eigenvectors. Lastly, our method achieves a better performance compared to other symmetrization methods, and the performance consistently improves with vision pretraining and more samples for testing. As a result, our model based on pre-trained ViT and 10 samples for testing achieves  $86.285\%$  test accuracy, surpassing all baselines.
<sup>4</sup>We note that careful engineering such as graph positional encoding improves performance of GNNs, but we have chosen simple and representative GNNs in the benchmark [27] to provide a controlled comparison.
Table 4: Results for real-world graph tasks. We report test performance at best validation performance.
method	Peptides-func	Peptides-struct	PCQM-Contact
AP ↑	MAE ↓	Hits@1 ↑	Hits@3 ↑	Hits@10 ↑	MRR ↑
GCN [48]	0.5930	0.3496	0.1321	0.3791	0.8256	0.3234
GCNII [15]	0.5543	0.3471	0.1325	0.3607	0.8116	0.3161
GINE [36]	0.5498	0.3547	0.1337	0.3642	0.8147	0.3180
GatedGCN [11]	0.5864	0.3420	0.1279	0.3783	0.8433	0.3218
GatedGCN+RWSE [11]	0.6069	0.3357	0.1288	0.3808	0.8517	0.3242
Transformer+LapPE [28]	0.6326	0.2529	0.1221	0.3679	0.8517	0.3174
SAN+LapPE [50]	0.6384	0.2683	0.1355	0.4004	0.8478	0.3350
SAN+RWSE [50]	0.6439	0.2545	0.1312	0.4030	0.8550	0.3341
GraphGPS [76]	0.6535	0.2500	-	-	-	0.3337
Exphormer [87]	0.6527	0.2481	-	-	-	0.3637
ViT-PS (Ours)	0.8311	0.2662	0.3268	0.6693	0.9524	0.5329
3.4 Real-World Graph Learning with Vision Transformer
Having observed that pre-trained ViT can learn graph tasks well when symmetrized with our method, we now provide a preliminary test of it in real-world graph learning. We use three real-world graph datasets from [28] that involve chemical and biological graphs. PCQM-Contact dataset contains 529,434 molecular graphs with 53 nodes in maximum, and the task is contact map prediction framed as link prediction  $(\mathbf{S}_n$  equivariant), on whether two atoms would be proximal when the molecule is in 3D space. Peptides-func and Peptides-struct are based on the same set of 15,535 protein graphs with 444 nodes in maximum and the tasks are property prediction  $(\mathbf{S}_n$  invariant), requiring multi-label classification for Peptides-func and regression for Peptides-struct. The tasks require complex understanding of how the amino acids of the proteins would interact in 3D space. We implement our method using a ViT-Base pre-trained on ImageNet-21k as the base model  $f_{\theta}$  and a 3-layer GIN as the equivariant distribution  $p_{\omega}(g|\mathbf{x})$ , following our model in Section 3.3. We use 10 samples for both training and testing. Further details can be found in Appendix A.3.6.
The results are in Table 4. In Peptides-func and PCQM-Contact, the pre-trained ViT symmetrized with our method achieves a significantly higher performance compared to both GNNs and graph transformers, improving previous best by a large margin $^5$  ( $0.6527 \to 0.8311$  for Peptides-func AP,  $0.1355 \to 0.3268$  for PCQM-Contact Hits@1). This demonstrates the scalability of our method as Peptides-func involves 444 maximum nodes, and also its generality as it performs well for both  $S_{n}$  invariant (Peptides-func) and equivariant (PCQM-Contact) tasks. We also note that, unlike some baselines, our method does not require costly Laplacian eigenvectors to compute positional encoding. On the Peptides-struct, our method achieves a slightly lower performance to SOTA graph transformers while still better than GNNs. We conjecture that regression is harder for the model to learn due to its stochasticity in predictions, and leave improving regression performance as future work.
We presented probabilistic symmetrization, a general framework that learns a distribution of group transformations conditioned on input data for symmetrization of an arbitrary function. By characterizing that the only condition for such distribution is equivariance to data symmetry, we instantiated models for a wide range of groups, including symmetric, orthogonal, Euclidean groups and their product combinations. Our experiments demonstrated that the proposed framework achieves consistent improvement over other symmetrization methods, and is competitive or outperforms equivariant networks on various datasets. We also showed that transferring pre-trained parameters across data in different symmetries can sometimes be surprisingly beneficial. Our approach has weaknesses such as sampling cost, which we further discuss in Appendix A.5; we plan to address these in future work.
Acknowledgements This work was supported in part by the National Research Foundation of Korea (NRF2021R1C1C1012540 and NRF2021R1A4A3032834) and IITP grant (2021-0-00537, 2019-0-00075, and 2021-0-02068) funded by the Korea government (MSIT).
<sup>5</sup>We note that the baseline architectures are constructed within 500k parameter budget as a convention [28], while we use an identical architecture to ViT-Base to leverage pre-trained representations.
A.1.1 Proof of Theorem 1 (Section 2.1)
Proof. We prove  $\phi_{\theta, \omega}(\rho_1(g')\mathbf{x}) = \rho_2(g')\phi_{\theta, \omega}(\mathbf{x})$  for all  $\mathbf{x} \in \mathcal{X}$  and  $g' \in G$ . From Eq. (4), we have:
$$
\phi_ {\theta , \omega} \left(\rho_ {1} \left(g ^ {\prime}\right) \mathbf {x}\right) = \mathbb {E} _ {p _ {\omega} \left(g \mid \rho_ {1} \left(g ^ {\prime}\right) \mathbf {x}\right)} \left[ \rho_ {2} (g) f _ {\theta} \left(\rho_ {1} (g) ^ {- 1} \rho_ {1} \left(g ^ {\prime}\right) \mathbf {x}\right) \right]. \tag {11}
$$
Let us introduce transformed random variable  $h = g'^{-1}g \in G$  such that  $g = g'h$ . Since the distribution  $p_{\omega}$  is  $G$  equivariant, we can see that  $p_{\omega}(g|\rho_1(g')\mathbf{x}) = p_{\omega}(g'^{-1}g|\rho_1(g'^{-1})\rho_1(g')\mathbf{x}) = p_{\omega}(g'^{-1}g|\mathbf{x}) = p_{\omega}(h|\mathbf{x})$ . Thus, we can rewrite the above expectation with respect to  $h$  as follows: