Research map / Large language models / Reasoning and chain of thought
Learning to Route LLMs with Confidence Tokens
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Self-REF, fine-tuning LLMs with confidence tokens labeled by answer correctness, so the model's confidence score drives routing to a larger LLM and rejection of "none of the above" questions.
Problem
LLM confidence estimates (token probabilities, verbalized confidence) are poorly calibrated and unstable, which limits their use for routing queries to stronger models and for abstaining from answering.
Approach
Attach trainable confident and unconfident tokens to the LLM's own generated answers, labeled by whether the answer is correct, fine-tune with gradients masked on incorrect answers, and extract a confidence score from the normalized probability of the confident token for routing and rejection.
Key results
- On MMLU, Self-REF on Llama3-8B-Instruct routes 39% of queries to match Llama3-70B-Instruct accuracy (0.739), with up to 2.03x latency speed-up.
- Self-REF on Llama3-8B-Instruct needs routing rates of 49%, 65%, and 40% on OpenbookQA, GSM8K, and MedQA to match the 70B model.
- Self-REF achieves the lowest routing rate across datasets, though it does not always have the best calibration scores (e.g., Finetuned Logits has lower BS and CE in many cases).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use confidence-token scores for routing instead of baseline confidence scores | Llama3-8B-Instruct routing to Llama3-70B-Instruct · MMLU, OpenbookQA, GSM8K, and MedQA | route rate to match Llama3-70B-Instruct performance: 39% vs 90%; 49% vs 70%; 65% vs 70%; 40% vs 80% | Helped |
| Use confidence-token scores for routing instead of baseline confidence scores | Mistral-7B-Instruct routing to Llama3-70B-Instruct · MMLU, OpenbookQA, GSM8K, and MedQA | route rate to match Llama3-70B-Instruct performance: 70% vs 95%; 40% vs 50%; 75% vs 100%; 70% vs 90% | Helped |
| Route using Self-REF confidence rather than using Llama3-70B-Instruct for every query | Llama3-8B-Instruct routing · MMLU | latency: 0.292 -> 0.145 (2.03$\times$) | Helped |
| Route using Self-REF confidence rather than using Llama3-70B-Instruct for every query | Mistral-7B-Instruct routing · OpenbookQA | latency: 0.292 -> 0.147 (2.00$\times$) | Helped |
| Use confidence-token scores to threshold rejection decisions | Self-REF · MMLU and OpenbookQA rejection learning | ROC performance | Helped |
| Apply gradient masking to incorrect-answer samples with unconfident tokens (ablation) | Self-REF · OpenbookQA routing | routing accuracy | Helped |
| Transfer Self-REF LoRA weights trained on OpenbookQA to MMLU | Llama3-8B-Instruct and Mistral-7B-Instruct · MMLU routing | route rate to match Llama3-70B-Instruct performance: 75% and 70%; direct fine-tuning: 70% and 40% | Mixed |
| Increase the proportion of confident samples relative to unconfident samples (ablation) | Self-REF training data · Downstream tasks | downstream task performance | Mixed |
Surprising findings
- Routing a fraction of queries to the local fine-tuned model can outperform the non-fine-tuned 70B model alone, which the authors attribute to knowledge injection and ensembling effects.
- Better calibration (ECE/Brier) does not guarantee better downstream routing utility.
Limitations
- Routing and rejection evaluated only on four multiple-choice and math/medical QA datasets with two small local models.
- Network transmission costs between models excluded from cost analysis.
- Some calibration metrics are not best, and gains in routing do not map to calibration quality.
Experimental setup
Datasets: MMLU OpenbookQA GSM8K MedQA
Benchmarks: MMLU OpenbookQA GSM8K MedQA
Metrics: accuracy vs routing rate routing rate to match 70B latency ECE Brier score cross-entropy ROC TPR/FPR for rejection
Baselines: Verbalizing Uncertainty Verbalizing Yes/No Tokens Zero-shot Logits Fine-tuned Logits
Builds on: LoRA temperature scaling Self-RAG
Closest papers by meaning
- SteerConf: Steering LLMs for Confidence Elicitation (NeurIPS 2025 · similarity 0.91)
- Can LLMs Express Their Uncertainty? An Empirical Evaluation of Confidence Elicitation in LLMs (ICLR 2024 · similarity 0.91)
- Diagnosing Multi-step Reasoning Failures in Black-box LLMs via Stepwise Confidence Attribution (ICML 2026 · similarity 0.90)
- ConfTuner: Training Large Language Models to Express Their Confidence Verbally (NeurIPS 2025 · similarity 0.89)
- HonestLLM: Toward an Honest and Helpful Large Language Model (NeurIPS 2024 · similarity 0.89)
- Who Knows the Answer? Finding the Best Model and Prompt for Each Query Using Confidence-Based Search (AAAI 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.