Multi-Dimensional Conformal Prediction

Thank you for your helpful and constructive review.

On demonstration over other data types

Thank you for your suggestion. Following your comment, we added an experiment with a text classification task in Appendix D.4. We use the 20 Newsgroups dataset, which comprises newsgroup posts on $20$ topics. We use a BERT-base model, and attach additional classification heads in a similar way to the other models. The results are reported on Table D.7:

Similarly to the other dataset, we can see the benefit of the proposed method over the baselines.

On providing more insights into the computational cost? Specifically, how does it compare to other baseline methods?

Thank you for pointing this out. Our proposed self-ensemble model requires attaching multiple classification heads to the model. Let $P$ denote the number of parameters per head and $H$ denote the number of heads, than the model size is increased by $P\cdot(H-1)$ parameters. in practice, we used $2$ fully connected layers for each classification head. Our results indicate that there is already an improvement with only $2$ classification heads, and that it is usually unnecessary to add more than $4$ heads. Thus, the size of the model is not largely affected by this modification. In terms of inference time, this is exactly as using a standard model with a single head, as all heads can be computed in parallel.

It should be emphasized that the main idea in the paper is how we can use multiple nonconformity scores to improve conformal prediction efficiency, where we propose to perform the calibration in the high-dimensional space instead of averaging the scores and performing standard one-dimensional conformal prediction. This idea can be applied in the case we have multiple different scores, which can be obtained in different ways: standard ensemble of multiple different models, test-time augmentation where scores are computed for the output of different augmented versions of the input image, multiple types of scores computed from a single output (e.g. Thr, APS, RAPS, SAPS) and our proposed self-ensemble that extracts scores from multiple classification heads. Note that test time augmentation does not require any modifications to the base model, and a single head with different type of scores is the least expensive option. Our main results are shown for the self-ensemble approach, while we also provide results for the other cases.

On the sensitivity to hyperparameter choices, such as the number of classification heads or diversity regularization.

Sensitivity to the number of heads - We can see on Figs. 3, D.1, D.2 and D.3 that as the number of heads increases we obtain smaller set sizes. However, an improvement is already observed for two heads, and most of the time $4-5$ heads are enough to get the optimal results.

Sensitivity to the diversity regularization - Following your comment, we added an experiment analyzing the influence of the parameter controlling the weight of the diversity regularization in Appendix D.3.:

"We conducted an ablation study to examine the affect of adding diversity regularization to the loss used for training the classification heads, defined in Eq. (12). We use the same fine-tuned model in the first training stage and change only the second stage to optimize the regularized loss in Eq. (12) with different values of $\lambda$, controlling the strength of the diversity regularization. Figure D.10. shows the similarity between heads with and without the diversity regularization. As expected, the similarity between heads is smaller for the model trained with the regularized loss. Figure D.11. compares the set sizes obtained for different values of $\lambda$, demonstrating that adding the regularization results in smaller sets. The optimal value of $\lambda$ is around $0.8$, after which an increase in set size is observed, apparently due to the fact that increasing head diversity comes at the cost of decreasing the accuracy of each head."

Dear reviewer 9Lus,

Thank you very much again for your helpful and constructive feedback. Your insights have been invaluable in refining our work, and we are grateful for the time and effort you have dedicated to reviewing our manuscript.

As the review process deadline approaches, we kindly request that you revisit our submission in light of the revisions and corrections we have made in response to your valuable feedback. We have carefully addressed the points you raised and believe these changes have strengthened the paper significantly.

If you find the revisions satisfactory, we would greatly appreciate it if you could consider raising your score to reflect the improvements. We would be more than happy to provide further clarifications if you have any additional concerns.

Best, Authors

Thank you for your helpful and constructive review.

On reporting worst slice coverage as a metric of adaptivity

We report the worst slice coverage (the size-stratified coverage violation (SSCV) metric defined in Eq. (18)) in Table D.1:

We can see that all methods achieve similar SSCV scores, with no method showing a clear advantage over the others.

Note that our cell score, defined by Eq. (8), puts an emphasis on the efficiency in terms of set size, however we could prioritize cells according to different metrics, for example using a score that takes into account some grouping of the input $x$ to encourage adaptivity. This is indeed an interesting direction for future work.

On using unnecessarily large calibration set

Thank you. We added an experiment comparing the performance with varying sizes of the calibration set in Appendix D.3. We vary the proportion $p$ of samples from the defined calibration split that are actually used. Then, the this $p$-th subset is split into half for cell computation, and re-calibration. Figure D.12 presents the set sizes for different values of $p$. It can be seen that the proposed method is almost always preferable, with its advantage becoming more pronounced as the size of the calibration data increases.

In addition, following the comment of reviewer 4GpN, we added a leave-one-out version of the proposed method that is more data efficient, as it does not require splitting the calibration data into two subsets. For more details, see Section 4.2. and Algorithm B.1.

On comparison with the normal best method where the tuning set is also added to the calibration set

Thank you for pointing this out. We agree that a fair comparison should include the normal method where the tuning set is also added to the calibration set, and this is indeed the baseline that we use. We added a clarification in the description of the baseline methods.

On reporting result with APS and threshold prediction sets

Thank you for pointing this out. Following your comment, we added experiments, evaluating our method over APS and threshold prediction sets in Appendix D.1:

“Our proposed method can be applied over any type of score. However, we have seen that it is not fully optimal for Thr and APS. The reason is that for these scores the values are more concentrated on specific levels, while for SAPS and RAPS the values are more spread. In order to improve this behavior we use temperature scaling, i.e. we divide the logits by $T$ before applying Softmax(). As $T$ increases the entropy of the probabilities increases and they become more spread, as illustrated in Fig. D.5. Figure D.6. presents the set sizes obtained for all methods with respect to different temperatures. We see that the efficiency the proposed method greatly improves as $T$ increases for both Thr and APS, while the performance of the baseline methods is less effected by $T$. For RAPS and SAPS the proposed method outperforms the baselines regardless of the temperature due to the inherent spread of these scores.”

On trying soft probability assignment to each region and optimize for that

Thank you for your suggestion. We added a new version of our method based on your proposal in Section 4.2:

"We present a more general variant of our approach, where instead of considering a hard assignment of each point into the closest center, we consider a softer assignment to $b$ nearest neighbors. This variant is summarized in Algorithm B.2. Here we detect the $b$ nearest neighbors, and include the point in the prediction set if at least half of its neighbors are in the chosen region."

We report the results in Table D.12:

We observe that in almost all cases $b=1$ (Algorithm 1) is preferable.

On an ablation study over the number of centroids.

We extended the ablation study over the number of centroids to show the influence of a very limited number of centroids (Appendix D. 3., Table D.3.):

As expected, the set size increases as $p$ decreases. However, the overall behavior remains stable, with only $9.4\\%-13.3\\%$ increase in set size for $50\\%$ reduction in the number of cells. In addition, we compare to the Optimized baseline, where the same set is used for optimizing the weights for combining the scores. We see for all $p$ values Multi-score method is advantageous over the Optimized baseline.

I think the authors misunderstood my question (Q1). Although I do not decrease my score, this is an important concern. Upon justification and a report on this metric, I will revisit the paper to see if I can increase the score again.

Here as a measure of conditional coverage, I referred to worst-slice coverage not set-size stratified coverage. For the worst slice coverage, I refer to [1] (section S.1.2 in the appendix). To quantify that, we create $k$ random vectors (tensors) with the dimensionality as the data. Then we sample $a, b$ randomly but among the sorted values of dot product results. This basically captures that on a random projection, within a bound what is the worst coverage. This metric captures if a local region in the distribution is miscovered (regardless of the set size).

Why is this metric important? The size-stratified coverage (worst coverage among sets of the same size) does not capture the information stored in conformal scores and how it quantifies uncertainty. Assume a dataset in which all points are aleatorically uncertain. Here small set sizes are misleading since if the score captures this uncertainty for every point all sets are large.

Another justification is to assume that our dataset has a minority subgroup for which the conformal prediction is miscovering. In that group the set size could be anything, and maybe reporting the coverage over sets of size 1, 2, ... has the same averaging effect that the marginal coverage has. Because the group is treated unfairly w.r.t. coverage not the average set size. While the worst slice coverage as defined in [1] is more likely to capture this notion.

The SAPS score function (Huang et al., 2024 in your references) also reports the set-size stratified coverage as a notion of conditional coverage quantification which is completely wrong as it does not incorporate any information regarding the location of $x$ in any data manifold.

Even if your method on the worst-slice coverage is imbalanced I would not decrease my score, but I expect it to be mentioned in the limitations.

[1] Classification with Valid and Adaptive Coverage, Yaniv Romano, Matteo Sesia, Emmanuel J. Candès, 2020.

On the relation of the proposed method multi-score setting vs CP for ensemble models

Thank you for pointing this out. The main idea in the paper is how we can use multiple nonconformity scores to improve CP efficiency, where we propose to perform the calibration in the high-dimensional space instead of averaging the scores and performing standard one-dimensional CP. This idea can be applied in the case we have multiple different scores, which can be obtained in different ways: standard ensemble of multiple different models, test-time augmentation where scores are computed for the output of different augmented versions of the input image, multiple types of scores computed from a single output (e.g. Thr, APS, RAPS, SAPS) and our proposed self-ensemble that extracts scores from multiple classification heads, trained with diversity-regularized loss. Our main results are shown for the self-ensemble approach, while we also provide results for the other two cases: test time augmentation and standard ensemble.

Following your comment, we added results for combining multiple types of scores. Specifically, we examined a setting where instead of considering multiple classification heads, we use a single head and compute different conformity scores: Thr, APS, RAPS and SAPS. Table D.2 summarizes the results for all datasets:

We see that combining multiple scores improves the results compared to the best single score, and that the proposed method obtains the smallest prediction sets in almost all cases. The advantage of this score fusion is that it does not require any additional modifications to the model or further fine-tuning iterations.

On the extensiveness of the empirical comparison, and including ImageNet in the experiments

Following your suggestion, we added an experiment on the ImageNet dataset to demonstrate the effectiveness of the proposed method with a large number of classes. The results are summarized in Table D.8 (Appendix D.4):

Furthermore, we tested the multi-score approach on the 20 Newsgroups dataset to make the empirical comparison more extensive by considering different data domains, as suggested also by reviewer 9Lus. The results are summarized in Table D.7 (Appendix D.4):

In both cases, we can see the benefit of the proposed method over the baselines.

On figures clarity and providing tables to show the results

Thank you for your remark. Trying to replace Figures 3, D.1, D.2 and D.3. with tables resulted in overly large and confusing layouts, therefore the figures were enlarged instead for better clarity. We took your suggestion into account and replaced other experimental plots with tables throughout the paper.

On putting some results from Appendix B in the main part

Following the reviewers comments, we added new parts to the paper (a demonstrating example for motivation, and additional variants of our method) . In addition, the figures were enlarged for clarity. Thus, although we wanted to include some experiments from Appendix D in the main section, as you suggested, there was not enough space left.

On Figure 2 - are all the compared methods employed on the multi-head model trained by the loss in Equation (11)?

Yes, all the compared methods on Figure 3 were employed on the multi-head model trained by the loss in Equation (11). We added results for a vanilla approach with a single head for comparison in Appendix D.5 (Table D.10).

Thank you for your helpful and constructive review.

On explaining why the multi-score conformal calibration can work better

Thank you. Following your and reviewer 4GpN comments, we added a motivating example to highlight the ideas behind using a multi-dimensional calibration score (Section 4.1):

Assume a binary classification problem with $\mathcal{Y}=\{-1,1\}$ and prior probabilities $\mathbb{P}(Y=1)$ and $\mathbb{P}(Y=-1)$. The input $X$ is generated from a mixture of two Gaussians, with $\mathbb{P}(X|Y)=\mathcal{N}(Y,\sigma_y^2)$. In this example, we can compute the posterior $\mathbb{P}(Y|X)$ using Bayes rule. Consider the following three classifiers:

(i) $\pi_0(x):=\mathbb{P}(y|x)$

(ii) $\pi_1(x) := \mathbb{P}(y|x) \text{ if } x<0$ and $(\epsilon,1-\epsilon) \text{ if } x>0$

(iii) $\pi_2(x) := (\epsilon,1-\epsilon) \text{ if } x<0$ and $\mathbb{P}(y|x) \text{ if } x>0$

where $\epsilon\sim\mathcal{N}(0,1)$. Here, $\pi_0(x)$ represents the ideal classifier, while $\pi_1(x)$ and $\pi_2(x)$ are ideal classifiers over half the range of $x$, but uninformative over the remaining half. Let $s_i(x,y)$ denote the nonconformity score computed over the $i$-th classifier, it is clear that performing conformal prediction over $s_0(x,y)$ would be the most efficient, however, using either $s_1(x,y)$ or $s_2(x,y)$ will lead to suboptimal results in the uniformative region. In this case, performing conformal prediction in the 2-dimensional space defined by $\mathbf{s}(x,y)=[s_1(x,y),s_2(x,y)]^T$ is advantageous as the individual scores provide complementary information for different ranges of the input $x$. This can be seen from Fig.2(b), where taking both axes into account can help to identify regions with high density of true points versus false points. Setting $\alpha=0.1$, we obtain the following average set sizes: $1.05$ for $s_0(x,y)$, $1.48$ for $s_1(x,y)$, $1.47$ for $s_2(x,y)$ and $1.24$ for our proposed method. Figure 2(c) compares the set sizes per $x$-domain. As expected, $s_0(x,y)$ obtains sets of size $1$ for the entire range, except for $x\in[-1.5,1.5]$, where the two Gaussians overlap. In contrast, $s_1(x,y)$ and $s_2(x,y)$ produce larger proportions of 2-element sets in the noisy regions, on the right for $s_1(x,y)$ and on the left for $s_2(x,y)$. Our method, utilizes both $s_1(x,y)$ and $s_2(x,y)$, obtaining similar behavior to the ideal case provided by $s_0(x,y)$.

More details are provided in Appendix A.

In general, the intuition behind our approach is that higher-dimensional spaces can better separate correct from incorrect labels, potentially leading to more efficient prediction sets with fewer false labels.

Thank you for your helpful and constructive review.

On the comparison to a vanilla conformal prediction baseline trained with a single head.

Following your comment, we added an experiment using a vanilla conformal prediction baseline in Appendix D.5. In the vanilla baseline, the conformal prediction procedure is performed over the original classification head, without the addition of multiple classification heads. We observe that the results are similar to the Best Head baseline (Table D.10):

On the practical usage - instead of using multidimensional nonconformity scores, a practitioner would likely focus on improving the underlying model generating the nonconformity scores.

Thank you. It is true that improving the underlying model would, in general, contribute to improving the efficiency of the conformal prediction procedure, and this will be a main focus of the practitioner to first try improving the underlying model. However, our proposed method presents an additional avenue for improvement that is orthogonal to the improvement of the underlying model. In many practical cases, the search for a better architecture that yields preferable performance is limited by the size or the quality of the data, as well as cost considerations that often impose constraints on model size and inference time. Our approach proposes to enhance the efficiency of any given model with a small addition of few classification heads, or by utilizing an already available ensemble of several models. This approach of improving the efficiency of conformal prediction over a given model was investigated in other works, as detailed in the related work section “Combining nonconformity scores”.

To further strengthen this claim we conducted an additional study where we used an improved model - ViT instead of ResNet50. We show that even for the stronger model we can improve conformal prediction efficiency, as can be seen from Table D.6 :

On the benefit of the proposed method for different datasets, and the dependency on the distribution of nonconformity scores in the multidimensional space.

Thank you for pointing this out. Indeed, the benefit of the proposed method relies on the distribution of the nonconformity scores in the multidimensional space. We added a synthetic example (Section 4.1) that demonstrates that in case the nonconformity scores are noisy for part of the input space, we can benefit from considering multiple scores instead of a single one. This aligns with previous findings, demonstrating that a weighted combination of various scores enhances the efficiency of conformal prediction. In this paper, we extend these ideas and treat the case of multiple scores in a more general way, inspecting it as a multidimensional score and performing the selection in the high dimensional space, while taking into account the distribution of true versus false labels. Our extensive evaluations over different datasets, scores, coverage levels, and data sizes, give strong evidence for the benefit of our method under diverse settings. Moreover, we show that we can modify the distribution of scores using temperature scaling in order to obtain better separation between true and false scores (see Figs. D.5 and D.6).

On the difference between L1 and L2 norm baselines Following tour comment, we added an experiment comparing L1 and L2 baselines in Appendix D.5. Table D.9 summarizes the results:

We observe that both baselines obtain similar performance. We can explain this by noting that the resemblance to the square-like appearance of the nonconformity scores is more prominent at lower scores. The desired coverages we aim for requires a higher threshold, and as a result, the difference between the L1 and L2 norms becomes less prominent. This is further illustrated in Fig. 4, where it is evident that the L2 norm threshold is far from the square-like distribution.

On the effect of increasing dimensionality

Our experimental results include different datasets with varied numbers of classes (between 9 and 1000) and different sizes of calibration points. We do not observe any strong effect in terms of the dimensionality of the problem or the data. In terms of the score dimensionality, we can see that the efficiency improves with increasing the number of classification heads, appearing to plateau around four to five heads, with occasional slight increases thereafter. This behavior might be linked to the curse of dimensionality, suggesting that the current approach is well-suited for moderately sized score spaces. From a practical standpoint, this aligns with typical scenarios, as ensembling methods are generally constrained to combining only a few models due to size and cost considerations.

On the case of duplicate cells (Algorithm 1, line 6)

It might happen that two or more points have the same score, thus can cause duplicate cells. In practice, this rarely occurs in dimensions higher than 1. In such cases, we remove the duplicate cells but account for the multiplicity in the cell score computation in Eq. (9).

On the cells in figure 1

Following your comment, we updated the figure to improve clarity, zooming in to the right bottom corner.

On increasing the size of Figures 2 and 3

Thank you for your remark. As you suggested, the figures have been enlarged.

On moving the part on background work to another section

We moved the part of the conformal prediction background to a dedicated section.

Dear reviewer 3FMy,

Thank you very much again for your helpful and constructive feedback. Your insights have been invaluable in refining our work, and we are grateful for the time and effort you have dedicated to reviewing our manuscript.

As the review process deadline approaches, we kindly request that you revisit our submission in light of the revisions and corrections we have made in response to your valuable feedback. We have carefully addressed the points you raised and believe these changes have strengthened the paper significantly.

If you find the revisions satisfactory, we would greatly appreciate it if you could consider raising your score to reflect the improvements. We would be more than happy to provide further clarifications if you have any additional concerns.

Best, Authors

Thank you for your helpful and constructive review.

On a leave-one-out version of the proposed method

Thank you very much for your helpful suggestion. We implemented a variant of your suggested approach following [1] and [2]. We added the exact implementation in Algorithm B.1. The core idea is to consider the entire calibration data but each time exclude the $i$th point from the set of centers and from the score computation in Eq. (9). We sweep over all possible labels $y\in\mathcal{Y}$ and assign it to the closest center while removing the $i$th center. We perform a similar assignment for the $i$th score $s(X_i,Y_i)$. Then, we compare the rank of the chosen cells, and include $y$ in the final prediction if its rank is smaller than $(1-\alpha)(m + 1)$ hold-out ranks evaluated on the true labeled data for every $i\in\\{1,\ldots,m\\}$. We examined the performance of the Jacknife+ and the split approaches. In both settings, the required 1-$\alpha$ coverage is achieved, but the Jackknife+ version obtained smaller sets, as expected:

However, the improvement appears to be small in this case and may not justify the additional computational cost.

[1] Barber, R. F., Candes, E. J., Ramdas, A., & Tibshirani, R. J. (2021). Predictive inference with the jackknife+.‏

[2] Romano, Y., Sesia, M., & Candes, E. (2020). Classification with valid and adaptive coverage. Advances in Neural Information Processing Systems, 33, 3581-3591.

On the theoretical underpinnings of the multi-dimensional score.

Thank you very much for your important insights. We can indeed think of scenarios where the optimal prediction sets are actually characterized by a multi-dimensional score. We added such an example to the paper, along with figures for illustration (Figures 2 and A.1.):

Assume a binary classification problem with $\mathcal{Y}=\\{-1,1\\}$ and prior probabilities $\mathbb{P}(Y=1)$ and $\mathbb{P}(Y=-1)$. The input $X$ is generated from a mixture of two Gaussians, with $\mathbb{P}(X|Y)=\mathcal{N}(Y,\sigma_y^2)$. In this example, we can compute the posterior $\mathbb{P}(Y|X)$ using Bayes rule. Consider the following three classifiers:

(i) $\pi_0(x):=\mathbb{P}(y|x)$

(ii) $\pi_1(x) := \mathbb{P}(y|x) \text{ if } x<0$ and $(\epsilon,1-\epsilon) \text{ if } x>0$

(iii) $\pi_2(x) := (\epsilon,1-\epsilon) \text{ if } x<0$ and $\mathbb{P}(y|x) \text{ if } x>0$

where $\epsilon\sim\mathcal{N}(0,1)$. Here, $\pi_0(x)$ represents the ideal classifier, while $\pi_1(x)$ and $\pi_2(x)$ are ideal classifiers over half the range of $x$, but uninformative over the remaining half. Let $s_i(x,y)$ denote the nonconformity score computed over the $i$-th classifier, it is clear that performing conformal prediction over $s_0(x,y)$ would be the most efficient, however, using either $s_1(x,y)$ or $s_2(x,y)$ will lead to suboptimal results in the uniformative region. In this case, performing CP in the 2-dimensional space defined by $\mathbf{s}(x,y)=[s_1(x,y),s_2(x,y)]^T$ is advantageous as the individual scores provide complementary information for different ranges of the input $x$. This can be seen from Fig.2(b), where taking both axes into account can help to identify regions with high density of true points versus false points. Setting $\alpha=0.1$, we obtain the following average set sizes: $1.05$ for $s_0(x,y)$, $1.48$ for $s_1(x,y)$, $1.47$ for $s_2(x,y)$ and $1.24$ for our proposed method. Figure 2(c) compares the set sizes per $x$-domain. As expected, $s_0(x,y)$ obtains sets of size $1$ for the entire range, except for $x\in[-1.5,1.5]$, where the two Gaussians overlap. In contrast, $s_1(x,y)$ and $s_2(x,y)$ produce larger proportions of 2-element sets in the noisy regions, on the right for $s_1(x,y)$ and on the left for $s_2(x,y)$. Our method, utilizes both $s_1(x,y)$ and $s_2(x,y)$, obtaining similar behavior to the ideal case, provided by $s_0(x,y)$.

More details are provided in Appendix A.

Thank you very much for reviewing our response and for your feedback.

Following your suggestion, we have included in Appendix A.1. a more detailed discussion on how the multi-dimensional perspective enhances set size optimization, relating to the above mentioned references:

"When $\mathbb{P}(Y|X)$ is known, it was shown that the optimal set with minimal size under coverage constraint is given by $\Gamma^*(x)=\\{y\in\mathcal{Y}|p(y|x)>q_\alpha\\}$ (Lei & Wasserman, 2014; Sadinle et al., 2019; Kim et al., 2021). This implies that the optimal set is a level set of the distribution $p(y|x)$. Thus, in our example, thresholding $s_0(x,y)=1-p(y|x)$ results in the optimal set. Using only $s_1(x,y)$ and $s_2(x,y)$ the optimal set can be equivalently defined as:

$

\Gamma^*(x)&=\{y\in\mathcal{Y}|(\mathbf{1}\{x\leq 0\}\cdot s_1(x,y)+\mathbf{1}\{x>0\}\cdot s_2(x,y))<1-q_\alpha\} =\{y\in\mathcal{Y}|\mathbf{s}^T(x,y)\cdot\mathbf{i}x(x)<1-q\alpha\},

$

where $\mathbf{i}_x(x)=[\mathbf{1}\\{x\leq 0\\}, \mathbf{1}\\{x>0\\}]^T$. Thus, we obtain that the optimal set is a function of the 2-dimensional nonconformity score $\mathbf{s}(x,y)$. In this example, each classifier specializes on a different subdomain of the input space $\mathcal{X}$. Another practical case is when classifiers specialize on different parts of the output space $\mathcal{Y}$. For example, consider $\mathcal{Y} = \\{0, 1, 2\\}$, and the following three classifiers:

$\pi_a(x) := \mathbb{P}(y=a|x) \text{ if } y=a$, or $\epsilon \text{ if } y=(a+1) \text{ mod } 3$ or $1-\mathbb{P}(y=a|x)-\epsilon, \text{ if } y=(a+2) \text{ mod } 3,\text{ } a\in\\{0,1,2\\},$

where $\epsilon\sim\mathcal{N}(0,1)$. In this case, the optimal set is given by:

$\Gamma^*(x)=\\{y\in\mathcal{Y}|\mathbf{s}^T(x,y)\cdot\mathbf{i}_y(y)<1-q\\}$

where $\mathbf{i}_y(y)=[\mathbf{1}\\{y=1\\}, \mathbf{1}\\{y=2\\},\mathbf{1}\\{y=3\\}]^T$.

We conclude that whenever $p(y|x)=\phi(\mathbf{s}(x,y);x,y)$, where $\phi:\mathcal{S}\times\mathcal{X}\times\mathcal{Y}\rightarrow [0,1]$ is a non-degenerate function of the multi-score vector, the optimal set relies on $\mathbf{s}(x,y)$, i.e.:

$\Gamma^*(x)=\\{y\in\mathcal{Y}|\phi(\mathbf{s}(x,y);x,y)<1-q_\alpha\\}\text{ }$ (17)

In contrast, relying solely on a single score will result in a suboptimal solution. Note that, according to Eq. (17), the ideal set corresponds to a level set of $\phi(\mathbf{s}(x, y); x, y)$, rather than $\mathbf{s}(x, y)$ itself. This implies that, in general, the decision boundaries in the multi-dimensional score space can be arbitrarily complex, depending on the properties of $\phi$.

In practice, we have access to neither the conditional distribution $\mathbb{P}(Y|X)$ nor the mapping function $\phi$. Instead, we aim to solve the problem of minimizing the set size subject to a coverage constraint, similarly to (Stutz et al., 2021; Bai et al., 2021; Kiyani et al., 2024). Since the space of all possible prediction sets is overly complex, the problem must be relaxed. Bai et al., 2021 proposed to optimize an arbitrary class of prediction sets $\Gamma_{\theta}$ parametrized by $\theta$, while Stutz et al., 2021 optimize a parametrized score $s_\theta(x,y)$. In (Kiyani et al., 2024), structured prediction sets of the form $\Gamma_h^s(x) = \\{y \in Y | s(x, y) \leq h(x)\\}$ were considered, where $h:\mathcal{X}\rightarrow \mathbb{R}$ is a learned adaptive threshold. In contrast, we work in the multi-score domain and consider sets defined as a union of a subset $I\subseteq 2^k$ of cells in $D_\textrm{cells}$, i.e. $\Gamma_{I}^\mathbf{s}(x)=\\{y\in\mathcal{Y}|\mathbf{s}(x,y)\in\cup_{i\in I}\mathcal{C}_{i}\\}$. Accordingly, the relaxed optimization problem can be written as:

$arg\min_{I\subseteq 2^k} \mathbb{E}\left[\textrm{size}(\Gamma_{I}^\mathbf{s}(X))\right]$

$\text{s.t. } \mathbb{E}\left[\mathbf{1}\\{Y\in \Gamma_{I}^\mathbf{s}(X)\\}\right]\geq 1-\alpha$

where $\textrm{size}(\Gamma_{I}^\mathbf{s}(X))=\sum_{q=1}^Q\mathbf{1}\\{Y\in\Gamma_{I}^\mathbf{s}(X)\\}=\sum_{i\in I}\sum_{q=1}^Q\mathbf{1}\\{Y\in\mathcal{C}_{i}(X)\\}$.

A finite sample approximation of the expected set size is given by:

$\mathcal{E}=\frac{1}{k}\sum_{j=1}^k\sum_{i\in I}\sum_{q=1}^Q \mathbf{1}\\{\mathbf{s}(X_j, q) \in \mathcal{C}_i\\}.$

Note that this is equivalent to summing the cell scores $D_i$ defined in Eq. (9) (without removing duplicate cells):

$\mathcal{E} = \frac{1}{k}\sum_{i\in I}D_i.$

Therefore, solving the optimization problem in Eq. (18) does not require enumerating all possible sets $I$, which is computationally infeasible. Instead, we can rank the cells according to $D_i$ and take the $(1-\alpha)$ proportion of cells with the lowest score values. To obtain exact coverage we perform a recalibration over $\mathcal{D}_\textrm{re-cal}$.

I have read the authors response and other reviewers opinions. I have no further concerns and vote for acceptance.