SlowGuess/NeurIPS_2023_OCR huggingface dataset:
and collect them into an embedding vector  $z^y \in \mathbb{R}^{|\mathcal{T}|}$ . If  $|\mathcal{I}_1^y| < n_\alpha$ , where  $n_\alpha = (1 / \min\{\alpha, 0.1\}) - 1$ , then the uppermost quantile in  $z^y$  will not be finite, so we simply assign  $y$  to the null cluster. For a pre-chosen number of clusters  $M$ , we run  $k$ -means clustering with  $k = M$  on the data  $\{z^y\}_{y \in \mathcal{Y} \setminus \mathcal{Y}_{\mathrm{null}}}$ , where  $\mathcal{Y}_{\mathrm{null}}$  denotes the set of labels assigned to the null cluster. More specifically, we use a weighted version of  $k$ -means where the weight for class  $y$  is set to  $| \mathcal{I}_1^y |^{1/2}$ ; this allows the class embeddings computed from more data to have more influence on the cluster centroids. We denote the cluster mapping that results from this procedure by  $\hat{h}: \mathcal{Y} \to \{1, \dots, M\} \cup \{\mathrm{null}\}$ .
Of course, we cannot generally recover an oracle clustering function  $h^*$  with the above (or any practical) procedure, so Proposition 2 does not apply. However, if the score distributions for the classes that  $\tilde{h}$  assigns to the same cluster are similar enough, then we can provide an approximate class-conditional coverage guarantee. We measure similarity in terms of the Kolmogorov-Smirnov (KS) distance, which is defined between random variables  $X$  and  $Y$  as  $D_{\mathrm{KS}}(X,Y) = \sup_{\lambda \in \mathbb{R}} |\mathbb{P}(X \leq \lambda) - \mathbb{P}(Y \leq \lambda)|$ .
Proposition 3. Let  $S^y$  denote a random variable sampled from the score distribution for class  $y$ , and assume that the clustering map  $\hat{h}$  satisfies
$$
D _ {\mathrm {K S}} \left(S ^ {y}, S ^ {y ^ {\prime}}\right) \leq \epsilon , \quad f o r a l l y, y ^ {\prime} s u c h t h a t \hat {h} (y) = \hat {h} \left(y ^ {\prime}\right) \neq \mathsf {n u l l}.
$$
Then, for  $\mathcal{C} = \mathcal{C}_{\mathrm{CLUSTERED}}$  and for all classes  $y$  such that such that  $\hat{h}(y) \neq \mathrm{null}$ ,
$$
\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 - \epsilon .
$$
To summarize, the two main tuning parameters used in the proposed method are  $\gamma \in [0,1]$ , the fraction of points to use for the clustering step, and  $M \geq 1$ , the number of clusters to use in  $k$ -means. While there are no universal fixed values of these parameters that serve all applications equally well, we find that simple heuristics for setting these parameters often work well in practice, which we describe in the next section.
We evaluate the class-conditional coverage of STANDARD, CLASSWISE, and CLUSTERED conformal prediction on four large-scale image data sets using three conformal score functions on each. Code for reproducing our experiments is available at https://github.com/tiffanyding/class-conditional-conformal.
3.1 Experimental setup
We run experiments on the ImageNet (Russakovsky et al., 2015), CIFAR-100 (Krizhevsky, 2009), Places365 (Zhou et al., 2018), and iNaturalist (Van Horn et al., 2018) image classification data sets, whose characteristics are summarized in Table 1. The first three have roughly balanced classes, the fourth, iNaturalist, has high class imbalance. We consider three conformal score functions: softmax, one minus the softmax output of the base classifier; APS, a score designed to improve  $X$ -conditional coverage; and RAPS, a regularized version of APS that often produces smaller sets. Precise definitions of the score functions are provided in Appendix B.1; we refer also to Romano et al. (2020b); Angelopoulos et al. (2021) for the details and motivation behind APS and RAPS. Throughout, we set  $\alpha = 0.1$  for a desired coverage level of  $90\%$ .
Table 1: Description of data sets. The class balance metric is described precisely in Appendix B.3.
Data set	ImageNet	CIFAR-100	Places365	iNaturalist
Number of classes	1000	100	365	663*
Class balance	0.79	0.90	0.77	0.12
Example classes	mitten	orchid	beach	salamander
triceratops	forest	sushi bar	legume
guacamole	bicycle	catacomb	common fern
*The number of classes in the iNaturalist data set can be adjusted by selecting which taxonomy level (e.g., species, genus, family) to use as the class labels. We use the species family as our label and then filter out any classes with  $< {250}$  examples in order to have sufficient examples to properly perform evaluation.
Our experiments all follow a common template. First, we fine-tune a pre-trained neural network as our base classifier (for details on model architectures, see Appendix B.2) on a small subset  $D_{\mathrm{fine}}$  of the original data, leaving the rest for calibration and validation purposes. We construct calibration sets of varying size by changing the average number of points in each class, denoted  $n_{\mathrm{avg}}$ . For each  $n_{\mathrm{avg}} \in \{10, 20, 30, 40, 50, 75, 100, 150\}$ , we construct a calibration set  $D_{\mathrm{cal}}$  by sampling  $n_{\mathrm{avg}} \times |\mathcal{Y}|$  examples without replacement from the remaining data  $D_{\mathrm{fine}}^{c}$  (where  $c$  denotes the set complement). We estimate the conformal quantiles for STANDARD, CLASSWISE, and CLUSTERED on  $D_{\mathrm{cal}}$ . The remaining data  $(D_{\mathrm{fine}} \cup D_{\mathrm{cal}})^{c}$  is used as the validation set for computing coverage and set size metrics. Finally, this process—splitting  $D_{\mathrm{fine}}^{c}$  into random calibration and validation sets—is repeated ten times, and the reported metrics are averaged over these repetitions.
Details about clustering. For CLUSTERED, we choose  $\gamma \in [0,1]$  (the fraction of calibration data points used for clustering) and  $M \geq 1$  (the number of clusters) in the following way. First, we define  $n_{\mathrm{min}} = \min_{y \in \mathcal{Y}} |\mathcal{I}^y|$ , the number of examples in the rarest class in the calibration set, and  $n_\alpha = (1 / \alpha) - 1$ , the minimum sample size needed so that the finite-sample adjusted  $(1 - \alpha)$ -quantile used in conformal prediction is finite (e.g.,  $n_\alpha = 9$  when  $\alpha = 0.1$ ). Now define  $\tilde{n} = \max(n_{\mathrm{min}}, n_\alpha)$  and let  $K$  be the number of classes with at least  $\tilde{n}$  examples. We then set  $\gamma = K / (75 + K)$  and  $M = \lfloor \gamma \tilde{n} / 2 \rfloor$ . These choices are motivated by two goals: we want  $M$  and  $\gamma$  to grow together (to find more clusters, we need more samples for clustering), and we want the proper calibration set to have at least 150 points per cluster on average; see Appendix B.4 for details on how the latter is achieved. Clustering is carried out by running  $k$ -means on the quantile-based embeddings, as described in Section 2.2; we use the implementation in sklearn.cluster.KMeans, with the default settings (Pedregosa et al., 2011).
3.2 Evaluation metrics
Denote the validation data set by  $\{(X_i', Y_i')\}_{i=1}^{N'}$  (recall this is separate from the fine-tuning and calibration data sets) and let  $\mathcal{J}^y = \{i \in [N'] : Y_i' = y\}$  be the indices of validation examples with label  $y$ . For any given conformal method, we define  $\hat{c}_y = \frac{1}{|\mathcal{J}^y|} \sum_{i \in \mathcal{J}^y} \mathbb{1}\{Y_i' \in \mathcal{C}(X_i')\}$  as the empirical class-conditional coverage of class  $y$ . Our main evaluation metric is the average class coverage gap (CovGap):
$$
\operatorname {C o v G a p} = 1 0 0 \times \frac {1}{| \mathcal {Y} |} \sum_ {y \in \mathcal {Y}} | \hat {c} _ {y} - (1 - \alpha) |.
$$
This measures how far the class-conditional coverage is from the desired coverage level of  $1 - \alpha$ , in terms of the  $\ell_1$  distance across all classes (multiplied by 100 to put it on the percentage scale). We also measure the sharpness of the prediction sets by average set size (AvgSize):
$$
\operatorname {A v g S i z e} = \frac {1}{N ^ {\prime}} \sum_ {i = 1} ^ {N ^ {\prime}} | \mathcal {C} \left(X _ {i} ^ {\prime}\right) |.
$$
For a given class-conditional coverage level (a given CovGap), we want a smaller average set size.
To begin, we investigate the CovGap of the methods on each data set, and display the results in Figure 2. In brief, CLUSTERED achieves the best or comparable performance across all settings. Restricting our attention to the baseline methods, note that the CovGap of STANDARD does not change much as we vary  $n_{\mathrm{avg}}$ ; this is as expected, because  $n_{\mathrm{avg}} \times |\mathcal{Y}|$  samples are being used to estimate the conformal quantile  $\hat{q}$ , which will be stable regardless of  $n_{\mathrm{avg}}$ , provided  $|\mathcal{Y}|$  is large. Conversely, the CovGap of CLASSWISE decreases significantly as  $n_{\mathrm{avg}}$  increases, because the classwise conformal quantiles  $\hat{q}^y$  are volatile for small  $n_{\mathrm{avg}}$ .
For the softmax score (left column of the figure), we see that CLUSTERED clearly outperforms STANDARD, and the gap widens with  $n_{\mathrm{avg}}$ ; further, CLUSTERED outperforms CLASSWISE for small enough values of  $n_{\mathrm{avg}}$  (that is,  $n_{\mathrm{avg}} < 75$  for ImageNet and Places365,  $n_{\mathrm{avg}} < 150$  for CIFAR-100, and  $n_{\mathrm{avg}} < 50$  for iNaturalist). The comparison between CLUSTERED and CLASSWISE is qualitatively similar for APS score (right column of the figure), with the former outperforming the
Figure 2: Average class coverage gap for ImageNet, CIFAR-100, Places365, and iNaturalist, for the softmax (left) and APS (right) scores, as we vary the average number of calibration examples per class. The shaded regions denote  $\pm 1.96$  times the standard errors (often, the standard errors are too small to be visible).
latter for small enough values of  $n_{\mathrm{avg}}$ . However, the behavior of STANDARD changes notably as we move from softmax to APS: it becomes comparable to CLUSTERED and only slightly worse for large  $n_{\mathrm{avg}}$ . Lastly, the results for RAPS (not shown, and deferred to Appendix C.1) are similar to APS but the CovGap is shifted slightly higher.
To examine the potential tradeoffs between class-conditional coverage and average set size, we focus on the  $n_{\mathrm{avg}} \in \{10, 20, 50, 75\}$  settings and report CovGap and AvgSize in Table 2 for all data sets and score functions. We see that CLUSTERED achieves the best or near-best CovGap in any experimental combination, and its improvement over the baselines is particularly notable in the regime where data is limited, but not extremely limited ( $n_{\mathrm{avg}} = 20$  or 50). Meanwhile, AvgSize for STANDARD and CLUSTERED is relatively stable across values of  $n_{\mathrm{avg}}$ , with the latter being generally slightly larger; AvgSize for CLASSWISE decreases as we increase  $n_{\mathrm{avg}}$ , but is still quite large relative to STANDARD and CLUSTERED, especially for iNaturalist.
We finish with two more remarks. First, we find CovGap tends to behave quite similarly to the fraction of classes that are drastically undercovered, which we define as having a class-conditional coverage at least  $10\%$  below the desired level. Results for this metric are given in Appendix C.2. Second, as is the case in any conformal method, auxiliary randomization can be applied to STANDARD, CLASSWISE, or CLUSTERED in order to achieve a coverage guarantee (marginal, class-conditional, or cluster-conditional, respectively) of exactly  $1 - \alpha$ , rather than at least  $1 - \alpha$ . These results are included in Appendix C.3. In terms of CovGap, we find that randomization generally improves CLASSWISE, and does not change STANDARD or CLUSTERED much at all; however, the set sizes from randomized CLASSWISE are still overall too large to be practically useful (moreover, injecting auxiliary randomness is prediction sets is not always practically palatable).
We summarize our practical takeaways, in an effort towards creating guidelines for answering the question: for a given problem setting, what is the best way to produce prediction sets that have good class-conditional coverage but are not too large to be useful?
- Extremely low-data regime. When most classes have very few calibration examples (say, less than 10), this is not enough data to run CLUSTERED or CLASSWISE unless  $\alpha$  is large, so the only option is to run STANDARD. With this method, softmax and RAPS are both good score functions. RAPS tends to yield better class-conditional coverage, while softmax tends to have smaller sets.
- Low-data regime. When the average number of examples per class is low but not tiny (say, around 20 to 75), CLUSTERED tends to strike a good balance between variance and granularity and often achieves good class-conditional coverage and reasonably-sized prediction sets with either softmax or RAPS. The STANDARD method with RAPS is also competitive towards the lower-data end of this regime.
- High-data regime. When the average number of examples per class is large (say, over 75), CLUSTERED conformal with either softmax or RAPS continues to do well, and CLASSWISE with these same scores can also do well if the classes are balanced. In settings with high class imbalance, CLASSWISE is unstable for rare classes and produces excessively large prediction sets, whereas CLUSTERED is more robust due to the data-sharing it employs.
- Extremely high-data regime. When the calibration dataset is so large that the rarest class has at least, say, 100 examples, then CLASSWISE with softmax or RAPS will be a good choice, regardless of any class imbalance.
The boundaries between these regimes are not universal and are dependent on the data characteristics, and the above guidelines are based only on our findings from our experiments. We also note that the precise boundary between the low-data and high-data regimes is dependent on the particular score function that is used: the transition happens around 20-40 examples per class for softmax and 50-100 for RAPS. This serves as further motivation for CLUSTERED, which performs relatively well in all regimes.
As a possible direction for future work, it might be interesting to generalize our approach to the broader problem of group-conditional coverage with many groups. In the present setting, the groups
are defined by class labels, but our clustering methodology could also be applied to other group structures (e.g., defined by the input features or components of a mixture distribution).
We thank Margaux Zaffran and Long Nguyen for helpful suggestions that improved our paper. This work was supported by the National Science Foundation (NSF) Graduate Research Fellowship Program under grant no. 2146752, the European Research Council (ERC) Synergy Grant Program, the Office of Naval Research (ONR) the Mathematical Data Science Program, and the ONR Multi-University Research Initiative (MURI) Program under grant no. N00014-20-1-2787.
Anastasios N. Angelopoulos and Stephen Bates. A gentle introduction to conformal prediction and distribution-free uncertainty quantification. Foundations and Trends in Machine Learning, 16(4): 494-591, 2023.
Anastasios N. Angelopoulos, Stephen Bates, Jitendra Malik, and Michael I. Jordan. Uncertainty sets for image classifiers using conformal prediction. In Proceedings of the International Conference on Learning Representations, 2021.
Maurizio Sadinle, Jing Lei, and Larry Wasserman. Least ambiguous set-valued classifiers with bounded error levels. Journal of the American Statistical Association, 114(525):223-234, 2019.
Fan Shi, Cheng Soon Ong, and Christopher Leckie. Applications of class-conditional conformal predictor in multi-class classification. In International Conference on Machine Learning and Applications, 2013.
Jiangming Sun, Lars Carlsson, Ernst Ahlberg, Ulf Norinder, Ola Engkvist, and Hongming Chen. Applying Mondrian cross-conformal prediction to estimate prediction confidence on large imbalanced bioactivity data sets. Journal of Chemical Information and Modeling, 57(7):1591-1598, 2017.
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The iNaturalist species classification and detection dataset. In IEEE Conference on Computer Vision and Pattern Recognition, 2018.
Vladimir Vovk. Conditional validity of inductive conformal predictors. In *Asian Conference on Machine Learning*, 2012.
Vladimir Vovk, Alex Gammerman, and Glenn Shafer. Algorithmic Learning in a Random World. Springer, 2005.
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(6):1452-1464, 2018.
Proof of Proposition 1. For each  $m \in \{1, \ldots, M\}$ , denote by  $G^{(m)}$  the distribution of the score  $s(X,Y)$  conditioned on  $Y$  being in cluster  $m$ . Consider a randomly sampled test example  $(X_{\mathrm{test}}, Y_{\mathrm{test}})$  with a label in cluster  $m$ ; the test score  $s_{\mathrm{test}} = s(X_{\mathrm{test}}, Y_{\mathrm{test}})$  then follows distribution  $G^{(m)}$ . Next consider  $\{s_i\}_{i \in \mathcal{I}_2(m)}$ , the scores for examples in the proper calibration set with labels in cluster  $m$ ; these also follow distribution  $G^{(m)}$ . Furthermore,  $s_{\mathrm{test}}$  and the elements of  $\{s_i\}_{i \in \mathcal{I}_2(m)}$  are all mutually independent, so the result follows by the standard coverage guarantee for conformal prediction with exchangeable scores (e.g., see Vovk et al. (2005) or Lei et al. (2018)).
Proof of Proposition 2. This is a direct result of exchangeability and Proposition 1.
Proof of Proposition 3. Let  $S = s(X,Y)$  denote the score of a randomly sampled example  $(X,Y) \sim F$ . Fix an  $m \in \{1,\dots,M\}$ . Define  $\mathcal{Y}^{(m)} = \{y \in \mathcal{Y} : \hat{h}(y) = m\}$  as the set of classes that  $\hat{h}$  assigns to cluster  $m$ . Without a loss of generality, we treat both  $\hat{h}$  and  $\hat{q}(m)$  as fixed for the remainder of this proof. This can be done by conditioning on both the clustering and proper calibration sets, leaving only the randomness in the test point  $(X,Y) \sim F$ , and then integrating over the clustering and proper calibration sets in the end.
Let  $G^{(m)}$  denote the distribution of  $S$  conditional on  $Y \in \mathcal{Y}^{(m)}$ , and let  $S^{(m)} \sim G^{(m)}$ . Similarly, let  $G^{y}$  denote the distribution of  $S$  conditional on  $Y = y$ , and let  $S^{y} \sim G^{y}$ . Since we assume that the KS distance between the score distribution for every pair of classes in cluster  $m$  is bounded by  $\epsilon$ , and  $G^{(m)}$  is a mixture of these distributions (that is,  $G^{(m)} = \sum_{y \in \mathcal{Y}^{(m)}} \pi_y \cdot G^{y}$  for some fixed probability weights  $\pi_y$ ,  $y \in \mathcal{Y}^{(m)}$ ), it follows by the triangle inequality that
$$
D _ {\mathrm {K S}} \left(S ^ {y}, S ^ {(m)}\right) \leq \epsilon , \quad \text {f o r a l l} y \in \mathcal {Y} ^ {(m)}.
$$
By definition of KS distance, this implies
$$
\left| \mathbb {P} (S \leq \hat {q} (m) \mid Y = y) - \mathbb {P} (S \leq \hat {q} (m) \mid Y \in \mathcal {Y} ^ {(m)}) \right| \leq \epsilon .
$$
Since the CLUSTERED procedure includes the true label  $Y$  from the prediction set  $C$  when  $S \leq \hat{q}(m)$ , these probabilities can be rewritten in terms of coverage events:
$$
\left| \mathbb {P} (Y \in \mathcal {C} (X) \mid Y = y) - \mathbb {P} (Y \in \mathcal {C} (X) \mid Y \in \mathcal {Y} ^ {(m)}) \right| \leq \epsilon .
$$
Combining the result from Proposition 1 gives the desired conclusion.
B Experiment details
B.1 Score functions
We perform experiments using three score functions:
- softmax: The softmax-based conformal score at an input  $x$  and a label  $y$  is defined as
$$
s _ {\text {s o f t m a x}} (x, y) = 1 - f _ {y} (x),
$$
where  $f_{y}(x)$  is entry  $y$  of the softmax vector output by the classifier  $f$  at input  $x$ .
- APS: The Adaptive Prediction Sets (APS) score of Romano et al. (2020b) is designed to improve  $X$ -conditional coverage as compared to the more traditional softmax score. This score is computed at an input  $x$  and label  $y$  as follows. Let
$$
f _ {(1)} (x) \leq f _ {(2)} (x) \leq \dots \leq f _ {(| \mathcal {Y} |)} (x)
$$
denote the sorted values of the base classifier softmax outputs  $f_{y}(x), y \in \mathcal{V}$ . Let  $k_{x}(y)$  be the index in the sorted order that corresponds to class  $y$ , that is,  $f_{(k_x(y))} = f_y(x)$ . The APS score is then defined as
$$
s _ {\mathsf {A P S}} (x, y) = \sum_ {i = 1} ^ {k _ {x} (y) - 1} f _ {(i)} (x) + \operatorname {U n i f} ([ 0, f _ {(k _ {x} (y))} (x) ]).
$$
- RAPS: The regularized APS (RAPS) score of Angelopoulos et al. (2021) is a modification of the APS score that adds a regularization term designed to reduce the prediction set sizes (which can often be very large with APS). The RAPS score is defined as
$$
s _ {\mathsf {R A P S}} (x, y) = s _ {\mathsf {A P S}} (x, y) + \max  \left(0, \lambda \left(k _ {x} (y) - k _ {\mathrm {r e g}}\right)\right),
$$
where  $k_{x}(y)$  is as defined above, and  $\lambda$  and  $k_{\mathrm{reg}}$  are user-chosen parameters. In all of our experiments, we use  $\lambda = 0.01$  and  $k_{\mathrm{reg}} = 5$ , which Angelopoulos et al. (2021) found to work well for ImageNet.
B.2 Model training
An important consideration when we fine-tune and calibrate our models is that we must reserve sufficient data to evaluate the class-conditional coverage of the conformal methods. This means we aim to exclude at least 250 examples per class from the fine-tuning and calibration sets so that we can then use this data for validation (applying the conformal methods and computing coverage and set size metrics).
For all data sets except ImageNet, we use a ResNet-50 model as our base classifier. We initialize to the IMAGENET1K_V2 pre-trained weights from PyTorch (Paszke et al., 2019), and then fine-tune all parameters by training on the data set at hand. For ImageNet, we must do something different, as explained below.
ImageNet. Setting up this data set for our experiments is a bit tricky because we need sufficient data for performing validation, but we also need this data to be separate from the fine-tuning and calibration sets. The ImageNet validation set only contains 50 examples per class, which is not enough for our setting. The ImageNet training set is much larger, with roughly 1000 examples per class, but if we want to use this data for validation, then we cannot use the ResNet-50 initialized to theImagenet1K_V2 pre-trained weights, as these weights were obtained by training on the whole ImageNet training set. We therefore instead use a SimCLR-v2 model (Chen et al., 2020), which was trained on the ImageNet training set without labels, to extract feature vectors of length 6144 for all images in the ImageNet training set. We then use  $10\%$  of these feature vectors for training a linear head (a single fully connected neural network layer). After training for 10 epochs, the model achieves a validation accuracy of  $78\%$ . We then apply the linear head to the remaining  $90\%$  of the feature vectors to obtain softmax scores for the calibration set.
CIFAR-100. This data set has 600 images per class (500 from the training set and 100 from the validation set). We combine the data and then randomly sample  $50\%$  for fine-tuning, and we use the remaining data for calibrating and validating our procedures. After training for 30 epochs, the validation accuracy is  $60\%$ .
Places365. This data set contains more than 10 million images of 365 classes, where each class has 5000 to 30000 examples. We randomly sample  $90\%$  of the data for fine-tuning, and we use the remaining data for calibrating and validating our procedures. After training for one epoch, the validation accuracy is  $52\%$ .