How Confident Is the First Token? An Uncertainty-Calibrated Prompt Optimization Framework for Large Language Model Classification and Understanding

Wei Chen; Guoyang Ju; Yuanyuan Qi

cs.CLFeb 23, 2026

How Confident Is the First Token? An Uncertainty-Calibrated Prompt Optimization Framework for Large Language Model Classification and Understanding

Authors: Wei Chen, Guoyang Ju, Yuanyuan Qi

Organizations: China Jiliang University, Hangzhou, China

Abstract

With the widespread adoption of large language models (LLMs) in natural language processing, prompt engineering and retrieval-augmented generation (RAG) have become mainstream to enhance LLMs' performance on complex tasks. However, LLMs generate outputs autoregressively, leading to inevitable output uncertainty. Since model performance is highly sensitive to prompt design, precise uncertainty measurement is crucial for reliable prompt optimization. For multi-class multiple-choice (understanding) tasks, conventional uncertainty measures (e.g., entropy) based on output probabilities treat all classes equally and ignore class prior differences in pretraining corpora. This failure to distinguish spurious confidence (from priors) from true certainty (from contextual understanding) results in poor confidence calibration. To address this, we propose Log-Scale Focal Uncertainty (LSFU), a first-token-based metric inspired by focal loss. LSFU incorporates label prior probabilities as a risk-modulation factor to suppress noise from high-frequency classes and emphasize risk for low-frequency long-tail classes, with a dynamic weighting mechanism unifying the measurement scale. Based on LSFU, we further propose the uncertainty-calibrated prompt optimization framework (UCPOF), which leverages the first token of model outputs to select high-quality exemplars and dynamically optimize prompts. Comprehensive evaluations show UCPOF improves average accuracy by 6.03% over few-shot baselines, surpasses always-on full RAG by 5.75% in overall average accuracy, and reduces the average retrieval trigger rate by 50.66%. By adaptively triggering RAG only for high-uncertainty samples, our framework significantly lowers computational costs while maintaining state-of-the-art performance.

Figures & tables

Figure 1: Schematic overview of the UCPOF workflow (offline preparation and online inference).

Figure 2: Structure of the reflection prompt for conditional correction.

Figure 3: Distribution density of LSFU values for correct and incorrect predictions.

Panel of Figure 4 — Figure 4: Risk-coverage curves of LSFU-based selective prediction on CASIE and ACE datasets.

Figure 5: Relationship between error coverage and RAG trigger rate for threshold selection.

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
Task $\rightarrow$ Example $\rightarrow$ Input	42.81	82.13	78.80	50.50	69.64	91.73	69.27
Example $\rightarrow$ Task $\rightarrow$ Input	51.07	76.40	82.42	51.45	78.57	92.46	72.06

Table 1: Impact of Prompt Component Ordering on Accuracy %

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	51.07	76.40	82.42	51.45	78.57	92.46	72.06
gold shot (Ours)	50.67	76.73	84.90	51.90	83.93	94.75	73.81
Full RAG	61.47	84.20	91.42	51.40	48.21	97.35	72.34
UCPOF (Ours)	62.08	85.40	89.99	51.54	82.14	97.40	78.09

Table 2: Accuracy Comparison Experiment (Accuracy %)

Strategy	ACE	Casie	AgNews	SST5	CB	R8	AVG
Random	51.07	76.40	82.42	51.45	78.57	92.46	72.06
Highest LSFU	45.87	74.87	87.62	53.98	69.64	90.95	70.49
Lowest LSFU (ours)	50.67	76.73	84.90	51.90	83.93	94.75	73.81

Table 3: Gold Shot Selection Strategy Comparison (Stage 1 Accuracy %)

Strategy	ACE	Casie	AgNews	SST5	CB	R8	AVG
UCPOF	62.08	85.40	89.99	51.54	82.14	97.40	78.09
w/o Gold Shot	61.47	84.93	89.08	51.53	71.43	96.90	75.89

Table 4: Ablation of Gold Shot Selection in UCPOF

model		ACE	Casie	AgNews	SST5	CB	R8	AVG
Main model		0.8104	0.7492	0.7761	0.5242	0.7568	0.9161	0.7555
	w/o $P_{\mathrm{prior}}$	0.8004	0.7338	0.7762	0.5235	0.7154	0.9083	0.7429

Table 5: Ablation of

P_{\mathrm{prior}}

(AUC Metric)

Method		ACE	Casie	AgNews	SST5	CB	R8	AVG
UCPOF		62.08	85.40	89.99	51.54	82.14	97.40	78.09
	w/o $P_{\mathrm{prior}}$	61.39	84.80	89.08	52.44	60.71	97.85	74.37

Table 6: Ablation of

P_{\mathrm{prior}}

in UCPOF (ACC Metric)

Metric		ACE	Casie	AgNews	SST5	CB	R8	AVG
ECE		0.2240	0.1354	0.0978	0.3545	0.1245	0.1422	0.1797
	w/o $P_{\mathrm{prior}}$	0.2271	0.1365	0.1015	0.3547	0.2046	0.1421	0.1944
Brier score		0.2604	0.1368	0.0978	0.3896	0.1246	0.1431	0.1920
	w/o $P_{\mathrm{prior}}$	0.2634	0.1374	0.1023	0.3896	0.2031	0.1426	0.2064

Table 7: Ablation of

P_{\mathrm{prior}}

(Calibration Metrics; lower is better)

Figure 6: Efficiency–Accuracy Pareto Frontier (ACE; Qwen2.5-7B-Instruct).

Figure 7: RAG Trigger Rate vs. Test Accuracy Trend (ACE; Qwen2.5-7B-Instruct).

Method	Metric	ACE	Casie	AgNews	SST5	CB	R8	AVG
UCPOF	Trigger rate (%)	77.06	50.60	20.87	95.48	16.07	35.95	49.34
UCPOF	Acc	62.08	85.40	89.99	51.54	82.14	97.40	78.09
Full RAG	Trigger rate (%)	100.00	100.00	100.00	100.00	100.00	100.00	100.00
Full RAG	Acc	61.47	84.20	91.42	51.40	48.21	97.35	72.34

Table 8: Efficiency comparison between UCPOF and Full RAG: retrieval trigger rate (lower is better) and test accuracy (higher is better).

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 8 — Figure 8: Efficiency analysis on CASIE.

Panel of Figure 9 — Figure 9: Efficiency analysis on AgNews.

Panel of Figure 10 — Figure 10: Efficiency analysis on CB.

Panel of Figure 11 — Figure 11: Efficiency analysis on R8.

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	33.94	62.73	80.47	48.10	64.29	93.19	63.79
gold shot (Ours)	37.61	65.53	82.24	46.52	62.50	93.51	64.65
Full RAG	60.55	82.73	90.92	50.63	82.14	96.03	73.83
UCPOF (Ours)	55.96	81.73	90.80	47.19	83.93	96.39	76.00

Appendix

Table 9: Cross-Model Experimental Results - Llama-3-8B-Instruct

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	52.60	83.07	80.25	53.08	87.50	92.51	74.82
gold shot (Ours)	47.09	86.53	86.40	53.26	81.36	93.42	74.68
Full RAG	65.44	89.47	91.50	54.12	82.14	97.94	80.10
UCPOF (Ours)	66.97	89.33	89.87	54.21	83.93	97.99	80.38

Appendix

Table 10: Cross-Model Experimental Results - Qwen2.5-14B-Instruct

model		ACE	Casie	AgNews	SST5	CB	R8	AVG
Qwen2.5-14B-instruct		0.6742	0.7295	0.7763	0.5409	0.8528	0.9058	0.7466
	w/o $P_{\mathrm{prior}}$	0.6666	0.7240	0.7229	0.5427	0.8364	0.8941	0.7311
Llama-3-8B-Instruct		0.7138	0.7603	0.8083	0.4869	0.9271	0.7976	0.7490
	w/o $P_{\mathrm{prior}}$	0.6183	0.7136	0.7876	0.4852	0.8954	0.7579	0.7097

Appendix

Table 11: Ablation of

P_{\mathrm{prior}}

(AUC Metric) on other models.

model		ACE	Casie	AgNews	SST5	CB	R8	AVG
Qwen2.5-14B-instruct		0.6742	0.7295	0.7763	0.5409	0.8528	0.9058	0.7466
	w/o $P_{\mathrm{prior}}$	0.6666	0.7240	0.7229	0.5427	0.8364	0.8941	0.7311
Llama-3-8B-Instruct		0.7138	0.7603	0.8083	0.4869	0.9271	0.7976	0.7490
	w/o $P_{\mathrm{prior}}$	0.6183	0.7136	0.7876	0.4852	0.8954	0.7579	0.7097

Appendix

Table 12: Ablation of

P_{\mathrm{prior}}

(AUC Metric) on other models.

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	42.81	82.13	78.80	50.50	69.64	91.73	69.27
gold shot (Ours)	42.81	76.20	85.27	53.26	73.21	92.69	70.57
Full RAG	61.77	83.47	91.71	52.26	50.00	96.67	72.65
UCPOF (Ours)	62.69	84.33	89.72	52.35	70.86	95.84	75.97

Appendix

Table 13: Cross-Model Experimental Results-Qwen2.5-7B-Instruct (Task→Example→Input)

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	49.54	83.53	83.72	55.61	87.50	95.29	75.87
gold shot (Ours)	42.53	87.40	84.98	51.90	80.36	95.80	73.82
Full RAG	67.58	90.00	91.75	53.67	85.71	97.99	81.12
UCPOF (Ours)	67.58	89.93	88.27	53.98	89.29	97.99	81.17

Appendix

Table 14: Cross-Model Experimental Results-Qwen2.5-14B-Instruct (Task→Example→Input)

method	ACE	Casie	AgNews	SST5	CB	R8	AVG
baseline	41.28	80.33	77.15	44.89	53.21	95.34	65.37
gold shot (Ours)	39.49	73.27	80.16	48.51	67.86	89.17	66.41
Full RAG	57.19	82.53	90.71	51.72	82.14	91.41	75.95
UCPOF (Ours)	57.79	81.16	88.68	51.72	82.14	96.16	76.27

Appendix

Table 15: Cross-Model Experimental Results-LLama-3-8B-Instruct (Task→Example→Input)

K	ACE	Casie	AgNews	SST5	CB	R8
50	0.443084	0.187070	0.026404	0.188524	0.104495	0.042302
100	0.451381	0.187354	0.026403	0.188524	0.104495	0.042302
500	0.451371	0.187555	0.026412	0.188525	0.104495	0.042343
ALL	0.456750	0.187555	0.026422	0.188537	0.104438	0.042370

Appendix

Table 16: Top-K Entropy Values with Different K (LSFU Calculation Stability Validation)

LLM Input Optimization

See all 24 papers

Confidence Before Answering: A Paradigm Shift for Efficient LLM Uncertainty Estimation
Mar 6, 2026Changcheng Li, Jiancan Wu, Hengheng Zhang +5Confidence Estimation in Language Models LLM Uncertainty Estimation
University of Science and Technology of China · Huawei Inc.
LogitScope: A Framework for Analyzing LLM Uncertainty Through Information Metrics
Mar 26, 2026Farhan Ahmed, Yuya Jeremy Ong, Chad DeLucaLLM Uncertainty Estimation Uncertainty Quantification
IBM Research · Plastic Labs
Probability is Not Enough: Exploring and Counting Divergent Tokens for Reasoning Uncertainty Quantification in LLMs
Sep 29, 2026Feiyang Li, Shengjing Liu, Qi Zhan +7Uncertainty Quantification Confidence Estimation in Language Models
College of Computer Science and Software Engineering, Shenzhen University · Tsinghua University · Behavioral and Spatial AI Lab, Peking University & Tongji University

Mar 6, 2026J/K move · Enter open · S save

Confidence Before Answering: A Paradigm Shift for Efficient LLM Uncertainty Estimation

Changcheng Li, Jiancan Wu, Hengheng Zhang, Zhengsu Chen+4

University of Science and Technology of China · Huawei Inc.

Confidence Estimation in Language Models LLM Uncertainty Estimation Uncertainty Calibration RL for Language Models

Open paper AskPDF arXiv

Reliable deployment of large language models (LLMs) requires accurate uncertainty estimation. Existing methods are predominantly answer-first, producing confidence only after generating an answer, which measure the correctness of a specific response and limits practical usability. We study a confidence-first paradigm, where the model outputs its confidence before answering, interpreting this score as the model's probability of answering the question correctly under its current policy. We propose CoCA(Co-optimized Confidence and Answers), a GRPO reinforcement learning framework that jointly optimizes confidence calibration and answer accuracy via segmented credit assignment. By assigning separate rewards and group-relative advantages to confidence and answer segments, CoCA enables stable joint optimization and avoids reward hacking. Experiments across math, code, and factual QA benchmarks show improved calibration and uncertainty discrimination while preserving answer quality, thereby enabling a broader range of downstream applications.

Abstract

LLM Input Optimization

Explore similar work

Confidence Before Answering: A Paradigm Shift for Efficient LLM Uncertainty Estimation

LogitScope: A Framework for Analyzing LLM Uncertainty Through Information Metrics

Probability is Not Enough: Exploring and Counting Divergent Tokens for Reasoning Uncertainty Quantification in LLMs