Learning to Route LLMs with Confidence Tokens
Authors: Yu-Neng Chuang
Prathusha Kameswara Sarma
Parikshit Gopalan
John Boccio
Sara Bolouki
Xia Hu
Helen Zhou
Venue: ICML 2025 poster (2025)
Keywords: Large Language Model
LLM Routing
Model Confidence
Primary area: social_aspects->accountability_transparency_and_interpretability
TL;DR: This work introduces Self-REF, a training strategy that teaches large language models to reliably express confidence in their answers, leading to improved performance in downstream tasks like routing and rejection learning.

Abstract: Large language models (LLMs) have demonstrated impressive performance on several tasks and are increasingly deployed in real-world applications. However, especially in high-stakes settings, it becomes vital to know when the output of an LLM may be unreliable. Depending on whether an answer is trustworthy, a system can then choose to route the question to another expert, or otherwise fall back on a safe default behavior. In this work, we study the extent to which LLMs can reliably indicate confidence in their answers, and how this notion of confidence can translate into downstream accuracy gains. We propose Self-Reflection with Error-based Feedback (Self-REF), a lightweight training strategy to teach LLMs to express confidence in whether their answers are correct in a reliable manner. Self-REF introduces confidence tokens into the LLM, from which a confidence score can be extracted. Compared to conventional approaches such as verbalizing confidence and examining token probabilities, we demonstrate empirically that confidence tokens show significant improvements in downstream routing and rejection learning tasks.

Learning to Route LLMs with Confidence Tokens
Yu-Neng Chuang" *1 Prathusha K. Sarma?
Parikshit Gopalan"
John Boccio
Sara Bolouki? 2
Xia Hu
Helen Zhou?
Abstract
2023; Xiao et al., 2023), and provide customer support (Ko- lasani, 2023). As LLMs are given more agency in settings of increasing consequence, becomes crucial to know when an output is reliable. Given knowledge of the trustworthi- ness of an output, systems can proactively prevent faulty behavior by seeking second opinion (e.g. by routing to more costly LLM), or choosing to abstain from answering (e.g, instead defaulting to safe behavior).
Large language models (LLMs) have demon- strated impressive performance several tasks
and are increasingly deployed in real-world ap- plications. However, especially in high-stakes settings, becomes vital to know when the out- put of an LLM may be unreliable. Depending on whether an answer is trustworthy, system can then choose route the question another expert, safe default behavior.
However, state-of-the-art LLMs face challenges with pro- viding an accurate estimate of confidence whether a prediction is correct. Several works identify that token probabilities (derived from softmax over logits) not well-aligned, or calibrated, with the actual probabilities of correctness (Huang et al., 2023; Detommaso et al., 2024; Wightman et al., 2023). Others have found that modifying the prompt by asking the LLM to verbalize its confidence yields slightly better results (Tian et al., 2023; Turpin et al., 2024; Mahaut et al., 2024); however, these results may be subject to the dataset and prompt engineering, often leading to unstable or unreliable results. Since LLMs are typically trained using a cross-entropy loss, they can overfit on accu- racy rather than calibration, often leading to overconfidence and misalignment with real-world distributions. Thus, a nat- ural question arises: how can we train LLMs to accurately estimate their own confidence levels, while preserving their performance on tasks of interest?
otherwise fall back In this work, we study the extent to which LLMs can reliably indicate confidence in their answers, and how this notion of confidence can translate into downstream accuracy gains. We propose Self- Reflection with Error-based Feedback (Self-REF). lightweight training strategy teach LLMs express confidence in whether their answers are correct in reliable manner. Self-REF introduces confidence tokens into the L.M,, from which a confidence score can be extracted. Compared to conventional approaches such as verbalizing confidence and examining token probabilities, we demonstrate empirically that confidence tokens show significant improvements in downstream routing and rejection learning tasks.
1. Introduction
To answer this question, we introduce Self- Reflection with Error- based Feedback (Self-REF), lightweight framework for fine-tuning LLMS to accurately assess confidence in their predictions while preserving their performance on relevant downstream tasks. Self-REF integrates naturally with any LLM backbone, allowing the model to condition its assess- ment of confidence its generated outputs. It consists of three steps: (i) Confidence token annotation: the training dataset is generated by labeling each instance with confi- dence tokens based on the correctness of the base LLM. When the base LLM provides correct responses to the in- put text, we augment the instances with confident (<CN>) tokens. Conversely, if it responds incorrectly, we augment the instances with unconfident (<UN>) tokens. (ii) Self-REF fine-tuning: The base LLM is fine-tuned on the augmented training data, so that one would expect confidence token to be generated following its predicted output. (iii) Confidence
Recent years have seen an explosive growth in the deploy- ment of large language models, in forms ranging from help- ful conversational agents (Achiam et al., 2023; Touvron al., 2023; Yang et al., 2024), frameworks for automat- ing software workflows (Harrison, 2022; Wu et al., 2023), to tools tailored to domain-speciff tasks. LLMs have been used to summarize doctor- doctor-patient interactions (Krishna et al., 2021; Jeong et al., 2024), teach students (Schmucker et al.,
'Work done during the internship at Apple. 'Rice University, Houston, TX, USA Apple, Inc., Cupertino, CA, USA. Correspon- dence to: Yu-Neng Chuang <ynchuang @rice. edu,, Helen Zhou chelen_Lhhou @ app...com..
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning to Route LLMs with Confidence Tokens
small fine-tuned language model
<User> A football stadium has 50, 50,000 seats. If the lower tier has 20% of the seats, how many seats are in the lower tier?
if confident,
answer
50, 50,000
<UN>
<EOS>
"50,000"
query
if unconfident, route
downstream setting #1: route to larger LM (e.g., with API access)
downstream setting #2:
reject abstain (default sale behavior)
LM
"10,000"
not sure"
Figure 1: Self-REF attaches confidence token each prediction. Routing proceeds based on the confidence.
score extraction: To obtain continuous confidence score, one computes the probability of the <CN> token, normalized by the sum of the probabilities of the <UN> and <CN> tokens.
2. Related Work
In this section, we briefly introduce related studies on un- certainty quantification, LLM routing, and LLM rejection learning. More discussion and details about related work are provided in Appendix A.
Given estimates of confidence, a natural follow-up question how to utilize such estimates to yield desirable system- level behavior. Often, one may have limited computational resources and only be able to locally fine-tune smaller mod- els. However, one may have API access to a larger model with higher base performance for some cost. In this work, we focus on two downstream tasks that utilize this notion of confidence for practical applications: (1) routing, where queries with low confidence are routed stronger but more costly LLMs, and (2) rejection, where the model may ab- stain from making choice, valuable for preventing unsafe behaviors. Our contributions include:
2.1. Uncertainty Quantification in LLMs
Recent studies in LLM uncertainty have explored leveraging internal features like logits and attention scores to estimate uncertainty without additional training (Zhou et al., 2023; Mahaut et al., 2024; Turpin et al., 2024). Post Post-processing methods such as temperature scaling (Guo et al., 2017), prompting LLMs to verbalize confidence scores (Tian et al., 2023), iterative prompting (Abbasi Yadkori al., 2024), and instruction learning/pretraining (Zhang al., 2021; Neeman et al., 2023; Li et al., 2023; Liu et al., 2024; Cohen et al., 2024) have also shown improvements in calibration. In con- trast to these methods, Self-REF allows the LLM to learn to estimate its confidence at fine-tuning time. Furthermore, this internal notion of confidence is aligned against correctness rather than the consistency of responses after re-sampling, as is studied in (Yona et al., 2024; Manakul al., 2023; Lin et al., 2023). As noted in (Huang et al., 2023; Wang & Zhou, 2024; Chuang et al., 2025), uncertainty scores derived from L.Mss may exhibit a weak correlation with actual prediction correctness, and the most consistent answers are often not the most correct. Finally, while much of the prior work focuses on calibration, our work goes beyond it and focuses the downstream utility of confidence scores for LLM routing and rejection learning.
1. Introducing Selr-REF, lightweight fine-tuning strat- egy that helps LLMs better learn when they should be confident in their predictions.
2. Introducing the confidence-based routing setting, where queries that a small local L.M uncertain about can be optionally routed to more powerful LLM.
Studying the confidence-based rejection setting, where the answer may be "none of the above" (e.g., if the LLM is not confident in any of the actions, it should defer to fallback mechanism rather than choosing arbitrarily).
Demonstrating that Self-REF outperforms baselines in both the routing and rejection learning settings on four public datasets.
2

---

Learning to Route LLMs with Confidence Tokens
2.2. LLM Routing
split of the dataset, vaal the validation split, and Duest the test split, where DD trrain sDaainDDDaaUDaaaa The goal in learn- ing confidence tokens is to fine-tune M on Dtrain to use confidence tokens, two special tokens with trainable embed- dings: as assessed on Dtest, the confident token <CN> should be generated when a model is confident that its prediction is correct (i.e.. 77(*) = y(")), and the unconfident token <UN> should be generated when model is not confident that its prediction is correct (i.e., g( # g)
In confidence-based query-routing scenarios, one line of re- search trains additional classifiers to route queries to LLMs based observed performance metrics and routing data, forming routing pipelines with trained multi-agent systems or extra routers to boost performance (Ding et al., 2024;
Ong et al., 2024; Stripelis et al., 2024; Jiang et al., 2024; Zhang et al., 2023). Another approach leverages LLM cas- cades, invoking increasingly complex LLMs depending on trained scoring functions which predict whether answer is correct, and consistency of responses upon resampling (Chen et al., 2023b; Yue et al., 2023). Unlike prior work in LLM cascades, Self-REF trains the estimate of the probabil- ity of correctness into the LLM itself instead of an external model. Here, the estimates of confidence from Self-REF are aligned against the correctness of the predictions, rather than the consistency. Additionally, Self-REF computes the confi- dence in the autoregressive decoding step, allowing for LLM itself to condition on its generated output. This dynamic approach allows for more accurate and adaptive routing decisions based on real-time confidence assessments.
3.2. Self-REF
The Self-REF framework fine-tunes a base model M to use confidence tokens and also learn their embeddings. It involves (i) confidence token annotation, where data are augmented with confidence tokens, (ii Self-REFfiinetuniin finitt on the augmented data, and (iii) confidence score extraction.
Confidence Token Annotation To incorporate model feedback into our training process, we first use the base model M to generate predictions gi for each sample = ,,...,Nyyy in the dataset, where g() E V. For instances where y * g( (the prediction is incorrect), we construct an augmented dataset Dirain, SUN> = {(x:10, <U))))
2.3. LLM Rejection Learning
Learning to reject (Cortes et al., 2016) was originally in- troduced to equip the Bayes classification model with a rejection option for predictions. Despite advancements in LLMs, their predictions can still be prone to uncertainty due to inherent knowledge limitations. This highlights the need for LLMs to have the ability to refuse to answer, es- pecially in situations where it is infeasible to fallback upon more advanced models. LLMs can be trained for the re- jection task through rejection knowledge injection (Chen et al., 2024; Zhang et al., 2024), an additional rejection
where ((*)> denotes the concatenation of the incorrect prediction g(t) with an unconfident token <UN>, Conversely, for instances where the prediction is correct (g)()) = 7(0),,
we create the set D'rain, = {(x(1) ((0))))))))) (()<)))) with <<>>>>>>>>>>>> representing the correct prediction concatenated with a confident token <CN>, The augmented training dataset is then formed by combining these two datasets. This aug- mentation strategy enriches the training data by explicitly encoding the correctness of the model's 's predictions, thereby providing additional supervision that can guide the model to better distinguish between correct and incorrect responses. The detailed steps of confidence token annotation are shown in Algorithm 1.
LLM agent (Mao al., 2023; 2024), and rejection-oriented loss adaptations (Li et al., 2024; Mohri et al., 2024). How- ever, such approaches involving extra knowledge injection or loss adaptation may degrade the base performance of the LLM. In this work, we aim to equip L.MSs with the capa- bility for rejection learning based on their confidence levels, without the need for designing new loss functions, thereby preventing any potential degradation in the performance of downstream tasks.
Self-REF Fine-tuning Next, the base model .M is fine- tuned using Pirsin- The confidence tokens <UN>and <CN>are special tokens during training, initialized the average of other existing token embeddings. Gradients from the incorrect answers in the unconfident samples are masked out, to avoid increasing the probability of an incorrect answer given a query P(7 * y())) z..)). Instead, the model learns to increase the probabilities of confident tokens given correct answers P(<CN>| 7 = ,,(0), ()()),, probabilities of unconfident tokens given incorrect answers P(<UN> | ((0) # y(")) (.)))), and probabilities of correct answers given the query P( y (1)))
3. Learning Confidence Tokens
3.1. Problem Setup and Notation
Consider any local large language model M X - y with input query domain X' and output prediction domain ), Here, local refers to a model that one has the ability to fine- tune (as opposed to API access). Let D = (xx(0) ) denote dataset, where 2((i) are queries and ( are ground truth answers to the queries. Let Derain denote the training
Note that the unconfident and confident tokens are gener- ated end-to-end conditioned on both the prefix query and the answer that the LLM itself generated, in contrast to prior work (Ong et al., 2024; Ding et al., 2024; Stripelis et al.,
3

---

Learning to Route LLMs with Confidence Tokens
Algorithm 1 Data Augmentation for Self-REF
Confidence-based Rejection A larger LLM may not be accessible in every situation. If one does not have access larger more capable model, one may still want to utilize measures of uncertainty in order to decide whether to rely on the model's prediction. This experimental setup studies the following scenario: "Sometimes, none of the provided options are good. How good am at saying that I am not confident in any of the answers?"
input Base LLM M(.) and original training data Pursin ((()), p)
output Augmented data D' train with unconfident samples and confident samples. Use MM-) to make predictions g( for the given input query ,.(()
1: Create a set of unconfident samples:
Dtrain,aune auNe {xx1(0 g() <UN>): y0) + T0) ) Nuis
To study this question, we create an evaluation set where half of the samples do not contain the ground truth, i.e., we remove all ground-truthh information from x()), and replace the label with "none of the above," " y = 0. Ideally, the model's 's confidence in any of the remaining answers should be low, and we can evaluate how good policy for abstaining from answering the question against this dataset. Thus, if M(-) has low confidence, i.e., (x((z()) (() < t for chosen threshold / E [0, 1], we treat it as if the model were to abstain from answering (i.e., g( = 0).
Ji=l
2:
Create set of confident samples:
D Pirmin,cos <CN> {xx(1) g(0<CN))) y(t) g(t) Ji=l"
3: Mix the unconfident and confident data to form the com- bined augmented dataset, subsampling the unconfident samples with tunable proportion E [0, 1]:
D Dirain- subsample, (Pirsin, ) U D' Pirain,<oN=
4. Experiment Setup
4.1. Datasets and Baselines
2024) which typically routes based the query alone with an extra router. By fine-tuning with the LLM's own predic- tions on the augmented training set, Self-REF customizes the notion of uncertainty as the indicator for routing and re- jection to the LLM itself, rather than relying on the inherent uncertainty of a query.
Datasets
All experiments are conducted on the following four public datasets (more details in Appendix B):
MMLU (Hendrycks et al., 2021a;b): The Massive Multitask Language Understanding (MMLU) dataset consists of multiple-choice questions covering a wide range of knowledge domains with 57 distinct tasks.
Confidence Score Extraction After training, one can extract a continuous confidence score ( E [0, 1] by getting the token probability of <CN>, normalized over the sum of the <UN> and <CN> token probabilities: (M() )() = P(<nD)+ P(cun) This continuous score gives us some control over navigating different tradeoffs in downstream settings.
OpenbookQA (Mihaylov et al., 2018): The Open- BookQA dataset is a quuestion-answerme dataset de- signed to test deep understanding through multi-step reasoning, common knowledge, and text comprehen- sion, similar to open-book exams.
P(<De)
GSM8K (Cobbe et al., 2021): GSM8K is dataset containing graduate school math questions, where each question is collected from the Math World Problem Repository (Roy & Roth, 2015).
3.3. Learning to Route and Reject
Confidence-based Routing to a Larger Model Often, constraints on computing resources and data may prevent researchers from fully fnn-ttnnin very large LLMs. How- ever, suppose one has the ability to fine-tune smaller local LLM, as well as API access to very large LLM. Given the limited capabilities of smaller L.MSs and the high cost of very large LLMs, can we effectively route queries to the larger model based on the confidence level of the smaller model? How does one trade-off between cost and accuracy?
MedQA (Jin et al., 2021): MedQA is multiple-choice open domain question -arioonn--nwwringgg dataset for problems collected from professional medical board exams.
Baselines We compare Self-REF four state-of-the-an baselines for measuring model confidence, following the settings from existing work. Details about the prompts and hyper-parameters utilized in these baselines are given in Appendices B and C. The baselines are listed as follows:
In confidence- based routing, answers where the smaller fine- tuned LLM is uncertain, i.e., (M) ()) < L for some chosen threshold & E [0, 1]. will be routed to a more pow- erful LLM Mragge (.). For answers with cM (z((), g(t) 2
(i.c., certain), the answers from the local LLM M(-) will be directly returned to the user.
Verbalizing Uncertainty (Tian et al., 2023): Prompt the LLM with in-context learning to yield confidence scores

---

Learning to Route LLMs with Confidence Tokens
Liama3-4B 4B MMLU
Lbmaa2-0 OpenbookgA
Llama3-0p a-8 6SMIK
Llama3-08 Medgn
3.22
3.000
2.75
0.723
2.75
6.4
0.43
Rack
3.50
035
12
Routing Rate
Routing Rate
Scuring Rate
Routing Rare
(a) MMLU
(b) OpenbookQA
(c) GSM8K
(d) McdQA
Mistral 78 MHLU
Mistral 78 OpenbookOA
Mistral- 78 GSMOK
Mistral 78 MedoA
0.45
077
1.60
2.5
0.55
LONK
0.15
Sal--gT LONK
Leik
0.55
Zaroahot Laght
025
3.0
02
24
64
10
3.0
0.2
244
0%6
0.0
10
0.2
04
26
08
0.10
245
0.8
Bousing
(e) MMLU
(f) OpenbookQA
(g) GSM8K
(h) MedQA
Figure 2: Overall accuracy vs. routing rate in confidence- routing task. Routing is from local LLMs Llama3-8B- Instruct (row 1) and Mistral-7B-Instrues (row 2) to Llama3-70B-Insteen on MMLU (column 1), OpenbookQA (column 2). GSM8K (column 3), and MedQA (column 4). Routing rate of O gives pure local LLM performance, and routing rate of 1 gives pure Llama3-70B-Instee uniformly at random from the local fine-tuned LLM to the large LLM. Self-REF consistently outperforms baselines.
performance. The solid black line corresponds to baseline which chooses to route
between 0 and 1. This encourages the LLM to generate its response while simultaneously estimating the confidence level of the response.
4.2. Experiment Settings
Confidence-basser Routing Task In the routing task, we aim to accurately route instances that the smaller LLMs (Llama3-8B-
Verbalizing Yes/No Tokens (Tian et al., 2023): Prompt the LLM with a question of whether the model is confi- dent in its answers. The confidence score is calculated by normalizing the probabilities of the "Yes" and "No" tokens by their sum.
and Mistral-7B-IInstrucl struggle with to the larger LLMs -Jlama3-7B--Insrrude for improved performance. All experiments utilize Llama3-70B Insttrtt with only its strong in-context learning capabilities dur- ing instance routing (no further fine-tuning). The routing decisions are determined by thresholding the normalized probability of the <CN> confidence token over the sum of the probabilities of <CN> and <UN> generated from Self-REF:
Zero-shot Logits (Mahaut et al., 2024): Extract the prob- ability of the generated prediction from the LLM without fine-tuning. In free-text generation questions, we average the probabilities of the tokens from the output answers. In the closed-form generation questions with single-token swers (e.g., multiple-choice questions), we directly select the probability of the predicted token.
To characterize the trade-off between cost and performance, we analyze performance along several possible routing thresholds t. In particular, for a set of confidence scores CM.D = (cm(((()) g ) we set at quantiles, where = 2p (CM,D) for p 1, 203. Instances are routed to the larger LLM when (2 g()) so these thresholds correspond to routing upproximately 0%, 5%, 10%, .... 100% of total queries to the larger L.L.M
Fine-tuned Logits (Liu et al., 2024): Extract the proba- bility of the generated prediction from a fine-tuned LLM. All methods for computing the probabilities of the out- puts identical to the "Zero-shot Logits" baseline. The confidence score is determined by the probabilities.
Confidence-basee Rejection Learning Task The rejec- tion learning task is evaluated on the MMLU and Open- bookQA multiple-choice datasets. Self-REF does not ex- plicitly train for the rejection learning task, as no additional samples are added where "None of the above" is a choice. This experiment tests whether the notion of confidence en-
Taking these continuous- confidence scores from each baseline, one can apply thresholds and compare the utility of these scores versus those of Self-REF in the downstream routing and rejection learning settings.
5

---

Learning to Route LLMs with Confidence Tokens
Llama3-88 MMLU
Rejaction: Llama3-88 OpenbookOA
Rejection: Mistral-7B MMLU
Rejaction: Mistral-78 OpanbookOA
Thue Postive Rata
Thee Rositvve Rate
Rosttive Rote
0.8-
6.6
Gread (46b
Thue
sattiie parr
0.0
0.2
0.
0.6
Positive
False Positive
False Positive Rate
(a) MMLU
(b) OpenbookQA
(c) MMLU
(d) OpenbookQA
Figure 3: ROC curves of Llama3-8B-Instruee -Bs-ruutt on MMLU (a) and OpenbookQA (b); and Vistral-7B-Insteee on MMLU (c)
and OpenbookQA (d), for the rejection learning task. The true positive rate (recall) indicates the proportion of samples
where "none of the above" is the ground truth that is correctly rejected. The false positive rate is the proportion of original
samples that have been falsely rejected. Across all settings, Self-REF outperforms other baselines in learning to reject
coded by <UN> and <CN> may nevertheless be useful for determining whether to reject all provided options.
5.1. Routing Performance (RQ1)
We analyze the routing performance from both a system- level accuracy and efficiency perspective.
To assess performance on the rejection learning task, we construct specialized evaluation set where 50% of the time, the correct answer choice is removed from the choice list in the input question. Ideally, for any question where the ground truth choice was removed, the Self-REF fine-tuned LLM should append the <UN> token to any selected answer, since it would be incorrect. In order to examine the trade- off between rejecting too many samples vs. providing too many incorrect answers, we again utilize the (z g()) confidence score with several possible thresholds t [0, 1].
Overall Accuracy Confidence-based routing using SeIf-REF consistently achieves the best accuracy vs. routing rate trade-off for all four datasets (MMLU, OpenbookQA, GSMSK, MedQA) and both local LLMs (Llama3-8B. Instruct model and Mistral-7B-Insttenen for MMLU, confidence-
(Figure 2). Using
Self-REF on Alama3--B-Instte based routing of just 39% of queries can achieve comparable performance to routing to the Llama3-70pP 70B-Instruct model alone. Similar observations hold for the OpenbookQA, GSM8K, and MedQA datasets, where routing the Llama3- IB-trutt model with rates 49%. 65%, and 40%, respec- tively, can yield comparable performance to Llama3-70B- Instruct alone. For Self-REF on Mistral-7B- BB-Inntrutt the minimum routing rates to achieve parity with Llama3-70B- Instruct are 70%, 50%, 75%, and 70% for MMLU, Open- BookQA, GSM8K, and MedQA, respectively. Among the baselines, the fine-tuned logits tend to perform second-best.
Implementation Details We select Llama-3-8B- Instruct (Touvron et al., 2023) and -Bsvral-7B-v0.3- Instruct (Jiang et al., 2023) models as the pre-trained local base LLMs for Self-REF. The 1amma3-70B-Instruee model (Touvron et al., 2023) is the larger, more powerful LLM to which unconfident queries are routed.
Parameter-efficient fine-tuning is performed using LoRA adapters with ranks 2, 4, and 8 following the observations from (Hayou et al., 2024). Under at most eight epochs of training, all other hyperparameters on model training are decided by grid search with the perplexity loss on valida- tion sets. For more details about the hyper-paramcters, see Appendix C.
Efficiency Gains Smaller local models often have lower latency (sec) during inference as well as lower monetary cost. Thus, routing only fraction of queries to larger, more powerful models can significantly improve the efficiency of the overall system compared to relying entirely on the more powerful models. For each small LLM (Llama3-8B- Instruct and Mistral-7B-Instruee rate as the minimum routing rate that achieves comparable performance Llama3-70B -nn-Itrruttt (see Figure 2 for the trade-off curves). Without sacrificing performance com- Instruct, -nnsstut observe that Self-REF achieves as much as 2.03x latency improvement for Camma--BB-nnstruee on MMLU and 2.00 latency im- provement for Mistral-7B. IB-Inttuuct on OpenbookQA (as shown in Table 1). This indicates the efficiency advantage of leveraging Self-REF in routing tasks.
5. Experiment Results
we select the best routing
We conduct experiments to evaluate the performance of Self-REF, centered around the following three research questions: RQ1: Compared with state-of-the-ard baselines, how does Self-REF perform confidence-basee routing? RQ2: How reliable are confidence scores from Self-REF for the rejection learning task? RQ3: How well-aligned are confidence-tokent sereeeettoen-bbbeeeeee baske scores of Self-REF and the actual probabilities of correctness?
pared to Llama3-70B
6

---

Learning to Route LLMs with Confidence Tokens
Table 1: Per-token latency (sec) of smaller LLMs with routing using Self-REF vs. entirely using Llama3-70B-Inssmn while maintaining comparable overall accuracy (Acc.). Parentheses indicate the multiplicative speed-up.
MMLU
OpenbookQA
GSM8K
MedQA
Acc.
Latency
Acc.
Lateney
Acc.
Latency
Acc.
Latency
All in Llama3-708-Instren
0.739
0.292
0.781
0.292
0.819
0.292
0.622
0.292
Self-REF - Llama3- 1e133-8-Inne 3B-Instruct
0.739
0,145
(2.03x)
0.781
0.172 (1.69% x)
0.816
0.220
(1.33x)
0.619
0.147 (2.00x)
Self-REF- - Mistral-7B-Insttee 0.735
0.234 (1.25x)
0.778
0.147 (2.00x)
0.815
0.240 (1.20x)
0.621
0.234 (1.25x)
Table 2: Calibration metrics and routing rates of Self-REF and baselines using Llama3-8B Instrrcct and -B-nsrrlt on the MMLU, OpenbookQA, GSM8K, and MedQA datasets. "Route" denotes the routing rate or proportion of queries that must be routed to achieve performance comparable Llama3-70B- Inntrucct The gray block refers to the best routing rate, "bold format" is the best calibration score, and "*"" is the second best calibration score. Self-REF consistently achieves the best routing rate, but does not always achieve the best calibration score
Llama3-SB- -nsttrrtu
MMLU
OpesbookQA
GSMSK
McdQA
Route
RCE
BS
CF
Route
ECE
B8
CE
Route
BCB
BS
CF
Routee
ECE
BS
CF
Verhalizing Uncertainty
100%:
0.217
0.333
7,383
75%
0.148
0.311
7,688
92%
0.047
0.307
7.999
100%
0.069
0.331
1,206
Verbalizing Yes/No Token
100%
0.466
0.397
10.281
95%
0.291
0.415
5.319
100%
0.107
0.384
8.369
100%
0.090
0.470
11,276
Zero-shot Logits
100%
0.347
1.733*
90%
0.496
1,705
0.027
0.292*
3.312
100%
0.108
3.767
Finetuned Logits
90%
0.081
0,206
0.602
70%
0.095
0.238
0.676
70%
0.055
0.256
0.930
80%
0.059
0.265
0.844
Self-RF Llama3 Instracct
39%
0,040
0.320
11.534
49%
0.090
0.254
1.118*
65%
0.019
0.400
3.088
40%
0.066
0.278
1.085
Mistral-7BB 7B-Istruttt
MMLU
OpenbookOA
GSMSK
McdQA
Roste
CE
Route
ECE
BS
CE
Route
ECE
Route
BCE
CE
Verhalizing Uncertainry
100%
0.126
0.313
8.725
100%
0.221
0.434
10.846
100%
0.163
0.669
3,822
100%
0.087
0.564
2,480
Verbalizing YeuNo Token
0.297
4.380
100%
0.252
0.394
5.863
100%
0.146
0.821
17.14
Zero-shot Logins
95%
0.311*
0.353
1.643
100%
0,232
0.442
2.498
100%
0,046
0.398
1.957
100%
0.087
0.507
2.409
Finetuned Logits
0.263
0.751
50%
0.143
0.207
0.627
100%
0.050*
0.323
0.633
Sel--RF- Mistral-7P 7B- TB-snsttct
70%
0,369
0.293
1.021*
10%
0,157
0.193
0.869-
75%
0.068
0,273
0.968
70%
0.037*
0.226
0.653*
5.2. Rejection Learning Performance (RQ2)
5.3. Calibration with Utility (RQ3)
In this section, we analyze the utility of confidence scores from Self-REF for the purpose of learning to abstain from answering. Without incorporating additional knowledge of the rejection task during training. Self-REF and base- lines rely solely on thresholding the confidence scores to determine whether to reject a prediction. In safety-critical applications where one may hope to detect cases that should be abstained from answering (i.e., high recall on the reject option) while retaining performance on the task at hand (i.e., low false positive rate), it is useful to examine the ROC curve, which makes this trade-off. Across all thresholds, we observe that on both MMLU and OpenbookQA dataset, Self-REF outperforms all baselines on the rejection learning task (Figure 3). Because the confidence tokens are trained and conditioned on the correctness of the predicted answers, Self-REF has a higher likelihood of being aware of the cor- rectness of its generated responses using these tokens. In this manner, Self- REF may better distinguish when there is no correct answer to choose, resulting in better rejection learning capabilities.
Here, we analyze the correlation between calibration, or how well predicted probabilities align with actual proba- bilities, and utility for downstream routing. Following the settings from Tian et al. (2023); Blasiok et al. (2023), three calibration metrics are used to assess the calibration: (1) expected calibration error, ECE (Guo et al., 2017); (2) Brier Score, BS (Brier, 1950); and (3) the cross-entropy score, CE. Alongside these calibration metrics, we also show the least rates of routing required for each technique to attain com- parable performance Llama3-70B- in Table 2, although Self-REF often achieves the best or second-best calibration scores, better-calibrated methods do not guarantee optimal results for routing. This aligns with observations from prior work (Huang et al., 2023), which found that well-calibrated confidence scores do not neces- sarily imply a strong correlation with the correctness of the predictions. Overall, Self-REF achieves superior rout- ing performance across all datasets and local LLMs, even though it only outperforms baseline methods under most of the calibration assessments
As shown
7

---

Learning to Route LLMs with Confidence Tokens
Lama3-38
0.80
0.00
Question-1:
What is the difference between male and a
0.75
female catheter?
0.73
0.70
0.99
0.95
Subject:
Clinical Knowledge
0.02
Choices: [A] Male and female catheters are different colors. [B]
SeFREE 186,
0.00
Hork 1
0.90
Male catheters are longer than female catheters. [C] Male catheters are bigger than female catheters. [D) Female catheters are longer than male catheters
Nansaan Reals 728
0.50
0.0
0.2
0.1
0.5
03 1.3
5.0
222
0.4
0.5
0.0
Ground Truth Answer: [B] Male catheters are longer than female catheters.
Figure 4: Impact of gradient masking during Relf-REF fine-
tuning versus not (No Mask), assessed on the routing task.
Self-REF-8B5 (X) [D] Female catheters are longer than male
catheters. <UN> Route
Larger LLM-70B: (V) [B] Male catheters are longer than female
Transfer: Midrall-7 HMLL
Trandferi Liama HN LU
catheters.
Question-2: A certain pipelined RISC machine has 8 general- purpose registers RO, RI, operations. [Truncated 50 words]. If the contents of these three registers must not be modified, what is the minimum number of clock cycles required for an operation sequence that computes the value of AB + ABC + BC?
R7 and supports the following
to
Pouking hais
1ou Rate
Figure 5: Transferability of Relf-REF, assessed on the rout- ing task. Local LLMs are fine-tuned with Self-REF on OpenbookQA and inference is run on the MMLU dataset.
Subject: College Computer Science
Choices: [A] 5 [B] 6 [C] 7 [D] 8
Ground Truth Answer: [B] 6.
Self-REF-SB: (X) [C] 7. <CN> - Overconfident
Larger LLM-70B: (V) [B] 6.
5.4. Ablation Study on Gradient Masking
This question refers to the following information. Albeit the king's Majesty justly and rightfully is and ought to be the supreme head of the Church of England, [Truncated 230 words]; From the passage, one may infer that the English Parliament wished to argue that the Act of Supremacy would
Question-3:
To evaluate the impact of gradient masking on downstream fine-tuning performance, we consider Self-REF with and without gradient masking on the OpenbookQA dataset. For both Blama--B--Insruuee and Mistral-7B-Instruce we find that applying gradient masking consistently improves rout- ing task accuracy (Figure 4). This supports the intuition behind gradient masking: by masking gradients on sam- ples with <UN> tokens, we prevent the model from learning spurious patterns based on uncertain inputs.
Subject: High School European History
Choices: [A] give the English king new position of authority [B] give the position of head of the Church of Fngland to Henry VIII alone and exclude his heirs [C] establish Calvinism as the one true theology in England [D] end various forms of corruption plaguing the Church in England Ground Truth Answer: [D] end various forms of corruption plaguing the Church in England
5.5. Transferability of Confidence Token
Self-REF-8B: (V) [D] end various forms of corruption plaguing
the Church in England. <UN> - Underconfident
We further explore the transfcrability of Self-REF across unseen datasets (i.e., not included in the training data). With fine-tuned LoRA weights on OpenbookQA, we test on MMLU for Llama3-8BB
Larger LLM-70B: (z) [D] end various forms of corruption plagu- ing the Church in England.
and Mistral-7B-IInstnu Self-REF outperforms the baselines with good routing per- formance on the MMLU dataset even when using the trans- ferred LoRA weights of the OpenbookQA dataset (Figure 5). To achieve parity with the 70B model, we observe that rout- ing rates of 75% and 70% are required on MMLU under 7Bstral-BB-Instruee and Llama3-8BB ----3-88---ntcutt respectively. Furthermore, compared the direct fine-tuning on MMLU (i.e.. 70% routing rate on Mistral-7B- Instruct and 40% routing rate on Alama3-8B-Instruuee Instrucct, Figure 2), transferred Self-REF achieves competitive routing performance with only a slight degradation in routing rates. This demonstrates the potential of Self-REF to be trained for general, multi- task purposes due to its transferability.
Figure 6: Examples of routing from Llama-8B --IIttruct to Hamaa-77B-Instrrue on the MMLU dataset.
5.6. Case Studies
While Self-REF significantly outperforms other methods in routing accuracy, some failure cases still occur, particularly on niche subjects or tasks requiring nuanced reasoning. In underconfident cases, the model often hesitates on context- heavy (i.e., High School Europe History in Figure 6) or ambiguous subjects, highlighting its sensitivity to topics requiring broader reasoning. Overconfident failure cases oc- cur in technical or highly structured domains (i.c., College Computer Science in Figure 6). where the model appears
8

---

Learning to Route LLMs with Confidence Tokens
familiar but fails grasp nuanced or edge-case reason- ing. This phenomenon is also observed in the jurisprudence domain, where subtle distinctions in legal language and context-specific reasoning often lead to overconfident yet incorrect predictions. Conversely, correctly predicted con- fident cases typically involve broad conceptual topics, due to training coverage and clearer correctness signals. Mean- while, incorrect but uncertain cases often correspond to detail-intensive subjects, where the model appropriately ex- presses hesitation due to limited recall or understanding. For more details see Appendix E.
and monetary cost of routing a larger LLM on a server Earth is enormous, and one may want avoid routing all costs. If one has a limited budget for API calls to more powerful LLM, then one may also want to route sparingly, On the other hand, if monetary cost is not an issue, and latency is not critical, then one can have the entire trade-off curve available to them, and simply choose the point that obtains the highest accuracy.
When navigating the rejection learning setting, one must consider the risk of an incorrect answer. For example, if one is using an LLM to inform medical treatment, then one may choose to abstain from answering more often. On the other hand if one is using LLM for creative endeavors and correctness is not vital, then one may choose to let the LLM make its predictions and rarely abstain. Overall, these settings are highly customizable and controllable depending on the parameters that one is faced with in reality.
6. Discussion
Self-REF delivers comparable or even superior system-level routing performance compared to the non-finetuned 70B model. Note that fineffftumedd 70B model is used to mimic the common setting where practitioners may not have sufficient resources to fine-tune such large model, but can instead fine-tune a smaller model on the task of interest. In this section, we dive into the following questions: (1) Why can Self-REF after confidence-ber based routing achieve better performance than the non-finetuned 70B model? and (2) How can one navigate the tradeoff' in routing and rejection?
Through confidence-based routing and rejection, Self-REF enhances the overall performance and safety of a system while reducing cost. In our analysis, we measure cost primar- ily in terms of GPU hours and latency, excluding network transmission time between models which could be hosted on different machines (which is subject to varying signal qual- ity, setup costs, and network congestion). Future research in downstream applications could seek to better characterize these networking costs, which may contribute substantially to savings gained from confidence-based routing.
For question (I), we believe that one possible reason behind this phenomenon is that fine-tuming with Self-REFonn does not degrade the model's performance and effectively introduces correct knowledge, such as niche astronomy knowledge or complicated moral disputes, that the model does not inher- ently possess. Additionally, since the error distributions of two independent models are unlikely to overlap perfectly, the ability to route can serve as form of ensembling of models. In Appendix D, we provide some case studies from MMLU that the powerful LJama3-70B- model answered incorrectly, but Self-REF with the smaller Llama- 8B model answered confidently and correctly. As shown in Appendix Figure& 8, the Baama3-70B-Instruree may lack spe- eific knowledge to judge question of complicated moral disputes without fine-tuning and injecting the knowledge directly. More case studies are provided in Appendix F.
7. Conclusion
In this work, we introduce Self-REF, a lightweight fine- tuning strategy that incorporates confidence tokens, enabling LLMs to better express their confidence levels in a way that closely correlates with prediction correctness. These confi- dence scores are valuable for confidence-basee routing and rejection learning, which can improve system performance, efficiency, and output reliability for both LLM routing and LLM rejection learning purposes. Our findings suggest that integrating confidence tokens into LLMs is promising step toward improving their reliability and safety, particularly in real-word world applications where errors carry significant costs.
For question (2), Self-REF provides a continuous-valle valued confidence score which enables downstream decision- makers a granular level of control. By contrast, we observe that certain methods for extracting confidence scores, such as verbalizing uncertainty, can suffer from mode collapse, where the scores cluster around limited set of numbers (Figure 2, red dotted triangle lines). To appropriately lever- age the granular control provided by Self-REF, one must consider the parameters of the downstream setting.
More broadly, Self-REF opens new avenues for research into confidence-based routing and rejection. For instance, Self-REF could be extended route queries multiple spe- cialized LLMs rather than single powerful model. Another direction involves finer-grained annotation schemes, where individual sentences are labeled as confident or uncertain. Based on confidence, the system could route to different models, stop early, or re-invoke the reasoning processes to refine outputs. For confidence-baie based rejection, one could expand Self-REF incorporate explanations of rejections or to suggest actionable clarifying information.
When navigating the accuracy vs. routing rate tra one must consider the setting's 's cost of routing. For example, if one is utilizing LLM on rover on Mars, then the latency
9

---

Learning to Route LLMs with Confidence Tokens
Impact Statement
Cohen, R., Dobler, K., Biran, E., and de Melo, G. don't know: Explicit modeling of uncertainty with an [idk] to- ken. Advances in Neural Information Processing Systems, 37:10935- 2024.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be
specifically highlighted here.
Cortes, C., DeSalvo, G., and Mohri, M. Learning with rejec- tion. In Algorithmic I Learning Theory: 27th International Conference, ALT 2016, Bari, Italy, October 19-21, 2016, Proceedings 27, pp. 67-82. 82. Springer, 2016.
References
Abbasi Yadkori, Y., Kuzborskij. I., György, A., and Szepes- vari, C. To believe or not to believe your IIm: Iterative prompting for estimating epistemic uncertainty. Advances in Neural Information Processing Systems, 37:58077- 58117, 2024.
Detommaso, G., Bertran, M., Fogliato, R., and Roth, A. Multicalibration for confidence scoring in lmss. In Pro- ceedings of the 41st International Conference on Machine Learning, pp. 10641, 2024.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., al. Gpt-4 technical report. arXiv preprint arXiv:2303.08721 087744 2023.
Ding, D., Mallick, A., Wang, C., Sim, R., Mukherjee, S., Rühle, V., Lakshmanan, L. V., and Awadallah, A. H. Hy- brid Ilm: Cost-efficient and quality-aware query routing. In The Twelfth International Conference on Learning Representations, 2024.
Asai, A., Wu, Z., Wang, Y., Sil, A., and Hajishirzi, H. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreviee nett/forum?ichhssssgowwy
Duan, J., Cheng, H., Wang, S., Zavalny, A., Wang, C., Xu, R., Kailkhura, B., and Xu, K. Shifting attention to rele- vance: Towards the predictive uncertainty quantification of free-form large language models. In Proceedings of the 62nd Annual Meeting of the Association for Compu- tational Linguistics (Volume I: Long Papers), pp. 5050- 5063, 2024.
Blasiok, J., Gopalan, P., Hu, L., and Nakkiran, P. unifying theory of distance from calibration. Proceedings of the 55th Annual ACM Symposium Theory of Computing, pp. 1727-1740, 2023.
Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q. On calibration of modern neural networks. In International conference on machine learning, pp. 1321 1330. PMLR, 2017.
Brier, G. W. Verification of forecasts expressed in terms of probability. Monthly weather review, 18,1) 1950.
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023a.
Harrison, C. Langchain-ai/langchaii build context-aware reasoning applications. Oct 2022. URL https:// gi github.ccom Satuubbcoo//angcaiiraii/L//
Chen, G., Li, X., Sun, C., and Wang. H. Learning to make adherence-aware advice. In The Twelfth International Conference on Learning Representations, 2024.
Hayou, S., Ghosh, N., and Yu, B. Lorat: Efficient low rank adaptation of large models. In International Conference on Machine Learning, pp. 17783-17806. PMIR, 2024.
Chen, L., Zaharia, M., and Zou, J. Frugalgpt: How to use large language models while reducing cost and improv- ing performance. Transactions on Machine Learning Research, 2023b.
Hendrycks, D., Burns, C., Basart, S., Critch, A., Li. J., Song. D., and Steinhardt, Aligning ai with shared human values. Proceedings of the International Conference Learning Representations (ICLR), 2021a.
Chuang. Y.-N., Yu, L.. Wang, G., Zhang, L., Liu, Z., Cai, X. Sui, Y., Braverman, V., and Hu, X. Confident or seek stronger: Exploring uncertainty-based on-device Ilm routing from benchmarking generalization. arXiv preprint wwXiv:2502.04228 2025.
Hendrycks, D.,, Burns, C., Basart, S., Zou, A., Mazeika, M., Song. D., and Steinhardt, Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021b.
Cobbe, K. Kosaraju, V. Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 14168, 2021.
Hou, B., Liu, Y., Qian, K., Andreas, J., Chang, S., and Zhang, Y. Decomposing uncertainty for large language models through input clarification ensembling. In Inter- national Conference on Machine Learning, pp. 19023 19042. PMLR, 2024.
10

---

Learning to Route LLMs with Confidence Tokens
Hu, Q. J., Bieker, J., Li, X., Jiang, N., Keigwin, B., Ran- ganath, G., Keutzer, K., and Upadhyay, S. K. Router- bench: A benchmark for multi-llm routing system. In Agentic Markets Workshop at ICML 2024, 2024.
the 59th Annual Meeting of the Association for Com- putational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: 2021.
Long Papers), pp.
Huang, J., Parthasarathi, P., Rezagholizadeh, M., and Chan- dar, S. Context-aware assistant selection for improved inference acceleration with large language models. In Proceedings of the 2024 Conference on Empirical Meth- ods in Natural Language Processing, pp. 5817-5830, 2024.
Kuhn, L.,, Gal, Y., and Farquhar, S. Semantic uncertainty: Linguistic invariances for uncertainty estimation in nat- ural language generation. In The Eleventh International Conference on Learning Representations, 2023.
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Inter- national Conference on Machine Learning, pp. 19274- 19286. PMLR, 2023.
Huang, Y., Song, J.. Wang, Z., Zhao, S.. Chen, H.. Juefei- Xu, F., and Ma, L. Look before you leap: An exploratory study uncertainty measurement for large language mod- els. arXiv preprint arXiv:2307. 10236, 2023.
Li, D., Rawat, A. S., Zaheer, M., Wang, X., Lukasik, M., Veit, A., Felix, X. Y., and Kumar, S. Large language mod- els with controllable working memory. In ACL ACL (Findings), 2023.
Jeong, H., Jabbour, S., Yang, Y. Thapta, R., Mozannar, H., Han, W. J., Mehandru, N., Wornow, M., Lialin, V., Liu, X., et al. Recent advances, applications, and open chal- lenges in machine learning for health: Reflections from research roundtables ml4h 2023 symposium. arXiv preprint arXiv:2403.01622 2024.
Li, X., Liu, S., Sun, C., and Wang, H. When no-rejection learning is consistent for regression with rejection. In AISTATS, 2024.
Jangg A. Q. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.- A., Stock, P., Scao, T. L., Lavril, T., Wang. T., Lacroix, T., and Sayed, W. E. Mistral 7b. arXiv, 2023.
Lin, Z., Trivedi, S., and Sun, J. Generating with confidence: Uncertainty quantification for black-box large language models. Transactions on Machine Learning Research, 2023.
Liu, L., Pan, Y., Li, X., and Chen, G. Uncertainty esti- mation and quantification for llms: A simple supervised approach. arXiv preprint arXiv:2404. 15993, 2024.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. 1., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv
preprint ar arXiv:2401.0400 2024.
Mahaut, M., Aina, L., Czarnowska, P., Hardalov, M..
Mueller, T., and Mårquez, L. Factual confidence of llms: on reliability and robustness of current estimators. In Pro- ceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 4554 4570, 2024.
Jin, D., Pan, E., Oufattole, N., Weng, W.-H., Fang, H., and Szolovits, P. What disease does this patient have? large-scale open domain question answering dataset from medical exams. Applied Sciences, 11(14):422,, 2021.
Jordan, M. I. and Jacobs, R. A. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2): 181-214, 1994.
Manakul, P., Liusie, A., and Gales, M. Selfcheckggt: Zero- resource black-box hallucination detection for genera- tive large language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 9004-9017, 2023.
Kamath, A., Jia, R., and Liang, P. Selective question an- swering under domain shift. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp.
2020.
Mao, A., Mohri, M., and Zhong, Y. Structured prediction with stronger consistency guarantees. Advances in Neural Information Processing Systems, 36:46903 46937, 2023.
Kolasani, S. Optimizing natural language processing, large language models (Ilms) for efficient customer service, and perrooaitttnnnnntioionn to enable sustainable growth and revenue. Transactions on Latest Trends Artificial Intelligence, 4(4), 2023.
Mao, A., Mohri, C.,, Mohri, M., and Zhong, Y. Two-stage learning to defer with multiple experts. Advances in neural information processing systems, 36, 2024.
Krishna, K., Khosla, S., Bigham, J. P., and Lipton, Z. C. Generating soap notes from doctor-patient conversations using modular summarization techniques. Proceedings
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering. EMNLP, 2018.

---

Learning to Route LLMs with Confidence Tokens
Mohri, C., Andor, D., Choi, E., Collins, M., Mao, A., and Zhong, Y. Learning reject with a fixed predictor: Ap- plication to decontextualization. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
models by prompt agreement. In Proceedings of the 3rd Workshop on Trustworthy Natural Language Processing (TrusINLP 2023), pp. 326-362, 2023.
Wu, Q., Bansal, G., Zhang, J., Wu, Y., Li, B., Zhu, E., Jiang, L., Zhang, X., Zhang, S., Liu, J., al. Autogen: Enabling next-gen Ilm applications via multi-agent conversation. In ICLR 2024 Workshop Large Language Model (LLM) Agents, 2023.
Mozannar, H. and Sontag, D. Consistent estimators for learning to defer an expert. In International conference machine learning, pp. 7076-7087. PMLR, 2020.
Neeman, E.,, Aharoni, R., Honovich, O., Choshen, I.., Szpek- tor, I., and Abend, O. Disentqa: Disentangling parametric and contextual knowledge with counterfactual question answering. In ACL (1), 2023.
Xiao, Xu, S. X., Zhang, K., Wang, Y., and Xia, L. Evalu- ating reading comprehension exercises generated by Ilms: A showcase of chatgpt in education applications. In Pro- ceedings of the 18th Workshop on Innovative Use of NLP for Building Educational Applications (BEA 2023), pp. 610-625, 2023.
Ong. I., Almahairi, A., Wu, V., Chiang, W.-L.,, Wu, T., Gon- zalez, E., Kadous, M. W., and Stoica, I. Routellm: Learning to route Ilms from preference data. In The Thirteenth International Conference on Learning Repre- sentations, 2024.
Yang, J., Jin, H., Tang, R., Han, X., Feng, Q., Jiang, H., Zhong, S., Yin, B., and Hu, X. Harnessing the power of Ilms in practice: A survey on chatgpt and beyond. ACM Transactions on Knowledge Discovery from Data, 18(6); 1-32, 2024.
Roy, S. and Roth, D. Solving general arithmetic word problems. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing. As- sociation for Computational Linguistics, 2015.
Yona, G., Aharoni, R., and Geva, M. Can large lan- guage models faithfully express their intrinsic uncer- tainty in words? In Al-Onaizan, Y., Bansal, M., and Chen, Y.NN, (eds.), Proceedings of the 2024 Confer- ence on Empirical Methods in Natural Language Pro- cessing, pp. 7752-7764, Miami, Florida, USA, Novem- ber 2024. Association for Computational Linguisties. doi: 10. wrg.:/ggg..rg//22////// wttp://rr/.22....mmmmaaaaaa. //aclanthology.orge
Schmucker, R., Xia, M., Azaria, A., and Mitchell, T. Ruf- fle&riley: Towards the automated induction of conversa- tional tutoring systems. arXiv preprint 11000223,0001400, 2023.
Stripelis, D., Hu, Z., Zhang, J., Xu, Z., Shah, A. D., Jin, H., Yao, Y., Avestimehr, S., and He, C. Polyrouter: A multi-lim querying system. CORR, 2024.
443. URL https:
Tian, K., Mitchell, E., Zhou, A., Sharma, A., Rafailov, R., Yao, H., Finn, C., and Manning, C. D. Just ask for calibration: Strategies for eliciting calibrated confi- dence scores from language models fine-tuned with hu- man feedback. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5433-5442, 2023.
Yue, M., Zhao, J., Zhang, M., Du, I.., and Yao, 7. Large language model cascades with mixture of thought rep- resentations for cost-efficient reasoning. In The Twelfth International Conference on Learning Representations, 2023.
Zhang, H., Diao, S., Lin, Y., Fung, Y.. Lian, Q., Wang, X., Chen, Y., Ji, H., and Zhang, T. R-tuning: In- structing large language models to say II don't ' know'. In Duh, K., Gomez, H., and Bethard, S. (eds.), Pro- ceedings of the 2024 Conference of the North Ameri- can Chapter of the Association for Computational Lin- guistics: Human Language Technologies (Volume 1: Long Papers), pp. 7113-7139, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 11216653111/222222acmmlooooo //aclanthology 2egbantthlooyy.org/2244.4. www.:o//////.ee/ee
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F,, et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2300. 2023.
Turpin, M., Michael, J., Perez, E., and Bowman, Lan- guage models don't always say what they think: unfaith- ful explanations in ofain-of-thought prompting. Advances in Neural Information Processing Systems, 36, 2024.
URL https:
Wang, X. and Zhou, D. Chain-of-thouugh reasoning without prompting. In The Thirty-eigtth Annual Conference on Neural Information Processing Systems, 2024.
-accl---nnnncllnnao 394/
Zhang, J., Krishna, R., Awadallah, A. H., and Wang, C. Ecoassistant: Using Ilm assistants more affordably and accurately. In ICLR 2024 Workshop on Large Language Model (LLM) Agents, 2023.
Wightman, G. P., Delucia, A., and Dredze, M. Strength in numbers: Estimating confidence of large language
12

---

Learning to Route LLMs with Confidence Tokens
Zhang, S., Gong, C., and Choi, E. Knowing more about questions help: Improving calibration in question answering. In Findings of the Association for Computa- tional Linguisties: ACL-IJCNLP 2021, pp. 1958-1970, 2021.
Zhou, K., Jurafsky, D., and Hashimoto, T. B. Navigating the grey area: How expressions of uncertainty and over- confidence affect language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5506-5524, 2023.
13

---

Learning to Route LLMs with Confidence Tokens
Appendix
A. Related Work
A.1. Uncertainty Quantification in LLMs
There a growing body of recent works examining the calibration of LLMs. One line of research (Zhou et al., 2023; Mahaut et al., 2024; Turpin et al., 2024) leverages the internal features of LLMs (e.g. logits, attention scores) to assess and estimate uncertainty without requiring additional training fine-tuning. Improvements in calibration have also been observed from post-processing approaches such as temperature scaling Platt scaling (Guo et al., 2017), as well as prompting approaches asking the LLM to verbalize its predicted confidence scores (Tian et al., 2023). Another line of work (Zhang et al., 2021; Neeman et al., 2023; Li et al., 2023; Liu et al., 2024) focuses on RLHF or instruction learning for achieving better prediction calibration.
In addition to codel-centric analysis, another school of approaches (Hou et al., 2024; Kuhn et al., 2023; Duan et al., 2024) focuses on examining the semanties and ambiguity in the given input questions to estimate uncertainty in LLMs. On the other hand, (Huang et al., 2023) have found that uncertainty scores derived from logits and verbalized scores exhibit weak correlation with the correctness of LLM predictions across several models and datasets. Thus, while calibration may be a desirable characteristic of uncertainty estimates, we note that this work goes beyond calibration as goal, focusing more on the utility of model confidence scores downstream, for tasks such as instance routing and rejection learning.
A.2. LLM Routing
LLM routing, where all or part of query is directed from one LLMs to different LLMs, has been studied in the literature in a few capacities: (1) routing entire queries to different LLMs, (2) -uummette generation (RAG) systems that route part of the query to an external retrieval system, and (3) speculative decoding which utilizes outputs from models of varying complexity in the decoding step.
(1) In the query query-routing scenario, one line of work trains additional classification models route queries to other LLMs for answers based on designated performance metrics (Ding et al., 2024; Hu et al., 2024; Huang al., 2024) and routing data (Ong al., 2024; Stripelis et al., 2024), which eventually form routing pipeline with multi-agent systems (Jordan & Jacobs, 1994; Jiang et al., 2024) to boost the performance. Another line of work leverages routers to measure utility- efficiency trade-offs, aiming to reconstruct model architectures to reduce inference overhead in terms of cost and execution time while maintaining utility (Chen et al., 2023b; Yue et al., 2023). (2) In RAG systems, routing provides effective and efficient navigation to specific data sources, offering efficient and effective indicators for RAG systems to retrieve. A notable approach is SeIf-RAG (Asai al., 2024), which fine-tunes an LLM retriever using augmented data with predetined routing indicators optimize retrieval performance. This approach improves RAG systems by enabling more efficient retrieval through dynamic routing decisions of each given sentence in the corpus based on predicted routing indicators generated by the fine-tuned L.M retriever. (3) In speculative decoding, the entire decoding process is carried out through an iterative collaboration between a smaller LLM and larger LLM. Each final output token sequence can be generated by either a smaller LLM or a larger LLM. Specifically, the token decoding process will route to a larger language model when the smaller model either declines to generate or encounters a failure in token generation (Chen et al., 2023a; Leviathan et al., 2023). Different from the routing concepts the work of training extra query routers, we route the queries based on the model confidence of each LLM output answer. The confidence estimations in our work are predicted given the output answers during the autoregressive decoding process, which makes the model confidence more aligned with the correctness output answer.
A.3. LLM Rejection Learning
Learning to reject (Cortes al., 2016) has been initially proposed to offer the Bayes classification model rejection option as the final prediction. However, with the rapid advancement of LLMss, the issue of uncertainty in L.LMs' predictions may stem from their limited inherent knowledge, necessitating an option or mechanism for LLMs to refuse to answer, especially when no advanced LLMs are helpers for two-step verification or answering in action. The rejection decisions made by LLMs can enhance the reliability of their generated predictions. One group of work offers LLMs a rejection option while they generate the responses by finetuning or instruction learning (Chen al., 2024) with newly designed rejection loss. Another group of work trains the additional rejection LLM agents issue rejection instruction after the predictor LLMs have made their prediction (Kamath et al., 2020; Mozannar & Sontag, 2020; Mao et al., 2023; 2024). Unlike the training paradigm of
14

---

Learning to Route LLMs with Confidence Tokens
traditional machine learning models, establishing new rejection loss with new knowledge in fine-tuning LLMs may result in sub-optimal performance and potentially degrade the abilities of pre-trained L.Ms in next-token prediction. A more effective approach is to adjust LLMs using the original training data and loss functions (Li et al., 2024; Mohri et al., 2024) while incorporating the capability for rejection learning without altering the foundational training dynamics. Moreover, training additional agents to handle rejection learning can introduce significant latency and may be overfitted, making the approach impractical for real-world applications. In this work, we aim to equip predictive LLMs with the capability for rejection learning by enabling them to reveal their confidence levels, without the need for designing new loss functions, thereby preventing any potential degradation in the performance of downstream tasks.
B. Details about Baseline and Datasets
B.1. Verbalizing Confidence Prompt Usage
We provide a listing of the evaluation prompts in Table 3 utilized in assessing the performance of the baselines among all four datasets. The first sections reveal the evaluation prompt for "Verbalizing Uncertainty:" the second sections demonstrate the evaluation prompt for "Verbalizing Yes/No Tokens;" and third batches demonstrate the evaluation prompt for "Zero-shoo/Fine-tunee Logits." "
Confidence Baselines
Dataset
Evaluation Prompts
Verbalizing Uncertainty
MMLU
Answer the input question. For each of the answer choices A, B. C. and D. Output your confidence that the correct answer. The total values must equal ONE, <Task Description>. You must follow the format for your final answer: "##Answer=Yoor Answer> ##Confidence:<<oou confidence>'
Verbalizing Uncertainty
OpenbookQA MedQA
Answer the input question. For each of the answer choices A, B, C, and D. Output your confidence that it is the correct answer. The total values must equal ONE You must follow the format for your final answer: *#Anwwer<<Yoor Answer> ##Confidence:<<Ymm confidences'
Verbalizing Uncertainty
GSM8K
Output your confidence that it is the correct answer. The total values must equal ONE. You must follow the format for your final answer: Answer> ##Confidencc:< confidence>". Let's think step by step and oulput the answers and confidence score following the format only." "
Verbalizing Yes/No Tokens
MMLU
Please answer the questions first. After providing your answers, indicate your confidence level in your responses. <Tassk Description>. Respond with 'Yes" you are confident and 'No" if you are not confident.
Verbalizing Yes/No Tokens
OpenbookQA MedQA
Please answer the questions first. After providing your answers, indicate your confidence level your responses. Respond with 'Yes" if you are confident and 'No' if you are not confident.
Verbalizing Yes/No Tokens
GSM8K
Please answer questions first. After providing your answers, indicate your confidence level in your responses. Respond with 'Yes' if you are confident and 'No' if you are not confident. You must follow the format for your final answer: "##Anser:<Yoou Answer> ##Confidence:< Your confidence>>. Let's think step by step and output the answers and confidence score following the format only."
MMLU
Zero-shot Logits Fine-tuned Logits
Answer the question by choosing one of the answer choices A, B, C, or D.
OpenbookQA MedQA
Zero-shot Logits Fine-tuned Logits
GSMSK
Let's think step by step.
Table 3: Evaluation prompts for verbalizing uncertainty on local LLMs under different dataset.
15

---

Learning to Route LLMs with Confidence Tokens
B.2. Finetuned Logits Training
The hyperparameter settings (shown in Table 5) used for LoRA training of the "Finetuned Logits" baseline are follows. The LoRA dimension set from either 4, and 8, and the reported performance of "Finetuned Logits" is the with highest downstream task accuracy.
Table 4: Hyper-parameter settings in "Fine-tuned Logits" baseline.
Dataset
MMLU
OpenbookQA
GSM8K
MedQA
Optimizer
Adam
Adam
Adam
Adam
Llama3-8B-Instue
Warm Start
200
200
200
2000
Total Training Data #
9,985
4.957
7,473
10,000
Optimizer
Adam
Adam
Adam
Adam
Mistral-7B-Inntruet
Warm Start
200
200
200
200
Total Training Data #
9,985
4,957
7,473
10,000
B.3. Details of Datasets
The experiment are conducted on four different open- sourced datasets. The details of the datasets are provided as follows:
MMLU (Hendrycks al., 2021b;a). A diverse benchmark covering 57 subjects across humanities, social sciences, STEM, and professional domains. It evaluates model's ability to perform zero-shot multiple-choice question answering using knowledge acquired during pretraining. In the experiment, we evaluate the model on val set.
OpenbookQA (Mihaylov et al., 2018). A 5,957 multiple-choice question answering dataset focused on elementary science. It requires combining small "open book" of core science facts with broader common knowledge and reasoning.
GSM8K (Cobbe et al., 2021). A dataset of grade-school level math word problems designed to test arithmetic reasoning. It includes both questions and detailed step-by-step solutions, making it ideal for evaluating models' reasoning and intermediate computation skills. We assess the model on 1,331 test set.
MedQA (Jin et al., 2021). A large-scale multiple-choice medical question answering dataset based on real-world medical licensing exams with 12,723 questions. It challenges models with clinically relevant questions requiring specialized medical knowledge and multi-step reasoning.
16

---

Learning to Route LLMs with Confidence Tokens
C. Experiment and Model Training Details
C.1. Hyper-parameter Settings
Our experiments on the dataset are conducted under Self-REF. We focus on the confidence-based routing and rejection learning tasks, following the pipeline of Base Model Training and Confidence Token Inference. The details of each step shown as follows:
Base Model Training In this work, Self-REF is fine-tuned local LLMs (i.e., Llama3-8B-Instrue and Mistral-7B- Instruct) under the following hyper-pun arper-parameters. We apply LoRA adapters to every query, key, and value (Q-K-V layers, the token embedding layers, and the final linear layers in the local LLMss with the batch size of 4 and learning rate of 1e-4. Fixing the overall dataset size, the parameter a is tuned based on performance on the validation set, selecting from ratio of unconfident to confident data of 1:1, 1:2, 1:3, 1:4, and 1:5. On one hand, including more confident samples helped fine-tune the model to perform better on the downstream task; on the other hand, a sufficient number unconfident samples was necessary to teach the model when express uncertainty. For the ratio of training dataset, we only adjust the OpenbookQA and MedQA dataset in Mistral-7B-Instteen model to prevent the overfitting situation. We retain the original dataset settings with the correctness and incorrectness directly inferred by local L.MS.
Dataset
MMLU
OpenbookQA
GSM8K
McdQA
LoRA Dimension
2, 8
2, 4, 8 Adam 10
2, 8 Adam
2, 4, Adam
Optimizer
Adam
Weight Decay
10 ~ 10-1
10-1 200
5 x 10 10-2
10-2 ~ 7 10-1 200
Llama3
Warm Start
200
200
Ratio <CN>*: <UN>
2.79
1.01
4.47
2.74
Total Training Data
9,985
4,957
7,473
10,000
Training Time (Hours)
~1
~2.
~4
LoRA Dimension
2, 4, 8 Adam
2, 4, 8 Adam
2, 4, 8
2, 4, Adam
Optimizer
Adam
Weight Decay
10T2 x 10-1 200
9 10 10 200
x 10 10"1
10
10
Mistral
Warm Start
200
200
Ratio <CND*; <UN>
1.37
2.25
1.01
2.05
Total Training Data
9,985
3,395
7,473
7,557
Training Time (Hours)
005
Table 5: Tlyper-parametert and model structures settings in Self-REF. <CN> denotes the amount of <CN> data, which depends on the number of questions that local LLMs answer correctly. The underlined ratios indicate the proportion of queries with <UN> that are downsampled to prevent overfitting.
Self-REF Inference During the inference time, the temperature is set O, top-p as 1.0, and all decoding processes greedy search. All other sampling strategies are forbidden with the fixed random seed 2024 for the reproducibility.
C.2. Computation Infrastructure
The experiments are conducted based on the following physical computing infrastructure in Table 6. This setup provided the necessary computational resources to efficiently train and evaluate our models.
Device Attribute
Value
Computing infrastructure
GPU
GPU model
Nvidia-A100
GPU number
8
GPU Memory
80G
Table 6: Computing infrastructure for the experiments.
17

---

Learning to Route LLMs with Confidence Tokens
D. Case studies of Self-REF Routing
We now perform a case study of Self-REF on the MMLU dataset. Recall from the results that Self-REF can effectively learn confidence tokens without degrading performance. We here showcase two types of examples: (I) cases that Self-REF finetuning on Llama3-8B-Insiee 8B-Instruct initially got incorrect but were corrected by routing to Llama3-705B 70B-Instruct using Self-REE confidence scores (Figure 6).
got correct but Llama3-70B-Instee got incorrect (Figure 6), and (2) cases that Llama3-
Question-1: 1: Baron admits that the versions of ticking bomb hypothetical she discusses are "stunningly stupid," she claims this actually evidence of?
Subject: Moral disputes
[A] the stupidity of most traditional philosophical examples. a general lack of intelligence among people with advanced degrees. [C] the wrongness of torture. [D] the readiness on the part of many intelligent people to see torture as the best solution to deal with terrorism.
Choices:
Groun Truth Answer: [D] the readiness on the part of many intelligent people to see torture the best solution to deal with terrorism.
8B Response: (V) [D] the readiness on the part of many intelligent people to see torture as the best solution to deal with terrorism. <CN>
70B Response: (X) [A] the stupidity most traditional philosophical examples.
Question-2:
Functions of the law include all but which of the following?
Subject: Jurisprudence
Choices: [A] maximizing individual freedom. [B] providing basis for compromise, [C] keeping the peace. [D] promoting the principles of the free enterprise system.
Ground Truth Answer: [D] promoting the principles of the free enterprise system.
8B Response: (V) [D] the readiness the part of many intelligent people torture the best solution to deal with terrorism.
<CN>
70B Response: (X) [A] maximizing individual freedom.
Question-3: Qucstion-263: What is the Second Gem in Buddhism??
Subject: World Religions
Choices: [A] The Dharma [B] The Sangha [C] The Buddha [D] The Bodhisativa
Ground Truth Answer: [A] promoting the principles of the free enterprise system.
8B Response: (V) [A] The Dharma, <CN>
70B Response: (X) [B] The Sangha.
Figure/ 7: The examples on MMLU dataset, where the Llama-8B model correctly answers a query after routing, while the non-fine-tuned -lama3-70B-Innee model answers incorrectly.
18

---

Learning to Route LLMs with Confidence Tokens
Question-1: Statement Every function from finite onto itself must be Statement Every subgroup of abelian group is abelian.
Subject: Abstract Algebra
Choices: [A] True, True [B] False, False [C] True, False [D) False, True
Ground Truth Answer: [A] True, True.
8B Response: (X) [C] True, False. <UN> - Route
70B Response: (V) [A] True, True.
Question-2: A model airplane flies slower when flying into the wind and faster with wind at its back. When launched right angles to the wind a cross wind its groundspeed compared with flying in still air is?
Subjeet: Conceptual Physics
Choices: [A] the same [B] greater [C] less [D] either greater or less depending on wind speed
Ground Truth Answer: [B] greater.
Response: (X) [D] either greater or less depending on wind speed. <UN> - Route
70B Response:
(/) [B] greater.
Question-3: Two long parallel conductors carry 100 A. If the conductors separated by 20 mm, the force per meter of length of each conductor will be?
Subject: Electrical Engincering
Choices: [A] 100 N. [B]0.1 0.1 N. [C] N. 0.] N.
Ground Truth Answer: [B] 0.1 N.
Response: (X) [C] N. <UN>
Route
70B Response: (V) N.
Figure 8: The routing examples with <UN> MMLU dataset, where the Llama-8B model incorrectly answers query initially. After routing to Llama3-70B Instructt the queries answer correctly.
19

---

Learning to Route LLMs with Confidence Tokens
E. Case studies by Subject Area
To better understand the model's s calibration behavior, we analyze the top five subject categories across four groups based on prediction correctness and confidence levels:
Correct Predictions <UN> (Underconfident):
Computer Security
High School Biology
High School European History
Human Sexuality
Miscellaneous
Incorrect Predictions + <CN> (Overconfident):
College Computer Science
Conceptual Physies
High School Computer Science
High School Microeconomics
Jurisprudence
Correct Predictions <ON> (Confident and Accurate):
International Law
College Biology
Moral Disputes
Philosophy
U.S. Foreign Policy
Incorrect Predictions <UN> (Unconfident and Incorrect):
Abstract Algebra
Anatomy
College Chemistry
College Medicine
Econometrics
This categorization reveals several noteworthy trends:
The model tends to be underconfident on subjects that involve ambiguity, contextual reasoning, or non-standard formats. For example, topics like human sexuality and European history often require nuanced interpretation. The presence of "Miscellaneouss" suggests increased uncertainty when the question content falls outside well-defined training distributions.
Correct + Underconfident:
Incorrect + Overconfident: Overconfidence is particularly common in structured, technical domains such as computer seience and physics. This behavior likely arises from overfitting to frequent patterns seen in pretraining, while failing on nuanced reasoning or edge cases (e.g., legal intricacies in jurisprudence).
Correct + Confident: In domains such as philosophy and international law, the model demonstrates reliable perfor- mance and appropriate confidence. These topics often involve conceptual reasoning with clear evaluative criteria, which may contribute to better calibration.
Incorrect Underconfident: In contrast, the model appropriately exhibits low confidence in subjects requiring exact recall and specialized knowledge-e.g. college-level STEM disciplines like chemistry and medicine. These topics demand precise understanding, and the model's hesitation reflects its awareness knowledge limitations.
20


Reviews:
[{"review_id": "p5h4HDFuAr", "reviewer": "Reviewer_kmXF", "summary": "The paper proposes Self-REF, a training strategy that adopts LoRA to fine-tune an LM on a dataset augmented with confidence tokens, based on prediction correctness. Self-REF enhances downstream applications like model routing and answer rejection by leveraging the learned confidence token scores.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**[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.**\n\n**[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. \nThus, 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.\n\n**[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.**\n\n**[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.\n\n**[Q-3] How was the optimal parameter $\\alpha$ for the <UN> data determined?** \n\n**[A-3]** $\\alpha$ is a hyperparameter in our work, and we select the optimal $\\alpha$ based on the validation set.\n    \n**[Q-4] How does the ratio of <CN>:<UN> in dataset construction affect the results?**\n\n**[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.\n   \n**[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.**\n\n**[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.\n\n**[Q-6] Missing reference, formats, and Typos.**\n\n**[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.", "claims_and_evidence": "Please refer to the \"Other Strengths And Weaknesses\" section.", "methods_and_evaluation_criteria": "Please refer to the \"Other Strengths And Weaknesses\" section.", "theoretical_claims": "N/A. No proofs or theoretical claims that require checking.", "experimental_designs_or_analyses": "Please refer to the \"Other Strengths And Weaknesses\" section.", "relation_to_broader_scientific_literature": "Please refer to the \"Other Strengths And Weaknesses\" section.", "essential_references_not_discussed": "Please refer to the \"Other Strengths And Weaknesses\" section.", "other_strengths_and_weaknesses": "**Strengths**\n- The proposed Self-REF method is straightforward yet effective in learning more calibrated confidence scores, utilizing a simple data augmentation strategy.\n\n- Overall, the paper is well-written and easy to follow, with comprehensive and detailed experiments that demonstrate strong empirical results across multiple datasets.\n\n**Weaknesses**\n- Missing reference to R-Tuning [1], which similarly constructs an augmented dataset consisting of certain and uncertain sets based on the correctness of LM predictions, and appends corresponding tokens as supervised signals to fine-tune a more calibrated LM capable of refraining from answering unknown questions. This work is not discussed in the paper, and could be considered as a baseline.\n\n- Self-REF requires a dedicated training/validation set and a fine-tuning stage for different downstream tasks/datasets, which limits its practical usages comparing to zero-shot baselines.\n\n- 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.\n \n[1] Zhang, Hanning, et al. *R-tuning: Instructing large language models to say ‘I don’t know’.* NAACL 2024", "other_comments_or_suggestions": "Presentation: The figure size and font size (e.g., in ```Figure 2```) might be too small, potentially affecting readability.", "questions_for_authors": "1. How was the optimal parameter $\\alpha$ for the <UN> data determined?\n2. How does the ratio of <CN>:<UN> in dataset construction affect the results?\n3. 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.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u7LnFdG9Y0", "reviewer": "Reviewer_w1rW", "summary": "Authors proposed a lightweight training strategy to teach LLMs to express confidence in whether their answers are correct in a reliable manner. Using this, the authors build a router algorithm that reduces latency and improves overall QA performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable feedback and ideas. We have updated our discussion accordingly.\n\n**[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?**\n\n**[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.", "claims_and_evidence": "The claims are well stated and supported.", "methods_and_evaluation_criteria": "All three RQ1, RQ2, RQ3 benchmarks and datasets are make sense and well analyzed.", "theoretical_claims": "Not applicable", "experimental_designs_or_analyses": "Experimental designs and analyses are valid. In particular, significantly reducing the latency metric, improves the validity of the results provided.", "relation_to_broader_scientific_literature": "The routing problem and Confidence Tokens are very novel research area. Every related literature is well covered.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Please see other sections", "other_comments_or_suggestions": "Paper is well written.", "questions_for_authors": "Your framework is only able to predict the confidence token at the end of the answer? Is that correct?\nIf 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "isOwVM9wl3", "reviewer": "Reviewer_gYcY", "summary": "Self-REF is a lightweight method for training a language model to show when its answers are correct or incorrect by using “confidence tokens.” The approach starts with a base model that generates predictions and labels each instance as “confident” or “unconfident” based on the answer’s correctness, creating an augmented dataset. The model is then fine-tuned on these labeled samples and learns both to provide the correct answer and to tag it appropriately. Finally, continuous confidence scores are computed by comparing the probabilities of the “confident” vs. “unconfident” token at the end of each response.\n\nThey propose two applications. First, low-confidence queries can be routed to a larger model to cut costs without sacrificing accuracy. Second, when no larger model is available, the system can reject responses it deems untrustworthy. Across MMLU, OpenbookQA, GSM8K, and MedQA, routing only uncertain queries matches the strong model’s accuracy while reducing latency and cost; rejecting low-confidence outputs also helps avoid incorrect claims for questions lacking a valid answer.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed feedback. We respond to each point below, and will update the discussion accordingly:\n\n**[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...**\n\n**[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.\n\n**[Q-2] If the base model has substantial weaknesses or biases, these may propagate through fine-tuning, limiting improvements or perpetuating biases.**\n\n**[A-2]** We respond using two perspectives:\n- The primary goal of Self-REF is to help the base model identify its existing weaknesses while maintaining performance, thereby enabling more effective routing.\n- 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.\n\n**[Q-3] Explore other methods for labeling confidence (consistency across sampling and external calibration).**\n\n**[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. \n\n```\nBins = 2\npredicted prob. of \"1\" = [0.5, 0.5, 0.5, 0.5]\nground truth = [0, 0, 1, 1]\n\nExample 1: \npredicted prob. of \"1\" = [0.5, 0.5, 0.5, 0.5]\n-> ECE(↓)= 0%\n-> accuracy (↑)= 0%\n\nExample 2: \npredicted prob. of \"1\" = [0.4, 0.6, 0.9, 0.9]\n-> ECE(↓) = 20%\n-> accuracy (↑) = 75%\n```\n[1] Huang, et al. \"Look before you leap: An exploratory ...\" arXiv 2023.\n\n[2] Yona, et al. \"Can Large Language ... Uncertainty in Words?.\" arXiv 2024.\n\n**[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.**\n\n**[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.\n\nWe 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.\n\n|Route_Ratio|0%|20%|40%|60%|80%|100%|least_routing_ratio|\n|-|-|-|-|-|-|-|-|\n|OOD-Probe|0.50|0.61|0.66|0.68|0.74|0.74|80%|\n|Self-REF|0.55|0.64|0.68|0.72|0.74|0.74|70%|\n\n[1] Mahaut, et al. \"Factual confidence of LLMs:...\" ACL 2024.", "claims_and_evidence": "* Claim that the model’s confidence tokens always align with correctness:\nThe paper states that once fine-tuned, the model should consistently output <CN> whenever it is correct and <UN> otherwise. In practice, even well-calibrated methods can yield discrepancies between predicted confidence and true correctness.\n\n* Claim of consistently “lightweight” overhead\nThe resulting method would yield a cascading method, where first we have to monitor the smaller model's output then route to the bigger one to get the final answer.", "methods_and_evaluation_criteria": "Overall, the methods and evaluation criteria do generally make sense given their focus on practical improvements in downstream tasks. However, I have a few concerns:\n* 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, though confidence in practice might not always reflect correctness accurately, potentially limiting generalizability.\n* The paper partially addresses imbalance by subsampling unconfident tokens with a tunable parameter, but does not elaborate on handling extreme cases where the base model might be consistently correct or incorrect, which could severely skew training data. In fact, the approach depends heavily on the correctness of the base model to annotate confidence tokens. If the base model has substantial weaknesses or biases, these may propagate through fine-tuning, limiting improvements or perpetuating biases.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "Yes, I would suggest studying additional baselines: \n* I'm curious if authors explored a baseline where the model generates a confidence token before producing an answer. This can reduce computation by preventing full generation when confidence is low. Also a learned router baseline would provide insight\n* Explore other methods for labeling confidence:\n    * Estimating confidence based on consistency across multiple stochastic samples from the model.\n    * External calibration techniques, e.g. entropy-based uncertainty", "relation_to_broader_scientific_literature": "The work addresses an important problem of allowing LLMs to self-assess the confidence of their predictions, and their applications to LLM routing and rejections.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PE3yX0FLfE", "reviewer": "Reviewer_JjAL", "summary": "The paper introduces Self-REF, a lightweight fine-tuning framework designed to teach large language models (LLMs) to express confidence in their answers through confidence tokens. These learned tokens, indicate whether the model is confident or uncertain about its prediction, improving reliability and performance in downstream routing and rejection learning tasks.\nThe authors demonstrate that Self-REF outperforms traditional approaches, such as verbalizing confidence or using token probabilities, on multiple public datasets by enabling more accurate routing to stronger models and better rejection of uncertain answers. The method achieves improved system efficiency and calibration while maintaining model performance.\n\n## update after the rebuttal\n\nDuring the rebuttal the authors have addressed all my concerns. I have decided to increase my score from 3 to 4", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**[Q-1] More analysis on a systematic investigation of failure cases to identify patterns in when and why self-REF misclassifies confidence.**\n\n**[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:\n- Predict Correctly+<UN> (underconfident in its predictions): (1) computer_security; (2) high_school_biology; (3) high_school_european_history; (4) human_sexuality; (5) miscellaneous\n- 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\n- Predict Correctly+<CN>: (1) international_law; (2) college_biology; (3) moral_disputes; (4) philosophy; (5) us_foreign_policy\n- Predict Wrongly+<UN>: (1) abstract_algebra; (2) anatomy; (3) college_chemistry; (4) college_medicine; (5) econometrics\n\nAnalyzing these categories, we note:\n- 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.\n- 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.\n- 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.\n- 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.\n\n**[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 ...**\n\n**[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.\n\n**[Q-3] In figure 2 I see random route to 70B baseline that I can't find it described. Can you elaborate on it?**\n\n**[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.\n\n**[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.**\n\n**[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.\n\n**[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...\"**\n\n**[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.\n\n**[Q-6] Missing reference, formats, and Typos.**\n\n**[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.", "claims_and_evidence": "The claims in the submission are supported by clear and convincing evidence. The paper presents quantitative results across four public datasets (MMLU, OpenbookQA, GSM8K, MedQA), showing that Self-REF improves routing efficiency by reducing the number of queries sent to larger models while maintaining accuracy (e.g., Llama3-8B routes only 39% of queries to match Llama3-70B's performance). Additionally, ROC curves demonstrate superior rejection learning, and calibration metrics (ECE, Brier Score, Cross-Entropy) confirm that confidence tokens align well with correctness. However, while the authors discuss trade-offs in routing and rejection, they do not extensively analyze potential failure cases, such as when Self-REF misidentifies confidence leading to incorrect rejections or unnecessary escalations.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria are well-aligned with the problem of confidence-based routing and rejection learning in LLMs. The use of confidence tokens is a novel yet intuitive approach that integrates seamlessly into autoregressive models, and the evaluation on standard QA and reasoning benchmarks (MMLU, OpenbookQA, GSM8K, MedQA) allows for fair comparison with prior art. The routing and rejection tasks are practical and relevant, as they reflect real-world scenarios where LLMs need to manage uncertainty efficiently. \nAs stated above, an error analysis of model behvior in confidence-based routing would help gain additional insights on the effectiveness of Self-REF.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experimental design is generally strong. The routing experiments are well-structured, using multiple confidence thresholds to analyze trade-offs between accuracy and efficiency. The rejection learning task is also appropriately tested with artificially modified datasets where the correct answer is removed.  Also the calibration analysis (ECE, Brier Score, CE)  help measure model confidence alignment with correctness.`\n\nOn the other hand, a much stronger analysis would include a systematic investigation of failure cases to identify patterns in when and why Self-REF misclassifies confidence. While the paper provides overall accuracy and calibration metrics, it does not explore whether certain types of questions, knowledge domains, or reasoning patterns lead to systematic overconfidence or underconfidence. For instance, if the model consistently misroutes some type of quesstions questions, this could highlight fundamental limitations in its confidence estimation. Identifying such failure patterns would not only improve interpretability but also inform targeted improvements to Self-REF, such as adjusting confidence token fine-tuning strategies or incorporating adversarial training.\n\nAlso, an analyses of the relation between confidence threshold (more below in the questions section) could help get a better understanding of model behavior when self-REF is used.\n\nWith the analysis above, I consider this paper to deserve a score of 4 instead of its current 3", "relation_to_broader_scientific_literature": "The paper builds on prior work in uncertainty quantification, LLM routing, and rejection learning, but introduces confidence tokens as a mechanism for end-to-end confidence estimation. Previous methods, such as logit-based calibration  and verbalized uncertainty prompts, shows worse alignment between confidence scores and correctness, whereas Self-REF fine-tunes the LLM itself to embed confidence directly. Unlike external classifiers for routing, Self-REF integrates confidence estimation within the autoregressive decoding process, making routing decisions more adaptive and model-aware. Additionally, it improves on LLM rejection learning by enabling confidence-based abstention without requiring a separate rejection model or new loss functions.", "essential_references_not_discussed": "All the relevant references needed to understand the paper are discussed.\n\n--\n\nI found this paper that came out after ICLR submission deadline but it could be a nice reference to add to the related work section: \nDhananjay Ashok, Jonathan May, Language Models Can Predict Their Own Behavior", "other_strengths_and_weaknesses": "The paper is generally well written and easy to follow.", "other_comments_or_suggestions": "Typos:\n\n- line 117: retarded? should be a typo\n- line 349: bed is s typo", "questions_for_authors": "- What data did you use in the setup when you mention you trained on MMLU? I looked at table 4 in the appendix but I'm not sure about the numbers given that, from the MMLU paper, I get \"The few-shot development set has 5 questions per subject, the validation set may\nbe used for selecting hyperparameters and is made of 1540 questions, and the test set has 14079 questions.\". In general, i'd move some information about training the confidence tokens from the appendix to the main body of the work, or at least better discuss them in the paper.\n\n- In figure 2 I see random route to 70B baseline that I can't find it described. Can you elaborate on it?\n\n- 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), what was the exact threshold value (t) that resulted in a routing rate of 0.4? More generally, it would be helpful to explicitly discuss the relationship between t and routing rate, as this would give practitioners a clearer understanding of how to tune Self-REF for different trade-offs between accuracy and efficiency.\n\n- 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. The routing decisions are determined by the probabilities\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: **Paper summary:**

This paper proposes a framework to fine-tune an LLM to express its confidence for an input query through a custom confidence token that is annotated in the fine-tuning data set. The framework has three phases: 

1. Confidence token annotation. If the model gives a correct answer, append <CN> to the ground-truth response. Else, append <UN>.
2. LoRA fine-tuning. Fine-tune the model on the augmented data.
3. (Inference time) Confidence score extraction. After decoding, the probability of <CN> (appropriately normalized) is a proxy for model confidence.

Two applications are considered: 1. learning to reject, 2. model cascading.

**Reviews**

The reviewers noted that the proposed approach is simple and yet effective. The paper is well written and easy to follow (kmXF, w1rW). At a high level, there are two main claims: 1. confidence tokens align with correctness, 2. small overhead. These claims are well supported by the provided empirical results  (gYcY, JjAL). Experiments are thorough.

Most of the questions raised were mainly clarification questions  (e.g., ratio of <CN> and <UN>), to which the rebuttal addressed them well. Among the remaining concerns, as may be expected, a number of reviewers (w1rW, gYcY) asked why the special confidence token is appended at the end (as opposed to at the beginning, which can save some decoding time). While the rebuttal states that  conditioning on the generated response can help improve uncertainty quantification, no further results were provided. That the work does not extensively analyze potential failure cases (JjAL) is another concern. In the rebuttal, the authors listed a few instances from a few tasks where the method tends to fail, partially addressing this concern. It would be good if the authors could please analyze failure cases in more detail and provide general insights (as opposed to a case-by-case analysis).  Owing to the extensive and strong empirical results, the AC thinks the merits outweigh the drawbacks. 

AC recommendation: accept.
Decision: Accept (poster)