SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuzhe Lu\*, Yilong Qin\*, Runtian Zhai\*, Andrew Shen\*, Ketong Chen\*  
Zhenlin Wang\*, Soheil Kolouri\*, Simon Stepputtis\*, Joseph Campbell\*, Katia Sycara\*  
 $^{1}$ Carnegie Mellon University  $^{2}$ Vanderbilt University
{yuzhelu, yilongq, rzhai, andrews, ketongc, zhenlinw}@cs.cmu.edu soheil.kolouri@vanderbilt.edu, {sstepput, jcampbell, sycara}@cs.cmu.edu
Out-of-distribution (OOD) data poses serious challenges in deployed machine learning models, so methods of predicting a model's performance on OOD data without labels are important for machine learning safety. While a number of methods have been proposed by prior work, they often underestimate the actual error, sometimes by a large margin, which greatly impacts their applicability to real tasks. In this work, we identify pseudo-label shift, or the difference between the predicted and true OOD label distributions, as a key indicator of this under-estimation. Based on this observation, we introduce a novel method for estimating model performance by leveraging optimal transport theory, Confidence Optimal Transport (COT), and show that it provably provides more robust error estimates in the presence of pseudo label shift. Additionally, we introduce an empirically-motivated variant of COT, Confidence Optimal Transport with Thresholding (COTT), which applies thresholding to the individual transport costs and further improves the accuracy of COT's error estimates. We evaluate COT and COTT on a variety of standard benchmarks that induce various types of distribution shift – synthetic, novel subpopulation, and natural – and show that our approaches significantly outperform existing state-of-the-art methods with up to 3x lower prediction errors. Our code can be found at https://github.com/luyuzhe111/COT.
Machine Learning methods are largely based on the assumption that test samples are drawn from the same distribution as training samples, providing a basis for generalization. However, this i.i.d. assumption is often violated in real-world applications where test samples are found to be out-of-distribution (OOD) – sampled from a different distribution than during training. This may result in a significant negative impact on model performance [15, 34, 25, 38]. A common practice for alleviating this issue is to regularly gauge the model's performance on a set of labeled data from the current target data distribution, and update the model if necessary. When labeled data is not available, however, one needs to predict the model's performance on the target distribution with unlabeled data, a task known as OOD performance prediction.
Performance prediction on unlabeled data has previously been shown to be impossible without imposing additional constraints over the unknown target distribution [7, 11, 5, 26], due to the fact that target samples may take any label. Thus, the feasibility of this task is dependent on what assumptions we make regarding the shift between the train and target distributions. Prior works often make the assumption that the conditional density  $P(y|x)$  remains fixed in the presence of covariate shift [37].
However, this tells us little when  $x$  falls outside the support of the train distribution. Despite the theoretical difficulty, prior works have proposed a number of heuristic methods to estimate the performance of a model based on unlabeled target samples. For instance, Average Confidence (AC) [20, 16] estimates error based on the average maximum softmax score for target samples, assuming the model has been calibrated for the train distribution. This method was further improved with the addition of a learned threshold [13], for which the error is predicted as the fraction of samples with confidence falling below it. Other approaches estimate model performance based on a disagreement score computed between the predictions of two models trained over the same dataset [22, 3]. Some works have found that applying a transformation over target samples and estimating the effect of the transformation leads to a reliable prediction of model performance in vision tasks [9, 10]. However, many of these prior methods have been shown to underestimate the model's error when it is miscalibrated in the target distribution [13, 22]; that is, the predicted softmax scores differ from the true class likelihoods. We empirically observe that this miscalibration is strongly positively correlated with pseudo-label shift, which is the difference between the predicted target label distribution  $P_T(\tilde{y})$  and the true label distribution  $P_T(y)$ . Thus, we treat pseudo-label shift as a key indicator of error underestimation.
In this work, we propose an approach which provides robust error estimates in the presence of pseudolabel shift. Our approach, Confidence Optimal Transport (COT), leverages the optimal transport framework to predict the error of a model as the Wasserstein distance between the predicted target class probabilities and the true source label distribution. We theoretically derive lower bounds for COT's predicted error. This results in a more provably conservative error prediction than AC, which is crucial for safety in many real-world machine learning applications. In addition, we introduce a variant of COT, Confidence Optimal Transport with Thresholding (COTT), which introduces a learned threshold over the optimal transportation costs in line with prior work [13] and empirically improves upon COT's performance. In this work, we mainly target covariate shifts and compare our proposed methods to existing state-of-the-art approaches in extensive empirical experiments over eleven datasets exhibiting distribution shift from both vision and language domains. These distribution shifts include: visual corruptions (e.g., blurred image data), novel subpopulation shifts (e.g., novel appearances of a category), and natural shifts in the wild (e.g., different stain colors for pathology images), all of which are frequently experienced in the real world. We find that COT and COTT consistently avoid the significant error underestimations suffered by previous methods. In particular, COTT achieves significantly lower prediction errors than existing methods for most models and datasets (up to 3x better), establishing new state-of-the-art results.
2 Preliminaries and Motivation
In this section, we introduce the problem setup of OOD error prediction and show how a popular baseline, Average Confidence (AC), can be understood as the Wasserstein Distance (WD) between a reference label distribution and the softmax output distribution from an Optimal Transport (OT) perspective. We then demonstrate why the reference label distribution AC uses is problematic by explaining pseudo-label shift and its correlation with miscalibration. After establishing the relation and utility of OT to OOD error prediction, we formally introduce our proposed methods in Sec. 3.
2.1 OOD Performance Prediction
In this work, we address the OOD problem in the domain of classification tasks. Let  $\mathcal{X} \subseteq \mathbb{R}^d$  be the input space of size  $d$ , and  $\mathcal{V} = \{1, \dots, k\}$  be the label space, where  $k$  is the number of classes. Let the source distribution over  $\mathcal{X} \times \mathcal{V}$  be  $P_S(x, y)$  and the target distribution be  $P_T(x, y)$ . A classifier  $\vec{f}: \mathcal{X} \to \Delta^{k-1}$  maps an input to a confidence vector (i.e., the output of the softmax layer), where  $\Delta^{k-1} = \{(z_1, \dots, z_k): z_1 + \dots + z_k = 1, z_i \geq 0\}$  is the  $k$ -dimensional unit simplex. Let the training samples be  $\{(x_{\mathrm{train}}^{(i)}, y_{\mathrm{train}}^{(i)})\} \sim P_S(x, y)$ , and the validation samples be  $\{(x_{\mathrm{val}}^{(i)}, y_{\mathrm{val}}^{(i)})\} \sim P_S(x, y)$ . The validation set is used in some previous methods and will also be used in COT and COTT to perform calibration [17].
The OOD performance prediction problem is formally stated as follows: Given an unlabeled test set  $\{x_T^{(i)}\} \sim P_T(x)$ , and a classifier  $\vec{f}$  trained over the training samples, predict the accuracy of
$\vec{f}$  over the test set  $\alpha = \frac{1}{n}\sum_{i=1}^{n}\mathbb{1}[y_T^{(i)} = \arg \max_j \vec{f}_j(x_T^{(i)})]$ , where  $y_T^{(i)}$  is the ground truth label. Equivalently, we can also predict the error  $\epsilon = 1 - \alpha$  with an estimate  $\hat{\epsilon}$ .
In this work, we are further investigating the distribution of confidence vectors. While a confidence vector  $\vec{f}(x)$  itself is a distribution of labels in the  $k$ -dimensional simplex  $\Delta^{k-1}$ , the distribution of confidence vectors  $\vec{f}_{\#}P(\vec{c})$  is a distribution of distributions, where  $\vec{c} \in \Delta^{k-1}$ .  $\vec{f}_{\#}P(\vec{c})$  is defined to be the pushforward of a covariate distribution  $P(x)$  using  $\vec{f}$ :  $\vec{f}_{\#}P(\vec{c}) = P(\vec{f}^{-1}(\vec{c}))$ . Consider the following example: Suppose we have a uniform covariate distribution  $P(x)$  on  $x \in \{A, B, C\}$  and  $\vec{f}$  maps  $A, B$  to  $[0.5, 0.5]^{\top}$  and  $C$  to  $[0.7, 0.3]^{\top}$ . Then,  $f_{\#}P(\vec{c})$  will assign  $\frac{2}{3}$  mass to  $[0.5, 0.5]^{\top}$  and  $\frac{1}{3}$  mass to  $[0.7, 0.3]^{\top}$ . To facilitate easy comparison between confidence vectors and labels, we denote  $\vec{y} \in \{0, 1\}^k \cap \Delta^{k-1}$  to be the one-hot representation of  $y \in \mathcal{V}$ , where the only non-zero element in  $\vec{y}$  is the  $y$ -th element, i.e.  $\vec{y}_j = \mathbb{1}[y = j]$ . We denote the distribution of one-hot labels as  $P(\vec{y})$ . For a covariate distribution  $P(x)$ , we will refer to the distribution of predicted labels from classifier  $\vec{f}$  as the pseudo-label distribution  $P_{\mathrm{pseudo}}(\vec{y})$ . We reuse the notation to denote the probability mass of  $\vec{y}$  also as which is given by the mass of the inputs that give this prediction, i.e.  $P_{\mathrm{pseudo}}(\vec{y}) = P(\{x \in \mathcal{X} | \arg \max_j \vec{f}_j(x) = y\})$ .
2.2 Wasserstein Distance and Optimal Transport
In recent years, optimal transport theory has found numerous applications in the field of machine learning [2, 4, 1, 27, 29]. Optimal transport aims to move one distribution of mass to another as efficiently as possible under a given cost function. In the Kantorovich formulation of optimal transport, we are given two distributions  $\mu(x)$  over  $\mathcal{X}$  and  $\nu(y)$  over  $\mathcal{Y}$  and a cost function  $c(x,y)$  that tells us the cost of transporting from location  $x$  to location  $y$ . Here, we aim to find a transport plan  $\pi(x,y)$  that minimizes the total transport cost. The transport plan is a joint distribution with marginal  $\pi(\cdot,\mathcal{Y}) = \mu(\cdot)$  and  $\pi(\mathcal{X},\cdot) = \nu(\cdot)$ . The conditional distribution  $\frac{\pi(x,y)}{\mu(x)}$  of the transport plan informs us how much mass is moved from  $x$  to  $y$ . Let  $\Pi(\mu,\nu)$  be the set of all transport plans. More formally, the Wasserstein Distance is defined as:
$$
W (\mu , \nu) = \inf _ {\pi \in \Pi (\mu , \nu)} \int_ {\mathcal {X} \times \mathcal {Y}} c (x, y) d \pi (x, y)
$$
The Wasserstein distance satisfies the definition of a metric (non-negativity, symmetry, and subadditivity), inducing a metric space over a space of probability distributions. Unlike other metrics, such as total variation, the Wasserstein metric induces a weaker topology and provides a robust framework for comparing probability distributions that respect the underlying space geometry [2].
For discrete distributions  $\mu, \nu$  such as the empirical distributions, the Wasserstein distance simplifies to the following linear programming problem:
$$
W (\mu , \nu) = \min _ {\mathbf {P} \in \Pi (\mu , \nu)} \langle \mathbf {P}, \mathbf {C} \rangle = \min _ {\mathbf {P} \in \Pi (\mu , \nu)} \sum_ {i, j} \mathbf {P} _ {i j} \mathbf {C} _ {i j}
$$
where  $\mathbf{C},\mathbf{P}\in \mathbb{R}^{m\times n}$  are the cost matrix and the plan matrix respectively and  $\mathbf{C}_{ij}$  is the transport cost from sample  $i$  to sample  $j$ . When  $n = m$ , the optimal transport problem reduces to the optimal matching problem (Proposition 2.1 in Peyre et al. [32]), i.e. the optimal transport plan  $\mathbf{P}^* = \arg \min_{\mathbf{P}\in \Pi (\mu ,\nu)}\langle \mathbf{P},\mathbf{C}\rangle$  is a permutation matrix. Not only does this constraint enable efficient algorithms like the Hungarian algorithm, but it will also help draw the connection between pseudolabel shift and target error that is central to our Confidence Optimal Transport method.
2.3 Average Confidence as Wasserstein Distance
We use Average Confidence, a popular OOD performance prediction method, as the starting point of our analysis. Average Confidence with Max Confidence (AC-MC) estimates the target accuracy by taking the empirical mean of maximum confidence of the classifier over all target samples  $x^{(i)} \sim P_T(x)$ , i.e.  $\frac{1}{n} \sum_{i=1}^{n} \max_j \vec{f}_j(x^{(i)})$ . Its corresponding target error estimate is therefore  $\hat{\epsilon}_{\mathrm{AC}} = 1 - \frac{1}{n} \sum_{i=1}^{n} \max_j \vec{f}_j(x^{(i)})$ . By definition,  $|\epsilon - \hat{\epsilon}_{\mathrm{AC}}|$  measures the miscalibration of the model, i.e. how far away the model's confidence is from its actual accuracy. In the following proposition, we connect the AC-MC estimates with distances in the Wasserstein Metric Space:
Figure 1: The AC-COT Triangle in the Wasserstein Space  $(\mathcal{P}(\mathbb{R}^k), W_\infty)$ , where  $\delta_0$  (Dirac delta over the zero vector) represents the origin. Red line: AC-MC error estimate. Blue line: (our) COT error estimate (assuming  $P_T(\vec{y}) \approx P_S(\vec{y})$ ). Green line: Pseudo-label shift. Left: AC-MC predicts the target error as the distance between the distribution of confidence vectors and its projection on the unit sphere, the smallest among all label distributions. This makes AC-MC prone to underestimating the target error. Middle: Mild pseudo-label shift. Right: Severe pseudo-label shift.
Proposition 1 ( $\hat{\epsilon}_{\mathrm{AC}} - W_{\infty}$  Equivalence). Let  $(\mathcal{P}(\mathbb{R}^k),W_{\infty})$  be the metric space of all distributions over  $\mathbb{R}^k$ , where  $W_{\infty}$  is the Wasserstein distance with  $c(x,y) = \| x - y\|_{\infty}$ . Then, the estimated error of AC-MC is given by  $\hat{\epsilon}_{AC} = W_{\infty}(\vec{f}_{\#}P(\vec{c}),P_{pseudo}(\vec{y}))$ .
We defer all proofs to the supplementary material. Here, we appeal to pictorial intuition in Figure 1, where  $\delta_0$  (Dirac delta over the zero vector) represents the origin. All one-hot label distributions, including  $P_T(\vec{y})$  and  $P_{\mathrm{pseudo}}(\vec{y})$ , are represented as points on the (dashed) unit sphere around  $\delta_0$ . A distribution of confidence vectors  $\vec{f}_{\#}P(\vec{c})$  is simply a point within the unit ball around  $\delta_0$ . The AC-MC accuracy estimate measures how far the point is to the origin  $\delta_0$ . Additionally, we can project the point to the unit sphere, resulting in the projection point  $P_{\mathrm{pseudo}}(\vec{y})$ , which is the pseudo-label distribution. The AC-MC error estimate measures the length of the projection.
Corollary 1 ( $P_{\mathrm{pseudo}}(\vec{y})$  is closest to  $\vec{f}_{\#}P(\vec{c})$ ). Let  $P^{\prime}(\vec{y}) \in \mathcal{P}(\{0,1\}^{k} \cap \Delta^{k-1})$  be a one-hot label distribution. Then  $W_{\infty}(\vec{f}_{\#}P(\vec{c}), P^{\prime}(\vec{y})) \geq W_{\infty}(\vec{f}_{\#}P(\vec{c}), P_{\mathrm{pseudo}}(\vec{y}))$ .
This suggests that AC-MC, among all possible reference one-hot label distributions, selects the closest one and reports the distance to be the predicted error. Thus, we can see that AC-MC is a very optimistic prediction strategy, so it is not surprising that AC-MC is widely reported to be over-estimating the performance on real tasks [13, 16, 22, 3].
2.4 Pseudo-Label Shift and its Correlation with Miscalibration
The pseudo-label shift is defined as  $W_{\infty}(P_{\mathrm{pseudo}}(\vec{y}), P_T(\vec{y}))$ , i.e. the distance from the pseudo label distribution  $P_{\mathrm{pseudo}}(\vec{y})$  to the ground truth target label distribution  $P_T(\vec{y})$ , which is also the length of the green line segment in Figure 1. An important property is  $W_{\infty}(P_{\mathrm{pseudo}}(\vec{y}), P_T(\vec{y})) \leq \epsilon \leq 1$ , i.e. the pseudo-label shift is a lower bound of the true target error of the model (as true target error can be interpreted as the average transport cost of a suboptimal matching). We can clearly see how a large pseudo-label shift could potentially destroy the prediction of AC-MC from Figure 1 (right), where the red line segment is much shorter than the green one which should be a lower bound of the true error. This lower bound, however, can be loose when it is small, as shown in Figure 1 (middle).
Most existing prediction methods heavily rely on the model being well-calibrated, as pointed out by Jiang et al. [22], Garg et al. [13], yet the underlying difficulty is a lack of precise understanding of when and by how much neural networks become miscalibrated. This is evident in the large-scale empirical study conducted in [30], showing that the variance of the error becomes much larger among the different types of shift studied, despite some positive correlation between the expected calibration error and the strength of the distribution shift. In this section, we empirically show that there is a strong positive correlation between the pseudo-label shift and  $|\epsilon - \hat{\epsilon}_{\mathrm{AC}}|$ , which we define as the model's miscalibration. Because of this correlation, we consider pseudo-label shift as the key signal to why existing methods have undesirable performance.
We evaluate the pseudo label shift and the prediction error of AC-MC  $|\epsilon - \hat{\epsilon}_{\mathrm{AC}}|$  on CIFAR-10 and CIFAR-100 under different distribution shifts, and plot the results in Figure 2 (left). The plots demonstrate a strong positive correlation between these two quantities, which not only means that the performance of AC-MC worsens as the pseudo-label shift gets larger, but also implies that the performance of existing methods [16, 13, 22] that depend on model calibration will drop.
Figure 2: Left: Absolute difference between the model's Average Confidence and its true error under different levels of pseudo-label shift. This difference measures the degree of miscalibration. We see a strong correlation between the miscalibration and pseudo-label shift on the common corruption benchmarks (CIFAR-10-C, CIFAR-100-C). Right: Absolute difference between GDE error estimate and true error under different levels of pseudo-label shift. The strong correlation is also observed.
Figure 3: We further compare the sensitivity of COT and COTT's estimation error to the degree of pseudo-label shift. Compared to Figure 2, we can clearly see that the correlation between the prediction error and pseudo-label shift weakens significantly. Moreover, COTT is even more robust to pseudo-label shift compared to COT.
Our exploration of miscalibration and pseudo-label shift reveals a previously unexplored tradeoff: When the pseudo-label shift is small, prior work has given us some reassurance to trust the calibration of neural networks, which motivated methods such as AC [20] and ATC [13]. More critically, as miscalibration worsens, the pseudo-label shift becomes a tighter low bound of the error and thus a more trustworthy estimate of the error.
Motivated by this observation, we leverage the information of the pseudo-label shift to improve the performance of confidence-based prediction methods for miscalibrated models. The problem, as mentioned earlier, is that without any information about the target label distribution, it is theoretically impossible to estimate the pseudo-label shift when the test set contains data from unseen domains. Thus, the only option is to make assumptions on the target label distribution.
In our proposed method, we make a natural assumption: the target label distribution is close to the source label distribution. This assumption aligns with most natural shifts that can be observed in real-world datasets. This extra assumption allows us to develop COT and COTT, which perform much better than existing methods in most cases (See Table 1), especially when the pseudo-label shift is large. Given this assumption, Figure 3 (Left), shows the prediction error of COT versus the pseudo-label shift. We can see that equipped with the extra information, COT is able to maintain a low prediction error even under very large pseudo-label shift, and the correlation between COT performance and the pseudo-label shift is weak. Since the pseudo-label shift is strongly correlated with miscalibration, this implies that COT is much more robust to miscalibration than existing miscalibration-sensitive prediction methods.
In this section, we formally introduce our proposed method – Confidence Optimal Transport – and propose an additional variation, COTT, that utilizes a threshold over the optimal transportation costs between two distributions, instead of taking a simple average.
3.1 Confidence Optimal Transport
Let  $\hat{P}_S(\vec{y})$  denote the empirical source label distribution. The predicted error of COT, which is the length of the blue line segment (assuming  $P_T(\vec{y}) = P_S(\vec{y})$ ) in Figure 1, is given by
$$
\hat {\epsilon} _ {\mathrm {C O T}} = W _ {\infty} \left(\vec {f} _ {\#} \hat {P} _ {T} (\vec {c}), \hat {P} _ {S} (\vec {y})\right).
$$
By Corollary 1,  $\hat{\epsilon}_{\mathrm{COT}} \geq \hat{\epsilon}_{\mathrm{AC}}$ , which means that COT provably predicts a larger error than AC-MC, which empirically tends to produce more accurate predictions as we find overestimation of the error far less common than underestimation. As mentioned earlier, in the case where the pseudo-label shift is large, such as in Figure 1 (Right), AC-MC can have arbitrarily large prediction error, while COT always has the following guarantee:
Proposition 2 (Calibration independent lower bound of COT). Under the assumption that  $P_T(\vec{y}) = P_S(\vec{y})$ , we always have  $\hat{\epsilon}_{\mathrm{COT}} \geq 0.5W_{\infty}(P_{\mathrm{pseudo}}(\vec{y}), P_T(\vec{y}))$ .
Thus, we can see that COT is by nature different from existing methods because it is a miscalibration-robust method. Its success is dependent on the difference between  $P_T(\vec{y})$  and  $P_S(\vec{y})$ , which is relatively small in most real-world scenarios, while the dependency on calibration of existing methods can not always be controlled. The geometric explanation of this is that COT measures a fundamentally different length in the Wasserstein Space  $(\mathcal{P}(\mathbb{R}^k), W_\infty)$ , allowing for the above guarantee which does not exist in previous methods. Moreover, as we will empirically demonstrate in the next section, large pseudo-label shift is prevalent in real models and datasets. Consequently, COT performs much better than miscalibration-sensitive methods in most cases.
3.2 Thresholding as a Robust Transport Cost Statistic
Computationally, COT (or Wasserstein Distance in general) is implemented as a two-step process: 1) calculating individual transport costs for all samples; 2) returning the mean across all transport costs as the error estimate. While we have seen that COT has some protection against miscalibration, it is not completely immune. Another outstanding issue lies in the second step - computing the mean of all transport costs. In statistics, it is well known that the mean is less robust to outliers than the median [21]. In the supplemental material, we show that the empirical transport cost distribution is also heavy-tailed and therefore COT is susceptible to outlier costs. A large fluctuation in the outliers impacts the mean much more than it does to the median. Therefore, a more robust statistic is desired to protect COT against these outliers.