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

arXiv: https://arxiv.org/abs/2603.18009 · HTML: https://mlsanity.com/paper/2603.18009/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in RAG robustness and adaptive retrieval, one of 15 directions in Retrieval-augmented generation.
- Builds on 32 of its 39 references that are in the map.
- First posted on arXiv 23 Feb 2026.

## Placed in

- Retrieval-augmented generation / [RAG robustness and adaptive retrieval](https://mlsanity.com/rag/rag-robustness-and-adaptive-retrieval/)

## Builds on (most cited first, 30 of 32)

- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Language Models (Mostly) Know What They Know](https://mlsanity.com/paper/2207.05221/) (2022)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Selective Classification for Deep Neural Networks](https://mlsanity.com/paper/1705.08500/) (2017)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [True Few-Shot Learning with Language Models](https://mlsanity.com/paper/2105.11447/) (2021)
- [Generating with Confidence: Uncertainty Quantification for Black-box Large Language Models](https://mlsanity.com/paper/2305.19187/) (2023)
- [LLM2LLM: Boosting LLMs with Novel Iterative Data Enhancement](https://mlsanity.com/paper/2403.15042/) (2024)
- [A Systematic Survey of Automatic Prompt Optimization Techniques](https://mlsanity.com/paper/2502.16923/) (2025)
- [Insights into LLM Long-Context Failures: When Transformers Know but Don't Tell](https://mlsanity.com/paper/2406.14673/) (2024)
- [Instance-adaptive Zero-shot Chain-of-Thought Prompting](https://mlsanity.com/paper/2409.20441/) (2024)
- [On Uncertainty Calibration and Selective Generation in Probabilistic Neural Summarization: A Benchmark Study](https://mlsanity.com/paper/2304.08653/) (2023)
- [Think Just Enough: Sequence-Level Entropy as a Confidence Signal for LLM Reasoning](https://mlsanity.com/paper/2510.08146/) (2025)
- [Emergent Response Planning in LLMs](https://mlsanity.com/paper/2502.06258/) (2025)
- [Beyond Prompt Content: Enhancing LLM Performance via Content-Format Integrated Prompt Optimization](https://mlsanity.com/paper/2502.04295/) (2025)
- [Where to show Demos in Your Prompt: A Positional Bias of In-Context Learning](https://mlsanity.com/paper/2507.22887/) (2025)
- [GAAPO: Genetic Algorithmic Applied to Prompt Optimization](https://mlsanity.com/paper/2504.07157/) (2025)
- [Improve LLM-based Automatic Essay Scoring with Linguistic Features](https://mlsanity.com/paper/2502.09497/) (2025)
- [Improving LLM First-Token Predictions in Multiple-Choice Question Answering via Output Prefilling](https://mlsanity.com/paper/2505.15323/) (2025)
- [Event Extraction in Large Language Model](https://mlsanity.com/paper/2512.19537/) (2025)
- [The Fair Language Model Paradox](https://mlsanity.com/paper/2410.11985/) (2024)
- [First Hallucination Tokens Are Different from Conditional Ones](https://mlsanity.com/paper/2507.20836/) (2025)

## Built on by (most cited first, 1 of 1)

- [TRACE: Single-Pass Decoding-Trace Risk Localization for Generation Calibration](https://mlsanity.com/paper/2609.35387/) (2026)

## Abstract (arXiv)

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.