SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 2: Performance comparison of large pretrained models finetuned on different vision datasets. Both classification accuracy and  $\mathcal{G}$ -averaged classification accuracies are reported. Acc refers to the accuracy on the original test set, and  $C8$ -Avg Acc refers to the accuracy on the augmented test set obtained using the group  $\mathcal{G} = C8$ .  $C8$ -Aug. refers to fine-tuning the pre-trained model with rotation augmentations restricted to  $C8$ .
Pretrained Large Prediction Network →	ResNet50	ViT
Datasets ↓	Model	Acc	C8-Avg Acc	Acc	C8-Avg Acc
CIFAR10 [31]	Vanilla	96.97 ± 0.01	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
Learned Canonicalization (LC) [30]	93.29 ± 0.01	92.96 ± 0.09	95.00 ± 0.01	94.80 ± 0.02
C8-Aug.	95.76 ± 0.07	94.36 ± 0.09	96.36 ± 0.02	94.17 ± 0.08
Prior-Regularized LC (ours)	96.19 ± 0.01	95.31 ± 0.17	96.14 ± 0.14	95.08 ± 0.10
CIFAR100 [31]	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
Learned Canonicalization (LC) [30]	78.50 ± 0.15	77.52 ± 0.07	80.86 ± 0.17	80.48 ± 0.20
C8-Aug.	83.00 ± 0.09	79.72 ± 0.10	83.45 ± 0.09	80.08 ± 0.38
Prior-Regularized LC (ours)	83.44 ± 0.02	82.09 ± 0.09	84.27 ± 0.10	83.61 ± 0.01
STL10 [42]	Vanilla	98.30 ± 0.01	73.87 ± 1.43	98.31 ± 0.09	76.66 ± 0.93
Rotation Augmentation	98.08 ± 0.06	94.97 ± 0.08	97.85 ± 0.17	94.07 ± 0.11
Learned Canonicalization (LC) [30]	95.30 ± 0.19	93.92 ± 0.10	95.11 ± 0.01	94.67 ± 0.02
C8-Aug.	98.31 ± 0.01	96.31 ± 0.13	97.83 ± 0.08	94.45 ± 0.35
Prior-Regularized 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 a potential novel research direction to explore. However, in the domain of point clouds, we discovered that combining the implementation from [30], which outputs 3D rotation matrices or elements of  $SO(3)$ , with the regularization loss to enforce the canonicalization prior leads to remarkably effective results. This approach is demonstrated in our model for rotation-robust point cloud classification and part segmentation, leveraging pretrained PointNet [3] and DGCNN [41] architectures (see Section 4.2).
In this section, we present experimental results on images and point clouds to evaluate our method of achieving equivariance with minimal modifications to pretrained networks. The key benefit of our approach is showcased by demonstrating its robustness when evaluating out-of-distribution data that arise from known transformations applied to the test set.
Our code is available at https://github.com/arnab39/EquivariantAdaptation
4.1 Image domain
Experiment Setup. We use ResNet50 [32] and ViT-Base [5], which are pretrained on ImageNet-1K dataset [43] for our image experiments. These are widely used for image classification tasks, and their pretrained checkpoints are publicly available  $^{1,2}$ . We finetune these models on several benchmark natural image classification datasets, including CIFAR10 [31], CIFAR100 [31], and STL10 [42]. Moreover, we incorporate four different strategies to finetune the pretrained models, namely: 1) Vanilla, 2) Rotation Augmentation, 3) Learn Canonicalization, and 4) Prior-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 is given by  $\mathcal{L}_{total} = \mathcal{L}_{fine - tune} + \beta \cdot \mathcal{L}_{prior}$ , where  $\mathcal{L}_{prior}$  is defined in Eq. 3,  $\mathcal{L}_{fine - tune}$  refers to the cross-entropy loss for classification, and  $\beta$  is a 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,
$^{1}$ https://pytorch.org/vision/stable/models/generated/torchvision.models.resnet50.html
$^{2}$ https://pytorch.org/vision/stable/models/generated/torchvision.models.vit_b_16.html
our goal was to establish an useful baseline to measure the gain in generalization of our proposed model to the out-of-distribution test dataset resulting from rotating the images. Further, we also report results on a 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 C8. We employ the Learned Canonicalization method from [30], where the canonical orientation is learned with the task signal only, which in our experiment is the classification. Finally, our proposed Prior-Regularized 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 as LC and Prior-Regularized LC.
Evaluation Protocol. In order to test the robustness of the models on rotations, we introduce  $\mathcal{G}$ -averaged test set that refers to an expanded dataset obtained by applying all group elements to each image, resulting in a test dataset of size  $|\mathcal{G}|$  times the original test set. In this section, we consider  $\mathcal{G}$  as  $C8$  (cyclic group with 8 rotations), thus  $|\mathcal{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 out-of-distribution performance of the models.
Results. We report the results of finetuning ResNet50 and ViT on CIFAR10, CIFAR100, and STL10 with various strategies in Table 2. As anticipated, we found that large pretrained networks for images are not robust to rotation transformations, as indicated by the large drop in performance from the accuracy to its  $C8$ -averaged 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  $C8$ -Average Accuracy as 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  $C8$ -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-Regularized LC method, we are able to reduce the gap between the accuracy on the original test set while still being equivariant to 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 out-of-distribution samples resulting from transformations like rotation.
Ideally, the accuracy of the original test set should be nearly identical for both the Vanilla setup and our Prior-Regularized LC method. However, we observed a slight difference between their corresponding accuracies. This disparity arises from the fact that the canonicalization model is unable to map all data points (images) perfectly to 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.
4.1.1 Instance Segmentation
Experiment Setup. We use MaskRCNN [45] and Segment Anything Model (SAM) [21], which are pretrained on Microsoft COCO [22] and SA-1B [21] 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 are publicly available  $^{3}$ . In our experiments, we evaluate these models on COCO 2017 dataset, i.e., report zero-shot performance on validation (val) set and use ground truth bounding boxes as prompts for SAM. For the canonicalization function, we use a  $C4$ -equivariant WideResNet architecture. The details of the architecture are available in Appendix B. As MaskRCNN and SAM are already trained for the instance segmentation task, we only train the canonicalization function using the prior loss  $\mathcal{L}_{prior}$  in Eq. 3 to make them equivariant to  $C4$  group.
<sup>3</sup>https://pytorch.org/vision/main/models/generated/torchvision.models.detectionmaskedrcnn_resnet50_fpn_v2.html <sup>4</sup>https://github.com/facebookresearch/segment-anything
Evaluation Protocol. To test the generalization capabilities of these models, we introduce and utilize the  $C4$ -averaged val set ( $C4$  refers to cyclic 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$ -Avg mAP in Table 3.
Results. Owing to its 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, we observe that there is a difference between these two 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  $C4$ ) 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 val set in Table 3. We compare a  $C4$ -equivariant convolutional network (G-CNN) with a  $C4$ -equivariant WideResNet (G-WRN) architecture. We observe that although equivariance is guaranteed, a 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 of parameters in canonicalization and prediction network, we report mAP and  $C4$ -averaged mAP values.  $\dagger$  indicates G-CNN and  $\ddagger$  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
COCO [22]	Zero-shot (0 M)	45.57	27.67	62.34	58.78
Prior-Regularized LC† (0.2 M)	35.77	35.77	59.28	59.28
Prior-Regularized LC‡ (1.9 M)	44.51	44.50	62.13	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.  $x / y$  here stands for training with  $x$  augmentation and testing with  $y$  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 →	ModelNet40	ShapeNet
Method ↓	z/z	z/SO(3)	SO(3)/SO(3)	z/SO(3)	SO(3)/SO(3)
PointNet [3]	85.9	19.6	74.7	38.0	62.3
DGCNN [41]	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
LC-PointNet	79.9 ± 1.3	79.6 ± 1.3	79.7 ± 1.3	73.5 ± 0.8	73.6 ± 1.1
LC-DGCNN	88.7 ± 1.8	88.8 ± 1.9	90.0 ± 1.1	78.4 ± 1.0	78.5 ± 0.9
Ours (with pretrained PointNet and DGCNN for each task)
no-aug/z	no-aug/SO(3)	no-aug/SO(3)
PRLC-PointNet [30]	84.1 ± 1.1	84.3 ± 1.2	82.6 ± 1.3
PRLC-DGCNN [30]	90.2 ± 1.4	90.2 ± 1.3	84.3 ± 0.8
Datasets. For our experiments involving point clouds, we utilized the ModelNet40 [46] and ShapeNet [47] datasets. The ModelNet40 dataset comprises 40 classes of 3D models, with a total of 12,311 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-part subset, which encompasses 16 object categories and over 30,000 models. We only train the canonicalization function using the prior loss  $\mathcal{L}_{prior}$  in Eq. 7.
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  $x / y$  indicates that transformation  $x$  is applied during training, while transformation  $y$  is applied during testing. Typically, three settings are employed:  $z / z$ ,  $z / SO(3)$ , and  $SO(3) / SO(3)$ . Here,  $z$  denotes data augmentation with rotations around the  $z$ -axis during training, while  $SO(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  $z$  and  $SO(3)$  augmented test datasets.
Results. We present our results on Table 4. 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 in PointNet's performance compared to DGCNN. This observation aligns with our analysis in Section 3.1, 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  $SO(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.
Out-of-distribution generalization is an important weakness of state-of-the-art deep models, where an important source of shift in distribution is due to application-specific 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 a significant impact on 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 canonicalization ensures equivariance and, therefore, adaptability to a family of out-of-distribution 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.
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  $E(2)$ -Steerable Networks [15]. By doing so, we aspire to incorporate 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 a crucial role in determining the overall performance of our model.
Another promising future direction we hope to explore is more flexible adaptation through canonicalization using examples from a target domain. Such examples from a 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.
This project was in part supported by CIFAR AI 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.
[1] Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, 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 J Guibas. Pointnet: Deep learning on 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 scale-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. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[11] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
[12] Omri Puny, Matan Atzmon, Edward J. Smith, Ishan Misra, Aditya Grover, Heli Ben-Hamu, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=zIUyj55nXR.
[13] Arnab Kumar Mondal, Pratheeksha Nair, and Kaleem Siddiqi. Group equivariant deep reinforcement learning. arXiv preprint arXiv: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 e (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems, 32, 2019.
[49] Carlos Esteves, Christine Allen-Blanchette, Ameesh Makadia, and Kostas Daniilidis. 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)-equivariant steerable CNNs. In International Conference on Learning Representations, 2022.
[51] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146, 2016.
[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.
Proof. The cross-entropy is given by
$$
\mathbb {E} _ {\mathbf {R} \sim q} \left[ \log p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) \right] = \int_ {S O (n)} q \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {q}, s _ {q}\right) \log p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) \mathrm {d} \mathbf {R}, \tag {8}
$$
where  $\mathrm{d}\mathbf{R}$  is the invariant Haar measure on  $SO(n)$ . Here we assume that it is scaled such that  $\int_{SO(n)}\mathrm{d}\mathbf{R} = 1$ .
$$
\mathbb {E} _ {\mathbf {R} \sim q} \left[ \log p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) \right] = \int_ {S O (n)} q \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {q}, s _ {q}\right) \left(s _ {p} \operatorname {T r} \left[ \hat {\mathbf {R}} _ {p} ^ {T} \mathbf {R} \right] - \log c \left(s _ {p}\right)\right) \mathrm {d} \mathbf {R}, \tag {9}
$$
$$
\mathbb {E} _ {\mathbf {R} \sim q} \left[ \log p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) \right] = s _ {p} \operatorname {T r} \left(\hat {\mathbf {R}} _ {p} ^ {T} \mathbb {E} _ {\mathbf {R} \sim q} [ \mathbf {R} ]\right) - \log c (s _ {p}). \tag {10}
$$
From Theorem 2.2 and Lemma 2.2 of [39], we have
$$
\mathbb {E} _ {\mathbf {R} \sim q} [ \mathbf {R} ] = \frac {d \log c (s _ {q})}{d s _ {q}} \hat {\mathbf {R}} _ {q}. \tag {11}
$$
Therefore, we find
$$
\mathbb {E} _ {\mathbf {R} \sim q} \left[ \log p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) \right] = \frac {d \log c (s _ {q})}{d s _ {q}} s _ {p} \hat {\mathbf {R}} _ {q} - \log c (s _ {p}), \tag {12}
$$
B Architecture Details
We extensively use esccnn library [15, 50] to design equivariant convolutional architectures.
Each layer of  $C8$ -equivariant convolutional network 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 [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 [51]). The rest of the architecture details and hyperparameters are identical to the design of  $C8$ -equivariant convolutional network.
C Additional analysis experiments