Length Optimization in Conformal Prediction (2406.18814v3)

Abstract: Conditional validity and length efficiency are two crucial aspects of conformal prediction (CP). Conditional validity ensures accurate uncertainty quantification for data subpopulations, while proper length efficiency ensures that the prediction sets remain informative. Despite significant efforts to address each of these issues individually, a principled framework that reconciles these two objectives has been missing in the CP literature. In this paper, we develop Conformal Prediction with Length-Optimization (CPL) - a novel and practical framework that constructs prediction sets with (near-) optimal length while ensuring conditional validity under various classes of covariate shifts, including the key cases of marginal and group-conditional coverage. In the infinite sample regime, we provide strong duality results which indicate that CPL achieves conditional validity and length optimality. In the finite sample regime, we show that CPL constructs conditionally valid prediction sets. Our extensive empirical evaluations demonstrate the superior prediction set size performance of CPL compared to state-of-the-art methods across diverse real-world and synthetic datasets in classification, regression, and LLM-based multiple choice question answering. An Implementation of our algorithm can be accessed at the following link: https://github.com/shayankiyani98/CP.

• The paper introduces CPL, a novel framework that minimizes prediction set lengths while ensuring conditional validity.
• It formulates the design of prediction sets as a minimax problem leveraging level set estimation for strong duality results.
• Empirical tests across regression, text, and image tasks demonstrate CPL's superiority over state-of-the-art CP methods.

Length Optimization in Conformal Prediction

Abstract and Introduction

Shayan Kiyani, George Pappas, and Hamed Hassani tackle two central challenges in conformal prediction (CP): conditional validity and length efficiency. Their paper introduces Conformal Prediction with Length-Optimization (CPL), a framework designed to construct prediction sets that are near-optimal in length while ensuring conditional validity under various covariate shifts. This work extends the literature by providing strong duality results in the infinite sample regime and constructing conditionally valid prediction sets in the finite sample regime. Their empirical evaluations demonstrate CPL’s superior performance across multiple domains, including classification, regression, and text-related tasks.

Conditional Validity and Length Efficiency

Conformal prediction aims to create prediction sets $C(x)$ for each input $x$ such that the true label $y$ is included with high probability. Traditionally, CP guarantees marginal coverage, but practical applications often require conditional validity to ensure accurate predictions across different subpopulations, such as various patient demographics in healthcare. However, full conditional coverage is impossible with finite data, leading to the use of relaxations such as group-conditional coverage and coverage under covariate shifts.

Simultaneously, length efficiency pertains to the need for prediction sets to be as small as possible while maintaining their coverage properties. Large prediction sets are less informative, limiting the utility of CP methods. The intertwining of these two aspects raises the question of how to achieve length-optimal prediction sets that still meet conditional validity requirements.

Framework and Theoretical Insights

Kiyani et al. address the primary problem of CP—minimizing the average length of prediction sets subject to conditional validity constraints—by formulating it as a minimax problem. This dual formulation navigates the space of conditionally valid prediction sets and identifies the length-optimal ones. The theoretical underpinning is the connection between conformal prediction and level set estimation, where the optimal prediction sets are characterized as specific level sets of the distribution functions.

Key results include:

• Propositional Findings: The optimal prediction sets $C^*(x)$ for the primary problem are level sets of the form $\{y \in \mathcal{Y} \;|\; f^*(x) p(y|x) \ge 1 \}$.
• Relaxed Minimax Problem: Practical constraints necessitate a relaxed version of the minimax problem, wherein structured prediction sets defined via a conformity score are optimized. Despite this relaxation, the framework guarantees conditional validity and near-optimality in set size.
• Theoretical Guarantees: The authors provide rigorous proofs to establish the strong duality between the relaxed minimax problem and the primary problem, validating that solutions remain close to the optimal even with finite data.

Empirical Evaluations

The authors conducted extensive empirical evaluations on multiple real-world and synthetic datasets, showcasing CPL's improvements over state-of-the-art methods in terms of prediction set size and maintaining proper coverage:

1. Regression Data: CPL outperformed traditional split conformal methods and sophisticated state-of-the-art techniques like CQR, reducing interval lengths while ensuring near-perfect marginal coverage across 11 diverse regression datasets.
2. Text Data: For the task of multiple-choice question answering using Llama 2, CPL achieved significantly smaller prediction sets compared to methods like Kumar et al.'s 2023 approach, maintaining proper coverage levels. This demonstrates CPL's effectiveness in quantifying the uncertainty of LLMs.
3. Synthetic Group-Conditional Data: In a controlled environment requiring group-conditional coverage, CPL provided superior length efficiency compared to BatchGCP while ensuring valid coverage across predefined groups.
4. General Classes of Covariate Shifts: Using the RxRx1 dataset for image classification under covariate shifts, CPL maintained valid conditional coverage with respect to predefined covariate shifts, outperforming the Conditional Calibration method in terms of prediction set size.

Practical and Theoretical Implications

The practical implications of this research are significant, given the demands for small, accurate prediction sets in various applications such as healthcare, decision-making, robotics, and LLMs. Theoretically, the paper builds a strong foundation linking level set estimation to the design of length-optimal conformal prediction sets, opening avenues for future research in the optimization of CP methods under different structural constraints of the data.

Future Directions

Future research can further explore duality results in optimization problems with infinitely-many constraints and the stability of prediction set length when restricting solution classes. Additionally, extending CPL to handle infinite-dimensional covariate shift classes and developing more detailed theoretical frameworks for length optimality in finite-sample regimes would significantly advance the field.

Conclusion

This paper presents a robust framework for optimizing the length of prediction sets while ensuring conditional validity in conformal prediction. Through rigorous theoretical developments and comprehensive empirical evaluations, Kiyani, Pappas, and Hassani provide a valuable contribution to CP literature, offering a practical solution with significant improvements over existing methods.