SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Test-Time Distribution Normalization for Contrastively Learned Vision-language Models
Yifei Zhou*  
University of California, Berkeley  
yifei_zhou@berkeley.edu
Juntao Ren* Cornell University j1r429@cornell. edu
Fengyu Li* Cornell University f1334@cornell.edu
Ramin Zabih Cornell University rdz@cs.cornell.edu
Ser-Nam Lim  
University of Central Florida  
sernam@ucf.edu
Advances in the field of vision-language contrastive learning have made it possible for many downstream applications to be carried out efficiently and accurately by simply taking the dot product between image and text representations. One of the most representative approaches proposed recently known as CLIP [50] has garnered widespread adoption due to its effectiveness. CLIP is trained with an InfoNCE loss that takes into account both positive and negative samples to help learn a much more robust representation space. This paper reveals that the common downstream practice of taking a dot product is only a zeroth-order approximation of the optimization goal, resulting in a loss of information during test-time. Intuitively, since the model has been optimized based on the InfoNCE loss, test-time procedures should also be in alignment. The question lies in how one can retrieve any semblance of negative samples information during inference in a computationally efficient way. To this end, we propose Distribution Normalization (DN), where we approximate the mean representation of a batch of test samples and use such a mean to represent what would be analogous to negative samples in the InfoNCE loss. DN requires no retraining or fine-tuning and can be effortlessly applied during inference. Extensive experiments on a wide variety of downstream tasks exhibit a clear advantage of DN over the dot product on top of other existing test-time augmentation methods. Our code is available at https://github.com/fengyuli2002/distribution-normalization.
Contrastive Learning [20, 3, 60, 6, 52] has become an important area of research in recent years due to its widespread success in self-supervised representation learning. The key idea behind contrastive learning is to group positive samples that share similar semantics (e.g., images resulting from applying different augmentations to the same image), while separating semantically different negative samples (e.g., different images with different data augmentations). On the heels of this progress, contrastive learning has also been shown to be very effective for creating a joint representation space between images and text [50, 46, 75, 10]. The CLIP approach proposed in [50] is at the forefront of this research effort and has been shown to be the state-of-the-art in multiple downstream tasks. The model is generally pre-trained on a large amount of image-text pairs that form the positive samples, while the negative samples are composed from random pairing of images and text. After the model is trained, similarities between text and images can be measured by simply taking a dot product between their corresponding representations. Such a paradigm has opened the doors to state-of-the-art performance
in many cross-modal alignment tasks such as zero-shot classification [73, 21], cross-modal retrieval [71, 40, 35, 32], and evaluation of machine-generated captions [17, 22, 30].
Figure 1: An illustration of how Distribution Normalization better aligns with infoNCE pre-training objective.  $\mu_{img}$  and  $\mu_{txt}$  are mean encoded representations for images and text respectively.
Despite CLIP's successes, we observe a "mismatch" between the pre-training objective of such cross-modal representation models and their downstream usage. Most if not all of these models are trained taking into account of both positive and negative samples, of which one of the most popular optimization objectives is the InfoNCE loss [58, 62, 74]. The InfoNCE loss considers negative samples from the data distribution per positive sample and has been shown to be very effective for cross-modal representation learning. In contrast, during test-time, the similarity is often measured using the dot product between the image and the text representations, which does not utilize any information with regard to the data distribution.
This paper presents our findings that rectifying such misalignment can boost performance consistently across a variety of downstream tasks. However, naively applying InfoNCE loss for downstream tasks requires iterating over all negative samples for every test sample, which is not a tractable operation. Instead, we found that a first-order approximation of the InfoNCE loss, consisting of simply subtracting the distribution mean in the representation space before taking the dot product, is able to achieve a similar effect. We call this proposed approach Distribution Normalization (DN) - DN is very easy to implement and does not require any retraining, fine-tuning or labeled data.
To investigate how this simple modification affects the performance of CLIP and its state-of-the-art variants in practice, we conduct extensive experiments on a wide variety of cross-modal alignment tasks, including zero-shot classification, cross-modal retrieval, and evaluation of image captions. In all the settings that we have studied, DN displays a clear benefit over direct dot product and can be combined with other existing test-time augmentation methods [55], which we will show in our experiments. Our contributions are: (1) We present an insightful analysis of why the dot product as a similarity measure is only a zeroth-order approximation; (2) We propose DN to overcome this shortcoming by utilizing an approximation of the test-time distribution, and also provide guidance on how DN can be conveniently implemented in practice; (3) We provide extensive supporting empirical evidence and ablation studies that demonstrate the fidelity of DN for practical use.
2 Background & Related Works
Learning via Natural Language Supervision Several recent studies [7, 12, 34, 35] leverage large-scale datasets of image-text pairs to supervise vision-language pre-training (VLP) with a contrastive objective. CLIP [50] is one such example that is contrastively trained to create a joint vision-language representation space, which has produced empirical success in downstream tasks such as zero-shot classification and cross-modal retrieval over previous state-of-the-art models [40, 65, 68, 33]. Since CLIP's introduction, several other works have utilized a similar contrastive objective to create good vision-language representations and achieve state-of-the-art performance on many downstream tasks involving datasets that the model was not pre-trained on [33, 21, 7, 11, 26]. For example, ALIGN scales joint pre-training to minimally-filtered datasets [21], while ALBEF [33] and TCL [68] focus on learning more semantically meaningful multi-modal interactions. It is noteworthy that these SOTA
models all leverage the InfoNCE loss as their key training objective, and apply a simple dot product in downstream tasks.
CLIP Augmentation/Adaptation for Downstream Applications Researchers have explored numerous methods to adapt CLIP for downstream tasks, capitalizing on its remarkable transferability. CoOp [75, 76] employs learnable textual prompts, inspired by prompt learning [36], for few-shot training data, while VT-CLIP [49] incorporates visual-guided texts for enhanced vision-language alignment. However, these adaptations demand extra training data and fine-tuning resources, undermining CLIP's core strength of efficient zero-shot recognition. To address this issue, CALIP [14] presents a training-free attention module to boost zero-shot image classification, and SuS-X [61] offers training-free name-only transfer by adapting CLIP with a curated support set. However, both CALIP and Sus-X are tailored for zeroshot classification only, while in this paper, we propose DN, a training-free adaptation that is conceptually simpler, easier to implement, and applicable to a wide range of downstream tasks.
Test-time Augmentation/Adaptation Test-time adaptation and test-time augmentation strategies share a common goal to adapt to out-of-distribution scenarios [53, 43] during test time. Many test-time adaptation techniques make adjustments to the batch normalization layers to normalize feature statistics [37, 41], while test-time augmentation techniques [64, 25, 42, 29, 15, 19, 57] aim to ameliorate performance by taking some aggregation [55] of the model's prediction over a batch of augmented images. Though DN is also applied during test-time, DN's task is to align the mismatch between training and inference objectives and not to deal with out of distribution test samples. This renders DN orthogonal to TTA. In our experiments, we use a representative TTA model from Shanmugam et al. [55] and show that while TTA improves model performance, adding DN to TTA gives an additional increase.
This section is organized as follows: in Section 3.2, we show that a dot product is only a zeroth-order approximation of the pre-training objective and how a better approximation arises naturally from the InfoNCE objective, and in Section 3.3, we explain how this idea can be efficiently implemented in practice.
In our analysis,  $\mathcal{D}_S$  is the training distribution and  $\mathcal{D}_T$  is the test distribution,  $x_0, y_0$  are images and text belonging to the same pair,  $x_1, \ldots, x_n$  and  $y_1, \ldots, y_n$  are random images and text in the same batch,  $\phi, \psi$  are image and text encoders respectively, and  $\tau$  is a temperature hyperparameter used in InfoNCE loss while training. Since in most applications, only relative comparisons of the distances matter (e.g., retrieving the best candidate caption describing the image), we consider two similarity measures  $S_1$  and  $S_2$  to be equivalent if one is a monotone function of the other, namely
$$
\forall x _ {1}, y _ {1}, x _ {2}, y _ {2}, S _ {1} \left(x _ {1}, y _ {1}\right) \geq S _ {1} \left(x _ {2}, y _ {2}\right) \Leftrightarrow S _ {2} \left(x _ {1}, y _ {1}\right) \geq S _ {2} \left(x _ {2}, y _ {2}\right). \tag {1}
$$
The goal here is to find a test-time similarity measure that is similar to the train-time optimization goal through careful simplification of InfoNCE loss.
3.2 Analysis of InfoNCE Loss
We start from the original form of InfoNCE loss [58, 62, 74]:
$$
\mathcal {L} _ {N C E} (\mathcal {D} _ {S}) = \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {S}} [ \mathbb {E} _ {y _ {1}, \dots , y _ {n} \sim \mathcal {D} _ {S}} \log \frac {1}{1 + \frac {\sum_ {i = 1} ^ {n} e ^ {\phi (x _ {0}) \top \psi (y _ {i}) / \tau}}{e ^ {\phi (x _ {0}) \top \psi (y _ {0}) / \tau}}} + \mathbb {E} _ {x _ {1}, \dots , x _ {n} \sim \mathcal {D} _ {S}} \log \frac {1}{1 + \frac {\sum_ {i = 1} ^ {n} e ^ {\phi (x _ {i}) \top \psi (y _ {0}) / \tau}}{e ^ {\phi (x _ {0}) \top \psi (y _ {0}) / \tau}}} ]. \tag {2}
$$
Since  $\frac{\sum_{i=1}^{n} e^{\phi(x_i)^\top \psi(y_0) / \tau}}{e^{\phi(x_0)^\top \psi(y_0) / \tau}}$  is usually quite small ( $< 0.1$ ) due to the exponentiation, we can Taylor expand the above term around 0 and discard second and higher order terms, resulting in
$$
\mathcal {L} _ {N C E} \left(\mathcal {D} _ {S}\right) \approx \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {S}} \left[ \mathbb {E} _ {y _ {1}, \dots , y _ {n} \sim \mathcal {D} _ {S}} \frac {\sum_ {i = 1} ^ {n} e ^ {\phi \left(x _ {0}\right) ^ {\intercal} \psi \left(y _ {i}\right) / \tau}}{e ^ {\phi \left(x _ {0}\right) ^ {\intercal} \psi \left(y _ {0}\right) / \tau}} + \mathbb {E} _ {x _ {1}, \dots , x _ {n} \sim \mathcal {D} _ {S}} \frac {\sum_ {i = 1} ^ {n} e ^ {\phi \left(x _ {i}\right) ^ {\intercal} \psi \left(y _ {0}\right) / \tau}}{e ^ {\phi \left(x _ {0}\right) ^ {\intercal} \psi \left(y _ {0}\right) / \tau}} \right], \tag {3}
$$
$$
\mathcal {L} _ {N C E} \left(\mathcal {D} _ {S}\right) \approx n \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {S}} \left[ \mathbb {E} _ {y _ {1} \sim \mathcal {D} _ {S}} e ^ {\phi \left(x _ {0}\right) ^ {\intercal} \left[ \psi \left(y _ {1}\right) - \psi \left(y _ {0}\right) \right] / \tau} + \mathbb {E} _ {x _ {1} \sim \mathcal {D} _ {S}} e ^ {\left[ \phi \left(x _ {1}\right) - \phi \left(x _ {0}\right) \right] ^ {\intercal} \psi \left(y _ {0}\right) / \tau} \right]. \tag {4}
$$
A successful generalization to the test distribution means that the infoNCE loss  $\mathcal{L}_{NCE}(\mathcal{D}_T)$  is small when the expectation over training distribution for  $x_0, y_0, x_1, y_1$  is substituted with the expectation over test distribution, where
$$
\mathcal {L} _ {N C E} \left(\mathcal {D} _ {T}\right) \approx n \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {S}} \left[ \mathbb {E} _ {y _ {1} \sim \mathcal {D} _ {T}} e ^ {\phi (x _ {0}) ^ {\intercal} [ \psi (y _ {1}) - \psi (y _ {0}) ] / \tau} + \mathbb {E} _ {x _ {1} \sim \mathcal {D} _ {T}} e ^ {[ \phi (x _ {1}) - \phi (x _ {0}) ] ^ {\intercal} \psi (y _ {0}) / \tau} \right]. \tag {5}
$$
3.2.1 Reduction to Zero-Order Dot Product
We argue that Eqn. 5 can be reduced to straightforwardly comparing the dot products. When we do not have of the test distribution  $\mathcal{D}_T$ , the only thing that we can do is to perform a zeroth-order approximation of Eqn. 5. The most intuitive way of doing so is to reduce the distribution of  $\phi(x_1)$  and  $\psi(y_1)$  to be 0 deterministically, so that:
$$
\mathcal {L} _ {N C E} ^ {(0)} \left(\mathcal {D} _ {T}\right) = 2 n \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {T}} e ^ {- \phi \left(x _ {0}\right) ^ {\top} \psi \left(y _ {0}\right) / \tau}. \tag {6}
$$
This shows that we can essentially just output  $e^{-\phi (x_0)^\top \psi (y_0)}$  as the distance between  $x_0$  and  $y_{0}$ . So we get the zeroth-order similarity:
$$
S _ {(0)} \left(x _ {0}, y _ {0}\right) = \phi \left(x _ {0}\right) ^ {\top} \psi \left(y _ {0}\right). \tag {7}
$$
3.2.2 Efficient First-Order Approximation
The common practice of taking a naive zeroth-order approximation can result in a loss of some important information of the distribution of negative samples. On the other hand, computing Eqn.5 for each pair  $x_0, y_0$  involves iterating over all samples in a distribution  $\mathcal{D}_T$ , and is computationally inefficient. We therefore turn to more efficient approximations using statistics of higher-order moments of the distribution. In our experiments, we found that this intuitive modification preserves most of the useful information as shown in Appendix D.2.
If we can estimate the first moment of the distribution, i.e., the mean of the unlabeled distribution  $\mu_x$  and  $\mu_y$  for  $\phi(x_1)$  and  $\psi(y_1)$ , we can perform a first-order approximation of the distribution with  $\widehat{P}(\phi(x_1)) = \mathbb{I}\{\phi(x_1) = \mu_x\}$  and  $\widehat{P}(\psi(y_1)) = \mathbb{I}\{\psi(y_1) = \mu_y\}$ , so that  $\widehat{P}(\phi(x_1)), \widehat{P}(\psi(y_1))$  matches the true distribution in terms of the first-order moment. This simplifies Eqn. 5 to be:
$$
\mathcal {L} _ {N C E} ^ {(1)} \left(\mathcal {D} _ {T}\right) = \mathbb {E} _ {x _ {0}, y _ {0} \sim \mathcal {D} _ {T}} \left[ e ^ {\phi \left(x _ {0}\right) ^ {\intercal} \left(\mu_ {y} - \psi \left(y _ {0}\right)\right) / \tau} + e ^ {\left(\mu_ {x} - \phi \left(x _ {0}\right)\right) ^ {\intercal} \psi \left(y _ {0}\right) / \tau} \right]. \tag {8}
$$
Therefore, to evaluate the distance of any  $x_0,y_0$ , we can simply output  $e^{\phi (x_0)^\top (\mu_y - \psi (y_0)) / \tau} + e^{(\mu_x - \phi (x_0))^{\top}\psi (y_0) / \tau}$  without having to calculate an expectation by iterating through all samples in the distribution.
3.3 Distribution Normalization
Considering that  $e^{\phi(x_0)^\top(\mu_y - \psi(y_0)) / \tau} + e^{(\mu_x - \phi(x_0))^\top\psi(y_0) / \tau}$  is equivalent to taking an algorithmic mean of  $e^{\phi(x_0)^\top(\mu_y - \psi(y_0)) / \tau}$  and  $e^{(\mu_x - \phi(x_0))^\top\psi(y_0) / \tau}$ , to simplify this further, we can convert this algorithmic mean to a geometric mean  $\sqrt{e^{\phi(x_0)^\top(\mu_y - \psi(y_0)) / \tau}e^{(\mu_x - \phi(x_0))^\top\psi(y_0) / \tau}}$  which achieves a similar effect as averaging the distance measure of  $e^{\phi(x_0)^\top(\mu_y - \psi(y_0)) / \tau}$  and  $e^{(\mu_x - \phi(x_0))^\top\psi(y_0) / \tau}$ . An additional empirical verification of the correlation between algorithmic mean and geometric mean in this step is provided in Appendix D.2. Since  $\sqrt{e^{\phi(x_0)^\top(\mu_y - \psi(y_0)) / \tau}e^{(\mu_x - \phi(x_0))^\top\psi(y_0) / \tau}}$  is a monotone function of  $-( \phi(x_0) - \frac{1}{2}\mu_x)^\top (\psi(y_0) - \frac{1}{2}\mu_y)$  after removing the square root and exponentiation, we can simply use  $-( \phi(x_0) - \frac{1}{2}\mu_x)^\top (\psi(y_0) - \frac{1}{2}\mu_y)$  to calculate the distance, resulting in a simplified first-order similarity measure:
$$
S _ {(1)} \left(x _ {0}, y _ {0}\right) = \left(\phi \left(x _ {0}\right) - \frac {1}{2} \mu_ {x}\right) ^ {\intercal} \left(\psi \left(y _ {0}\right) - \frac {1}{2} \mu_ {y}\right). \tag {9}
$$
Note that the new similarity measure is equivalent to subtracting the means immediately after the original representations are calculated, a method we call Distribution Normalization (DN), after
which the similarity is just their dot product. Clearly, our proposal is extremely easy to implement when access to the mean of the distribution of interest is available.
Fortunately, we observed that the mean of the distribution can usually be robustly estimated from very few unlabeled samples as shown in Section 4.6.3. This is easily accessible in all the downstream tasks that we have considered in this paper and is much cheaper than fine-tuning where not only more data is required but also the data needs to be annotated. For example, in the task of cross-modal retrieval, the mean can be estimated from the gallery data and past queries (without knowing the correct answer to the queries); in the task of zero-shot classification, the mean can be estimated from a small and unlabeled validation set; and in the task of image caption metrics, the mean can be estimated from the images to be captioned and the references provided.
Our experiments are designed to answer the following questions: 1) whether our proposed DN can uniformly improve a wide range of cross-modal alignment tasks for different kinds of cross-modal representation models and whether this gain is larger than that achieved by other zeroshot CLIP augmentations, 2) whether DN can be used in parallel with other common test-time adaptation methods compatible with CLIP, 3) how robust DN is when only scarce, unlabeled data is available to estimate the mean from, and 4) whether DN can improve the performance of fine-tuned models in addition to pre-trained models.
4.1 Downstream Tasks
Cross-modal Retrieval includes two subtasks where images are used to query corresponding text and vice versa. Following [68, 33], for each subtask, we consider both the zero-shot setting and the fine-tuning setting on COCO [39] and Flickr30K [48]. In the zero-shot setting, we directly apply the pre-trained model on the test set of each dataset while in the fine-tuned setting we first fine-tune each pre-trained model on the training set of each dataset with the InfoNCE loss and apply the fine-tuned model to the test set. For each setting, we report Recall@k ( $k = 1,5$ ) whereby the retrieval is considered a success if its corresponding image/text is in the top-k.
Zero-shot Classification is the task where we directly apply the pre-trained models to predict the category of the images in the test set. Each class is represented with a fixed textual template of "a photo of {class_name}". and the corresponding representation is obtained by encoding those sentences with the text encoder of the pre-trained model. For each image, we encode it with the image encoder and output the top-k most similar class representations as predictions. We benchmark all methods with top-1 and top-5 accuracy on ImageNet1K [9], Cifar100 [28], SUN397 [67], Stanford Cars [27], Caltech 101 [31], and Flowers102 [44] datasets.
Image Captioning Metric is the task that measures the correlation between human ratings and similarity metrics on machine-generated captions. Following [17], we consider both the reference-free setting and the reference-based setting. In the reference-free setting, the metric does not have access to human-written example captions (references) and has to give a score based on the source image and the generated caption alone while in the reference-based setting the metric can additionally take advantage of the information from the provided references.
This paper compares the augmentations of Distribution Normalization and other baselines on CLIP [50], ALBEF [33], and TCL [68]. More details of the base representation models can be found in Appendix B. In particular, the following are the main baselines being examined in the experiments.
TTA [55] aims to better adapt vision encoders to out-of-distribution test samples by augmenting test images and taking a weighted reduction of the various augmentations. However, since DN is meant for the zero-shot setting, training these test-time augmenting layers is not applicable in our setting. Instead, we apply the standard test-time augmentation policy from Shanmugam et al. [55] which consists of a combination of flips, crops and scales, and take a mean over the model outputs on these augmented images.
CALIP [14] replaces CLIP's dot-product similarity with a parameter-free attention module. It improves performance on zero-shot image classification by enabling visual and textual representations
to better explore cross-modal features. This method results in improved adaptive zero-shot alignment by blending images with textual-aware signals. We implement CALIP using their official code and report our results on the "ViT-B/32"-based CLIP architecture, consistent with our approach in other experiments. Since the original paper focuses solely on zero-shot image classification, CALIP serves as our baseline for this task.
TPT [56] is another CLIP zeroshot method that performs prompt tuning on each single test image to promote the model's prediction consistency on different augmented views of the single test image. TPT is mainly proposed for classification tasks and it is unclear how it can be used for general cross-modal alignment tasks. We implement TPT using their official code and report our results on the "ViT-B/32"-based CLIP architecture, consistent with our approach in other experiments. Importantly, TPT is much more computationally demanding than vanilla CLIP dot product because it requires taking gradient steps for each single test image over its 64 different augmentations, while other methods do not suffer from such computation issues.
DN is our proposed plug-in augmentation using Eqn.9. Estimated image mean and text mean are respectively subtracted from image and text embeddings immediately after the embeddings are calculated from base encoder models. Dot products are then taken between text and image representations as the similarity. Different ways to estimate the image mean and text mean from few validation samples or references in different downstream tasks are detailed in the following sections.
DN* Considering the motivation for DN is to better align the downstream similarity with the pre-training goal, it might introduce instability when there is a significant mismatch between the pre-raining and testing distribution. We therefore examine this variant called DN* where the only difference from DN is that DN* averages the similarity output by DN and a vanilla zero-order dot product for improved robustness.
4.3 Cross-modal Retrieval Results
We conducted experiments on both image-to-text and text-to-image retrieval on Microsoft COCO [39] and Flicker30K [48] (see Appendix C.1 for datasets details). For all the retrieval tasks, we estimated the mean with 100 random unlabeled samples from the validation set and calculated standard deviations and average recalls with 5 random seeds.