SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbf {X} ^ {\prime} = \left[ \mathbf {X}; \mathbf {v} \right], \quad \mathbf {A} ^ {\prime} = \left[ \begin{array}{c c} \mathbf {A} & \mathbf {1} \\ \mathbf {1} ^ {\top} & 0 \end{array} \right], \tag {44}
$$
where the feature of the virtual node  $\mathbf{v} \in \mathbb{R}^{d_{\mathrm{in}}}$  is a trainable parameter. Then, we prepared the input node features  $\mathbf{H} \in \mathbb{R}^{(n+1) \times d_{\mathrm{in}}}$  to the GIN as  $\mathbf{H} = \mathbf{X}' + \epsilon$  where the noise  $\epsilon \in \mathbb{R}^{(n+1) \times d_{\mathrm{in}}}$  is i.i.d. sampled from Unif[0,  $\eta$ ] with scale hyperparameter  $\eta$ . Then, we employ following 3-layer GIN with 64 hidden dimensions to obtain processed node features  $\mathbf{H}' \in \mathbb{R}^{(n+1) \times 1}$ :
$$
\mathbf {H} ^ {\prime} = \operatorname {G I N C o n v} _ {6 4, 6 4, 1} \circ \operatorname {G I N C o n v} _ {6 4, 6 4, 6 4} \circ \operatorname {G I N C o n v} _ {d _ {\text {i n}}, 6 4, 6 4} (\mathbf {H}), \tag {45}
$$
where each  $\mathrm{GINConv}_{d_1,d_2,d_3}$  computes below with a two-layer elementwise MLP:  $\mathbb{R}^{n\times d_1}\to \mathbb{R}^{n\times d_3}$  with hidden dimension  $d_{2}$ , ReLU activation, batch normalization, and trained scalar  $e$ :
$$
\mathbf {H} \mapsto \operatorname {M L P} \left(\left(\mathbf {A} ^ {\prime} + (1 + e) \mathbf {I}\right) \mathbf {H}\right). \tag {46}
$$
<sup>7</sup>We do not utilize edge attributes in equivariant distribution  $p_{\omega}$ , while we utilize them in base model  $f_{\theta}$ .
Then, from the processed node features  $\mathbf{H}' \in \mathbb{R}^{(n+1) \times 1}$ , we finally obtain the features  $\mathbf{Z} \in \mathbb{R}^n$  for postprocessing by discarding the feature of the virtual node. Then, postprocessing into a permutation matrix is done with  $\text{argsort} : \mathbf{Z} \mapsto \mathbf{P}_g \in \mathbb{R}^{n \times n}$  as in Eq. (8).
Training To backpropagate through  $\mathbf{P}_g$  for end-to-end training of  $p_{\omega}(g|\mathbf{x})$ , we use straight-through gradient estimator [6] with an approximate permutation matrix  $\hat{\mathbf{P}}_g \approx \mathbf{P}_g$ . For this, we first apply L2 normalization  $\mathbf{Z} \mapsto \bar{\mathbf{Z}}$  and use the below differentiable relaxation of the argsort operator [61, 33, 98]:
$$
\hat {\mathbf {P}} _ {g} = S \left(- | \bar {\mathbf {Z}} \mathbf {1} ^ {\top} - \mathbf {1} \operatorname {s o r t} (\bar {\mathbf {Z}}) ^ {\top} | / \tau\right), \tag {47}
$$
where  $S(\cdot/\tau)$  is Sinkhorn operator [61] with temperature hyperparameter  $\tau \in \mathbb{R}_+$  that performs elementwise exponential followed by iterative normalization of rows and columns. Following [61], we use 20 Sinkhorn iterations which worked robustly in all our experiments. For the correctness of straight-through gradients, it is desired that  $\hat{\mathbf{P}}_g$  closely approximates the real permutation matrix  $\mathbf{P}_g$  during training. For this, we choose the temperature  $\tau$  to be small, 0.01 in general, and following prior work [98], employ a regularizer on the mean of row- and column-wise entropy of  $\hat{\mathbf{P}}_g$  with a strength of 0.1 in all experiments. The  $S_n$  equivariance of the relaxed argsort  $\mathbf{Z} \mapsto \hat{\mathbf{P}}_g$  can be shown in a similar way to Proposition 3 from the fact that elementwise subtraction, absolute, scaling by  $-1/\tau$ , exponential, and iterative normalization of rows and columns all commute with  $\mathbf{P}_{g'} \in S_n$ .
A.3.2 Implementation Details of  $p_{\omega}$  for Product Group  $S_n \times \mathrm{E}(3)$  (Section 3.2)
In our  $n$ -body experiment on the product group  $\mathbf{S}_n\times \mathrm{E}(3)$ , we implement the  $\mathbf{S}_n\times \mathrm{O}(3)$  equivariant distribution  $p_{\omega}(g|\mathbf{x} - \bar{\mathbf{x}}\mathbf{1}^{\top})$ , i.e.,  $q_{\omega}:(\mathbf{x} - \bar{\mathbf{x}}\mathbf{1}^{\top},\boldsymbol {\epsilon})\mapsto (\mathbf{P}_{g},\mathbf{Q}_{g})$  based on a 2-layer Vector Neurons version of DGCNN with 96 hidden dimensions [23] that has around 7k parameters. Due to the architecture's complexity, we focus on describing input and output of the network and postprocessing, and guide the readers to the original paper [23] for further architectural details. In a high-level, the Vector Neurons receives position  $\mathbf{P}\in \mathbb{R}^{n\times 3}$  and velocity  $\mathbf{V}\in \mathbb{R}^{n\times 3}$  of the zero-centered input  $\mathbf{x} - \bar{\mathbf{x}}\mathbf{1}^{\top}$  with noises  $\epsilon_{1},\epsilon_{2}\in \mathbb{R}^{n\times 3}$  i.i.d. sampled from normal  $\mathcal{N}(0,\eta^2)$  with scale hyperparameter  $\eta$ , and produces features  $\mathbf{H}_{\mathrm{S}_n}\in \mathbb{R}^{n\times 3\times d_1}$  and  $\mathbf{H}_{\mathrm{O(3)}}\in \mathbb{R}^{n\times 3\times d_2}$  with  $d_{1} = 1$  and  $d_{2} = 3$  as follows:
$$
\mathbf {H} _ {\mathrm {S} _ {n}}, \mathbf {H} _ {\mathrm {O} (3)} = \operatorname {V N - D G C N N} (\mathbf {P} + \boldsymbol {\epsilon} _ {1}, \mathbf {V} + \boldsymbol {\epsilon} _ {2}). \tag {48}
$$
Then, we apply O(3) invariant pooling on  $\mathbf{H}_{\mathrm{S}_n}$  and  $\mathbf{S}_n$  invariant pooling on  $\mathbf{H}_{\mathrm{O}(3)}$ , both supported as a part of [23], to obtain features for postprocessing  $\mathbf{Z}_{\mathrm{S}_n} \in \mathbb{R}^{n \times 1}$  and  $\mathbf{Z}_{\mathrm{O}(3)} \in \mathbb{R}^{3 \times 3}$ , respectively:
$$
\mathbf {Z} _ {\mathrm {S} _ {n}} = \operatorname {P o o l} _ {\mathrm {O} (3)} \left(\mathbf {H} _ {\mathrm {S} _ {n}}\right), \quad \mathbf {Z} _ {\mathrm {O} (3)} = \operatorname {P o o l} _ {\mathrm {S} _ {n}} \left(\mathbf {H} _ {\mathrm {O} (3)}\right). \tag {49}
$$
Then, postprocessing with argsort:  $\mathbf{Z}_{\mathrm{S}_n} \mapsto \mathbf{P}_g \in \mathbb{R}^{n \times n}$  and Gram-Schmidt orthogonalization  $\mathbf{Z}_{\mathrm{O}(3)} \mapsto \mathbf{Q}_g \in \mathbb{R}^{3 \times 3}$  is performed identically as described in the main text (Section 2.2). For the straight-through gradient estimation of the argsort operator, we use relaxed argsort described in Appendix A.3.1, with the only difference of using the temperature  $\tau = 0.1$ .
A.3.3 Graph Isomorphism Learning with MLP (Section 3.1)
Base Model  $f_{\theta}$  For EXP and EXP-classify, the model is given adjacency matrix  $\mathbf{A} \in \mathbb{R}^{64 \times 64}$  and binary node features  $\mathbf{X} \in \mathbb{R}^{64}$  which are zero-padded to maximal number of nodes 64. For GRAPH8c, the input graphs are all of size 8 without node features, and the model is given adjacency matrix  $\mathbf{A} \in \mathbb{R}^{8 \times 8}$ . For EXP-classify, the prediction target is a scalar binary classification logit.
For the base model for EXP-classify, we use a 5-layer MLP  $f_{\theta} : \mathbb{R}^{64 \times 64 + 64} \to \mathbb{R}$  on flattened and concatenated adjacency matrix and node features, with an identical architecture to other symmetrization baselines (MLP-GA and MLP-FA [74]) as in below:
$$
f _ {\theta} = \mathrm {F C} _ {1, 1 0} \circ \mathrm {F C} _ {1 0, 2 0 4 8} \circ \mathrm {F C} _ {2 0 4 8, 4 0 9 6} \circ \mathrm {F C} _ {4 0 9 6, 2 0 4 8} \circ \mathrm {F C} _ {2 0 4 8, 4 1 6 0}, \tag {50}
$$
where  $\mathrm{FC}_{d_2,d_1}:\mathbb{R}^{d_1}\to \mathbb{R}^{d_2}$  denotes a fully-connected layer and ReLU activation is omitted. For EXP, we drop the last layer to obtain 10-dimensional output. For GRAPH8c, we use the following architecture  $f_{\theta}:\mathbb{R}^{8\times 8}\rightarrow \mathbb{R}^{10}$  that takes flattened adjacency to produce 10-dimensional output [74]:
$$
f _ {\theta} = \mathrm {F C} _ {1 0, 6 4} \circ \mathrm {F C} _ {6 4, 1 2 8} \circ \mathrm {F C} _ {1 2 8, 6 4}. \tag {51}
$$
In PyTorch [72], one can simply replace  $\mathbf{P}_g$  with  $(\mathbf{P}_g - \hat{\mathbf{P}}_g)$ . detach() +  $\hat{\mathbf{P}}_g$  during forward passes.
Training For EXP-classify, we train our models with binary cross-entropy loss using Adam optimizer [46] with batch size 100 and learning rate 1e-3 for 2,000 epochs, which takes around 30 minutes on a single RTX 3090 GPU with 24GB using PyTorch [72]. We additionally apply 200 epochs of linear learning rate warm-up and gradient norm clipping at 0.1, which we found helpful for stabilizing the training. For the equivariant distribution  $p_{\omega}$ , we use noise scale  $\eta = 1$ . Since EXP and GRAPH8c concern randomly initialized models, we do not train the models for these tasks.
A.3.4 Particle Dynamics Learning with Transformer (Section 3.2)
Base Model  $f_{\theta}$  The model is given zero-centered position  $\mathbf{P} \in \mathbb{R}^{5 \times 3}$  and velocity  $\mathbf{V} \in \mathbb{R}^{5 \times 3}$  of 5 particles at a time point with pairwise charge difference  $\mathbf{C} \in \mathbb{R}^{5 \times 5}$  and squared distance  $\mathbf{D} \in \mathbb{R}^{5 \times 5}$ . We set the prediction target as difference of position  $\Delta \mathbf{P} \in \mathbb{R}^{5 \times 3}$  after a certain time.
For the base model, we use a 8-layer transformer encoder  $f_{\theta} : \mathbb{R}^{25 \times 8} \rightarrow \mathbb{R}^{25 \times 3}$  that operates on sequences of length 25 with dimension 8. At each prediction, we first organize the input into a single tensor  $\in \mathbb{R}^{5 \times 5 \times 8}$  by placing  $\mathbf{P}$  and  $\mathbf{V}$  on the diagonals of  $\mathbf{C}$  and  $\mathbf{D}$ , and then turn the tensor into a sequence of 25 tokens  $\in \mathbb{R}^{25 \times 8}$  by flattening the first two axes. Analogously, we organize the output of the model into a tensor  $\in \mathbb{R}^{5 \times 5 \times 3}$  and take the diagonal entries as the predictions. For the transformer, we use the standard implementation provided in PyTorch [72, 96], with 64 hidden dimensions, 4 attention heads, GELU activation [34] in feedforward networks, PreLN [100], learnable 1D positional encoding, and an MLP prediction head with 1 hidden layer. The model has around 208k trainable parameters, around  $2.3 \times$  compared to the GNN backbones of E(3) symmetrization baselines in the benchmark (GNN-FA and GNN-Canonical.) with 92k parameters.
Training We train our models with MSE loss using Adam optimizer [46] with batch size 100 and learning rate 1e-3 for 10,000 epochs, which takes around 8.5 hours on a single RTX 3090 GPU with 24GB using PyTorch [72]. We use weight decay with strength 1e-12 and dropout on the distribution  $p_{\omega}$  with probability 0.08. For the equivariant distribution  $p_{\omega}$ , we use noise scale  $\eta = 1$ .
A.3.5 Graph Pattern Recognition with Vision Transformer (Section 3.3)
Base Model  $f_{\theta}$  The model is given adjacency matrix  $\mathbf{A} \in \mathbb{R}^{188 \times 188}$  and node features  $\mathbf{X} \in \mathbb{R}^{188 \times 3}$  zero-padded to maximal 188 nodes. The prediction target is node classification logs  $\mathbf{Y} \in \mathbb{R}^{188 \times 2}$ .
For the base model, we use a transformer with an identical architecture to ViT-Base [26] that operates on  $224 \times 224$  images with  $16 \times 16$  patch, using configuration from HuggingFace [99] model hub. We first remove the input patch projection and output head layers, which gives us a backbone transformer:  $\mathbb{R}^{(14 \times 14) \times 768} \to \mathbb{R}^{(14 \times 14) \times 768}$  on sequences of  $(224 / 16) \times (224 / 16) = 14 \times 14$  tokens. Then, we use the following as the base model  $f_{\theta}: (\mathbf{A}, \mathbf{X}) \mapsto \mathbf{Y}$ :
$$
f _ {\theta} (\mathbf {A}, \mathbf {X}) = \text {d e t o k e n i z e (t r a n s f o r m e r (t o k e n i z e} (\mathbf {A}, \mathbf {X})))}, \tag {52}
$$
where, for tokenize:  $\mathbb{R}^{188\times 188\times 1}\times \mathbb{R}^{188\times 3}\to \mathbb{R}^{(14\times 14)\times 768}$  we organize the input into a single tensor  $\in \mathbb{R}^{188\times 188\times 4}$  by placing  $\mathbf{X}$  on the diagonals of  $\mathbf{A}$  and apply 2D convolution with kernel size and stride 14, and for detokenize:  $\mathbb{R}^{(14\times 14)\times 768}\to \mathbb{R}^{188\times 2}$  we apply transposed 2D convolution with kernel size and stride 14 to obtain a tensor  $\in \mathbb{R}^{188\times 188\times 2}$  and take its diagonal entries as output.
Training We train our models with binary cross-entropy loss weighted inversely by class size [27] using AdamW [53] optimizer with batch size 128, learning rate 1e-5, and weight decay 0.01. We train the models for 25k steps under learning rate warm-up for 5k steps then linear decay to 0 with early stopping based on validation loss, which usually takes less than 5 hours on 8 RTX 3090 GPUs with 24GB using PyTorch Lightning [29]. For the equivariant distribution  $p_{\omega}$  we use noise scale  $\eta = 1$  and dropout with probability 0.1. For probabilistic symmetrization that involves sampling-based estimation, we use sample size 1 for training. For group averaging, sample size 1 for training led to optimization challenges, and therefore we use sample size 10 for training which yielded better results.
A.3.6 Real-World Graph Learning with Vision Transformer (Section 3.4)
Base Model  $f_{\theta}$  For Peptides-func/struct, the model is given adjacency matrix  $\mathbf{A} \in \mathbb{R}^{444 \times 444}$ , node features  $\mathbf{X} \in \mathbb{R}^{444 \times 64}$ , and edge features  $\mathbf{E} \in \mathbb{R}^{444 \times 444 \times 7}$ , zero-padded to maximal 444 nodes. The prediction target is binary classification logits  $\mathbf{Y} \in \mathbb{R}^{10}$  for Peptides-func, and regression targets
Table 7: Supplementary results for  $S_{n}$  invariant graph separation with  $S_{n}$  symmetrized GIN-ID base function. Baseline scores for GIN-ID-GA and GIN-ID-FA are taken from [74].
method	arch.	sym.	GRAPH8c ↓	EXP ↓	EXP-classify ↑
GIN-ID-GA	-	Sn	0	0	50%
GIN-ID-FA	-	Sn	0	0	100%
GIN-ID-Canonical.	-	Sn	0	0	84%
GIN-ID-PS (Ours)	-	Sn	0	0	100%
Table 8: Supplementary results for  $S_{n} \times E(3)$  equivariant  $n$ -body with  $E(3)$  symmetrized GNN base function. Baseline scores for GNN-FA and GNN-Canonical are from [74] and [41], respectively.
method	arch.	sym.	Position MSE ↓
GNN-FA	Sn	E(3)	0.0057
GNN-Canonical.	Sn	E(3)	0.0043
GNN-Canonical. (Reproduced)	Sn	E(3)	0.00457
GNN-GA	Sn	E(3)	0.00408 ± 0.00002
GNN-PS (Ours)	Sn	E(3)	0.00386 ± 0.00001
$\mathbf{Y} \in \mathbb{R}^{11}$  for Peptides-struct. For PCQM-Contact, the model is given adjacency matrix  $\mathbf{A} \in \mathbb{R}^{53 \times 53}$ , node features  $\mathbf{X} \in \mathbb{R}^{53 \times 68}$ , and edge features  $\mathbf{E} \in \mathbb{R}^{53 \times 53 \times 6}$ , zero-padded to maximal 53 nodes. The prediction target is binary edge classification logit  $\mathbf{Y} \in \mathbb{R}^{53 \times 53 \times 1}$ .
For the base model, we use a transformer with an identical architecture to ViT-Base that operates on  $14 \times 14$  tokens, same as in Appendix A.3.5. For Peptides-func and Peptides-struct, we use the following as the base model  $f_{\theta}:(\mathbf{A},\mathbf{X},\mathbf{E})\mapsto \mathbf{Y}$ :
$$
f _ {\theta} (\mathbf {A}, \mathbf {X}, \mathbf {E}) = \text {d e t o k e n i z e} _ {[ c 1 s ]} \left(\text {t r a n s f o r m e r} \left(\text {t o k e n i z e} _ {2 D} (\mathbf {A}, \mathbf {E}) + \text {t o k e n i z e} _ {1 D} (\mathbf {X}))\right)\right), \tag {53}
$$
where  $\mathrm{tokenize}_{\mathrm{2D}}: \mathbb{R}^{444 \times 444 \times (1 + 7)} \to \mathbb{R}^{(14 \times 14) \times 768}$  is 2D convolution with kernel size and stride 32,  $\mathrm{tokenize}_{\mathrm{1D}}: \mathbb{R}^{444 \times 64} \to \mathbb{R}^{196 \times 768}$  is 1D convolution with kernel size and stride 3, and  $\mathrm{detokenize}_{[\mathrm{cls}]}$  performs linear projection of the global [c1s] token [26] to the target dimensionality. For PCQM-Contact, we use the following as base model  $f_{\theta}: (\mathbf{A}, \mathbf{X}, \mathbf{E}) \mapsto \mathbf{Y}$ :
$$
f _ {\theta} (\mathbf {A}, \mathbf {X}, \mathbf {E}) = \operatorname {d e t o k i n i z e} _ {2 \mathrm {D}} \left(\text {t r a n s f o r m e r (t o k i n i z e} _ {2 \mathrm {D}} (\mathbf {A}, \mathbf {E}) + \operatorname {t o k i n i z e} _ {1 \mathrm {D}} (\mathbf {X}))\right), \tag {54}
$$
where  $\mathrm{tokenize}_{\mathrm{2D}}: \mathbb{R}^{53 \times 53 \times (1 + 6)} \to \mathbb{R}^{(14 \times 14) \times 768}$  is 2D convolution with kernel size and stride 4,  $\mathrm{tokenize}_{\mathrm{1D}}: \mathbb{R}^{53 \times 64} \to \mathbb{R}^{196 \times 768}$  is 1D convolution with kernel size and stride 1, and  $\mathrm{detokenize}_{\mathrm{2D}}: \mathbb{R}^{(14 \times 14) \times 768} \to \mathbb{R}^{53 \times 53 \times 1}$  is transposed 2D convolution with kernel size and stride 4.
Training We train our models with cross-entropy for classification and L1 loss for regression using AdamW [53] optimizer with batch size 128, learning rate 1e-5 except for PCQM-Contact where we use 5e-5, and weight decay 0.01. We train the models for 50k steps under learning rate warm-up for 5k steps then linear decay to 0 with early stopping based on validation loss, which usually takes less than 12 hours on 8 RTX 3090 GPUs with 24GB using PyTorch Lightning [29]. For the equivariant distribution  $p_{\omega}$ , we use noise scale  $\eta = 1$ , and use dropout with probability 0.1 except for PCQM-Contact where we do not use dropout. We use 10 samples for estimation during training.
A.4 Supplementary Experiments (Continued from Section 3)
In this section, we present additional experimental results that supplement the experiments in Section 3 but could not be included in the main text due to space constraints.
A.4.1 Graph Isomorphism Learning (Section 3.1)
In our experiments on graph isomorphism learning in Section 3.1, we mainly experimented for  $S_{n}$  symmetrization of an MLP. Here, we provide supplementary results on  $S_{n}$  symmetrization of a GIN base model with node identifiers, following [74]. The results can be found in Table 7. In accordance with Section 3.1, our approach successfully performs  $S_{n}$  symmetrization of GIN-ID.
Figure 4: Test accuracy of MLP  $f_{\theta}$  symmetrized by equivariant distribution  $p_{\omega}(g|\mathbf{x})$  trained on EXP-classify dataset across a range of training sample sizes. Inference sample size is set to 10.
Figure 5: Row- and column-wise entropy of aggregated permutation matrices  $\mathbf{P}_g\sim p_{\omega}(g|\mathbf{x})$  after trained on EXP-classify dataset across a range of training sample sizes. Dashed line indicates entropy measured with random permutation matrices from Unif  $(G)$
A.4.2 Particle Dynamics Learning (Section 3.2)
In our experiments on  $n$ -body dataset in Section 3.2, we experimented for  $\mathrm{S}_n\times \mathrm{E}(3)$  symmetrization using a 1D sequence transformer architecture which has  $2.3\times$  parameters compared to baselines. To provide parameter-matched comparison against baselines in literature, we apply our approach for  $\operatorname {E}(3)$  symmetrization of  $\mathbf{S}_n$  equivariant GNN base model that is widely used in literature [74, 41]. We faithfully follow [74, 41] on the experimental setups including training hyperparameters and the configuration of GNN base model, and only add  $\operatorname {E}(3)$  equivariant distribution  $p_{\omega}(g|\mathbf{x} - \bar{\mathbf{x}}\mathbf{1}^{\top})$  i.e.,  $q_{\omega}:(\mathbf{x} - \bar{\mathbf{x}}\mathbf{1}^{\top},\epsilon)\mapsto \mathbf{Q}_{g}$  by utilizing the 2-layer Vector Neurons architecture described in Appendix A.3.2 using only its O(3) prediction head. We use 20 samples for training and testing. The results can be found in Table 8. In accordance with the results in Section 3.2, our approach outperforms other symmetrization approaches and achieves a new state-of-the-art of 0.00386 MSE.
A.4.3 Effect of Sample Size on Training and Inference
In this section, we provide additional analysis on how the sample size for estimation of symmetrized function (Eq. (4)) affects training and inference. We use the experimental setup of EXP-classify (Section 3.1;  $S_n$  invariance) and analyze the behavior of MLP-PS with identical initialization and hyperparameters, only controlling sample sizes  $\in \{1,2,5,10,20,50\}$  for training. Specifically, we analyze (1) variance of permutation matrices  $\mathbf{P}_g \sim p_{\omega}(g|\mathbf{x})$  measured indirectly by the entropy of their aggregation  $\bar{\mathbf{P}} = \sum \mathbf{P}_g / N$  as in Section 3.1, (2) sample variance of the unbiased estimator  $g \cdot f_{\theta}(g^{-1} \cdot \mathbf{x})$  of the symmetrized function  $\phi_{\theta,\omega}(\mathbf{x})$  as in Eq. (4), and (3) sample mean and variance
Figure 6: Variance of estimation of MLP  $f_{\theta}$  symmetrized by equivariant distribution  $p_{\omega}(g|\mathbf{x})$  and trained on EXP-classify dataset for a range of training and inference sample sizes.
of the estimated task loss  $\mathcal{L}(\mathbf{y},g\cdot f_{\theta}(g^{-1}\cdot \mathbf{x}))$  where  $\mathcal{L}$  is binary cross entropy. All measurements are repeated 100 times and averaged over the inputs and labels  $(\mathbf{x},\mathbf{y})$  of the validation dataset.
Observations are as follows. First, models trained with smaller sample sizes need more iterations to converge, but after sufficient training (2000 epochs), all achieve  $>95\%$  test accuracy when evaluated with 10 samples (Figure 4). Second, models trained with smaller sample sizes tend to be more sample efficient, i.e., tend to perform a lower variance estimation. Their distribution  $p_{\omega}(g|\mathbf{x})$  tend to learn more low-variance permutations (Figure 5), and the models tend to learn low-variance estimation of output and loss (left and center panels of Figure 6). This indicates that small sample size may serve as a regularizer that encourages lower variance of the estimator. However, this regularization effect is not always beneficial in terms of task loss (right panel of Figure 6), as training sample size 1 achieves a poor task loss for all sample sizes presumably due to the optimization challenge caused by over-regularization. In other words, the sample size for training introduces a tradeoff; a small sample size takes more training iterations to converge, but serves as a regularizer that encourages lower variance of the estimator and thus a better inference time sample efficiency. On the other hand, larger sample sizes for inference consistently benefits all models (Figure 6).
Interestingly, this observed tendency is consistent with the theoretical claims in literature [67, 68] on the sampling based training of symmetrized models, which we reprise here. When training the symmetrized model  $\phi_{\theta ,\omega}(\mathbf{x})$  in Eq. (4), we cannot directly observe  $\phi_{\theta ,\omega}(\mathbf{x})$ , but observe samples of its unbiased estimator  $g\cdot f_{\theta}(g^{-1}\cdot \mathbf{x})$ . Thus, it can be questionable what objective we are actually optimizing during the sampling-based training. Based on [67, 68], it turns out that minimizing a convex loss function  $\mathcal{L}$  on the estimated output  $g\cdot f_{\theta}(g^{-1}\cdot \mathbf{x})$  is equivalent to minimizing an upper bound to the true objective on the symmetrized output  $\phi_{\theta ,\omega}(\mathbf{x})$ . This is because our estimation is no longer unbiased when computing loss, as we have the following from Jensen's inequality:
$$
\mathbb {E} _ {p _ {\omega} (g | \mathbf {x})} [ \mathcal {L} (\mathbf {y}, g \cdot f _ {\theta} (g ^ {- 1} \cdot \mathbf {x})) ] \geq \mathcal {L} (\mathbf {y}, \mathbb {E} _ {p _ {\omega} (g | \mathbf {x})} [ g \cdot f _ {\theta} (g ^ {- 1} \cdot \mathbf {x}) ]) = \mathcal {L} (\mathbf {y}, \phi_ {\theta , \omega} (\mathbf {x})). \tag {55}
$$
That is, minimizing the sampling-based loss is minimizing an upper-bound surrogate to the true objective. It has been claimed that optimizing this upper bound has an implicit low-variance regularization effect [67, 68], which is consistent with our observations. This also roughly explains why our distribution  $p_{\omega}(g|\mathbf{x})$  does not collapse to uniform distribution although we do not impose any low-variance regularization explicitly; training to directly minimize the task loss with samples implicitly nudges the distribution towards low-variance solutions.