SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Michael Kirchhof
michael.kirchhof@uni-tuebingen.de
University of Tübingen, Tübingen AI Center
Bálint Mucśányi
Enkelejda Kasneci
TUM University
Representation learning has significantly driven the field to develop pretrained models that can act as a valuable starting point when transferring to new datasets. With the rising demand for reliable machine learning and uncertainty quantification, there is a need for pretrained models that not only provide embeddings but also transferable uncertainty estimates. To guide the development of such models, we propose the Uncertainty-aware Representation Learning (URL) benchmark. Besides the transferability of the representations, it also meaExamplesusses the zero-shot transferability of the uncertainty estimate using a novel metric. We apply URL to evaluate eleven uncertainty quantifiers that are pretrained on ImageNet and transferred to eight downstream datasets. We find that approaches that focus on the uncertainty of the representation itself or estimate the prediction loss directly outperform those that are based on the probabilities of upstream classes. Yet, achieving transferable uncertainty quantification remains an open challenge. Our findings indicate that it is not necessarily in conflict with traditional representation learning goals. Code is available at https://github.com/mkirchhof/url.
Pretrained models are a vital component of many machine learning applications. The driving force behind their development has been representation learning benchmarks, e.g. Roth et al. (2020); Chen et al. (2020): They task models to output representations  $e(x)$  of input data  $x$  that generalize across datasets in a zero-shot manner. These pretrained representations provide a valuable starting point for downstream applications, requiring less supervised data to be fine-tuned for specific tasks.
At the same time, uncertainty quantification remains a major challenge in the recent efforts towards reliable machine learning (Collier et al., 2023; Tran et al., 2022). Uncertainty quantification refers to estimating the degree of uncertainty or risk  $u(x) \in \mathbb{R}$  in a model's prediction. This is particularly important in high-stakes applications such as medical image classification. Here, the model can refrain from making predictions if the uncertainty, e.g.,  $u(x) \coloneqq 1 - \max_y P(Y = y|x)$ , is too high (Zou et al., 2023; Bouvier et al., 2022). Beyond classification, uncertainty is an inherent property of vision and language (e.g., low image resolution or ambiguous text inputs) that cannot be learned away even with large amounts of data (Chun et al., 2022; Kendall and Gal, 2017). Consequently, recent literature suggests representing images not as points  $e(x)$ , but as probabilistic embeddings (Kirchhof et al., 2023; Collier et al., 2023; Chun et al., 2021). Here,  $u(x)$  is the variance parameter of a distribution around  $e(x)$  in the embedding space, representing the input's inherent ambiguity. This can then be utilized for uncertainty-aware retrieval.
A major hurdle on the way to reliable uncertainty estimates is that  $u(x)$  needs to be trained from the ground up for each specific task, requiring substantial labeled data. Replicating the successes of representation learning promises to reduce this burden by pretraining a  $u(x)$  which can be transferred to downstream tasks in a zero- or finetuned few-shot manner. Yet, this transferability of  $u(x)$  to new datasets has not been tested in literature, with previous benchmarks evaluating on the same datasets they trained on (Detommaso et al., 2023; Nado et al., 2021). Thus, we propose a novel Uncertainty-Aware Representation Learning (URL) benchmark. Models that output both embeddings  $e(x)$  and uncertainty estimates  $u(x)$  of any form are pretrained on large collections of upstream data and evaluated on unseen downstream datasets. The transferability of their embeddings  $e(x)$  is evaluated in terms of the Recall@1 (R@1), as in established representation learning benchmarks (Roth et al., 2020; Chen et al., 2020). The transferability of their uncertainty estimates  $u(x)$  is evaluated with a novel metric, the Recall@1 AUROC (R-AUROC). It naturally extends R@1-based benchmarks and can be seamlessly integrated in as little as four lines of code, without requiring any new ground-truth labels. Nonetheless, it is not only an abstract metric but has practical significance: Models with higher R-AUROC are also more aligned with human uncertainties and react better to uncertainty-inducing interventions like image cropping.
On this benchmark, we reimplement and train eleven state-of-the-art uncertainty estimators, from class-entropy baselines over probabilistic embeddings to ensembles, with ResNet (He et al., 2016) and ViT (Dosovitskiy et al., 2021) backbones on ImageNet-1k (Deng et al., 2009). Our main findings are:
1. Transferable uncertainty estimation is an unsolved challenge (Section 4.2),
2. MCInfoNCE and direct loss prediction generalize best (Section 4.3),
3. Uncertainty estimation is not always in conflict with embedding estimation (Section 4.4),
4. Models with good uncertainties upstream are not necessarily good downstream (Section 4.5),
5. URL captures how aligned a model is with human uncertainty (Section 4.6).
These findings demonstrate that pretraining models for downstream uncertainty estimation is an important yet unsolved challenge. We hope that our benchmark will serve as a valuable resource in guiding the field towards pretrained models with reliable and transferable uncertainty estimates.
Our benchmark connects recent uncertainty quantification benchmarks with representation and zero-shot learning for unseen data, which we introduce below. Specific datasets and methods for uncertainty quantification are described in the experiments section when they are benchmarked.
Uncertainty benchmarks. Uncertainty quantification has become an essential consideration for reliable machine learning, and so several libraries have been recently developed to guide its advancement (Detommaso et al., 2023; Nado et al., 2021). These libraries provide various metrics for evaluating and improving uncertainty estimates on in-distribution data. Galil et al. (2023b) and Galil et al. (2023a) benchmarked over 500 large vision models trained on ImageNet from the timm (Wightman, 2019) library and reported that Vision Transformers (ViT) provide the best uncertainty estimates. Further, scaling of these ViTs to up to 22B parameters and pretraining on a large corpus of upstream data results in very accurate uncertainty estimates (Dehghani et al., 2023; Tran et al., 2022). However, when moving away from in-distribution data, the quality of uncertainty estimates deteriorates (Tran et al., 2022) and we can only expect that they will be generally higher and allow for out-of-distribution detection (Ovadia et al., 2019). This motivates our benchmark: We aim to develop pretrained models that can generalize their uncertainty estimates and discriminate certain from uncertain examples even within unseen datasets. While some works applied their uncertainty estimates to unseen datasets (Cui et al., 2023; Collier et al., 2023; Ardeshir and Azizan, 2022; Karpukhin et al., 2022), their downstream evaluations focused on embeddings, leaving the uncertainty estimates untested. Our benchmark intends to bridge this gap and assess the transferability of uncertainty estimates, with the goal of enhancing large pretrained models towards zero-shot uncertainty estimation. To design a benchmark for transferability, we connect the upper benchmarking techniques to paradigms from representation and zero-shot learning below.
1 def url_benchmark(pretrained_model, downstreamloader):   
2 # Predict embeddings and uncertainties on downstream data   
3 labels  $=$  embeddings  $\equiv$  uncertainties  $\equiv$  list()   
4 for (image, label) in downstreamloader:   
5 embedding, uncertainty  $\equiv$  pretrained_model(image)   
6 (labels, embeddings, uncertainties).append labelled, embedding, uncertainty)   
7   
8 # Calculate R@1 and R-AUROC   
9 next-neighboridx  $\equiv$  search_most_similar(embeddings)   
10 is_same_class  $\equiv$  labels  $= =$  labels[next_neighboridx]   
11 r_at_1  $\equiv$  mean(is_same_class)   
12 r_auroc  $\equiv$  compute_auroc(uncertainties, not is_same_class)   
13   
14 return r_at_1, r_auroc
Algorithm 1: Adding URL to existing representation learning benchmarks takes only the four highlighted lines of code.
Representation and zero-shot learning. Transferability is generally evaluated by testing whether models can make sensible decisions on unseen data. In zero-shot learning (Xian et al., 2017), the model is tasked to give class-predictions on new downstream classes. This requires both learning a transferable representation space on upstream data and creating classifier heads for the new classes from auxiliary information. Representation learning benchmarks (Roth et al., 2020; Khosla et al., 2020; Bengio et al., 2013) focus on the former. To this end, they use a metric similar to class accuracy, the Recall@1 (R@1) (Mikolov et al., 2013). It calculates the model's embeddings of all unseen downstream data and compares whether each embedding's next neighbor is in the same class or not. This tells whether the embeddings are semantically meaningful, such that the pretrained model can be successfully transferred to downstream tasks. We extend representation learning benchmarks to additionally judge the transferability of uncertainty estimates. To this end, we propose a metric that can be implemented on top of the R@1 in four lines of code in the next section.
3 Uncertainty-aware representation learning (URL) benchmark
3.1 Evaluating uncertainty about representations
To quantify its uncertainty, a model  $f: \mathcal{X} \to \mathcal{E} \times \mathcal{U}$  is assumed to predict both an embedding  $e(x) \in \mathcal{E}$  and a scalar uncertainty value  $u(x) \in \mathcal{U} \subset \mathbb{R}$  for each input image  $x \in \mathcal{X}$ . We do not impose restrictions on how  $u(x)$  is calculated, e.g., it could be the negative maximum probability of a softmax classifier, a predicted variance from a dedicated uncertainty module, or the disagreement between ensemble members. The predicted uncertainty  $u(x)$  is commonly benchmarked in terms of its expected calibration error (ECE), negative log-likelihood, area under the receiver-operator characteristics curve (AUROC), or abstained prediction curves. All of these measures are w.r.t. the correctness of a classification decision. Hence,  $u(x)$  can only be evaluated in-distribution with known classes. Our setup involves unseen datasets and classes, so we need to develop a fitting measure.
To this end, let us take Lahlou et al. (2023)'s decision-theoretic perspective on uncertainty quantification: Uncertainty quantification is loss prediction. The uncertainty expresses the expected loss of a model's decision on a specific datapoint. In Gaussian regression with an  $L_{2}$  loss, the expected loss is the target's variance, so an uncertainty quantifier  $u(x)$  should be proportional to it. In classification with a 0-1 loss,  $u(x)$  should be proportional to the probability of returning the correct class.
In representation learning, the model's decision is the embedding  $e(x)$  and the loss is the R@1. The uncertainty quantifier's goal is then to report the loss attached to the embedding, i.e.,  $u(x)$  should be proportional to whether the R@1 will be correct or not. This demonstrates the use case of  $u(x)$ : Telling whether an embedding  $e(x)$  can be trusted or could be misplaced in the embedding space. This is an important property as models of the form  $x \to e \to y$  have an information bottleneck in the quality of the embedding  $e(x)$  due to the data-processing inequality (Cover, 1999). For every downstream task, a higher uncertainty  $u(x)$  about  $e(x)$  monotonically increases the loss of  $y(e(x))$ . In other words, if the embedding is wrong, then the prediction in any downstream task will be wrong.
We measure whether the uncertainty quantifier  $u(x)$  is proportional to the correctness of the embedding  $e(x)$  via the AUROC with respect to whether the R@1 is correct (one) or not (zero), named R-AUROC. As the R@1 is a 0-1 loss, the R-AUROC can be interpreted as the probability that an incorrect embedding will receive a higher uncertainty score than a correct embedding (Fawcett, 2006). An R-AUROC close to 1 means that  $u(x)$  clearly separates correct from incorrect embeddings, while an R-AUROC of 0.5 means that it has no more predictive power than a random guess.
A positive trait of the R-AUROC is that it is indicative of how well-aligned the model is with human uncertainties and how well the model reacts to uncertainty-inducing interventions (see Fig. 5 and Section 4.6). It also does not require uncertainty ground-truths and takes only four lines of code to implement into existing representation learning benchmarks, as shown in Algorithm 1. We choose the AUROC over other calibration measures such as the ECE because it accepts uncertainties  $u(x) \in \mathbb{R}$  (instead of  $u(x) \in [0,1]$ ) and because it avoids some loopholes of the ECE. We discuss these and more design choices behind this metric in Appendix A.
3.2 URL benchmark protocol
The R-AUROC can be evaluated on any downstream dataset that allows calculating the  $\mathbf{R}@\mathbf{1}$ , i.e., has class labels. Yet, in order to keep future results comparable, we propose a benchmark protocol for uncertainty-aware representation learning (URL). Our code is based on timm (Wightman, 2019) and available at https://github.com/mkirchhof/url.
Datasets. We train each model on ImageNet-1k (Deng et al., 2009) as upstream dataset. We note that future works may use larger-scale upstream datasets (Collier et al., 2023; Tran et al., 2022) or auxiliary information (Han et al., 2023; Ortiz-Jimenez et al., 2023), as long as they stay disjoint to the downstream datasets. As downstream datasets, we follow the standardized representation learning protocol of Roth et al. (2020) and use CUB-200-2011 (Wah et al., 2011), CARS196 (Krause et al., 2013), and Stanford Online Products (Song et al., 2016). We follow the original splits that divide their classes into equally sized train and test sets. Following Roth et al. (2020), we further divide the classes in the train set equally into a train and a validation split. In our zero-shot transferring setup, all models are trained only on the upstream ImageNet dataset and do not use the downstream train splits. All results report the performance on the test sets, averaged across the three datasets and three seeds.
Hyperparameters. We use the downstream validation split to select the best learning rate, early stopping, and further hyperparameters of each model individually, see also Appendix B. Each model is tuned for 10 search iterations via Bayesian Active Learning (Biewald, 2020). The best model is chosen based on the R-AUROC on validation data, where models with an  $R@1$  below 0.1 on the validation splits are filtered out. The best model is replicated on three seeds.
Architectures. Following uncertainty quantification and representation learning benchmarks (Wen et al., 2021; Dusenberry et al., 2020; Roth et al., 2020), we use a ResNet-50 (He et al., 2016) with an embedding space dimension of 2048 as a backbone. We further study ViT-Medium (Dosovitskiy et al., 2021) backbones due to their performance (Galil et al., 2023a,b) and increasing number of large-scale uncertainty quantifiers built on top of them (Collier et al., 2023; Tran et al., 2022). Methods that predict  $u(x)$  with explicit modules use a 3-layer MLP head attached to the embeddings.
Training infrastructure. Each model is trained with an aggregated batch size of 2048, as recent studies indicate higher batch sizes might benefit uncertainty quantification (Galil et al., 2023b). We use the Lamb optimizer (You et al., 2020) with cosine annealing learning rate scheduling (Loshchilov and Hutter, 2017) for all models since it performed best in preliminary experiments. The ResNets and ViTs are trained on NVIDIA RTX 2080 Ti and A100 GPUs, respectively, for 32 epochs from a checkpoint pretrained on ImageNet to reduce the computational costs. In total, the experiments took 3.2 GPU years of runtime.
Further metrics. Uncertainty estimates aim to assess the errors made by individual models, so that they are necessarily model- and performance-dependent. To provide a comprehensive view, we not only evaluate the quality of the uncertainty estimate using R-AUROC but also consider the model's representation learning performance using R@1.
4.1 Uncertainty estimators
We apply URL to benchmark two baselines (CE, InfoNCE), five probabilistic embeddings approaches (MCInfoNCE, ELK, nivMF, HIB, HET-XL), two direct variance models (Losspred, SNGP), and two ensembles (Ensemble, MCDropout). We introduce each approach below and explain further details on their implementations and hyperparameters in Appendix B and runtimes in Appendix C.7.
Cross Entropy (CE) is a supervised baseline which trains under a cross-entropy loss. It uses the entropy of the upstream class probabilities  $u(x) \coloneqq \mathcal{H}(P(Y|x))$  as uncertainty estimate.
InfoNCE (Oord et al., 2018) is an unsupervised baseline. Following SIMCLR (Chen et al., 2020), it takes two random transforms of each image and pulls their embeddings towards each other and repels them from the remaining batch. InfoNCE itself does not estimate  $u(x)$ , so we use the embedding norm  $u(x) \coloneqq \| e(x)\|_2$  as a heuristic (Kirchhof et al., 2022; Scott et al., 2021; Li et al., 2021).
MCInfoNCE (Kirchhof et al., 2023) follows the unsupervised setup of InfoNCE, but predicts a certainty  $\kappa(x) \eqqcolon 1 / u(x)$  along with each embedding to define a distribution in the embedding space, so called probabilistic embeddings. It draws samples from them and applies InfoNCE on each.
Expected Likelihood Kernel (ELK) (Kirchhof et al., 2022; Shi and Jain, 2019) also predicts certainties  $\kappa(x) \eqqcolon 1 / u(x)$  to define probabilistic embeddings. The probabilistic embeddings are compared to class distribution via an expected likelihood distribution-to-distribution kernel (Jebara and Kondor, 2003). This makes it a supervised probabilistic embedding-based loss.
Non-isotropic von Mises-Fisher (nivMF) (Kirchhof et al., 2022) is analogous to ELK, but models class distributions as non-isotropic von Mises-Fisher distributions, thereby allowing different variances along each embedding space axis. Image certainties are still scalars  $\kappa(x) \eqqcolon 1 / u(x)$ .
Hedged Instance Embeddings (HIB) (Oh et al., 2019) predicts variances  $\sigma(x) \coloneqq u(x)$  of probabilistic embeddings. Samples are drawn to compute match probabilities between two images. It aims to increase the match probabilities of same-class pairs and decrease that of different-class ones.
Heteroscedastic Classifier (HET-XL) (Collier et al., 2023) differs from the above probabilistic embeddings approaches in that it predicts full covariance matrices  $\Sigma(x)$  in the embedding space. It draws samples from these probabilistic embeddings to calculate the expected  $P(Y|x)$ . We test both  $u(x) \coloneqq \det \Sigma(x)$  and the class entropy  $u(x) \coloneqq \mathcal{H}(P(Y|x))$  as possible uncertainty estimates.
Spectral-normalized Neural Gaussian Processes (SNGP) (Liu et al., 2020) model class logits as Gaussian Processes with a predicted mean and a heteroscedastic variance. They are pooled into class probabilities  $P(Y|x)$  and trained under a CE loss. The entropy of these probabilities serves as uncertainty value  $u(x) \coloneqq \mathcal{H}(P(Y|x))$ .
Loss Prediction (Losspred) approaches (Upadhyay et al., 2023; Lahlou et al., 2023; Levi et al., 2022; Laves et al., 2020; Yoo and Kweon, 2019) in regression treat uncertainty quantification as secondary regression task. We apply the same principle to classification, where we task the uncertainty module  $u(x)$  to predict the (gradient-detached) CE loss at each sample via an  $L_{2}$  loss added to the train loss.
Deep Ensembles (Lakshminarayanan et al., 2017) train multiple randomly initiated networks under a CE loss to obtain several predictions. They are pooled one class distribution  $P(Y|x)$ . We define the uncertainty either via its entropy  $u(x) \coloneqq \mathcal{H}(P(Y|x))$  or the Jensen-Shannon divergence between the ensemble members' class probability distributions. Following (Lee et al., 2015), we only train multiple output heads and share the backbone to reduce computational complexity.
MCDropout (Gal and Ghahramani, 2016) applies Dropout (Srivastava et al., 2014) at inference time. This gives multiple predictions per input, imitating the upper Ensemble. We use both the entropy  $u(x) \coloneqq \mathcal{H}(P(Y|x))$  and the Jensen-Shannon divergence between the ensemble members' class probability distributions as possible uncertainty metrics.
4.2 Transferable uncertainty estimation is an unsolved challenge
Fig. 1 presents the URL benchmark results, i.e., the R-AUROC calculated for all above approaches on ResNet and ViT backbones. The barplot shows the minimum, average, and maximum performance across three seeds of each hyperparameter-tuned approach.
Figure 1: Zero-shot uncertainty estimates of pretrained models (bars) do not reach the performance of many-shot models yet (dashed line). The URL benchmark aims to guide the field to close this gap. Minimum, average, and maximum R-AUROC across three seeds.
Before comparing the models, we first investigate whether the transferability problem URL addresses is already solved by any of the existing methods. To obtain an upper reference, we additionally train a ResNet 50 with standard cross-entropy loss and entropy of the class probabilities as uncertainty prediction on the downstream test classes (split into a train and test split for this experiment only). This many-shot performance of 0.68 is not reached by any of the methods that transfer their uncertainty in a zero-shot way, marking URL as an open challenge. In the standard R@1, this gap has already been closed in representation learning (see Appendix C.1) and we hope that URL guides the field towards the same for transferable uncertainty estimation.
4.3 MCInfoNCE and direct loss prediction generalize best
To compare the approaches in detail, in addition to Fig. 1, Fig. 2 reports both the downstream uncertainty and R@1 performance. The overall best method is Losspred, with the second-best average R-AUROC of 0.568, close to the best method, MCInfoNCE, with an average R-AUROC of 0.569, while maintaining the second-best average R@1 of 0.53, close to the best R@1 of 0.57 achieved by nivMF. MCInfoNCE marks the best performance in both metrics within the ResNet models, closely followed by nivMF. This is remarkable as it is the only unsupervised method aside from the InfoNCE baseline. One final noteworthy mention is ELK which provides decent uncertainty estimates on both ResNets and ViTs, whereas most other models vary in their performance depending on the backbone.
When grouping the approaches, those that directly model the variance (Losspred, SNGP) appear to have an edge on the ViTs, especially Losspred, which is the only method that disentangles variance estimation from how the class logits are calculated. Such disentanglement via having two losses could be added to other approaches in future works. Probabilistic embeddings, especially MCInfoNCE, nivMF and ELK, also show promising performance both on the bigger ViTs and the smaller ResNets. Ensembles fail to provide transferable uncertainty estimates. The baselines unsurprisingly fail, indicating that well-calibrated class probabilities on the upstream dataset do not serve as good uncertainties on downstream data. We investigate this further in Section 4.5.
4.4 Uncertainty estimation is not always in conflict with embedding estimation
A commonly raised concern is whether or not uncertainty quantification deteriorates the prediction, or, in the representation learning setup, the embedding quality. In the previous section, we have