Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance
Authors: Jinwoo Kim
Dat Tien Nguyen
Ayhan Suleymanzade
Hyeokjun An
Seunghoon Hong
Venue: NeurIPS 2023 spotlight (2023)
Keywords: equivariant machine learning
transformers
graphs
general-purpose architectures
TL;DR: A framework for equivariant learning with non-equivariant architectures, e.g., vision transformer for graph recognition.

Abstract: 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.

Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance
Jinwoo Kim
Tien Dat Nguyen
Ayhan Suleymanzade
Hyeokjun An Seunghoon Hong
KAIST
Abstract
We present novel framework overcome the limitations of equivariant archi- tectures 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 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 symmetrization. We show that this approach ensures not only equivariance 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 symme- try 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 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 available at ftthttt://// pttttp:://iithute cww cotmcb/33////33/pppp
Introduction
Many perception problems in machine learning involve functions that invariant equivariant to certain symmetry group of transformations of data. Examples include learning on sets and graphs, point clouds, molecules, proteins, and physical data, name few [10 [2). Equivariant architecture design has emerged as a successful approach, where every building block of a model is carefully restricted to equivariant symmetry group of interest [2] 23). However, equivariant architecture design faces fundamental limitations, as individual construction of models for each group can be laborious or computationally expensive 99 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 others of different symmetries the architecture would be incompatible.
This motivates us to seek a symmetrization solution that can achieve group invariance and equivari- ance with general-purpose, group-agnostic architectures such as MLP transformer. As basic form of symmetrization, any parameterized function fa : :X' + y on vector spaces X', ) can be made invariant equivariant by group averaging [T02 671. i.e., averaging over all possible transformations of inputs X' and outputs E y by symmetry group G = {
de (x = G
g. folg x), Σg gEG
where is equivariant invariant the group G. An important advantage is that the symmetrized function Φy can leverage the expressive power the base function fo: it has been shown that Φy is a universal approximator of invariant equivariant functions if fo is universal approximator [102]
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Probabilistic Symmetrization
B(3) equivariant
S. x B(3) equivariant
S2 equivariont
(,,c(
g fo(g-1 x)
G= B(3)
G= S. 8, B(3)
Base Fonccton fs
9
Pooggox
G Equivcrraant Distribution
motion
goometric graph
graph
Figure Overview of probabilistic symmetrization. We symmetrize an unconstrained base func- tion fo into an equivariant function 6eww for group G using learned equivariant distribution pw (g]x).
which includes MLP 35
transformer [044 On the other hand, an immediate challenge is that for many practical groups involving permutation and rotations, the cardinality of the group |G] is large infinite, 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., frame [74] average over [67 68,. or implement relaxed version of equivariance [1] 83
An alternative method for tractability to interpret Eq. TT as an expectation with uniform distribution Unif(G) over the compact group G (62). and use sampling- based average to estimate it [102] 67 68):
(x(x) (g--ii(( [g fe x)]
where g. folg-1 1 x) serves an unbiased estimator of do (x).. While simple and general, this approach has practical issues that the base function fo 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 Unif(G) for the expectation in Eq. (2) with parameterized distribution Pa (g)x)) in a way that equivariance and expressive power are always guaranteed, and train it end-to-end with the base function fo to directly minimize task loss. We show that the distribution Pw (g(x) only needs to satisfy one simple condition to guarantee equivariance and expressive power: it has to be probabilistically equivariant (9]. This allows us generally implement (g x) as noise-outsourced map (x, E) + g with an invariant noise e and a small equivariant network qwo which enables gradient- based training with reparameterization 47) can enhance the learning of fo by producing group transformations of lower variance compared to Unif(G) so that fo is less burdened, and coordinating with fe maximize task performance. We refer to our approach probabilistic symmetrization. An overview is provided in Figure 1
Pa is trained,
We implement and test method with MI.P and transformer. In particular, our transformer backbone is architecturally identical to patch- based vision transformers [26]. which allows to initialize most 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.. Pw.(g] (g|x) for a wide range of practical symmetry groups including permutation (S,,) and Euclidean groups (O (OSO(d)) and E. E/SE(d)) and their product combi- nations (e.g., S, x 0()).. all of which are combinatorial or infinite groups Empirical tests on 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 744 4T 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 symmetrie modality (graphs).
ceneral-purposs architectures the base function fo:

---

2 Probabilistic Symmetrization for Equivariance
We introduce and analyze our approach called probabilistic symmetrization. which involves an equivariant distribution Pa and group-agnostic base function fo. in Section 2.1 implementation of Pa for practical groups including permutations and rotations Section 2.2 Then, we describe our choice of base function fo focusing on MLP and, transformers in particular, in All proofs can be found in Appendix A..
2.1 We then describe
Section 2.3
Problem Setup In general, our goal is to construct function ó : X + y on finite vector spaces X. that is invariant or equivariant to symmetry specified by group G {g I This is formally described by specifying how the group act as transformations on the input and output. A group representation p : + GL.(X), where (x' ) is the set of all invertible matrices
X, associates
each group element y E G to an invertible matrix p(g) that transforms a given vector x E X' through
x 9 x x(g)x. Given that, a function @ X + y G equivariant if:
(p)(g)x) p2(g) Vx E Kg E G,
where the representations and P2 are on the input and output, respectively. G invariance is a special case of equivariance when the output representation is trivial, p2(g) I.
2.1 Probabilistic Symmetrization
To construct a G equivariant function do. group averaging symmetrizes an arbitrary base function fo X y' by taking expectation with uniform distribution the group (Eq. (2). Instead, propose to use an input-conditional parameterized distribution Pa x) and symmetrize fo as follows:
@9 w (x) = Ep x) (p2())le(pr(o) "x)],
where the distribution Pa (glx) itself satisfies probabilistic G' equivariance:
P= (g (g' (g')xx; Vx E X, E G.
(5)
Importantly, show that probabilistic symmetrization with equivariant Pw guarantees equivariance as well as expressive power of the symmetrized 60...
Theorem 1. If Pws is G equivariant, then so.w G equivariant for arbitrary fo-
Proof. The proof can be found in Appendix
A.I.I
Theorem 2. If Pa. is G equivariant and fo is a universal approximator, then do is a universal approximator of G equivariant functions.
Proof. The proof can be found in Appendix A.1.2
While the base function fe that guarantees universal approximation can be chosen in a group- agnostic manner e.g., an MLP B35 20 or a transformer on token sequences [104], the distribution Pa needs to be instantiated group specifically to satisfy G equivariance. A simplistic choice using uniform distribution Unif(G) (G) for all inputs x with no parameterization (reducing to group averaging), which is technically equivariant and therefore guarantees equivariance and universality. However, appropriately parameterizing and learning Pa, can provide distinguished advantages compared to the uniform distribution, as it can (1) learn from data to collaborate with (pre-trained) base function fo to maximize task performance and (2) learn to produce more consistent samples g ~ Pa (=[g)x) (x) that can offer more stable gradients for the base function fe during early training.
We now provide a generic blueprint of G equivariant distribution Pa (gix) for any compact group G. Our goal is to sample g (g x) to obtain group representation p(g) for symmetrization (Eq. (4) in a differentiable manner so that pw can be trained end-to-end. Since we only need sampling and there no need to evaluate likelihoods, we simply implement pw (g|x) as a nouse-outsourced, differentiable transformation 4 (x, e) of a noise variable e E E that directly outputs group representation p(g):
p(g) qu (x, E), e~ p(E),
(6)
In this paper, assume the group G to be compact.
3

---

where is G equivariant and p(c) is G invariant under an appropriate representation p':
I ((9(())x,''())) (x,) e), p(c) = p(o))((())))), Vx VxEX,CEE, Given above implementation, we can show the G equivariance of pui
(7)
Theorem 3. q is G equivariant and p(E) is G invariant under representation of that det p'(g) 1Vg E G, the distribution Pu x) characterized by qu : (x, E) + p(g) is G equivariant.
Proof. The proof can be found in Appendix A.1.3 A.I..
In practice, one can use any available G equivariant neural network to implement qw e.g., graph neural network for the symmetric group S,, or an equivariant MLP which can be constructed for any matrix group [3) Since we expect most of the reasoning to be done by the base function fa. the equivariant network qu 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
p(c),
simple choices often suffice for G invariance. For example, standard normal e ~ (O0, I,) provides invariance for the symmetric group S, as well as the (special) orthogonal group ( and SO(n).
One important detail in designing q. constraining its output be a valid group representation p(g). For this, we apply appropriate postprocessing to refine neural network features into group representa- tions, e.g., Sram-Schmidt orthogonalization to obtain a representation p(g) =R"*" of the orthogonal group E (n). Importantly, to not break the G equivariance of 9w. this postprocessing needs to be equivariant itself, e.g., Cram-Schmic process itself ((n) equivariant [4T]. Implementations of pa for range of practical symmetry groups are provided in detail in Section 2.2
2.2 Equivariant Distribution Pw
We present implementations of the G equivariant distribution (glx) for range of practical groups demonstrated in our experiments (Section 3). Formal proof's of correctness are in Appendix A.I.4
Symmetric Group S,, The symmetric group S,, over finite set of elements contains all permu- tations of the set, which describes symmetry to ordering desired for learning set and graph data. The base representation is given by p(g) Ps where P, {0, 1]nx= a permutation matrix for g.
To implement S, equivariant distribution (glx) that provides permutation matrices Pg from graph data x. we use the following design. We first sample invariant noise e E Rnxd from i.i.d. uniform Unif(0, with noise scale 7). For the Sn equivariant map 9. (x, E) p(g). first neural network (GNN) as an equivariant map (x, E) + that outputs nodewise scalar E R". Then, use below argsort operator Z - P, 9 obtain permutation matrix 773 98
graph
assuming Z is tie- ?
P3 = eq sort ))
(8)
where eq denotes elementwise equality indicator. The argsort operator is S,, equivariant, i.e., it maps Z - P.p for all P.'' E S.. To backpropagate through P, during training, we use straight- through gradient estimator 66] with approximate permutation matrix P,, P, obtained from differentiable relaxation of argsort operator 61
33
98]
details can be found in Appendix A.3.
Orthogonal Group O(n), SO(n) The orthogonal group O(n) contains all roto-reflections in R" around origin, and the special orthogonal group SO(n) contains all rotations without reflections. These groups describe rotation symmetries desirable in learning geometric data. The base group representation for O(n) is given by p(g) = Q2 where Q2 is the orthogonal matrix for g. For SO(n). the representation is p(g) Q where + the orthogonal matrix of g with determinant +1.
To implement ((n)/SO(n) equivariant distribution (g(xx) that provides orthogonal matrices given input data x, we use the following design. We first sample invariant noise e E Raad from i.i.d. normal A (1 n²2) with noise scale 77. For the (n))SO(()) equivariant map Pa : (x, E) + p(g). we first use an equivariant neural network as map (x, E) + Z that outputs n features Z E R=*x Then, assuming Z is full-rank we use Gram-Schmidt process for orthogonalization, which differentiable
This can assumed since the invariant noise e E Ra serves as tiebreaker between the n nodes. In practice, we add random Gaussian matrix of tiny magnitude to Z to prevent rank collapse.

---

and (n) equivariant [4. This completes the postprocessing Z - Qg for O(n). For (O(n), we can further use a simple scale operator Q - Q3 to set the determinant of orthogonalized matrix +1:
scale:
Q1
Qn
det(Q) Q1
Q₂
(9)
The scale operator is differentiable and SO(n) equivariant, iee, it maps Q Q + Qg
Q for all
Q E (O(n), thereby completing the postprocessing Z + + for SO(n).
Euclidean Group B(n), SE(n) The Euclidean group E(n) contains all roto-translations and re- flections in IX" and their combinations, and the special Euclidean group SE(n) contains all roto- translations without reflections. These groups are desired in learning physical systems such particle in motion. Formally, the Euclidean group is given as combination of orthogonal group and translation group E(n) (n)) T(n), and similarly SE(n) = SO(n) x T(n). As the translation group (nn) is non-compact which violates our assumption for symmetrization, we handle it separately, Following prior work [74[[]] we subtract the centroid from the input data x as x X, and add it to the rotation symmetrized output as X + fo(g-1 (x x))) where is sampled from (nn)SS((n) equivariant Pw (g|x x). This makes the overall symmetrized function E(n)/SE((n) equivariant.
Product Group H K While we have described several groups individually, in practice often encounter product combinations of groups G И x K that describe joint symmetry to each group H and K. For example, S,, x 0(3) describes joint symmetry to permutations and rotations, which is desired in learning point clouds, molecules, and particle interactions. In general, an element applied
of H K is given g (h, k) where h E H and k E K, and group operations elementwise gg' = (h, (h', k²) = (hh', k:k). The base group representation is accordingly given as pair of representations p(g)
(p(h), p(k)). While common approach to handling H x 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,, 0(3) equivariance [74] 4T), extend to full symmetrization on H x K since our goal is not imposing any constraint on the base function fo-
To implement H x K equivariant distribution P.. (g that gives p(g) (p(h), p(k) from data use the following design. We first sample invariant noise e from i.i.d. normal NN n7²) with scale For the H x K equivariant map qu (x, c) -> (p(h), p(k)), we employ H x K equivariant neural network as a map (x, c) + [Z Zx) such that the postprocessing for each group H and K provides p(h) and Z p(k) respectively, leading full representation p(g) = (p(h), P())).
maps NU For this whole procedure to be H x K equivariant, it is sufficient to have Zμ be K invariant and Ze be H invariant. These are special cases of H K equivariance, and is supported by range of equivariant neural networks especially regarding S. x 0(3) or Sn x SO(3) equivariances 23
2.3 Base Function fo
We now describe the choice of handled by the equivariant distribution P., P_ (g(x), any symmetry concern is hidden from fa, allowing the inputs and outputs y to be treated as plain multidimensional arrays. This allows us to implement fo with powerful ceneral-purpose architectures, namely as an MLP that operates on flattened vectors of inputs and outputs, or a transformer as we describe below.
base function fo x xx y. As group symmetry is
Let inputs x E Ra Xn, and outputs y y = xn/jxe be multidimensional arrays with and c' channels, respectively. Our transformer base function fo x + y is given as:
fe detokenize o transformer o tokenize,
(10)
where tokenize X - g parses input array a sequence of m tokens, transformer Ro Im is a standard transformer encoder on tokens used in language and vision [24] 26), and detokenize m - y' decodes encoded tokens to output array. For the tokenizer and detokenizer, we can use linear projections on flattened chunks of the array, which directly extends flattened patch projections in vision transformers higher dimensions 26 85 86). This enables mapping between different dimensional inputs and outputs, e,g, for graph node classification with x E Raxxxe and y nax'' it is possible to use 2D patch projection for the input and ID projection for the output.
5

---

Above choice of fo offers important advantages including universal approximation [T04] and ability share and transfer the learned knowledge in 9 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. our experiments, we only replace the tokenizer and detokenizer of vision transformer pre-trained on ImageNet-21k [99] (6) and fine-tune it to perform diverse S, equivariant tasks such as graph classification, node classification, and link prediction.
2.4 Relation Other Symmetrization Approaches
We discuss relation of our symmetrization method to prior ones, specifically group averaging 102 4
frame averaging 74. and canonicalization [4 An extended discussion on broader related work can be found in Appendix A.2 Since these symmetrization methods share common formalization y = E, {g x)], one can expect close theoretical relationship between them. We observe that probabilistic symmetrization is quite general; based on particular choices of the G' equivariant distribution Pw. ()]()))) it can become most of the related symmetrization methods as special cases. This can be easily seen for group averaging [102], as the distribution Pa can reduce to the uniform distribution Unif(G) over the group. Frame averaging [74] also takes an average, but over a subset of group given by frame F X' + 2 \ 0; importantly, it is required that the frame itself is G equivariant F F(p(g)x) = (x). We can make the following connection between our method and frame averaging. by adopting the concept of stabilizer subgroup Gx = {g G: p(g)x = x}:
Proposition 1. Probabilistic symuettrizatio with G equivariant distribution pw ( x) can become frame averaging (44 by assigning uniform density to a set of orbits Gxg for some group elements
Proof. The proof can be found in Appendix A.I.3
Canonicalization [41] uses a single group element for symmetrization, produced by a trainable canonicalizer Cuu : X + G. Here, is required that the canonicalizer itself satisfies relaxed G equivariance C(p(g)x) = 9g' 9g' C(x) up arbitrary action from the stabilizer g' Gx. We now show:
Proposition 2. Probabilistic symmetrization with G equivariant distribution xx) can become canonicalization [41 by assigning uniform density to a single orbit Gxg of some group element g.
Proof. The proof can be found in Appendix A.1.5
Assuming that stabilizer Gx is trivial this can be implemented with our method by removing random noise e, which reduces pa to deterministic map p(g) = qu (x). We use this approach to implement canonicalizer for the S, group, while [41] only provides canonicalizers for Fuclidean groups.
3 Experiments
We empirically demonstrate and analyze prohabilistic symmetrization on a range of symmetry groups S,, E(3) (O(3)) and the product S,, xB(3) (S,, 0(3)) on a variety of invariant and equivariant tasks, with general-purper
base functions fe 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 (S,,) has been considered important and challenging. as simple and efficient GNNs are often limited in expressive power to certain Weisfeiler- isomorphism tests like 1-WL [TOT] [56]. Since an MLP equipped with probabilistic symmetrization is in theory universal and S, 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 (S,, invariant). GRAPH8c [2] consists of all non-isomorphh -somoooo connected graphs with 8 nodes, and EXP [T] 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 [I02). frame averaging 774,, and canonicalization [41. Our method uses the same ML. architecture to symmetrization baselines, and its S,, equivariant distribution p., for symmetrization is implemented

---

Table Results for Sn 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 aceuracy at best validation accuracy. The columns arch. and sym. denote architectural and symmetrized equivariance, respectively. The results for baselines are from 74 except for MIP MILP-Canonical which is tested by us.
method
arch.
sym.
GRAPH8c
EXP=classify 1
GCN
Sn
4755
600
50%
GAT
Sn
1828
600
50%
GIN
386
600
50%
ChebNet 22
Sn
44
71
82%
PPGN 56
0
0
100%
GNNML3 2
Sn
0
0
100 %
MLP-GA 102
S
0
0
50%
MLP-FA
S.
%
MI P.nnnnnic
0
0
50%
MLP.PS (Ouns), lixed
79.5%
MI P.PS (Ours)
0
100 %
3.8 1
0.7
perautatious
2.8
Entropy
Loss
0.5
2.6
0.3
2.4
0.1
0
500 Epochs
1000
500 Epochs
1000
Epoch 1
Epoch 1000
Figure 2: Learned (glxx) x)) over time. The entropy of aggregated permutation matrices P P from Pp (glx) for each input drops in early training, indicating that the distribution learns to produce lower-variance permutations in below visualizations.
using a 3-layer GIN 101 which is I-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 different base model can be found in Appendix A.4.
The results are in Table I 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, naïve symmetrization with group averaging fails, presumably because the MLP fails to adjust 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 pa ( x) can learn to provide more consistent permutations during early training, the consistency of samples from (g x) over training progress by sampling N 50 permutation matrices P, ~ Pw (g[x) and measuring the row-wise entropy of their average P = Σ , for each input x. The more consistent the sampled permutations, the sharper their average, and lesser the entropy. As training progresses, Pw: learns 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 Pw (g]x) be consistent from the first place, would it work? To answer this, we also tested coopoobbailiittcc version of our model that uses single permutation per input p(g) = q. (x). which is a canonicalizer under relaxed equivariance 41 as described in Section 2.44 As Table [ canonicalization fails, suggesting that probabilistic nature of pa, (g xx) can be beneficial for learning. For another deterministic version of our model made by fixing the noise e (Fq. (6) at initialization, the performance drops to 79.5%, further implying that stochasticity of (_. has role.
measured illustrate in Figure 2 In the figure,
7

---

Table 2: Results for S,, x 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 MSF +
SE(3) Transformer
S, x SE(3)
0.0244
TFN [9]]
S, SB(3)
0,0155
Radial Field 49
E(3)
0.0104
EGNN 84
S, E(3)
0.0071
GNN-FA 774
S
F(3)
0.0057
GNN-Canonical.
S.
E(3)
0.0043
S., x E(3)
0.00508
Transformer-Ci
S,, H(3)
0.00414 + 0.00001
Transformer-Pi (Ours)
S, B(3)
0.00401 = 0.00001
3.2 Particle Dynamics Learning with Transformer
Learning sets or graphs attributed with position and velocity in 3D 3D xE(3) practically significant as they universally appear in physics, chemistry, and biology applications. While prior symmetrization methods employ an already S,, equivariant base function and partially symmetrize the E(3) part, we attempt to symmetrize the entire product group S, x E(3) and choose the base model fo as a sequence transformer leverage its expressive power. For empirical demonstration, we adopt the experimental setup of [4]] and use the n-body dataset [84 31 where the task is predicting the position of n charged particles after certain time given their initial position and velocity in RS (S,, x B(3) equivariant). We compare our method to S,, E(3) equivariant neural networks and partial symmetrization methods applying B(3) symmetrization to GNNs. We also test prior symmetrization methods the full group S,, E(3) along method, but could not test for frame averaging since equivariant frames for the full group S, x E(3) was not available in current literature. Our method is implemented using transformer with sequence positional encodings with around 2.3x parameters of the baselines, and the Sn x E(3) equivariant distribution Pw for symmetrization is implemented using 2-layer Vector Neurons [3 that has around 0.03x of parameters to the transformer. We use 20 samples for symmetrization during training, and use 10x sample size for testing since the task is regression where appropriate variance reduction is necessary guarantee reliable performance. Further details can be found in Appendix [A.3.4 and supplementary results on E(3) partial symmetrization of GNN base model can be found 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 = particles interacting with each other, such that combining with expressive base model fo (a sequence transformer) can adjust to = 120 equally possible permutations and their rotations in 3D. Nevertheless, our method outperforms group averaging and achieves MSE, presumably as the parameterized distribution Pa learns to further maximize task performance. On the other hand, the canonicalization approach, implemented by eliminating noise variable from our method (Section 2.4), performs relatively poorly. We empirically observe that So memorizes the per-input canonical orientations provided by p(g) qu (x) that do not generalize to test inputs. This again shows that probabilistic nature of p_ x) can be beneficial for performance.
state of the art 0.00401
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 fo. such that we can leverage knowledge learned from other symmetries by transferring the parameters 0. 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,, 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-18 188 nodes, whose task is finding certain subgraph pattern by binary node classification.
8

---

Table Results for S, equivariant node classification on PATTERN. We report test accuracy 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 1
GCN 48 16 layers
85.614
GAT 97.. 16 layers
78.271
GatedGCN M. 16 layers
85,568
GIN [001. 16 layers
85,387
RingGNN 16 2 layers
86.245
RingGNN 16 8 layers
diverged
PPGN 56
layers
85 661
PPGN 56 8 layers
diverged
VIT-GA, I-sample
76.776 + 0.137
VIT-GA, 10-sample
83.119 + 0.048
VIT-GA, I-sample
ImageNet-21k
81,407 = 0.101
VIT-GA, 10-sample
ImageNet-21k
84.351 + 0.053
VrT-FA
70.063
VIT-FA
ImageNet-2Ik
79.637
VT-Canonical.
85,460
VTT VT--Cannical.
ImageNet-21k
86. 166
VIT-PS (Ours), 1-sample
85.542 0.012
VIT-PS (Ours), 10-sample
85.635 0.021
VYT-PS (Ours), 1-sample
ImageNet:21k
86.226 + 0.028
ViT-PS (Ours). 10-sumplc
ImageNet-21k
86.285 0.015
Based on pre-trained ViT 126 89 we construct the base model fe by modifying only the input and output layers take the flattened patches of 2D zero-padded adjacency matrices of size 188 188 and produce output as ID per-node classification logits of length 188 with 2 channels. In addition to standard GNN we compare group averaging, frame averaging. and canonicalization our method, and also test whether pre-trained representations from ImageNet-21k is beneficial for the task. For the Sn equivariant distribution P., our method and canonicalization, we use 3-layer GIN [00]] 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, observe that transferring pre-trained VIT parameters consistently improves node classification for all symmetrization methods. 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 + S, equivariance) and the data generating process (natural images > 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 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 good initialization point or transferable computation motifs.
On the other hand, frame averaging shows 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 p(g) q (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. result, our model based on pre-trained ViT and 10 samples for testing achieves 86.285% test accuracy, surpassing all baselines.
'We note that careful engineering such graph positional encoding improves performance of GNNs, but we have chosen simple and representative GNNs in the benchmark [27] to provide controlled comparison.

---

Table 4: Results for real-world graph tasks.
We report test performance at best validation performance.
method
Peptides func
Peptides-stme struct
PCOM-Contar
AP f
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
GaledGCN
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, now provide preliminary test of it in real-world graph learning. We three real-world graph datasets from (88 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 (S,, equivariant,, on whether two atoms would be proximal when the molecule is in 3D space. Peptides-fune and Peptides-struct with 444 nodes in maximum and the tasks are property prediction (S,, invariant), requiring multi- label classification for Peptides-fung 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 VīT-Base pre-trained on ImageNet-21 21k as the base model fo and 3-layer GIN as the equivariant distribution (g)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.366
based on the same set of 15,535 protein graphs
The results are in Table In Peptides-fume 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 (0.6527 - 0.8311 1 for Peptides-fune func AP, 0.1355 0.3268 for PCQM-Contace Hits @ 1). This demonstrates the scalability of our method as Peptides-fune involves 444 maximum nodes, and also its generality as it performs well for both S, invariant (Peptides-funde and equivariant (PCOM-Contace baselines, our method does not require costly Laplacian eigenvectors 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 predictions, and leave improving regression performance as future work.
tasks. We also note that, unlike some
Conclusion
We presented probabilistic symmetrization, general framework that learns distribution of group transformations conditioned on input data for symmetrization an arbitrary function. By character- izing that the only condition for such distribution is equivariance to data symmetry, 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 consis- tent improvement over other symmetrization methods, and is competitive or outperforms equivariant networks 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 further discuss in Appendix A.5 we plan to address these future work.
Acknowledgements This work was supported in part by the National Research Foundation of Korea (NRF202IRIA WTFF2202IIIII101225555 and NRF202 IR AAA4333328)34 and ITTP grant (2021-0-00537, 2019-0-00075, and 2021-0-02068) funded by the Korea government (MSIT).
'we note that the baseline architectures are constructed within 500k parameter budget as convention 28 while identical architecture to ViT-Base leverage pre-trained representations.
10

---

References
[1] R. Abboud, I. i. Ceylan. M. Grohe, and T. Lukasiewicz. The surprising power of graph neural networks
with random node initialization. In IJCAI, 2021. (Cited on 6
23
[2] M. Balcilar, Héroux, B. Gaüzère, Vasseur, Adam, and Honeine. Breaking the limits of message
passing graph neural networks. In ICML, 2021. (Cited on 6 7 (3)
[3] S. Basu, P. Katdare, P. Sattigeri, V. Chenthamarakshan, K. Equivariant few-shot learning from pretrained models. arXiv, 2023. (Cited on (2
P. Das, and L. R. Varshney.
[4] S. Basu, P. Sattigeri, K. N. Ramamurthy, V. Chenthamarakshan, K. R. Varshney, I.. R. Varshney, and Group equivariant fine-tuning of pretrained models. arXiv, 2022. (Cited on 2
P. Das.
[5] W. Battaglia, J. B. Hamrick, V. Bapst, A. Ganzhez-Gonzalez,
Zambaldi, M. Malinowski, A. Tac- chetti, D. Raposo, A. Santoro, R. Faulkner, Ç. Gülçchre, H. Song, A. J. Ballard, Gilmer, G. E. Dahl, A. Vaswani, K. R. Allen, C. Nash, V. Langston, C. Dyer, N. Heess, D. Wierstra, P. Kohli, M. M. Botvinick,
O. Vinyals, Y. Li, and R. Pascanu. Relational inductive biases, deep learning, and graph networks. arXiv,
2018. (Cited on 22
[6] Y. Bengio, N. Léonard, and A. C. Courville, Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv, abs/1308 3,332, 2013. (Cited on 24
[7] G. W. Benton, Finzi, P. Izmailov, and G. Wilson, Learning invariances in neural networks from training data. In NeurIPS, 2020. (Cited on 22 23p
[8] B. Bevilacqua, F. Frasca, D. Lim, B. Srinivasan, C. Cai, G. Balamurugan, M. M. Bronstein, and H. Maron.
Equivariant subgraph aggregation networks. ICLR, 2022. (Cited on 22)
[9] B. Bloem-Reddy and Y. W. Teh. Probabilistic symmetries and invariant neural networks. J. Mach. Learn. Res., 2020. (Cited on 2
[10] A. Bogatskiy, Ganguly, T. Kipf, Kondor, D. W. Miller, D. Murnane, J. T. Offermann, M. Pettee,
P. Shanahan, C. Shimmin, and S. Thais. Symmetry group equivariant architectures for physics. arXiv,
2022. (Cited on 22
[11] X. Bresson and T. Laurent. Residual gated graph convnets. arXiv, 2019. (Cited on 10
[12) M. M. Bronstein, J. Bruna, T. Cohen, and P. Velickovic, Geometric deep learning: Grids, groups, graphs,
geodesics, and gauges, arXiv, 2021. (Cited on
22
[13] T. Brown, B. Mann, N. Ryder, M. Subbiah, Kaplan, Dhariwal, A. Neelakantan, P. Shyam, G. Sastry,
A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess. J. Clark, C. Berner.
S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In NeurIPS, 2020. (Cited on 22)
[14] C. Cai and Y. Wang. A note on over-smoothing for graph neural networks. arXiv, 2020. (Cited on 22
[15] M. Chen, 7. Wei, Z. Huang, B. Ding, and Y. Li. Simple and deep graph convolutional networks. In ICMI.,
2020. (Cited on [Tp
[16] Z. Chen, S. Villar. L. Chen, and J. Bruna. On the equivalence between graph isomorphism testing and function approximation with gnns. In NeurlPS, 2019. (Cited on [
[17) E. Chien, C. Pan, J. Peng, and O. Milenkovic. You are allset: A multiset function framework for hypergraph neural networks. In ICLR, 2022. (Cited on 22
[18] T. Cohen and M. Welling. Group equivariant convolutional networks. In ICML,, 2016. (Cited 22
[19] T. S. Cohen and M. Welling, Stecrable cnns. In ICLR, 2017. (Cited on 22
[20] G. Cyhenko. Approximation by superpositions of a sigmoidal function. Math. Control. Signals Syst.,
1989. (Cited on 3 22
[21] G. Dasoulas. L. D. Santos. K. Scaman. and A. Virmaux. Coloring graph neural networks for node disambiguation. In IJCAI, 2020. (Cited 22
[22] M. Defferrard, X. Bresson, and P. Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In NeurIPS, 2016. (Cited on
[23] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L.. J. Guibas. Vector neurons: A general framework for squui-qqui networks. In ICCV, 2021. (Cited on
[24] J. Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLLT, 2019. (Cited
[25] T. Dinh, Y. Zeng. R. Zhang, Z. Lin, M. Gira, S. Rajput, J. Sohn, D. S. Papailiopoulos, and K. Lee. LIFT: language-interfacee mie-tuning for non non-language machine learning tasks. In NewrIPS, 2022. (Cited on
22
11

---

[26] A. Dosovitskiy, L. Beyer. A. Kolesnikov. D. Weissenborn, X. Zhai. T. Unterthiner, M. Dehghani.
M. Minderer. G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words:
Transformers for image recognition at scale. In ICLR, 2021. (Cited on 2
26
[27] V. P. Dwivedi, C. K. Joshi, T. Laurent, Y. Bengio, and X. Bresson. Benchmarking graph neural networks.
arXiv, 2020. (Cited on g
9
23 25
[28] V. P. Dwivedi, L. Rampásek, M. Galkin, A. Parviz, G. Wolf, A. T. Luu, and D. Beaini. Long range graph
benchmark. In NeurlPS, 2022. (Cited on 10 23
[29]
W. Falcon. Pytorch lightning. https://gitiin http://ggttubb..cooooo httpab...mm/iiggtttnnngg///aaaia
2019. (Cited on
25
26
[30]
M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer
perceptrons for arbitrary matrix groups. ICML, 2021. (Cited on 19 22
[3]] F. Fuchs. D. E. Worrall, V. Fischer, and M. Welling, transftransorr::rs:: 3d roto- tro-translation equivariant
attention networks. In NewrlPS, 2020. (Cited on 8 22
23
[32] J. Gilmer, S. S. Schoenholz, P. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum
chemistry. In ICML, 2017. (Cited on 23
[33] A. Grover, F. Wang. Zweig, and S. Ermon. Stochastic optimization of sorting networks via continuous relaxations. In ICLR, 2019. (Cited on (4)
[34] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv, 2018. (Cited on 25)
[35] K. Hornik, M. B. Stinchcombe, and H. White. Multilayer feedforward networks are universal approxima- tors. Neural Networks, 1989. (Cited on
[36] W. Hu, B. Liu, J. Gomes, M. Zitnik, Liang, V. Pande, and J. Leskovec. Strategies for pre-training graph neural networks. In ICLR, 2020. (Cited TO
[37] A. Immer, T. F. A. van der Ouderaa, G. Rätsch, V. Fortuin, and M. van der Wilk. Invariance learning in
deep neural networks with differentiable laplace approximations. In NeurIPS, 2022. (Cited on 22
[38] A. Jaegle, S. Borgeaud, J. Alayrac C. Doersch, C. Jonescu, D. Ding, S. Koppula, D. Zoran, A. Brock,
E. Shelhamer, O. Hénaff, M. M. Botvinick, A. Zisserman, O. Vinyals, and Carreira. Perceiver IO: general architecture for structured inputs & outputs. In ICLR, 2022. (Cited on 22
[39] A. Jaegle, F. Gimeno, A. Brock, O. Vinyals, A. Zisserman, and J. Carreira. Perceiver: General perception with iterative attention. In ICML, 2021. (Cited on 22)
22)
[40] C. K. Joshi, C. Bodnar, S. V. Mathis, T. Cohen, and P. Liò. On the expressive power of geometric graph neural networks. arXiv, 2023. (Cited on
[41] Kaba, A. K. Mondal, Zhang, Bengio, and S. Ravanbakhsh. Equivariance with learned canonical-
ization functions. arXiv, 2022. (Cited on
6
30
[42] P. Kicki, P. Skrzypezynski, and M. Ozay. A new approach to design symmetry invariant neural networks. IJCNN, 2021. (Cited on 2 22)
[43] J. Kim, D. Nguyen, S. Min, S. Cho, M. Lee, H. Lee, and S. Hong. Pure transformers are powerful graph
learners. In NeurlPS, 2022. (Cited on 22
[44] J. Kim, S. Oh, S. Cho, and S. Hong. Equivariant hypergraph neural networks. In ECCV, 2022. (Cited on
[45] J. Kim, Oh, and Hong. Transformers generalize deepsets and can be extended to graphs and hypergraphs. In NeurlPS, 2021. (Cited on (
[46] D. P. Kingma and Ba. Adam: A method for stochastic optimization. In ICLR, 2015. (Cited 25
[47] D. P. Kingma and M. Welling. Auto-encoding variational bayes. In ICLR, 2014. (Cited on 2
[48] T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In ICLR,
2017. (Cited on 7 g 10)
[49] J. Köhler, L. Klein, and F. Noć. Equivariant flows: Sampling configurations for multi-body systems with symmetric energies. arXiv, 2019. (Cited on
[50] D. Kreuzer, D. Beaini, W. L. Hamilton, V. Létoureau, and P. Tossou. Rethinking graph transformers with spectral attention. NewrIPS, 2021. (Cited on 10 22
[51] J. Lee, Y. Lec, J. Kim, A. R. Kosiorck, S. Choi, and Y. W. Tch. Set transformer: A framework for attention- based permutationn-invaiian neural networks. In ICML, 2019. (Cited on 22
[52] J. Li, D. Li. Savarese, and S. C. H. Hoi. BLIP-2: bootstrapping language-image pre-training with frozen image encoders and large language models. arXiv, 2023. (Cited on 22
[53] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In ICLR, 2019. (Cited 25 26
12

---

[54] K. Lu, A. Grover, P. Abbeel, and I. Mordatch. Frozen pretrained transformers universal computation engines. In AAAI, 2022. (Cited on 22
[55] S. Luo, T. Chen, Y. Xu, S. Zheng. T. Liu, D. He, and L. Wang. One transformer can understand both 2d &
3d molecular data. arXiv, 2022. (Cited on 22
[56] H. Maron, H. Ben- Hamu, Serviansky, and Lipman. Provably powerful graph networks. NewrlPS,
2019. (Cited on
1 6
[57] H. Maron, H. Ben- Hamu, N. Shamir, and Lipman. Invariant and equivariant graph networks. In ICLR,
2019. (Cited on
19
[58] H. Maron, E. Fetaya, N. Segol, and Y. Lipman. On the universality of invariant networks. In ICML, 2019. (Cited [
[59] H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In ICML, 2020. (Cited on 19
[60] B. D. McKay, M. A. Yirik, and C. Steinbeck. Surge: a fast open-source chemical graph generator. J. Cheminformaties, 2022. (Cited on 31
[61] E. Mena, D. Belanger, W. Linderman, and J. Snoek. Learning latent permutations with gumbel-
sinkhorn networks. In ICLR, 2018. (Cited on (24)
[62] F. Mezzadri. How generate random matrices from the classical compact groups. Notices of the American Mathematical Society, 2006. (Cited on
[63] E. Min, R. Chen, Y. Bian, T. Xu, K. Zhao, W. Huang, P. Zhao, J. Huang. S. Ananiadou, and Y. Rong. Transformer for graphs: An overview from architecture perspective. arXiv, 2022. (Cited on (2)
[64] C. Morris, M. Ritzert, M. Fey, W. L. Hamilton, J. E. Lenssen, G. Rattan, and M. Grohe Weisfeiler and
leman go neural: Higher-order graph neural networks. In AAAI, 2019. (Cited on I 22
[65] C. Mouli and B. Ribeiro. Neural networks for learning counterfactual g-invariances from single environments. In ICLR, 2021. (Cited on
o 22
[66] Müller, M. Galkin, C. Morris, and L. Rampásck. Attending to graph transformers. arXiv, 2023. (Cited
on
[67] R. L. Murphy, B. Srinivasan, V. Rao, and B. Ribeiro. Janossy pooling: Learning deep permutation-
invariant functions variable-size inputs. ICLR, 2019. (Cited on 2 22 28
[68] R. L. Murphy, B. Srinivasan, V. A. Rao, and B. Ribeiro. Relational pooling for graph representations. In ICML, 2019. (Cited on
2 22 28
[69] H. NT and T. Machara. Revisiting graph neural networks: All we have is low-pass filters. arXiv, 2019.
(Cited on 22
[70] K. Oono and Suzuki. Graph neural networks exponentially lose expressive power for node classification.
In ICLR, 2020. (Cited on
[71] F. Paischer, T. Adler, V. Patil, A. Bitto-Nemling, M. Holzleitner, S. Lehner, Eghhal-Zadch, and Hochreiter. History compression via language models reinforcement learning. ICML., 2022. (Cited on (2)
[72] A. Paszke, S. Gross, F. Massa, Lerer, Bradbury, Chanan, T. Killeen, Lin, N. Gimelshein, Antiga, A. Desmaison, A. Kopf, E. Yang. Z. DeVito, Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In NeurIPS, 2019. (Cited on 24 25
[73] Prillo and J. M. Fisenschlos. Softsort: continuous relaxation for the argsort operator. ICMI,, 2020.
(Cited ④
[74] O. Puny, M. Atzmon, E. J. Smith, I. Misra, Grover, H. Ben- and Y. Lipman. Frame averaging
for invariant and equivariant network design. ICLR, 2022. (Cited on 2
26 (7)
[75] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., 2020. (Cited
on 22
[76] Rampásek, M. Galkin,
P. Dwivedi, Luu, G. Wolf, and D. Beaini. Recipe for a general, powerful, scalable graph transformer. In NeurlPS, 2022. (Cited on T0
[77] M. Rath and A. P. Condurache, Boosting deep neural networks with geometrical prior knowledge: A survey, arXiv, 2020. (Cited on (2
[78] S. Ravanbakhsh, J. G. Schneider. and B. Póczos. Equivariance through srrameter-sharing In ICML,
2017. (Cited on 22
13

---

[79] L. F. R. Ribeiro, M. Schmitt, H. Schütze, and I. Gurevych. Investigating pretrained language models for graph-to-text generation. arXiv, 2020. (Cited on 22
[80] D. W. Romero and S. Lohit. Learning partial equivariances from data. In NewrIPS, 2022. (Cited on 22
[81] C. Rommel, T. Moreau, and A. Gramfort. Deep invariant networks with differentiable augmentation layers. In NewrlPS, 2022. (Cited on 222
[82] D. Rothermel, M. Li, T. Rocktäschel, and J. Foerster. Don't sweep your learning rate under the rug:
closer look at cross-modal transfer pretrained transformers. arXiv, 2021. (Cited on 22
[83] A. Sannai, M. Kawano, and W. Kumagai. Equivariant and invariant rcynolds networks. arXiv, 2021. (Cited on
[84] G. Satorras, E. Hoogeboom, and M. Welling. E(n) equivariant graph neural networks. In JCML, 2021.
(Cited on o [2 23)
[85] J. Shen, L. Li, L. M. Dery, C. Staten, M. Khodak, G. Neubig, and A. Talwalkar. Cross-modal fine-tuning: Align then refine. arXiv, 2023. (Cited on 22
[86] W. Shi, Caballero, F. Huszar, Totz, Aitken, R. Bishop, Rueckert, and Z. Wang. Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network. In CVPR, 2016. (Cited on
[87] H. Shirzad, A. Velingker, B. Venkatachalam, D. J. Sutherland, and A. K. Sinop. Exphormer: Sparse transformers graphs. arXiv, 2023. (Cited on [0
[88] B. Srinivasan and B. Ribeiro. On the equivalence between positional node embeddings and structural graph representations. In /CLRR, 2020. (Cited on 22
[89] A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and L. Beyer. How train your vit? data, augmentation, and regularization in vision transformers. arXiv, 2021. (Cited on 2
[90] E. H. Thiede, W. Zhou, and R. Kondor. Autobahn: Automorphism-beee based graph neural nets. In Newr!PS, 2021. (Cited on 30)
[91] N. Thomas, T. E. Smidt,, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and Riley. Tensor field networks: Rotation- and translation -utuiationnquivariiatt neural networks for 3d point clouds. arXiv, 2018. (Cited on [ 22
[92] J. Topping, F. D. Giovanni, B. P. Chamberlain, X. Dong, and M. M. Bronstein. Understanding over- squashing and bottlenecks on graphs via curvature. In ICLR, 2022. (Cited on 22
[93] T. F. A. van der Ouderaa and M.
der Wilk. Learning invariant weights in neural networks. In J. Cussens and K. Zhang, editors, UAI, 2022. (Cited on [2]
[94] E. van der Pol, D. E. Worrall, H. van Hoof, F. A. Oliehoek, and M. Welling. MDP homomorphic networks:
Group symmetries in reinforeement learning. In NeurlPS, 2020. (Cited on 22
[95] M. van der Wilk, M. Bauer, S. T. John, and J. Hensman. Learning invariances using the marginal
likelihood. In S. Bengio, M. Wallach, H. Larochelle, K. Grauman, Cesa- and R. Garnett,
editors, NewrlPS, 2018. (Cited on (2)
[96] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In NeurlPS, 2017. (Cited on 22 25
[97] Velickovic, G. Cucurull, A. Casanova, A. Romero, Liò, and Bengio. Graph attention networks. In
ICLR, 2018. (Cited on 7
[98] R. Winter, F. Noé, and D. Clevert. cermutation-ivvaiann variational autoencoder for graph-level represen- tation learning. In NeuriPS, 2021. (Cited 4 24
[99] Wolf, L. Debut, Sanh, Chaumond, C. Delangue, A. Moi, Cistac, T. Rault, Louf, M. Funtowicz, and J. Brew. Huggingface's 's transformers: State-of-the-art natural language processing. arXiv, 2019.
(Cited on 25
[100] R. Xiongg, Y. Yang, D. He, K. Zheng. S. Zheng, C. Xing, H. Zhang, Y. Lan, L. Wang, and T. Liu. On layer normalization in the transformer architecture. In ICML, 2020. (Cited 22()
[101] K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In ICLR, 2019. (Cited on [
[102]
D. Yarotsky. Universal approximations of invariant maps by neural networks. arXiv, 2018. (Cited
6 7 15
22)
[103] C. Ying. T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T. Liu. Do transformers really perform bad for graph representation? In NewrIPS, 2021. (Cited on 22
[104] C. Yun, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-scquence functions? In ICLR, 2020. (Cited on
[105] B. Zhang, S. Luo, L. Wang, and D. He. Rethinking the expressive power of gnns via graph biconnectivity. arXiv, 2023. (Cited [
14

---

A
Appendix
A.1 Proofs
A.1.1 Proof of (
Theorem
I
If Pi is G equivariant, then deww equivariant for arbitrary fo-
Proof. We prove (ew((i() p2(g') (x) for all X' and g' E G. From Eq. 14 we have:
60 (p P (olor (0=0
(11)
Let us introduce transformed random variable = = E G such that g = g'h. Since the distribution pa, is G equivariant, we can that pu(glpi(g')x 0 (g' (g')x) ( glx)) (h|x). Thus, we can rewrite the above expectation with respect to as follows:
(p)( (p)())))) ((lx)) (eck:))e(((ff66001())))1111 [x)]
p2 (pp)))) (p=(t))fl(p((t))===)
(12)
showing the G equivariance of dew for arbitrary fe-
A.1.2 Proof of Theorem 2 (Section 2.1)
Theorem 2
If Pu is G equivariant and fe is universal approximator, then de a universal approximator of G equivariant functions.
Proof. The proof is inspired by universality proofs of prior symmetrization approaches 102 74
Let v X + y/ be arbitrary G equivariant function. By equivariance of vh, we have:
(xx) (x)
(x) - Ep (alx) (1)) 'x)]]
(alx)) [x(xx) (sbx) (pe(o)fo(p1(0)
(olxx (2(())=(0)> xxx) Ep)(gix) (02(0)((((1(0)
(glx)))
(0=6)))((((0)'=)) - Epu(gix) (2()))((p(()))1 (x)=
(olxx))
(<=(s)((pr(>)) - (a (x))
(13)
As y is finite- diteesiimniiinl,, we can assume that the linear operators in GL(()) are bounded and so is the induced operator norm of group representation [| for g E G. Thus, we have:
[|x(x) - Φe.w (x)]l I max [I ([pp)o)))))) [|a)l [M(p((g) - xx) fe(p((9)~
bEG
<c | (p))))) (p)(r (g)
folpr (9) (2))()))))))))))))))))))))) -x=]
(14)
for some >>> O. If fo is universal approximator, for any compact set K. C X' and any 0, there exists some such that [|ı(() fe(x)] I e for all x E X. Consider the set Ksym = UpcGPi ( where P1 (g) denotes the image of the set K. under linear transformation by (1().. We use the fact that Kaym is also a compact set since is the image of the compact set G x K. under continuous map (g, x) + P1 (g)x. As consequence, for any compact set I. C X' and any c/c > there exists some such that maxgeg (pv (g)x) - fo (pi x]] & c/c for all x E K. Since group is closed under inverse, for any compact set K. S X and any 0, there exists some 0 such that:
[x*x()) 2e S c PP (axx) [(((p1(g)~ - fo(p)(9))
I cmax (p)( fo(p1(9) 9EG
=6,
(15)
for x E K, showing that de. is universal approximator of equivariant functions.
While we have assumed that the group G is compact in the proof, we conjecture that the results can be extended non-compact groups we make an alternative assumption that the distribution Pw (g x) is compactly supported for all x e K. We leave proving this a future work.
15

---

A.1.3 Proof of Theorem 3 (Section 2.1
Theorem 3 qw is G equivariant and p(c) is G invariant under representation of that det p'(g)
1Vg G. the distribution pa (g) x) characterized by qu (x, 6) + p(g) G equivariant.
Proof. We prove (g'gll (g') )x) = (glx) for x E X and g.g' E G. In general, we are interested in obtaining faithful representation p. i.e., such that p(g) is distinct for each We can interpret the probability (g]x, E) as a delta distribution centered at the group representation p(g):
(gl c) 8(p(g) q (x, E)).
(16)
To obtain (g] we marginalize over p(c):
/ pulg)x e)p(e)de
pw(g)x)
b(p()) = 4x, c))p(c)de.
(17)
Let us consider Pa (d''oll)()')x))
Puld ( ((p()')) ( (e)eed e))ple)de.
(18)
Using the G' equivariance JNs we have:
qu (pi = (o
= (x,
(19)
which leads to the following:
(_(g')) I ((plg'9) (ll')qn((x p'g (c))))c
[ ((p(g) qx و' - ()))p())de.
(20)
Note that the second equality follows from invertibility of p(g'). We now introduce change of DE that E = ('g'))e':
variables e' -p(
)=( / 8(p(g) (x, (())(((())))
1 det ('( de'
(21)
With | det (( and G invariance of p(E) which gives p(p'(g') =(EE), we get:
P=(s'slo)(60)))) = 5(p(g) = (x(xxx'))pll (')ee'dd
(g)x),
(22)
showing the G equivariance of "w (g)x).
A.1.4 Proof of Validity for Implemented Equivariant Distributions p.. (Section 2.2
We formally show G cquivariance of the implemented distributions (g(x) presented in Section 2.2
All implementations have a form of noise-outsoureed function [x, E) + p(g) using distribution
~ p(c) and map q which is composed of G equivariant neural network and postprocessing to p(g). From Theorem 3 for G equivariance of (g)x), it is sufficient to show G invariance of p(c) under representation pf such that det pg(g) along with G equivariance of gws which we show below.
Symmetric Group Sn, We recall that (ll( for the symmetric group S,, is implemented as below:
1. Sample node-level noise _xad from i.idd. uniform Unif[0, n).
2. Use GNN obtain node-level scalar features (x,e) + Z F
16

---

3. Assuming Z is tie-free, use argsort [98 to obtain group representation - P3 = =p(g).
Pg eq (Z1', [z) )
(23)
where eq denotes elementwise equality indicator.
We now show the following:
Proposition 3. The proposed distribution Pw (g)x for the symmetric group Sn, is equivariant.
Proof. Given p(e) is elementwise i.i.d., it is Sn invariant under the base representation p'(g) P which satisfies det p,l = 1 from orthogonality. As GNN is S, equivariant, we only need to show Sn equivariance of argsort: : + Pg. This can be shown by transforming with any permutation matrix Py Since sort operator and any row replicated matrices are invariant to P Pg' we have:
eq(Pg g 1sort(P p = eq(Pgr Z1 1sort(Z)
= eq(P, g' 21 T P. 1sor( T.
(24)
Since eq commutes with P Pg,, we have:
eq(P, g. (.PP, Z1 Isort(P, ())
eq 72'1'' [ysoo(())) (Z)
Pg Z1 T Issot(Z)]] )) P.
(25)
showing that argsort is S,, equivariant, ie.,, it maps Pg. P.P for all Pg E Sn. Combining the above, by Theorem 3 the distribution Pw (g(xx) is S,, equivariant.
Orthogonal Group O(n), SO(n) We recall that (glx) for the orthogonal group (n)) or special orthogonal group (n is implemented as follows:
1. Sample noise E Rx from i.i.d. normal NN r7²).
2. Use ((n)/SO(n) equivariant neural network to obtain features (x, c) + Z E R=x"
3. Assuming Z is full-rank, use Gram-Schmin process 41 obtain an orthogonal matrix Z - Q.
4. For the O(n) group, the obtained matrix group representation Q Qg = =
5. For the SO(n) group, use below scale operator obtain group representation Q - Q p(g)
scale:
Q1
Qn
det Q) Qi ... Qn
(26)
We now show the following:
Proposition 4. The proposed distribution pu(g)x) x) for the orthogonal group O(n) is equivariant.
Proof. Without loss of generality, let us omit the scale r for brevity, which gives that each column 6; E R of the noise independently follows multivariate standard normal ~ (N I,). Then, the density p(e;) = (2/2 exp (-1|e:1l//2) /2) is invariant under orthogonal transformation Q since | (Q,,) Qe; e. Q Qe; E; E; = =/E=:1 Therefore, the distribution p(e) is invariant under the base representation p'(g) = Q which satisfies det (g() from orthogonality. As we use an equivariant neural network to obtain Z, and Srram-Schmidt procedure Z + Q, is O(n)
equivariant (Theorem 5 of 41...om by Theorem 3 the distribution Pw (g(g) O(n) equivariant.
Proposition 5. The proposed distribution Pw pg)x)) for special orthogonal group SO(n) is equivariant.
Proof. From the proof of Proposition it follows that the distribution p(E) is invariant under the base representation p'(g) = Q which satisfies det أ'() = due to orthogonality.
we
use an equivariant neural network to obtain Z, and Gram-Schmidt procedure + Q has O(n) equivariance which implies SO(n) equivariance because of SO(z) O(n), we only need to show
17

---

(n equivariance of scale Q + Q This can be done by transforming Q with an orthogonal
Q g' of determinant +1. Since det Q Q)
det Q det(( det(Q), have the following:
scale scale(Q,) Q7*Q
det d Q) (Q7 Q)
( Q)
det(Q) (Q Q)
(Q Qn
(27)
Also, scaling the first column of the product 9 Q with det( Q) equivalent to scaling the first
column of Q with det( Q) then computing the product since (Q, Q)
De Qf Qaj- This gives:
scale (Q = Q det(Q) Q
Q
Q =f scalee())). QQ).
(28)
showing that scale operator is SO(n) equivariant. We also note that scale (Q) gives orthogonal matrix of determinant +1, as it returns Q if det(Q) +1, otherwise (det( Q) -1 since is orthogonal) scales the first column by -1 which flips determinant to +1 while not affecting orthogonality. Combining the above, by Theorem 3 the distribution Pa (g)x) is SO(n) equivariant.
Euclidean Group B(n), SE(n) We recall that, unlike the other groups, we handle the Euclidean group E(n) and special Euclidean group SE(n) at symmetrization level as the translation component T(n) in B(n) = O(n) x T(n) and SE(n) = SO(n) к T(n) is non-compact. This is done as follows:
(x) E (slx T
x1 g fo(g- (x x1 x1]
(29)
where X E R is centroid (mean over channels) of data x E Raxd and distribution Pw is (Sn)SSO(m)
equivariant for B(n)/SE(n) equivariant symmetrization, respectively. We now show the following: Proposition The proposed symmetrization 60 for the Euclidean group E(n) equivariant.
Proof. We prove (g x) = g' ( (x) for all x E X' and g' E E((n). From Eq. 29
we have:
(g x) E x
+g. fo (g' .
(30)
In general, an element of Euclidean group y' E B(n) acts on data x E Radd via g' x where Q E On)) its rotation component and ty' 1² is its translation component 74
Q
x
With
this, the centroid of the transformed data g' x is given as follows:
xx== Qg. + +t, = xxxtttt = Qg
(31)
which leads to the following:
g'
g x1 - 1 Qx.x1 - -t.
= (x - (x-x1
(32)
Above shows that subtracting centroid eliminates the translation component the problem and leaves (n)) equivariance component. Based that, we have the following:
de x) = Ep. p. (x [Q, x1 + tzg 1 + folg (Q (x x1 M)]]
=R (ell""((xxx111 [g' x1 + folg 'og (x x1 x1]))] + +t,
(33)
Note that, inside the expectation, we interpret the rotation component of g' as element of the orthogonal group Onn). Similar as the proof of Theorem I we introduce transformed random variable h g E O(n) that g g' h. Since the distribution p. is Orz) equivariant, we can see
18

---

that ( (xp--((x-) (x--x =
pu g/x x1 ) = (h x1
Thus we can rewrite the above expectation with respect as follows:
00 x)
Bp (h(x-x11 [g'
x1
+ g'h fo((g')) 'd' (x- x1 x1x)) }]] + +tg.1
Ep (Nix (A1x-x17) x)
(g
x1
+ gh so -1 -x- x11))] + ttg.
Qg Bn (N/x (N/x x11x-xx xxxxx1)
[I]] h so(h -1 (x- x1 x1]))] t
Qg (x) + Eyg
= showing the B(n) equivariance of 0e.w. Proposition 7. The proposed symmetrization vu...for special Euclidean group SE(n) is equivariant.
g' (x),
(34)
O
Proof. We prove de (g' x g" (x) for all x X' and g' E +SE(n), in an analogous manner
to the proof of Proposition 6 From Eq. (29). we have:
(g' x) E pu
g x1 + g folg (g' x x- x'x'xxx1 x1]
(35)
In general, an element of special Euclidean group o' E SE(n) acts on data x E Rnxd via g' x
Q , t.. where Q E SO(n) is rotation component and ty' E R" is translation 74
With
this, the centroid of the transformed data g' x is given as follows:
g x -x= 0 x + tz. 1
Q yf
x + to = - tt''
(36)
which leads to the following:
g' x x1 T Q tg/ 1
tg'
Q (x
x1 ))
(37)
Similar as Proposition 6 subtracting centroid only leaves SO(n) component. We then have:
0o x) F (s (x-x1 (x-xx Q x1 + t.g. IT g. folg (Q7 (x x1 M>]]
= F (slo" )) [o' x1 + Solo "g (x-x1 x]] + 1,,
(38)
where, inside the expectation, we interpret the rotation component of g' as an element of the special orthogonal group SO(n). Similar in Theorem | introduce h -g' SO(n) that g g'h. As the distribution Pw is SO(n) equivariant, we have (glg'' (x x1
Pw 1 gl''' 'g' (x-xI)
x1 )) pw (h|x - x1 ). We then rewrite the expectation with respect to h:
(0 x) = Pe (11x-x17) [g' x1 T
+ g'h. (o((g'))= (x x1 xI D)] + tt.. IT
_P (1)x-x1
[g' x1 T + g'h fo(h- -1. 1..1..1x-- -x--xI xx-x11m
x1 ))] + +t,1
O.t.Ba. x1T)
[x1] T
+ h so -xx x1 T) })} + 1
Q (0 + tg 1
= g' (x),
(39) L
showing the SE(n) equivariance of do.w.
Product Group x K For the product group H x K, we assume that the base representation for
each element y (h, k) is given pair of representations p(g) (p(h), p(k)). Without loss of generality, we further assume that the representation p(g) can be expressed as the Kronecker product p(g) p(h) & p(k) that acts flattened data vec(x)
+ vee (p(g vec(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 (S,, x (Sdd)/Soaf)) D
and sets of symmetric elements (S,, x H) in general [59] For
Thii is after handling the translation component of the Euclidean group B(d) (d)))E(d) in Eq. 29
19

---

Gl
Gan
Geth
Gult
Guil
Gett
Dx27
Gush
Group
Orbits Gxg
Distribution po(s)x)
Frame F(x)
Canonicalizer Co(x)
Figure 3: Visual illustration of the symmetrization methods based on probabilities assigned upon the partitioning of the group G into orbits Gxg. Note that, while we use concentric circles of different
perimeters to illustrate each orbit, all orbits actually have an identical cardinality [Gxg] |Gx]-
example, for the group S, O(d) on data x E Raxd, an element g = (h, k) has representation P(g) p(h) @ p(k) E mndxnd combined from permutation p(h) E R=*= and rotation p(k) e Rddd which acts by x vee (elg) (p(g)ve(xx)) or more simply p(h)xp(k))
Now we recall that the p.. (g(x)) for the product group x K implemented follows:
Sample noise E & from i.idd normal (0, r7²) such that p(E) is invariant under representations of H and K that satisfy det (h)) = and det o' pk(k) = 1, respectively. For example, for Sn x O(d). the noise E " that follows i.id. normal 170,,²²) is invariant under base representations of both S,, and O(d) which orthogonal.
Use a H K equivariant neural network obtain features (x, E) (ZH Zx) where ZH K invariant and Z is H invariant. For example, for Sn x O(d), we expect node-level scalar features Zs.. E R" to be O(d) invariant and global rotary features Zora) E Rddd to S,, invariant.
3. Apply postprocessing for H and K groups onto Zμ and Z respectively to obtain representations Z + p(h) and Z + plk) H and K groups respectively. For example, for S,, O(d), we use argsort in Eq. 23 to obtain Zs., + p(h) and Gram-Schmidt process to obtain Zoça) + p(k).
Combine the representations p(g) (p(h), p (k)) obtain a representation for the H x K group.
We now show the following: Proposition 8. The proposed distribution Pw (g)g) for the product group H K is equivariant.
Proof. By assumption, p(E) is invariant under representations of Hl and K that satisfy det p')h) = and det pl()) p(k <) for all E, H, k E K gives p(k: h e) p(k (h- Kronecker product of matrices of determinant gives a matrix of determinant 1. Furthermore, the map (x, e) - (p(h), p(k)) p(g) is overall H K equivariant, since input transformed with ( is first mapped by the equivariant neural network as (g' x, o' .E) + (h'- -Zu, K Zk), then postprocessed
respectively. This implies x K invariance as well, since p(e) = p(h-
p(h
p(e), and
g' = (h', (h',k')
(h' Bz K' ZR) + (p(h') (h)), p(k') ((k)) (p(h'), ((k)) (p(h), ((g)))()) Combining the above, by Theorem the distribution (g) x) is x equivariant. L.
A.1.5 Proof of Proposition T Jnd Proposition 2 (Section 2.4
Before proceeding to proofs, we recall that the stabilizer subgroup Gy group G for x is defined as {g' G g' x x} and acts on given group element G through left multiplication g g'g. For some g E G. by Gx9 we denote its orbit under the action by Gx. i.e., the set of elements G to which g be moved by the action of elements g' Gx. Importantly, can show the following:
Property 1. Any group G is union of disjoint orbits Gxg of equal cardinality.
Proof. Let us consider the equivalence relation ~ on G induced by the action of the stabilizer Gx- defined as
< E Gxg. The orbits Gx9 are the equivalence classes under this relation, and the set of all orbits of G under the action of Gx forms partition of G (i.., the quotient G/ )x Furthermore, since Gx S G and right multiplication by some
a faithful action G on itself, have [Gxgl |Gx] for all g G, which shows that all orbits Gxg have equal cardinality.
D
20

---

The partition of group G into disjoint orbits Gx9 is illustrated in the first and second panel of Figure We now show the following:
Property G equivariant (glx) assigns identical probability to all elements on each orbit Gx9-
Proof. With equivariance, we have p P (g'ggg glg'-l).'...)) Since g' x for all g' E Gx. we have pg x) - Pw glx) for all g' E Gxi all elements on orbit Gxy have an identical probability.
Property 2 characterizes probability distributions over G that can be expressed with pa (g)x), which we illustrate in the third panel of Figure Intuitively, Pa, x) assigns constant probability densities over each of the orbit Gx that partitions G as shown in Property | We now prove Proposition and Proposition 2 by showing that (g can become frame and canonicalizer as special cases:
Proposition Probabilistic symmetrization with G equivariant distribution Pu(g)x x) can become frame averaging (74) by assigning uniform density set of orbits Gx9 for some group elements g.
Proof. A frame is defined as a set-valued function F: X L + 2 \. 0 that satisfies G equivariance F(g- x) =gF(xx 774 For some frame F, frame averaging is defined follows:
Σ gEF(X)
Solg
(40)
[(()]
which can be equivalently written as the below expectation:
E Eg-)Um((F()) [g fo(g .x]
(41)
From Theorem of 744 we have that Fxx) is a disjoint union of equal size orbits Gxg. Therefore, Unif(F(x)) uniform probability distribution over the union of the orbits. This can expressed by G' equivariant distribution Pu (g)x) by assigning identical probability over all orbits in the frame F and zero probability all orbits Therefore, probabilistic symmetrization can become frame averaging.
the frame (illustrated in the fourth panel of Figure B
Proposition 2 Probabilistic symmetrization with G equivariant distribution (u(g)>) can become canonicalization [41] by assigning uniform density single orbit Gx9 of some group element g.
Proof. A canonicalizer is defined as a (possibly stochastic) parameterized map C. X + G that satisfies relaxed G equivariance Culg x) (x)) for some g' E Gx 41). For some canonicalizer Cw. canonicalization is defined as follows:
g folg x)),
x), 9 Cx().
(42)
From relaxed equivariance, we have (x) = gxC for some g' E Gx. valid choice for the canonicalizer Cu is stochastic map that samples from the uniform distribution over a frame Cu (x) ~ (E, where the frame is assumed to always provide single orbit Fx(x) = Gxg. In this case, canonicalization is equivalent to I-sample estimation of the below expectation:
Fgg-U [g fo(g-1.x)) x)]
(43)
Furthermore, uniform distribution over the oingle-orbit frame (ni((F.(x)) can be expressed by G equivariant distribution (g)xx) by assigning nonzero probability to the single orbit Gx9 and assigning zero probability to the rest (illustrated in the last panel of Figure Therefore, probabilistic symmetrization can become canonicalization.
A.2 Extended Related Work (Continued from Section 2.4
Our work draws inspiration from extensive array of prior research, ranging from equivariant architectures and symmetrization to general-purpose deep learning with transformers. This section outlines comprehensive review of these fields, spotlighting ideas specifically relevant to our work.
21

---

Equivariant Architectures Equivariant architectures. defined by the group equivariance of their building blocks, have been a prominent approach for equivariant deep learning [2] T0]. 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] 157 g and geometric
graphs [2 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 [01 (56] 64 [05 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 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 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 B5 20 or a transformer [[04] 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 [6
68
88
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-bans based estimation for certain worst-case inputs. Canonicalization [41] utilizes 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- pre-trained models for small finite groups 4 3. and to our knowledge, we are the first to investigate symmetrization of 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 do- and proven its ability to learn diverse modalities
mains 196 24
13 75
[39
38 or transfer knowledge
across them 85|
54
25 71 52). Although transformer-style architectures have been developed
for symmetric data modalities like sets [5], graphs [103]
45
43
50
66
hypergraphs [1 44).
and geometric graphs 53). they often require specific architectural modifications to achieve equiv- ariance 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 ( x) on a group for symmetrization of form (x) = Eyg0 one may find similarity distributions over data augmentations (e.g., (g)) for a similar symmetrization. The key difference that, while these approaches aim 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 (x) given the known symmetry group G of data (e.g., G S,, for graphs). Because of this, the symmetrizing distribution has to be designed differently. In our case, we parameterize the distribution ( xx) itself to be equivariant to a specific given group G. while for augmentation learning approaches, the distribution (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 Pa (glxx) per input data while keeping the symmetrized function da (x) exactly G' equivariant, while augmentation
[o .x) and learns it from data,
Augerino 7 and related approaches 777 8T
93
80 37 that learn
22

---

Table Overview of the datasets.
Dataset
Symmetry
Domain
Task
Feat. (dim)
GRAPH8c
Graph Separation
EXP
S, Invariant
Graph Isomorphism
Adj. (1)
EXP-classify
Graph Classification
n-body
S. x H(3) Hquivariant
Physics
Position Regression
Pos. (3) Vel. (3) Charge (1)
PATTERN
S, Equivariant
Mathematicul Modeling
Node Classificution
Rand. Node Attr. (3) + Adj. (1)
Peptides-func
Graph Classification
Peptides-struct
S, Invariant
Chemistry
Graph Regression
Atom (9) Bond (3) + Adj. (1)
PCQM-Contact
S, 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
EXP
1,200
44.44
110.21
EXP-clussify
n--ody
7,000
Fully Connected
PATTERN
14,000
188
117.47
4749.15
Peptides func
15,535
150.94
307.30
Peptides-struct
529,434
30.14
61.09
learning does not guarantee equivariance for given group in general and often has to reduce to trivial group averaging pa Unif(G) to be exactly G equivariant, and (2) while augmentation learning has to employ regularization 7 model selection [37 to prevent collapse 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)
We provide details of the datasets and models used in experiments in Section The details of the can be found in Table and Table
datasets from the original papers [2] 27 28 31 84]
A.3.1 Implementation Details of p. for Symmetric Group S,, (Section 3.1 3.3 3.4
In all experiments regarding the symmetric group Sn, we implement the S,, equivariant distribution Pa x). i.e., qw (x, + P, as a 3-layer GIN with 64 hidden dimensions [101] that has around 25k parameters. Specifically, given a graph x with node features X E xgn and adjacency matrix A E RR we first augment virtual node [2]] which is connected to all nodes facilitate global interaction while retaining S, equivariance, as follows:
A
1
= [Xv], A'
(44)
1
0
where the feature of the virtual node v E Rdd is a trainable parameter. Then, we prepared the input node features E ()x+1)xxde the GIN H X' + c where the noise c E R(*+1)xd, is i.i.d. sampled from Unif(O, 7/ with scale hyperparameter 7). Then, we employ following 3-layer GIN with 64 hidden dimensions to obtain processed node features H' E {(*+1)x1,
H' GINConve4.64. o GINConv64.64.6 o GINConW M.NConwd,,64444o 64.64 (H),
(45)
where each GINConv .d,d,, computes below with a two-layer elementwise MLP Rexd + x with hidden dimension dg. ReLU activation, batch normalization, and trained scalar e:
H MLP((A' + (1+ H))
(46)
We do utilize edge attributes in equivariant distribution pv. while utilize them in base model fo.
23

---

Then, from the processed node features H' [x+++)xx we finally obtain the features Z E R" for postprocessing by discarding the feature of the virtual node. Then, postprocessing into permutation matrix is done with argsort: Z + Pg E T as in Eq.
Training To backpropagate through Po for end-to-end training of pa (g(x), we use straight-through gradient estimator [ with an approximate permutation matrix P. R Pg
8 For this, we first apply L2
normalization Z + Z and use the below differentiable relaxation of the argsort operator [6 33 98
P3 = 5(-111111 1sort(Z) (Z) 1/T),
(47)
where S(//T) is Sinkhorn operator [6]] with temperature hyperparameter 7 E R, that performs elementwise exponential followed by iterative normalization of rows and columns. Following [6]] we use 20 Sinkhorn iterations which worked robustly in all our experiments. For the correctness of straight-through gradients, it is desired that P. closely approximates the real permutation matrix P, during training. For this, we choose the temperature T 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 Pg with strength of 0.1 in all experiments. The S, equivariance of the relaxed argsort Z + P, can be shown in similar way to Proposition 3 from the fact that elementwise subtraction, absolute, scaling by -1 -1/T, exponential, and iterative normalization of rows and columns all commute with Pg E Sn.
A.3.2 Implementation Details of Pa for Product Group S,, x E(3) (Section 3.2
In our n--ody experiment the product group S,, E(3), we implement the S,, x 0(3) equivariant distribution (g - x1 ). i.e., qw (x - x1 e)) + (P,, g Q₂) based on 2-layer Vector Neurons version of DGCNN with 96 hidden dimensions 23 that has around 7k parameters. Due to the architecture's 's complexity, and guide the readers to the original paper [23] for further architectural details. In a high-level, the Vector Neurons receives position P E RR*3 and velocity V R=x3 of the zero-centered input x1 T with noises E1, E2 E Rax3 ii.d. sampled from normal (O, r7²)) with scale hyperparameter 1). and produces features Hs, E Rnx33xd and H0,3) E IRxx3xd2 with d1 = 1 and d2 = 3 as follows:
focus describing input and output of the network and postprocessing,
x-x1
Hs.. Ho(3) = VN-DGCNN(P + E1, V +2))
(48)
Then, apply 0(3) invariant pooling Hs., and S,, invariant pooling on H0r3)- both supported as a part of 23 to obtain features for postprocessing Zs,, E =x and Zo(3) E R3x3 respectively:
Zs.- = Poolo(3) (Hs.), Zo(3) Pools, (Hor3)).
(49)
Then, postprocessing with argsort: Zs.. + P. E Rux= and Gram-Schmidt orthogonalization Zo(3) + Q2 E RR33 is performed identically as described in the main text (Section 2.2). For the straight-through gradient estimation of the argsort operator, we relaxed argsort described in Appendix A with the only difference of using the temperature 7 = 0.1.
A.3.3 Graph Isomorphism Learning with MLP (Section 3.1
Base Model fo For EXP and EXP-classify, the model is given adjacency matrix A 166x66 and binary node features X E R&4 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 A E R 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 fo R64x64464 + R flattened and con- catenated adjacency matrix and node features, with an identical architecture to other symmetrization baselines (MLP-GA and MLP-FA (74)) as in below:
fo = FC FC10.2018 FC2018, 1096 FC1096,2018 2018 FC2048, 1160-
(50)
where FC RRd + R denotes a fully-connecte 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 so 8 10 that takes flattened adjacency produce 10-dimensional output 774
fo = FC10.61 o FC61,128 o FC128.61-
(51)
In PyTocc (72). simply replace P, with (P, P.)) detach() + P. during forward passes.
24

---

Training For Ela-classify. we train our models with binary cross cross-entropy loss using Adam optimizer [46] with batch size 100 and learning rate le-3 for 2,000 epochs, which takes around 30 minutes on a single RTX 3090 GPU with 24GB using PyTorch 722. 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 equívariant distribution Pw: we use noise scale 1) = 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 fa The model is given zero-centered position P E R and velocity V E 5R5x of particles at a time point with pairwise charge difference C E IR555 and squared distance D E R5x5 We set the prediction target difference of position AP R5x3 after certain time.
For the base model, we use 8-layer transformer encoder fo I225x8 - 25xx that operates on sequences of length 25 with dimension 8. At each prediction, we first organize the input into single tensor xx by placing P and V on the diagonals of C and D, and then turn the tensor into sequence of 25 tokens 2 by flattening the first two axes. Analogously, we organize the output of the model into a tensor E R5x6xx3 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 ID positional encoding, and an MLP prediction head with I hidden layer. The model has around 208k trainable parameters, around 2.3 x compared to the GNN backbones of E(3) symmetrization baselines in the benchmark (GNN-FA and
with 92k parameters.
Training We train our models with MSE loss using Adam optimizer 46 with batch size 100 and learning rate le-3 for 10,000 epochs, which takes around 8.5 hours single RTX 3090 GPU with 24GB using Py TyTorch (72). We use weight decay with strength le- 12 and dropout on the distribution Pa with probability 0.08. For the equivariant distribution Pws use noise scale r) = 1.
A.3.5 Graph Pattern Recognition with Vision Transformer (Section 3.3
Base Model Jo The model is given adjacency matrix A E R188x188 and node features X E IR 188883 zero-padded maximal 188 nodes. The prediction target is node classification logits Y R188x2
For the base model, we use a transformer with an identical architecture to ViT-Base [26] that operates on 224x 224 images with 16x 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 [(14x14)x76s + ((14x14)x766 on sequences of (224/16) x (224/16) = 14 x 14 tokens. Then, we use the following as the base model fa (A, + Y:
fe = detokenize (transformer (tokenize A, X))),
(52)
where, for tokenize R 18 x R 188x3 + [114x14)x768 we organize the input into a single tensor E IR 188x 188x4 by placing X on the diagonals of A and apply 2D convolution with kernel size + IR x88x2 we apply transposed 2D convolution
and stride 14, and for detokenize : [[I4x14)xX with kernel size and stride 14 obtain a tensor E 1188x18822 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 Ic-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 Py PyTorch Lightning [29]. For the equivariant distribution Pa we use noise scale 1) = 1 and dropout with probability 0.1. For probabilistic symmetrization that involves sampling-bast 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 fo For Peptides-func/sttuue the model is given adjacency matrix A E RM1xx444 node features X E 114611 and edge features E e R44x111xx zero-padded to maximal 444 nodes. The prediction target is binary classification logits Y e IR for Peptides-fung and regression targets
25

---

Table Supplementary results for S, invariant graph separation with S,, symmetrized GIN-ID base function. Baseline scores for GIN-ID-GA ID-GA and GIN-ID- FA are taken from 77..
method
arch.
sym.
GRAPHSc
EXP
EXP-classity 1
GIN-ID-GA
S,
0
0
50%
GIN-ID-FA
Sn 0
0
100%
MIN-ID-Canonicn
S,
84%
GIN-ID-PS (Ours)
S.
100%
Table Supplementary results for S., x E(3) equivariant n-hody with E(3) symmetrized GNN base
function. Baseline scores for GNN-FA and GNN-Canonical. are from 74 and [11] respectively.
method
arch.
sym.
Position MSE
GNN-FA
Sn
D(3)
0.0057
GNN- Canonical.
S,,
E(3)
0.0043
ONN-Canonical (Reproduced)
S,
E(3)
0.00457
GNN- ONN-GA
S, E(3)
0.00408 + 0.00002
GNN-PS (Ours)
Sn E(3)
0.00386 + 0.00001
Y E R for Peptides-struct. For PCQM-Contact, the model given adjacency matrix A R33x53 node features X E 553668 and edge features E E 1233x53x6 zero- parro--pdeed to maximal 53 nodes. The prediction target is binary edge classification logit Y R53x53x1
For the base model, we use a transformer with an identical architecture to VīT-Base that operates on 14 14 tokens, same as Appendix A.3.5 For Peptides-func and Peptides-struct, use the following as the base model fe (A, X, E) + Y:
fo(A, X, E) detokenize(caa) (transformer (tokenizezo A, E) tokenize:D(X)) (X)))
(53)
where tokenizegp : [444x444(147) 1(14x14)x768 is 2D convolution with kernel size and stride 32, tokenizenp R444X64 + R 196x 768 is ID convolution with kernel size and stride 3, and
detokenizeteaa performs linear projection of the global [cls] token (26) the target dimensionality.
For PCQM-Contact we use the following base model fo (A, X, E) Y:
fe(A. E) detokenizezp (transformer (tokenizezo A, E) tokenize (X))),
(54)
where tokenizegp 553x53x(1+6) + (14x14))x76 is 2D convolution with kernel size and stride 4, tokenizenp R 53 R 796 is ID convolution with kernel size and stride 1, and detokenize [114114)8768 +
1 53x53x1 is transposed 2D convolution with kernel size and stride 4.
Training We train our models with cross-entropy for classification and LI loss for regression using AdamW [53] optimizer with batch size 128, learning rate lc-5 except for PCQM-Contac 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 Py Torch Lightning [29]. For the equivariant distribution Par we use noise scale rj = 1, and use dropout with probability except for PCQM-Contace 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 but could not be included in the main text due to space constraints.
A.4.1 Graph Isomorphism Learning (Section 3.1
we mainly experimented for S,, symmetrization of an MLP. Here, we provide supplementary results on S,, symmetrization of a GIN base model with node identifiers, following (74). The results can found in Table In accordance with Section
In our experiments on graph isomorphism learning in Section 3.1
3.1 our approach successfully performs S,, symmetrization of GIN-ID.
26

---

Test accuracy (%)
100
90
1
10
50
Train sample size
Figure 4: Test accuracy of MLP fe symmetrized by equivariant distribution (g]x) trained on EXP-c classify dataset across range of training sample sizes. Inference sample size is set to 10.
Entropy sampled permutation
7.0
6.5
6.0
5.5
5.0
10
Train sample size
Figure 5: Row- and column- wise entropy of aggregated permutation matrices Pg (g)x)) after trained on EXP-classify dataset across range of training sample sizes. Dashed line indicates entropy measured with random permutation matrices from Uni
A.4.2 Particle Dynamics Learning (Section
3.2
experimented for S,, x E(3) symmetrization using a ID sequence transformer architecture which has 2.3x parameters compared to basclines. To comparison against baselines in literature, E(3) symmetrization of S, equivariant GNN base model that is widely used in literature [74, 4T]. We faithfully follow [74] 41 the experimental setups including training hyperparameters and the configuration of GNN base model, and only add B(3) equivariant distribution (g x1 ),
In our experiments n-body dataset in Section 3.2
provide parameter-matee
apply approach for
i.e, Q. (x x1 .6) + Q 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 g In accordance with the results in Section 3.2 our approach outperforms other symmetrization approaches and achieves new state-of- oftth--r----rttt-- of 0.00386 MSE.
A.4.3 Effect of Sample Size Training and Inference
In this section, we provide additional analysis how the sample size for estimation of symmetrized function (Eq. affects training and inference. We use the experimental setup of EXP-classify (Section 3.T S invariance) and analyze the behavior of MLP-PS with identical initialization and hyperparameters. only controlling sample sizes E {1, 2, 5, 10, 20, 50} for training. Specifically, we analyze (1) variance of permutation matrices P 9g Pw measured indirectly by the entropy of their aggregation P = P.P/N in Section 3.1 (2) sample variance of the unbiased estimator
g fo x) of the symmetrized function de x as in Eq. and (3) sample mean and variance
27

---

Output varfance
Loss variance
Loss mean
10
10
10a
Inference sample sizee
10:
10
10:4
10
10-3
10-4
10a5
10-3
10-1
10
10
Train sample size
Train sample size
Train samplie size
Figure 6: Variance of estimation of MLP fo symmetrized by equivariant distribution (glx) and
trained on EXP-classify dataset for a range of training and inference sample sizes.
of the estimated task loss Cly,g g fo(g-1. are repeated 100 times and averaged over the inputs and labels (x,y) y) of the validation dataset.
x)) where L is binary cross entropy. All measurements
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 D.. Second, models trained with smaller sample sizes tend to be more sample efficient, i.e., tend to perform a lower variance estimation. Their distribution Pw (g(x) tend to learn more vow-variance permutations (Figure 5 and the models tend to learn low-variance estimation of output and loss (left and center panels of Figure o). This indicates that small sample size may serve as regularizer that encourages lower variance of the estimator. However, this regularization effect not always beneficial in terms of task loss (right panel of Figure 6). as training sample size 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; small sample size takes more training iterations to converge, but serves as 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 consistent with the thcoretical claims in literature 67 68 on the sampling based training of symmetrized models, which we reprise here. When training the symmetrized model dw. (x) in Eq. (4). we cannot directly observe (o,, but observe samples of actually optimizing during the sampling-bawe -based training. Based on [67 68], it turns out that minimizing convex loss function L on the estimated output g fo bound to the true objective on the symmetrized output de ~ (x). This is because our estimation is no longer unbiased when computing loss, as we have the following from Jensen's inequality:
its unbiased estimator g fo(g -1 x). Thus, it can be questionable what objective
x) equivalent minimizing an upper
(p [A(y, g. folg-1 1 x))] 2 L(y, Bp x))) fo(g-
(x))))))
(55)
That is, minimizing the sampling- based loss minimizing an upper-bound surrogate to the true objective. It has been claimed that optimizing this upper bound has an implicit low-variance reg- ularization effect (67 68) which is consistent with our observations. This also roughly explains why our distribution P.. (g(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.
A.4.4 Additional Comparison to Group Averaging
In this section, we provide additional analysis of our approach in comparison to sampling-based group averaging in terms of sample variance and convergence. We use the experimental setup of EXP-classify (Section 3.1 S, invariance) and experiment with MLP-PS and MLP-GA.
We first analyze whether using the equivariant distribution (g for symmetrization offers a lower variance estimation, i.e., a better sample efficiency, compared to group averaging with Gn). This supplements the results in Figure that (g(x) learns to produce low-variance permutations compared to Uni Specifically, we fix a randomly initialized MLP So and symmetrize it using our approach and group averaging. We then measure (1) the sample variance of the unbiased estimator
28

---

Oulput variance
Loss variance
GA
10
GA
PS (Ours)
PS (Ours)
10-1
10-7-
10-4
10-h
10-1
200
100 200
Inference sample size
Inference sample size
Figure 7: Sample variance of output g fo(g- x) (left) and loss L(y.g Soly-t x)) (right) of an identical MLP so symmetrized by equivariant distribution (PS) and uniform distribution (GA).
Gradient direction norm
10
10a
10 -1
GA
10-2
P5 (Ours)
10 3
10
100
1000
2000
Train epoch
Figure 8: Norm of full gradient over training epochs with respect to the parameters of an identically initialized MLP symmetrized by equivariant distribution (PS) and uniform distribution (GA).
g. fo(g-1 of the symmetrized function (Eq. ① and Eq. (4)) and (2) the sample variance of the estimated task loss L(y.9 "o(o-'-x) where L is binary cross entropy. All measurements are repeated 100 times and averaged over the inputs and labels (x.y) of the validation dataset. The results are in Figure [ showing that symmetrization with equivariant distribution p. (glx) consistently offers lower variance estimation than group averaging across inference sample sizes.
In addition, we analyze whether the equivariant distribution (w |x) for symmetrization offers more stable gradients for the base function fo during training compared to group averaging, conjectured in Section 2 For this, we fix a randomly initialized MLP fa and symmetrize it using our approach and group averaging. For every few training epochs, we measure the full gradient of the task loss over the entire training dataset with respect to the parameters of the base MLP fa. This averages out the variance from individual data points and provides the net direction of the gradient on the base function offered by Pw (g)x) or Gn. The results are in Figure [] showing that that symmetrization with equivariant distribution Pw. (g)x) offers a consistently larger magnitude of the net gradient, while group averaging with Uni((G) leads to near-zen zero net gradients. This indicates, for Unif(G). the gradients from each training data instances are oriented in a largely divergent manner and therefore the training signal is collectively not very informative, while using (olx)) for symmetrization leads to more consistent gradient across training data instances, i.e., it offers more stable training signal.
A.4.5 Additional Comparison to Canonicalization
In this section, we provide additional analysis of our approach in comparison to canonicalization 41] that uses a single group element from an equivariant canonicalizer C. x p(g). The main claim that there always exist certain inputs that canonicalization fails to guarantee exact G equivariance, while our approach guarantees equivariance for all inputs in expectation as in Theorem [
29

---

0,05
-0000
5,00
-1015
52
-0023
Inner-symmetric nodes
(b) MLP-PS (Ours)
MLP-Canonical.
Figure 9: A graph x with stabilizer subgroup Gx = Sg.
-101
45
-6662
(a) Inner-symmetric nodes
(b) MLP-PS (Ours)
(c) MLP-Canonical.
Figure 10: A graph x with stabilizer subgroup Gx Sy x S2-
DVRES
2.01
0O00
-863
(a) Inner-symmetric nodes
(b) MLP-PS (Ours)
(c) MLP MLP-Canonical.
Figure 11: A graph x with stabilizer subgroup Gx D4.
More specifically, let us recall the definition of G equivariant canonicalizer from 41 A canonicalizer Cu equivariant (g x) = p(g) (x) for g G and x E X. Consider an input x which has non-trivial stabilizer Gx h G h x x, i.e., has inner symmetries. It can be shown that equivariant canonicalizers are ill-defined for these inputs. To see this, let 91 ghi and g2 ghzz for some G' and any h1, h2 E Gx where h * l22 Then we have Cu(gi x) = Ca x)
(g x) = Cu (gh2 x) = Cu(g2 x), implying that p(g Cx) = p(g2) (x). Since 91 * g2- this contradicts the group axiom, and thus equivariant canonicalizer cannot exist for inputs with non-trivial mnner-yymmetries. To handle all inputs, canonicalization 41 adopts relaxed cquivariance: canonicalizer C. satisfies relaxed equivariance if Cg x) = p(gh)Cu (x) up arbitrary action from the stabilizer Gx- As a result, the symmetrization Φ9... (x)
folg x) performed using relaxed canonicalizer C., only guarantees relaxed equivariance Φ to arbitrary action from the stabilizer h E Gx (Theorem A.2 of (4T).. In other words, canonicalization does guarantee equivariance for inputs with inner symmetries.
gh ) (x) up
To visually demonstrate this, we perform minimal experiment using several graphs x with non- trivial stabilizers Gx. i.e., inner symmetries, taken from [90]. We fix randomly initialized MLP
30

---

fo R - R" and symmetrize it using our approach and canonicalization. When symmetrized, the MLPis is expected to provide scalar embedding of each node, which we color-code for visualization. The results are in Figures 9 TO [ For each graph, we illustrate three panels: the leftmost one illustrates the color-coding of the inner symmetry of nodes (automorphism), the middle one illustrates node embedding from MLP-PS, and the rightmost one illustrates embedding from ML.PCCanoncal a method is G cquivariant, it is expected to give identical embeddings for automorphic nodes, since an equivariant model cannot distinguish them in principle (88) As in the Figures 9 10 and II in the presence of inner symmetry (left panels), MLP with probabilistic symmetrization (middle panels) achieves G equivariance and produces close embeddings for automorphic nodes. However, the same MLP with canonicalization produces relatively unstructured embeddings (right panels). The result illustrates potential advantage of probabilistic symmetrization over canonicalization when learning data with inner symmetries, which is often found in applications such as molecular graphs [60].
A.5 Limitations and Broader Impacts (Continued from Section
While the equivariance, universality, simplicity, and scalability of our approach offers potential for positive impact for deep learning for chemistry, biology, physics, and mathematics, it also has limitations and potential negative impacts. The main limitation of our work is that it trades off certain desirable traits in equivariant deep learning in favor of achieving architecture agnostic equivariance. For example, (1) our approach is less interpretable compared to equivariant architectures due to less structured computations in the base model, (2) our approach is presumably less parameter and data efficient compared equivariant architectures due to less imposed prior knowledge parameterization, and (3) our approach is expected to be challenged when input size generalization is required, partially because the maximum input size has to be specified in advance. Another genuine weakness compared to canonicalization is that, our method is stochastic and therefore incurs O(N) cost when using N samples for estimation. These limitations might lead to potential negative environmental impacts, since less interpretability and lower efficiency implies higher reliance on larger models with more computation cost. We acknowledge the aforementioned limitations and impacts of our work, and will make effort to address them in follow-up research. For example, we believe data efficiency of our approach could improve with pretrained knowledge transfer from other domains, as it would imposc a strong prior on the hypothesis space and may work similarly architectural priors that benefit data efficiency. Also, for the sampling cost, since the sampling completely parallelizable and analogous to using larger batch size, we believe it can be overcome to some degree by leveraging parallel computing techniques developed for scaling batch size.
31


Reviews:
[{"review_id": "4g6RNvQtxn", "reviewer": "Reviewer_SgC2", "summary": "This work presents a generic approach to symmetrize wide range of base models. Instead of relying on uniform average sampling, the approach introduces a trainable transformation to model the group equivariant distribution. The framework theoretically encompasses existing approaches such as group averaging, frame averaging, and canonical function approaches. Furthermore, the author provides empirical evidence demonstrating competitive performance. \n\n\n", "questions": "1. Is it feasible to explore the convergence of the probabilistic equivariant distribution $p_{\\omega}$ and assess its dissimilarity to a uniform distribution? This analysis would provide insights into the nature of the mechanism. \n\n2. What is the sensitivity of the proposed method to the architecture scales of MLP and transformer? Can the efficacy of the proposed method be maintained when the base models have a large number of parameters?\n\n3. Is there a general framework or guideline for designing the $G$ equivariant neural network $q_{\\omega}$? How significantly does the expressivity of $q_{\\omega}$ impact the performance?\n\n4. Why some experimental comparisons between [1-3] are missing? For instance, in ``3.1 Graph Isomorphism Learning with MLP'', what limits the comparison with [3]? \n\n\n[1] S. Basu, P. Sattigeri, K. N. Ramamurthy, V. Chenthamarakshan, K. R. Varshney, L. R. Varshney, and P. Das. Equi-tuning: Group equivariant fine-tuning of pretrained models\n\n[2] S. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonicalization functions\n\n[3] O. Puny, M. Atzmon, E. J. Smith, I. Misra, A. Grover, H. Ben-Hamu, and Y. Lipman. Frame averaging for\ninvariant and equivariant network design\n\n", "limitations": "The authors have discussed limitations and potential societal impact associated with this work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method demonstrates the ability to symmetrize architectures in a group-agnostic manner for general purposes, supported by sound theory.\n\n2. The theoretical analysis presented in this work showcases the ability of the proposed approach to encompass interesting literature that assigns distributions on the compact group $G$.\n\n3. The paper is written in a reader-friendly manner, making it easy to follow.\n\n\n\n\n", "weaknesses": "1. A thorough discussion of the empirical and theoretical advantages and disadvantages of literatures, specifically group average, frame average, and canonicalization, would greatly enhance our understanding, given the shared problem setup with the proposed method.\n\n\n2. I have reservations about the assertion made in Line 3 that  \n>we use an arbitrary base model (such as an MLP or a transformer)...\n\nWhile the argument of this work suggests that the base model $f_{\\theta}$ can be arbitrary, in the experiments, only MLP and transformer architectures were explored and evaluated.\n\n\n3. Typically [1,2], showcasing improved performance on image classification problems is one of the common applications used to demonstrate the effectiveness of a newly proposed equivariant network. It is encouraged to include such experiments in order to comprehensively validate and demonstrate the efficacy of the proposed equivariant network.\n\n\n[1] S. Basu, P. Sattigeri, K. N. Ramamurthy, V. Chenthamarakshan, K. R. Varshney, L. R. Varshney, and P. Das. Equi-tuning: Group equivariant fine-tuning of pretrained models\n\n[2] S. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonicalization functions\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. A thorough discussion of advantages and disadvantages of symmetrization approaches would enhance our understanding.\n\nA1. We will add related work section with added thorough comparison. A tentative discussion is in Table 3. Please note that we performed theoretical analysis in Section 2.4 and Appendix A.2.5, and also included discussions on performances in Sections 3.1-3.3.\n\nTable 3. Overview of symmetrization approaches.\n\n|  | GA | FA | Canonicalization | PS (Ours) |\n|---|---|---|---|---|\n| Symmetrization | g ~ Unif(G) | g ~ Unif(F(x)) | g = C(x) | g ~ p(g\\|x) |\n| \\|Sample space\\| | \\|G\\| | \\|G\\_x\\| ≤ \\|F(x)\\| ≤ \\|G\\| | 1 | \\|G\\_x\\| ≤ \\|supp p(g\\|x)\\| ≤ \\|G\\| |\n| Advantages | Simple | Efficient if F(x)is small | Efficient, X(x) can be learned | p(g\\|x) can be learned, empirically works well |\n| Disadvantages | The group G has to be compact, training challenge, costs linearly to sample size | F(x) requires manual solving, F(x) cannot be learned, costs linearly to sample size if F(x) is large | No exact equivariance for x with \\|G\\_x\\| > 1, training challenge compared to PS | The group G has to be compact, costs linearly to sample size |\n\nQ2. While it is suggested that the base model f can be arbitrary, only MLP and transformer were evaluated.\n\nA2. In principle, we can indeed use an arbitrary base model f, but focused on MLP and transformers as they have universality guarantees as in Theorem 2, and are widely used as general-purpose architectures. We will revise Line 3 to “we use a non-equivariant, general-purpose backbone...” to clarify our contributions.\n\nTo demonstrate that a more wide range of base models can be used, we added two experiments on graph isomorphism (Section 3.1) with S_n symmetrization of a GIN base model with node identifiers (GIN-ID) [7], and on particle dynamics (Section 3.2) with E(3) symmetrization of a GNN base model. The base architectures (GIN-ID and GNN) are directly adopted from [3]. In both, we achieved state-of-the-art performance; the results are in Tables 1 and 2 of the common response. The results indicate that our approach works for a wide range of base models.\n\nQ3. In [1,2], showing improved performance on image classification is a common application.\n\nA3. While rotated image datasets are useful to highlight equivariance, we think the groups are overly simple to compare symmetrization methods comprehensively. This is because the involved groups are very small, e.g., c4 contains four elements representing 90-degree rotations. In such cases, even full group averaging is computable [1]. As our method is designed to overcome the sample complexity of group averaging, we believe that it is more suitable to show our approach in more challenging groups, such as combinatorial, infinite groups, and their products, as we presented in experiments (Section 3). Please note that [3] was also not demonstrated for rotated images, likely for similar reasons.\n\nThus, instead of image classification, we additionally conducted an experiment on E(3) symmetrization of GNN, which is another widely used setup [2, 3]. We have obtained state-of-the-art result, which is in Table 2 of the common response.\n\nQ4. Convergence of the p(g|x) and dissimilarity to uniform distribution?\n\nA4. We have provided an analysis of the convergence of p(g|x) for S_n in Figure 1 of the main text, with comparison to uniform distribution (random permutations). In Figure 1, it can be seen that the distribution learns to produce consistent samples g ~ p(g|x) in early training, diverging from uniform distribution. In later training, the samples' variance g ~ p(g|x) slightly increases, but validation loss keeps improving, indicating that the model is leveraging stochasticity to improve performance.\n\nQ5. Sensitivity to the scales?\n\nA5. Our approach is robust to a wide range of scales, as our smallest base architecture involves 107,523 parameters (Section 3.2), and the largest involves 86M parameters (ViT-Base in Sections 3.3-3.4).\n\nQ6. A guideline for designing the G equivariant q? How significantly does its expressivity impact performance?\n\nA6. We found that having global interactions across input in q is helpful. For example, for S_n, the network q is expected to output a permutation by assigning score to each node. It is expected that it has to perform a kind of comparison over the graph, where global interaction could be useful and was implemented by virtual node and batch normalization. Other than this, the performance was robust across setups for a reasonable q. In all of our experiments on S_n, we use 3-layer GIN, and in all experiments on S_n x E(3) or E(3), we use 2-layer Vector Neurons, and found no issues.\n\nQ7. Some experimental comparisons between [1-3] missing?\n\nA7. We would like to clarify that our approach is compared to [1-3] for all experiments in Sections 3.1-3.3, except in one case where [3] was unavailable. Note that Equi-Tuning [1] is group averaging [5], as clarified in [1]. In this regard, in all experiments in Sections 3.1-3.3, we made comparisons with GA [1, 6], canonicalization [2], and FA [3]. One exception is in Table 2, where we could not experiment FA since frame for the full group S_n x E(3) has not been identified in current literature.\n\nTo further provide experimental comparisons to [1-3], we added experiments as mentioned in A1. The results can be found in Tables 2 and 3 of the common response, and they are consistent with our main findings as PS outperforms [1-3].\n\n[1] S. Basu et al., Equi-tuning: Group equivariant fine-tuning of pretrained models\n\n[2] S. Kaba et al., Equivariance with learned canonicalization functions\n\n[3] O. Puny et al., Frame averaging for invariant and equivariant network design\n\n[4] S. Basu et al., Equivariant few-shot learning from pretrained models\n\n[5] D. Yarotsky. Universal approximations of invariant maps by neural networks\n\n[6] B. Sturmfels. Algorithms in invariant theory\n\n[7] R. L. Murphy et al., Relational pooling for graph representations"}, {"review_id": "dyQ7FUPFu2", "reviewer": "Reviewer_Ry8Q", "summary": "This paper proposes a new symmetrization method for achieving equivariance for any base function. It absorbs previous methods with the same objective as special cases, including group averaging, frame averaging, and canonicalization. The method uses a learnable, equivariant map to generate the group representation necessary for achieving equivariance from an invariant, external noise variable, followed by necessary post-processing to ensure the validity of the group representation. The authors empirically demonstrate that the learned stochasticity in symmetrization leads to improved model performance than the previous methods, and showcase the benefits of general pre-training for equivariant models.", "questions": "* What's the computational cost of this method, including training and testing, compared to the canonicalization method?\n* Analyzing the sampling efficiency for estimating expectations would be beneficial, especially in comparison to the uniform distribution, as it is a major claim of the effectiveness of this method.", "limitations": "The authors have not discussed the limitations.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The proposed method is general, providing a unified perspective for previous methods and may lay the foundation for further theoretical exploration of the \"optimal\" conditional group distribution.\n* The presentation is clear and easy to read.\n* The finding of the benefit of non-symemtric pretraining for equivariant models seems interesting and new.", "weaknesses": "* My primary concern is that the supriorty of the proposed method compared to the canonicalization method is not clear enough.\n    - The advantages of the former over the latter seem to be more evident only in the first experiment, as the canonicalization method appears to have difficulties during optimization (although it is unknown whether this issue can be mitigated through some warm-up techniques). In the other two experiments, the performance of both methods is quite similar.\n    - On the other hand, compared to the canonicalization method, the proposed method incurs more computational cost in practice, including the additional parameters and hyperparameters in the learnable module, as well as the sampling cost (the normalization method only requires a single \"sample\", whereas here we need 10-20 samples).\n    \n+ This paper lacks thorough empirical and theoretical analyses for the two claimed superiorities over previous methods. Specifically,\n    - \"Learn to collaborate with base function to maximize task performance\": No sensitivity analysis is provided to demonstrate how this method adapts better to different base functions.\n    - \"learn to produce low-variance samples that can offer more stable gradients for the base function\": No analysis of gradient magnitudes is provided to indicate its stability.\n\n* The limitation is not discusses, including the extra computational cost.\n* The post-process seems interesting but also laborious, as it requires different designs for each different (compact) groups.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. The superiority compared to canonicalization is not clear.\n\nA1. Theoretically, we can show that there always exist certain input data that canonicalization fails to guarantee exact G equivariance, while PS guarantees G equivariance to all input data (as in Theorem 1).\n\nTo see this, let us recall the definition of G equivariant canonicalizer C: x ↦ ρ(g) [1]. A canonicalizer C is G equivariant if C(g⋅x) = ρ(g)C(x) for all g ∈ G and x ∈ X. Consider an input x which has a non-trivial stabilizer subgroup G_x = {h ∈ G | h⋅x = x}, i.e., has inner symmetries. It can be seen that G equivariant canonicalizer is ill-defined for these inputs. Specifically, let g1 = gh1 and g2 = gh2 for some g ∈ G and any h1, h2 ∈ G_x where h1 ≠ h2. Then we have C(g1⋅x) = C(gh1⋅x) = C(g⋅x) = C(gh2⋅x) = C(g2⋅x), implying that ρ(g1)C(x) = ρ(g2)C(x). Since g1 ≠ g2, this contradicts the group axiom, and therefore a G equivariant canonicalizer C cannot exist for inputs x with non-trivial G_x. For more detailed discussion, please see [1].\n\nTo handle all inputs, canonicalization [1] adopts relaxed equivariance: a canonicalizer C satisfies relaxed equivariance if C(g⋅x) = ρ(gh)C(x) up to arbitrary action from the stabilizer h ∈ G_x. As a consequence, the symmetrization ϕ(x) = g⋅f(g-1⋅x) performed using a relaxed canonicalizer C only guarantees relaxed equivariance ϕ(g⋅x) = gh⋅ϕ(x) up to arbitrary action from the stabilizer h ∈ G_x (proof is in [1]). Intuitively, this means canonicalization does not guarantee G equivariance for input data x with inner symmetries G_x.\n\nTo visually demonstrate this, we performed a minimal experiment using graphs x with non-trivial stabilizers G_x (inner symmetries [2]). The results are in Figures 1-3 of the common response. We fixed a randomly initialized MLP f: Rn x n → Rn and symmetrized it using our approach and canonicalization. When symmetrized, the MLP is expected to provide scalar embedding of each node, which we color-code for visualization. For each graph, we illustrate three panels: the leftmost one illustrates the color-coding of the inner symmetry of nodes (automorphic nodes), the middle one illustrates node embedding from MLP-PS, and the rightmost one illustrates embedding from MLP-Canonical. If a method is G equivariant, it is expected to give identical embeddings for automorphic nodes as an equivariant model cannot distinguish between the identities of individual automorphic nodes in principle [3]. \n\nAs in Figures 1-3, in the presence of inner symmetry (left panels), MLP with PS (middle panels) achieves G equivariance and produces almost identical embeddings for automorphic nodes. However, the same MLP with canonicalization fails and produces unstructured embeddings (right panels). The result illustrates a clear advantage of PS over canonicalization when learning data with inner symmetries, which is often found in applications such as molecular graphs [4].\n\nQ2. The advantages over canonicalization are evident only in the first experiment.\n\nA2. Please note that empirical advantage is evident not only in the first experiment (Section 3.1) but also in the second one (n-body; Section 3.2) where Transformer-PS achieves 0.00417 MSE, which significantly improves over Canonicalization with 0.00779 MSE.\n\nTo further demonstrate advantage over canonicalization, we added two experiments on graph isomorphism (Section 3.1) with S_n symmetrization of GIN base model with node identifiers (GIN-ID) [7], and on n-body (Section 3.2) with E(3) symmetrization of GNN. In both, we achieved state-of-the-art performance; the results are in Tables 2 and 3 of the common response. Our approach consistently improves over canonicalization as well as other symmetrizations.\n\nQ3. Compared to canonicalization, the method incurs more cost, including parameters and hyperparameters in learnable module, as well as sampling.\n\nA3. The O(N) cost with N samples is a genuine weakness, and we will add clarification in the limitations (Appendix A.4). As the sampling is parallelizable, we believe this can be overcome to some degree by leveraging parallelism [5] developed for scaling batch size.\n\nFor the parameters and hyperparameters, we think it adds a negligible overhead to canonicalization. This is because canonicalization also uses a equivariant module C: x ↦ ρ(g), so our approach does not add parameters in principle. For the hyperparameters, our approach only adds scale hyperparameter ŋ of noise, and we find a simple choice ŋ = 1 works robustly in all our experiments.\n\nQ6. No thorough analyses for the claimed superiorities. No analysis on sensitivity and gradients. Analyzing sampling efficiency would be beneficial.\n\nA6. In Lines 90-97, we are making a specific comparison against group averaging rather than claiming superiorities in general. We will revise the main text for clarification.\n\nWe are planning a controlled experiment to understand the training dynamics, including behavior of p(g|x), stability of gradients, and sample complexity in comparison to group averaging. We will post the results during the discussion period.\n\nQ7. The post-processing seems laborious.\n\nA7. The post-processing indeed has to be implemented for each group. However, it is not an issue specific to our approach, as it is an issue of canonicalization as well [1]. Also, we think designing it can often be more straightforward compared to alternatives, e.g., a frame which is G equivariant function that has to produce a set of group elements.\n\n[1] Kaba et al., Equivariance with learned canonicalization functions (2022)\n\n[2] Thiede et al., Autobahn: Automorphism-based graph neural nets (2022)\n\n[3] Srinivasan et al., On the equivalence between positional node embeddings and structural graph representations (2019)\n\n[4] MaKey et al., Surge: a fast open-source chemical graph generator (2022)\n\n[5] You et al., Scaling SGD batch size to 32k for ImageNet training (2017)\n\n[6] Puny et al., Frame averaging for invariant and equivariant network design (2021)"}, {"review_id": "E5GXXhjk5F", "reviewer": "Reviewer_wgQS", "summary": "The paper suggests a probabilistic approach to symmetrization, where an input conditional distribution is used to replace the untractable haar measure distribution of infinite groups. In turn, the paper identifies what are the conditions on the conditional distribution under which the symmetrization yields an equivariant function.  The method is evaluated on several benchmarks involving permutation and rotation symmetries.", "questions": "I would appreciate any response from the authors regarding the weakness stated above.", "limitations": "Could not find a discussion on limitations", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall the paper is well-written and easy to follow.\nThe formulation of the method seems to be adequate.\nThe proposed idea is simple, and a natural extension of previous methods.", "weaknesses": "Relation to previous works\n\nThe paper in the introduction states that existing approaches focus on either manually deriving smaller subsets [1], or implementing a relaxed version of equivariance [2]. It is not clear what manually means in this context. In fact, [1] identifies a similar condition to eq(5) when the function p(g|x) is a delta function. I noticed section 2.4 relates to these previous works more thoroughly. More importantly, I feel [2] and the suggested approach are very similar, as [2] does not only suggest an approximate equivariant but rather suggests a learnable frame function (group subset), satisfying the equivariance constraint (similar to eq (5)). Thus, I feel this work should convince us of the necessity of the probabilistic model for learnable frames. To this end, I would have expected an ablation study testing “apples-to-apples” architectures where the only difference is the addition of the input noise term.  \nIt would also be beneficial to separate the test of averaging over S_n to O(3) in the experiment in table (2) and use symmetrization only to the O(3) symmetries with S_n equivariant networks.\n\nMissing discussion\n\nIn fact, the network prediction in eq(4) is only approximated. Thus I assume, the loss terms used for training only estimate the true loss (gradients). How stable are these estimates? How well do they approximate the true loss?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. The introduction states that [1] focuses on manually deriving smaller subsets, but it is not clear.\n\nA1. We intended to describe that one needs to manually solve G equivariant set-valued frame, and it cannot be discovered from data. In contrast, the G equivariant distribution p(g|x) requires less hand design and can be learned entirely from data.\n\nQ2. [2] and the suggested approach is very similar.\n\nA2. Theoretically, we can show that there always exist certain input data that canonicalization fails to guarantee exact G equivariance, while PS guarantees G equivariance to all input data (as in Theorem 1).\n\nTo see this, let us recall the definition of a G equivariant canonicalizer C: x ↦ ρ(g) [2]. In the paper, it is stated that a canonicalizer C is G equivariant if C(g⋅x) = ρ(g)C(x) for all g ∈ G and x ∈ X. Now consider an input x which has a non-trivial stabilizer subgroup G_x = {h ∈ G | h⋅x = x}, i.e., has inner symmetries. It can be seen that G equivariant canonicalizer is ill-defined for these inputs. Specifically, let g1 = gh1 and g2 = gh2 for some g ∈ G and any h1, h2 ∈ G_x where h1 ≠ h2. Then we have C(g1⋅x) = C(gh1⋅x) = C(g⋅x) = C(gh2⋅x) = C(g2⋅x), implying that ρ(g1)C(x) = ρ(g2)C(x). Since g1 ≠ g2, this contradicts the group axiom, and therefore a G equivariant canonicalizer C cannot exist for inputs x with non-trivial G_x. For a more detailed discussion of this problem, please see Appendix A of [1].\n\nTo handle all possible inputs, canonicalization [1] adopts relaxed equivariance: a canonicalizer C satisfies relaxed equivariance if C(g⋅x) = ρ(gh)C(x) up to arbitrary action from the stabilizer h ∈ G_x. As a result, the symmetrization ϕ(x) = g⋅f(g-1⋅x) performed using a relaxed canonicalizer C only guarantees relaxed equivariance ϕ(g⋅x) = gh⋅ϕ(x) up to arbitrary action from the stabilizer h ∈ G_x (proof is given in [1]). Intuitively, this means canonicalization does not guarantee G equivariant processing for input data x with inner symmetries G_x.\n\nTo visually demonstrate this, we performed a minimal experiment using several graphs x with non-trivial stabilizers G_x (inner symmetries [2]). The results are in Figures 1-3 of the common response. We fixed a randomly initialized MLP f: Rn x n → Rn and symmetrized it using our approach and canonicalization. When symmetrized, the MLP is expected to provide scalar embedding of each node, which we color-code for visualization. For each graph, we illustrate three panels: the leftmost one illustrates the color-coding of the inner symmetry of nodes (automorphic nodes), the middle one illustrates node embedding from MLP-PS, and the rightmost one illustrates node embedding from MLP-Canonical. If a method is G equivariant, it is expected to give identical embeddings for automorphic nodes because an equivariant model cannot distinguish between the identities of individual automorphic nodes in principle [3]. \n\nAs  in Figures 1-3, in the presence of inner symmetry (left panels), MLP with PS (middle panels) is able to perform G equivariant processing and produce almost identical embeddings for automorphic nodes. However, the same MLP with canonicalization fails and produces unstructured embeddings (right panels). The result illustrates a potential advantage of PS over canonicalization when learning data with inner symmetries, which is found in applications such as molecular graph processing [4].\n\nQ3. I expected an ablation study where the only difference is the addition of the input noise.\n\nA3. All our main experiments in Sections 3.1-3.3 are conducted in the way the reviewer suggested. Please see Lines 214-216 in the main text; all canonicalization models are constructed by deleting the noise term from the probabilistic mode using identical architectures.\n\nQ4. It would be beneficial to use symmetrization only to the O(3) symmetries with S_n equivariant networks.\n\nA4. We additionally conducted an experiment on E(3) symmetrization of an S_n equivariant GNN for n-body task (Section 3.2). The results are in Table 2 of the common response, and our approach obtains state-of-the-art performance and significantly improves over all other symmetrization approaches.\n\nQ5. The loss terms used for training only estimate the true loss (gradients). How stable are these estimates? How well do they approximate the true loss?\n\nA5. While Eq. (4) gives a G equivariant symmetrized function ϕ(x), we cannot directly observe ϕ(x), but can obtain samples of the unbiased estimator g⋅f(g-1⋅x). Given that, it can be questionable what objective we are actually optimizing when we use estimations of ϕ(x). Fortunately, a theoretical framework for training of S_n symmetrized models has been established [3-4], from which we generalize to the general group. The main message is that minimizing a convex loss function based on the sampled g⋅f(g-1⋅x) is equivalent to minimizing an upper bound to the true objective that involves ϕ(x).\n\nGiven a training set D = {(x1, y1), …, (xn, yn)} for a G equivariant task, our true objective would be minimizing the empirical loss L(D; θ, ω) = sum_i l(y_i, ϕ(x_i)) where l is a convex loss. However, in practice, ϕ(x) cannot be observed, and we observe g⋅f(g-1⋅x). While g⋅f(g-1⋅x) serves as an unbiased estimator of ϕ(x), the estimation is no longer unbiased when g⋅f(g-1⋅x) is used to compute the loss: from Jensen’s inequality, we have E_g[l(y, g⋅f(g-1⋅x))] ≥ l(y, E_g[g⋅f(g-1⋅x)]). That is, minimizing the sampling-based loss is minimizing an upper-bound surrogate to the true objective [3].\n\nQ6. Could not find a discussion on limitations.\n\nA6. We will move Appendix A.4 to the main text for visibility.\n\n[1] O. Puny et al., Frame averaging for invariant and equivariant network design (2021)\n\n[2] S. Kaba et al., Equivariance with learned canonicalization functions (2022)\n\n[3] R. L. Murphy et al., Relational pooling for graph representations\n\n[4] R. L. Murphy et al., Janossy pooling: Learning deep permutation invariant functions for variable-sized inputs"}, {"review_id": "yYpuJ37X5q", "reviewer": "Reviewer_GZXj", "summary": "The paper presents a method to learn to symmetrize a neural network using data. The method considers a learnable probability distribution over the group and uses group averaging to enforce (relaxed) equivariance w.r.t. the learned distribution. ", "questions": "a) How does the method differ from equivariant Augerino (Sec. 3.1 of [1])?\n\nb) What objective is used, or how is inference over the probability distribution being performed?\n\nc) If a maximum likelihood objective is used, what mechanism prevents the symmetrization from collapsing into the identity? Did authors notice such a collapse in experiments, or is this mitigated somehow?\n\n\n\n[1] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" Advances in neural information processing systems 33 (2020): 17605-17616.\n", "limitations": "- The paper seems to miss prior works that consider learning similar parameterizations over group structures.\n- The proposed method seems to be proposed already (equivariant Augerino, Sec. 3.1 of [1])\n- It is not clear what objective is being used to train the model and what prevents collapse of the symmetrization distribution.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Over recent years there has been a growing interest in the community to learn relaxed symmetries (invariance and equivariance) from data. The paper addresses this important problem and provides a practical method that scales to interesting model classes and data problems (e.g. the use of Transformers, particle dynamics, and graphs is noteworthy). ", "weaknesses": "* Missed prior work\n\nThere have been several approaches that aim to learn symmetrizing distributions (invariance or equivariance). Placing a probability distribution on the group over which a function is averaged and learning this distribution from data is not new. For instance, see approaches listed in survey Sec. 6 of [2].\n\n* Proposed method already exists \n\nThe paper misses this literature and it seems that the method has been published before (Equivariant Augerino, Sec. 3.1 of [1]).\nAn argument could be made that although the Augerino paper describes the method for general groups, it only considers very simple affine groups in practice. This paper does consider more interesting group structures and domains (e.g. graphs). However, if the same method has been published before and is applied to new domains/data, this should at least be credited and discussed. This would make it more clear what the contributions of the work are.\n\n* Objective function\n\nA well-known difficulty with learning symmetry distribution is not necessarily describing a probability distribution on the group but rather the objective to learn the distribution over the group. Although the method claims to be probabilistic, it is not clear what objective is used to train the model (or whether inference is being performed). If regular maximum likelihood loss is used, it seems there is no encouragement in the objective that prevents the collapse of the symmetrization distribution into a delta peak at the identity. In such a case, it would result in not learning the symmetrization that generalises well. Prior works often consider more sophisticated losses, such as regularization [1], lower bounds or model selection, to prevent this. The paper seems to skip over this issue.\n\n[1] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" Advances in neural information processing systems 33 (2020): 17605-17616.\n\n[2] Rath, Matthias, and Alexandru Paul Condurache. \"Boosting deep neural networks with geometrical prior knowledge: A survey.\" arXiv preprint arXiv:2006.16867 (2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. The proposed method already exists.\n\nA1. Since our approach parameterizes a distribution p(g|x) on a group G for symmetrization ϕ(x) = E_g[g⋅f(g-1⋅x)] and learns it from data, one may find some similarity to Augerino [1] and other approaches in [2] such as [3-7] that learn distributions over augmentations (e.g., p(g)) for a similarl symmetrization. However, we would like to clarify that our method is distinguished from all these approaches [1-7], including Augerino [1], as our method is the only one that guarantees G equivariance of symmetrized ϕ(x) while able to learn useful distribution p(g|x) from data to maximize performance. Let us elaborate below.\n\nFirst, it has to be noted that the objectives of the approaches are different. Our approach aims to obtain an exact G equivariant symmetrization ϕ(x) given the known symmetry group G of data (e.g., G=S_n for graphs), while [1-7] aim to discover underlying (approximate) symmetry constraint from data. Because of this, the symmetrizing distribution on G has to be designed differently: In our case, we parameterize the distribution p(g|x) itself to be G equivariant (see Theorem 1), while for [1-7], the distribution p(g) is parameterized for a completely different purpose of covering a range of different symmetry constraints and their approximations (e.g., a set of 2D affine transformations [1]).\n\nImportantly, this allows our approach to learn non-trivial and useful distribution p(g|x) per data x while keeping the symmetrized ϕ(x) exactly G equivariant. This is a key advantage distinguished from [1-7] if the group G is given, since [1-7] does not guarantee G equivariance, nor can they learn useful p(g|x) in the case they learn to be G equivariant. For example, Augerino [1] puts an unconditional distribution p(g) on group G; in order to achieve exact G equivariance, p(g) has to learn a right invariant distribution, i.e., p(gh) = p(g) for all h ∈ G (see Sec. 3.1 of [1]). Since such p(g) is a Haar distribution, i.e., the uniform distribution Unif(G), it can be seen that Augerino has to reduce to basic group averaging [8] in order to achieve exact G equivariance. Please note that we have already extensively discussed the advantages of our approach compared to group averaging in the main text; for example, see Lines 42-45 and 90-97. The same argument also applies to [3-6]: see Sec. 3.1 of [3], Appendix A of [4], Sec. 3 of [5], and Sec. 3.4 of [6]. LILA [7] is closer to our approach as it defines an input-conditional augmented distribution p(x’|x) (see Eq. (1) of [7]), but still does not guarantee G equivariance of ϕ(x) as the distribution p(x’|x) is unconstrained.\n\nWe will revise the main text and add the above citations and discussion to make this more clear.\n\nQ2. It is not clear what objective is being used to train the model and what prevents the collapse of the symmetrization distribution.\n\nA2. In symmetrization approaches that aim to discover symmetry or its extent from data [1-7], encouragement in the objective, e.g., through regularization [1] or model selection [7], has to be used to prevent the collapse of the symmetry to the trivial group (i.e., p(g|x) = \\delta(g=id) for all x). At a high level, this is because the symmetrized ϕ(x) = E_g[g⋅f(g-1⋅x)] is allowed to search over the space of symmetries, which includes the trivial group; if one uses the maximum likelihood objective, the model would likely favor the trivial symmetry because it is the least constrained and would fit the training data most easily.\n\nHowever, in our case, the goal is not to search over the space of symmetry groups but to build an exact G equivariant function for a given symmetry group G (e.g., G=S_n for graphs). For this, as in Theorem 1, we constrain the symmetrizing distribution p(g|x) to be G equivariant, which leads to a strong provable guarantee that the symmetrized function ϕ(x) = E_g[g⋅f(g-1⋅x)] is always G equivariant regardless of how the (G equivariant) distribution p(g|x) is trained. In other words, the symmetrized model ϕ(x) cannot collapse to trivial symmetry, as it is enforced to be equivariant for the given symmetry group G. This allows us to use regular maximum likelihood objective for training without the need to address symmetry collapse.\n\nAs a more bottom-up interpretation, we can show that if a symmetrizing distribution p(g|x) is G equivariant (as in Theorem 1), it cannot technically collapse to a delta peak at the identity p(g|x) = \\delta(g=id) for all x. To prove this, assume that a G equivariant p(g|x) has collapsed. Recall G equivariance p(g|x) = p(g’g|g’⋅x) for all x and g’ ∈ G (see Eq. (5)); transforming input has to transform the distribution accordingly. This yields a contradiction, as transforming an input x to g’⋅x has to transform the distribution p(g|x) with g’ as well, and as a result, p(g|x’) for x’ = g’⋅x is no longer a delta peak at identity.\n\nQ3. How is inference performed?\n\nA3. Since we use regular maximum likelihood objective, we only need to perform sampling g ~ p(g|x) for MC estimation of ϕ(x) to obtain loss, and inference or density estimation on the symmetrizing distribution p(g|x) is not necessary.\n\n[1] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" (2020)\n\n[2] Rath, Matthias, and Alexandru Paul Condurache. \"Boosting deep neural networks with geometrical prior knowledge: A survey.\" (2020).\n\n[3] ​​Rommel, Cédric, et al. “Deep invariant networks with differentiable augmentation layers.” (2022).\n\n[4] Wilk, Mark van der, et al. “​​Learning invariances using the marginal likelihood.” (2018).\n\n[5] Wilk, Mark van der, et al. “Learning invariant weights in neural networks.” (2022).\n\n[6] Romero, David W., et al. “Learning partial equivariances from data.” (2021).\n\n[7] Immer, Alexander, et al. “Invariance learning in deep neural networks with differentiable Laplace approximations.” (2022).\n\n[8] Yarotsky, Dmitry. ”Universal approximations of invariant maps by neural networks.” (2018).\n"}, {"review_id": "RvP7H3prC9", "reviewer": "Reviewer_SdiQ", "summary": "The paper proposes a probabilistic version of the canonization method to construct equivariant architectures from generic non-equivariant backbones.\nThis approach is inspired by the symmetrization solution, which involves averaging a non-equivariant function to obtain an equivariant one, but replaces the uniform averaging distribution with an input-conditional one to achieve a lower sample complexity.\n\n", "questions": "\nWhy is it necessary to assume the group is compact?\n\nline 96-97:  \"2) learn to produce low-variance samples $g \\sim p_\\omega(g|x)$ that can offer more stable gradients for the base function $f_\\theta$ during training.\"\nI think this claim requires some additional comments to justify it. Why should the network learn to produce low-variance samples?\n\nOn a related note, do you train the models with some low-variance regularization? what prevents the model from outputting always a uniform distribution?\n\nlines 114-117: this seems more a problem of parameterizing the group G itself, not its representations. The representation $\\rho(g)$ can be computed in a differentiable way from $g$ when $G$ is a Lie group at least.\n\nEq. 8: it is not really clear how you output a permutation matrix. You use the notation $P_g \\approx \\hat{P}_g$ but you don't explain what it means. The matrix $\\hat{P}_g$ is a real valued matrix and doesn't belong to $(0, 1)^{n \\times n}$. Do you apply rounding on it? \n\nIf so, isn't it possible to just use the matrix $\\hat{P}_g$? Is there any benefit in using the boolean $P_g$ instead?\n\n\nYou compare your probabilistic method with the deterministic canonicalization method.\nIt is not clear to me if the benefit is coming from the fact the backbone is averaged over $N>1$ samples or if the stochasticity of $p_\\omega$ plays a role.\nCould you try a version of your model which outputs $N$ samples but deterministically (e.g. sample the noise variables only once before training and keep them fixed)?\n\n\nWhile this is already mentioned by the authors among the limitations as a future work, it would be interesting to understand if this kind of architectures enjoys the same data-efficiency property of other equivariant networks. \nDo the authors have any insights about the effect of the training set size on the ability of the model to learn the symmetrising distribution?\n\n\nThe paper is missing a related works section. Moreover, I think an important part of the literature about equivariant networks using group-convolution is not cited.\n", "limitations": "Some limitations are discussed in the supplementary material.\nI encourage the authors to move the main points mentioned there in the main paper.\nSee also my comments under Weaknesses.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is clearly written and well motivated.\n\nThe proposed idea overcomes some important limitations of both the symmetrisation and the canonicalisation methods.\n\nI liked the idea of using a general-purpose transformer architecture, which makes the proposed method very flexible and easy to apply in more generic contexts.\nThe surprising performance of pre-trained vision models on non-vision tasks is also a very interesting result and I wonder if it can tell us something more general about these architectures.\n\nDespite the generality and flexibility of the approach, the experiments show it can also achieve competitive performance.\n", "weaknesses": "\nWith respect to the canonicalization method from [26], the proposed approach essentially averages over $N>1$ samples and, therefore, is $N$ times more expensive than canonicalization.\nSince this work is partially motivated by computational efficiency, I think this aspect has not been sufficiently discussed in the manuscript.\n\nMoreover, I imagine the number of samples chosen can play an important role on 1) the model's final performance, 2) the training dynamic, 3) the variance of the learned distributions and 4) the training time.\nI think this point also requires some additional discussion.\nI would encourage the authors to include a short experiment to study the effect of the sample size on (at least some of) these aspects.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Q1. The proposed approach essentially averages over N > 1 samples and, therefore, is N times more expensive than canonicalization.\n\nA1. The O(N) cost with N samples is a genuine weakness, and we will add clarification n the main text. Nevertheless, as the sampling is completely parallelizable and analogous to using a larger batch size, we believe this can be overcome to some degree by leveraging parallelism [5] developed for scaling batch size.\n\nQ2. I would encourage the authors to include a short experiment to study the effect of the sample size on training.\n\nA2. We agree that controlling the number of samples used at training time would offer interesting results. We are planning to add a small empirical analysis on how the sample size affects the optimization dynamics and will post the result during the discussion period.\n\nQ3. Why is it necessary to assume the group is compact?\n\nA3. The compactness of group G is necessary for the proof in Theorem 2. Except that, the group does not need to be compact. We can extend Theorem 2 to non-compact groups if we make additional assumptions, such as the support of p(g|x) being always compact.\n\nQ4. Why should the network learn to produce low-variance samples? Do you train the models with some low-variance regularization? what prevents the model from always outputting a uniform distribution?\n\nA4. At some point, certain group elements g ~ p(g|x) can be favored more than others in terms of task loss; this is likely due to the specific configuration of the parameters of the base function f coming from random initialization or pre-training. In either case, it will signal the distribution p(g|x) to favor generating these group elements to minimize task loss, which eventually leads to learning to produce those samples consistently.\n\nIn all experiments in the paper, we directly minimize the task loss and do not use low-variance regularization. In our early exploration, we experimented with two variance regularization techniques, which are simple equivariance loss [3] and entropy estimation on the distribution p(g|x) [4]. However, equivariance loss often resulted in collapsed trivial predictions, and entropy estimator did not bring performance gain. We conjecture that the task loss provides a strong signal to nudge the distribution p(g|x) to favor specific group elements g ~ p(g|x) and this leads to a reduced variance of group elements without collapsing to Unif(G).\n\nQ6. Lines 114-117: this seems more like a problem of parameterizing the group G itself, not its representations.\n\nA6. Indeed, the distribution p(g|x) does not fundamentally need to output a representation, i.e., it may produce alternative specifications of g (e.g., a Lie algebra element) from which representation is computed (e.g., with exponential map). In our work, we produce group representation as it provides a convenient medium for sampling and backpropagation through p(g|x) while guaranteeing the G equivariance of p(g|x). Alternatives, e.g., based on Lie algebra [5], is possible. We will revise the main text to clarify this.\n\nQ7. Eq. 8: it is not really clear how you output a permutation matrix.\n\nA7. We apologize for the confusion. Given nodewise scalar Z ∈ Rn from GNN, we first perform argsort to directly obtain the hard permutation matrix Pg [6]. And then, we compute the soft permutation matrix hat{P}_g following Eq. (8) [6]. hat{P}_g does not affect how the hard permutation Pg is computed, but it affects training as it is used for straight-through gradient.\n\nQ8. Isn't it possible to just use the matrix hat{P}_g? Is there any benefit in using the boolean P_g instead?\n\nA8. The first reason we use P_g is, since hat{P_g} is not a valid representation, symmetrization using it would not be G equivariant. The second reason is, using hard permutation is often critical for convergence. We conjecture this is due to the smoothing of soft permutations; in early training, hat{P}_g is close to a dense matrix, and applying it to input would smooth the information of nodes. Such smoothing is known to underlie oversmoothing of GNNs [8]. We added an experimental result of using the soft permutation for training, please see Table 4 below.\n\nQ9. Could you try a version of your model that outputs N samples but deterministically?\n\nA9. We ran a variant of our approach in the EXP-classify (Section 3.1), where the noise are sampled once before training and fixed. The results are in Table 4. The model quickly reaches 82% test accuracy, but then severely overfits the training data. This implies that the effect of averaging is insufficient to explain our approach's performance, and the stochasticity of p(g|x) plays an important role.\n\nTable 4. Ablation-extended results for EXP-classify.\n\n| method | EXP-classify ↑ |\n|:---:|:---:|\n| MLP-GA [7] | 50% |\n| MLP-FA [7] | 100% |\n| MLP-Canonical. | 50% |\n| MLP-PS (Ours) | 100% |\n| MLP-PS (Ours; soft permutation) | 50% |\n| MLP-PS (Ours; fixed noise) | 82% |\n\nQ11. The paper is missing a related works section. I encourage the authors to move the limitation sections in the main paper.\n\nA11. We will revise the main text accordingly.\n\n[1] Kaba et al., Equivariance with learned canonicalization functions (2022)\n\n[2] You et al., Scaling SGD batch size to 32k for ImageNet training (2017)\n\n[3] Murphy et al., Janossy pooling: Learning deep permutation invariant functions for variable-sized inputs (2019)\n\n[4] Seo et al., State entropy maximization with random encoders for efficient exploration (2021)\n\n[5] Finz et al., Generalizing convolutional neural networks for equivariance to Lie groups on arbitrary continuous data (2020)\n\n[6] Prillo et al., Softsort: A continuous relaxation for the argsort operator (2020)\n\n[7] Puny et al., Frame averaging for invariant and equivariant network design (2021)\n\n[8] Cai et al., A note on over-smoothing for graph neural networks (2020)"}]

Author rebuttal:
The common response PDF contains the following:

- Table 1, including additional results on $\textnormal{S}\_n$ invariant graph isomorphism learning with $\textnormal{S}\_n$ symmetrized GIN-ID base function.

- Table 2, including additional results on $\textnormal{S}\_n\times\textnormal{E}(3)$ equivariant $n$-body problem with $\textnormal{E}(3)$ symmetrized GNN base function.

- Figure 1 to 3, including visualizations of scalar node embeddings of graphs produced by randomly initialized MLP-PS (Ours) and MLP-Canonical. under the presence of inner-symmetries.

We have added the responses to each reviewer individually.

Meta-review: Summary

This work presents an approach to make deep learning models equivariant to a given group. The experiments of this paper on patch based transformer show competitive performance on a variety of equivariant groups such as permutation and rotation. This is valuable given the interest in using transformer based architectures for a variety of recognition tasks. Moreover, the paper also shows applicability of the method on particle dynamics and graphs. 

Reviews & Justification

This paper received positive reviews and all the reviewers were unanimous in recognizing its acceptance. The wide applicability of this method across architectures, and domains (vision, particle dynamics, graphs) makes it interesting to a wider audience. The AC also appreciates the in-depth author responses and additional experiments to further bolster the reviewer’s confidence. Please add in the related work suggested by the reviewers.
Decision: Accept (spotlight)