Title: Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance
Venue: NeurIPS 2023
Accepted: yes
Keywords: equivariant machine learning

transformers

graphs

general-purpose architectures
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

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 .

1 Introduction

Many perception problems in machine learning involve functions that are invariant or equivariant to certain symmetry group of transformations of data. Examples include learning on sets and graphs, point clouds, molecules, proteins, and physical data, to name a few [10, 12]. Equivariant architecture design has emerged as a successful approach, where every building block of a model is carefully restricted to be equivariant to a symmetry group of interest [5, 12, 23]. However, equivariant architecture design faces fundamental limitations, as individual construction of models for each group can be laborious or computationally expensive [91, 58, 64, 45], the architectural restrictions often lead to limited expressive power [101, 56, 105, 40], and the knowledge learned from one problem cannot be easily transferred to others of different symmetries as the architecture would be incompatible.

This motivates us to seek a symmetrization solution that can achieve group invariance and equivariance with general-purpose, group-agnostic architectures such as an MLP or a transformer. As a basic form of symmetrization, any parameterized function f θ : X → Y on vector spaces X , Y can be made invariant or equivariant by group averaging [102, 67], i.e. , averaging over all possible transformations of inputs x ∈ X and outputs y ∈ Y by a symmetry group G = { g } :

where ϕ θ is equivariant or invariant to the group G . An important advantage is that the symmetrized function ϕ θ can leverage the expressive power of the base function f θ ; it has been shown that ϕ θ is a universal approximator of invariant or equivariant functions if f θ is a universal approximator [102],

G

Probabilistic Symmetrization

-1

-1

x

g

∈

G

ϕθ,ω

(

x

)

g

-1 ·

x

g

·

f

θ

(

g

-1 ·

x

)

f

θ

-1

G

= S

n

...

=

S

n

×

E(3)

=   E(3)

graph

geometric graph

motion

...

...

S

n

equivariant

S

n

×

E(3)

equivariant

E(3)

equivariant

G

Equivariant

Distribution

p

ω

(

g

|

x

)

Base Function

f

θ

G

which includes an MLP [35] or a transformer [104]. On the other hand, an immediate challenge is that for many practical groups involving permutation and rotations, the cardinality of the group | G | is large or infinite, so the exact averaging is intractable. Due to this, existing symmetrization approaches often focus on small finite groups [4, 65, 94, 42], manually derive smaller subsets of the entire group e.g. , a frame [74] to average over [67, 68], or implement a relaxed version of equivariance [41, 83].

An alternative method for tractability is to interpret Eq. (1) as an expectation with uniform distribution Unif ( G ) over the compact group G [62], and use sampling-based average to estimate it [102, 67, 68]:

where g · f θ ( g -1 · x ) serves as an unbiased estimator of ϕ θ ( x ) . While simple and general, this approach has practical issues that the base function f θ 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 a parameterized distribution p ω ( g | x ) in a way that equivariance and expressive power are always guaranteed, and train it end-to-end with the base function f θ to directly minimize task loss. We show that the distribution p ω ( 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 to generally implement p ω ( g | x ) as a noise-outsourced map ( x , ϵ ) ↦→ g with an invariant noise ϵ and a small equivariant network q ω , which enables gradient-based training with reparameterization [47]. As p ω is trained, it can enhance the learning of f θ by producing group transformations of lower variance compared to Unif ( G ) so that f θ is less burdened, and coordinating with f θ to maximize task performance. We refer to our approach as probabilistic symmetrization. An overview is provided in Figure 1.

We implement and test our method with two general-purpose architectures as the base function f θ : MLP and transformer. In particular, our transformer backbone is architecturally identical to patchbased vision transformers [26], which allows us to initialize most of its parameters from ImageNet-21k pretrained weights [89] and only replace the input and output projections to match task dimensions. We implement the conditional distribution p ω ( g | x ) for a wide range of practical symmetry groups including permutation (S n ) and Euclidean groups (O / SO ( d ) and E / SE ( d ) ) and their product combinations ( e.g. , S n × O (3) ), all of which are combinatorial or infinite groups. Empirical tests on a wide range of invariant and equivariant tasks involving graphs and motion data show competitive results against tailored equivariant architectures as well as existing symmetrization methods [74, 41]. This suggests the potential for learning invariant or equivariant functions for diverse groups with a groupagnostic general-purpose backbone. We further show evidence that pretraining from non-symmetric modality (vision) leads to enhanced learning in symmetric modality (graphs).

2 Probabilistic Symmetrization for Equivariance

We introduce and analyze our approach called probabilistic symmetrization, which involves an equivariant distribution p ω and group-agnostic base function f θ , in Section 2.1. We then describe implementation of p ω for practical groups including permutations and rotations in Section 2.2. Then, we describe our choice of base function f θ focusing on MLP and, transformers in particular, in Section 2.3. All proofs can be found in Appendix A.1.

Problem Setup In general, our goal is to construct a function ϕ : X → Y on finite vector spaces X , Y that is invariant or equivariant to symmetry specified by a group G = { g } 1 . This is formally described by specifying how the group act as transformations on the input and output. A group representation ρ : G → GL ( X ) , where GL ( X ) is the set of all invertible matrices on X , associates each group element g ∈ G to an invertible matrix ρ ( g ) that transforms a given vector x ∈ X through x ↦→ g · x = ρ ( g ) x . Given that, a function ϕ : X → Y is G equivariant if:

where the representations ρ 1 and ρ 2 are on the input and output, respectively. G invariance is a special case of equivariance when the output representation is trivial, ρ 2 ( g ) = I .

2.1 Probabilistic Symmetrization

To construct a G equivariant function ϕ θ , group averaging symmetrizes an arbitrary base function f θ : X → Y by taking expectation with uniform distribution over the group (Eq. (2)). Instead, we propose to use an input-conditional parameterized distribution p ω ( g | x ) and symmetrize f θ as follows:

where the distribution p ω ( g | x ) itself satisfies probabilistic G equivariance:

Importantly, we show that probabilistic symmetrization with equivariant p ω guarantees equivariance as well as expressive power of the symmetrized ϕ θ,ω .

Theorem 1. If p ω is G equivariant, then ϕ θ,ω is G equivariant for arbitrary f θ .

Proof. The proof can be found in Appendix A.1.1.

Theorem 2. If p ω is G equivariant and f θ is a universal approximator, then ϕ θ,ω is a universal approximator of G equivariant functions.

Proof. The proof can be found in Appendix A.1.2.

While the base function f θ that guarantees universal approximation can be chosen in a groupagnostic manner e.g. , an MLP [35, 20] or a transformer on token sequences [104], the distribution p ω needs to be instantiated group specifically to satisfy G equivariance. A simplistic choice is using uniform distribution Unif ( 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 p ω can provide distinguished advantages compared to the uniform distribution, as it can (1) learn from data to collaborate with (pre-trained) base function f θ to maximize task performance and (2) learn to produce more consistent samples g ∼ p ω ( g | x ) that can offer more stable gradients for the base function f θ during early training.

We now provide a generic blueprint of G equivariant distribution p ω ( g | x ) for any compact group G . Our goal is to sample g ∼ p ω ( g | x ) to obtain group representation ρ ( g ) for symmetrization (Eq. (4)) in a differentiable manner so that p ω can be trained end-to-end. Since we only need sampling and there is no need to evaluate likelihoods, we simply implement p ω ( g | x ) as a noise-outsourced, differentiable transformation q ω ( x , ϵ ) of a noise variable ϵ ∈ E that directly outputs a group representation ρ ( g ) :

where q ω is G equivariant and p ( ϵ ) is G invariant under an appropriate representation ρ ′ :

Given above implementation, we can show the G equivariance of p ω :

Theorem 3. If q ω is G equivariant and p ( ϵ ) is G invariant under representation ρ ′ that | det ρ ′ ( g ) | = 1 ∀ g ∈ G , the distribution p ω ( g | x ) characterized by q ω : ( x , ϵ ) ↦→ ρ ( g ) is G equivariant.

In practice, one can use any available G equivariant neural network to implement q ω , e.g. , a graph neural network for the symmetric group S n , or an equivariant MLP which can be constructed for any matrix group [30]. Since we expect most of the reasoning to be done by the base function f θ , the equivariant network q ω 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 ( ϵ ) , simple choices often suffice for G invariance. For example, standard normal ϵ ∼ N (0 , I n ) provides invariance for the symmetric group S n as well as the (special) orthogonal group O ( n ) and SO ( n ) .

One important detail in designing q ω is constraining its output to be a valid group representation ρ ( g ) . For this, we apply appropriate postprocessing to refine neural network features into group representations, e.g. , Gram-Schmidt orthogonalization to obtain a representation ρ ( g ) ∈ R n × n of the orthogonal group g ∈ O ( n ) . Importantly, to not break the G equivariance of q ω , this postprocessing needs to be equivariant itself, e.g. , Gram-Schmidt process is itself O ( n ) equivariant [41]. Implementations of p ω for a range of practical symmetry groups are provided in detail in Section 2.2.

2.2 Equivariant Distribution p ω

We present implementations of the G equivariant distribution p ω ( g | x ) for a range of practical groups demonstrated in our experiments (Section 3). Formal proofs of correctness are in Appendix A.1.4.

Symmetric Group S n The symmetric group S n over a finite set of n elements contains all permutations of the set, which describes symmetry to ordering desired for learning set and graph data. The base representation is given by ρ ( g ) = P g where P g ∈ { 0 , 1 } n × n is a permutation matrix for g .

To implement S n equivariant distribution p ω ( g | x ) that provides permutation matrices P g from graph data x , we use the following design. We first sample invariant noise ϵ ∈ R n × d from i.i.d. uniform Unif [0 , η ] with noise scale η . For the S n equivariant map q ω : ( x , ϵ ) ↦→ ρ ( g ) , we first use a graph neural network (GNN) as an equivariant map ( x , ϵ ) ↦→ Z that outputs nodewise scalar Z ∈ R n . Then, assuming Z is tie-free 2 , we use below argsort operator Z ↦→ P g to obtain permutation matrix [73, 98]:

where eq denotes elementwise equality indicator. The argsort operator is S n equivariant, i.e. , it maps P g ′ Z ↦→ P g ′ P g for all P g ′ ∈ S n . To backpropagate through P g during training, we use straightthrough gradient estimator [6] with an approximate permutation matrix ˆ P g ≈ P g obtained from a differentiable relaxation of argsort operator [61, 33, 98]; details can be found in Appendix A.3.1.

Orthogonal Group O ( n ) , SO ( n ) The orthogonal group O ( n ) contains all roto-reflections in R n 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 ρ ( g ) = Q g where Q g is the orthogonal matrix for g . For SO ( n ) , the representation is ρ ( g ) = Q + g where Q + g is the orthogonal matrix of g with determinant +1 .

To implement O ( n ) /SO ( n ) equivariant distribution p ω ( g | x ) that provides orthogonal matrices given input data x , we use the following design. We first sample invariant noise ϵ ∈ R n × d from i.i.d. normal N (0 , η 2 ) with noise scale η . For the O ( n ) /SO ( n ) equivariant map p ω : ( x , ϵ ) ↦→ ρ ( g ) , we first use an equivariant neural network as a map ( x , ϵ ) ↦→ Z that outputs n features Z ∈ R n × n . Then, assuming Z is full-rank 3 , we use Gram-Schmidt process for orthogonalization, which is differentiable

and O ( n ) equivariant [41]. This completes the postprocessing Z ↦→ Q g for O ( n ) . For SO ( n ) , we can further use a simple scale operator Q ↦→ Q + g to set the determinant of orthogonalized matrix to +1 :

The scale operator is differentiable and SO ( n ) equivariant, i.e. , it maps Q + g ′ Q ↦→ Q + g ′ Q + g for all Q + g ′ ∈ SO ( n ) , thereby completing the postprocessing Z ↦→ Q + g for SO ( n ) .

Euclidean Group E ( n ) , SE ( n ) The Euclidean group E ( n ) contains all roto-translations and reflections in R n and their combinations, and the special Euclidean group SE ( n ) contains all rototranslations without reflections. These groups are desired in learning physical systems such as a particle in motion. Formally, the Euclidean group is given as a combination of orthogonal group and translation group E ( n ) = O ( n ) ⋉ T ( n ) , and similarly SE ( n ) = SO ( n ) ⋉ T ( n ) . As the translation group T ( n ) is non-compact which violates our assumption for symmetrization, we handle it separately. Following prior work [74, 41], we subtract the centroid ¯ x from the input data x as x -¯ x , and add it to the rotation symmetrized output as ¯ x + g · f θ ( g -1 · ( x -¯ x ))) where g is sampled from O ( n ) /SO ( n ) equivariant p ω ( 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 we often encounter product combinations of groups G = H × K that describe joint symmetry to each group H and K . For example, S n × O (3) describes joint symmetry to permutations and rotations, which is desired in learning point clouds, molecules, and particle interactions. In general, an element of H × K is given as g = ( h, k ) where h ∈ H and k ∈ K , and group operations are applied elementwise gg ′ = ( h, k )( h ′ , k ′ ) = ( hh ′ , kk ′ ) . The base group representation is accordingly given as pair of representations ρ ( g ) = ( ρ ( h ) , ρ ( k )) . While a common approach to handling H × K is partially symmetrizing on H and imposing K equivariance on the base architecture ( e.g. , rotational symmetrization of a graph neural network for S n × O (3) equivariance [74, 41]), we extend to full symmetrization on H × K since our goal is not imposing any constraint on the base function f θ .

To implement H × K equivariant distribution p ω ( g | x ) that gives ρ ( g ) = ( ρ ( h ) , ρ ( k )) from data x , we use the following design. We first sample invariant noise ϵ from i.i.d. normal N (0 , η 2 ) with scale η . For the H × K equivariant map q ω : ( x , ϵ ) ↦→ ( ρ ( h ) , ρ ( k )) , we employ a H × K equivariant neural network as a map ( x , ϵ ) ↦→ ( Z H , Z K ) such that the postprocessing for each group H and K provides maps Z H ↦→ ρ ( h ) and Z K ↦→ ρ ( k ) respectively, leading to full representation ρ ( g ) = ( ρ ( h ) , ρ ( k )) . For this whole procedure to be H × K equivariant, it is sufficient to have Z H be K invariant and Z K be H invariant. These are special cases of H × K equivariance, and is supported by a range of equivariant neural networks especially regarding S n × O (3) or S n × SO (3) equivariances [23].

2.3 Base Function f θ

We now describe the choice of group-agnostic base function f θ : x ↦→ y . As group symmetry is handled by the equivariant distribution p ω ( g | x ) , any symmetry concern is hidden from f θ , allowing the inputs x and outputs y to be treated as plain multidimensional arrays. This allows us to implement f θ with powerful general-purpose architectures, namely as an MLP that operates on flattened vectors of inputs and outputs, or a transformer as we describe below.

Let inputs x ∈ X = R n 1 × ... × n a × c and outputs y ∈ Y = R n ′ 1 × ... × n ′ b × c ′ be multidimensional arrays with c and c ′ channels, respectively. Our transformer base function f θ : x ↦→ y is given as:

where tokenize : X → R m × d parses input array to a sequence of m tokens, transformer : R m × d → R m × d is a standard transformer encoder on tokens used in language and vision [24, 26], and detokenize : R m × d →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 to higher dimensions [26, 85, 86]. This enables mapping between different dimensional inputs and outputs, e.g , for graph node classification with x ∈ R n × n × c and y ∈ R n × c ′ , it is possible to use 2D patch projection for the input and 1D projection for the output.

Above choice of f θ offers important advantages including universal approximation [104] and ability to share and transfer the learned knowledge in θ over different domains of different group symmetries. Remarkably, this allows us to directly leverage large-scale pre-trained parameters from data-abundant domains for learning on symmetric domains. In our experiments, we only replace the tokenizer and detokenizer of a vision transformer pre-trained on ImageNet-21k [99, 26] and fine-tune it to perform diverse S n equivariant tasks such as graph classification, node classification, and link prediction.

2.4 Relation to Other Symmetrization Approaches

We discuss relation of our symmetrization method to prior ones, specifically group averaging [102, 4], frame averaging [74], and canonicalization [41]. An extended discussion on broader related work can be found in Appendix A.2. Since these symmetrization methods share common formalization y = E g [ g · f θ ( g -1 · x )] , one can expect a close theoretical relationship between them. We observe that probabilistic symmetrization is quite general; based on particular choices of the G equivariant distribution p ω ( g | x ) , it can become most of the related symmetrization methods as special cases. This can be easily seen for group averaging [102], as the distribution p ω 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 a frame F : X → 2 G \ ∅ ; importantly, it is required that the frame itself is G equivariant F ( ρ ( g ) x ) = gF ( x ) . We can make the following connection between our method and frame averaging, by adopting the concept of stabilizer subgroup G x = { g ∈ G : ρ ( g ) x = x } :

Proposition 1. Probabilistic symmetrization with G equivariant distribution p ω ( g | x ) can become frame averaging [74] by assigning uniform density to a set of orbits G x g for some group elements g .

Proof. The proof can be found in Appendix A.1.5.

Canonicalization [41] uses a single group element for symmetrization, produced by a trainable canonicalizer C ω : X → G . Here, it is required that the canonicalizer itself satisfies relaxed G equivariance C ( ρ ( g ) x ) = gg ′ C ( x ) up to arbitrary action from the stabilizer g ′ ∈ G x . We now show:

Proposition 2. Probabilistic symmetrization with G equivariant distribution p ω ( g | x ) can become canonicalization [41] by assigning uniform density to a single orbit G x g of some group element g .

Proof. The proof can be found in Appendix A.1.5.

Assuming that stabilizer G x is trivial, this can be implemented with our method by removing random noise ϵ , which reduces p ω to deterministic map ρ ( g ) = q ω ( x ) . We use this approach to implement canonicalizer for the S n group, while [41] only provides canonicalizers for Euclidean groups.

3 Experiments

We empirically demonstrate and analyze probabilistic symmetrization on a range of symmetry groups S n , E (3) (O (3) ), and the product S n × E (3) (S n × O (3) ), on a variety of invariant and equivariant tasks, with general-purpose base functions f θ 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 n ) has been considered important and challenging, as simple and efficient GNNs are often limited in expressive power to certain Weisfeiler-Lehman isomorphism tests like 1-WL [101, 56]. Since an MLP equipped with probabilistic symmetrization is in theory universal and S n equivariant, it has potential for graph learning that require high expressive power. To explicitly test this, we adopt the experimental setup of [74] and use two datasets on graph separation task (S n invariant). GRAPH8c [2] consists of all non-isomorphic connected graphs with 8 nodes, and EXP [1] consists of 3-WL distinguishable graphs that are not 2-WL distinguishable. We compare our method to standard GNNs as well as an MLP symmetrized with group averaging [102], frame averaging [74], and canonicalization [41]. Our method uses the same MLP architecture to symmetrization baselines, and its S n equivariant distribution p ω for symmetrization is implemented

Entropy

method

arch.

sym.

GRAPH8c

↓

EXP

↓

EXP-classify

↑

GCN [48]

S

n

-

4755

600

50%

GAT [97]

S

n

-

1828

600

50%

GIN [101]

S

n

-

386

600

50%

ChebNet [22]

S

n

-

44

71

82%

PPGN [56]

S

n

-

0

0

100%

GNNML3 [2]

S

n

-

0

0

100%

MLP-GA [102]

-

S

n

0

0

50%

MLP-FA [74]

-

S

n

0

0

100%

MLP-Canonical.

-

S

n

0

0

50%

MLP-PS (Ours), fixed

ϵ

-

S

n

0

0

79.5%

MLP-PS (Ours)

-

S

n

0

0

100%

Epochs

Loss

0

500

1000

2.4

2.6

2.8

3.8

Random permutations

Epochs

0

500

1000

0.1

0.7

0.3

0.5

Epoch 1000

using a 3-layer GIN [101] which is 1-WL expressive. We use 10 samples for symmetrization during both training and testing. Further details can be found in Appendix A.3.3, and supplementary results on symmetrization of a different base model can be found in Appendix A.4.1.

The results are in Table 1. At random initialization, all symmetrization methods can provide perfect separation of all graphs, similar to PPGN [56] and GNNML3 [2] that are equivariant neural networks carefully designed to be 3-WL expressive. However, when trained with gradient descent to solve classification problem, naïve symmetrization with group averaging fails, presumably because the MLP fails to adjust to equally possible 64! permutations of 64 nodes in maximum. On the other hand, our method is able to learn the task, achieving the same accuracy to frame averaging that utilizes costly eigendecomposition of graph Laplacian [74]. What makes our method work while group averaging fails? We conjecture this is since the distribution p ω ( g | x ) can learn to provide more consistent permutations during early training, as we illustrate in Figure 2. In the figure, we measured the consistency of samples from p ω ( g | x ) over training progress by sampling N = 50 permutation matrices P g ∼ p ω ( g | x ) and measuring the row-wise entropy of their average ¯ P = ∑ P g /N for each input x . The more consistent the sampled permutations, the sharper their average, and lesser the entropy. As training progresses, p ω learns to produce more consistent samples, which coincides with the initial increase in task performance. Given that, a natural question would be: if we enforce the samples g ∼ p ω ( g | x ) to be consistent from the first place, would it work? To answer this, we also tested a non-probabilistic version of our model that uses a single permutation per input ρ ( g ) = q ω ( x ) , which is a canonicalizer under relaxed equivariance [41] as described in Section 2.4. As in Table 1, canonicalization fails, suggesting that probabilistic nature of p ω ( g | x ) can be beneficial for learning. For another deterministic version of our model made by fixing the noise ϵ (Eq. (6)) at initialization, the performance drops to 79.5%, further implying that stochasticity of p ω ( g | x ) has a role.

3.2 Particle Dynamics Learning with Transformer

Learning sets or graphs attributed with position and velocity in 3D (S n × E (3) ) is practically significant as they universally appear in physics, chemistry, and biology applications. While prior symmetrization methods employ an already S n equivariant base function and partially symmetrize the E (3) part, we attempt to symmetrize the entire product group S n × E (3) and choose the base model f θ as a sequence transformer to leverage its expressive power. For empirical demonstration, we adopt the experimental setup of [41] and use the n -body dataset [84, 31] where the task is predicting the position of n = 5 charged particles after certain time given their initial position and velocity in R 3 (S n × E (3) equivariant). We compare our method to S n × E (3) equivariant neural networks and partial symmetrization methods applying E (3) symmetrization to GNNs. We also test prior symmetrization methods on the full group S n × E (3) along our method, but could not test for frame averaging since equivariant frames for the full group S n × E (3) was not available in current literature. Our method is implemented using a transformer with sequence positional encodings with around 2.3 × parameters of the baselines, and the S n × E (3) equivariant distribution p ω for symmetrization is implemented using a 2-layer Vector Neurons [23] that has around 0.03 × of parameters to the transformer. We use 20 samples for symmetrization during training, and use 10 × sample size for testing since the task is regression where appropriate variance reduction is necessary to guarantee a reliable performance. Further details can be found in Appendix A.3.4, and supplementary results on E(3) partial symmetrization of GNN base model can be found in Appendix A.4.2.

The results are in Table 2. We observe simple group averaging exhibits a surprisingly strong performance, as it achieves 0.00414 MSE and already outperforms previous state of the art 0.0043 MSE. This is because the permutation component of the symmetry is fairly small, with n = 5 particles interacting with each other, such that combining it with an expressive base model f θ (a sequence transformer) can adjust to 5! = 120 equally possible permutations and their rotations in 3D. Nevertheless, our method outperforms group averaging and achieves a new state of the art 0.00401 MSE, presumably as the parameterized distribution p ω 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 f θ memorizes the per-input canonical orientations provided by ρ ( g ) = q ω ( x ) that do not generalize to test inputs. This again shows that probabilistic nature of p ω ( g | x ) can be beneficial for performance.

3.3 Graph Pattern Recognition with Vision Transformer

One important goal of our approach, and symmetrization in general, is to decouple the symmetry of problem from the base function f θ , such that we can leverage knowledge learned from other symmetries by transferring the parameters θ . We demonstrate an extreme case by transferring the parameters of a vision transformer [26] trained on large-scale image classification (translation invariant) to solve node classification on graphs (S n equivariant) for the first time in literature. For this, we use the PATTERN dataset [27] that contains 14,000 purely topological random SBM graphs with 44-188 nodes, whose task is finding certain subgraph pattern by binary node classification.

Based on pre-trained ViT [26, 89], we construct the base model f θ by modifying only the input and output layers to take the flattened patches of 2D zero-padded adjacency matrices of size 188 × 188 and produce output as 1D per-node classification logits of length 188 with 2 channels. In addition to standard GNNs 4 , we compare group averaging, frame averaging, and canonicalization to our method, and also test whether pre-trained representations from ImageNet-21k is beneficial for the task. For the S n equivariant distribution p ω in our method and canonicalization, we use a 3-layer GIN [101] with only 0.02% of the base model parameters. Further details can be found in Appendix A.3.5.

The results are in Table 3. First, we observe that transferring pre-trained ViT parameters consistently improves node classification for all symmetrization methods. It indicates that some traits of the pre-trained visual representation can benefit learning graph tasks which vastly differ in both the underlying symmetry (translation invariance → S n 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 in this dataset is impossible. We conjecture that group averaged ViT can in some way learn meaningful graph representation internally to solve the task, and vision pretraining helps in acquiring the representation by providing a good initialization point or transferable computation motifs.

On the other hand, frame averaging shows a low performance, 79.637% accuracy with vision pretraining, which is also surprising considering that frames vastly reduce the sample space of symmetrization in general; in fact, the size of frame of each graph in PATTERN is exactly 1. We empirically observe that, unlike group averaging, ViT with frame averaging memorizes frames of training graphs rather than learning generalizable graph representations. In contrast, canonicalization that also uses a single sample per graph successfully learns the task with 86.166% accuracy. We conjecture that the learnable orderings provided by an equivariant neural network ρ ( 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. As a result, our model based on pre-trained ViT and 10 samples for testing achieves 86.285% test accuracy, surpassing all baselines.

3.4 Real-World Graph Learning with Vision Transformer

Having observed that pre-trained ViT can learn graph tasks well when symmetrized with our method, we now provide a preliminary test of it in real-world graph learning. We use three real-world graph datasets from [28] that involve chemical and biological graphs. PCQM-Contact dataset contains 529,434 molecular graphs with 53 nodes in maximum, and the task is contact map prediction framed as link prediction (S n equivariant), on whether two atoms would be proximal when the molecule is in 3D space. Peptides-func and Peptides-struct are based on the same set of 15,535 protein graphs with 444 nodes in maximum and the tasks are property prediction (S n invariant), requiring multilabel classification for Peptides-func and regression for Peptides-struct. The tasks require complex understanding of how the amino acids of the proteins would interact in 3D space. We implement our method using a ViT-Base pre-trained on ImageNet-21k as the base model f θ and a 3-layer GIN as the equivariant distribution p ω ( 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.3.6.

The results are in Table 4. In Peptides-func and PCQM-Contact, the pre-trained ViT symmetrized with our method achieves a significantly higher performance compared to both GNNs and graph transformers, improving previous best by a large margin 5 (0.6527 → 0.8311 for Peptides-func AP, 0.1355 → 0.3268 for PCQM-Contact Hits@1). This demonstrates the scalability of our method as Peptides-func involves 444 maximum nodes, and also its generality as it performs well for both S n invariant (Peptides-func) and equivariant (PCQM-Contact) tasks. We also note that, unlike some baselines, our method does not require costly Laplacian eigenvectors to compute positional encoding. On the Peptides-struct, our method achieves a slightly lower performance to SOTA graph transformers while still better than GNNs. We conjecture that regression is harder for the model to learn due to its stochasticity in predictions, and leave improving regression performance as future work.

4 Conclusion

We presented probabilistic symmetrization, a general framework that learns a distribution of group transformations conditioned on input data for symmetrization of an arbitrary function. By characterizing that the only condition for such distribution is equivariance to data symmetry, we instantiated models for a wide range of groups, including symmetric, orthogonal, Euclidean groups and their product combinations. Our experiments demonstrated that the proposed framework achieves consistent improvement over other symmetrization methods, and is competitive or outperforms equivariant networks on various datasets. We also showed that transferring pre-trained parameters across data in different symmetries can sometimes be surprisingly beneficial. Our approach has weaknesses such as sampling cost, which we further discuss in Appendix A.5; we plan to address these in future work.

Acknowledgements This work was supported in part by the National Research Foundation of Korea (NRF2021R1C1C1012540 and NRF2021R1A4A3032834) and IITP grant (2021-0-00537, 2019-0-00075, and 2021-0-02068) funded by the Korea government (MSIT).

References

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)

M. Balcilar, P. Héroux, B. Gaüzère, P. Vasseur, S. Adam, and P. Honeine. Breaking the limits of message passing graph neural networks. In ICML , 2021. (Cited on 6, 7, 23)

S. Basu, P. Katdare, P. Sattigeri, V . Chenthamarakshan, K. Driggs-Campbell, P. Das, and L. R. Varshney. Equivariant few-shot learning from pretrained models. arXiv , 2023. (Cited on 22)

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. arXiv , 2022. (Cited on 2, 6, 22)

P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. F. Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, Ç. Gülçehre, H. F. Song, A. J. Ballard, J. 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 1, 22)

Y. Bengio, N. Léonard, and A. C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv , abs/1308.3432, 2013. (Cited on 4, 24)

G. W. Benton, M. Finzi, P. Izmailov, and A. G. Wilson. Learning invariances in neural networks from training data. In NeurIPS , 2020. (Cited on 22, 23)

B. Bevilacqua, F. Frasca, D. Lim, B. Srinivasan, C. Cai, G. Balamurugan, M. M. Bronstein, and H. Maron. Equivariant subgraph aggregation networks. In ICLR , 2022. (Cited on 22)

B. Bloem-Reddy and Y. W. Teh. Probabilistic symmetries and invariant neural networks. J. Mach. Learn. Res. , 2020. (Cited on 2)

A. Bogatskiy, S. Ganguly, T. Kipf, R. 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 1, 22)

X. Bresson and T. Laurent. Residual gated graph convnets. arXiv , 2019. (Cited on 9, 10)

M. M. Bronstein, J. Bruna, T. Cohen, and P. Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv , 2021. (Cited on 1, 22)

T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. 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)

C. Cai and Y. Wang. A note on over-smoothing for graph neural networks. arXiv , 2020. (Cited on 4, 22)

M. Chen, Z. Wei, Z. Huang, B. Ding, and Y. Li. Simple and deep graph convolutional networks. In ICML , 2020. (Cited on 10)

Z. Chen, S. Villar, L. Chen, and J. Bruna. On the equivalence between graph isomorphism testing and function approximation with gnns. In NeurIPS , 2019. (Cited on 9)

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)

T. Cohen and M. Welling. Group equivariant convolutional networks. In ICML , 2016. (Cited on 22)

T. S. Cohen and M. Welling. Steerable cnns. In ICLR , 2017. (Cited on 22)

G. Cybenko. Approximation by superpositions of a sigmoidal function. Math. Control. Signals Syst. , 1989. (Cited on 3, 22)

G. Dasoulas, L. D. Santos, K. Scaman, and A. Virmaux. Coloring graph neural networks for node disambiguation. In IJCAI , 2020. (Cited on 22)

M. Defferrard, X. Bresson, and P. Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. In NeurIPS , 2016. (Cited on 7)

C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neurons: A general framework for so(3)-equivariant networks. In ICCV , 2021. (Cited on 1, 5, 8, 22, 24)

J. Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT , 2019. (Cited on 5, 22)

T. Dinh, Y. Zeng, R. Zhang, Z. Lin, M. Gira, S. Rajput, J. Sohn, D. S. Papailiopoulos, and K. Lee. LIFT: language-interfaced fine-tuning for non-language machine learning tasks. In NeurIPS , 2022. (Cited on 22)

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, 5, 6, 8, 9, 25, 26)

V. P. Dwivedi, C. K. Joshi, T. Laurent, Y. Bengio, and X. Bresson. Benchmarking graph neural networks. arXiv , 2020. (Cited on 8, 9, 23, 25)

V. P. Dwivedi, L. Rampásek, M. Galkin, A. Parviz, G. Wolf, A. T. Luu, and D. Beaini. Long range graph benchmark. In NeurIPS , 2022. (Cited on 10, 23)

W. Falcon. Pytorch lightning. https://github.com/Lightning-AI/lightning , 2019. (Cited on 25, 26)

M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In ICML , 2021. (Cited on 4, 19, 22)

F. Fuchs, D. E. Worrall, V. Fischer, and M. Welling. Se(3)-transformers: 3d roto-translation equivariant attention networks. In NeurIPS , 2020. (Cited on 8, 22, 23)

J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In ICML , 2017. (Cited on 23)

A. Grover, E. Wang, A. Zweig, and S. Ermon. Stochastic optimization of sorting networks via continuous relaxations. In ICLR , 2019. (Cited on 4, 24)

D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). arXiv , 2018. (Cited on 25)

K. Hornik, M. B. Stinchcombe, and H. White. Multilayer feedforward networks are universal approximators. Neural Networks , 1989. (Cited on 2, 3, 22)

W. Hu, B. Liu, J. Gomes, M. Zitnik, P. Liang, V. S. Pande, and J. Leskovec. Strategies for pre-training graph neural networks. In ICLR , 2020. (Cited on 10)

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, 23)

A. Jaegle, S. Borgeaud, J. Alayrac, C. Doersch, C. Ionescu, D. Ding, S. Koppula, D. Zoran, A. Brock, E. Shelhamer, O. J. Hénaff, M. M. Botvinick, A. Zisserman, O. Vinyals, and J. Carreira. Perceiver IO: A general architecture for structured inputs & outputs. In ICLR , 2022. (Cited on 22)

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)

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 1, 4, 22)

S. Kaba, A. K. Mondal, Y. Zhang, Y. Bengio, and S. Ravanbakhsh. Equivariance with learned canonical-

ization functions. arXiv , 2022. (Cited on 2, 4, 5, 6, 7, 8, 15, 17, 18, 19, 21, 22, 26, 27, 29, 30)

P. Kicki, P. Skrzypczynski, and M. Ozay. A new approach to design symmetry invariant neural networks. In IJCNN , 2021. (Cited on 2, 22)

J. Kim, D. Nguyen, S. Min, S. Cho, M. Lee, H. Lee, and S. Hong. Pure transformers are powerful graph learners. In NeurIPS , 2022. (Cited on 22)

J. Kim, S. Oh, S. Cho, and S. Hong. Equivariant hypergraph neural networks. In ECCV , 2022. (Cited on 22)

J. Kim, S. Oh, and S. Hong. Transformers generalize deepsets and can be extended to graphs and hypergraphs. In NeurIPS , 2021. (Cited on 1, 22)

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR , 2015. (Cited on 25)

D. P. Kingma and M. Welling. Auto-encoding variational bayes. In ICLR , 2014. (Cited on 2)

T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In ICLR , 2017. (Cited on 7, 9, 10)

J. Köhler, L. Klein, and F. Noé. Equivariant flows: Sampling configurations for multi-body systems with symmetric energies. arXiv , 2019. (Cited on 8)

D. Kreuzer, D. Beaini, W. L. Hamilton, V. Létourneau, and P. Tossou. Rethinking graph transformers with spectral attention. NeurIPS , 2021. (Cited on 10, 22)

J. Lee, Y. Lee, J. Kim, A. R. Kosiorek, S. Choi, and Y. W. Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In ICML , 2019. (Cited on 22)

J. Li, D. Li, S. 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)

I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In ICLR , 2019. (Cited on 25, 26)

K. Lu, A. Grover, P. Abbeel, and I. Mordatch. Frozen pretrained transformers as universal computation engines. In AAAI , 2022. (Cited on 22)

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)

H. Maron, H. Ben-Hamu, H. Serviansky, and Y. Lipman. Provably powerful graph networks. In NeurIPS , 2019. (Cited on 1, 6, 7, 9, 22)

H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. In ICLR , 2019. (Cited on 19, 22)

H. Maron, E. Fetaya, N. Segol, and Y. Lipman. On the universality of invariant networks. In ICML , 2019. (Cited on 1)

H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In ICML , 2020. (Cited on 19)

B. D. McKay, M. A. Yirik, and C. Steinbeck. Surge: a fast open-source chemical graph generator. J. Cheminformatics , 2022. (Cited on 31)

G. E. Mena, D. Belanger, S. W. Linderman, and J. Snoek. Learning latent permutations with gumbelsinkhorn networks. In ICLR , 2018. (Cited on 4, 24)

F. Mezzadri. How to generate random matrices from the classical compact groups. Notices of the American Mathematical Society , 2006. (Cited on 2)

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 22)

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 1, 22)

S. C. Mouli and B. Ribeiro. Neural networks for learning counterfactual g-invariances from single environments. In ICLR , 2021. (Cited on 2, 22)

L. Müller, M. Galkin, C. Morris, and L. Rampásek. Attending to graph transformers. arXiv , 2023. (Cited on 22)

R. L. Murphy, B. Srinivasan, V. Rao, and B. Ribeiro. Janossy pooling: Learning deep permutationinvariant functions for variable-size inputs. In ICLR , 2019. (Cited on 1, 2, 22, 28)

R. L. Murphy, B. Srinivasan, V. A. Rao, and B. Ribeiro. Relational pooling for graph representations. In ICML , 2019. (Cited on 2, 22, 28)

H. NT and T. Maehara. Revisiting graph neural networks: All we have is low-pass filters. arXiv , 2019. (Cited on 22)

K. Oono and T. Suzuki. Graph neural networks exponentially lose expressive power for node classification. In ICLR , 2020. (Cited on 4, 22)

F. Paischer, T. Adler, V. P. Patil, A. Bitto-Nemling, M. Holzleitner, S. Lehner, H. Eghbal-Zadeh, and S. Hochreiter. History compression via language models in reinforcement learning. In ICML , 2022. (Cited on 22)

A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. 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)

S. Prillo and J. M. Eisenschlos. Softsort: A continuous relaxation for the argsort operator. In ICML , 2020. (Cited on 4)

O. Puny, M. Atzmon, E. J. Smith, I. Misra, A. Grover, H. Ben-Hamu, and Y. Lipman. Frame averaging for invariant and equivariant network design. In ICLR , 2022. (Cited on 2, 5, 6, 7, 8, 15, 18, 19, 21, 22, 24, 26, 27)

C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 2020. (Cited on 22)

L. Rampásek, M. Galkin, V. P. Dwivedi, A. T. Luu, G. Wolf, and D. Beaini. Recipe for a general, powerful, scalable graph transformer. In NeurIPS , 2022. (Cited on 10)

M. Rath and A. P. Condurache. Boosting deep neural networks with geometrical prior knowledge: A survey. arXiv , 2020. (Cited on 22)

S. Ravanbakhsh, J. G. Schneider, and B. Póczos. Equivariance through parameter-sharing. In ICML , 2017. (Cited on 22)

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)

D. W. Romero and S. Lohit. Learning partial equivariances from data. In NeurIPS , 2022. (Cited on 22)

C. Rommel, T. Moreau, and A. Gramfort. Deep invariant networks with differentiable augmentation layers. In NeurIPS , 2022. (Cited on 22)

D. Rothermel, M. Li, T. Rocktäschel, and J. Foerster. Don't sweep your learning rate under the rug: A closer look at cross-modal transfer of pretrained transformers. arXiv , 2021. (Cited on 22)

A. Sannai, M. Kawano, and W. Kumagai. Equivariant and invariant reynolds networks. arXiv , 2021. (Cited on 2)

V. G. Satorras, E. Hoogeboom, and M. Welling. E(n) equivariant graph neural networks. In ICML , 2021. (Cited on 8, 22, 23)

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 5, 22)

W. Shi, J. Caballero, F. Huszar, J. Totz, A. P. Aitken, R. Bishop, D. 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 5)

H. Shirzad, A. Velingker, B. Venkatachalam, D. J. Sutherland, and A. K. Sinop. Exphormer: Sparse transformers for graphs. arXiv , 2023. (Cited on 10)

B. Srinivasan and B. Ribeiro. On the equivalence between positional node embeddings and structural graph representations. In ICLR , 2020. (Cited on 22, 31)

A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and L. Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv , 2021. (Cited on 2, 9)

E. H. Thiede, W. Zhou, and R. Kondor. Autobahn: Automorphism-based graph neural nets. In NeurIPS , 2021. (Cited on 30)

N. Thomas, T. E. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds. arXiv , 2018. (Cited on 1, 8, 22)

J. Topping, F. D. Giovanni, B. P. Chamberlain, X. Dong, and M. M. Bronstein. Understanding oversquashing and bottlenecks on graphs via curvature. In ICLR , 2022. (Cited on 22)

T. F. A. van der Ouderaa and M. van der Wilk. Learning invariant weights in neural networks. In J. Cussens and K. Zhang, editors, UAI , 2022. (Cited on 22)

E. van der Pol, D. E. Worrall, H. van Hoof, F. A. Oliehoek, and M. Welling. MDP homomorphic networks: Group symmetries in reinforcement learning. In NeurIPS , 2020. (Cited on 2, 22)

M. van der Wilk, M. Bauer, S. T. John, and J. Hensman. Learning invariances using the marginal likelihood. In S. Bengio, H. M. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, NeurIPS , 2018. (Cited on 22)

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In NeurIPS , 2017. (Cited on 22, 25)

P. Velickovic, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio. Graph attention networks. In ICLR , 2018. (Cited on 7, 9)

R. Winter, F. Noé, and D. Clevert. Permutation-invariant variational autoencoder for graph-level representation learning. In NeurIPS , 2021. (Cited on 4, 17, 24)

T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, and J. Brew. Huggingface's transformers: State-of-the-art natural language processing. arXiv , 2019. (Cited on 6, 25)

R. Xiong, 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 on 25)

K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In ICLR , 2019. (Cited on 1, 4, 6, 7, 9, 22, 23)

D. Yarotsky. Universal approximations of invariant maps by neural networks. arXiv , 2018. (Cited on 1, 2, 6, 7, 15, 22)

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 NeurIPS , 2021. (Cited on 22)

C. Yun, S. Bhojanapalli, A. S. Rawat, S. J. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequence functions? In ICLR , 2020. (Cited on 2, 3, 6, 22)

B. Zhang, S. Luo, L. Wang, and D. He. Rethinking the expressive power of gnns via graph biconnectivity. arXiv , 2023. (Cited on 1, 22)

A Appendix

A.1 Proofs

A.1.1 Proof of Theorem 1 (Section 2.1)

Theorem 1. If p ω is G equivariant, then ϕ θ,ω is G equivariant for arbitrary f θ .

Proof. We prove ϕ θ,ω ( ρ 1 ( g ′ ) x ) = ρ 2 ( g ′ ) ϕ θ,ω ( x ) for all x ∈ X and g ′ ∈ G . From Eq. (4), we have:

Let us introduce transformed random variable h = g ′-1 g ∈ G such that g = g ′ h . Since the distribution p ω is G equivariant, we can see that p ω ( g | ρ 1 ( g ′ ) x ) = p ω ( g ′-1 g | ρ 1 ( g ′-1 ) ρ 1 ( g ′ ) x ) = p ω ( g ′-1 g | x ) = p ω ( h | x ) . Thus, we can rewrite the above expectation with respect to h as follows:

showing the G equivariance of ϕ θ,ω for arbitrary f θ .

A.1.2 Proof of Theorem 2 (Section 2.1)

Theorem 2. If p ω is G equivariant and f θ is a universal approximator, then ϕ θ,ω is a universal approximator of G equivariant functions.

Proof. The proof is inspired by universality proofs of prior symmetrization approaches [102, 74, 41]. Let ψ : X → Y be an arbitrary G equivariant function. By equivariance of ψ , we have:

As Y is finite-dimensional, we can assume that the linear operators in GL ( Y ) are bounded and so is the induced operator norm of group representation ∥ ρ 2 ( g ) ∥ for all g ∈ G . Thus, we have:

for some c > 0 . If f θ is a universal approximator, for any compact set K ⊆ X and any ϵ > 0 , there exists some θ such that ∥ ψ ( x ) -f θ ( x ) ∥ ≤ ϵ for all x ∈ K . Consider the set K sym = ∪ g ∈ G ρ 1 ( g ) K where ρ 1 ( g ) K denotes the image of the set K under linear transformation by ρ 1 ( g ) . We use the fact that K sym is also a compact set since it is the image of the compact set G ×K under continuous map ( g, x ) ↦→ ρ 1 ( g ) x . As a consequence, for any compact set K ⊆ X and any ϵ/c > 0 , there exists some θ such that max g ∈ G ∥ ψ ( ρ 1 ( g ) x ) -f θ ( ρ 1 ( g ) x ) ∥ ≤ ϵ/c for all x ∈ K . Since a group is closed under inverse, for any compact set K ⊆ X and any ϵ > 0 , there exists some θ such that:

for all x ∈ K , showing that ϕ θ,ω is a universal approximator of G equivariant functions.

While we have assumed that the group G is compact in the proof, we conjecture that the results can be extended to non-compact groups if we make an alternative assumption that the distribution p ω ( g | x ) is compactly supported for all x ∈ K . We leave proving this as a future work.

A.1.3 Proof of Theorem 3 (Section 2.1)

Theorem 3. If q ω is G equivariant and p ( ϵ ) is G invariant under representation ρ ′ that | det ρ ′ ( g ) | = 1 ∀ g ∈ G , the distribution p ω ( g | x ) characterized by q ω : ( x , ϵ ) ↦→ ρ ( g ) is G equivariant.

Proof. We prove p ω ( g ′ g | ρ 1 ( g ′ ) x ) = p ω ( g | x ) for all x ∈ X and g, g ′ ∈ G . In general, we are interested in obtaining a faithful representation ρ , i.e. , such that ρ ( g ) is distinct for each g . We can interpret the probability p ω ( g | x , ϵ ) as a delta distribution centered at the group representation ρ ( g ) :

To obtain p ω ( g | x ) , we marginalize over p ( ϵ ) :

Let us consider p ω ( g ′ g | ρ 1 ( g ′ ) x ) :

Using the G equivariance of q ω , we have:

which leads to the following:

Note that the second equality follows from invertibility of ρ ( g ′ ) . We now introduce a change of variables ϵ ′ = ρ ′ ( g ′-1 ) ϵ that ϵ = ρ ′ ( g ′ ) ϵ ′ :

With | det ρ ′ ( g ′-1 ) | = 1 , and G invariance of p ( ϵ ) which gives p ( ρ ′ ( g ′ ) ϵ ′ ) = p ( ϵ ′ ) , we get:

showing the G equivariance of p ω ( g | x ) .

A.1.4 Proof of Validity for Implemented Equivariant Distributions p ω (Section 2.2)

We formally show G equivariance of the implemented distributions p ω ( g | x ) presented in Section 2.2. All implementations have a form of noise-outsourced function q ω : ( x , ϵ ) ↦→ ρ ( g ) using distribution ϵ ∼ p ( ϵ ) and map q ω which is composed of G equivariant neural network and postprocessing to ρ ( g ) . From Theorem 3, for G equivariance of p ω ( g | x ) , it is sufficient to show G invariance of p ( ϵ ) under a representation ρ ′ such that | det ρ ′ ( g ) | = 1 along with G equivariance of q ω , which we show below.

Symmetric Group S n We recall that p ω ( g | x ) for the symmetric group S n is implemented as below:

Sample node-level noise ϵ ∈ R n × d from i.i.d. uniform Unif [0 , η ] .

Use a GNN to obtain node-level scalar features ( x , ϵ ) ↦→ Z ∈ R n .

Assuming Z is tie-free, use argsort [98] to obtain group representation Z ↦→ P g = ρ ( g ) .

where eq denotes elementwise equality indicator.

We now show the following:

Proposition 3. The proposed distribution p ω ( g | x ) for the symmetric group S n is equivariant.

Proof. Given p ( ϵ ) is elementwise i.i.d., it is S n invariant under the base representation ρ ′ ( g ) = P g which satisfies | det P g | = 1 from orthogonality. As a GNN is S n equivariant, we only need to show S n equivariance of argsort : Z ↦→ P g . This can be shown by transforming Z with any permutation matrix P g ′ . Since sort operator and any row replicated matrices are invariant to P g ′ , we have:

Since eq commutes with P g ′ , we have:

showing that argsort is S n equivariant, i.e. , it maps P g ′ Z ↦→ P g ′ P g for all P g ′ ∈ S n . Combining the above, by Theorem 3, the distribution p ω ( g | x ) is S n equivariant.

Orthogonal Group O ( n ) , SO ( n ) We recall that p ω ( g | x ) for the orthogonal group O ( n ) or special orthogonal group SO ( n ) is implemented as follows:

Sample noise ϵ ∈ R n × d from i.i.d. normal N (0 , η 2 ) .

Use an O ( n ) /SO ( n ) equivariant neural network to obtain n features ( x , ϵ ) ↦→ Z ∈ R n × n .

Assuming Z is full-rank, use Gram-Schmidt process [41] to obtain an orthogonal matrix Z ↦→ Q .

For the O ( n ) group, use the obtained matrix as group representation Q = Q g = ρ ( g ) .

For the SO ( n ) group, use below scale operator to obtain group representation Q ↦→ Q + g = ρ ( g ) .

We now show the following:

Proposition 4. The proposed distribution p ω ( g | x ) for the orthogonal group O ( n ) is equivariant.

Proof. Without loss of generality, let us omit the scale η for brevity, which gives that each column ϵ i ∈ R n of the noise ϵ independently follows multivariate standard normal ϵ i ∼ N (0 , I n ) . Then, the density p ( ϵ i ) = (2 π ) -n/ 2 exp( -∥ ϵ i ∥ 2 2 / 2) is invariant under orthogonal transformation Q since ∥ Q ϵ i ∥ 2 2 = ( Q ϵ i ) ⊤ Q ϵ i = ϵ ⊤ i Q ⊤ Q ϵ i = ϵ ⊤ i ϵ i = ∥ ϵ i ∥ 2 2 . Therefore, the distribution p ( ϵ ) is invariant under the base representation ρ ′ ( g ) = Q g which satisfies | det ρ ′ ( g ) | = 1 from orthogonality. As we use an equivariant neural network to obtain Z , and Gram-Schmidt procedure Z ↦→ Q g is O ( n ) equivariant (Theorem 5 of [41]), by Theorem 3, the distribution p ω ( g | x ) is O ( n ) equivariant.

Proposition 5. The proposed distribution p ω ( g | x ) for special orthogonal group SO ( n ) is equivariant.

Proof. From the proof of Proposition 4, it follows that the distribution p ( ϵ ) is invariant under the base representation ρ ′ ( g ) = Q + g which satisfies | det ρ ′ ( g ) | = 1 due to orthogonality. As we use an equivariant neural network to obtain Z , and Gram-Schmidt procedure Z ↦→ Q has O ( n ) equivariance which implies SO ( n ) equivariance because of SO ( n ) ≤ O ( n ) , we only need to show

SO ( n ) equivariance of scale : Q ↦→ Q + g . This can be done by transforming Q with an orthogonal Q + g ′ of determinant +1 . Since det ( Q + g ′ Q ) = det ( Q + g ′ ) det ( Q ) = det ( Q ) , we have the following:

Also, scaling the first column of the product Q + g ′ Q with det ( Q ) is equivalent to scaling the first column of Q with det ( Q ) then computing the product since ( Q + g ′ Q ) ij = ∑ k Q + g ′ ik Q kj . This gives:

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 Q 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 p ω ( g | x ) is SO ( n ) equivariant.

Euclidean Group E ( 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 E ( n ) = O ( n ) ⋉ T ( n ) and SE ( n ) = SO ( n ) ⋉ T ( n ) is non-compact. This is done as follows:

where ¯ x ∈ R n is centroid (mean over channels) of data x ∈ R n × d and distribution p ω is O ( n ) /SO ( n ) equivariant for E ( n ) /SE ( n ) equivariant symmetrization, respectively. We now show the following:

Proposition 6. The proposed symmetrization ϕ θ,ω for the Euclidean group E ( n ) is equivariant.

Proof. We prove ϕ θ,ω ( g ′ · x ) = g ′ · ϕ θ,ω ( x ) for all x ∈ X and g ′ ∈ E ( n ) . From Eq. (29), we have:

In general, an element of Euclidean group g ′ ∈ E ( n ) acts on data x ∈ R n × d via g ′ · x = Q g ′ x + t g ′ 1 ⊤ where Q g ′ ∈ O ( n ) is its rotation component and t g ′ ∈ R n is its translation component [74, 41]. With this, the centroid of the transformed data g ′ · x is given as follows:

which leads to the following:

Above shows that subtracting centroid eliminates the translation component of the problem and leaves O ( n ) equivariance component. Based on that, we have the following:

Note that, inside the expectation, we interpret the rotation component of g ′ as an element of the orthogonal group O ( n ) . Similar as in the proof of Theorem 1, we introduce transformed random variable h = g ′-1 g ∈ O ( n ) that g = g ′ h . Since the distribution p ω is O ( n ) equivariant, we can see

that p ω ( g | g ′ · ( x -x1 ⊤ )) = p ω ( g ′-1 g | g ′-1 g ′ · ( x -x1 ⊤ )) = p ω ( g ′-1 g | x -x1 ⊤ ) = p ω ( h | x -x1 ⊤ ) . Thus we can rewrite the above expectation with respect to h as follows:

showing the E ( n ) equivariance of ϕ θ,ω

Proposition 7. The proposed symmetrization ϕ θ,ω for special Euclidean group SE ( n ) is equivariant.

Proof. We prove ϕ θ,ω ( g ′ · x ) = g ′ · ϕ θ,ω ( x ) for all x ∈ X and g ′ ∈ SE ( n ) , in an analogous manner to the proof of Proposition 6. From Eq. (29), we have:

In general, an element of special Euclidean group g ′ ∈ SE ( n ) acts on data x ∈ R n × d via g ′ · x = Q + g ′ x + t g ′ 1 ⊤ where Q + g ′ ∈ SO ( n ) is rotation component and t g ′ ∈ R n is translation [74, 41]. With this, the centroid of the transformed data g ′ · x is given as follows:

which leads to the following:

Similar as in Proposition 6, subtracting centroid only leaves SO ( n ) component. We then have:

where, inside the expectation, we interpret the rotation component of g ′ as an element of the special orthogonal group SO ( n ) . Similar as in Theorem 1, we introduce h = g ′-1 g ∈ SO ( n ) that g = g ′ h . As the distribution p ω is SO ( n ) equivariant, we have p ω ( g | g ′ · ( x -x1 ⊤ )) = p ω ( g ′-1 g | g ′-1 g ′ · ( x -x1 ⊤ )) = p ω ( h | x -x1 ⊤ ) . We then rewrite the expectation with respect to h :

showing the SE ( n ) equivariance of ϕ θ,ω .

Product Group H × K For the product group H × K , we assume that the base representation for each element g = ( h, k ) is given as a pair of representations ρ ( g ) = ( ρ ( h ) , ρ ( k )) . Without loss of generality, we further assume that the representation ρ ( g ) can be expressed as the Kronecker product ρ ( g ) = ρ ( h ) ⊗ ρ ( k ) that acts on flattened data vec ( x ) as x ↦→ vec -1 ( ρ ( 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 n × O ( d ) / SO ( d ) 6 ) and sets of symmetric elements (S n × H ) in general [59]. For

)

G

G

x

g

1

G

x

g

2

G

x

g

3

G

x

g

4

G

x

g

5

G

x

g

2

G

x

g

3

G

x

g

4

G

x

g

2

G

x

g

1

G

x

g

2

G

x

g

3

G

x

g

4

G

x

g

5

Group

G

Orbits

G

x

g

Distribution

p

ω

(

g

|

x

)

Frame

F

(

x

)

C

ω

(

x

example, for the group S n × O ( d ) on data x ∈ R n × d , an element g = ( h, k ) has representation ρ ( g ) = ρ ( h ) ⊗ ρ ( k ) ∈ R nd × nd combined from permutation ρ ( h ) ∈ R n × n and rotation ρ ( k ) ∈ R d × d , which acts by x ↦→ vec -1 ( ρ ( g ) vec ( x )) or more simply x ↦→ ρ ( h ) x ρ ( k ) ⊤ .

Now we recall that the p ω ( g | x ) for the product group H × K is implemented as follows:

Sample noise ϵ ∈ E from i.i.d. normal N (0 , η 2 ) such that p ( ϵ ) is invariant under representations of H and K that satisfy | det ρ ′ ( h ) | = 1 and | det ρ ′ ( k ) | = 1 , respectively. For example, for S n × O ( d ) , the noise ϵ ∈ R n × d that follows i.i.d. normal N (0 , η 2 ) is invariant under base representations of both S n and O ( d ) which are orthogonal.

Use a H × K equivariant neural network to obtain features ( x , ϵ ) ↦→ ( Z H , Z K ) where Z H is K invariant and Z K is H invariant. For example, for S n × O ( d ) , we expect node-level scalar features Z S n ∈ R n to be O ( d ) invariant and d global rotary features Z O ( d ) ∈ R d × d to be S n invariant.

Apply postprocessing for H and K groups onto Z H and Z K respectively to obtain representations Z H ↦→ ρ ( h ) and Z K ↦→ ρ ( k ) of H and K groups respectively. For example, for S n × O ( d ) , we use argsort in Eq. (23) to obtain Z S n ↦→ ρ ( h ) and Gram-Schmidt process to obtain Z O ( d ) ↦→ ρ ( k ) .

Combine the representations ρ ( g ) = ( ρ ( h ) , ρ ( k )) to obtain a representation for the H × K group.

We now show the following:

Proposition 8. The proposed distribution p ω ( g | x ) for the product group H × K is equivariant.

Proof. By assumption, p ( ϵ ) is invariant under representations of H and K that satisfy | det ρ ′ ( h ) | = 1 and | det ρ ′ ( k ) | = 1 , respectively. This implies H × K invariance as well, since p ( ϵ ) = p ( h · ϵ ) = p ( k · ϵ ) for all ϵ ∈ E , h ∈ H,k ∈ K gives p ( k · h · ϵ ) = p ( k · ( h · ϵ )) = p ( h · ϵ ) = p ( ϵ ) , and Kronecker product of matrices of determinant 1 gives a matrix of determinant 1 . Furthermore, the map ( x , ϵ ) ↦→ ( ρ ( h ) , ρ ( k )) = ρ ( g ) is overall H × K equivariant, since an input transformed with g ′ = ( h ′ , k ′ ) is first mapped by the equivariant neural network as ( g ′ · x , g ′ · ϵ ) ↦→ ( h ′ · Z H , k ′ · Z K ) , then postprocessed as ( h ′ · Z H , k ′ · Z K ) ↦→ ( ρ ( h ′ ) ρ ( h ) , ρ ( k ′ ) ρ ( k )) = ( ρ ( h ′ ) , ρ ( k ′ )) · ( ρ ( h ) , ρ ( k )) = ρ ( g ′ ) ρ ( g ) . Combining the above, by Theorem 3, the distribution p ω ( g | x ) is H × K equivariant.

A.1.5 Proof of Proposition 1 and Proposition 2 (Section 2.4)

Before proceeding to proofs, we recall that the stabilizer subgroup G x of a group G for x is defined as { g ′ ∈ G : g ′ · x = x } and acts on a given group element g ∈ G through left multiplication g ↦→ g ′ g . For some g ∈ G , by G x g we denote its orbit under the action by G x , i.e. , the set of elements in G to which g can be moved by the action of elements g ′ ∈ G x . Importantly, we can show the following:

Property 1. Any group G is a union of disjoint orbits G x g of equal cardinality.

Proof. Let us consider the equivalence relation ∼ on G induced by the action of the stabilizer G x , defined as g ∼ h ⇐⇒ h ∈ G x g . The orbits G x g are the equivalence classes under this relation, and the set of all orbits of G under the action of G x forms a partition of G ( i.e. , the quotient G/G x ). Furthermore, since G x ≤ G and right multiplication by some g ∈ G is a faithful action of G on itself, we have | G x g | = | G x | for all g ∈ G , which shows that all orbits G x g have equal cardinality.

The partition of group G into disjoint orbits G x g is illustrated in the first and second panel of Figure 3. We now show the following:

Property 2. G equivariant p ω ( g | x ) assigns identical probability to all elements on each orbit G x g .

Proof. With equivariance, we have p ω ( g | x ) = p ω ( g ′ g | g ′ · x ) . Since g ′ · x = x for all g ′ ∈ G x , we have p ω ( g | x ) = p ω ( g ′ g | x ) for all g ′ ∈ G x ; all elements on orbit G x g have an identical probability.

Property 2 characterizes probability distributions over G that can be expressed with p ω ( g | x ) , which we illustrate in the third panel of Figure 3. Intuitively, p ω ( g | x ) assigns constant probability densities over each of the orbit G x g that partitions G as shown in Property 1. We now prove Proposition 1 and Proposition 2 by showing that p ω ( g | x ) can become frame and canonicalizer as special cases:

Proposition 1. Probabilistic symmetrization with G equivariant distribution p ω ( g | x ) can become frame averaging [74] by assigning uniform density to a set of orbits G x g for some group elements g .

Proof. A frame is defined as a set-valued function F : X → 2 G \ ∅ that satisfies G equivariance F ( g · x ) = gF ( x ) [74]. For some frame F , frame averaging is defined as follows:

which can be equivalently written as the below expectation:

From Theorem 3 of [74], we have that F ( x ) is a disjoint union of equal size orbits G x g . Therefore, Unif ( F ( x )) is a uniform probability distribution over the union of the orbits. This can be expressed by a G equivariant distribution p ω ( g | x ) by assigning identical probability over all orbits in the frame F and zero probability to all orbits not in the frame (illustrated in the fourth panel of Figure 3). Therefore, probabilistic symmetrization can become frame averaging.

Proposition 2. Probabilistic symmetrization with G equivariant distribution p ω ( g | x ) can become canonicalization [41] by assigning uniform density to a single orbit G x g of some group element g .

Proof. A canonicalizer is defined as a (possibly stochastic) parameterized map C ω : X → G that satisfies relaxed G equivariance C ω ( g · x ) = gg ′ C ω ( x ) for some g ′ ∈ G x [41]. For some canonicalizer C ω , canonicalization is defined as follows:

From relaxed G equivariance, we have C ω ( x ) = g ′ C ω ( x ) for some g ′ ∈ G x . A valid choice for the canonicalizer C ω is a stochastic map that samples from the uniform distribution over a frame C ω ( x ) ∼ Unif ( F ω ( x )) where the frame is assumed to always provide a single orbit F ω ( x ) = G x g . In this case, canonicalization is equivalent to a 1-sample estimation of the below expectation:

Furthermore, uniform distribution over the single-orbit frame Unif ( F ω ( x )) can be expressed by a G equivariant distribution p ω ( g | x ) by assigning nonzero probability to the single orbit G x g and assigning zero probability to the rest (illustrated in the last panel of Figure 3). Therefore, probabilistic symmetrization can become canonicalization.

A.2 Extended Related Work (Continued from Section 2.4)

Our work draws inspiration from an extensive array of prior research, ranging from equivariant architectures and symmetrization to general-purpose deep learning with transformers. This section outlines a comprehensive review of these fields, spotlighting ideas specifically relevant to our work.

Equivariant Architectures Equivariant architectures, defined by the group equivariance of their building blocks, have been a prominent approach for equivariant deep learning [12, 10]. These architectures have been primarily developed for data types associated with permutation and Euclidean group symmetries, including images [18, 19], sets, graphs, and hypergraphs [5, 57, 8], and geometric graphs [23, 84, 91]. Additionally, they have been extended to more general data types under arbitrary finite group [78] and matrix group symmetries [30]. However, they face challenges such as limited expressive power [101, 56, 64, 105, 40] and architectural issues like over-smoothing [70, 14, 69] and over-squashing [92] in graph neural networks. Our work aims to develop an equivariant deep learning approach that relies less on equivariant architectures, to circumvent these limitations and enhance parameter sharing and transfer across varying group symmetries.

Symmetrization Our approach is an instance of symmetrization for equivariant deep learning which aims to achieve group equivariance using base models with unconstrained architectures. This is in general accomplished by averaging over specific group transformations of the input and output such that the averaged output exhibits equivariance. This allows us to leverage the expressive power of the base model e.g. , achieve universal approximation using an MLP [35, 20] or a transformer [104], and potentially share or transfer parameters across different group symmetries. Existing literature has explored the choices of group transformations and base models for symmetrization. A straightforward approach is to average over the entire group [102], which is suitable for small, finite groups [4, 65, 42, 94] and requires sampling-based estimation for large groups such as permutations [67, 68, 88, 21]. Recent studies have attempted to identify smaller, input-dependent subsets of the group for averaging. Frame averaging [74] employs manually discovered set-values functions called frames, which still demand sampling-based estimation for certain worst-case inputs. Canonicalization [41] utilizes a single group transformation predicted by a neural network, but sacrifices strict equivariance. Our approach jointly achieves equivariance and end-to-end learning by utilizing parameterized, inputconditional equivariant distributions. Furthermore, our approach is one of the first demonstrations of symmetrization for the permutation group in real-world graph recognition task. Concerning the base model, previous work mostly examined small base models like an MLP or partial symmetrization of already equivariant models like GNNs. Few studies have explored symmetrizing pre-trained models for small finite groups [4, 3], and to our knowledge, we are the first to investigate symmetrization of a pre-trained standard transformer for permutation groups or any large group generally.

Transformer Architectures A significant motivation of our work is to combine the powerful scaling and transfer capabilities of the standard transformer architecture [96] with equivariant deep learning. The transformer architecture has driven major breakthroughs in language and vision domains [96, 24, 13, 75], and proven its ability to learn diverse modalities [39, 38] or transfer knowledge across them [85, 54, 82, 25, 71, 52]. Although transformer-style architectures have been developed for symmetric data modalities like sets [51], graphs [103, 45, 43, 50, 66, 63], hypergraphs [17, 44], and geometric graphs [31, 55], they often require specific architectural modifications to achieve equivariance to the given symmetry group, compromising full compatibility with transformer architectures used in language and vision domains. Apart from a few studies on linguistic graph encoding with language models [79], we believe we are the first to propose a general framework that facilitates full compatibility of the standard transformer architecture for learning symmetric data. For example, we have shown that a pre-trained vision transformer could be repurposed to encode graphs.

Learning Distribution of Data Augmentations Since our approach parameterizes a distribution p ω ( g | x ) on a group for symmetrization of form ϕ θ,ω ( x ) = E g [ g · f θ ( g -1 · x )] and learns it from data, one may find similarity to Augerino [7] and related approaches [77, 81, 95, 93, 80, 37] that learn distributions over data augmentations ( e.g. , p ω ( g ) ) for a similar symmetrization. The key difference is that, while these approaches aim to discover underlying (approximate) symmetry constraint from data and searches over a space of different group symmetries, our objective aims to obtain an exact G equivariant symmetrization ϕ θ,ω ( x ) given the known symmetry group G of data ( e.g. , G = S n for graphs). Because of this, the symmetrizing distribution has to be designed differently. In our case, we parameterize the distribution p ω ( g | x ) itself to be equivariant to a specific given group G , while for augmentation learning approaches, the distribution p ω ( g ) is parameterized for a different purpose of covering a range of different group symmetry constraints and their approximations ( e.g. , a set of 2D affine transformations [7]). This leads to advantages of our approach if the symmetry group G is known, as (1) our approach can learn non-trivial and useful distribution p ω ( g | x ) per input data x while keeping the symmetrized function ϕ θ,ω ( x ) exactly G equivariant, while augmentation

learning does not guarantee equivariance for a given group in general and often has to reduce to trivial group averaging p ω = Unif ( G ) to be exactly G equivariant, and (2) while augmentation learning has to employ regularization [7] or model selection [37] to prevent collapse to trivial symmetry that is the least constrained and would fit the training data most easily [37], our approach fixes and enforces equivariance for the given symmetry group G by construction, which allows us to use regular maximum likelihood objective for training without the need to address symmetry collapse.

A.3 Experimental Details (Section 3)

We provide details of the datasets and models used in our experiments in Section 3. The details of the datasets from the original papers [2, 1, 27, 28, 31, 84] can be found in Table 5 and Table 6.

A.3.1 Implementation Details of p ω for Symmetric Group S n (Section 3.1, 3.3, 3.4)

In all experiments regarding the symmetric group S n , we implement the S n equivariant distribution p ω ( g | x ) , i.e. , q ω : ( x , ϵ ) ↦→ P g as a 3-layer GIN with 64 hidden dimensions [101] that has around 25k parameters. Specifically, given a graph x with node features X ∈ R n × d in and adjacency matrix A ∈ R n × n , 7 we first augment a virtual node [32] which is connected to all nodes to facilitate global interaction while retaining S n equivariance, as follows:

where the feature of the virtual node v ∈ R d in is a trainable parameter. Then, we prepared the input node features H ∈ R ( n +1) × d in to the GIN as H = X ′ + ϵ where the noise ϵ ∈ R ( n +1) × d in is i.i.d. sampled from Unif [0 , η ] with scale hyperparameter η . Then, we employ following 3-layer GIN with 64 hidden dimensions to obtain processed node features H ′ ∈ R ( n +1) × 1 :

where each GINConv d 1 ,d 2 ,d 3 computes below with a two-layer elementwise MLP : R n × d 1 → R n × d 3 with hidden dimension d 2 , ReLU activation, batch normalization, and trained scalar e :

Then, from the processed node features H ′ ∈ R ( n +1) × 1 , we finally obtain the features Z ∈ R n for postprocessing by discarding the feature of the virtual node. Then, postprocessing into a permutation matrix is done with argsort : Z ↦→ P g ∈ R n × n as in Eq. (8).

Training To backpropagate through P g for end-to-end training of p ω ( g | x ) , we use straight-through gradient estimator [6] with an approximate permutation matrix ˆ P g ≈ P g . 8 For this, we first apply L2 normalization Z ↦→ ¯ Z and use the below differentiable relaxation of the argsort operator [61, 33, 98]:

where S ( · /τ ) is Sinkhorn operator [61] with temperature hyperparameter τ ∈ R + that performs elementwise exponential followed by iterative normalization of rows and columns. Following [61], we use 20 Sinkhorn iterations which worked robustly in all our experiments. For the correctness of straight-through gradients, it is desired that ˆ P g closely approximates the real permutation matrix P g during training. For this, we choose the temperature τ 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 ˆ P g with a strength of 0 . 1 in all experiments. The S n equivariance of the relaxed argsort Z ↦→ ˆ P g can be shown in a similar way to Proposition 3 from the fact that elementwise subtraction, absolute, scaling by -1 /τ , exponential, and iterative normalization of rows and columns all commute with P g ′ ∈ S n .

A.3.2 Implementation Details of p ω for Product Group S n × E (3) (Section 3.2)

In our n -body experiment on the product group S n × E (3) , we implement the S n × O (3) equivariant distribution p ω ( g | x -¯ x1 ⊤ ) , i.e. , q ω : ( x -¯ x1 ⊤ , ϵ ) ↦→ ( P g , Q g ) based on a 2-layer Vector Neurons version of DGCNN with 96 hidden dimensions [23] that has around 7k parameters. Due to the architecture's complexity, we focus on describing input and output of the network and postprocessing, and guide the readers to the original paper [23] for further architectural details. In a high-level, the Vector Neurons receives position P ∈ R n × 3 and velocity V ∈ R n × 3 of the zero-centered input x -¯ x1 ⊤ with noises ϵ 1 , ϵ 2 ∈ R n × 3 i.i.d. sampled from normal N (0 , η 2 ) with scale hyperparameter η , and produces features H S n ∈ R n × 3 × d 1 and H O (3) ∈ R n × 3 × d 2 with d 1 = 1 and d 2 = 3 as follows:

Then, we apply O (3) invariant pooling on H S n and S n invariant pooling on H O (3) , both supported as a part of [23], to obtain features for postprocessing Z S n ∈ R n × 1 and Z O (3) ∈ R 3 × 3 , respectively:

Then, postprocessing with argsort : Z S n ↦→ P g ∈ R n × n and Gram-Schmidt orthogonalization Z O (3) ↦→ Q g ∈ R 3 × 3 is performed identically as described in the main text (Section 2.2). For the straight-through gradient estimation of the argsort operator, we use relaxed argsort described in Appendix A.3.1, with the only difference of using the temperature τ = 0 . 1 .

A.3.3 Graph Isomorphism Learning with MLP (Section 3.1)

Base Model f θ For EXP and EXP-classify, the model is given adjacency matrix A ∈ R 64 × 64 and binary node features X ∈ R 64 which are zero-padded to maximal number of nodes 64 . For GRAPH8c, the input graphs are all of size 8 without node features, and the model is given adjacency matrix A ∈ R 8 × 8 . For EXP-classify, the prediction target is a scalar binary classification logit.

For the base model for EXP-classify, we use a 5-layer MLP f θ : R 64 × 64+64 → R on flattened and concatenated adjacency matrix and node features, with an identical architecture to other symmetrization baselines (MLP-GA and MLP-FA [74]) as in below:

where FC d 2 ,d 1 : R d 1 → R d 2 denotes a fully-connected layer and ReLU activation is omitted. For EXP, we drop the last layer to obtain 10-dimensional output. For GRAPH8c, we use the following architecture f θ : R 8 × 8 → R 10 that takes flattened adjacency to produce 10-dimensional output [74]:

Training For EXP-classify, we train our models with binary cross-entropy loss using Adam optimizer [46] with batch size 100 and learning rate 1e-3 for 2,000 epochs, which takes around 30 minutes on a single RTX 3090 GPU with 24GB using PyTorch [72]. We additionally apply 200 epochs of linear learning rate warm-up and gradient norm clipping at 0.1, which we found helpful for stabilizing the training. For the equivariant distribution p ω , we use noise scale η = 1 . Since EXP and GRAPH8c concern randomly initialized models, we do not train the models for these tasks.

A.3.4 Particle Dynamics Learning with Transformer (Section 3.2)

Base Model f θ The model is given zero-centered position P ∈ R 5 × 3 and velocity V ∈ R 5 × 3 of 5 particles at a time point with pairwise charge difference C ∈ R 5 × 5 and squared distance D ∈ R 5 × 5 . We set the prediction target as difference of position ∆ P ∈ R 5 × 3 after a certain time.

For the base model, we use a 8-layer transformer encoder f θ : R 25 × 8 → R 25 × 3 that operates on sequences of length 25 with dimension 8 . At each prediction, we first organize the input into a single tensor ∈ R 5 × 5 × 8 by placing P and V on the diagonals of C and D , and then turn the tensor into a sequence of 25 tokens ∈ R 25 × 8 by flattening the first two axes. Analogously, we organize the output of the model into a tensor ∈ R 5 × 5 × 3 and take the diagonal entries as the predictions. For the transformer, we use the standard implementation provided in PyTorch [72, 96], with 64 hidden dimensions, 4 attention heads, GELU activation [34] in feedforward networks, PreLN [100], learnable 1D positional encoding, and an MLP prediction head with 1 hidden layer. The model has around 208k trainable parameters, around 2.3 × compared to the GNN backbones of E (3) symmetrization baselines in the benchmark (GNN-FA and GNN-Canonical.) with 92k parameters.

Training We train our models with MSE loss using Adam optimizer [46] with batch size 100 and learning rate 1e-3 for 10,000 epochs, which takes around 8.5 hours on a single RTX 3090 GPU with 24GB using PyTorch [72]. We use weight decay with strength 1e-12 and dropout on the distribution p ω with probability 0.08. For the equivariant distribution p ω , we use noise scale η = 1 .

A.3.5 Graph Pattern Recognition with Vision Transformer (Section 3.3)

Base Model f θ The model is given adjacency matrix A ∈ R 188 × 188 and node features X ∈ R 188 × 3 zero-padded to maximal 188 nodes. The prediction target is node classification logits Y ∈ R 188 × 2 .

For the base model, we use a transformer with an identical architecture to ViT-Base [26] that operates on 224 × 224 images with 16 × 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 : R (14 × 14) × 768 → R (14 × 14) × 768 on sequences of (224 / 16) × (224 / 16) = 14 × 14 tokens. Then, we use the following as the base model f θ : ( A , X ) ↦→ Y :

where, for tokenize : R 188 × 188 × 1 × R 188 × 3 → R (14 × 14) × 768 we organize the input into a single tensor ∈ R 188 × 188 × 4 by placing X on the diagonals of A and apply 2D convolution with kernel size and stride 14, and for detokenize : R (14 × 14) × 768 → R 188 × 2 we apply transposed 2D convolution with kernel size and stride 14 to obtain a tensor ∈ R 188 × 188 × 2 and take its diagonal entries as output.

Training We train our models with binary cross-entropy loss weighted inversely by class size [27] using AdamW [53] optimizer with batch size 128, learning rate 1e-5, and weight decay 0.01. We train the models for 25k steps under learning rate warm-up for 5k steps then linear decay to 0 with early stopping based on validation loss, which usually takes less than 5 hours on 8 RTX 3090 GPUs with 24GB using PyTorch Lightning [29]. For the equivariant distribution p ω we use noise scale η = 1 and dropout with probability 0.1. For probabilistic symmetrization that involves sampling-based estimation, we use sample size 1 for training. For group averaging, sample size 1 for training led to optimization challenges, and therefore we use sample size 10 for training which yielded better results.

A.3.6 Real-World Graph Learning with Vision Transformer (Section 3.4)

Base Model f θ For Peptides-func/struct, the model is given adjacency matrix A ∈ R 444 × 444 , node features X ∈ R 444 × 64 , and edge features E ∈ R 444 × 444 × 7 , zero-padded to maximal 444 nodes. The prediction target is binary classification logits Y ∈ R 10 for Peptides-func, and regression targets

Y ∈ R 11 for Peptides-struct. For PCQM-Contact, the model is given adjacency matrix A ∈ R 53 × 53 , node features X ∈ R 53 × 68 , and edge features E ∈ R 53 × 53 × 6 , zero-padded to maximal 53 nodes. The prediction target is binary edge classification logit Y ∈ R 53 × 53 × 1 .

For the base model, we use a transformer with an identical architecture to ViT-Base that operates on 14 × 14 tokens, same as in Appendix A.3.5. For Peptides-func and Peptides-struct, we use the following as the base model f θ : ( A , X , E ) ↦→ Y :

where tokenize2D : R 444 × 444 × (1+7) → R (14 × 14) × 768 is 2D convolution with kernel size and stride 32, tokenize1D : R 444 × 64 → R 196 × 768 is 1D convolution with kernel size and stride 3, and detokenize [cls] performs linear projection of the global [cls] token [26] to the target dimensionality. For PCQM-Contact, we use the following as base model f θ : ( A , X , E ) ↦→ Y :

where tokenize2D : R 53 × 53 × (1+6) → R (14 × 14) × 768 is 2D convolution with kernel size and stride 4, tokenize1D : R 53 × 64 → R 196 × 768 is 1D convolution with kernel size and stride 1, and detokenize2D : R (14 × 14) × 768 → R 53 × 53 × 1 is transposed 2D convolution with kernel size and stride 4.

Training We train our models with cross-entropy for classification and L1 loss for regression using AdamW [53] optimizer with batch size 128, learning rate 1e-5 except for PCQM-Contact where we use 5e-5, and weight decay 0.01. We train the models for 50k steps under learning rate warm-up for 5k steps then linear decay to 0 with early stopping based on validation loss, which usually takes less than 12 hours on 8 RTX 3090 GPUs with 24GB using PyTorch Lightning [29]. For the equivariant distribution p ω , we use noise scale η = 1 , and use dropout with probability 0.1 except for PCQM-Contact where we do not use dropout. We use 10 samples for estimation during training.

A.4 Supplementary Experiments (Continued from Section 3)

In this section, we present additional experimental results that supplement the experiments in Section 3 but could not be included in the main text due to space constraints.

A.4.1 Graph Isomorphism Learning (Section 3.1)

In our experiments on graph isomorphism learning in Section 3.1, we mainly experimented for S n symmetrization of an MLP. Here, we provide supplementary results on S n symmetrization of a GIN base model with node identifiers, following [74]. The results can be found in Table 7. In accordance with Section 3.1, our approach successfully performs S n symmetrization of GIN-ID.

1

10

50

Train sample size

90

95

100

Test accuracy (%)

1

10

50

Train sample size

5.0

5.5

6.0

6.5

7.0

Entropy of sampled permutation

A.4.2 Particle Dynamics Learning (Section 3.2)

In our experiments on n -body dataset in Section 3.2, we experimented for S n × E (3) symmetrization using a 1D sequence transformer architecture which has 2 . 3 × parameters compared to baselines. To provide parameter-matched comparison against baselines in literature, we apply our approach for E (3) symmetrization of S n equivariant GNN base model that is widely used in literature [74, 41]. We faithfully follow [74, 41] on the experimental setups including training hyperparameters and the configuration of GNN base model, and only add E (3) equivariant distribution p ω ( g | x -¯ x1 ⊤ ) , i.e. , q ω : ( x -¯ x1 ⊤ , ϵ ) ↦→ Q g by utilizing the 2-layer Vector Neurons architecture described in Appendix A.3.2 using only its O (3) prediction head. We use 20 samples for training and testing. The results can be found in Table 8. In accordance with the results in Section 3.2, our approach outperforms other symmetrization approaches and achieves a new state-of-the-art of 0.00386 MSE.

A.4.3 Effect of Sample Size on Training and Inference

In this section, we provide additional analysis on how the sample size for estimation of symmetrized function (Eq. (4)) affects training and inference. We use the experimental setup of EXP-classify (Section 3.1; S n invariance) and analyze the behavior of MLP-PS with identical initialization and hyperparameters, only controlling sample sizes ∈ { 1 , 2 , 5 , 10 , 20 , 50 } for training. Specifically, we analyze (1) variance of permutation matrices P g ∼ p ω ( g | x ) measured indirectly by the entropy of their aggregation ¯ P = ∑ P g /N as in Section 3.1, (2) sample variance of the unbiased estimator g · f θ ( g -1 · x ) of the symmetrized function ϕ θ,ω ( x ) as in Eq. (4), and (3) sample mean and variance

1

10

50

Train sample size

10

7

10

6

10

5

10

4

10

3

10

2

10

1

Output variance

1

10

50

Train sample size

10

5

10

3

10

1

10 1

Loss variance

1

10

50

Train sample size

10

2

10

1

10 0

Loss mean

1

10

50

200

Inference sample size

of the estimated task loss L ( y , g · f θ ( g -1 · 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.

Observations are as follows. First, models trained with smaller sample sizes need more iterations to converge, but after sufficient training (2000 epochs), all achieve > 95% test accuracy when evaluated with 10 samples (Figure 4). Second, models trained with smaller sample sizes tend to be more sample efficient, i.e. , tend to perform a lower variance estimation. Their distribution p ω ( g | x ) tend to learn more low-variance permutations (Figure 5), and the models tend to learn low-variance estimation of output and loss (left and center panels of Figure 6). This indicates that small sample size may serve as a regularizer that encourages lower variance of the estimator. However, this regularization effect is not always beneficial in terms of task loss (right panel of Figure 6), as training sample size 1 achieves a poor task loss for all sample sizes presumably due to the optimization challenge caused by over-regularization. In other words, the sample size for training introduces a tradeoff; a small sample size takes more training iterations to converge, but serves as a regularizer that encourages lower variance of the estimator and thus a better inference time sample efficiency. On the other hand, larger sample sizes for inference consistently benefits all models (Figure 6).

Interestingly, this observed tendency is consistent with the theoretical claims in literature [67, 68] on the sampling based training of symmetrized models, which we reprise here. When training the symmetrized model ϕ θ,ω ( x ) in Eq. (4), we cannot directly observe ϕ θ,ω ( x ) , but observe samples of its unbiased estimator g · f θ ( g -1 · x ) . Thus, it can be questionable what objective we are actually optimizing during the sampling-based training. Based on [67, 68], it turns out that minimizing a convex loss function L on the estimated output g · f θ ( g -1 · x ) is equivalent to minimizing an upper bound to the true objective on the symmetrized output ϕ θ,ω ( x ) . This is because our estimation is no longer unbiased when computing loss, as we have the following from Jensen's inequality:

That is, minimizing the sampling-based loss is minimizing an upper-bound surrogate to the true objective. It has been claimed that optimizing this upper bound has an implicit low-variance regularization effect [67, 68], which is consistent with our observations. This also roughly explains why our distribution p ω ( 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 n invariance) and experiment with MLP-PS and MLP-GA.

We first analyze whether using the equivariant distribution p ω ( g | x ) for symmetrization offers a lower variance estimation, i.e. , a better sample efficiency, compared to group averaging with Unif ( G ) . This supplements the results in Figure 2 that p ω ( g | x ) learns to produce low-variance permutations compared to Unif ( G ) . Specifically, we fix a randomly initialized MLP f θ and symmetrize it using our approach and group averaging. We then measure (1) the sample variance of the unbiased estimator

1

10

100 200

Inference sample size

10

9

10

8

10

7

Output variance

GA

PS (Ours)

1

10

100 200

Inference sample size

10

8

10

7

10

6

Loss variance

GA

PS (Ours)

100

1000

2000

Train epoch

10

4

10

3

10

2

10

1

10 0

10 1

Gradient direction norm

GA

PS (Ours)

g · f θ ( g -1 · x ) of the symmetrized function (Eq. (1) and Eq. (4)), and (2) the sample variance of the estimated task loss L ( y , g · f θ ( g -1 · 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 7, showing that symmetrization with equivariant distribution p ω ( g | x ) consistently offers a lower variance estimation than group averaging across inference sample sizes.

In addition, we analyze whether the equivariant distribution p ω ( g | x ) for symmetrization offers more stable gradients for the base function f θ during training compared to group averaging, as conjectured in Section 2. For this, we fix a randomly initialized MLP f θ 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 f θ . This averages out the variance from individual data points and provides the net direction of the gradient on the base function offered by p ω ( g | x ) or Unif ( G ) . The results are in Figure 8, showing that that symmetrization with equivariant distribution p ω ( g | x ) offers a consistently larger magnitude of the net gradient, while group averaging with Unif ( G ) leads to near-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 p ω ( g | x ) for symmetrization leads to more consistent gradient across training data instances, i.e. , it offers a 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 g from an equivariant canonicalizer C ω : x ↦→ ρ ( g ) . The main claim is 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 1.

0

1

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

(a) Inner-symmetric nodes

0

1

2

3

4

5

6

0.025

0.020

0.015

0.010

0.005

0.000

0.005

(b) MLP-PS (Ours)

0

1

2

3

4

5

6

0.02

0.01

0.00

0.01

0.02

(c) MLP-Canonical.

0

1

2

3

4

5

6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

(a) Inner-symmetric nodes

0

1

2

3

4

5

6

0.02

0.01

0.00

0.01

(b) MLP-PS (Ours)

0

1

2

3

4

5

6

0.02

0.01

0.00

0.01

0.02

(c) MLP-Canonical.

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

0

2

4

6

8

(a) Inner-symmetric nodes

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

0.010

0.005

0.000

0.005

(b) MLP-PS (Ours)

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

0.03

0.02

0.01

0.00

0.01

0.02

(c) MLP-Canonical.

̸

More specifically, let us recall the definition of G equivariant canonicalizer from [41]. 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 G x = 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 g 1 = gh 1 and g 2 = gh 2 for some g ∈ G and any h 1 , h 2 ∈ G x where h 1 = h 2 . Then we have C ω ( g 1 · x ) = C ω ( gh 1 · x ) = C ω ( g · x ) = C ω ( gh 2 · x ) = C ω ( g 2 · x ) , implying that ρ ( g 1 ) C ω ( x ) = ρ ( g 2 ) C ω ( x ) . Since g 1 = g 2 , this contradicts the group axiom, and thus an equivariant canonicalizer cannot exist for inputs with non-trivial inner-symmetries. To handle all inputs, canonicalization [41] 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 (Theorem A.2 of [41]). In other words, canonicalization does not guarantee equivariance for inputs with inner symmetries.

̸

To visually demonstrate this, we perform a minimal experiment using several graphs x with nontrivial stabilizers G x , i.e. , inner symmetries, taken from [90]. We fix a randomly initialized MLP

f θ : R n × n → R n and symmetrize 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. The results are in Figures 9, 10, and 11. 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 MLP-Canonical. If a method is G equivariant, 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 11, 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 a 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 4)

While the equivariance, universality, simplicity, and scalability of our approach offers a 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 to equivariant architectures due to less imposed prior knowledge on 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 impose a strong prior on the hypothesis space and may work similarly to architectural priors that benefit data efficiency. Also, for the sampling cost, since the sampling is completely parallelizable and analogous to using a larger batch size, we believe it can be overcome to some degree by leveraging parallel computing techniques developed for scaling batch size.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\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\nQUESTIONS:\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\nLIMITATIONS:\nSome 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOver 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).\n\nWEAKNESSES:\n* 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).\n\nQUESTIONS:\na) 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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOverall 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.\n\nWEAKNESSES:\nRelation 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?\n\nQUESTIONS:\nI would appreciate any response from the authors regarding the weakness stated above.\n\nLIMITATIONS:\nCould not find a discussion on limitations", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe authors have not discussed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThe authors have discussed limitations and potential societal impact associated with this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-14"}