SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sourya Basu *
Pulkit Katdare *
Prasanna Sattigeri
Vijil Chenthamarakshan
Katherine Driggs-Campbell
Payel Das
Lav R. Varshney
Efficient transfer learning algorithms are key to the success of foundation models on diverse downstream tasks even with limited data. Recent works of Basu et al. (2023) and Kaba et al. (2022) propose group averaging (equitune) and optimization-based methods, respectively, over features from group-transformed inputs to obtain equivariant outputs from non-equivariant neural networks. While Kaba et al. (2022) are only concerned with training from scratch, we find that equitune performs poorly on equivariant zero-shot tasks despite good finetuning results. We hypothesize that this is because pretrained models provide better quality features for certain transformations than others and simply averaging them is deleterious. Hence, we propose  $\lambda$ -equitune that averages the features using importance weights,  $\lambda s$ . These weights are learned directly from the data using a small neural network, leading to excellent zero-shot and finetuned results that outperform equitune. Further, we prove that  $\lambda$ -equitune is equivariant and a universal approximator of equivariant functions. Additionally, we show that the method of Kaba et al. (2022) used with appropriate loss functions, which we call equizero, also gives excellent zero-shot and finetuned performance. Both equitune and equizero are special cases of  $\lambda$ -equitune. To show the simplicity and generality of our method, we validate on a wide range of diverse applications and models such as 1) image classification using CLIP, 2) deep Q-learning, 3) fairness in natural language generation (NLG), 4) compositional generalization in languages, and 5) image classification using pretrained CNNs such as Resnet and Alexnet.
Group-equivariant deep learning leverages group equivariance as an inductive bias to design efficient and reliable neural networks. Popular examples include convolutional neural networks (CNNs) equivariant to translations (LeCun et al., 1989), group convolutional neural networks (GCNNs) equivariant to general discrete groups (Cohen & Welling, 2016), and recently Alphafold2 equivariant to 3D rotations (Jumper et al., 2021). But these methods cannot leverage pretrained models.
With the increase in open sourced large pretrained models, it is now crucial to develop efficient transfer learning algorithms that can leverage group equivariance. Basu et al. (2023) proposed equitune, an equivariant finetuning method that uses group averaging over features extracted from pretrained
Figure 1: Implementation of  $\lambda$ -equitune on CLIP. Weighted average of image features corresponding to transformed inputs are computed, which is further used for computing text-image similarity scores.
models. Several other methods were proposed to get equivariant output from non-equivariant backbone architectures, e.g. some form of averaging (Puny et al., 2021; Atzmon et al., 2022) or optimization over certain proxy loss functions (Kaba et al., 2022). But these latter methods were originally designed for training from scratch and not much is known about their finetuning abilities.
Here, we find that equitune performs poorly on zero-shot tasks. Our first main contribution is to show that the optimization method of Kaba et al. (2022) when used with appropriate proxy loss functions provides excellent zero-shot and finetuning performance. We call this method equizero.
The results from equizero suggest that pretrained models provide better quality features for some group transformations than others. Thus, we propose  $\lambda$ -equitune, which learns importance weights directly from the data and uses them to perform a weighted group averaging. We show  $\lambda$ -equitune outperforms equitune and is competitive with equizero for zero-shot tasks. Moreover, for finetuning,  $\lambda$ -equitune often outperforms both equitune and equizero. This constitutes our second main contribution.
To validate our methods, we provide experiments on a diverse set of pretrained models and datasets. We show zero-shot performance of equizero on: 1) image classification using CLIP, 2) deep Q-learning, 3) fairness in natural language generation (NLG), and 4) compositional generalization in languages. For CLIP and deep Q-learning, we used the naturally available loss functions, namely, similarity scores and  $Q$ -values, respectively. For fairness in NLG and compositional generalization, we closely follow the setup of Basu et al. (2023), and we use regard scores Sheng et al. (2019) and the negative of the maximum of the output probability distribution, respectively, as the loss function.
We first show results of  $\lambda$ -equitune for image classification using CLIP, finding that  $\lambda$ -equitune performs competitively with equizero and outperforms equitune for both zero-shot and finetuning tasks. Then, we show a simple case where finding a good proxy loss function for equizero is nontrivial: image classification using pretrained CNNs. Here, we find that equizero performs even worse than equitune but  $\lambda$ -equitune easily outperforms equitune and equizero. The organization of the paper is summarized below:
- §3 provides details of  $\lambda$ -equitune and equizero and proves a number of their properties.
- §4 provides overview of the applications considered in the experiments and how equivariance methods are used there.
- §5 provides experimental details and results for all the applications.
Here we discuss relevant concepts in group equivariance and group equivariant transfer learning.
Group Equivariance A group  $(G,\cdot)$  is a set  $G$  accompanied by a binary operator  $\cdot$  that satisfy the axioms of a group, namely i) closure:  $g\cdot h\in G$  for all  $g,h\in G$ ; ii) associativity:  $(g\cdot h)\cdot k = g\cdot (h\cdot k)$ ;
iii) identity: there exists  $e \in G$ , such that  $g \cdot e = e \cdot g$  for all  $g \in G$ ; iv) inverse: for every element  $g \in G$ , there exists  $g^{-1}$ , such that  $g \cdot g^{-1} = g^{-1} \cdot g = e$ . We write  $g \cdot h$  as  $gh$  for brevity.
Given a set  $\mathcal{X}$ , we define a group action of  $G$  on  $\mathcal{X}$  as  $\Gamma_{\mathcal{X}}: G \times \mathcal{X} \mapsto \mathcal{X}$  such that it satisfies two axioms, namely i) identity:  $\Gamma_{\mathcal{X}}(e, x) = x$  for all  $x \in \mathcal{X}$ , where  $e \in G$  is the identity; ii) compatibility:  $\Gamma_{\mathcal{X}}(g, \Gamma_{\mathcal{X}}(h, x)) = \Gamma_{\mathcal{X}}(gh, x)$ , for all  $g, h \in G, x \in \mathcal{X}$ . We write  $\Gamma_{\mathcal{X}}(g, x)$  simply as  $gx$  for brevity.
A model  $\mathbf{M}:\mathcal{X}\mapsto \mathcal{Y}$  is equivariant to  $G$  under the group action of  $G$  on  $\mathcal{X}$  and  $\mathcal{V}$  if  $\mathbf{M}(gx) = g\mathbf{M}(x))$  for all  $g\in G,x\in \mathcal{X}$ . This essentially means that any group transformation  $g$  to the input  $\Gamma_{\mathcal{X}}(g,x)$  should reflect with an equivalent group transformation of the output  $\Gamma_{\mathcal{Y}}(g,\mathbf{M}(x))$ .
Equivariant Finetuning Recently, Basu et al. (2023) proposed a finetuning method called equituning that starts with potentially non-equivariant model  $\mathbf{M}$  and produces a model  $\mathbf{M}_{\mathbf{G}}$  that is equivariant to  $G$ . Equituning converts a pretrained model into an equivariant version by minimizing the distance of features obtained from pretrained and equivariant models. The output of an equituned model is given by
$$
\mathbf {M} _ {\mathbf {G}} (x) = \frac {1}{| G |} \sum_ {g \in G} g ^ {- 1} \mathbf {M} (g x). \tag {1}
$$
While the averaging in equation 1 is shown to be useful for finetuning, we find it leads to poor equivariant zero-shot learning. This could be because the pretrained model outputs high quality features only for some of the transformed inputs. Hence, averaging them directly leads to low quality zero-shot performance. This can be avoided by using weighted averaging as discussed in Sec. 3.
Optimization-Based Canonical Representation On the other hand, Kaba et al. (2022) show that group equivariant output,  $\mathbf{M}_{\mathbf{G}}^{\mathrm{OPT}}(x)$ , can be obtained by optimizing a (non-equivariant) loss function  $l(\mathbf{M}(gx))$  with respect to group elements  $g \in G$  for any  $x \in \mathcal{X}$  as shown below.
$$
\mathbf {M} _ {\mathbf {G}} ^ {\mathrm {O P T}} (x) = g _ {*} ^ {- 1} \mathbf {M} \left(g _ {*} x\right), \tag {2}
$$
where  $g_{*} = \arg \min_{g\in G}l(\mathbf{M}(gx))$  with  $l:\mathcal{V}\mapsto \mathbb{R}$  being an injective proxy loss function and assuming the minima is unique. However, the purpose of this formulation in Kaba et al. (2022) is only to obtain an equivariant representation for training from scratch. Moreover, no competitive zero-shot or finetuning performance is obtained in Kaba et al. (2022) using this method. We show that the choice of  $l$  plays an important role in its zero-shot and finetuning performance, even outperforming equituning. This method is obtained as a special case of  $\lambda$ -equituning introduced next.
Additional Related Works Group equivariance plays a key role in geometric deep learning (Bronstein et al., 2021) for designing efficient domain specific neural networks. Several elegant architectures have been proposed for equivariant image classification (Cohen & Welling, 2016, 2017; Romero & Cordonnier, 2020), reinforcement learning (Mondal et al., 2020; van der Pol et al., 2020; Mondal et al., 2022; Wang et al., 2022), graph (Satorras et al., 2021; Keriven & Peyré, 2019; Gasteiger et al., 2021) and mesh (De Haan et al., 2020; He et al., 2021; Basu et al., 2022) processing, natural language processing (Gordon et al., 2019; Li et al., 2022), and data generation (Dey et al., 2021). These architectures need to be trained from scratch, which is not always desirable.
Frame averaging produces equivariant output from non-equivariant architecture backbones (Puny et al., 2021; Atzmon et al., 2022; Duval et al., 2023). Most work here focuses on finding good frames, which are equivariant subsets of groups, for specific groups, and not for general groups. And not much is known about their performance with pretrained models. Kaba et al. (2022) also give a canonicalization-based method that uses an equivariant auxiliary network for constructing equivariant networks out of non-equivariant backbones and is used for training from scratch. But this work requires an additional equivariant network and appropriate trainable parameterization of the group actions, which is presented only for certain groups of interest. Further, this work is not concerned with equivariant performance of pretrained models. Zero-shot group equivariance was also recently used by Muglich et al. (2022) for zero-shot coordination in partially observable Markov decision processes (POMDPs). In contrast, our work aims to be provide efficient equivariant transfer learning algorithms that are general in terms of considered tasks and groups, and does not require additional equivariant networks.
3  $\lambda$ -Equitune
We propose  $\lambda$ -equitune, where unequal weights are assigned to features obtained from transformed inputs. This is a simple generalization of equitune in equation 1 and the optimization-based approach in equation 2, where the goal is to assign higher values to better features. Like these previous methods,  $\lambda$ -equitune is equivariant and a universal approximator of equivariant functions.
The main idea of  $\lambda$ -equitune is that given a pretrained model  $\mathbf{M}$ , the features  $\mathbf{M}(gx)$  for any fixed  $x$  are not all equally important for all  $g \in G$ . We denote by  $\lambda(gx)$  the importance weight of feature  $\mathbf{M}(gx)$  for  $g \in G, x \in \mathcal{X}$ . We assume  $G$  is finite, just as in Basu et al. (2023). Suppose  $\lambda: \mathcal{X} \mapsto \mathbb{R}^{+}$  is known a priori, and denote the  $\lambda$ -equituned model as  $\mathbf{M}_{\mathbf{G}}^{\lambda}$ . Then we want to minimize
$$
\min  _ {\mathbf {M} _ {\mathbf {G}} ^ {\lambda} (x)} \sum_ {g \in G} \left\| \lambda (g x) \mathbf {M} (g x) - \mathbf {M} _ {\mathbf {G}} ^ {\lambda} (g, x) \right\| _ {2} ^ {2} \tag {3}
$$
s.t.  $\mathbf{M}_{\mathbf{G}}^{\lambda}(gx) = g\mathbf{M}_{\mathbf{G}}^{\lambda}(x)$  for all  $g\in G$
We call the solution to equation 3 as  $\lambda$ -equitune, given by
$$
\mathbf {M} _ {\mathbf {G}} ^ {\lambda} (x) = \frac {1}{\sum_ {g \in G} \lambda (g x)} \sum_ {g \in G} ^ {| G |} g ^ {- 1} \lambda (g x) \mathbf {M} (g x). \tag {4}
$$
3.1 Special Cases
When  $\lambda(gx) = 1$  for  $g \in G$ , then equation 4 is the same as equation 1. Further, we get equation 2 when  $\lambda$  is an indicator function  $\lambda(gx) = \mathbb{1}_{\{g = g_*\}}$ , where  $g_* = \arg \min_{g \in G} l(\mathbf{M}(gx))$  with  $l: \mathcal{V} \mapsto \mathbb{R}$  such that the minimization is well defined. We use  $\mathbf{M}_{\mathbf{G}}^0$  to denote the equizero model.
The first main contribution of our work is experimental. We show there are good loss functions for equizero for several diverse tasks that easily outperform equitune by choosing the best features.
The second main contribution is to show that  $\lambda$ -equitune outperforms equitune and is competitive with equizero even when good loss functions are known. Moreover, when the loss functions are not trivial, equizero performs even worse than equitune, but  $\lambda$ -equitune easily outperforms both.
3.2 Canonical-  $\lambda$  -Equitune
Here, we provide an extension of the  $\lambda$ -equitune algorithm of equation 4 to continuous groups. The idea is to combine the canonicalization method of Kaba et al. (2022) with  $\lambda$ -equitune leading to an expressive equivariant network with weighted averaging over features with different group actions applied to them.
Definition 1 (Canonical- $\lambda$ -equitune). Given a (continuous) group  $G$ , a non-equivariant function  $M: X \mapsto Y$ , and an equivariant auxiliary function (from the setting of Kaba et al. (2022))  $h: X \mapsto G$ , lambda functions  $\lambda: X \mapsto R^{+}$ , and a set of group elements  $\Theta = \{\theta_{1}, \ldots, \theta_{k}\}$ , i.e.  $\theta_{i} \in G$ , we define the canonical- $\lambda$ -equitune operators as
$$
M _ {G, e q u i} ^ {\lambda} (x) = \frac {1}{\sum_ {\theta \in \Theta} \lambda \left(\theta h (x) ^ {- 1} x\right)} \sum_ {\theta \in \Theta} \lambda \left(\theta h (x) ^ {- 1} x\right) h (x) M \left(\theta h (x) ^ {- 1} x\right) \tag {5}
$$
$$
M _ {G, i n v} ^ {\lambda} (x) = \frac {1}{\sum_ {\theta \in \Theta} \lambda \left(\theta h (x) ^ {- 1} x\right)} \sum_ {\theta \in \Theta} \lambda \left(\theta h (x) ^ {- 1} x\right) M \left(\theta h (x) ^ {- 1} x\right). \tag {6}
$$
In Thm. 1, we show that the canonical-  $\lambda$ -equivariant network is equivariant to the group  $G$
Theorem 1.  $M_{G, \text{equiv}}^{\lambda}(x)$  and  $M_{G, \text{inv}}^{\lambda}(x)$  are, respectively, equivariant and invariant to  $G$ .
3.3 Properties
Now we show in Thm. 2 that equation 4 is equivariant with respect to  $G$ .
Theorem 2 (Equivariance).  $\mathbf{M}_{\mathbf{G}}^{\lambda}$  defined in equation 4 is equivariant to  $G$
Table 1: Inference times of equitune, equizero, and  $\lambda$ -equitune for the c4 group for various CLIP models on CIFAR100. We use batch size 32 on a single Nvidia A100 GPU.
Model	Time (sec.)	Memory (MB)
Equitune	Equizero	λ-Equitune	Equitune	Equizero	λ-Equitune
RN50	14.15	14.10	23.5	3703	3703	3941
ViT-B/32	11.00	10.27	16.08	2589	2587	2915
We define a universal approximator in Def. 2. Then, Thm. 3 proves that  $\lambda$ -equitune is a universal approximator of equivariant functions for groups where  $\|g\| = 1$  for  $g \in G$ . This includes a wide range of groups including the permutation group, the  $SO(n)$  groups of special orthogonal groups, etc. This condition is the same as the class of groups considered in Basu et al. (2023).
Definition 2 (Universal approximator). A model  $\mathbf{M}:\mathcal{X}\mapsto \mathcal{Y}$  is a universal approximator of a continuous function  $f:\mathcal{X}\mapsto \mathcal{Y}$  if for any compact set  $\kappa \subset \mathcal{X}$  and  $\epsilon >0$ , there exists a choice of parameters for  $\mathbf{M}$  such that  $\| f(x) - \mathbf{M}(x)\| \leq \epsilon$  for all  $x\in \mathcal{K}$ .
Theorem 3 (Universality). Let  $f_{G}:\mathcal{X}\mapsto \mathcal{Y}$  be any continuous function equivariant to group  $G$  and let  $\lambda :\mathcal{X}\mapsto \mathbb{R}^{+}$  be any positive scalar function. And let  $\mathbf{M}:\mathcal{X}\mapsto \mathcal{Y}$  be a universal approximator of  $\frac{f_G}{\lambda}$ . Here  $\mathcal{X}$  is such that if  $x\in \mathcal{X}$ , then  $gx\in \mathcal{X}$  to ensure the equivariance of  $f_{G}$  is well-defined. Then,  $\mathbf{M}_{\mathbf{G}}^{\lambda}$  is a universal approximator of  $f_{G}$ .
Computational Complexity Note that equitune, equizero, and  $\lambda$ -equitune have the same compute complexity. In practice, in Tab. 1 we find that equitune and equizero have exactly the same time and memory consumption, whereas  $\lambda$ -equitune takes a little more time and memory because of the additional  $\lambda$  network. We illustrate on RN50 and ViT-B/32 models of CLIP using the same text encoder, but different image encoders. RN50 uses a Resnet50 based image encoder, whereas ViT-B/32 uses a vision transformer based image encoder.
Beyond Zero-Shot Learning Let us emphasize that even though the equizero model in equation 2 is not differentiable due to the arg max, we can still use simple gradient estimators known in the literature. One popular estimator is the straight-through estimator (Bengio et al., 2013), where the equizero output in equation 2 would be written as  $\mathbf{M}_{\mathbf{G}}^{0}(x) = \mathbf{M}(x) + (\mathbf{M}_{\mathbf{G}}^{0}(x) - \mathbf{M}(x))$ . detach(), where detach() indicates that no gradient flows through the term  $(\mathbf{M}_{\mathbf{G}}^{0}(x) - \mathbf{M}(x))$ . In practice, we found it to be slightly better to use  $\mathbf{M}_{\mathbf{G}}(x)$  instead of  $\mathbf{M}(x)$  and write  $\mathbf{M}_{\mathbf{G}}^{0}(x) = \bar{\mathbf{M}}_{\mathbf{G}}(x) + (\mathbf{M}_{\mathbf{G}}^{0}(x) - \mathbf{M}_{\mathbf{G}}(x))$ . detach(). §5.1.3 illustrates few-shot learning and finetuning using equizero and compares with equituning.
First we present several applications in §4.1 where finding a loss function for equizero is easy. This naturally leads to excellent equivariant zero-shot results outperforming equitune. Then, in §4.2 we provide two applications in image classification to show 1) the benefits and drawbacks of equizero compared to equitune and  $\lambda$ -equitune, and 2) that  $\lambda$ -equitune consistently performs well avoiding the drawbacks of equizero.
4.1 Equizero Applications
Here we provide applications where equizero achieves excellent zero-shot performance, namely: 1) deep Q-learning, 2) fairness in NLG, and 3) compositional generalization in languages.
Equizero Reinforcement Learning Recent works, such as van der Pol et al. (2020); Mondal et al. (2020), have developed RL algorithms that leverage symmetries in the environments that helps improve robustness and sample efficiency. But no existing work efficiently uses group equivariance on top of pretrained RL models. Here we apply equizero and equitune on top of pretrained models inspired from the group symmetries found in van der Pol et al. (2020). We find that equitune outperforms non-equivariant pretrained models but equizero outperforms both. We simply use the  $Q$ -values as the proxy loss function as described in §C.1 with more details.
Group-Theoretic Fairness in NLG We seek to reduce the social biases inherent in language models (LMs), focusing on GPT2 (Radford et al., 2019). We consider the group-theoretic fairness setup of Sheng et al. (2019) and Basu et al. (2023). We take the sets of demographics, namely [‘man’, ‘woman’], [‘straight’, ‘gay’], and [‘black’, ‘white’]. For each demographic group, Sheng et al. (2019) proposed two tasks, called respect task and occupation task, where each task consists of five context phrases. A language model (LM) is given these contexts to generate sentences. These generated sentences are then classified as ‘positive’, ‘negative’, ‘neutral’, or ‘other’ by a regard classifier also proposed by Sheng et al. (2019). These outputs are called regard scores. A regard classifier is a BERT-based model similar to a sentiment classifier but more specific for fairness tasks. We use equizero using the regard scores as the proxy loss function, so as to maximize positivity in the generated texts, while guaranteeing group-theoretic fairness. To that end, we propose two algorithms that help in reducing existing biases across demographic groups.
EquizeroLM and R-EquizeroLM: Basu et al. (2023) define EquiLM and R-EquiLM, that use a sequential version of equitune to perform group transformed averaging to achieve fairness across demographic groups. While EquiLM and R-EquiLM generate debiased outputs, they do not produce positive regard scores, which is desirable to reduce toxicity in generated text. We propose EquizeroLM and R-EquizeroLM which use equizero to maximize regard scores, ensuring both fair and less toxic. Further details provided in §C.2.
Zero-Shot Compositional Generalization We show compositional generalization capabilities of equizero on the SCAN dataset (Lake & Baroni, 2018). SCAN measures compositionality using a language-to-action translation task. E.g., if the model learns that the phrase "Jump", "Run", "Run Twice" translate to the actions "JUMP", "RUN", "RUN RUN" from the train set, then, SCAN tests whether the model also learns that "Jump Twice" translates to "JUMP JUMP". Such a reasoning however common in human beings is hard to find in language models.
We apply equizero on two tasks in SCAN, Add Jump and Around Right. Gordon et al. (2019) solved these tasks by constructing sophisticated group equivariant networks from scratch and training them. Basu et al. (2023) used the same group actions as Gordon et al. (2019), but used equituning on pretrained non-equivariant models for a few iterations and obtained comparable results. But, as we note, equitune has poor zero-shot performance. We show that equizero using negative of maximum probability from the output as the loss function gives much better zero-shot performance. Using gradient estimators described in §3.3, we also compare the finetuning performance of equizero against equitune. Details of group actions used are given in §C.3
4.2  $\lambda$ -Equitune Applications
Here we consider two important applications: CLIP-based and CNN-based image classification. For CLIP, it is easy to find a loss function for equizero that provides better results than equitune. But for the simple case of CNN-based classification it is non-trivial to find such a loss function. Since  $\lambda$ -equitune does not require a loss function, it performs better than equitune in both cases. Equizero only performs well for CLIP, but fails miserably for CNN-based classification.
CLIP-Based Image Classification CLIP is a pretrained model consisting of image and text encoders that give impressive zero-shot classification performance across variety of unseen datasets. But, in Fig. 4a and 7, we find that CLIP is not robust to simple transformations such as rotation by multiples of  $90^{\circ}$  or random flips. This trend is seen across different image encoders like RN50, RN101 (Radford et al., 2021), ViT-B/32 and ViT-B/16 (Dosovitskiy et al., 2021). This can be avoided by making the model in/equi-variant to such transformations, e.g., by using equitune. But we show in §5.2.1 that equitune does not produce good zero-shot performance. We show in §5.2.1 that using equizero with image-text similarity score as loss function provides much better zero-shot results than equitune. Later, in §5.2.1, we show that  $\lambda$ -equitune achieves better zero-shot performance than equitune without any loss function. Finally, when finetuned, we find that  $\lambda$ -equitune tends to perform the best, possibly because it does not need gradient estimators like equizero, and because it uses weighted averaging to obtain better features than equitune.