SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tiffany Ding
University of California, Berkeley  
tiffany_ding@berkeley.edu
Anastasios N. Angelopoulos
University of California, Berkeley angelopoulos@berkeley.edu
Stephen Bates
MIT s_bates@mit.edu
Michael I. Jordan
University of California, Berkeley  
jordan@cs.berkeley.edu
Ryan J. Tibshirani
University of California, Berkeley  
ryantibs@berkeley.edu
Standard conformal prediction methods provide a marginal coverage guarantee, which means that for a random test point, the conformal prediction set contains the true label with a user-specified probability. In many classification problems, we would like to obtain a stronger guarantee—that for test points of a specific class, the prediction set contains the true label with the same user-chosen probability. For the latter goal, existing conformal prediction methods do not work well when there is a limited amount of labeled data per class, as is often the case in real applications where the number of classes is large. We propose a method called clustered conformal prediction that clusters together classes having "similar" conformal scores and performs conformal prediction at the cluster level. Based on empirical evaluation across four image data sets with many (up to 1000) classes, we find that clustered conformal typically outperforms existing methods in terms of class-conditional coverage and set size metrics.
Consider a situation in which a doctor relies on a machine learning system that has been trained to output a set of likely medical diagnoses based on CT images of the head. Suppose that the system performs well on average and is able to produce prediction sets that contain the actual diagnosis with at least 0.9 probability. Upon closer examination, however, the doctor discovers the algorithm only predicts sets containing common and relatively benign conditions, such as {normal, concussion}, and the sets never include less common but potentially fatal diseases, such as intracranial hemorrhage. In this case, despite its high marginal (average) performance, the doctor would not want to use such an algorithm because it may lead to patients missing out on receiving critical care. The core problem is that even though the average performance of the algorithm is good, the performance for some classes is quite poor.
Conformal prediction (Vovk et al., 2005) is a method for producing set-valued predictions that serves as a wrapper around existing prediction systems, such as neural networks. Standard conformal prediction takes a black-box prediction model, a calibration data set, and a new test example  $X_{\mathrm{test}} \in \mathcal{X}$  with unknown label  $Y_{\mathrm{test}} \in \mathcal{Y}$  and creates a prediction set  $\mathcal{C}(X_{\mathrm{test}}) \subseteq \mathcal{Y}$  that satisfies marginal coverage:
$$
\mathbb {P} \left(Y _ {\text {t e s t}} \in \mathcal {C} \left(X _ {\text {t e s t}}\right)\right) \geq 1 - \alpha , \tag {1}
$$
for a coverage level  $\alpha \in [0,1]$  chosen by the user. However, as the example above shows, the utility of these prediction sets can be limited in some real applications. In classification, which we study in
this paper, the label space  $\mathcal{V}$  is discrete, and it is often desirable to have class-conditional coverage:
$$
\mathbb {P} \left(Y _ {\text {t e s t}} \in \mathcal {C} \left(X _ {\text {t e s t}}\right) \mid Y _ {\text {t e s t}} = y\right) \geq 1 - \alpha , \quad \text {f o r a l l} y \in \mathcal {Y}, \tag {2}
$$
meaning that every class  $y$  has at least  $1 - \alpha$  probability of being included in the prediction set when it is the true label. Note that (2) implies (1). Predictions sets that only satisfy (1) may neglect the coverage of some classes, whereas the predictions sets in (2) are effectively "fair" with respect to all classes, even the less common ones.
Standard conformal prediction, which we will refer to as STANDARD, does not generally provide class-conditional coverage. We present a brief case study to illustrate.
ImageNet case study. Running STANDARD using a nominal coverage level of  $90\%$  on 50,000 examples sampled randomly from ImageNet (Russakovsky et al., 2015), a large-scale image data set described later in Section 3, yields prediction sets that achieve very close to the correct marginal coverage  $(89.8\%)$ . However, this marginal coverage is achieved by substantially undercovering some classes and overcovering others. For example, water jug is severely uncovered: the prediction sets only include it in  $50.8\%$  of the cases where it is the true label. On the other hand, flamingo is substantially uncovered: it achieves a class-conditional coverage of  $99.2\%$ . This underscores the need for more refined methods in order to achieve the class-conditional coverage defined in (2).
In principle, it is possible to achieve (2) by splitting the calibration data by class and running conformal prediction once for each class (Vovk, 2012). We refer to this as CLASSWISE. However, this procedure fails to be useful in many real applications since data-splitting can result in very few calibration examples for each class-level conformal prediction procedure. This typically happens in problem settings where we have many classes but limited data; in such settings, the classwise procedure tends to be overly conservative and produces prediction sets that are too large to be practically useful. We note that previous papers (Shi et al., 2013; Lofstrom et al., 2015) on class-conditional conformal have not focused on the many classes regime and have instead studied binary or at most 10-way classification tasks.
In this work, we focus on the challenging limited-data, many-class classification setting and we develop a method targeted at class-conditional coverage. Our method mitigates the problems that arise from data-splitting by clustering together classes that have similar score distributions and combining the calibration data for those classes. Figure 1 illustrates how the method we propose strikes a balance between STANDARD and CLASSWISE. As we will later show in our experiments, this can improve class-conditional coverage in many situations.
Low variance but insufficient granularity
Clustered conformal prediction
Sufficient granularity but high variance
Figure 1: A schematic comparison of conformal prediction methods, including the CLUSTERED method we propose. Each colored circle represents the calibration data for a particular class. Existing methods fall on extremes of the spectrum: STANDARD is very stable because it groups data for all classes together, but it is not able to treat classes differently when needed; CLASSWISE has high granularity, but it splits all of the data by class and consequently has high variance in limited-data settings; CLUSTERED strikes a balance by grouping together data for "similar" classes.
Our contributions. We make three advances in tackling the problem of class-conditional coverage.
- We propose an extension of conformal prediction called clustered conformal prediction that often outperforms standard and classwise conformal, in terms of class-conditional coverage, when there is limited calibration data available per class.
- We present a comprehensive empirical evaluation of class-conditional coverage for conformal methods on four large-scale classification data sets, each with many (100 to 1000) classes.
- We provide general guidelines to practitioners for how to choose an appropriate conformal prediction method in order to achieve class-conditional coverage in the problem setting at hand.
Mondrian conformal prediction. We work in the split-conformal prediction framework (Papadopoulos et al., 2002; Lei et al., 2018) in which the training data set (to train the base classifier) and calibration data set are disjoint. Mondrian conformal prediction (MCP) is a general procedure that encompasses many kinds of conditional conformal prediction (Vovk et al., 2005). For any chosen grouping function  $G: \mathcal{X} \times \mathcal{Y} \rightarrow \mathcal{G}$  where  $\mathcal{G}$  denotes the set of all groups, MCP provides coverage guarantees of the form  $\mathbb{P}(Y_{\mathrm{test}} \in \mathcal{C}(X_{\mathrm{test}}) \mid G(X_{\mathrm{test}}, Y_{\mathrm{test}}) = g) \geq 1 - \alpha$  for all groups  $g \in \mathcal{G}$ . The high-level idea behind MCP is to split the calibration data by group and then run conformal prediction on each group. The CLASSWISE conformal procedure is a special case of MCP where each value of  $\mathcal{Y}$  defines a group (this is also sometimes referred to as label-conditional conformal prediction).
To the best of our knowledge, previous work has not focused on class-conditional coverage of conformal methods in the many-classes, limited-data regime that is common to many real classification applications. Lofstrom et al. (2015) present an empirical study of STANDARD and CLASSWISE on binary, 3-way, and 4-way classification data sets (with a focus on the binary setting). For class-imbalanced problems, they find that STANDARD tends to overcover the majority class, and undercover the minority class. Shi et al. (2013) consider CLASSWISE on the MNIST and USPS data sets, which are 10-way classification problems. Sun et al. (2017) use a cross-fold variant of CLASSWISE for binary classification on imbalanced bioactivity data. Hechtlinger et al. (2018) run class-conditional experiments on the 3-class Iris data set, and provide preliminary ideas on how to incorporate interactions between classes using density estimation. Guan and Tibshirani (2022) use CLASSWISE in an outlier detection context. Sadinle et al. (2019) consider CLASSWISE with a modification to avoid empty prediction sets and perform experiments on 3-way or 10-way classification tasks. To reiterate, the aforementioned papers all focus on the regime that is generally favorable to CLASSWISE, where data is abundant relative to the number of classes. Our work focuses on the limited-data regime that often arises in practice.
Other types of conditional conformal prediction. "Conditional coverage" is a term that, within the conformal prediction literature, often refers to coverage at a particular input value  $X = x$ . It is known to be impossible to achieve coverage conditional on  $X = x$ , for all  $x$ , without invoking further distributional assumptions (Lei and Wasserman, 2014; Vovk, 2012; Barber et al., 2021). However, approximate  $X$ -conditional coverage can be achieved in practice by designing better score functions (Romano et al., 2019, 2020b; Angelopoulos et al., 2022) or by modifying the conformal procedure itself (Romano et al., 2020a; Guan, 2023; Gibbs et al., 2023). Our work draws inspiration from the latter camp, but our aim is to achieve class-conditional ( $Y$ -conditional) coverage rather than  $X$ -conditional coverage. Note that  $X$ -conditional coverage is hard to interpret in settings in which  $Y$  is not random given  $\bar{X} = x$  (e.g., in image classification, if  $x$  is an image of a dog, then the true label is always  $y = \mathrm{dog}$ , with no intrinsic randomness after conditioning on  $X = x$ ). In such settings, it is more natural to consider class-conditional coverage.
We work in the classification setting in which each input  $X_{i}\in \mathcal{X}$  has a class label  $Y_{i}\in \mathcal{V}$ , for some discrete set  $\mathcal{V}$ . Let  $\{(X_i,Y_i)\}_{i = 1}^N$  denote a calibration data set, where each  $(X_{i},Y_{i})\stackrel {\mathrm{~id~}}{\sim}F$ . Given a new independent test point  $(X_{\mathrm{test}},Y_{\mathrm{test}})\sim F$ , our goal is to construct (without knowledge of  $Y_{\mathrm{test}}$ ) a prediction set  $\mathcal{C}(X_{\mathrm{test}})$  that satisfies (2) for some user-specified  $\alpha \in [0,1]$ .
Let  $s: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$  denote a (conformal) score function, where we take  $s(x,y)$  to be negatively oriented, which means that lower scores indicate a better agreement between the input  $x$  and the proposed class label  $y$ . The score function is typically derived from a pre-trained classifier  $f$  (a simple example to keep in mind is  $s(x,y) = 1 - f_y(x)$ , where  $f_{y}(x)$  represents the  $y^{\mathrm{th}}$  entry of the softmax vector output by  $f$  for the input  $x$ ). For brevity, we denote the score of the  $i^{\mathrm{th}}$  calibration data point as  $s_i = s(X_i,Y_i)$ . For  $\tau \in [0,1]$  and a finite set  $A \subseteq \mathbb{R}$ , let Quantile(  $\tau ,A$ ) denote the smallest  $a\in A$  such that  $\tau$  fraction of elements in  $A$  are less than or equal to  $a$ . For  $\tau >1$ , we take  $\operatorname {Quantile}(\tau ,A) = \infty$ .
With this notation, the STANDARD conformal prediction sets are given by
$$
\mathcal {C} _ {\text {S T A N D A R D}} \left(X _ {\text {t e s t}}\right) = \{y: s \left(X _ {\text {t e s t}}, y\right) \leq \hat {q} \},
$$
$$
\hat {q} = \operatorname {Q u a n t i l e} \left(\frac {\lceil (N + 1) (1 - \alpha) \rceil}{N}, \{s _ {i} \} _ {i = 1} ^ {N}\right).
$$
These prediction sets are guaranteed to satisfy (1) (Vovk et al., 2005). We can interpret  $\hat{q}$  as a finite-sample adjusted  $(1 - \alpha)$ -quantile of the scores in the calibration data set. Note that all  $N$  data points are used for computing a single number,  $\hat{q}$ .
In contrast, the CLASSWISE procedure computes a separate quantile for every class. Let  $\mathcal{I}^y = \{i\in [N]:Y_i = y\}$  be the indices of examples in the calibration data set that have label  $y$ . The CLASSWISE conformal prediction sets are given by
$$
\mathcal {C} _ {\text {C L A S S W I S E}} \left(X _ {\text {t e s t}}\right) = \left\{y: s \left(X _ {\text {t e s t}}, y\right) \leq \hat {q} ^ {y} \right\},
$$
$$
\hat {q} ^ {y} = \operatorname {Q u a n t i l e} \left(\frac {\lceil (| \mathcal {I} ^ {y} | + 1) (1 - \alpha) \rceil}{| \mathcal {I} ^ {y} |}, \{s _ {i} \} _ {i \in \mathcal {I} ^ {y}}\right).
$$
These prediction sets are guaranteed to satisfy (2) (Vovk, 2012). However, the quantile  $\hat{q}^y$  is only computed using a subset of the data of size  $|\mathcal{I}_y|$ , which might be quite small. Importantly, for any class  $y$  for which  $|\mathcal{I}^y| < (1 / \alpha) - 1$ , we will have  $\hat{q}^y = \infty$ , hence any prediction set generated by CLASSWISE will include  $y$ , no matter the values of the conformal scores.
The advantage of STANDARD is that we do not need to split up the calibration data to estimate  $\hat{q}$ , so the estimated quantile has little noise even in limited-data settings; however, it does not in general yield class-conditional coverage. Conversely, CLASSWISE is guaranteed to achieve class-conditional coverage, since we estimate a different threshold  $\hat{q}^y$  for every class; however, these estimated quantiles can be very noisy due to the limited data, leading to erratic behavior, such as large sets.
When we assume the scores are almost surely distinct, the exact distribution of the CLASSWISE class-conditional coverage of class  $y$  given a fixed calibration set is (Vovk, 2012; Angelopoulos and Bates, 2023):
$$
\mathbb {P} \big (Y _ {\mathrm {t e s t}} \in \mathcal {C} _ {\mathrm {C L A S S W I S E}} (X _ {\mathrm {t e s t}}) \mid Y _ {\mathrm {t e s t}} = y, \{(X _ {i}, Y _ {i}) \} _ {i = 1} ^ {N} \big) \sim \mathrm {B e t a} (k _ {\alpha} ^ {y}, | \mathcal {I} ^ {y} | + 1 - k _ {\alpha} ^ {y}),
$$
where  $k_{\alpha}^{y} = \lceil (|\mathcal{I}^{y}| + 1)(1 - \alpha)\rceil$  . Note that the probability on the left-hand side above is taken with respect to the test point  $X_{\mathrm{test}}$  as the only source of randomness, as we have conditioned on the calibration set. The beta distribution  $\mathrm{Beta}(k_{\alpha}^{y},|\mathcal{I}^{y}| + 1 - k_{\alpha}^{y})$  has mean  $k_{\alpha}^{y} / (|\mathcal{I}^{y}| + 1)$  (which, as expected, is always at least  $1 - \alpha$  ) and variance
$$
\frac {k _ {\alpha} ^ {y} (| \mathcal {I} ^ {y} | + 1 - k _ {\alpha} ^ {y})}{(| \mathcal {I} ^ {y} | + 1) ^ {2} (| \mathcal {I} ^ {y} | + 2)} \approx \frac {\alpha (1 - \alpha)}{| \mathcal {I} ^ {y} | + 2},
$$
which can be large when  $|\mathcal{I}^y|$  is small. For example, if class  $y$  only has 10 calibration examples and we seek  $90\%$  coverage, then the class-conditional coverage given a fixed calibration set is distributed as  $\mathrm{Beta}(10,1)$ , so there is probability  $\approx 0.107$  that the coverage of class  $y$  will be less than  $80\%$ . Somewhat paradoxically, the variance of the class-conditional coverage means that on a given realization of the calibration set, the CLASSWISE method can exhibit poor coverage on a substantial fraction of classes if the number of calibration data points per class is limited.
2 Clustered conformal prediction
With the goal of achieving the class-conditional coverage guarantee in (2), we introduce clustered conformal prediction. Our method strikes a balance between the granularity of CLASSWISE and the data-pooling of STANDARD by grouping together classes according to a clustering function. For each cluster, we calculate a single quantile based on all data in that cluster. We design the clustering algorithm so that clustered classes have similar score distributions, and therefore, the resulting cluster-level quantile is applicable to all classes in the cluster. Next, in Section 2.1, we formally describe the clustered conformal prediction method; then, in Section 2.2, we describe the clustering step in detail.
2.1 Meta-algorithm
To begin, we randomly split the calibration data set into two parts: the clustering data set  $D_{1} = \{(X_{i},Y_{i}):i\in \mathcal{I}_{1}\}$ , for performing clustering, and a proper calibration data set  $D_{2} = \{(X_{i},Y_{i}):i\in \mathcal{I}_{2}\}$ , for computing the conformal quantiles, where  $|\mathcal{I}_1| = \lfloor \gamma N\rfloor$  and  $|\mathcal{I}_2| = N - |\mathcal{I}_1|$  for some tuning parameter  $\gamma \in [0,1]$ . Then, we apply a clustering algorithm to  $D_{1}$  to obtain a clustering function  $\hat{h}:\mathcal{V}\to \{1,\dots ,M\} \cup \{\mathrm{null}\}$  that maps each class  $y\in \mathcal{V}$  to one of  $M$  clusters or the "null" cluster (denoted null). The reason that we include the null cluster is to handle rare classes that do not have sufficient data to be confidently clustered into any of the  $M$  clusters. Details on how to create  $\hat{h}$  are given in the next subsection.
After assigning classes to clusters, we perform the usual conformal calibration procedure within each cluster. Denote by  $\mathcal{I}_2^y = \{i\in \mathcal{I}_2:Y_i = y\}$  the indices of examples in  $D_{2}$  with label  $y$ , and by  $\mathcal{I}_2(m) = \cup_{y:\hat{h} (y) = m}\mathcal{I}_2^y$  the indices of examples in  $D_{2}$  with labels in cluster  $m$ . The CLUSTERED conformal prediction sets are given by
$$
\mathcal {C} _ {\text {C L U S T E R E D}} (X _ {\text {t e s t}}) = \{y: s (X _ {\text {t e s t}}, y) \leq \hat {q} (\hat {h} (y)) \},
$$
where, for  $m = 1,\ldots ,M$
$$
\hat {q} (m) = \text {Q u a n t i l e} \left(\frac {\lceil (\left| \mathcal {I} _ {2} (m) \right| + 1) (1 - \alpha) \rceil}{\left| \mathcal {I} _ {2} (m) \right|}, \left\{s _ {i} \right\} _ {i \in \mathcal {I} _ {2} (m)}\right),
$$
$$
\hat {q} (\mathsf {n u l l}) = \operatorname {Q u a n t i l e} \left(\frac {\lceil (| \mathcal {I} _ {2} | + 1) (1 - \alpha) \rceil}{| \mathcal {I} _ {2} |}, \{s _ {i} \} _ {i \in \mathcal {I} _ {2}}\right).
$$
In words, for each cluster  $m$ , we group together examples for all classes in that cluster and then estimate a cluster-level quantile  $\hat{q}(m)$ . When constructing prediction sets, we include class  $y$  in the set if the conformal score for class  $y$  is less than or equal to the quantile for the cluster that contains  $y$ . For classes assigned to the null cluster, we compare the conformal score against the quantile that would be obtained from running STANDARD on the proper calibration set.
We now consider the properties of the CLUSTERED prediction sets. For all classes that do not belong to the null cluster, we have the following guarantee (the proof of this result, and all other proofs, are deferred to Appendix A).
Proposition 1. The prediction sets  $\mathcal{C} = \mathcal{C}_{\mathrm{CLUSTERED}}$  from CLUSTERED achieve cluster-conditional coverage:
$$
\mathbb {P} \left(Y _ {\text {t e s t}} \in \mathcal {C} \left(X _ {\text {t e s t}}\right) \mid \hat {h} \left(Y _ {\text {t e s t}}\right) = m\right) \geq 1 - \alpha , \quad \text {f o r a l l c l u s t e r s} m = 1, \dots , M.
$$
This coverage result comes from the exchangeability between a test point drawn from cluster  $m$  and all of the calibration points that belong to cluster  $m$ . We get this exchangeability for free from the assumption that the calibration point and test points are sampled i.i.d. from the same distribution. Cluster-conditional coverage is a stronger guarantee than marginal coverage, but it is still not as strong as the class-conditional coverage property that we hope to achieve. However, cluster-conditional coverage implies class-conditional in an idealized setting: suppose we have access to an "oracle" clustering function  $h^*$  that produces  $M$  clusters such that, for each cluster  $m = 1, \ldots, M$ , the scores for all classes in this cluster are exchangeable (this would hold, for example, if all classes in the same cluster have the same score distribution). In this case, we have a guarantee on class-conditional coverage.
Proposition 2. If  $\hat{h} = h^*$ , the "oracle" clustering function as described above, then the prediction sets from CLUSTERED satisfy class-conditional coverage (2) for all classes  $y$  such that  $h^*(y) \neq \emptyset$ .
This coverage result arises because the oracle clustering function ensures exchangeability between a test point drawn from class  $y$  and all of the calibration points that belong to the cluster to which  $y$  belongs. To try to achieve this exchangeability, we need to construct the clusters carefully. Specifically, we want to design  $\hat{h}$  so that the scores within each cluster  $m = 1, \dots, M$  are as close to identically distributed as possible, an idea we pursue next.
2.2 Quantile-based clustering
We seek to cluster together classes that have similar score distributions. To do so, we first summarize the empirical score distribution for each class via a vector of score quantiles evaluated at a discrete set of levels  $\tau \in \mathcal{T} \subseteq [0,1]$ . In this embedding space, a larger distance between classes means their score distributions are more different. While there are various options for defining such an embedding, recall that to achieve class-conditional coverage, we want to accurately estimate the (finite-sample adjusted)  $(1 - \alpha)$ -quantile for the score distribution of each class. Thus, we want to group together classes with similar quantiles, which is what our embedding is designed to facilitate. After obtaining these embeddings, we can then simply apply any clustering algorithm, such as  $k$ -means.
In more detail, denote by  $\mathcal{I}_1^y = \{i\in \mathcal{I}_1:Y_i = y\}$  the indices of examples in  $D_{1}$  with label  $y$ . We compute quantiles of the scores  $\{s_i\}_{i\in \mathcal{T}_1^y}$  from class  $y$  at the levels
$$
\mathcal {T} = \left\{\frac {\lceil (| \mathcal {I} ^ {y} | + 1) \tau \rceil}{| \mathcal {I} ^ {y} |}: \tau \in \{0. 5, 0. 6, 0. 7, 0. 8, 0. 9 \} \cup \{1 - \alpha \} \right\}.
$$