Equivariance with Learned Canonicalization Functions

Symmetry-based neural networks often constrain the architecture in order to achieve invariance or equivariance to a group of transformations. In this paper, we propose an alternative that avoids this architectural constraint by learning to produce a canonical representation of the data. These canonicalization functions can readily be plugged into non-equivariant backbone architectures. We offer explicit ways to implement them for many groups of interest. We show that this approach enjoys universality while providing interpretable insights. Our main hypothesis is that learning a neural network to perform canonicalization is better than using predefined heuristics. Our results show that learning the canonicalization function indeed leads to better results and that the approach achieves excellent performance in practice.


Introduction
The problem of designing machine learning models that properly exploit the structure and symmetry of the data is becoming more important as the field is broadening its scope to more complex problems. In multiple applications, the transformations with respect to which we require a model to be invariant or equivariant are known and provide a strong inductive bias (Bronstein et al., 2021;Bogatskiy et al., 2022;van der Pol et al., 2020;Mondal et al., 2020;Celledoni et al., 2021).
As is often the case, taking a step back and drawing analogies with human cognition is fruitful here. Human pattern recognition handles some symmetries with relative ease. When data is transformed in a way that preserves its essential characteristics, we precisely know if and how we should adapt our response. One context in which this has been particularly well-studied in cognitive science is visual shape recognition. Experiments have shown that subjects can accurately make the difference between different orientations of an object and actual modifications to the structure of an object (Shepard and Metzler, 1971;Carpenter and Eisenberg, 1978).
There are multiple ways in which this could be achieved. According to Tarr and Pinker (1989), theories of invariant shape recognition broadly fall into three categories: viewpointindependent models, in which object representations depend only on invariants features, multiple-view models in which an object is represented as a set of representations corresponding to different orientations, and single-view-plus-transformation models in which an object is converted to a canonical orientation by a transformation process.
Correspondingly, similar ideas have been explored in deep learning to achieve equivariance. Models that impose equivariance through constraints in the architecture (Shawe-Taylor, 1989;Cohen and Welling, 2016; or that only use invariants as inputs (Villar et al., 2021) can be seen as belonging to the viewpoint-independent type. The multiple-view approach includes models that symmetrize the input by averaging over all the transformations or a subset of them (Manay et al., 2006;Benton et al., 2020;Yarotsky, 2022;Puny et al., 2022). By contrast, the transformation approach has seen less interest and, to our knowledge, has yet to be used to achieve exact equivariance. This is all the more surprising considering that evidence from cognitive science suggests that this approach is used in human visual cognition (Shepard and Metzler, 1971;Carpenter and Eisenberg, 1978;Hinton and Parsons, 1981). When presented with a rotated version of an original pattern, the time taken by humans to do the association is proportional to the rotation angle, which is more consistent with the hypothesis that we perform a mental rotation.
Some works (Jaderberg et al., 2015;Qi et al., 2017) have proposed to learn transformations of input to facilitate processing in a downstream task, but these approaches are closer to regularizers and provide no guarantees. Another method is to use heuristics to standardize inputs (Yüceer and Oflazer, 1993;Lowe, 2004;Aslan et al., 2022), but this approach requires significant hand-engineering and is hardly generalizable.
Present work We introduce a systematic and general method for equivariant machine learning based on learning mappings to canonical samples. We hypothesize that among all valid canonicalization functions, some will lead to better downstream performance than others. Rather than trying to hand-engineer these functions, we may as well let them be learned in an end-to-end fashion with a prediction neural network. Our method can readily be used as an independent module that can be plugged into existing architectures to make them equivariant to arbitrary transformation groups, discrete or continuous. Our approach enjoys similar expressivity advantages to methods like frame averaging by (Puny et al., 2022), but has several added benefits. It is simpler, more efficient, and replaces hand-engineered frames for each group by a systematic end-to-end learning approach.
Our contributions are as follows: • Novel Framework: We introduce a general framework for equivariance to arbitrary groups based on mappings to canonical samples. This framework can be plugged into any existing non-equivariant architecture.
• Theoretical Guarantees: We prove that in some settings, such models are universal approximators of equivariant functions.
• Efficient Implementations: We provide multiple variants of efficient implementations of this framework to specific domains.
• Practical Performance: We perform experiments that show that the proposed method achieves excellent results in the image domain. We also support our hypothesis that learning the canonicalization function is a better strategy than fixing it.

Problem Setting
We are interested in learning functions φ : X → Y with inputs x ∈ X and outputs y ∈ Y belonging to arbitrary finite-dimensional vector spaces. We will consider a set of linear symmetry transformations T ⊂ GL (X ), where GL (X ) is the set of invertible matrices over the vector space X . This is described by a group representation ρ : G → T , where G is an abstract group. Without loss of generality, we can assume that ρ is a group isomorphism. Therefore, the inverse ρ −1 : where the group action ρ on the input and the group action ρ on the output will be clear from the context. In particular, when ρ (g) = I, we say that φ is invariant. We use ρ (G) to denote the image of ρ. We call ρ (G) x = {ρ (g) x | ∀ g ∈ G} the orbit of the element x. It is the set of elements to which x can be mapped to by the group action. The set of orbits, denoted X /G forms a partition of the set X .

General Formulation
The invariance requirement on a function φ amounts to having all the members of a group orbit mapped to the same image by φ. It is thus possible to achieve invariance by appropriately mapping all elements to a canonical sample from their orbit before applying any function. For equivariance, elements can be mapped to a canonical sample and, after a function is applied, transformed back according to their original position in the orbit. This can be formalized by writing the equivariant function φ in canonicalized form as where the function f : X → Y is called the prediction function and the function h : X → ρ (G) is called the canonicalization function. Here h (x) −1 is the inverse of the representation matrix and h (x) = ρ ρ −1 (h (x)) is the counterpart of h (x) on the output. Equivariance in equation 2 is obtained for any prediction function if the canonicalization function is itself G-equivariant 1 , h (ρ (g) x) = ρ (g) h (x) ∀ g, x ∈ G × X .
It may seem like the problem of obtaining an equivariant function has merely been transferred in this formulation. This is however not the case: in equation 2, the equivariance and prediction components are effectively decoupled. The canonicalization function h can therefore be chosen as a simple and inexpressive equivariant function, while the heavy-lifting is done by the prediction function f .
A more general condition can be formulated, such that the decoupling is partial.
and the prediction function f is K-equivariant, then φ defined in equation 2 is G-equivariant.
If K is a normal subgroup such that G J K, this can be realized with a canonicalization function that has image ρ (J), and that is J-equivariant and K-invariant.
The proof follows in Appendix B. This is equivalent to saying that the canonicalization function should output a representation of a coset in G/K in an equivariant way, the applied transformation being chosen arbitrarily within the coset. When K = {e}, only the canonicalization function is constrained, which is the case described above. The other extreme, given by K = G, corresponds to transforming the input in an arbitrary way and constraining the prediction function as is usually done in equivariant architectures. These are respectively the single-view-plus-transformation and the viewpointindependent implementations described in the introduction. Subgroups {e} < K < G offer intermediary options; the lattice of subgroups of G therefore defines a family of models. Since equivariance to a smaller group is less constraining for the prediction function, set inclusion in the subgroup lattice corresponds to increased expressivity.

Universality Result
We can now introduce a more formal result on the expressivity of equivariant functions obtained with canonicalization functions. A parameterized function φ is a universal approximator of G-equivariant functions if for any G-equivariant continuous function ψ, any compact set K ⊆ X and any > 0, there exists a choice of parameters such that ψ (x) − φ (x) < ∀ x ∈ K. We make the additional assumption that the set K is closed under the group action.
Theorem 2 A G-equivariant parameterized function φ written in canonicalized form (equation 2) and satisfying equation 3 with K ≤ G, is a universal approximator of G-equivariant functions if the prediction function is a universal approximator of K-equivariant functions.
The proof follows in Appendix C. The following corollary is especially relevant.
Corollary 3 A G-equivariant parameterized function φ written in canonicalized form with a G-equivariant canonicalization function and a multilayer perceptron (MLP) as a prediction function is a universal approximator of G-equivariant functions.
This result can significantly simplify the design of universal approximators of equivariant functions since a non-universal equivariant architecture for the canonicalization function can be combined with an MLP. solves this issue. However, because the subset of symmetric inputs is of measure zero, using standard equivariance is not expected to be a problem in practice.

Model Design for Canonicalization Functions
In the next section, we elaborate on how suitable canonicalization functions can be obtained in different settings.
The canonicalization function can be chosen as any existing equivariant neural network architecture with the output being a group element; we call this the direct approach. For permutation groups and Lie groups, an equivariant multilayer perceptron (Shawe-Taylor, 1989;Finzi et al., 2021) can be used. We provide examples of implementations in the next section.
We also introduce an alternative method that can prove especially useful when dealing with continuous groups. The canonicalization function can be defined as where s : X → R is a neural network. Such a mapping is G-equivariant provided that s has a unique minimum in each orbit. We call this the optimization approach. Intuitively, the function s represents the unnormalized distance between the input and the canonical sample of the orbit and is therefore minimized when ρ (g) is the transformation that takes to the canonical sample. For continuous groups, a function with unique minimum can be learned with Input Convex Neural Networks (ICNN) (Amos et al., 2017), but satisfying this condition for each orbit is more challenging. The minimization in equation 4 can then be performed using any standard convex optimization method. For practical application, we are typically interested in low-dimensional Lie groups, which makes second-order optimization methods usable. The gradient of the neural network with respect to the parameters can be obtained efficiently (e.g. without having to differentiate through all the optimization steps) with implicit differentiation (Blondel et al., 2022). The optimization implementation presents a close analogy with the mental rotation phenomenon described in the introduction, as humans try to minimize the distance between their representation of an object and the canonical one. As such, it is expected that the optimization process will take more iterations when the input sample is farther away in orbit from the canonical sample. This is consistent with the experimental evidence for mental rotation (Shepard and Metzler, 1971;Carpenter and Eisenberg, 1978).

Euclidean Group
The Euclidean group E (d) is used to capture symmetry to rotations, translations, and reflections. Domains in which this type of symmetry is especially relevant include point cloud modelling, image processing and applications in physics. Below we give the design principles to get an equivariant model for image and point cloud inputs.
Image Input. In order to build networks equivariant to E(2), we need a canonicalization function that outputs an element of E(2) given an image. This can be achieved by using a G-CNN (Cohen and Welling, 2016) and taking an argmax of the output feature map. While G-CNN only considers four rotations, this can be easily extended to arbitrary rotations. This approach can be further simplified if we use a translation equivariant prediction network which is generally the case with most CNN-based architectures. As the translation group T (2) is a normal subgroup of the Euclidean group E(2), Theorem 1 indicates that we only require the canonicalization function to be equivariant to O(2). This means we can eliminate the spatial dimension in the output feature map of the canonicalization function and only need to take an argmax along the point group dimension to identify the correct orientation of the image. However, there are two problems with this approach. First, extending G-CNNs to higher-order rotations will lead to more computation and artifacts due to the finer rotation of filters. Second, we can not backpropagate through the canonicalization function as the argmax operation is not differentiable.
We can avoid the first problem by using a shallower network with a larger filter size. We empirically show why this is a sound choice for canonicalization function in Section 4. Furthermore, to solve the second problem, we can use the straight-through estimator trick (Bengio et al., 2013). Appendix E contains a PyTorch code snippet to perform the canonicalization function of images in a differentiable way using a G-CNN. Similarly, the output of the canonicalization function can be used to invert the feature maps from the prediction network back to their original orientation in a differentiable way.
Since CNNs are universal approximators of T (2)-equivariant functions (Yarotsky, 2022), it follows from Theorem 2 that a CNN augmented with an O(2) equivariant canonicalization function is a universal approximator of E(2)-equivariant functions.
Point Cloud Input. Elements of the Euclidean group can be written as (O, t), where O ∈ R n×n is an orthogonal matrix and t ∈ R n is an arbitrary translation vector. The vector representation on n + 1 dimensional vectors (defined by concatenating a constant 1 to the original vectors) is then defined in the following way This representation is bounded We seek to define an E (d)-equivariant canonicalization function. This can be done by defining it as h (x) = ρ h O (x) , h t (x) , where the functions h O : X → R n×n and h t : X → R n output the point transformation and the translation respectively. By the equivariance condition,  (Zaheer et al., 2017) as a backbone architecture would result in a universal approximator of E (d) and permutation equivariant functions, following Theorem 2 and Theorem 1 of (Segol and Lipman, 2020).

Symmetric Group
The symmetric group S n over a finite set of n elements contains all the permutations of that set. This group captures the inductive bias that input order should not matter. Domains for which S n -equivariance is desirable include object modelling and detection, graph representation learning, and applications in language modelling. S n -equivariant canonicalization functions can be obtained with a direct approach. In particular, the methods introduced by Mena  We propose an alternative of S n -equivariance for the canonicalization function that circumvents these issues. Instead of producing a permutation of the input, the canonicalization function can instead be designed to obtain a clustering of the input, with a function c : X → {1, . . . , k}, where k is the number of possible clusters. This can be made differentiable using the straight-through estimator described for the images above. Then, a corresponding permutation is obtained by sorting elements by cluster index, with elements belonging to the same cluster being sorted arbitrarily. The canonicalization function would therefore respect equation 3 with respect to the subgroup K = S l 1 × · · · × S l k , where l i is the number of elements assigned to cluster i, the product of symmetric groups within clusters. Heuristically, this enforces the inductive bias that set elements should be mapped to high-level classes, across which interactions will be evaluated in an expressive way. This is close in spirit to what is done in the ClusterFormer (Wang et al., 2022) and Reformer (Kitaev et al., 2020) transformer architectures.
The prediction function f would then have to be S l 1 × · · · × S l k equivariant to obtain permutation equivariance. Such a function can readily be obtained with an equivariant multilayer perceptron, as derived by Wang et al. (2020).

Method
Error % ↓ We perform a preliminary empirical analysis of the proposed framework in the image domain. Experiments on other domains will be the subject of more extensive work. We selected the Rotated MNIST dataset (Larochelle et al., 2007), which is used as a benchmark dataset that uses a classification task to test equivariant architectures in prior work (Cohen and Welling, 2016). In all our experiments, we only change the network architecture. We train the models by minimizing the cross entropy loss for 100 epochs using Adam (Kingma and Ba, 2014) with a learning rate of 0.001. We perform early stopping based on the classification performance of the validation dataset with patience of 20 epochs. We first compare our method with different CNN and G-CNN baselines. For CNN (base), we choose an architecture with 7 layers where layer 1 to 3 has 32, 4 to 6 has 64 and layer 7 has 128 channels, respectively. Instead of pooling, we use convolution filters of size 5 × 5 with a stride 2 at layers 4 and 7. The remaining convolutions have filters of size 3 × 3 and stride 1. All the layers are followed by batch-norm and ReLU activation with dropout(p=0.4) only at layers 4 and 7. For the canonicalization function, we choose a shallow G-CNN with three layers. We start with a lifting layer with a filter that is the same size as the input image. This is followed by ReLU and group equivariant layers with 1 × 1 filters.
We consider two variants: an untrained canonicalization function with frozen weights and a canonicalization function learned end-to-end with the prediction function. For a pure G-CNN-based baseline, we provide the value reported by Cohen and Welling (2016) and design a variant which has similar architecture to CNN (base) while matching the number of parameters of our model with canonicalization function. We call this G-CNN (p4 & =params). We consider two more variants of the CNN: one with the same number of parameters as our model and the other where the canonicalization function is done by finding the orientation of the digits using Principal Component Analysis (PCA). We see that using a fixed canonicalization function technique like PCA or canonicalization function with frozen parameters leads to an improvement in performance over the CNN baseline. However, learning the canonicalization function provides a significant performance improvement. Our approach outperforms all the CNN-based baselines and is comparable to G-CNNs. We further notice that adding a canonicalization function for reflections hurts performance. This is, in fact, not surprising as reflection augmentations were not used to build the rotated-MNIST dataset. Adding reflection invariance while the data does not possess it may introduce ambiguity between certain digits like 2 and 5. Next, we perform experiments to understand the role of different components in our model using the group of n discrete rotations (pn).

Ablation Study
First, we vary both the number of layers of the canonicalization network and the number of rotations it is equivariant to. For this, we extend the layers of GCNN to any arbitrary rotations. As we noticed that using a larger filter leads to better performance for higher order rotations we stick to architecture with a lifting layer that has image-sized filters followed by 1 × 1 filters. From Table 4, we notice that adding equivariance to higher order rotation in the canonicalization function leads to significant performance improvement in comparison to adding more layers. Figure 1 shows the canonical orientation resulting from the learnt canonicalization function with a single lifting layer on 110 randomly sampled images of class 7 from the test dataset. This suggests that having a shallow network is enough to learn the correct canonicalization function with a sufficiently high order of discrete rotations. For p64, we see that all the similar-looking samples are aligned in one particular orientation. In contrast, although techniques like PCA or freezing parameters of the canonicalization function finds the correct canonicalization function for simple digits like 1 (see Appendix 3) they struggle to find stable mappings for more complicated digits like 7. Next, we compare the inference time of our model with pure G-CNN-based architectures. For this experiment, we take the CNN architecture of our predictor network and replace the convolutions with group convolutions. As increasing the order of rotation in G-CNN requires more copies of rotated filters in the lifting layer and more parameters in the subsequent group convolution layers, we decreased the number of channels to keep the number of parameters the same as our model. Figure 2 shows that although G-CNN is slightly better for the p4 group increasing the order of discrete rotations improves our model's performance significantly compared to G-CNN. In addition to performance gain, our model's inference speed remains more or less constant while encoding invariance to higher-order rotations due to the shallow canonicalization network. This makes our approach more suitable for building equivariance to bigger groups.

Conclusion
In this work, we propose using learned canonicalization function to obtain equivariant machine learning models. These canonicalization functions can conveniently be plugged into existing architectures, resulting in highly expressive models. We have described general approaches to obtain canonicalization functions and specific implementations for the Euclidean group (for images and point clouds) and the symmetric group.
We performed preliminary experimental studies in the image domain to test our initial hypotheses regarding this framework. First, we show that our approach achieves comparable or better performance than baselines on invariant tasks. Importantly, learning the canonical network is a better approach than using a fixed mapping, either a frozen neural network or a heuristic approach like PCA. Our results also show that the canonicalization function can be realized with a shallow equivariant network without hindering performance. Finally, we show that this approach reduces inference time and is more suitable for bigger groups then G-CNNs.
Areas of future work include an experimental study of this framework on other domains, such as point clouds and graphs. In the image domain, we will also explore building canonicalization functions using steerable networks. The function would output an orientation fibre that transforms by the irreducible representation of the special orthogonal group. Understanding how design choices for canonicalization functions affect downstream performance would also be a potentially fruitful research direction.

Appendix A. Symmetric inputs and relaxed equivariance
We say that an input x ∈ X is symmetric if its stabilizer subgroup G x = {g ∈ G | ρ (g) x = x} is non-trivial. Given any g 1 , g 2 ∈ G, a necessary and sufficient condition for is that g 1 and g 2 are part of the same coset for the stabilizer, e.g. g 1 , g 2 ∈ gG x . This is follows from the well-known relation between orbits and stabilizers. Therefore, symmetric inputs are always fixed by multiple group elements. Symmetric inputs are problematic when using the standard definition of equivariance for the canonicalization function because for g 1 , g 2 ∈ gG x , we have For general group actions, there will not exist a h (x) ∈ ρ (G) such that the last equality is satisfied.
A relaxed version of equivariance can be defined to address this issue.
Definition 4 (Relaxed equivariance) Given group representations ρ : G → GL (X ) and ρ : G → GL (Y), a function h : X → Y satisfies the relaxed equivariance condition if ∀g 1 , x ∈ G × X there exists a g 2 ∈ g 1 G x such that h (ρ (g 1 ) x) = ρ (g 2 ) h (x) This is a generalization of the idea of multiset-equivariance introduced by  to arbitrary group representations. When G x = {e}, standard equivariance is recovered. Canonicalization functions satisfying this condition do not suffer from the aforementioned problem. In addition, this condition is sufficient to obtain relaxed equivariance for canonicalized functions (Eq. 2). This is because, for g 2 ∈ g 1 G x : Using the fact that g 2 −1 g 1 ∈ G x , The fact that we obtain relaxed equivariance for the canonicalized function is a feature rather than a bug. This captures the desideratum that a function should be able to output asymmetric outputs from symmetric inputs, which is not the case for standard equivariant functions.
Since any orbit member can conversely be written as x = ρ (k, j)x for some k ∈ K and j ∈ J, if the canonicalization function is J-equivariant and K-invariant, we have h (x) = ρ (e, j) h (x) h (x) = ρ (e, j) which completes the proof.

Appendix C. Proof of Theorem 2
Proof The proof is inspired by the symmetrization approach of Yarotsky (2022). Let ψ be an arbitrary G-equivariant function, and φ be defined by equation 2. By the equivariance of ψ, we have Since Y is finite-dimensional, we know that linear operators in GL (Y) are bounded. This means that for every representation matrix there exists a positive number c that bounds the induced operator norm, e.g. ∀g ∈ G, ∃c > 0, ρ (g) ≤ c.
We therefore obtain where c > 0. If f is a universal approximator of K-equivarant functions, then it is also a universal approximator of G-equivariant functions. We therefore have Using the fact that K is closed under the group action, we obtain the desired result