Length Optimization in Conformal Prediction

Addressing Reviewer's Comments on the Experiments

**Scaleability of Experimental Design** In our section on the experiments, we evaluate the performance of CPL in three different setups ranging in both regression and classification settings on real-world and synthetic datasets. Our experiments aim to compare CPL with other **state-of-the-art conformal prediction algorithms that treat the base model as a black box**. All three experimental setups are standard and have been used in the community to measure the performance of CP methods. In particular, section 5.1 setup is standard in comparing methods that provide marginal coverage as introduced by [2], section 5.2 is standard for group conditional methods as designed by [3], and finally, 5.3 setup is identical to that used by Gibbs et al. [1] to showcase performance under a class of covariate shifts. Also, in response to the reviewers request, we will add the method of "Split localized conformal prediction" to the baselines of section 5.1. **Additional Experiment on LLM Question Answering** To further demonstrate the applicability/scalability of our method, we will include a large-scale experiment involving LLM question-answering in the camera-ready version. This experiment uses multiple-choice question-answering datasets, including TruthfulQA, MMLU, OpenBookQA, PIQA, and BigBench. The calibration data for MMLU alone consists of approximately 100,000 prompts and answers, illustrating the large scale of this experiment. The goal is to quantify the uncertainty of Llama 2 and create prediction sets using this model. We follow a procedure similar to that proposed by prior work [4], adapting it as follows: for each dataset, the task is a multiple-choice question-answering. We pass the question to Llama 2 using a consistent prompt: "This is a 4-choice question that you should answer: {question} The correct answer to this question is:" We examine the logits of the first output token for options A, B, C, and D. Applying a softmax function gives us a probability vector, and we define the conformity score as $1 − \text{probability of the correct option}$ similar to $1-f(x)_y$ for classification. CPL is implemented using a linear head (as $h(x)$) on top of a pre-trained GPT-2. In more detail, with GPT-2 having 768-dimensional hidden layers, the inner maximization involves optimizing a 768-dimensional linear map from GPT-2’s last hidden layer representations to a scalar. We also implemented the baseline method that directly applies the split conformal method to the scores. Please **see the uploaded PDF for plots** and a caption for more details. **Computational Cost** We acknowledge and agree with the reviewer's comment on the importance of computational cost. Two points about this matter are in order. (i) Going beyond the simple solution of split conformal for black box base models, would require an adaptive threshold (as argued in our paper). This adaptive threshold is often achieved by an extra optimization on the calibration data (e.g. see [1] and [2]). We are introducing length optimization on top of conditional validity, hence it is expected that one would have to still do optimization to find the threshold. (ii) our computational cost is not more than training an ML model for the threshold, and it can effectively be solved by gradient descent. For example, computing the gradients for inner maximization amounts to computing the gradient with respect to the parameters of a neural network (hence is totally scalable). Also, the outer minimization step is computationally lightweight and can be viewed as a form of hyper-parameter tuning. For instance, in scenarios focused on marginal coverage, the outer minimization reduces to a simple scalar optimization. We will ensure to include a comprehensive discussion on scalability in the camera-ready version of the paper. **Description of the Experiments** We acknowledge the reviewer's comment on the importance of a detailed description and detailed report of the results. The brevity in some parts is due to space constraints. In the revised version, we will take advantage of the allowed extra page to include more details on the setups and plot captions. We also acknowledge the importance of higher-order statistics. Some of these reports are in the paper (e.g. see lines 283-284 or 297), and some others were dropped as the error bars were negligible. We will make sure to improve the quality of statistics reports in the revised manuscript. We thank the reviewer for this comment. We will also provide full details of implementation specifications, and include a link to our Python implementations in the camera-ready version. **References** [1] Gibbs, Isaac, et al.. "Conformal prediction with conditional guarantees." (2023). [2] Romano, Y., et al. "Conformalized quantile regression" (2019). [3] Jung, Christopher, et al. "Batch multivalid conformal prediction." (2022). [4] Kumar, Bhawesh, et al. "Conformal prediction with large language models ...." (2023).

Response to Individual Questions

“Which data is used to optimize...” The algorithm 1 uses all the calibration data to optimize both $f(\cdot)$ and $h(\cdot)$, through a minimax procedure, where we alternatively take gradient descent with respect to $f(\cdot)$ and gradient ascent with respect to $h(\cdot)$. In that process, the outer min ensures the conditional validity of the prediction sets and the inner max improves length efficiency. We will clarify this in the updated manuscript. “In Figure 3, right-hand-side plot, CPL achieves smaller ...” The optimal oracle has the smallest interval length **averaged over all the covariates (all the test samples)**, while achieving valid coverage with respect to each of the groups. You can see this by looking at the first bars on the right of the “Figure 3, right-hand-side plot”. The optimal oracle is achieving mean length interval of around 8.0 while CPL is clearly achieving more than 8. However, looking at the mean interval conditioned on each grouping of the data (which corresponds to the rest of the bars in the “Figure 3, right-hand-side plot”) CPL has smaller or larger mean interval lengths in different groups. We will elaborate on this matter in the revised version. “I am somewhat confused by the notation. Sometimes ...” We acknowledge the difficulty in the notation and we will make sure to improve the clarity of the notations in the revised version. In general, throughout the paper $\mathcal{F}$ denotes the class of covariate shift and $f\in\mathcal{F}$ a single covariate shift function in the class. However, sometimes $f(\cdot)$ is also used to address the base model (which will be fixed). In line 280, by a NN with two hidden layers we meant our neural network architecture consists of three fully connected layers, with ReLU nonlinearities between layers. The first layer takes as input the feature vector X and outputs 64 hidden variables. The second layer follows the same template, outputting another 64 hidden variables. Finally, a linear output layer returns a pointwise estimate of the response variable Y. "How are $f(.)$ and $h(.)$ are parametrized in the experiments? Could the ..." The function $h(\cdot)$ is parameterized by the natural parameters of the machine learning model used to implement it. For example, $h(\cdot)$ could be implemented as a neural network (NN). In each experimental setup, we specify the ML model we use for $h(\cdot)$. In our experiment setups, simple MLP models or fine-tuning the last layer of a pre-trained model suffices to get a good performance. For the function $f(\cdot)$, recall that we define $\mathcal{F} = \{\langle \boldsymbol{\beta}, \Phi(x)\rangle \mid \boldsymbol{\beta} \in \mathbb{R}^d\}$, where each $f \in \mathcal{F}$ can be represented as $f(x) = \langle \boldsymbol{\beta}, \Phi(x)\rangle \equiv f_{\boldsymbol{\beta}}(x)$. Therefore, iterating on $f$ can be implemented by updating $\boldsymbol{\beta} \in \mathbb{R}^d$. For instance, in the case of marginal coverage, $\boldsymbol{\beta} \in \mathbb{R}$, hence it becomes a scalar optimization. For the case of conditional coverage with respect to $m$ groups, $\boldsymbol{\beta}$ would be an $m$-dimensional vector (see section 2.1). In a nutshell, the number of parameters used for $h(\cdot)$ and $f(\cdot)$ is typically significantly less than the number of (learned) parameters at training time. Concerning a more powerful base model, as mentioned above, we treat the base model as a black-box and quantify its uncertainty. We have also added a CIFAR-10 experiment (described earlier) to showcase that CPL can also improve the length on top of training methods like conformal training. "Could the authors share some intuition or results on how tight..." The coverage results in conformal prediction are typically provided in two ways. One is over the expectation on the calibration data (e.g. the one in Gibbs et al. [1]). The others are PAC-style guarantees similar to what we provide in Theorems 4.2. and 4.3. PAC-style guarantees are very common in the CP literature (e.g. see [2], [3], [4]). The PAC-style guarantees are stronger in the sense that they are stated with high probability over the draw of the calibration (rather than over expectation). These bounds are also known as training-conditional guarantees in the literature. PAC-style guarantees of $O(1/\sqrt{n})$ are generally optimal (e.g. see [2], [3]). "The description of the covariate shift experiments in" We will respond to this question fully in the next official comment. **Regarding the minor comments, we thank the reviewer very much and we'll revise the text accordingly.** **References** [1] Gibbs, Isaac, et al.. "Conformal prediction with conditional guarantees." (2023). [2] Sangdon Park, et al. "Pac confidence sets for deep neural networks via calibrated prediction." 2019. [3] Vovk, Vladimir. "Conditional validity of inductive conformal predictors." 2012. [4] Jung, Christopher, et al. "Batch multivalid conformal prediction." (2022).

Experimental Setup in Section 5.3

Here we respond to the reviewer's concern regarding the clarity of our description of the experimental setup in Section 5.3. Below, we provide a more detailed explanation of our experimental setup for the RxRx1 dataset from the WILDS repository, highlighting key aspects and methodologies. Dataset Overview: The RxRx1 dataset consists of cell images obtained through fluorescent microscopy, with the task of predicting one of 1,339 genetic treatments. These images come from 51 different experiments, each representing covariate shifts due to varying execution and environmental conditions. Our goal is to create prediction sets that maintain valid coverage across these shifts, a challenge due to heterogeneity in the quality of predictions across different experiments and cell types. Experimental Design: 1. **Covariate Shift Characterization:** - We follow a data-driven approach to characterize covariate shifts, similar to the method in Gibbs et al. (2023), by splitting the calibration data in half and estimate the probabilities of experiment membership in the first half. We estimate these probabilities using $\ell_2$-regularized multinomial linear regression on the pre-trained ResNet50 model's feature map outputs. We then use these estimated probabilities to form the class of covariate shifts. 2. **Model and Data Splitting:** - We use a ResNet50 as the base model, pre-trained on 37 of the 51 experiments. These 37 experiments contain a total of approximately 91000 images. The remaining 14 experiments are split into calibration and test sets, each of which with approximately 16000 images. For CPL inner maximization (corresponds to $h(.)$), we train a linear head on the last-layer representations of the pre-trained ResNet50 model. 3. **Conformity Scores and Temperature Scaling:** - For each image $x$, the ResNet50 model outputs weights $\{f^i(x)\}_{i=1}^{1339}$ for the treatments. These are converted to probability weights $\pi^i(x)$ using softmax with a temperature parameter $T$ that is calibrated to adjust these probabilities. Temperature refines the accuracy of these probabilities. The conformity score $S(x, y)$ is computed as the sum of probabilities for treatments where the predicted probability $\pi^i(x)$ exceeds that of the correct treatment $\pi_y(x)$. 4. **Comparative Analysis:** - We compare CPL with the Conditional Calibration algorithm, as proposed by Gibbs et al. (2023), and the Split Conformal method. Our findings, depicted in Figure 4, show that while both CPL and Conditional Calibration maintain almost valid coverage across shifts, CPL significantly reduces the average prediction set size due to its length optimization, which enhances the efficiency of prediction sets. We will take advantage of the extra page in the revised manuscript to enhance the description of section 5.3.

Additional Experiment on LLM Question Answering

To further demonstrate the applicability/scalability of our method, we will include a large-scale experiment involving LLM question-answering in the camera-ready version. This experiment uses multiple-choice question-answering datasets, including TruthfulQA, MMLU, OpenBookQA, PIQA, and BigBench. The calibration data for MMLU alone consists of approximately 100,000 prompts and answers, illustrating the large scale of this experiment. The goal is to quantify the uncertainty of Llama 2 and create prediction sets using this model. We follow a procedure similar to that proposed by prior work [4], adapting it as follows: for each dataset, the task is a multiple-choice question-answering. We pass the question to Llama 2 using a consistent prompt: "This is a 4-choice question that you should answer: {question} The correct answer to this question is:" We examine the logits of the first output token for options A, B, C, and D. Applying a softmax function gives us a probability vector, and we define the conformity score as $1 − \text{probability of the correct option}$ similar to $1-f(x)_y$ for classification. CPL is implemented using a linear head (as $h(x)$) on top of a pre-trained GPT-2. In more detail, with GPT-2 having 768-dimensional hidden layers, the inner maximization involves optimizing a 768-dimensional linear map from GPT-2’s last hidden layer representations to a scalar. We also implemented the baseline method that directly applies the split conformal method to the scores. Please **see the uploaded PDF for plots** and a caption for more details.