Prune 'n Predict: Optimizing LLM Decision-making with Conformal Prediction

We are delighted with the positive feedback on our paper. We appreciate the recognition of strengths on the ideas, presentation and soundness of empirical analysis. Our response to the queries is as follows.

Why tanh for $g$ ? While the choice of the function class $\mathcal{G}$ is up to the user, in general it makes sense to use a flexible non-linear function class. A multi-layer neural network with any activation function could be a good fit here. We chose tanh as its range is (-1, 1) and is known to have nice properties such as resulting in efficient training [1].

[1] https://cseweb.ucsd.edu/classes/wi08/cse253/Handouts/lecun-98b.pdf

On (mis)match between distribution of training and test set. In this work we assume that the training, calibration and test data are i.i.d. (independent and identically distributed). Logit scores use calibration data to estimate the threshold for conformal prediction, which is used on the test data to create prediction sets. Thus, we evaluate both scores under the same conditions on the data.

If we anticipate distribution shift in the test set, we could modify the CP-OPT objective to use distributionally robust optimization (DRO) techniques, so that it is robust to distribution shifts. This is beyond the scope of our work and could be interesting to explore in future.

We hope our response resolves the queries. We are happy to answer any further questions you may have.

We appreciate the detailed feedback and recognition of the clarity and broad applicability of our work. Our response is as follows,

Relationship to other UQ methods. The reviewer correctly points out that there are other methods for quantifying uncertainty in the context of LLMs besides conformal prediction (CP), including methods for estimating and calibrating confidence. Our goal in this paper was to generate subsets of answer options with guaranteed coverage probability, which makes conformal prediction a natural framework. To the best of our knowledge, no other uncertainty quantification technique provides a similar guarantee. We emphasize that small prediction sets are valuable in and of themselves because they result in lower query costs, so this reduction of the space of answer options is an important feature of our procedure. [Please see the Small conformal prediction sets reduce costs section in the response to reviewer z4Sk.]

We believe the flexibility in the choice of score function is part of the appeal of CP and CROQ. We agree that it will be important to investigate how CROQ works with other choices of score functions, which we leave for future work.

Choice of $\alpha$. Please see the Choice of $\alpha$ section in our response to reviewer z4Sk.

Theoretical guarantees. We appreciate the reviewer pointing out that the relationship between Proposition 2.1 and our procedure may not be clear. We have reframed this proposition and its proof slightly to make it clear that our procedure enjoys this coverage guarantee. Regarding the optimization procedure for CP-OPT, we have pointed out (in lines 222-224, second column) that the empirical surrogates converge almost surely to their population counterparts. We have reworded the text to point out that this also holds for the cross-entropy term $\widehat{C}(g)$. We have also added intuition regarding the relationship between problems (P2) and (P1). We defer a more formal convergence analysis for future work.

Importance and Generality of the MCQ setting. Please see the same section in our response to the reviewer p4om.

Computational cost. As implemented in our paper, CROQ requires two queries to a given LLM. However, as noted by reviewer grWo, it's possible to cache the input such that the difference in cost between one query vs. two will be relatively minimal. Furthermore, as discussed in the Small conformal prediction sets reduce costs section in the response to reviewer z4Sk, it is possible to use a very cheap method to generate the scores such that the cost (both computational and literal) primarily derives from the query which produces the MCQ answer. For example, the user could use a pre-trained semantic embedding and then use the cosine similarity between the query and each response option as the conformal score. This cost is minimal compared to the cost of an MCQ query, and in general, we expect that it will be more than offset by the reduction in query cost due to the reduction in the number of answer options. Methods that involve self-consistency or self-refinement will, in general, require multiple queries, and therefore, we expect them to be more expensive, but we leave a full investigation for future work. However, we have added a discussion along the lines of the above to the paper appendix.

Magnitude of CP-OPT gains vs. logits. While the advantages of CP-OPT over logits in terms of set sizes and accuracy are numerically small, we believe that the real-world difference can be substantial at scale when large numbers of users are querying an LLM repeatedly. The figures and tables in the appendix aim to provide a finer-grained view of when and how CP-OPT improves over logits. For example, we observe that CP-OPT, in general, yields more sets of size 1 than logits and that the accuracies vary as a function of set size (which indicates that set size is a good measure of overall uncertainty). We hypothesize that CP-OPT improves over logits precisely where logits are poorly calibrated, i.e., where an LLM is over- or under-confident. We plan to test this hypothesis in future work.

We hope our response resolves the queries. We are happy to answer any further questions you may have.

We thank the reviewer for careful attention to the paper. Our response to the queries is as follows,

Choice of $\alpha$. We set $\alpha$ to a single (fairly arbitrary) value of $0.05$ in the main body of the paper for simplicity of exposition, but $\alpha$ can be treated as a hyperparameter and tuned to maximize accuracy. We discuss this in the Appendix in section B.1, with results in Figures 4 and 5. As discussed in lines 73-83 (first column), we agree that $\alpha$ represents a tradeoff whereby larger values result in smaller sets with some probability of excluding the correct answer. If we set $\alpha = 0$, then we'd include all the answer options each time, which recovers the original MCQ setting, so in some sense, we've simply parameterized a tradeoff that we can then optimize for a downstream criterion of interest like accuracy.

Evaluation criteria. We appreciate the reviewer's point. The additional figures and tables in the appendix provide a finer-grained view of (1) the distributions of answer set sizes produced by our procedure, and (2) accuracy conditional on set sizes. In addition, we have added tables that illustrate how often the CROQ procedure causes the LLM to "switch" from incorrect to correct answers and vice versa. We do indeed observe that our CROQ procedure causes the LLM to incorrectly answer some small proportion of questions which it initially answered correctly, but this is more than offset in general by answers which it initially got incorrect and which it gets correct after CROQ.

Small conformal prediction sets reduce costs. In addition to the goal of improving accuracy, we note that smaller set sizes are generally desirable in and of themselves because fewer tokens in the prompt means lower query costs. This difference can be substantial in settings like text-to-SQL or other tool usage/API selection settings, where the text describing a given answer option can be extremely large. While in our experiments we used the same LLM both to generate the conformal scores for the purposes of constructing the conformal prediction set and to generate the final MCQ answers, it's possible to use a model cascade such that a small/cheap model is used to generate the conformal prediction sets and then the MCQ is passed to a larger/more expensive LLM to generate an answer. In ongoing experiments in a text-to-SQL setting, we observe that by using a model cascade like this, we are able to substantially reduce the overall query cost while preserving or improving downstream accuracy. Once again, these tradeoffs can be optimized by tuning $\alpha$. We have added a section to the appendix and a small reference in the main text discussing this. (In addition, please see the existing section B.2 in the appendix for some discussion of model cascades and cost tradeoffs.)

Why Conformal Prediction (CP) is a useful framework. The idea of appending the conformal set to the original query is interesting and we believe would be worth investigating in future work. Our goal with the CROQ procedure however was to reduce the amount of uncertainty in MCQ-type queries, rather than to inform the LLM about its own uncertainty per se. We aim to reduce the likelihood that the LLM will be "distracted" by an available incorrect answer by pruning those answers (with high probability). This procedure is motivated by the simple empirical observation in Figure 1 that LLMs are more likely to answer correctly when there are fewer distractor options. Regarding theoretical justification, we emphasize that our procedure satisfies a coverage guarantee (Proposition 2.1), which means that the correct answer will be inadvertently removed at most a proportion $\alpha$ of the time. As discussed above, this is some sense simply a generalization of the vanilla MCQ setting with a parameterized tradeoff that can be optimized.

We hope our response resolves the queries. We are happy to answer any further questions.

We appreciate the thoughtful review and the noted strengths on presentation, problem motivation, empirical evaluation, and results of our paper. Our response to the queries and comments is as follows.

Input caching to make re-prompting efficient. We thank the reviewer for the suggestion to use input caching to reduce the inference time in the second round of CROQ. Since the part of the question without the options remains the same, we can definitely cache the inference output on tokens in the question and re-use this in the next round's inference with a reduced set of options. We plan to implement and release this soon.

CP-OPT requires training a small neural network on a dataset. As noted, logit scores can be obtained off-the-shelf from the language model. However, they can be unreliable and could be less effective for downstream use cases such as ours. To improve the quality of the scores for the application at hand, it might be necessary to tune the scores accordingly. Note that our procedure CP-OPT to tune scores is light-weight — it only requires training a small 2-layer neural network $g$ on features extracted from the LLM, and the inference overhead of $g$ is negligible compared to the LLM inference which is also needed for logit scores. Thus, while there is an extra compute cost with CP-OPT, it is insignificant in comparison to the LLM inference cost, and this cost pays off with improved performance. It also pays off with reduced set sizes, which means lower query costs. Please see the paragraph Small conformal prediction sets reduce costs in our response to reviewer z4Sk.

Re-use $g$ across models and datasets. We appreciate the thoughts to re-use the $g$ to mitigate this minor cost of re-training it for different datasets and models. On sharing $g$ across LLMs, we agree with your view on this. Since $g$ is trained on features on features from an LLM, using it on features from another LLM will likely not work. Moreover, different datasets may differ in features and thus a $g$ trained on one dataset may not work well on another dataset. We make two final remarks on this,

i) The cost of training $g$ is insignificant compared to the LLM compute costs.

ii) Reuse of $g$ could be achieved for instance by incorporating multiple models and datasets while training $g$ or first learning some model invariant representations that are fed to $g$. Exploring this could be an interesting future work.

W3. $\lambda_1$ corresponds to the weight decay hyperparameter.

Discussion on choices of hyperparameters for CP-OPT. Yes, there is a small cost of tuning the hyperparameters. We select the hyperparameters by observing the performance on the validation data. We have included a more detailed discussion on this in the paper.

Equation (2) clarification. Thanks for your careful attention. We believe there is some confusion due to the interpretation of scores. Some of the canonical works (Angelopoulos & Bates, 2022) in conformal prediction have used non-conformity scores for $g$, i.e., lower is better. In our work, we interpret $g$ as measuring conformity, i.e., higher is better. Thus, in our setting, the threshold will be the lowest $\alpha$ quantile of the scores, hence equation (2).

Analysis of CROQ in toy setting. This is indeed very intriguing. We appreciate the reviewer's enthusiasm. We can characterize the accuracy gain and $\alpha$ if we assume the LLM satisfies monotone accuracy property as in Figure 1.

Consider a predictor (LLM) that has accuracy $f(k)$ on questions with $k$ choices. It is fair to assume that as the number of choices $k$ decreases, the accuracy $f(k)$ increases, i.e. $f$ is a monotonically decreasing function of $k$. This is also confirmed in our experiments (Figure 1). We refer to this as the monotone accuracy property of the predictor.

Now, let the initial number of options in the questions be $M$ and after revising them with conformal prediction (CP) the questions have $m<M$ choices and it is guaranteed by CP that the true answer is still in the $m$ choices for $1-\alpha$ fraction of the questions. Then, the gain in accuracy after CROQ is as follows,

$\text{Gain} = \text{Accuracy After} - \text{Accuracy Before}$ The $\text{Accuracy After} = f(m)$ times the fraction of questions for which true choice is in the revised question = $f(m)(1-\alpha)$

$\Delta(M,m,\alpha) = f(m)(1-\alpha) - f(M) = f(m) - f(M) - \alpha f(m)$. Now we can make two claims,

• If $\alpha$ is fixed, then we should see improvements whenever $f(m) > \frac{f(M)}{1-\alpha}$.

• If $\alpha$ is not fixed, then the gain $\Delta(M,m,\alpha) > 0$, for any $\alpha < \frac{f(m) - f(M)}{f(m)}$. By the monotone accuracy property of the predictor $f(m) - f(M) >0$, that means any $\alpha \in (0, \frac{f(m) - f(M))}{f(m)})$ will yield a gain in accuracy.

---

We hope our response resolves the queries. We are happy to answer any further questions.

Thanks for the feedback. Our response to the queries is as follows,

Importance and generality of the MCQ setting. The multi-choice question-answering framework encompasses any setting in which an LLM must select from among a finite number of options. We believe that this describes many if not most steps in agentic workflows. This includes selecting function calls or APIs, selecting among UI elements, selecting databases or tables in text-to-SQL settings, selecting which agent to pass an output to next, selecting the next step in a plan, selecting a document from a database for RAG, selecting a set of in-context examples from a database of such examples, selecting diagnosis codes or labels, etc. The answer options do not necessarily need to be defined in advance: in open-ended response settings, for example, it's possible for example to have an LLM generate a shortlist of initial candidates which can then be pruned before being passed to other agents. We agree that it is an interesting open question how to extend our framework to the open-ended response setting, but we believe that the MCQ abstraction is widely applicable on its own.

Comparison with prompting-based methods. Our work's focus is to evaluate the hypothesis that conformal prediction based pruning can be effective in improving accuracy. Comparison with prompting-based methods would be interesting future work.

Other advantages of pruning answer choices. In addition to improving accuracy, pruning answer choices can lead to computational and dollar savings. This can occur when a score function is used in the conformal procedure that is cheap to compute relative to the cost of the MCQ query. For more details, please see the section Small conformal prediction sets reduce costs in our response to reviewer z4Sk, and the section Computational cost in our response to reviewer uJFE.

On conditional uncertainty given the prompt: In general conditional calibration is not possible in conformal prediction (Barber et al. 2020). Due to this limitation, conformal prediction with marginal coverage guarantee is widely used and it is suitable for our work where we aim to evaluate the hypothesis that conformal prediction based pruning can be effective in improving accuracy.

Further, the logit scores are generated by running the forward pass of the LLM on the entire prompt including all the answer options; these scores then serve as features, along with the other features described, for the CP-OPT procedure. In that sense, both the logit scores and our CP-OPT scores represent the uncertainty conditional on the prompt, and the resulting sets reflect this uncertainty.

Regarding realistic decision-making scenarios, our method can be applied to any setting in which an LLM must choose from among a finite set of response options, whether those options represent question answers, APIs, function calls, downstream agents which can receive input, etc. We refer the reviewer to the section Importance and generality of the MCQ setting above.

We hope our response resolves the queries. We are happy to answer any further questions you may have.