Multi-Dimensional Conformal Prediction

Authors: Yam Tawachi, Bracha Laufer-Goldshtein
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: classification, conformal prediction, uncertainty quantification

Abstract:
Conformal prediction has attracted significant attention as a distribution-free method for uncertainty quantification in black-box models, providing prediction sets with guaranteed coverage. However, its practical utility is often limited when these prediction sets become excessively large, reducing its overall effectiveness. In this paper, we introduce a novel approach to conformal prediction for classification problems, which leverages a multi-dimensional nonconformity score. By extending standard conformal prediction to higher dimensions, we achieve better separation between correct and incorrect labels. Utilizing this we can focus on regions with low concentrations of incorrect labels, leading to smaller, more informative prediction sets. To efficiently generate the multi-dimensional score, we employ a self-ensembling technique that trains multiple diverse classification heads on top of a backbone model. We demonstrate the advantage of our approach compared to baselines across different benchmarks.

Decision: accept (poster)
Reviewer scores: 6, 6, 3, 8, 5

Reviews and discussion:
<thread>
creation date:
1730000601638

thread_type:
official_review

<thread_object>
creator: reviewer_9lus

<content>
summary:
The paper introduces a multi-dimensional conformal prediction method for classification problems. A nonconformity measure is defined depending on the multiple scores, which leads to finer separation between correct and wrong labels. The proposed procedure mainly includes three steps: first, partitioning the score space into cells using $D_{cells}$; second, scoring and ranking the cells according to the ratio of false to true labels; and third, calibration-achieved by progressively adding top-ranked cells until the desired coverage level is reached.

strengths:
The paper is well-organized. The problematic and approach are well introduced, and the concept concisely presented. The proposed method successfully extends conformal prediction into a multi-dimensional setting with enhanced separation between correct and incorrect labels, and it appears sufficiently general.

weaknesses:
1. The experiments mainly focus on image classification datasets (e.g., CIFAR100, Tiny ImageNet, PathMNIST), limiting the demonstration of the other data types, such as text or tabular data.

2. While the paper is claimed to be cost-efficient, it still involves training and managing multiple classification heads, which could be computationally expensive for large-scale datasets or resource-constrained environments. Could you provide more insights into the computational cost? Specifically, how does it compare to other baseline methods?

questions:
How sensitive is the method to hyperparameter choices, such as the number of classification heads or diversity regularization? And how do different configurations affect set size?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730127724205

thread_type:
official_review

<thread_object>
creator: reviewer_ugxt

<content>
summary:
This paper tends to combine several different conformity score functions aiming to reduce the prediction set size. In short their approach is to combine different score functions per $(x, y)$ forming a vector of scores. Dividing the calibration set into two separate sets of re-cal and tuning set, they use the tuning set as the basis (centroids) in the mutli-dimensional scores space. They assign a rating to each of these centriods based on the proportion of false labels inside that area (this means 1 + set size if there is only one point at the place). During calibration they deactivate regions in decreasing order to the number of false predictions, as long as the $1 - \alpha$ coverage is satisfied. Each calibration point is assigned to the closest centriod. Similarly during test. they check the closest centiod and add the label to the prediction set if the centriod is active.

The propose different methods to build multi-dimensional score including using different classification heads, to simply using different score functions.

strengths:
The method is simple yet effective. The paper is also easy to read and “to the point”. I do not see any major flaw or discouraging point in it. I also see that the authors tried to include a comprehensive evaluation by comparing with best head, uniform average and optimized weights.

weaknesses:
1. Why the authors did not report worst slice coverage which is metric of adaptivity? 
2. Why the authors used an unnecessarily large calibration set? I don’t think the evaluation setup is realistic with 16500 calibration points to predict 11000 test points. What is the gain of this large calibration set? Is there any statistical efficiency results that are missing?
3. Related to the previous problem, a fair comparison is to report paper’s best method in comparison with the normal best method where the tuning set is also added to the calibration set. The results are not expected to be significantly different but still this is important to be reported.
4. Why the authors did not report their result with APS and threshold prediction sets? Although there are adaptivity and efficiency constraints, still those result are important to be reported. RAPS and SAPS are ignoring informations in the logit space either by penalty, or just reporting the rank, and intuitively in a perfect classifier this should cause inefficiencies.
5. Is there any preference to have 0-1 assignment to the zones?, the author could also try soft probability assignment to each region and optimize for that. It is not a must, but either trying soft centroid assignment, or showing that the optimal is happening with the current 0-1 assignment would make the work stronger.

I would increase my score if a convincing answer to these two questions is provided.

1. Why the calibration set is very large in all experiments. Are there any limitations? What would the result be with smaller calibration set (e.g. 2000 points)
2. Why threshold prediction sets (softmax score) and adaptive prediction sets are not evaluated. The authors could have similar results with these score functions as well.
3. An ablation study on the number of centroids is important. How the results would change with very limited number of centriod (e.g. 200)?

questions:
Please see Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730624828155

thread_type:
official_review

<thread_object>
creator: reviewer_52ya

<content>
summary:
In this paper, the authors aim to improve the efficiency of conformal prediction sets with multiple non-conformity scores. In particular, they propose a multidimensional framework of conformal predictions, which uses a single score in one of the dimensions. They provide a theory to show the desired coverage can be guaranteed through the multidimensional selection. To obtain multiple diverse nonconformity scores, they employ a multi-head model and train it with a diversity regularization. The proposed method is validated on a CIFAR100, Tiny ImageNet, and PathMNIST to show the effectiveness.

strengths:
1. The method is novel. To the best of my knowledge, this is the first work to utilize ensemble predictions as multidimensional score.
2. The empirical improvement is significant. The proposed method can produce smaller prediction sets compared to other baselines.

weaknesses:
Major:
1. Why the multi-score conformal calibration can work better is not explained. The authors only provide a theorem to ensure the valid coverage of the proposed method, but do not elaborate about how it leads to efficient prediction sets. It could be better if they can provide convincing explanation in this aspect.

2, The multi-score setting seems different from previous works. Previous multi-score works aim to utilize multiple score functions to obtain better performance, but this work proposes the method to generate multiple scores with ensemble models. In my view, this problem should be CP for ensemble models instead of the multi-score setting. So it might be better for authors to rewrite the motivation (subsection 3.2) for avoiding misunderstanding. (Or the authors can show the method can also improve the setting of multiple score functions.)

3. The empirical comparison is not extensive. The authors should include ImageNet in the experiments, or it is hard to evaluate the effectiveness of the proposed method in large number of classes. 

Minor:
1. presentation: it is hard to read the information from Figures 2 and 3. Maybe the authors can simply provide tables to show the results.
2. It might be better to put some results in Appendix B in the main part, as there is almost one page left as blank.

questions:
In figure 2, are all the compared methods employed on the multi-head model trained by the loss in Equation (11)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730816445123

thread_type:
official_review

<thread_object>
creator: reviewer_3fmy

<content>
summary:
This paper proposes an extension to conformal prediction for classification tasks. Standard conformal prediction constructs prediction sets with guaranteed coverage by thresholding a nonconformity score, which measures the disagreement between the model and the ground truth class label. This work has two main contributions: (a) a variant of conformal prediction that is compatible with multidimensional nonconformity scores, and (b) a self-ensembling method for producing these multidimensional nonconformity scores using diverse classification heads. The proposed method partitions the multidimensional space and then selects a set of cells to include in prediction sets by thresholding the cells based on the relative frequency of false labels. The proposed approach is shown experimentally to lead to smaller prediction sets.

strengths:
Strengths of the paper include the following:
- The main claims are supported with evidence. Theoretical justification for the proposed procedure is provided in Appendix A.1. Smaller prediction set sizes are also demonstrated in Figure 2 and Figure 3. However, see weaknesses below regarding multiple classification heads.
- There is a good amount of exploration performed, even beyond what is reported in the main paper. For example, results on standard ensembling, test-time augmentations, and size of $\mathcal{D}_\text{cells}$ are all useful.
- The paper is well-written. I thought the related work section in particular explored a good amount of relevant literature.

weaknesses:
Weaknesses of the paper include the following:
- The empirical evidence provided in the paper is overall good, but the role of the multiple classification heads is still unclear. The best head baseline uses a network trained with multiple heads. How would the proposed approach compare to a vanilla conformal prediction baseline trained with a single head (i.e. the network prior to attaching the six additional classification heads)?
- The contribution of the paper is fairly solid in terms of theoretical and empirical justification. However, from a practical perspective, it seems more likely that rather than investing time into using multidimensional nonconformity scores, a practitioner would likely simply focus on improving the underlying model generating the nonconformity scores. The cases in which the multidimensional approach succeed are those where the model is suboptimal, and fixing the model will remove the need to use multidimensional scores. This limits the significance of the proposed work.
- The benefits of the proposed approach are empirical in nature. There does not seem to be a way to know beforehand whether multidimensional nonconformity scores will help or not. In some sense, this is a modeling question that depends on the distribution of nonconformity scores in the multidimensional space. There will be some datasets where the modeling assumptions proposed here (analogous to a KNN variant with $K=1$) will be beneficial, and others where it will not be.

questions:
- Please see the first point under the weaknesses section about. How would the proposed approach compare to a vanilla RAPS/SAPS baseline with one classification head (i.e. before the six additional heads are attached)?
- In Figure 4, the nonconformity scores appear to be distributed in a pattern resembling a square. Would an L1 norm perform better than an L2 norm in this case?
- What is the effect on the algorithm as the dimensionality increases? Is there a curse of dimensionality effect here?
- In Algorithm 1 (line 6), when do the duplicate cells occur?

Minor comments:
- The cells in Figure 1 are difficult to understand.
- Figures 2 and 3 are small and therefore hard to read.
- Since Section 3.1 discusses background work, it should probably be moved out of Section 3, which is title "Proposed Method".
</content>
</thread_object>

</thread>
<thread>
creation date:
1731427864243

thread_type:
official_review

<thread_object>
creator: reviewer_4gpn

<content>
summary:
This paper studies a multi-dimensional non-conformity approach to conformal prediction. They introduce a novel method to identify high probability regions with low low false-to-true labels ratio, which requires to carefully design a calibration procedure with multiple stages, including partitioning the space, ranking different partitions, and selecting the high quality partitions.

strengths:
This paper is very well-written and has a very clear flow. The multi-dimensional CP looks like a promising avenue of research. They provide a novel perspective of finding conformalized regions by ranking the false-to-true ratio and they also show sufficient experimental results to back up their claim about improving efficiency.

weaknesses:
I believe this is a very good paper. Although I do have some suggestions to further improve both the theoretical and algorithmic perspectives.

- You break the dataset in two parts and then use the second part as a recalibration to ensure coverage. A natural way to improve the efficiency of your method and use all the data together is to use a leave-one-out method. At each time, you can use all the data but one (say X_i, Y_i), go through the partitioning and ranking and then find the smallest threshold (rank) such that the selected regions include the left-out sample (X_i, Y_i). Lets call that threshold (rank) tau_i. Then you can take (1-\alpha))-quantile of the ranks and use that as the final threshold (rank) to make the regions. There are some subtleties to make this work, you might want to look at https://arxiv.org/abs/1905.02928 for more details. 

- For the case of scalar scores there are some theoretical underpinnings that shows the optimal (think of it as the tightest possible) prediction sets are of the form of level-sets of a scalar function. (i.e., thresholding a scalar score) These results are considered classic in CP (look at https://www.stat.cmu.edu/~jinglei/LeiW14.pdf) and they have been also recently generalized to the more advanced case of CP with conditional coverage requirements (look at https://arxiv.org/pdf/2406.18814). These results and their insights also heavily contributed in the line of research around designing better scores. Therefore, I believe it is important to discuss whether a multi-dimensional score design has any fundamental theoretical underpinnings or not. For instance, it would be nice to think of a scenario where (in the population regime) the optimal prediction sets are actually characterized by a multi-dimensional score. Or maybe trying to show that in a (perhaps simplified) scenario the multi-dimensional perspective is provably a better approximation of the optimal prediction sets. These kind of results and discussions are more than just an effort for having some theory in the paper, but rather can bring useful insights for follow-up works in this line of research.

I am willing to increase my score in case of a satisfactory response.

questions:
I have no further questions. Might ask more once I see the authors response.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732459559182

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment

comment:
We would like to thank the reviewers for their time and effort in reviewing our paper. Their constructive comments helped us to improve the paper as reflected in the revised manuscript. Following their comments, we made major changes to the paper's content, clarifying important aspects, and providing additional details on the experimental setup. Moreover, we added new experiments and ablation studies that strengthen the empirical verification of our method.
 
We would like to emphasize the major modifications to the revised manuscript:

(i) We introduced an illustrative example that demonstrates the motivation for our proposed multi-dimensional approach (in Section 4.1). 

(ii) We added two additional versions of our proposed method (Section 4.2): (1) A statistically efficient version based on a leave-one-out procedure, known as Jackknife+. (2) A soft version that considers for each point $b$ neighboring cells, where our core method is a private case with $b=1$. 

(iii) We conducted extensive additional experimental studies, showing the robustness of the  proposed method with respect to: the data domain and size, the size of the calibration data and the number of center points, the strength of the diversity regularization and the choice of the nonconformity scores. These results further emphasize the advantage of the proposed method under varied settings.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732784005466

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We sincerely appreciate the thoughtful and constructive feedback provided by all the reviewers. In response, we have uploaded a significantly revised version of our paper, which carefully addresses the review comments. We believe these revisions have led to meaningful improvements in several key aspects, thanks to your valuable suggestions.

Please do not hesitate to share any additional questions or concerns.

Thank you once again for your time, understanding, and insightful feedback.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper extends the standard approach for conformal prediction to produce prediction sets (uses a single non-conformity score) to multi-dimensional non-conformity scores. The key idea behind the approach is to partition the multi-dimensional space and selecting the high-quality partitions with low false-to-true ratio of labels using the calibrated thresholds. The paper provides theoretical analysis and empirical evidence for smaller prediction sets from the proposed method.

The reviewers' were generally positive about the paper, but also raised a number of questions. The author rebuttal was thorough and answered almost all questions satisfactorily, and the paper is also revised. One reviewer argued for rejection for the following reason: "the explanation of why the method works is not convincing without exact validation". This seems unfair to me given the motivation and set-size optimization perspective, and empirical results. As reviewer UGXT suggested, showing that the rankings are more aligned with ground truth aleatoric uncertainty using synthetic experiments can be helpful and I encourage the authors' to consider it.

Therefore, I recommend accepting the paper and strongly encourage the authors' to incorporate all the discussion in the camera copy to further improve the paper.