SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jannik Kossen<sup>1</sup>  
Mark Collier<sup>2</sup>  
Basil Mustafa<sup>3</sup>  
Xiao Wang<sup>3</sup>  
Xiaohua Zhai<sup>3</sup>
Lucas Beyer $^{3}$  Andreas Steiner $^{3}$  Jesse Berent $^{2}$  Rodolphe Jenatton $^{3}$  Efi Kokiopoulou $^{2}$
$^{1}$  OATML, Department of Computer Science, University of Oxford
2 Google Research 3 Google DeepMind
We introduce Three Towers (3T), a flexible method to improve the contrastive learning of vision-language models by incorporating pretrained image classifiers. While contrastive models are usually trained from scratch, LiT [85] has recently shown performance gains from using pretrained classifier embeddings. However, LiT directly replaces the image tower with the frozen embeddings, excluding any potential benefits from training the image tower contrastively. With 3T, we propose a more flexible strategy that allows the image tower to benefit from both pretrained embeddings and contrastive training. To achieve this, we introduce a third tower that contains the frozen pretrained embeddings, and we encourage alignment between this third tower and the main image-text towers. Empirically, 3T consistently improves over LiT and the CLIP-style from-scratch baseline for retrieval tasks. For classification, 3T reliably improves over the from-scratch baseline, and while it underperforms relative to LiT for JFT-pretrained models, it outperforms LiT for ImageNet-21k and Places365 pretraining.
Approaches such as CLIP [58] and ALIGN [34] have popularized the contrastive learning of aligned image and text representations from large scale web-scraped datasets of image-caption pairs. Compared to image-only contrastive learning, e.g. [51, 9, 25], the bi-modal image-text objective allows these approaches to perform tasks that require language understanding, such as retrieval or zero-shot classification [42, 58, 34]. Compared to traditional transfer learning from supervised image representations [52, 68, 44, 38], contrastive approaches can forego expensive labelling and instead collect much larger datasets via inexpensive web-scraping [58, 11, 64]. A growing body of work seeks to improve upon various aspects of contrastive vision-language modelling, cf. related work in §5.
CLIP and ALIGN train the image and text towers from randomly initialized weights, i.e. 'from scratch'. However, strong pretrained models for either image or text inputs are often readily available, and one may benefit from their use in contrastive learning. Recently, Zhai et al. [85] have shown that pretrained classifiers can be used to improve downstream task performance. They propose LiT, short for 'locked-image text tuning', which is a variation of the standard CLIP/ALIGN setup that uses frozen embeddings from a pretrained classifier as the image tower. In other words, the text tower in LiT is contrastively trained from scratch to match locked and pretrained embeddings in the image tower. Incorporating knowledge from pretrained models into contrastive learning is an important research direction, and LiT provides a simple and effective recipe for doing so.
However, a concern with LiT is that it may be overly reliant on the pretrained model, completely missing out on any potential benefits the image tower might get from contrastive training. Zhai et al.
$\nabla$  Equal contribution.  $\triangle$  Work done while interning at Google.  $\square$  Equal advising.
Correspondence to jannik.kossen@cs.ox.ac.uk and markcollier@google.com.
Figure 1: CLIP and ALIGN do not make use of pretrained models, and LiT directly uses a frozen pretrained model as the image tower. With Three Towers (3T), we propose a flexible strategy to improve contrastive learning with pretrained models: in addition to a pair of CLIP-style from-scratch image and text towers, we introduce a third tower which contains fixed pretrained image embeddings; extra loss terms align the main towers to the third tower. Unlike for CLIP/ALIGN and LiT, the image tower can benefit from both contrastive learning and pretrained classifier embeddings.
[85] themselves give one example where LiT performs worse than standard contrastive training: when using models pretrained on Places365 [87]—a dataset relating images to the place they were taken—the fixed embeddings do not generalize to downstream tasks such as ImageNet-1k (IN-1k) [40, 61] or CIFAR-100 [40]. For their main results, Zhai et al. [83] therefore use models pretrained on datasets such as ImageNet-21k (IN-21k) [14, 61] and JFT [68, 84] which cover a variety of classes and inputs. However, even then, we believe that constraining the image tower to fixed classifier embeddings is not ideal: later, we will show examples where LiT performs worse than standard contrastive learning due to labels or input examples not covered by IN-21k, cf. §4.2. Given the scale and variety of contrastive learning datasets, we believe it should be possible to improve the image tower by making use of both pretrained models and contrastive training.
In this work, we propose Three Towers (3T): a flexible approach that improves the contrastive learning of vision-language models by effectively transferring knowledge from pretrained classifiers. Instead of locking the main image tower, we introduce a third tower that contains the embeddings of a frozen pretrained model. The main image and text towers are trained from scratch and aligned to the third tower with additional contrastive loss terms (cf. Fig. 1). Only the main two towers are used for downstream task applications such that no additional inference costs are incurred compared to LiT or a CLIP/ALIGN baseline. This simple approach allows us to explicitly trade off the
Table 1: For retrieval, 3T improves on LiT and the CLIP-style baseline (top-1 recall  $\uparrow$ ). Models are g scale, using Text-Filtered WebLI, and JFT pretraining for LiT/3T, cf. §4.1.
Method	Basel.	LiT	3T
Flickr img2txt	85.0	83.9	87.3
Flickr txt2img	67.0	66.5	72.1
COCO img2txt	60.0	59.5	64.1
COCO txt2img	44.7	43.6	48.5
main contrastive learning objective against the transfer of prior knowledge from the pretrained model. Compared to LiT, the image tower in 3T can benefit from both contrastive training and the pretrained model. We highlight the following methodological and empirical contributions:
- We propose and formalize the 3T method for flexible and effective transfer of pretrained classifiers into contrastive vision-language models (§3).
- 3T consistently improves over LiT and a from-scratch baseline for retrieval tasks (e.g. Table 1, §4.1).
- For classification tasks, 3T outperforms LiT and the baseline with IN-21k pretrained models; for JFT pretraining, 3T outperforms the baseline but not LiT (§4.2).
- We extend the evaluation of Zhai et al. [85] to additional tasks and pretraining datasets, showing that 3T is significantly more robust than LiT to deficits in the pretrained model ( $\S 4.2$  and  $\S 4.4$ ).
- We show that 3T benefits more from model size or training budget increases than LiT (§4.3).
- We introduce a simple post-hoc method that allows us to further improve performance by combining 3T- and LiT-like prediction (§4.5).
2 Background: Contrastive Learning of Vision-Language Models
Before introducing 3T, we recap contrastive learning of vision-language models as popularized by [58, 34, 86] and give a more formal introduction to LiT [85].
CLIP/ALIGN. We assume two parameterized models: an image tower  $f = f_{\theta}$  with parameters  $\theta$  and a text tower  $g = g_{\phi}$  with parameters  $\phi$ . Each input sample  $(I_i, T_i)$  consists of a pair of matching image  $I_i \in \mathcal{I}$  and text  $T_i \in \mathcal{T}$ , and the contrastive loss is computed over a batch  $i \in \{1, \dots, N\}$
of examples. The towers map the input modalities to a common  $D$ -dimensional embedding space,  $f:\mathcal{I}\to \mathbb{R}^{D}$  and  $g:\mathcal{T}\rightarrow \mathbb{R}^{D}$ . We further assume that  $f$  and  $g$  produce embeddings that are normalized with respect to their L2 norm,  $\| f(I)\| _2 = \| g(T)\| _2 = 1$  for any  $I\in \mathcal{I}$  and  $T\in \mathcal{T}$ . For a batch of input samples, the bi-directional contrastive loss [66, 77, 51, 9, 86] is computed as
$$
\mathcal {L} _ {f \leftrightarrow g} = \frac {1}{2} \left(\mathcal {L} _ {f \rightarrow g} + \mathcal {L} _ {g \rightarrow f}\right), \text {w h e r e} \tag {1}
$$
$$
\mathcal {L} _ {f \rightarrow g} = - \frac {1}{N} \sum_ {i = 1} ^ {N} \log \frac {\exp \left(f \left(I _ {i}\right) ^ {\top} g \left(T _ {i}\right) / \tau\right)}{\sum_ {j = 1} ^ {N} \exp \left(f \left(I _ {i}\right) ^ {\top} g \left(T _ {j}\right) / \tau\right)}, \tag {2}
$$
$$
\mathcal {L} _ {g \rightarrow f} = - \frac {1}{N} \sum_ {i = 1} ^ {N} \log \frac {\exp \left(f \left(I _ {i}\right) ^ {\top} g \left(T _ {i}\right) / \tau\right)}{\sum_ {j = 1} ^ {N} \exp \left(f \left(I _ {j}\right) ^ {\top} g \left(T _ {i}\right) / \tau\right)}. \tag {3}
$$
Here,  $\tau$  is a learned temperature parameter and  $f(I)^\top g(T) \in \mathbb{R}$  are dot products. The two directional loss terms,  $\mathcal{L}_{f \to g}$  and  $\mathcal{L}_{g \to f}$ , have a natural interpretation as standard cross-entropy objectives for classifying the correct matches in each batch. The parameters  $\theta$  and  $\phi$  of the two towers,  $f_\theta$  and  $g_\phi$ , are jointly updated with standard stochastic optimization based on Eq. (1).
Downstream Tasks. After training,  $f$  and  $g$  are treated as fixed representation extractors. For retrieval, the dot product  $f(I)^{\top}g(T) \in \mathbb{R}$  ranks similarity between inputs. For few-shot image classification, a linear classifier is trained atop the feature representations of  $f$  from few examples;  $g$  is not used. For zero-shot image classification,  $f$  embeds images and  $g$  all possible class labels (see [58]). For each image, one predicts the label with the largest dot product similarity in embedding space.
LiT. Zhai et al. [85] initialize the parameters  $\theta$  of the image tower from a pretrained classifier and then keep them frozen them during training. That is, only the parameters  $\phi$  of the text tower are optimized during contrastive training. As the image tower  $f_{\theta}$ , LiT uses the pre-softmax embeddings of large scale classifiers, such as vision transformers [17] trained on JFT-3B [84] or IN-21k. During contrastive training, the text tower is trained from scratch using the same objective Eq. (1).
Experimentally, Zhai et al. [85] investigate all combinations for 'training from scratch', locking and finetuning a pretrained model for both towers on a custom union of the CLIP-subset of YFCC-100M [69] and CC12M [7] (cf. Fig. 3 in [85]). For the image tower, locking gives a significant lead on IN-1k over finetuning and training from scratch, and performs similarly to finetuning and better than training from scratch for retrieval. For the text tower, a locked configuration performs badly, and finetuning gives small to negligible gains over training from scratch. Given these results, Zhai et al. [85] choose the 'locked image tower and from-scratch text tower' setup that they call LiT. At large scale, they show that a locked image tower outperforms from-scratch training and finetuning on zero-shot IN-1k, ImageNet-v2 (IN-v2) [59], CIFAR-100, and Oxford-IIIT Pet [53] classification tasks. They further show LiT outperforms CLIP/ALIGN on IN-R [31], IN-A [32], and ObjectNet [3].
While Zhai et al. [85] show strong classification performance with LiT on a wide range of datasets, locking the image tower is a drastic measure that introduces a severe dependency on the pretrained model, prohibiting the image tower from improving during contrastive training. We will later show that, if the embeddings in the frozen image tower are not suited to a particular downstream tasks, LiT underperforms compared to approaches that train the image tower on the varied contrastive learning dataset, see, for example, §4.2 and §4.4. The 3T approach seeks to address these concerns.
3 Three Towers: Flexible Contrastive Learning with Pretrained Models
With Three Towers (3T), we propose a simple and flexible approach to incorporate knowledge from pretrained models into contrastive learning. Instead of directly using the pretrained model locked as the main image tower, we instead add a third tower,  $h$ , which contains the fixed pretrained embeddings. The main image and text towers are trained from scratch, and we transfer representations from the third tower to the main towers with additional contrastive losses. In this setup, the main image tower benefits from both pretraining knowledge and contrastive learning.
More formally, in addition to the standard image and text towers,  $f_{\theta}$  and  $g_{\phi}$ , cf. §2, we now have access to fixed pretrained image embeddings  $p: \mathcal{I} \to \mathbb{R}^P$ . Because  $P$  can be different from the target dimension  $D$ , we define the third tower as  $h(I) = \text{linear}(p(I))$ , where linear:  $\mathbb{R}^P \to \mathbb{R}^D$  projects
Figure 2: Details of the 3T approach. (a) Linear adaptor heads (gray) align the representations between the main towers and the third tower. (b) For downstream tasks, 3T is used in the same way as CLIP/ALIGN and LiT. We discard the third tower, using only the main towers,  $f(I)$  and  $g(T)$ .
embeddings to the desired dimensionality. In principle, the 3T architecture is also compatible with pretrained text models. However, like Zhai et al. [85], we do not observe benefits from using pretrained text models, cf. §A.3, and so our exposition and evaluation focuses on pretrained image classifiers.
When computing loss terms involving the third tower, we make use of learned linear projection heads. These heads afford the model a degree of flexibility when aligning representations between towers. First, we define  $\mathrm{NL}(\cdot) = \mathrm{norm}(\mathrm{linear}(\cdot))$ , where  $\mathrm{norm}(x) = x / \|x\|_2$  normalizes with respect to L2 norm and, overloading notation, linear:  $\mathbb{R}^D \to \mathbb{R}^D$  now preserves dimensionality. We adapt the main image and text towers as  $f_h(I) = \mathrm{NL}(f(I))$  and  $g_h(T) = \mathrm{NL}(g(T))$ . We project the third tower embeddings  $h$  to  $h_f(I) = \mathrm{NL}(h(I))$  and  $h_g(I) = \mathrm{NL}(h(I))$  for computation of the loss with the image and text towers respectively. The linear layers introduced for  $f_h, g_h, h_f$ , and  $h_g$  are independently learned from scratch. Per input batch, the 3T approach then optimizes the following loss objective:
$$
\mathcal {L} _ {\mathrm {3 T}} = \frac {1}{3} \cdot \left(\mathcal {L} _ {f \leftrightarrow g} + \mathcal {L} _ {f _ {h} \leftrightarrow h _ {f}} + \mathcal {L} _ {g _ {h} \leftrightarrow h _ {g}}\right). \tag {4}
$$
Here,  $\mathcal{L}_{f\leftrightarrow g}$  is the original contrastive loss, cf. Eq. (1), and  $\mathcal{L}_{fh\leftrightarrow h_f}$  and  $\mathcal{L}_{g_h\leftrightarrow h_g}$  are additional contrastive losses between the image/text tower and the third tower projections. All loss terms share a global temperature  $\tau$ . We train both towers,  $f$  and  $g$ , and all linear layers from scratch by optimizing Eq. (4) over input batches. Figure 2 (a) visualizes the adaptor heads and loss computation for 3T.
After training, the third tower is discarded and we use only the main two towers, cf. Fig. 2 (b). Therefore, the inference cost of 3T is equal to the baseline methods. For training, the additional cost over the from-scratch CLIP/ALIGN baseline is negligible, as frozen embeddings from the third tower can be pre-computed and then stored with the dataset, as also done in Zhai et al. [83].
Intuitions for the 3T Architecture. Intuitively, the additional losses align the representations of the main towers to the pretrained embeddings in the third tower. In fact, Tian et al. [71] show that contrastive losses can be seen as distillation objectives that align representations between a teacher and a student model. They demonstrate that contrastive losses are highly effective for representation transfer, outperforming alternative methods of distillation. Thus, the additional terms,  $\mathcal{L}_{f_h\leftrightarrow h_f}$  and  $\mathcal{L}_{g_h\leftrightarrow h_g}$ , transfer representations from the pretrained model to the unlocked main towers, albeit without the usual capacity bottleneck between the student and teacher models. Of course, for 3T, we also need to consider the original objective  $\mathcal{L}_{f\leftrightarrow g}$  between the unlocked text and image towers. In sum, the unlocked towers benefit both from the pretrained model and contrastive training.
Design Choices. In §4.6, we ablate various design choices, such as our use of equal weights among the terms of Eq. (4), the contrastive loss for representation transfer, the shared global temperature  $\tau$ , the linear layers for  $h$  in the third tower, as well as our design of the adaptor heads.
Discussion. The 3T approach includes pretrained knowledge without suffering from the inflexibility of directly using the pretrained model as the main image tower. For example, unlike LiT, 3T allows for architectural differences between the unlocked image tower and pretrained model. Further, it seems plausible that 3T should generally be more robust than LiT: as the image tower learns from the highly-diverse contrastive learning datasets, the chances of encountering 'blindspots' in downstream applications, e.g. due to labels or examples not included in the pretraining dataset, should be lower. In a similar vein, LiT is most appropriate for pretrained models so capable that they need not adapt during contrastive training. For example, few-shot classification performance, which uses only the image tower, by design cannot improve at all during contrastive training with LiT. On the other hand, with 3T
we may be able to successfully incorporate knowledge from weaker models, too. Lastly, finetuning-- instead of locking--the main image tower from a pretrained classifier often has at most marginal positive effects over training from scratch after a high number of training steps, cf. §2. In contrast, the additional terms in Eq. (4) consistently ensure the main towers align with the pretrained model.
In this section, we compare 3T to LiT and to a standard CLIP/ALIGN baseline trained from scratch, which we will refer to as the 'baseline' for simplicity. Our experimental setup largely follows Zhai et al. [85]: for all methods, we use Vision Transformers (ViT) [17] for the image and text towers, replacing visual patching with SentencePiece encoding [41] for text inputs, further sharing optimization and implementation details with [85]. We rely on the recently proposed WebLI dataset [11], a large-scale dataset of 10B image-caption pairs (Unfiltered WebLI). We also explore two higher-quality subsets derived from WebLI: Text-Filtered WebLI, which uses text-based filters following Jia et al. [34], and Pair-Filtered WebLI (see §A.7), which retains about half of the examples with the highest image-text pair similarity. For image tower pretraining, we consider both proprietary JFT-3B [84] and the publicly available IN-21k checkpoints of Dosovitskiy et al. [17]. For IN-21k experiments, our largest model scale is L, with a  $16 \times 16$  patch size for ViT, and for JFT pretraining we go up to g scale at a patch size of  $14 \times 14$ . Unless otherwise stated, we train for 5B examples seen at a batch size of 14 336.
4.1 Retrieval
We study zero-shot retrieval performance on COCO [10] and Flickr [56]. Table 1 shows results for g scale models trained on Text-Filtered WebLI, with JFT pretraining for 3T and LiT. In Table 2, we report performance of L scale models trained on Unfiltered WebLI for JFT and IN-21k pretraining. We give results for additional WebLI splits for IN-21k and JFT pretraining in Tables A.4 and A.6.
Table 2: 3T outperforms LiT and the baseline for retrieval (top-1 recall ↑, L scale models, Unfiltered WebLI, see §4.1).
Pretraining Method	- Basel.	IN-21k	JFT
LiT	3T	LiT	3T
Flickr* img2txt	75.6	71.7	80.0	78.7	80.0
Flickr* txt2img	57.1	49.3	60.9	58.8	61.4
COCO img2txt	51.0	46.1	54.4	52.7	54.4
COCO txt2img	34.2	27.8	37.7	36.7	37.9
3T improves on LiT and the baseline for retrieval tasks across scales, datasets, and for both JFT and IN-21k pretraining. A rare exception to this are the Unfiltered WebLI results in Table A.6, where 3T beats LiT for retrieval on average and for txt2img, but not for img2txt. In general however, LiT underperforms for retrieval and regularly does not outperform the baseline: at L scale, LiT shows a strong dependence on the pretrained model and can only outperform the baseline with JFT pretraining. In contrast, 3T obtains similar improvements over the baseline for both pretraining datasets. We will continue to see this pattern in our experiments: 3T consistently improves over the baseline, while LiT results can vary wildly and depend strongly on the pretraining dataset. We discuss our retrieval results in the context of SOTA performance in §A.7: the SOTA method CoCA [81] achieves better results but uses about 4 times more compute; increasing the compute budget for 3T would likely reduce the gap.
Intuitively, while the fixed classifier embeddings in LiT can categorize inputs into tens of thousands of labels, they may not be fine-grained enough for retrieval applications. On the other hand, the contrastive training of the baseline is closely related to the retrieval task but misses out on knowledge from pretrained models. Only 3T is able to combine benefits from both for improved performance.
4.2 Few-Shot and Zero-Shot Image Classification
Next, we compare the approaches on few- and zero-shot image classification. See §D for citations of all the datasets that we use. For zero-shot classification, we follow the procedure described in §2. For few-shot tasks, we report 10-shot accuracy, more specifically, the accuracy of a linear classifier trained on top of fixed image representations, averaging over 3 seeds for the 10 random examples per class. As few-shot performance depends only on the image embeddings, LiT's few-shot accuracy is precisely the same as that of the pretrained model. Despite this, Zhai et al. [85] show that LiT outperforms the unlocked baseline on few-shot IN-1k and CIFAR-100 evaluations.
*For historical reasons, Flickr* results do not use the 'Karpathy' split [36]. However, they are available for all runs, and they are directly comparable. Our g scale runs do have Karpathy split results, denoted Flickr (no star).
For IN-21k pretraining, Table 3 reports the performance of L scale models trained on Unfiltered WebLI, and we give results on additional datasets in Table A.4. Here, 3T outperforms both LiT and the baseline. For JFT pretraining, Table A.5 gives results at L scale on Unfiltered WebLI and Table A.6 presents results at g scale across WebLI splits. In all JFT settings, LiT performs best on average for image classification tasks, despite few-shot performance being fixed for LiT. However, for both JFT and IN-21k pretraining, 3T improves over the baseline for almost all tasks. This is different for LiT, where performance heavily depends on the pretraining data.
Risk of Locking. There is a risk associated with LiT, both in a positive and negative sense. Using fixed classifier representations can result in excellent performance if the downstream task distribution and pretraining dataset are well-aligned: for example, IN-21k contains hundreds of labels of bird species, and IN-21k-LiT performs well on the Birds task, outperforming 3T by  $18\%$ . However, the IN-21k label set does not contain a single car brand, and thus, IN-21k-LiT does not perform well on Stanford Cars, underperforming relative to the baseline and 3T by almost  $40\%$ . ObjectNet, IN-A, and IN-R were created to be challenging for ImageNet models, and so IN-21k-LiT performs worse than the baseline and 3T here, too. For example, IN-R contains artistic renditions of objects that are challenging for IN-21k-based models as they have mostly been trained on realistic images. IN-21k-LiT also struggles with more specialized tasks, performing  $29\%$  worse than the baseline on the remote sensing RESISC dataset.