Learning to Route LLMs with Confidence Tokens

Thank you for the valuable feedback and reference. We've incorporated them into our related work and made the suggested expository improvements in our revised paper.

[Q-1] More analysis on a systematic investigation of failure cases to identify patterns in when and why self-REF misclassifies confidence.

[A-1] This is a good suggestion, we previously included some case studies of successful routing in Appendix D, and will add more failure cases as well. Analyzing the top 5 categories of correct/incorrect and overconfident/underconfident predictions, we have:

• Predict Correctly+<UN> (underconfident in its predictions): (1) computer_security; (2) high_school_biology; (3) high_school_european_history; (4) human_sexuality; (5) miscellaneous
• Predict Wrongly+<CN> (overconfident in its predictions): (1) college_computer_science; (2) conceptual_physics; (3) high_school_computer_science; (4) high_school_microeconomics; (5) jurisprudence
• Predict Correctly+<CN>: (1) international_law; (2) college_biology; (3) moral_disputes; (4) philosophy; (5) us_foreign_policy
• Predict Wrongly+<UN>: (1) abstract_algebra; (2) anatomy; (3) college_chemistry; (4) college_medicine; (5) econometrics

Analyzing these categories, we note:

• Predict Correctly+<UN> (underconfident in its predictions): The model often hesitates with context-heavy subjects due to their ambiguity and need for broader reasoning. "miscellaneous" reflects general uncertainty in non-standard topics.
• Predict Wrongly+<CN> (overconfident in its predictions): Many of these involve technical, structured domains, where the model shows overconfidence, likely due to familiarity from training. However, it may struggle with edge cases and nuanced reasoning, especially in areas like jurisprudence.
• Predict Correctly+<CN>: These topics rely on broad conceptual knowledge rather than strict calculations, and the model appears well-calibrated, likely due to strong training data coverage or clearer signals of correctness.
• Predict Wrongly+<UN>: These highly specialized, detail-heavy subjects require precise recall or deep understanding, and the model’s uncertainty may reflect an awareness of its limitations in recalling detail-heavy content to answer questions.

[Q-2] What data did you use in the setup when you mention you trained on MMLU? In general, i'd move some information about training the confidence tokens ...

[A-2] We use a randomly sampled subset of the official MMLU training set, where ground truth answer choices are augmented with confidence tokens when fine-tuning Self-REF on MMLU (Algorithm 1, Section 3.2). We will move more details about the training setup from the appendix into section 4.1.

[Q-3] In figure 2 I see random route to 70B baseline that I can't find it described. Can you elaborate on it?

[A-3] The random routing approach is a naive baseline that uniformly at random routes to the 70B model at the specified rate. When the random routing rate is 0.0, then it gives the performance of the small LM, and when the random routing rate is 1.0, it gives the performance of routing to the 70B model entirely. We will update our paper with these details.

[Q-4] I am slightly confused by the quantile thresholds in Section 4.2 and their relation to the routing rate in Figure 2. Specifically, how do these two quantities interact? for example, in Figure 2(a), ...? More generally, it would be helpful to explicitly discuss the relationship between t and routing rate, as ... trade-offs between accuracy and efficiency.

[A-4] To better analyze the routing performance, we set the thresholds t at 20-quantiles, as described in Section 4.2. For instance, the threshold t corresponding to a routing rate of 0.4 is the 40th percentile of confidence scores across all input queries. Practically, one could observe the 20-quantiles of the distribution, and choose thresholds that should empirically correspond to rough routing rates, and then monitor these thresholds over time. Note that Self-REF fine-tuning does not require committing to a particular tradeoff, as instead of relying directly on the sampled tokens, we instead extract confidence scores from the logits of the confidence tokens, thus giving us the ability to threshold this score and control how often to route.

[Q-5] Can you clarify what you mean with the following sentence about the in-context learning for llama3: "All experiments utilize Llama3-70B-Instruct with only its strong in-context learning capabilities during instance routing..."

[A-5] This refers to the fact that we did not fine-tune Llama3-70B-Instruct (the larger, more expensive LLM) in the routing setting.

[Q-6] Missing reference, formats, and Typos.

[A-6] Thank you for the valuable comments and references. We will add this to our related work section on rejection learning. The suggestions for formats and typos are fixed in our next version.

Thank you for the detailed feedback. We respond to each point below, and will update the discussion accordingly:

[Q-1] The assumption that <CN> tokens represent correctness and <UN> represent incorrectness lacks a clear theoretical grounding. The paper implicitly assumes confidence directly correlates with correctness...

[A-1] A core contribution of this paper is utilizing confidence scores for rejection and routing downstream. Unlike prior works that derive uncertainty scores from logits, verbalized outputs, or re-sampling, we train our notion of confidence into the model, teaching it to reflect on its answer after producing it. The notion of confidence is grounded in the probability that a prediction is correct, i.e., $P(\text{<CN>} |X,Z) = P(Y=Z|X,Z)$, where CN is the confidence token, $X$ is the question, $Z$ is the predicted answer, and $Y$ is the true answer. This is a more direct form of confidence than other techniques, which often instead reflect consistency in the answer. For example, logits would instead reflect $P(Z|X)$, with no relation to $Y$. As noted in our paper (Section 1 and 5.3), calibrated uncertainty metrics are not necessarily correlated with correctness. In our paper, we assess both the utility of the confidence scores as well as their calibration.

[Q-2] If the base model has substantial weaknesses or biases, these may propagate through fine-tuning, limiting improvements or perpetuating biases.

[A-2] We respond using two perspectives:

• The primary goal of Self-REF is to help the base model identify its existing weaknesses while maintaining performance, thereby enabling more effective routing.
• One possible approach for extreme incorrectness is as follows: first, fine-tune the base model on the downstream task to improve its task-specific capabilities; then, in a second stage, apply the Self-REF framework to teach the model to route effectively based on uncertainty. This is a useful extension of Self-REF we will add to the discussion.

[Q-3] Explore other methods for labeling confidence (consistency across sampling and external calibration).

[A-3] As mentioned in Q-1, the goal of Self-REF is to produce confidence tokens useful for downstream settings such as routing and rejection, boosting correctness of the overall system. Well-calibrated confidence scores do not always correlate with correctness [1, 2] (Section 1 and 5.3). Two toy examples to explain this misalignment intuitively: Assume we have a binary classification task with predicted probabilities (see below). Example 1 achieves lower accuracy but has a lower ECE score, whereas Example 2 achieves higher accuracy but a higher ECE score. This example demonstrates that calibration metrics are not correlated with the correctness of the prediction. Similarly, consistency-based uncertainty is not necessarily aligned with downstream correctness, e.g., when one has highly consistent incorrect answers. This misalignment can degrade routing performance when such signals are used for routing.

Bins = 2
predicted prob. of "1" = [0.5, 0.5, 0.5, 0.5]
ground truth = [0, 0, 1, 1]

Example 1: 
predicted prob. of "1" = [0.5, 0.5, 0.5, 0.5]
-> ECE(↓)= 0%
-> accuracy (↑)= 0%

Example 2: 
predicted prob. of "1" = [0.4, 0.6, 0.9, 0.9]
-> ECE(↓) = 20%
-> accuracy (↑) = 75%

[1] Huang, et al. "Look before you leap: An exploratory ..." arXiv 2023.

[2] Yona, et al. "Can Large Language ... Uncertainty in Words?." arXiv 2024.

[Q-4] I'm curious if authors explored a baseline where the model generates a confidence token before producing an answer. Also a learned router baseline would provide insight.

[A-4] To the best of our knowledge, there is no existing baseline that generates a confidence token prior to producing the answer. In multiple-choice QA, token generations are relatively short, making it feasible to predict a confidence token after producing the answer. This allows the model to condition its confidence estimation directly on the generated output, potentially resulting in more accurate confidence scores for routing. However, for tasks involving longer responses (e.g. reasoning), an alternative approach to improve efficiency could involve an early-stopping mechanism for confidence tokens. The model might produce a confidence token midway through its answer generation, allowing an earlier routing decision.

We have additional experiments with a learned confidence-based router baseline: OOD-Probe [1] to assess its routing performance from Mistral-7B to Llama3-70B in the MMLU dataset. The results are shown in the table below. We provide the accuracy in different routing ratios, and observe that Self-REF outperforms the new baseline.

[1] Mahaut, et al. "Factual confidence of LLMs:..." ACL 2024.

Thank you for the valuable feedback and ideas. We have updated our discussion accordingly.

[Q-1] Your framework is only able to predict the confidence token at the end of the answer? Is that correct? If so, I see the drawback in terms of routing capabilities. Since the model need to generate the full response before going to route the response to another larger LLM? Do authors see any way to get the confidence tokens before the full answer is generated using the proposed framework?

[A-1] Thank you for the insightful feedback. In question-answering tasks, the cost of generation is relatively low, and predicting a confidence token after generating the answer allows the model to condition its confidence on the output, potentially leading to more accurate uncertainty estimation for routing. However, in other long-response tasks, such as LLM-based reasoning, a promising direction could be to develop an early-stopping approach for confidence token generation. The model may emit a confidence token partway through answer generation in this setup, enabling an earlier routing decision to stronger models. This strategy could be particularly beneficial in complex reasoning scenarios, and we consider it a valuable direction for future work.

Thank you for the valuable feedback and reference. We've incorporated them into our related work and made the suggested expository improvements in our revised paper.

[Q-1] Self-REF requires a dedicated training/validation set and a fine-tuning stage for different downstream tasks/datasets, which limits its practical usages compared to zero-shot baselines.

[A-1] Our method targets the scenario where one would like to optimize performance for a particular downstream task while (1) being able to fine-tune smaller specialized models, (2) reducing monetary or latency cost, and (3) maintaining comparable performance on the downstream task. Thus, we agree that fine-tuning is necessary in self-REF. However, we argue that it is often the case that one may want to fine-tune a smaller model for one's specific task of interest. Exploring broader use cases is a valuable direction for future work. An interesting extension would be to explore how self-REF can extend to multi-task training for general use cases.

[Q-2] It is unclear how much transferability Self-REF holds beyond similar tasks such as OpenbookQA --> MMLU. It would be very interesting to see how Self-REF performs when fine-tuning in a multi-task setup with a collective of datasets.

[A-2] We agree that fine-tuning self-REF with a mixture of tasks augmented with confidence tokens is a promising extension. However, we would like to emphasize that the nature of uncertainty can vary significantly across tasks. For instance, uncertainty in question-answering may differ from that in code generation or other complex reasoning tasks. While QA tasks often have well-defined ground truth, other tasks may require soft confidence labels (tokens), which could affect the effectiveness of confidence-based routing. We value and acknowledge the potential of this direction and will consider incorporating a broader set of datasets under a multi-task framework in future work.

[Q-3] How was the optimal parameter $\alpha$ for the <UN> data determined?

[A-3] $\alpha$ is a hyperparameter in our work, and we select the optimal $\alpha$ based on the validation set.

[Q-4] How does the ratio of <CN>:<UN> in dataset construction affect the results?

[A-4] We fixed the overall dataset size and treated the ratio between <CN> and <UN> samples as a hyperparameter to optimize performance for a particular downstream task with fine-tuning. On one hand, including more <CN> samples helped fine-tune the model to better perform the downstream task; on the other hand, a sufficient number of <UN> samples was necessary to teach the model when to express uncertainty. An imbalance—either too many <UN> samples, which could degrade task performance, or too many <CN> samples, which could reduce the quality of uncertainty estimation—can negatively affect the model. To balance this trade-off, we experimented with ratios of 1:1, 1:2, 1:3, 1:4, and 1:5 and selected the optimal setting based on each task’s validation set.

[Q-5] In Figure 2: why do different models/baselines/datasets have different number of datapoints? For instance, Mistral only has 3 data points for the "Verbalizing uncertainty" baseline curve on MMLU.

[A-5] The points correspond to 20-quantiles, as described in Section 4.2, which ideally yield 20 distinct and equally spaced routing rates, assuming all confidence scores are unique. However, certain methods for extracting confidence scores, such as verbalizing uncertainty, can suffer from mode collapse, where the scores cluster around a limited set of numbers. As a result, a substantial portion of the confidence scores is identical, leading to fewer than 20 distinct routing rates in practice.

[Q-6] Missing reference, formats, and Typos.

[A-6] Thank you for the valuable comments and references. We will add this to our related work section on rejection learning. The suggestions for formats and typos are fixed in our next version.