Equivariant Adaptation of Large Pretrained Models
Authors: Arnab Kumar Mondal
Siba Smarak Panigrahi
Sékou-Oumar Kaba
Sai Rajeswar
Siamak Ravanbakhsh
Venue: NeurIPS 2023 poster (2023)
Keywords: deep learning
large pretrained models
symmetry
equivariance
group theory
computer vision
point clouds
foundation models
TL;DR: We present a technique to achieve equivariance to a specific known transformation group in any pre-trained large model, without compromising performance, computational speed, or the need for retraining from scratch.

Abstract: Equivariant networks are specifically designed to ensure consistent behavior with respect to a set of input transformations, leading to higher sample efficiency and more accurate and robust predictions. However, redesigning each component of prevalent deep neural network architectures to achieve chosen equivariance is a difficult problem and can result in a computationally expensive network during both training and inference. A recently proposed alternative towards equivariance that removes the architectural constraints is to use a simple canonicalization network that transforms the input to a canonical form before feeding it to an unconstrained prediction network. We show here that this approach can effectively be used to make a large pretrained network equivariant. However, we observe that the produced canonical orientations can be misaligned with those of the training distribution, hindering performance. Using dataset-dependent priors to inform the canonicalization function, we are able to make large pretrained models equivariant while maintaining their performance. This significantly improves the robustness of these models to deterministic transformations of the data, such as rotations. We believe this equivariant adaptation of large pretrained models can help their domain-specific applications with known symmetry priors.

Equivariant Adaptation of Large Pretrained Models
Arnab Kumar Mondal Mila, McGill University ServiceNow Research
Siba Smarak Panigrahi" Mila, McGill University
Sékou-Oumar Kaba Mila, McGill University
Sai Rajeswar
Siamak Ravanbakhsh Mila, McGill University
ServiceNow Research
Abstract
Equivariant networks are specifically designed to ensure consistent behavior with respect of input transformations, leading to higher sample efficiency and more accurate and robust predictions. However, redesigning each component of prevalent deep neural network architectures achieve chosen equivariance is a difficult problem and can result in computationally expensive network during both training and inference. recently proposed alternative towards equivariance that removes the architectural constraints is to use simple canonicalization network that transforms the input to a canonical form before feeding it to an unconstrained prediction nefwork. We show here that this approach can effectively be used to make large pretrained network equivariant. However, we observe that the produced canonical orientations can be misaligned with those of the training distribution, hindering performance. Using datasct-dependent priors to inform the canonicalization function, are able to make large pretrained models equivariant while maintaining their performance. This significantly improves the robustness of these models deterministic transformations of the data, such rotations. We believe this equivariant adaptation large pretrained models can help their comain-speciffi applications with known symmetry priors.
Introduction
Deep neural networks (DNNs) have demonstrated exceptional success in dealing with range of data modalities, such as image and point cloud 115. The majority these applications require working with input data that undergoes various transformations, such rotation, scaling, or translation in image data. Invariant and equivariant networks are "aware" of such transformations; they are specifically designed ensure that the network's behaviour is insensitive to input transformations, leading to more accurate and robust predictions and improved sample complexity [e.g., 6. 7 9- 13 14] As such, they have emerged as a promising solution wide range of computer vision and point cloud processing tasks [7] 9 10 15|20) 20), including classification, object recognition, and segmentation, among others.
In parallel to this, it has been shown that scaling models, both in the number of parameters and amount of training data, systematically improve their performance (1 28 Pre-training massive datasets has emerged as a practical method to harness this advantage. Several large models are now publicly available, and fine-tuning them specific applications has proven to be successful strategy in a variety of applications. However, such foundation models 29 are typically not equivariant
These authors contributed equally to this work
Corresponndece: arnab.mondal @ mila.quebec. Work done while at ServiceNow Research
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Ground Truth Mask
Prediction from SAM
Prediction from our equivariant SAM
Canonicalization cost: A purameters +0.3% inference time: +7.3%
Predicted masks from the Segment Anything Model (SAM) 21,, showcasing both the original model and our proposed equivariant adaptation for 90º coumter-clockwiss rotated input images taken from the COCO 2017 dataset [22]. Our method makes SAM equivariant to the group of 909 rotations while only requiring 0.3% extra parameters and modestly increasing the inference time by 7.3%. For further insights into this experiment, refer to Section 1.1
Figure I:
to most transformations except translations, as the currently known methods to achieve this are non-trivial to adapt
architectures and remain computationally expensive.
This work aims at bridging the gap between foundation models and the systematic generalization offered by equivariance. The key idea is to decouple equivariance from the main task-specific DNN. One way achieving this is through the idea of learned canonicalization function [30], where canonicalization network learns transform the data into canonical form. The pretrained prediction network then applied to this canonical form. Another way to achieve this decoupling is through the idea of symmetrization where all the transformations of datapoint are passed through the prediction network and the final output is computed as the average of these predictions [2 g However, symmetrization comes notably higher computational cost compared canonicalization, primarily due to the requirement of a forward pass through the prediction network for each data transformation. This problem is even more pronounced when employing large-scale prediction network. Hence, the canonicalization technique is a more efficient and practical choice for making large pretrained network equivariant.
Nevertheless, a naive application of the canonicalization idea fails in practice. This is because the canonicalization network's 's choice of canonical form result in change of distribution in the input the pretrained prediction network that is, canonicalization can be uninformed about the preference of the prediction network, undermining its performance.
As outlined in Figure 2 we resolve this misalignment by matching the distribution of predicted canonical forms with the dataset. We demonstrate empirically that imposing this prior is essential in an equivariant adaptation of pretrained models across different domains and datasets. Our approach offers practical solution to obtaining large-scale equivariant models by providing an independent module that can be plugged into existing large pretrained DNNs, making them equivariant to wide range of transformation groups.
2 Background
Groups and Transformations Groups capture the abstract structure of symmetry transformations, which is due to their compositional form - e.g., the composition of transformation with its inverse results in an identity transformation. Formally, a group G is set equipped with associative binary operation, such that the set closed under this operation, and each element E G has unique inverse. A group G can act set X by transforming its elements N' through bijection. Assuming X

---

is a vector space, the linear action of on x is given by p(g) x. The function 9 - GL(X) is called a group representation, and GL((X) the set of invertible linear transformations on X. For example, if g 50(2) is the group of 2D rotations, its action on any image x E X could rotate it around some center.
Equívariance in Deep Learning In the context of deep learning, we are interested in designing models that are invariant or equivariant to the input data's symmetries. Suppose that input data live in a space X' with some symmetry structure that we would like to preserve in our deep learning model.
An equivariant function f X y is function that commutes with the group action of G on X and y, i.e., for all g E G and x E X, we have:
(((gg) x) = p')) f(x)
where م': : T is another group representation that acts on the output space V In other words, an cquivariant function is a function whose output changes in a predictable way under transformations of the input induced by the group G. If p' is the trivial representation, i.e., (g(g) I or identity transformation for all E g, then we say that f is invariant to the group action of G.
Equivariance with Learned Canonicalization The concept of invariance involves ensuring that all elements within group orbit produce the same output when processed by a function f. To achieve this, elements be mapped to canonical sample from their orbit prior to applying the function. On the other hand, equivariance involves mapping elements to a canonical sample, applying the function, and then transforming them back to their original position within the orbit. This notion can be precisely expressed by representing the equivariant function f in a canonicalized form as follows:
f (x) - p'c((x) (p)
where the function p: X 3 is called the prediction function and the function c: :: + G is called the canonicalization function. In simpler terms, the function maps input data to group elements such that the inverse action these elements can return the data back to its canonical sample.
The function f in eq. (2) achieves equivariance for any prediction function long as the canonical- ization function also G-equivariant - that e (p (g) x) = ge (x) V g, x E G X. Previous work by (30] used a direct approach to design a canonicalization function using existing equivariant neural network architectures.
This formulation offers the advantage of removing the constraint on the main prediction network and placing instead on the network that learns the canonicalization function. As shown in (30), this decoupling can be partial, where the prediction network is itself equivariant to some symmetry transformations, and the canonicalization function used for equivariance to additional symmetries.
3 Method
The flexibility of the equivariance with canonicalization approach enables the conversion of any existing large pretrained Deep Neural Network (DNN) into an equivariant model with respect to certain known transformations. To achieve this, the pretrained DNN can be utilized as the prediction network in the formulation provided by eq. (2). Subsequently, a canonicalization function can be designed to produce the elements of the known group of transformation while maintaining equivariance with respect to this group.
One could learn the canonicalization function while optionally finetuning the prediction function using the same task objective- in our experiments, we consider both zero-shot and fine-tuning setup.
The performance of the model requires the alignment of these two networks. For example, the canonicalization network produces images that are upside-down, compared to those that the pretrained network is expecting, the overall performance is significantly degraded. In addition to alignment, there is augmentation effect that further muddies the water: during its training, the canonicalization network performs data augmentation. As we see shortly, one needs to consider both alignment and augmentation effects when analyzing the performance of this type of equivariant network.
When both networks are trained together from scratch, the alignment a non-issue, and (unwanted) augmentation degrade improve performance, depending the extent of symmetry in the
3

---

dataset. However, when dealing with pretrained prediction networks one needs to also consider the aligmment effect. One could then think of freezing the pretrained prediction network, therefore avoid- ing unwanted augmentation, and backpropagating the task loss through it to align the canonicalization network. However, this can become prohibitively expensive for large pretrained models, such as segment anything (SAM) considered in this work. We propose an alternative, where we directly regularize the canonicalization network to produce canonical forms consistent with the training data, which in turn aligns with the prediction network.
3.1 Learning Canonicalization Augmentation and Alignment
When learning the canonicalization function during training. the process implicitly performs dynamic augmentation of the prediction network. Consider a model designed to be equívariant to a certain group of transformations G by canonicalization. At the start of training, the randomly initialized weights of the canonicalization function will output random canonical orientations for each data point. This has similar effect to data augmentation using the group & for the prediction network. As training progresses, the canonical orientations for similar-looking images begin to converge, as demonstrated in Figure B] in Appendix [] causing the augmentation effect to diminish. Thus, in addition to guaranteeing equivariance, this formulation provides a free augmentation effect to the prediction network.
Table 1: Effect of augmentation on the Prediction network. Top-1 classification accuracy and G- Averaged classification accuracy for CIFAR10 and CIFAR 100 [31. C8-Avg Acc refers to the top-1 accuracy on the augmented test set obtained using the group G = C8, with each element of G applied on the original test set.
Dataset
+
CIFARIO
CIFARIO0
Prediction Network +
Rotation Augmentation
Acc
C8-Avg Ace
Acc
C8-Avg Acc
10 to + 10 degrees
90,96 0.41
44.87 1 0.60
74.83
0.15
37.14 1: 0.42
ResNe150 32
180 + 180 degrees
84.60 + 1.83
81.04 + 1.86
61.07
0.27
59.42 + 0.70
Learnd Canonicalization (LC) 30
83.11 + 0.35
82.89 1 0.41
59.84
1 0.67
59.45 : 0.49
ReaNet50 3
- 10 to + +10 degrees
96.97 : 0.01
$1.77 + 0,25
85.84 1: 0.10
44,86 = 0.12
(pretrained on ImtgeNet)
+ 180 degrees
94.91 0.07
90.11 1:
80.21 0.09
74.12 0.05
Learnd Canonicalization (LC)
30
93,29 = 0.01
92,96 + 0.09
77.52 = 0.07
However, there are two scenarios where the augmentation provided by learning the canonicalization function can be detrimental:
First, in cases where the training only requires small transformations as augmentations, providing all the transformations of a group G can actually hurt the performance of the prediction network during the start of training or fine-tuming. For example, in natural image datasets like CIFAR 10, small rotation augmentations (from -00 to 10 degrees) are beneficial, while a canonicalization function would output any rotation from - 180 to +180 degrees during the beginning phase of training. This can lead to unstable training of the prediction network and hinder the model's performance by training on additional data that is far outside the distribution of the train set. We show this in Table T by training a prediction network with different rotation augmentations, including the one due to learned canonicalization on both CIFAR datasets. Furthermore, we also observe that this effect is more pronounced when the prediction network trained from scratch, and the dataset more complicated with a larger number of classes. This effect can be understood as an example of the variance-inven invariance tradeoff introduced by [33]. Since, the test distribution is not perfectly symmetric under rotations, training with arbitrary augmentations biases the prediction function.
Second, we notice that relying solely on the task loss objective is not be sufficient for the canonical- ization function to learn the correct orientation. This could be due to the small size of the finetuning dataset compared to the pretraining dataset. We see experimentatlly that this leads to the canoni- calization function outputting inconsistent canonical orientations during inference, impacting the prediction network's performance. For instance, on CIFARIO (non-augme)ted), we expect the canon- ical orientation for every datapoint to be similar after training. However, from Figure in Appendix we can see that the canonical orientations for the test set are distributed uniformly from - 180 to + 1800 degrees even after training until convergence of the task objective. As result, during inference,

---

Predicted image orientations distribution
g
90121
(argmax)
Pg I.
p=
Pg...
Prior Regularisation (Minimising KL)
Canonicalization
Distribution of image
orientations dataset
Training
Inference
Figure Training and inference with our proposed regularized canonicalization method. The canonicalization function outputs distribution over image orientations used to canonicalize the input image. Additionally, during training, this predicted distribution is regularized to match the orientations seen in the dataset.
the prediction network will view images with different orientations and underperform. This issue arises from the fact that the prediction networks are not inherently robust to these transformations.
3.2 Canonicalization Prior
As have seen, the canonicalization network may induce shift in canonical orientations away from those present the pretraining datasets. To further encourage the canonicalization function to align inputs an orientation that will help the prediction network, we introduce a simple regularizer we call canonicalization prior (CP). It is motivated by noticing that the images or point clouds in the finetuning dataset provide useful information, of the training dataset. We, therefore. take as prior that the canonicalization should align inputs as closely as possible their original orientation in the finetuning dataset.
expect these orientations be similar those
We derive the regularizer by taking probabilistic point of view. The canonicalization function maps the group of transformations, denoted G. Let IP (o(x) denote
each data point distribution the distribution induced by the canonicalization function over 9 for given data point x. We assume the existence of canonicalization prior associated with the dataset D that has distribution Pp over G. To enforce the canonicalization prior, we seek to minimize the Kullback-Leibler (KL) divergence between PD and Fo(x) over the entire dataset DD that is Cprirr EX-D
[D (PPD || Pc(x)) )-
We assume that the canonicalization function e estimates the parameters of distribution over rotations with probability density function (R| cx)).. We denote the probability density function of the prior to be q (R). Since the prior distribution is kept constant, minimizing the KL divergence is equivalent to minimizing the cross-entropy, and the prior loss simplifies to:
Sprux
ExxD ER~g [log (R. (x))]
Hereafter, we derive the regularization for the discrete and continuous cases separately.
3.2.1 Discrete Rotations
We first consider the group of 2D discrete rotations, the cyclic group Cn. which can be seen discrete approximation to the full rotation group SO(2). In this case, we consider a categorical distribution over group elements, with the prior distribution having probability mass of 1 for the identity element. Then (R) ôR.I, where δR.1 is the Kronecker delta function and the cross entropy in eq. (3) becomes log Pe(a) (T). Hence, the loss becomes L = -I EEDD log Pe(w) (T). In other words, the regularization loss is simply the negative logarithm of the probability assigned by the canonicalization function to the identity element I of the group.
Practical Implementation For images, similar to 30
canonicalization network needs output
logits corresponding to every group element in the discrete cyclic group C,, This can be achieved by using a Group Convolutional Network [7] or an -tte--teerable Network (34) that produces outputs using regular representation. To design the canonicalization function, we take spatial average and
5

---

get logits corresponding to every element in the group along the fibre dimension. This is similar to the approach used [3] Now, we can get discrete distribution over the group elements by taking softmax and minimizing the prior loss along with the task objective. During training, to ensure consistency with the implementation in [30] for fair comparisons across all experiments, we utilize the argmax operation instead of sampling from this distribution using Gumbel Softmax [35] and employ the straight through gradient trick [36]. All our image-based experiments use this discrete rotation model.
3.2.2 Continuous rotations
When considering canonicalization with continuous rotations, it is natural to use the matrix Fisher distribution introduced by [37]. It the analogue of the Gaussian distribution the 5O(n) manifold and is defined as
p (R F)
n (F) exp (Tr [ [FTR])
where F E RR* is the parameter of the distribution and n (F) is normalization constant. Interpreta- tion of the parameter F and useful properties of the distribution are provided in [0). In particular, considering the proper singular value decomposition F USVT, we find that R UVT is the mode of the distribution and the singular values can be interpreted as concentration parameters in the different axes. We therefore S
sI obtain the isotropic version of the distribution,
p R R..))
1 n (s)
exp Tr ( RRRR)
(5)
where the normalization constant only depends (Theorem 2.1 of (39). Note that SO (2), this becomes the Von-Mises distribution as expected.
We introduce the following result, which will allow to derive the regularization.
Proposition 1. Let p and q be matrix Fisher distributions of R
p R Rp. sn) n (8p)
sp Tr RTR) q (R RR, Sq) n (s,))
exp (sy Tr R R
exp
The cross-entropy is given by ERog logg R Rp: 8p)] N ((5g 8p Tr RR)..) logo
(6)
where N (sy) only depends on Sg.
The proof follows in Appendix
Setting the location parameters of the estimated and prior distributions as Re(x)) and R. = respectively, find that the canonicalization prior eq. (3) is given
Cprior - (Re(x)) = Rc(()) - 1/1F
where we have eliminated terms that do not depend on Re(a)) and N (s²) sp. Following intuition, the strength of the regularization is determined by the concentrations of the distributions around their mode.
Practical Implementation For image domain, canonicalization network needs to output rotation matrices Re(s) (2) that equivariantly transforms with the input image. This can be achieved by using a ((2)-Steerable Network (34) that outputs two vector fields. design the canonicalization function we can take spatial average over both vector fields and Gram Schmidt orthonormalize the
vectors to get 2D rotation matrix. While this sounds promising in theory, in practice we found
empirically difficult to optimize using the loss to enforce canonicalization prior (see Appendix D). We

---

Table Performance comparison of large pretrained models finetuned on different vision datasets. Both classification accuracy and G-averaged classification accuracies are reported. Acc refers to the accuracy on the original test set, and C8-Avg Ace refers to the accuracy on the augmented test set obtained using the group G = C88 C8-Aug. refers to fine-tuning the pre-trained model with rotation augmentations restricted to C8.
Pretrained Large Prediction Network
ResNer50
VIT
Dutasets
Model
Acc
C8-Avg Acc
Acc
C8-Avg Acc
Vanilla
96.97 +
57.77 + 0.25
98.13
0.04
63.59 0.48
Rotation Augmentation
94.91 +: 0.07
90.11 + 0.19
96.26 + 0.15
93.67 + 0.39
CIFARIO
Learned Canonicalization (LC)
30
93.29 +: 0.01
92.96 +: 0.09
95.00 = 0.01
94.80
0.02
CS-Aug.
94.36
96.36
0.02
94.17
Prior- Regulaegglrr LC (ours)
96.19 + 0.01
95.31 = 0.17
96.14 0.14
95.08 = 0.10
Vanilla
85.84 +: 0.10
44.86 + 0.12
87.91 + 0.28
55.87 + 0.14
Rotation Augmentation
80.21 +: 0.09
74.12 0.05
82.59 - 0.44
78.39 + 0.89
CIFARTO0 31
Learned Canoniculization (LC)
30
78,50 +
77.52 -: 0.07
80.86 - 0.17
80.48 0.20
C8-Aug-
83.00 +
79.72 0.10
83.45 + 0.09
80.08 0.38
Prior Regullazgd LC (ours)
83.44 + 0.02
82.09 + 0.09
84.27 + 0.10
83.61 + 0.01
Vanilla
98,30 + 0.01
73.87 1.43
98.31 + 0.09
76.66 = 0.93
STLI0 22
Rotation Augmentation
98.08 +: 0.06
94.97 + 0.08
97.85 + 0.17
94.07 + 0.11
Learned Canonicalization (LC)
93.92 0.10
94.67 0.02
C8-Aug.
98.31 + 0.01
96.31 + 0.13
97.83 + 0.08
94.45 + 0.35
Prior- Regularigd LC (ours)
97.01 + 0.01
96,37 + 0.12
96.15 = 0.05
95.73 = 0.16
believe this warrants further investigation and present potential novel research direction explore. However, in the domain of point clouds, we discovered that combining the implementation from (30), which outputs 3D rotation matrices elements of SO(3), with the regularization loss to enforce the canonicalization prior leads to remarkably effective results. This approach is demonstrated in model for rotation-robust point cloud classification and part segmentation, leveraging pretrained PointNet 3 and DGCNN [4]] architectures (see Section 4.2p.
4 Experiments
In this section, we present experimental results on images and point clouds to evaluate our method of achieving equivariance with minimal modifications pretrained networks. The key benefit of our approach is showcased by demonstrating its robustness when evaluating out-of-disrribution data that arise from known transformations applied the test
Our code is available at
4.1 Image domain
Experiment Setup. We use ResNet50 32] and ViT-Base (5], which are pretrained on ImageNet-I
dataset 43 for our image experiments. These are widely used for image classification tasks, and their pretrained checkpoints are publicly available ?
We finetune these models on several benchmark natural image classification datasets, including CIFARI0 [3I] CIFAR100 B1] and STL10 [22. Moreover, we incorporate four different strategies to finetune the pretrained models. namely: Vanilla, 2) Rotation Augmentation, 3) Learn Canonicalization, and 4) Prior- Friur-Regularized Learned Canonicalization. For the canonicalization function, we use a C8-equivariant convolutional network. The details of the architecture are available in Appendix B We jointly train the canonicalization function and fine-tune the pretrained image classification networks. Our total loss given by Liotat Liine-tune + 3 Lpriors where Cyrior is defined in Eq. 3 L Lfine-tune tune refers to the cross-entropy loss for classification, and 3 is hyperparameter which is set to 100.
The Vanilla model refers to fine-tuning the pretrained checkpoints using data augmentations such as horizontal flips and small angle rotations, while the Rotation Augmentation model covers angles ranging from 180 to +180 degrees. Although Rotation Augmentation model is not equivariant,

---

our goal was to establish an useful baseline to measure the gain in generalization of our proposed test dataset resulting from rotating the images. Further, we also report
model to the results on discretized version of Rotation Augmentation, mentioned as C8-Aug. in Table 2 where the fine-tuning dataset is augmented with the application of group elements in C88. We employ the Learned Canonicalization method from [3), where the canonical orientation is learned with the task signal only, which in our experiment is the classification. Finally, our proposed Prior-Regularieed Learned Canonicalization approach adds a prior regularization loss that tries to map all images in the original training dataset to the identity group element e. We refer to the final two equivariant techniques LC and Prior-Rege figir-Regulrrized LC.
Evaluation Protocol. In order to test the robustness of the models on rotations, we introduce G-averaged test set that refers to an expanded dataset obtained by applying all group elements to each image, resulting in test dataset of size [g] times the original test set. In this section, we consider G as C8 (cyelic group with 8 rotations), thus |g] = 8. We report the top-1 classification accuracy achieved on the original as well as this augmented test set (referred to as C8-Average Accuracy) to evaluate both in distribution and cut-of-distribution performance of the models.
Results. We report the results of finetuning ResNet50 and ViT on CIFARIO, CIFARIO0, and STL10
with various strategies in Table 2 As anticipated, we found that large pretrained networks for images not robust to rotation transformations, as indicated by the large drop in performance from the accuracy to its C88aveeaaed counterpart for both ResNet50) and VIT. Nevertheless, we observe that ViT is more robust to rotations compared to ResNet50, which has also been observed by 44 We notice that augmenting with a full range of rotation angles during training improves the AveAaerage Accuracy demonstrated by our Rotation Augmentation baseline. However, it hurts the accuracy of the prediction network in the original test set and does not guarantee equivariance. Augmenting with necessary rotations in CE-Augmentation does not ensure equivariance to C8 but retains performance on the original test set and reduces the gap between original and C8-averaged accuracies.
LC guarantees equivariance, which can be seen from the minor difference between the accuracies of the original and augmented test sets. Nevertheless, in every dataset, we can observe a significant drop in accuracy for the original test set. We extensively discussed this issue in Section 3.1 However, with our Prior-Regulard test set while still being equivariant rotations. This demonstrates that this prior regularization on LC is a promising direction to improve the performance of large-pretrained models while guaranteeing robustness to -istr--distribution samples resulting from transformations like rotation.
LCC method, we are able to reduce the gap between the accuracy on the original
Ideally, the accuracy of the original test set should be nearly identical for both the Vanilla setup and our Prior-Ree Regularige LC method. However, corresponding accuracies. This disparity arises from the fact that the canonicalization model is unable map all data points (images) perfectly the identity element e, supported by our observations that the regularization loss for prior matching does not diminish to zero. We hypothesize that this stems from the intentional limitation in the expressivity of the canonicalization function, which is done on purpose in [30] to avoid adding excessive computational overhead to the overall architecture. Finally, we note that due to rotation artifacts, a small difference between C8-Average Accuracy and Accuracy on original test set is unavoidable.
observed a slight difference between their
4.1.1 Instance Segmentation
Experiment Sctup. We use MaskRCNN 45 and Segment Anything Model (SAM) 221. which are pretrained on Microsoft COCO 22 and SA-IB 211 datasets respectively. MaskRCNN is widely used for instance segmentation task, while SAM is a recently proposed foundational model that can leverage prompts (bounding boxes and key points) for instance segmentation, and their pretrained checkpoints publicly available In experiments, we evaluate these models on COCO 2017 dataset, i.c.. report zero-shot performance on validation (val) set and use ground truth bounding boxes as prompts for SAM. For the canonicalization function, we use C4-equivariant NeteResNet architecture. The details of the architecture are available in Appendix [ As MaskRCNN and SAM already trained for the instance segmentation task, we only train the canonicalization function using the prior loss Cprior in Eq. 3 to make them equivariant to C4 group.
ahchitectuur
8

---

Evaluation Protocol. To test the generalization capabilities of these models, we introduce and utilize the C4-averaged val set refers to eyclic group with 4 rotations for reducing the complexities with transforming ground truth boxes and masks) along with the original val set. We report the mask-mAP score for mask prediction on both these datasets, denoted respectively as mAP and C4-AAgg mAP in 3
Results. Owing to pre-training on larger datasets and the use of bounding boxes as prompts, SAM outperforms MaskRCNN in both mAP and C4-Avg mAP scores. The difference between mAP and C4-Avg mAP scores demonstrates that SAM is more robust than MaskRCNN in the case of these transformations.
However, observe that there a difference between these reported numbers for both models. With our prior regularized LC framework, we can achieve equivariance with any large pretrained model to the desired group (here C'4) while retaining the performance on the original val set. Further, we analyze the relationship between the expressivity of the canonicalization function and the downstream effect on mAP values on the original set in Table 3 We compare a C'4-equivariant convolutional network (G-CNN) with a C4-equivariant WideResNet (G-WRN) architecture. We observe that although equivariance guaranteed, less expressive canonicalization function leads to decreased performance in the original val set due to its inability to map complex images to identity,
Table 3: Zero-shot performance comparison of large pretrained segmentation models with and without trained canonicalization functions on COCO 2017 dataset [22]. Along with the number parameters in camonicalization and prediction network, report mAP and C'4aaverged mAP values. indicates G-CNN and indicates a more expressive G-WRN for canonicalization.
Pretrained Large Segmentation Network +
MaskRCNN (46.4 M)
SAM (641 M)
Datasets +
Model
mAP
C4-Avg mAP
mAP
C4-Avg mAP
Zero-shot M)
45.57
27.67
62.34
58.78
COCO 22
Prior-Reguee Prgol-Regglariized LC (0.2 M)
35.77
35.77
59.28
59.28
Prior-Regulare Regularigd LCl (1.9 M)
44.51
44.50
62.
62.13
4.2
Point Cloud Domain
Table 4: Classification accuracy of different pointcloud models on the ModelNet40 dataset 46 in different train/test scenarios and ShapeNet [47] Part segmentation mean IoUs over 16 categories in different train/test scenarios. z/9 here stands for training with x augmentation and testing with augmentation. z here stands for aligned data augmented by random rotations around the vertical/z axis and SO(3) indicates data augmented by random 3D rotations.
Task
Classification
Part Segmentation
Dataset +
ModelNer40
ShapeNet
Method
z/z
z/SO(3)
S0(3)/80(3)
z/SO(3)
S0(3)/50(3)
PointNet
85.9
19.6
74.7
38.0
62.3
DGCNN
90.3
33.8
88.6
49.3
78.6
VN-PointNet
77.5
77.5
77.2
72.4
72.8
VN-DGCNN
89.5
89.5
90.2
81.4
81.4
LCC-PintNeet
79.9
79.6
1.3
79.7
1.3
73.5 0.8
73.6 +
LC-DGCNN
88.7
88.8
1.9
90.0
1.1
78.4 + 1.0
78.5 + 0.9
Ours (with pretrained PointNet und DGCNN for each tusk)
no-aug/2
no-aug/SO(3)
no-aug/SO(3)
PRLC-PointNet
30
84.3 1.2
82.6
PRLC-DGCNN 30
90.2
90.2 1.3
84.3 + 0.8

---

Datasets. For our experiments involving point clouds, we utilized the ModelNet40 [6 and ShapeNet (47] datasets. The ModelNet40 dataset comprises 40 classes of 3D models, with a total of 12,311 I models. Among these, 9,843 models were allocated for training, while the remaining models were reserved for testing in the classification task. In the case of part segmentation, we employed the ShapeNet-par subset, which encompasses 16 object categories and over 30,000 models. We only train the canonicalization function using the prior loss Cprior in Eq.
Evaluation protocol. To ensure consistency and facilitate comparisons, we followed the established conventions set by [49] and adopted by [10] for the train/test rotation setup in the classification and segmentation tasks. The notation 2/ z/y indicates that transformation is applied during training, while transformation y is applied during testing. Typically, three settings are employed: 2/2, 2/ s3))(3). and SO(3) SO3)). Here, denotes data augmentation with rotations around the z-axis during training, while S0(3) represents arbitrary rotations. However, since we regularize the output of the canonicalization with the identity transformation, we trained our canonicalization function and fine-tuned our pretrained model without any rotation augmentation. During inference, we tested on both and S0(3) augmented test datasets.
Results. We present our results on Table Notably, our method showcased superior performance in terms of robustness, outperforming existing methods for point cloud tasks. Specifically, the inclusion of the prior loss has led to a significant improvement PointNet's performance compared to DGCNN. This observation aligns with our analysis in Section [31] where we highlight that training the prediction network with large rotations can hinder its performance and serve as a bottleneck for equivariance within the learnt canonicalization framework. The empirical evidence, particularly in the 50(3) (SO(3) results of vanilla PointNet and DGCNN, where we notice a more pronounced drop PointNet's performance, supports this and strengthens our findings.
Discussion
Out-of-disttibution generalization is important weakness of state-of-the-as deep models, where an important source of shift in distribution is due to application-speciin transformations of the data- from the change of colour palette, magnification and rotation in images to change of volume or pitch in sound. A mechanism for making pretrained models robust to such changes can have significant impact their adaptation to different domains. This paper takes the initial steps toward this goal by removing barriers to making the model equivariant through a lightweight canonicalization process. The proposed solution ensures that the pretrained model receives in-distribution data, while canoni- calization ensures equivariance and, therefore, adaptability to a family of ouiit--distiibutiin samples. Our extensive experimental results using different pretrained models, datasets, and modalities gives insight into this subtle issue and demonstrate the viability of our proposed solution.
Limitations and Future Work
An important limitation of our approach is the dependency on an equivariant canonicalization network, which imposes restrictions on the range of transformations to which we can adapt. Using the optimization approach of (30) offers more flexibility as it replaces the equivariant network with the outcome of an optimization process. However, this approach can raise efficiency concerns which merit further investigation. Another limitation of the current exposition is the limited group of transformations explored in experiments. In future, we hope to add a comprehensive evaluation of this approach to other groups of transformation. We also aim to address optimization challenges related to prior regularization for continuous rotations into our image framework, thereby expanding its capabilities. Moreover, this observation emphasizes the necessity to delve into effect of expressivity of the canonicalization function, as it plays crucial role in determining the overall performance of our model.
Networks [5]. By doing so, we aspire to incorporate
Another promising future direction we hope to explore more flexible adaptation through canonical- ization using examples from target domain. Such examples from target domain can, in principle, replace the prior knowledge of the transformation group assumed in equivariant networks, thereby bridging the gap between equivariant deep learning and domain adaptation.
10

---

Acknowledgements
This project was in part supported by CIFAR AT Chairs and NSERC Discovery grant. Computational resources are provided by Mila, ServiceNow Research and Digital Research Alliance (Compute Canada). The authors would like to thank Sebastien Paquet for his valuable feedback.
References
[1] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):277-23224 1998.
[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
[3] Charles R Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660, 2017.
[4] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross B. Girshick. Mask r-cnn. ICCV, 2017.
[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[6] Yann LeCun and Yoshua Bengio. Convolutional Networks for Images, Speech, and Time Series, page 255-258. MIT Press, Cambridge, MA, USA, 1998.
[7] Taco Cohen and Max Welling. Group equivariant convolutional networks. In Maria Florina Balcan and Kilian Q Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 2990-2999, New York, New York, USA, 20 22 Jun 2016. PMLR.
[8] Sourya Basu, Prasanna Sattigeri, Karthikeyan Natesan Ramamurthy, Vijil Chenthamarakshan, Kush R Varshney, Lav R Varshney, and Payel Das. Equi-tuning: Group equivariant fine-tuning of pretrained models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 6788-6796, 2023.
[9] Daniel Worrall and Max Welling. Deep spaces:: Equivariance over scale. Advances in Neural Information Processing Systems, 32, 2019.
[10] Congyue Deng, Or Litany, Yueqi Duan, Adrien Poulenard, Andrea Tagliasacchi, and Leonidas J Guibas. Vector neurons: A general framework for so (3)-equivariant networks. Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[11] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Velicković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint 2rX1v22104.13478, 2021.
[12] Omri Puny, Matan Atzmon, Edward J. Smith, Ishan Misra, Aditya Grover, Heli Ben- Ben-Hamu, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. URL /ttppss///oeerrevie wtt.ee/wwwe///orrmmme/e id=zIUy]55nXR
[13] Arnab Kumar Mondal, Pratheeksha Nair, and Kaleem Siddiqi. Group equivariant deep rein- forcement learning, arXiv preprint vrXiv:2007.03437, 2020.
[14] Arnab Kumar Mondal, Vineet Jain, Kaleem Siddiqi, and Siamak Ravanbakhsh. Eqr: Equivariant representations for data-efficient reinforcement learning. In International Conference on Machine Learning, pages 15908-15926. PMLR, 2022.
[15] Maurice Weiler and Gabriele Cesa. General (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems, 32, 2019.

---

[16] Hong-Xing Yu. Jiajun Wu, and Li Yi. Rotationally equivariant 3d object detection. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1456-1464 2022.
[17] Hai Wu, Chenglu Wen, Wei Li, Xin Li, Ruigang Yang, and Cheng Wang. Transformation- equivariant 3d object detection for autonomous driving. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 2795-2802, 2023.
[18] Erik Bekkers. Maxime W Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Roto-translation covariant convolutional networks for medical image analysis. In Medical Image Computing and Computer Assisted Intervention-MICCA. 2018: 21st In- ternational Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part I, pages 440-448 Springer, 2018.
[19] Dongdong Chen, Mike Davies, Matthias J. Ehrhardt, Prooia-Bibiane Schönlieb, Ferdia Sherry, and Julián Tachella. Imaging with equivariant deep learning: From unrolled network design to 147, 2023.
fully unsupervised learning. IEEE Signal Processing Magazine,
[20] Yilong Yang, Srinandan Dasmahapatra, and Sasan Mahmoodi. Scale-equivare histopathology image segmentation. In Erik Bekkers, Jelmer M. Wolterink, and Angelica Aviles-Rivero, editors, Proceedings of the First International Workshop on Geometric Deep Learning in Medical Image Analysis, volume 194 of Proceedings of Machine Learning Research, pages 130-148. PMLR, 18 Nov 2022.
unet for
[21] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything, 2023,
[22] Tsung-Yr Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740-755. Springer, 2014.
[23] Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pages 6105-61 14. PMLR, 2019.
[24] Ting Chen, Simon Korbllith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton. Big self-supervised models strong semi-supervin learners. Advances in neural information processing systems, 33:22243-22255, 2020.
[25] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33: 177711719 2020.
[26] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, PMLR, 2021.
pages
[27] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
2022.
[28] Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Yixuan Wei, Qi Dai, and Han Hu. On data scaling in masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10365-10374 2023.
[29] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bossclut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint Sa.ii::108.07255:: 2021.
[30] Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravan- bakhsh. Equivariance with learned canonicalization functions. In NewrIPS 2022 Workshop on Symmetry and Geometry in Neural Representations, 2022.
12

---

[31] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009.
[32] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[33] Shuxiao Chen, Edgar Dobriban, and Jane Lee. A group-theoretic framework for data augmenta- tion. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21321-21333, Curran Associates, Inc., 2020.
[34] Maurice Weiler and Gabriele Cesa. General e(2)eequivariant steerable cnns. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[35] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softme softmax. In International Conference on Learning Representations, 2017.
[36] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1308.3432, 3432, 2013.
[37] Thomas D Downs. Orientation statistics. Biometrika, 59(3):665-66) 676, 1972.
[38] Mohamed A Khamsi and William A Kirk. An introduction to metric spaces and fixed point theory. John Wiley & Sons, 2011.
[39] Taeyoung Lee. Bayesian attitude estimation with the matrix fisher distribution on so(3). IEEE Transactions on Automatic Control, 63(10):3377-3332 2018.
[40] David Mohlin, Josephine Sullivan, and Gérald Bianchi. matrix fisher distributions. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin. editors, Advances in Neural Information Processing Systems, volume 33, pages 4884-4893. Curran Associates, Inc., 2020.
Probabilistic orientation estimation with
[41] Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, and Justin M. Solomon. Dynamic graph enn for learning on point clouds. ACM Transactions on Graphics (TOG), 2019.
[42] Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsu- pervised feature learning. In Geoffrey Gordon, David Dunson, and Miroslav Dudík, editors, Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, volume 15 of Proceedings of Machine Learning Research, pages 215-223, PMLR, 2011.
[43] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. Ieee, 2009.
[44] Nate Gruver, Mare Finzi, Micah Goldblum, and Andrew Gordon Wilson. The lie derivative for measuring learned equivariance. arXiv preprint 2..022220000099: 2022.
[45] Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick. Mask r-enn. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[46] Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3d shapenets: A deep representation for volumetric shapes. pages 1912-1920, 2015.
[47] Angel X. Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li. Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, Jianxiong Xiao, Li Yi, and Fisher Yu. ShapeNet: An Information Rich 3D Model Repository. Technical Report arXiv:1512.03012 [cs.GR], 2015.
[48] Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph CNN for learning on point clouds. ACM Transactions on Graphics, 38,5):1-12, 2019.
13

---

[49] Carlos Esteves, Christine Allen- Blan-hlanhhete, Ameesh Makadia, and Kostas Danilidiis. Learning
SO(3) equivariant representations with spherical cnns. pages 52 68, 2018.
[50] Gabriele Cesa, Leon Lang, and Maurice Weiler. A program to build e(n)-qquivariant steerable CNNs. In International Conference on Learning Representations, 2022.
[51] Sergey Zagoruyko and Nikos Komodakis. arXiv:1605.071400 2016.
Wide residual networks. arXiv preprint
[52] Hugo Larochelle, Dumitru Erhan, Aaron Courville, James Bergstra, and Yoshua Bengio. An empirical evaluation of deep architectures on problems with many factors of variation. In Proceedings of the 24th international conference on Machine learning, pages 473-480, 2007.
14

---

A Proof of Proposition
Proof. The cross-entropy is given by
ER-g
logg
R. Rp: sp)]
o R R2 80) logp R Rpp 8p) dR,
sO(n)
where dR is the invariant Haar measure on (O(n). Here we assume that it is scaled such that [so(m) dR =8=1
We obtain
ER logp
R Rp s5)]
(R) Ro ss Tr R
log cogc(sp)) dR.
(9)
So(n)
ER logp R Rp. sp)
log e (sp) Sp Ir R ER RRI)
(10)
From Theorem 2.2 and Lemma 2.2 of 33 we have
dlog e (8g) Ry.
E R
R
(11)
dsg
Therefore, find
ER logp
R
Rp sp)]
d log c (s4),
sp R., log (sp),
(12)
d8g
which completes the proof.
B Architecture Details
We extensively use escnn library 11 50 to design equivariant convolutional architectures.
Each layer of C8-equivariant convolutional network consists of convolution with regular repre- sentation except the first layer, which maps the trivial representation of the C8 group to its regular representation. We use equivariant implementation of batch normalization, ReLU activation function, and dropout as proposed in [15] [50]. Major hyperparameters tuned include the number of layers, kernel sizes, dropout, and learning rates. C4-equivariant convolutional network utilized in Section 4.1.1 has an identical architecture, with each layer being equivariant to C4 instead of C8.
C4-equivariant WideResNet includes repetitive stacking of equivariant versions of basic residual blocks on several consecutive bottleneck residual blocks (for details on these residual blocks, we refer readers to Figure 1 in [5I). The rest of the architecture details and hyperparameters are identical to the design of C8-equivariant convolutional network.
C Additional analysis experiments
Augmentation effects
We provide the canonical orientation predicted by Learned Canonicalization
(LC) [30] for Rotated MNIST [52] at the start of the training and after training along with original images. The progression of canonized images at the start and after training demonstrates the diminishing effect of augmentation effects when trained with learned canonicalization.
Effect of prior regularization. We present a comprehensive analysis of the output distribution over 8 discrete angles predicted by the canonicalization function, both before and after training, the test set. These findings are depicted in Figure for Learned Canonicalization (LC) [30], and Figure 5 for Prior Regulaeegu LC. Here, the numbers 0 to 7 correspond to angles that multiples of 45 degrees, ranging from 0 to 315 degrees, respectively.
We demonstrate that incorporating prior regularization into the canonicalization function aids in mapping the images the identity prior (represented by angle). This improvement positively impacts the accuracy the original test set, evidenced by the results in Table 2 Converscly,
15

---

3
J
K
Y
3
4
Canonized images at the training
Ganonized images
Original Images
training
Figure Visualization of the diminishing augmentation effect introduced by learning canonicalization 33) during training for rotated MNIST dataset. In this visualization, the leftmost image represents the original training images. Moving towards the center, we present the canonized images at the beginning of the training process. Finally, the rightmost image unveils the transformation of the canonized images after training the model for 100 epochs.
relying solely the classification task loss yields significant alteration in the angle distribution, as the post-training test set distributions remain random.
Additionally, provide valuable insights into the fraction of images mapped to the identity element in Table It is important to note that the expressivity of the canonicalization function, specifically employing lightweight equivariant networks, contributes to the inability to map all images to the iden- tity elements. This observation calls for further exploration in understanding the role of expressivity and generalization within canonicalization networks with known prior orientations.
Table 5: Fraction of images mapped identity.
Training Completed
Dataset
Model
X
Learned Canonicalization (LC)
0.11
0.23
CIFAR1O 31
Prior- RigolReegulaizzed L.C. (ours)
0.11
0.76
listogrem of output angles
Histogram of output angles
7000
2000
1500
2000
1000
500
Anzzles
Anrzles
Distribution of angles at the start of training
Discribution of angles after training
Figure 4: Distribution of angles output from canonicalization function C8 for Learned Canon- icalization (30 for CIFARI0 31 before and after training. We use indices on the r-axis instead of angle values represent the corresponding multiple of 45", Frequency denotes the number of images mapped to particular multiple of 45"
16

---

Hittogram nutput angles
7050
1550
10.0
Iequeryy
2020
3020
2000
1020
Distribution of angles at the start of training
Distribution of angles after training
Figure 5: Distribution of angles output from canonicalization function in C8 for Learned Canonical- ization [3 with prior regularization for CIFARIO [I] before and after training. We use indices on the z-axis instead of angle values to represent the corresponding multiple of 45% Frequency denotes the number of images mapped to particular multiple of 45
Optimization challenges on continuous 2D rotations for Images
We investigate the prior introduced in eq. (7) specifically for images. To achieve this, employ E-2 steerable network (34 to devise a canonicalization function that generates rotation matrices. Our approach involves initially generating two vector fields, which are subsequently averaged across the spatial dimension to obtain two vectors. By performing Gram-Schmidr orthonormalization on these vectors, derive rotation matrix.
In Figure 6. we present the distributions predicted angles ranging from 180 to 1800 Through analysis, we observe that the mean and standard deviation of the predicted angles the test set are -0.54 and 80.23, respectively. This observation signifies that while the prior guides the canonicalization function to output angles close O, there exist instances where angles significantly deviate from this central value. This phenomenon warrants further investigation and a complete understanding of this left future work.
Histogram of output or argles
Histogrom of output angles
403
225
300
100
-130
-200
-150
-1.00
100
130
Ancles
And m
Distribution angles the start training
Distribution angles training
Figure 6: Distribution of angles output from steerable canonicalization function in SO2 for Learned Canonicalization (30) with prior regularization (eq. ( for CIFARIO 31 before and after training. rraxis denotes angles from - 180 to +180. Frequency denotes the number of images mapped to particular angle.
17


Reviews:
[{"review_id": "Ap0heeEzJZ", "reviewer": "Reviewer_qZrQ", "summary": "The work proposes a method for making large pre-trained models equivariant to specified group actions. The method uses the canonicalization approach, where the input is oriented to a specific orientation before being fed to the pre-trained model with the help of a trainable canonicalization network. The authors use a prior distribution over the group elements to regularize the output of the canonicalization network. The proposed technique shows robustness with respect to the group actions in different downstream tasks, such as classification and part segmentation.", "questions": "1. Should the Kronecker delta functions in line 172 be normalized by the number of group elements for discrete priors?\n2. What is the specific structure of the **canonicalization** module used in this work?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper pursues a novel approach to make pre-trained models scale equivariant.\n2. The proposed technique is simple and effective, making it well-suited for practical adaptation.\n3. The work is well-written and easy to follow.", "weaknesses": "1. Experimental Setup: One major drawback of this work is the experimental setup. The authors used priors over the discrete group elements (i.e., 8 discrete rotations of C8 for the image domain) for the experiments. During the evaluation on the test set, the images were augmented by the action of group C8. This approach is not appropriate because we are providing the model with information about the test set augmentation during training.\n    \n    On the other hand, for the data augmentation setup, the baseline models were fine-tuned using random rotations between -180 to 180 degrees (i.e., all possible rotations). In this case, no information about the test set augmentation was given to the baseline.\n    \n    This difference in augmentation approaches makes the comparisons inappropriate. An appropriate comparison should either:\n    \n    1. Use a uniform continuous prior over the group elements.\n    2. Perform data augmentation only with the discrete group actions while training the baseline.\n    \n    Option 1 is more practical because in the real world, we often do not know specific priors over the group actions.\n    \n2. Recent studies have shown that fine-tuning large pre-trained models on small datasets hampers downstream tasks [1] and makes the model vulnerable to out-of-distribution data. Experiments with frozen pre-trained backbones and trainable classification/segmentation modules would be more justifiable.\n\n1.Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful feedback. Below we address the reviewer's concerns. \n\n> Concerns on experimental setup: [...] appropriate comparison should [...] perform data augmentation only with the discrete group actions while training the baseline.\n\nWe thank the reviewer for pointing out this issue. We now provide the results for fine-tuning with C8 augmentations in Table 2 in the rebuttal pdf. Our proposed canonicalization method is better than the suggested baseline across all datasets and significantly so in the case of CIFAR100. Regarding uniform continuous prior, Appendix C details some of the optimization challenges when incorporating this with steerable networks. As mentioned in lines 314-316, a complete understanding of this and an extension of our method to continuous rotation in images is left as a future work. \n\n> Fine-tuning and frozen backbone architectures:  [...] experiments with frozen pre-trained backbones and trainable classification/segmentation modules would be more justifiable\n\nWe understand the justification of the reviewer’s suggestion to use frozen pre-trained backbones and trainable modules. To address this concern for more realistic settings, we provide the result for instance segmentation on COCO 2017 [1] (val set) with Segment Anything Model [2] (SAM) and MaskRCNN [3], where the pre-trained models were not fine-tuned (due to computational resource constraint). Our training includes only training the canonicalization function to map images in the train set of COCO 2017 to identity. As presented in Table 1 in the rebuttal pdf, we observe that our equivariant adaptation obtains mask-mAP scores on the original test set identical to the zero-shot performance of the model while being equivariant to the instance segmentation task. \n\n> Should the Kronecker delta functions in line 172 be normalized by the number of group elements for discrete priors?\n\nSince we are putting all probability mass of the prior on identity for natural images, we are using a Kronecker delta function. If we use the normalized Kronecker delta function as suggested, the prior will not be a probability distribution. \n\n> What is the specific structure of the canonicalization module used in this work?\n\nWe thank the reviewer for pointing out this omission, and we will add the following architectural details to the appendix of the final version of our paper:\n\nWe extensively use the $\\texttt{escnn}$ library [4, 5] to design equivariant convolutional architectures.\n- For classification experiments, we use a $C8$-equivariant convolutional network; each layer consists of convolution with regular representation except the first layer, which maps the trivial representation of the $C8$ group to its regular representation. We use equivariant implementation of batch normalization, ReLU activation function, and dropout as proposed in [4, 5].\n- For instance segmentation, we use a $C4$-equivariant WideResNet, which includes repetitive stacking of equivariant versions of *basic* residual blocks on several consecutive *bottleneck* residual blocks (for details on these residual blocks, we refer readers to Figure 1 in [6]). The rest of the architecture details and hyperparameters are identical to the design of the $C8$-equivariant convolutional network.\n\n$$$$\n$$$$\n\n[1] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740–755. Springer, 2014.\n\n[2] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything, 2023.\n\n[3] Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.\n\n[4] Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems, 32, 2019.\n\n[5] Gabriele Cesa, Leon Lang, and Maurice Weiler. A program to build e(n)-equivariant steerable CNNs. In International Conference on Learning Representations, 2022.\n\n[6] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146, 2016.\n"}, {"review_id": "t2WgfgW3of", "reviewer": "Reviewer_45sc", "summary": "Paper proposes to enforce equivariance/invariance in pretrained models by using the recently proposed canonicalization functions that transform a given input to its canonical form. This allows one to enforce symmetry during the finetuning stage without the need for retraining the model. Paper uses a prior so that the canonicalization network is biased to output identity transformation for the images in the finetuning dataset (i.e., a canonical orientation is defined as the orientation in the finetuning dataset). Empirically, this is shown to be more robust than without the prior. ", "questions": "Q1. Section 4.1: why is the rotation augmentation baseline applied over all angles from -180 to 180 (lines 235-236)? I believe the group is known to be C8 for the LC and prior-regularized LC methods, and should be used for the baseline as well. \n\nQ2. Section 4.1: Were rotations part of the data augmentations that were performed on the model during its pretraining on ImageNet?\n\nQ3. Section 4.2: VN/CN/PCN- in Table 3 are not defined; I am not sure what these baselines are. Which dataset were the PointNet and DGCNN models pretrained on? Do all the other methods in Table 3 use pretrained models?\n\nQ4. Line 247 says that proposed approach tries to map all images in the \"original training dataset\" to identity. Should it say \"finetuning dataset\" instead?\n", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Imposing invariances (via canonicalization) during finetuning rather than pretraining can be very beneficial as different finetuning datasets can have different symmetries.\nAuthors clearly motivate the challenges of using canonicalization naively for a pretrained model that do not exist when training from scratch (e.g., Figure 3) and propose a simple regularization to solve it. Experiments show better performance across multiple finetuning datasets and pretrained models. \n", "weaknesses": "\nW1. Details of the overall loss function is not presented or is assumed to be known to the reader. It is harder to evaluate the impact of the prior without these details.\n\nW2. In Section 3.1, the first drawback of Learned Canonicalization is its undesired augmentation effect during the initial phase of finetuning. \n- However, this should exist even for the proposed prior-regularized method because the network still begins with a random initialization.\n- To solve this, one probably has to encode the prior of outputting identity within the architecture (for example, by defining c(x):= I + r(x) where I is the identity and r(x) initialized with very small weights). \n\nW3. The assumption that the orientations in the finetuning and pretraining datasets are \"similar\" (line 153) is not well defined. It may be more reasonable to apply the prior using (a small subset of) the pretraining dataset. This will ensure that the finetuning images match the orientation that the model was pre-trained on.\n\nW4. Experiments only evaluate invariance and not equivariance in its general form. I think it is important to demonstrate the advantage of imposing end-to-end equivariance over a pretrained model (as this is different from imposing equivariance in its intermediate representations).\n\nW5. Paper can be strengthened with experiments on other transformations than rotations (but not strictly required). \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their useful comments on the paper. \n\n> Missing overall loss function equation: Details of the overall loss function is not presented or is assumed to be known to the reader [...]\n\nWe thank the reviewer for pointing out the absence of an overall loss function. The total loss function consists of task loss (which is cross-entropy loss) and prior loss. We will add an expression for total loss as $\\mathcal{L} = \\mathcal{L_{task}} + \\lambda * \\mathcal{L_{prior}}$, where $\\lambda$ is a hyperparameter. We refer the reviewer to eq. (3), which describes the prior loss, and further to line 171, which demonstrates that it reduces to cross-entropy loss over the number of discrete rotations. We set $\\lambda$ as 100 in our experiments.\n\n> Effect of initial undesired augmentation in canonicalization model: [...]\n\nWe agree with the reviewer that if the prediction network is not frozen when the network begins to train, the problem of the undesired augmentation effect exists in the case of our proposed prior-regularized canonicalization. However, this undesired effect quickly reduces with training with the regularization, as demonstrated in Figures 4 and 5 in the Supplementary material (Appendix C, Effect of prior regularization paragraph). Encoding the prior in the architecture, as suggested by the reviewer, breaks the equivariance of the canonicalization function, and therefore the entire setup is no longer guaranteed to be equivariant. \n\n> Similar orientations in pre-training and fine-tuning dataset: The assumption that the orientations in the finetuning and pretraining datasets are \"similar\" (line 153) is not well defined. It may be more reasonable to apply the prior using (a small subset of) the pretraining dataset.\n\nWe agree with the suggestion and will add a sentence to point out this possibility. While the proposed suggestion is interesting, in the natural image datasets we considered, the canonical orientations are aligned in the finetuning and pretraining datasets. Hence, we assume that this canonical orientation is close to the identity element (or upright images). However, learning or inferring this prior automatically using the pretraining dataset could be an exciting future research direction of this work.  \n\n> Experiments with equivariance tasks:  Experiments only evaluate invariance and not equivariance in its general form. \n\nWe appreciate the reviewer's concern that the current set of experiments does not evaluate equivariance and rather focuses primarily on invariance tasks. To address this, we provide the result for instance segmentation on COCO 2017 [1] (val set) with Segment Anything Model  [2] (SAM) and MaskRCNN [3], where the pre-trained models were not fine-tuned. Our training includes only training the canonicalization function to map images in the train set of COCO 2017 to identity. As presented in Table 1 in the rebuttal pdf, we observe that our equivariant adaptation obtains mask-mAP scores on the original test set identical to the zero-shot performance of the model while being equivariant to the instance segmentation task.\n\n> Results for fine-tuning with $C8$ augmentations: [...] the group is known to be C8 for the LC and prior-regularized LC methods, and should be used for the baseline as well [...]\n\nWe now provide the results for fine-tuning with C8 augmentations in Table 2 in the rebuttal pdf. Our proposed canonicalization method is better than the suggested baseline across all datasets.\n\n> Pre-training setup:  Were rotations part of the data augmentations that were performed on the model during its pretraining on ImageNet\n\nWe want to point out that we did not perform the pretraining of large models but rather utilized the widely available checkpoints, as mentioned in line 229. However, the training setup of ResNet-50 and ViT can be found in [4] and [5], which shows that large rotations were not part of data augmentations which explains the poor G-Average Accuracy.\n\n> Description of abbreviations: Expand VN/CN/PCN in Table 3 and do all the other methods in Table 3 use pre-trained models?\n\nWe thank the reviewer for pointing this out. Abbreviations VN and CN are taken from [6] and refer to Vector Neuron and Learned Canonicalization (LC). PCN stands for Prior Regularized LC. In the final version, we will change CN and PCN to LC and Prior Regularized LC to make the abbreviations consistent across all domains.\nFor classification, we took checkpoints trained on ModelNet40, and for instance segmentation, we took checkpoints trained on ShapeNet dataset. The other methods in the table are trained from scratch. Note that we weren't aware of any foundation model for the PointCloud domain but still wanted to test our idea and show prior regularized finetuning can improve learned canonicalization [6]. We will add these details on pretrained checkpoints in Section 4.2. \n\n> Typo in the submission: Line 247 says [...] Should it say \"finetuning dataset\" instead?\n\nThanks for pointing out the typo. \n\n$$$$\n$$$$\n\n[1] Lin, T. et al. Microsoft coco: Common objects in context. In Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740–755. Springer, 2014.\n\n[2] Kirillov, A. et al. Segment anything, 2023.\n\n[3] He, K. et al. Mask r-cnn. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.\n\n[4] He, K. et al. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016\n\n[5] Dosovitskiy, A. et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.\n\n[6] Kaba, S. et al. Equivariance with Learned Canonicalization Functions Proceedings of the 40th International Conference on Machine Learning, PMLR 202:15546-15566, 2023"}, {"review_id": "MIAmfi3Zbx", "reviewer": "Reviewer_vNov", "summary": "The paper proposed a fine-tuning based adaptation method for equivariance in pre-trained models. The core of the method is the canonicalization function, which is trained to normalize the input to a common orientation. The paper investigates challenges associated with training canonicalization functions and proposes a canonicalization prior regularization to alleviate those challenges. Experiments in image and point-cloud domains demonstrate that fine-tuning a backbone network with the proposed method improves equivariant properties of a backbone model.", "questions": "I suggest authors address weaknesses for the rebuttal.", "limitations": "Limitation section nicely highlights limitations of the proposed approach.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is very clearly written and easy to follow. All parts of the paper/method are nicely motivated.\n\n2. The paper tackles an important problem of efficient equivariance in neural networks. This problem is especially relevant with regard to large models.\n\n3. The paper provides valuable insights into a practical application of canonicalization functions to achieve equivariance/invariance.\n", "weaknesses": "In no particular order.\n\n1. The paper argues about the equivariant adaptation, but only invariance (and only to the rotation group) is demonstrated in the experiments.\n\n2. Although there is a conceptual novelty in the paper (in the part of analyzing the challenge of applying canonicalization functions), the proposed adaptation approach relies heavily on the prior work of [11]. With this, the methodological novelty of the proposed approach is rather limited.\n\n3. The proposed method requires fine-tuning the whole model together with the canonicalization function. This can be troublesome, especially for large models. With this, I am also wondering if the \"equivariant adaptation\" is the right name, because the essence of the method seems to be more in equivariant fine-tuning. It would be nice to have experiments, where the backbone model is frozen, and only the canonicalization function is trained.\n\n4. The paper claims out-of-distribution robustness as a key benefit of the method (L224). However, it is not clear how the canonicalization function is itself robust to the distribution shift. If the canonicalization function is trained on cifar10 and tested on cifar100, for example, will it still be able to deliver reasonable distribution over a group orbit? In that sense, it is important to note that the proposed method does not provide guaranteed equivariance as equivariant networks do. \n\n5. I am not sure that comparing methods based on raw Accuracy is a fair evaluation protocol. If the goal of equivariant adaptation is to robustify a model (to a symmetry group of interest), then it seems more appropriate to keep track of the relative change of accuracy from original to G-averaged test sets. Otherwise, it is hard to disentangle if the performance improvement is due to better equivariance or due to improved data augmentation (and a higher accuracy on the non-transformed dataset as a result). This seems to be an important part of the comparison, which is missing from the paper. For example, consider Table 2 results for VIT on Cifar10, the LC model loses 0.2% of its Accuracy, while Regularized LC loses 1% of its Accuracy. Which one would we call more robust?\n\nOthers:\n1. Possibly useful related work to resolve the issue discussed in L208 - 219. Moskalev et al. LieGG: Studying Learned Lie Group Generators. NeurIPS22 https://arxiv.org/abs/2210.04345\n2. Missed related work. Tai et al. Equivariant Transformer Networks. ICML19. https://arxiv.org/abs/1901.11399\n\nUPD: Authors addressed most of my concerns in the rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and suggestions.\n\n> Experiments on equivariant tasks: The paper argues about the equivariant adaptation, but only invariance (and only to the rotation group) is demonstrated in the experiments.\n\nWe appreciate the reviewer's concern. To address this, we provide results of instance segmentation on COCO 2017 [1] (val set) with Segment Anything Model [2] (SAM) and MaskRCNN [3], where the pre-trained models were not fine-tuned. Our training includes only training the canonicalization function $c$ to map images in the train set of COCO 2017 to identity. As presented in Table 1 in the rebuttal pdf, we observe that our equivariant adaptation obtains mask-mAP scores on the original test set identical to the zero-shot performance of the model while being equivariant to the instance segmentation task. The original models are far from being equivariant, as highlighted by our results. Finally, extending our framework to transformations beyond rotations with approximate equivariance will be an interesting future direction.\n\n> Experiments with frozen backbone:  fine-tuning the whole model together with the canonicalization function. This can be troublesome, especially for large models. [...] nice to have experiments, where the backbone model is frozen, and only the canonicalization function is trained.\n\nWe now provide results for instance segmentation in Table 1 in the rebuttal pdf, where the setup is identical to the reviewer’s suggestion, i.e., we freeze the Segment Anything Model and MaskRCNN and only trained the canonicalization function $c$.\n\n> Out of distribution generalization of canonicalization function, and equivariance:  [...] If the canonicalization function is trained on cifar10 and tested on cifar100, for example, will it still be able to deliver reasonable distribution over a group orbit? In that sense, it is important to note that the proposed method does not provide guaranteed equivariance as equivariant networks do [...]\n\nWe would like to highlight the fact that the out-of-distribution in the submission refers to the robustness to the rotation of the original images. Since the canonicalization function $c$ is equivariant to rotation groups, contrary to the statement in the review, it is guaranteed that $c$ and the entire pipeline is equivariant and robust to distribution shift. In the particular case of CIFAR 10 and CIFAR 100, since the datasets are identical, their canonicalization networks can be identical and can transfer. In general, out-of-distribution generalization to transformations beyond those handled through equivariance remains an interesting direction for future.  \n\n> Fair comparison metrics: I am not sure that comparing methods based on raw Accuracy is a fair evaluation protocol. [...]  it seems more appropriate to keep track of the relative change of accuracy from original to G-averaged test sets.[...]\n\nBoth the LC [3] and our Prior-Regularized LC are equivariant, which means, in principle, there should be no gap between accuracy and G-averaged accuracy. Therefore the suggested metric will not be informative. In practice, we observe small discrepancies between Accuracy and G-Averaged Accuracy, which we attribute to the rotation artifacts which destroy information in the image. We will add a sentence to the paper to clarify this. \n\n> Additional related work:\n\nWe thank the reviewer for their constructive feedback and suggestion of these related works, which we will incorporate into our work.\n\n$$$$\n$$$$\n[1] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything, 2023.\n\n[2] Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.\n\n[3] Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with Learned Canonicalization Functions Proceedings of the 40th International Conference on Machine Learning, PMLR 202:15546-15566, 2023\n\n[4] Moskalev et al. LieGG: Studying Learned Lie Group Generators. NeurIPS22."}, {"review_id": "3Q6E1YHT7R", "reviewer": "Reviewer_nPdb", "summary": "The authors tackle an important problem: equivariant NNs are very useful because of capturing important symmetries of the data that lead to higher sample efficiency and more accurate and robust predictions. Yet, EqNNs are very hard to scale because of their specialized architectures. In particular, they cannot benefit from large scale pre-traininig. \n\nThe authors propose to use the canonical localization network (an EqNN) from [11] to equip pre-trained model with equivariance. The authors train the localization NN jointly with the prediction network. They test both on the original and augmented test set (with various transformations). \n\nThe authors identify that directly applying the approach from [11] does not work because the localization network can transform the input into atypical views for the training dataset. Then they propose a KL regularization term that encourages a delta function at the identify transformation. Their approach works effectively both for discrete and continuous groups. ", "questions": "See the above Weaknesses. My comments should be treated as questions that could improve the quality of your work. I think you have an awesome paper, but I also think that addressing W1-3 could make your work even more impactful.", "limitations": "The authors adequately addressed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "S1: The idea of equipping pre-trained models with equivariance is very timely and potentially useful for large-scale applications.\n\nS2: The regularization term is well-motivated and effective.\n\nS3: The experiments are thorough, including both discrete and continuous groups.", "weaknesses": "W1: It would be useful to study your approach in more realistic setting where you do not have to augment the test dataset but rather take an existing test dataset. I'd recommend standard out-of-distribution evaluation on IN-{adversarial, sketch, renditions}, INv2 etc. Another domain for experimentation is in Transfer Learning from IN-1K pertaining. There are some standard datasets, such as Flowers-102, Food, etc. \n\nW2: Since you want to integrate equiavriance with pre-trained models, I think you should also consider SSL pre-trained checkpoints, since they are the foundation of SOTA pre-training (Pre-training on IN-1K may not give you the best features for downstream tasks). Some pre-trained checkpoints are available on github, e.g. https://github.com/rdangovs/essl/tree/main/imagenet/simclr.\n\nW3: As you acknowledged, your model is constrained by the EqNN architectures for the canonical localization network. I wonder if you could consider in your main text discussion more flexible (albeit non-exact) approaches for encouraging symmetry in predictions, such as MSE (https://arxiv.org/abs/2303.02484) and E-SSL (Equivariant Contrastive Learning).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of our work and constructive feedback to improve the quality of the paper.\n\n> More realistic settings without augmented test sets:  I'd recommend standard out-of-distribution evaluation on IN-{adversarial, sketch, renditions}, INv2 etc. Another domain for experimentation is in Transfer Learning from IN-1K pertaining. There are some standard datasets, such as Flowers-102, Food, etc. \n\nWe would like to highlight that the out-of-distribution in the submission refers to the robustness to transformations that form a group, like rotation of the original images. OOD generalization with other suggestions such as IN-{adversarial, sketch, renditions}, INv2 datasets involve potentially non-linear transformations for which we do not have the corresponding equivariant network required for canonicalization. Extending our framework to transformations beyond groups with approximate equivariance will be an interesting future direction of research.\n\n> Experiments with SSL pretrained checkpoints: [...] I think you should also consider SSL pre-trained checkpoints [...] \n\nWe now provide the full fine-tuning results for the E-SSL checkpoint (with best acc1 on the GitHub link shared by the reviewer) on CIFAR10 in Table 3 in the rebuttal pdf. Interestingly, while the Accuracy was similar to vanilla classification checkpoints, the C8-Average Accuracy for E-SSL was poorer. Using our method, we are able to observe similar Accuracy as well as equivariance. \n\n> Non-equivariant architecture for canonicalization: I wonder if you could consider in your main text discussion more flexible (albeit non-exact) approaches for encouraging symmetry in predictions, such as MSE [...] and E-SSL\n\nIn addition to the E-SSL results above, we now also provide results for training non-equivariant canonicalization function $c$ in Table 2 in the rebuttal pdf (row titled ConvNet Canonicalizer). We use a simple ConvNet as $c$ and train it to predict the angle for a rotated image, i.e., the prior loss is replaced with an MSE loss between the predicted and actual angle of rotation. As a result, we observe “approximate”  equivariance with a slight drop in the Accuracy over the original test sets."}]

Author rebuttal:
We thank all the reviewers for their valuable and constructive feedback. We are pleased to see that they found our work simple (reviewers **45sc**, **qZrQ**),  novel, and effective (reviewers **qZrQ**, **nPdb**)  in addressing a relevant, important (**vNov**) and timely (**nPdb**) problem.  In the general response, attached rebuttal PDF, and individual reviewer responses, we believe to have addressed most if not all, reviewer questions and concerns. In particular, we are including significant new results on the Segment Anything Model [1] (SAM) and MaskRCNN [2] for equivariant segmentation, as well as additional baselines and ablations for invariant classification. These new results address common concerns on 1) equivariant tasks, 2) adapting with completely frozen weights, 3) missing baseline with C8 augmentation, 4) using E-SSL embeddings 5) a new baseline using a non-equivariant canonicalization network.                 

Below, we give more context on these new experiments and refer the reviewers to the new and updated tables in the attached rebuttal PDF:

1. We provide results of instance segmentation on COCO 2017 [3] (val set) with Segment Anything Model (SAM) and MaskRCNN, where the pre-trained models were not fine-tuned, in part, due to the size of these pre-trained models. Our training includes only training the canonicalization function to map images in the train set of COCO 2017 to identity. As presented in Table 1 in the rebuttal pdf, we observe that our equivariant adaptation obtains mask-mAP scores on the original test set identical to the zero-shot performance of the model while being equivariant to the instance segmentation task. As this task is an equivariant task where we don’t finetune the large foundational prediction model, this addresses reviewers' concerns about only focussing on invariant tasks and fine-tuning the prediction network. We show that we can achieve equivariance by just training the canonicalization function (canonicalizer) and attaching it to a foundation model.  

2. Based on reviewers’ feedback, we provide some additional baselines for the classification tasks on the image datasets: 
    - we have added a baseline with C8 augmentation during fine-tuning, and the corresponding results are available in Table 2 in the rebuttal pdf (row titled C8-Aug.).   
    - CIFAR10 results using the suggested E-SSL [4] checkpoint for ResNet-50 also appear in Table 3 in the rebuttal pdf.
    - Finally, based on the suggestion of Reviewer **nPdb**, we explore a setup for pre-trained ResNet-50 where we replace an equivariant canonicalizer with a non-equivariant convolutional network (ConvNet Canonicalizer). We will add the architectural details in the appendix of the final version of our paper. We train the ConvNet Canonicalizer to predict the angle for a rotated image, i.e., the prior loss is replaced with an MSE loss between the predicted and actual angle of rotation. As a result, we observe “approximate”  equivariance with a slight drop in the Accuracy over the original test sets. The accuracy values are provided in Table 2 (row titled ConvNet Canonicalizer).

$$$$
We plan to include all these new results in the final version of the paper. We hope that the reviewers can reassess our work and its contributions in light of this.

$$$$
$$$$

[1] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything, 2023.

[2] Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.

[3] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740–755. Springer, 2014.

[4] Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Cheung, Pulkit Agrawal, and Marin Soljacic. Equivariant self-supervised learning: Encouraging equivariance in representations. In International Conference on Learning Representations, 2022.

Meta-review: Three out of four reviewers advocate for the acceptance of this paper. The reviewer that recommended rejection did not respond during the rebuttal discussion period. The AC have checked over the author's rebuttal and believe that the reviewer’s concerns are partially addressed. Overall, the AC finds this paper interesting and has practical implications. The AC trusts that the authors will include all the results and polish the writing.
Decision: Accept (poster)