SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Arnab Kumar Mondal*†  
Mila, McGill University  
ServiceNow Research
Siba Smarak Panigrahi*  
Mila, McGill University
Sékou-Oumar Kaba  
Mila, McGill University
Sai Rajeswar ServiceNow Research
Siamak Ravanbakhsh  
Mila, McGill University
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.
Deep neural networks (DNNs) have demonstrated exceptional success in dealing with a range of data modalities, such as image and point cloud [1-5]. The majority of these applications require working with input data that undergoes various transformations, such as rotation, scaling, or translation in image data. Invariant and equivariant networks are "aware" of such transformations; they are specifically designed to 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-11, 13, 14]. As such, they have emerged as a promising solution to a wide range of computer vision and point cloud processing tasks [7, 9, 10, 15-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 [21, 23-28]. Pre-training on massive datasets has emerged as a practical method to harness this advantage. Several large models are now publicly available, and fine-tuning them on specific applications has proven to be a successful strategy in a variety of applications. However, such foundation models [29] are typically not equivariant
† Correspondence: arnab.mondal@mila.quebec. Work done while at ServiceNow Research
Figure 1: Predicted masks from the Segment Anything Model (SAM) [21], showcasing both the original model and our proposed equivariant adaptation for  $90^{\circ}$  counter-clockwise rotated input images taken from the COCO 2017 dataset [22]. Our method makes SAM equivariant to the group of  $90^{\circ}$  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 4.1.1.
to most transformations except translations, as the currently known methods to achieve this are non-trivial to adapt to some 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 of achieving this is through the idea of learned canonicalization function [30], where a canonicalization network learns to transform the data into a canonical form. The pretrained prediction network is then applied to this canonical form. Another way to achieve this decoupling is through the idea of symmetrization where all the transformations of a datapoint are passed through the prediction network and the final output is computed as the average of these predictions [12, 8]. However, symmetrization comes at a notably higher computational cost compared to 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 a 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 choice of canonical form can result in a change of distribution in the input of 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 a 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 a wide range of transformation groups.
Groups and Transformations Groups capture the abstract structure of symmetry transformations, which is due to their compositional form - e.g., the composition of a transformation with its inverse results in an identity transformation. Formally, a group  $\mathcal{G}$  is a set equipped with an associative binary operation, such that the set is closed under this operation, and each element  $g\in \mathcal{G}$  has a unique inverse. A group  $\mathcal{G}$  can act on a set  $\mathcal{X}$  by transforming its elements  $\mathbf{x}\in \mathcal{X}$  through a bijection. Assuming  $\mathcal{X}$
is a vector space, the linear action of  $g$  on  $\mathbf{x}$  is given by  $\rho(g) \cdot \mathbf{x}$ . The function  $\rho: \mathcal{G} \to \mathrm{GL}(\mathcal{X})$  is called a group representation, and  $\mathrm{GL}(\mathcal{X})$  is the set of invertible linear transformations on  $\mathcal{X}$ . For example, if  $\mathcal{G} = SO(2)$  is the group of 2D rotations, its action on any image  $\mathbf{x} \in \mathcal{X}$  could rotate it around some center.
Equivalence 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  $\mathcal{X}$  with some symmetry structure that we would like to preserve in our deep learning model.
An equivariant function  $f: \mathcal{X} \to \mathcal{Y}$  is a function that commutes with the group action of  $G$  on  $\mathcal{X}$  and  $\mathcal{Y}$ , i.e., for all  $g \in G$  and  $\mathbf{x} \in \mathcal{X}$ , we have:
$$
f (\rho (g) \cdot \mathbf {x}) = \rho^ {\prime} (g) \cdot f (\mathbf {x}) \tag {1}
$$
where  $\rho': G \to T'$  is another group representation that acts on the output space  $\mathcal{V}$ . In other words, an equivariant function is a function whose output changes in a predictable way under transformations of the input induced by the group  $\mathcal{G}$ . If  $\rho'$  is the trivial representation, i.e.,  $\rho'(g) = I$  or identity transformation for all  $g \in \mathcal{G}$ , then we say that  $f$  is invariant to the group action of  $\mathcal{G}$ .
Equivariance with Learned Canonicalization The concept of invariance involves ensuring that all elements within a group orbit produce the same output when processed by a function  $f$ . To achieve this, elements can be mapped to a 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 (\mathbf {x}) = \rho^ {\prime} (c (\mathbf {x})) \cdot p \left(\rho (c (\mathbf {x}) ^ {- 1}) \cdot \mathbf {x}\right) \tag {2}
$$
where the function  $p: \mathcal{X} \to \mathcal{Y}$  is called the prediction function and the function  $c: \mathcal{X} \to \mathcal{G}$  is called the canonicalization function. In simpler terms, the function  $c$  maps input data to group elements such that the inverse action of these elements can return the data back to its canonical sample.
The function  $f$  in eq. (2) achieves equivariance for any prediction function as long as the canonicalization function is also  $\mathcal{G}$ -equivariant – that is  $c(\rho(g) \cdot \mathbf{x}) = g c(\mathbf{x}) \forall g, \mathbf{x} \in \mathcal{G} \times \mathcal{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 it 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 is used for equivariance to additional symmetries.
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, if 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 an 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 is a non-issue, and (unwanted) augmentation can degrade or improve performance, depending on the extent of symmetry in the
dataset. However, when dealing with pretrained prediction networks one needs to also consider the alignment effect. One could then think of freezing the pretrained prediction network, therefore avoiding 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 equivariant to a certain group of transformations  $\mathcal{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 a similar effect to data augmentation using the group  $\mathcal{G}$  for the prediction network. As training progresses, the canonical orientations for similar-looking images begin to converge, as demonstrated in Figure 3 in Appendix C, 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  $\mathcal{G}$ -Averaged classification accuracy for CIFAR10 and CIFAR100 [31].  $C8$ -Avg Acc refers to the top-1 accuracy on the augmented test set obtained using the group  $\mathcal{G} = C8$ , with each element of  $\mathcal{G}$  applied on the original test set.
Dataset →	CIFAR10	CIFAR100
Prediction Network ↓	Rotation Augmentation	Acc	C8-Avg Acc	Acc	C8-Avg Acc
ResNet50 [32]	-10 to +10 degrees	90.96 ± 0.41	44.87 ± 0.60	74.83 ± 0.15	37.14 ± 0.42
-180 to +180 degrees	84.60 ± 1.83	81.04 ± 1.86	61.07 ± 0.27	59.42 ± 0.70
Learned Canonicalization (LC) [30]	83.11 ± 0.35	82.89 ± 0.41	59.84 ± 0.67	59.45 ± 0.49
ResNet50 [32]	-10 to +10 degrees	96.97 ± 0.01	57.77 ± 0.25	85.84 ± 0.10	44.86 ± 0.12
(pretrained on ImageNet)	-180 to +180 degrees	94.91 ± 0.07	90.11 ± 0.19	80.21 ± 0.09	74.12 ± 0.05
Learned Canonicalization (LC) [30]	93.29 ± 0.01	92.96 ± 0.09	78.50 ± 0.15	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  $\mathcal{G}$  can actually hurt the performance of the prediction network during the start of training or fine-tuning. For example, in natural image datasets like CIFAR10, small rotation augmentations (from  $-10$  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 1 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 is trained from scratch, and the dataset is more complicated with a larger number of classes. This effect can be understood as an example of the variance-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 canonicalization 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 experimentatly that this leads to the canonicalization function outputting inconsistent canonical orientations during inference, impacting the prediction network's performance. For instance, on CIFAR10 (non-augmented), we expect the canonical orientation for every datapoint to be similar after training. However, from Figure 4 in Appendix C, we can see that the canonical orientations for the test set are distributed uniformly from  $-180$  to  $+180$  degrees even after training until convergence of the task objective. As a result, during inference,
Figure 2: Training and inference with our proposed regularized canonicalization method. The canonicalization function outputs a 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 we have seen, the canonicalization network may induce a shift in canonical orientations away from those present in the pretraining datasets. To further encourage the canonicalization function to align inputs in 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, as we expect these orientations to be similar to those of the training dataset. We, therefore, take as prior that the canonicalization should align inputs as closely as possible to their original orientation in the finetuning dataset.
We derive the regularizer by taking a probabilistic point of view. The canonicalization function maps each data point to a distribution over the group of transformations, denoted as  $\mathcal{G}$ . Let  $\mathbb{P}_{c(\mathbf{x})}$  denote the distribution induced by the canonicalization function over  $\mathcal{G}$  for a given data point  $\mathbf{x}$ . We assume the existence of a canonicalization prior associated with the dataset  $\mathcal{D}$  that has a distribution  $\mathbb{P}_{\mathcal{D}}$  over  $\mathcal{G}$ . To enforce the canonicalization prior, we seek to minimize the Kullback-Leibler (KL) divergence between  $\mathbb{P}_{\mathcal{D}}$  and  $\mathbb{P}_{c(\mathbf{x})}$  over the entire dataset  $\mathcal{D}$  that is  $\mathcal{L}_{\mathrm{prior}} = \mathbb{E}_{\mathbf{x} \sim \mathcal{D}}\left[D_{KL}(\mathbb{P}_{\mathcal{D}} \| \mathbb{P}_{c(\mathbf{x})})\right]$ .
We assume that the canonicalization function  $c$  estimates the parameters of a distribution over rotations with probability density function  $p(\mathbf{R} \mid c(\mathbf{x}))$ . We denote the probability density function of the prior to be  $q(\mathbf{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:
$$
\mathcal {L} _ {\text {p r i o r}} = - \mathbb {E} _ {\mathbf {x} \sim \mathcal {D}} \mathbb {E} _ {\mathbf {R} \sim q} [ \log p (\mathbf {R} \mid c (\mathbf {x})) ] \tag {3}
$$
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  $C_n$ , which can be seen as a 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  $p_{\mathcal{D}}(\mathbf{R}) = \delta_{\mathbf{R},\mathbf{I}}$ , where  $\delta_{\mathbf{R},\mathbf{I}}$  is the Kronecker delta function and the cross entropy in eq. (3) becomes  $-\log p_{c(x)}(\mathbf{I})$ . Hence, the loss becomes  $\mathcal{L}_{\mathrm{prior}} = -\mathbb{E}_{x\sim \mathcal{D}}\log p_{c(x)}(\mathbf{I})$ . In other words, the regularization loss is simply the negative logarithm of the probability assigned by the canonicalization function to the identity element  $\mathbf{I}$  of the group.
Practical Implementation For images, similar to [30], the canonicalization network needs to output logits corresponding to every group element in the discrete cyclic group  $C_n$ . This can be achieved by using a Group Convolutional Network [7] or an  $E(2)$ -Steerable Network [34] that produces outputs using regular representation. To design the canonicalization function, we take a spatial average and
get logits corresponding to every element in the group along the fibre dimension. This is similar to the approach used in [30]. Now, we can get a discrete distribution over the group elements by taking a 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 is the analogue of the Gaussian distribution on the  $SO(n)$  manifold and is defined as
$$
p (\mathbf {R} \mid \mathbf {F}) = \frac {1}{n (\mathbf {F})} \exp \left(\operatorname {T r} \left[ \mathbf {F} ^ {T} \mathbf {R} \right]\right) \tag {4}
$$
where  $\mathbf{F} \in \mathbb{R}^{n \times n}$  is the parameter of the distribution and  $n(\mathbf{F})$  is a normalization constant. Interpretation of the parameter  $\mathbf{F}$  and useful properties of the distribution are provided in [38-40]. In particular, considering the proper singular value decomposition  $\mathbf{F} = \mathbf{U}\mathbf{S}\mathbf{V}^T$ , we find that  $\hat{\mathbf{R}} \equiv \mathbf{U}\mathbf{V}^T$  is the mode of the distribution and the singular values  $\mathbf{S}$  can be interpreted as concentration parameters in the different axes. We therefore set  $\mathbf{S} = s\mathbf{I}$  to obtain the isotropic version of the distribution,
$$
p \left(\mathbf {R} \mid \hat {\mathbf {R}}, s\right) = \frac {1}{n (s)} \exp \left(s \operatorname {T r} \left[ \hat {\mathbf {R}} ^ {T} \mathbf {R} \right]\right) \tag {5}
$$
where the normalization constant only depends on  $s$  (Theorem 2.1 of [39]). Note that on  $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  $\mathbf{R}$
$$
p \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {p}, s _ {p}\right) = \frac {1}{n \left(s _ {p}\right)} \exp \left(s _ {p} \operatorname {T r} \left[ \hat {\mathbf {R}} _ {p} ^ {T} \mathbf {R} \right]\right), q \left(\mathbf {R} \mid \hat {\mathbf {R}} _ {q}, s _ {q}\right) = \frac {1}{n \left(s _ {q}\right)} \exp \left(s _ {q} \operatorname {T r} \left[ \hat {\mathbf {R}} _ {q} ^ {T} \mathbf {R} \right]\right).
$$
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] = N \left(s _ {q}\right) s _ {p} \operatorname {T r} \left(\hat {\mathbf {R}} _ {p} ^ {T} \hat {\mathbf {R}} _ {q}\right) + \log c \left(s _ {p}\right) \tag {6}
$$
where  $N(s_q)$  only depends on  $s_q$ .
The proof follows in Appendix A
Setting the location parameters of the estimated and prior distributions as  $\mathbf{R}_{c(x)}$  and  $\hat{\mathbf{R}}_q = \mathbf{I}$  respectively, we find that the canonicalization prior eq. (3) is given by
$$
\mathcal {L} _ {\text {p r i o r}} = - \lambda \operatorname {T r} \left(\mathbf {R} _ {c (x)}\right) = \frac {\lambda}{2} \left\| \mathbf {R} _ {c (x)} - \mathbf {I} \right\| _ {F} \tag {7}
$$
where we have eliminated terms that do not depend on  $\mathbf{R}_{c(x)}$  and  $\lambda = N(s_q)s_p$ . 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  $\mathbf{R}_{c(x)}\in SO(2)$  that equivariantly transforms with the input image. This can be achieved by using a  $E(2)$ -Steerable Network [34] that outputs two vector fields. To design the canonicalization function we can take a spatial average over both vector fields and Gram Schmidt orthonormalize the vectors to get a 2D rotation matrix. While this sounds promising in theory, in practice we found it empirically difficult to optimize using the loss to enforce canonicalization prior (see Appendix D). We