PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Route LLMs with Confidence Tokens

Yu-Neng Chuang, Prathusha Kameswara Sarma, Parikshit Gopalan, John Boccio, Sara Bolouki, Xia Hu, Helen Zhou

arXiv:2410.13284v3cs.CLcs.AIcs.LG

TL;DR

LLMs can struggle to estimate whether their predictions are correct, even though confidence can guide routing or abstention. Self-REF fine-tunes models to express confidence with tokens attached to predictions. The paper reports that it outperforms baselines in routing and rejection learning across four public datasets.

  • Problem

    LLMs may estimate confidence poorly: token probabilities can be misaligned with correctness, and verbalized confidence can be unstable across datasets and prompts.

    Source: §1 · p. 1

  • Method

    Self-REF fine-tunes LLMs to append confidence tokens to predictions, using confident tokens for correct answers and unconfident tokens for incorrect ones.

    Source: §1 · pp. 1–2

  • Results

    Self-REF outperforms baselines in routing and rejection learning on four public datasets.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report that confidence scores are useful for routing and rejection learning, with potential to improve system performance, efficiency, and output reliability.

    Source: §7 · p. 9

    The cost analysis measures GPU hours and latency but excludes network transmission time between models, which may contribute substantially to routing savings.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Self-REF addresses the challenge of estimating whether LLM predictions are correct, a capability needed to route uncertain answers or fall back safely. It trains models to append confidence tokens and uses their scores for routing and rejection.

  • Reliable confidence estimates matter because systems can route uncertain answers to another expert or abstain instead of relying on a potentially faulty output.Source: §1 · pp. 1–2
  • Token probabilities may not align with correctness, while verbalized confidence can be unstable across datasets and prompts.The paper asks how to train LLMs to estimate confidence while preserving performance on tasks of interest.Source: §1 · p. 1
  • Self-REF fine-tunes LLMs to append confident tokens to correct predictions and unconfident tokens to incorrect predictions.It extracts a continuous score by normalizing the confident-token probability over the confident- and unconfident-token probabilities.Source: §1 · pp. 1–2
  • The paper studies routing low-confidence queries from a smaller local LLM to a stronger model and rejection by deferring when none of the options is confident.Source: §1 · p. 2
  • Self-REF outperforms baselines in routing and rejection learning on four public datasets.Source: §1 · p. 2

2. Related Work

The related work surveys confidence estimation, LLM routing, and rejection learning. It positions Self-REF as learning confidence against prediction correctness within the LLM and applying it to downstream routing and rejection.

2. Related Work

  • The section situates the study across uncertainty quantification, LLM routing, and LLM rejection learning.Source: §2 · p. 2

2.1. Uncertainty Quantification in LLMs

  • Prior uncertainty methods use internal signals, post-processing, prompting, or training; Self-REF learns confidence during fine-tuning and aligns it with correctness rather than response consistency.The paper notes that uncertainty scores may correlate weakly with correctness and that consistent answers are not necessarily the most correct.Source: §2.1 · p. 2

2.2. LLM Routing

  • Self-REF estimates correctness inside the LLM during autoregressive decoding, unlike routing approaches that use an external model or query-only signals.The paper frames its downstream contribution as routing and rejection utility beyond calibration.Source: §2.2 · p. 3

2.3. LLM Rejection Learning

  • Prior rejection methods use extra knowledge, agents, or loss adaptations, which may degrade base performance; Self-REF aims to support rejection without designing a new loss.Source: §2.3 · p. 3

3. Learning Confidence Tokens

Self-REF fine-tunes a local LLM to associate its predictions with confidence tokens, then extracts a confidence score for routing or rejection. Its setup uses labeled query-answer data and defines confidence by whether the model prediction matches the ground truth.

3.1. Problem Setup and Notation

  • The setup defines a local LLM as one available for fine-tuning and uses query-answer examples split into training, validation, and test data.The confident token is intended for predictions matching the ground truth, and the unconfident token for predictions that do not.Source: §3.1 · p. 3

3.2. Self-REF

  • Self-REF annotates base-model predictions with <CN> when correct and <UN> when incorrect, then combines the augmented examples for fine-tuning.The annotation supplies correctness information alongside the model's own predictions.Source: §3.2 · p. 3
  • During fine-tuning, gradients for incorrect answers are masked so the model learns confidence-token probabilities without increasing the probability of an incorrect answer.Confidence tokens are generated end-to-end conditioned on both the query and the model's generated answer, unlike the cited prior approaches.Source: §3.2 · p. 3
  • The score normalizes the <CN> probability over the <CN> and <UN> probabilities, providing a continuous value for downstream trade-offs.Self-REF customizes uncertainty using the LLM's own predictions rather than relying on the inherent uncertainty of a query.Source: §3.2 · p. 4

3.3. Learning to Route and Reject

  • In routing, answers below a chosen confidence threshold go to a more powerful LLM, while sufficiently confident local-model answers are returned directly.Source: §3.3 · p. 4
  • For rejection evaluation, half the samples omit the ground truth and use “none of the above”; low-confidence predictions count as abstentions.Source: §3.3 · p. 4

4. Experiment Setup

The experiments evaluate confidence-based routing across four public datasets and two local LLMs, and test confidence-based rejection on MMLU and OpenbookQA against confidence-estimation baselines.

4.1. Datasets and Baselines

  • The four public evaluation datasets are MMLU, OpenbookQA, GSM8K, and MedQA, spanning broad knowledge, science reasoning, math word problems, and medical-board questions.Source: §4.1 · p. 4
  • The comparison includes four confidence-estimation baselines, including verbalized uncertainty, verbalized yes/no confidence, and zero-shot and fine-tuned output logits.Logit baselines derive confidence from the predicted answer’s token probability or, for free-text answers, the average output-token probabilities.Source: §4.1 · pp. 4–5

4.2. Experiment Settings

  • Routing decisions use the normalized probability of the <CN> token over the probabilities of <CN> and <UN>, routing queries below a threshold to Llama3-70B-Instruct.The larger model uses in-context learning only, without further fine-tuning.Source: §4.2 · p. 5
  • Routing thresholds are set at 21 confidence-score quantiles, corresponding to approximately 0%–100% of queries routed in 5% increments.Queries are routed when their confidence score falls below the threshold.Source: §4.2 · p. 5
  • The rejection evaluation removes the correct choice in half the questions, making “none of the above” the ground truth, and thresholds confidence to trade off rejection against false rejection.Self-REF receives no additional training samples with “none of the above” as a choice.Source: §4.2 · pp. 5–6
  • The local models are Llama-3-8B-Instruct and Mistral-7B-v0.3-Instruct; LoRA adapters use ranks 2, 4, and 8, with other training hyperparameters selected by validation perplexity grid search.Training runs for at most eight epochs.Source: §4.2 · p. 6

5. Experiment Results

The experiments evaluate Self-REF on routing, rejection, calibration utility, gradient masking, transfer, and case behavior. Self-REF leads routing and rejection comparisons, though calibration does not always predict routing utility and some failures remain.

5.1. Routing Performance (RQ1)

  • 2.03× latency improvement for Llama3-8B-Instruct on MMLU and 2.00× for Mistral-7B-Instruct on OpenbookQA, without sacrificing accuracy versus Llama3-70B-Instruct.Source: §5.1 · p. 6, Table 1 · p. 7
  • Self-REF achieves the best accuracy-versus-routing-rate trade-off across four datasets and both local LLMs.Source: §5.1 · p. 6, Table 2 · p. 7

5.2. Rejection Learning Performance (RQ2)

  • Self-REF outperforms all baselines on rejection learning across thresholds on MMLU and OpenbookQA.Source: §5.2 · p. 7

5.3. Calibration with Utility (RQ3)

  • Better calibration scores do not guarantee optimal routing, although Self-REF has superior routing performance across datasets and local models.Source: §5.3 · p. 7, Table 2 · p. 7

5.4. Ablation Study on Gradient Masking

  • On OpenbookQA, gradient masking consistently improves routing accuracy for both Llama3-8B-Instruct and Mistral-7B-Instruct.Source: Figure 4 · p. 8, §5.4 · p. 8

5.5. Transferability of Confidence Token

  • OpenbookQA-trained Self-REF transfers to MMLU, achieving parity with the 70B model at 70% routing for Llama3-8B-Instruct and 75% for Mistral-7B-Instruct.Source: Figure 5 · p. 8, §5.5 · p. 8

5.6. Case Studies

  • Failure cases include underconfidence on context-heavy or ambiguous subjects and overconfidence on technical, structured, or jurisprudence tasks.Source: §5.6 · pp. 8–9

6. Discussion

Self-REF can match or outperform a non-fine-tuned 70B model in system-level routing, and routing can act as a form of ensembling. Its confidence scores also support granular downstream choices whose trade-offs depend on routing costs and the risk of incorrect answers.

6. Discussion

  • Self-REF achieves comparable or sometimes superior system-level routing performance to a non-fine-tuned 70B model.The comparison reflects a setting where practitioners fine-tune a smaller model but may not be able to fine-tune the larger one.Source: §6 · p. 9
  • Routing can serve as a form of ensembling because independent models’ errors are unlikely to overlap perfectly.The authors suggest that Self-REF can preserve performance and introduce correct knowledge, including niche astronomy knowledge or complicated moral disputes.Source: §6 · p. 9
  • Self-REF provides continuous confidence scores for granular control, whereas verbalized uncertainty can suffer from mode collapse.The resulting score supports choices among downstream routing and rejection trade-offs.Source: §6 · p. 9
  • Routing decisions should reflect the setting’s cost of invoking a larger model, including latency and monetary cost.A constrained API budget or high Mars-to-Earth latency may favor routing sparingly; when cost and latency are not concerns, the highest-accuracy point can be chosen.Source: §6 · p. 9
  • Rejection thresholds should reflect the risk of incorrect answers, with higher abstention rates potentially appropriate for medical treatment than creative work.The settings are customizable according to real-world parameters.Source: §6 · p. 9
  • The authors report improved system performance and safety while reducing cost through confidence-based routing and rejection, but exclude inter-model network transmission time from cost measurements.They note that networking costs could contribute substantially to the savings from confidence-based routing.Source: §6 · p. 9

7. Conclusion

The conclusion presents Self-REF as confidence-token fine-tuning intended to align confidence with answer correctness, and identifies extensions for routing and rejection. The accompanying related-work and appendix passages position its focus alongside uncertainty estimation, routing, rejection learning, and prompt evaluation.

7. Conclusion

  • Self-REF uses confidence-token fine-tuning to align expressed confidence with prediction correctness; proposed extensions include routing among specialized models and finer-grained confidence or rejection signals.Possible extensions include sentence-level labels, stopping early or refining outputs, and explanations or clarifying information for rejection.Source: §7 · p. 9

A.1. Uncertainty Quantification in LLMs

  • Calibration work spans internal-feature estimates, post-processing, prompting, and instruction learning; this paper focuses on downstream routing and rejection utility beyond calibration.The related-work discussion also describes uncertainty estimates based on question semantics and ambiguity.Source: §A.1 · p. 14

A.2. LLM Routing

  • Routing literature covers query dispatch, retrieval-augmented generation, and speculative decoding; Self-REF instead routes according to confidence in each generated answer.The passage distinguishes this approach from routing based on extra query routers and describes confidence estimates produced from model outputs during autoregressive decoding.Source: §A.2 · p. 14

A.3. LLM Rejection Learning

  • Rejection approaches include training with rejection objectives or separate rejection agents; this work aims to enable rejection without new loss functions or additional agents.The passage notes possible degradation to next-token prediction and added latency from those alternatives.Source: §A.3 · pp. 14–15

B.1. Verbalizing Confidence Prompt Usage

  • Table 3 presents evaluation prompts across datasets for verbalized uncertainty, verbalized yes/no tokens, and zero-shot or fine-tuned logits.The table caption identifies prompts for verbalizing uncertainty on local LLMs under different datasets.Source: §B.1 · p. 15, Table 3 · p. 15

B.2. Finetuned Logits Training

The Finetuned Logits baseline uses LoRA training, with the reported performance selected by downstream task accuracy.

B.2. Finetuned Logits Training

  • LoRA dimensions of 2, 4, and 8 are listed for training the Finetuned Logits baseline.The passage says its hyperparameter settings are shown in Table 5.Source: §B.2 · p. 16
  • The reported Finetuned Logits performance is the configuration with the highest downstream task accuracy.The passage describes this selection among the LoRA training settings.Source: §B.2 · p. 16
  • Table 4 is captioned as giving hyperparameter settings for the Fine-tuned Logits baseline.Source: Table 4 · p. 16

B.3. Details of Datasets

The experiments use four open-sourced datasets spanning broad knowledge questions, elementary science, grade-school math, and medical question answering.

B.3. Details of Datasets

  • MMLU covers 57 subjects across humanities, social sciences, STEM, and professional domains; evaluation uses its validation set.It is a diverse zero-shot multiple-choice benchmark assessing knowledge acquired during pretraining.Source: §B.3 · p. 16
  • OpenbookQA contains 5,957 elementary-science multiple-choice questions requiring core facts and broader knowledge or reasoning.Source: §B.3 · p. 16
  • GSM8K evaluates grade-school arithmetic reasoning, and the experiments assess its 1,331-question test set.The dataset includes questions and detailed step-by-step solutions.Source: §B.3 · p. 16
  • MedQA contains 12,723 medical multiple-choice questions based on real-world licensing exams and requiring specialized knowledge and multi-step reasoning.Source: §B.3 · p. 16

C. Experiment and Model Training Details

The experiments describe Self-REF fine-tuning and inference settings, plus the computing infrastructure used for training and evaluation.

C.1. Hyper-parameter Settings

  • Self-REF fine-tuning uses LoRA on query, key, value, embedding, and final linear layers, with batch size 4 and learning rate 1e-4.The unconfident-to-confident data ratio is tuned from 1:1 to 1:5; the authors note that confident examples aid downstream performance while unconfident examples teach uncertainty.Source: §C.1 · p. 17, Table 5 · p. 17
  • Inference uses temperature 0, top-p 1.0, greedy decoding, no other sampling strategies, and fixed random seed 2024.Source: §C.1 · p. 17

C.2. Computation Infrastructure

  • The experiments use the computing infrastructure listed in Table 6, which the authors say supported model training and evaluation.Source: §C.2 · p. 17, Table 6 · p. 17

D. Case studies of Self-REF Routing

MMLU case studies illustrate two routing patterns: a local model answering correctly when the larger model answers incorrectly, and an initially incorrect answer corrected after routing.

E. Case studies by Subject Area

The subject-area case study groups predictions by correctness and confidence, reporting distinct subject patterns for each of the four groups.

  • Correct but underconfident predictions tend to involve ambiguity, contextual reasoning, or non-standard formats, including human sexuality, European history, and miscellaneous topics.The authors associate “Miscellaneous” with questions outside well-defined training distributions.Source: E. Case studies by Subject Area · p. 20
  • Incorrect but overconfident predictions are common in structured technical subjects, including computer science, physics, and jurisprudence.The authors suggest frequent pretraining patterns may contribute, while nuanced reasoning and edge cases remain challenging.Source: E. Case studies by Subject Area · p. 20
  • Correct and confident predictions appear in subjects including international law, college biology, moral disputes, philosophy, and U.S. foreign policy.The authors note these topics often involve conceptual reasoning with clear evaluative criteria.Source: E. Case studies by Subject Area · p. 20
  • Incorrect but underconfident predictions appear in exact-recall and specialized subjects, including abstract algebra, anatomy, chemistry, medicine, and econometrics.The authors interpret this hesitation as awareness of knowledge limitations.Source: E. Case studies by Subject Area · p. 20

Related papers

Loading 2410.13284v3…