PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2603.18009v1cs.CLcs.AI

TL;DR

Conventional uncertainty measures for classification can overlook label-prior differences, weakening confidence calibration for prompt optimization. The paper proposes first-token LSFU and builds UCPOF to select exemplars and conditionally correct prompts; UCPOF reports 78.09% average accuracy, 6.03% higher than Baseline, with lower inference cost than full RAG.

  • Problem

    For classification, entropy-based uncertainty ignores class-prior differences, while always-on RAG can add cost and distracting context.

    Source: §1 · p. 2

  • Method

    LSFU reweights first-token entropy using label priors, and UCPOF uses it for exemplar selection and two-stage static prompt optimization plus dynamic correction.

    Source: §1 · p. 3, §4.1 · p. 8, §7 · p. 20

  • Results

    78.09% average accuracy, 6.03% higher than Baseline, with inference cost lower than full RAG.

    Source: §7 · p. 20

  • Takeaways & Limitations

    The reported results present UCPOF as balancing prompt-optimization effectiveness and efficiency through static optimization and dynamic correction.

    Source: §7 · p. 20

    The method is designed only for classification tasks, and generalization of its dynamic threshold depends on domain-data distribution stability.

    Source: §7 · pp. 20–21

Abstract

from arXiv · show

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.

1 Introduction

For classification, conventional uncertainty measures can mistake confidence driven by class priors for contextual understanding, leaving prompt optimization without reliable guidance. The paper proposes first-token LSFU and uses it to select exemplars and gate dynamic prompt correction.

1 Introduction

  • Entropy ignores class-prior differences, so low uncertainty for frequent labels may not indicate the same mastery as low uncertainty for rare labels.This can lead entropy-based selection to favor redundant frequent-class examples and miss representative low-frequency ones.Source: §1 · p. 2
  • LSFU reweights first-token entropy with label priors, suppressing high-frequency-class noise while emphasizing prediction risk for low-frequency classes.The metric is proposed as a prompt-optimization calibration measure that requires no additional training.Source: §1 · pp. 2–3
  • Selecting few-shot exemplars with the lowest LSFU values provides more stable classification anchors than random or similarity-based selection.The paper reports that this confidence-based strategy significantly improves the in-context learning baseline.Source: §1 · p. 3
  • UCPOF uses LSFU to trigger retrieval and prompt correction only for high-risk samples, rather than retrieving for every sample.The paper presents this as a way to reduce inference cost and filter noise from always-on RAG while improving efficiency and accuracy.Source: §1 · p. 3

2 Related Work

The related work covers prompt optimization, uncertainty quantification, adaptive retrieval, and formal definitions for the paper’s setting. It identifies confidence calibration and label-prior awareness as gaps in existing approaches.

2 Related Work

  • The related-work section reviews in-context learning, uncertainty quantification, retrieval-augmented generation, adaptive inference, and formal definitions for the setting.The section’s stated scope is a review of these areas and definitions relevant to the paper.Source: §2 · p. 3

2.1 In-Context Learning and Automated Prompt Optimization

  • Prompt design affects in-context learning across exemplars, ordering, instructions, positional bias, and semantic realization, while most optimization methods focus on accuracy rather than confidence calibration.The cited gap leaves exemplar selection and dynamic adjustment without reliable confidence signals, particularly for long-tail or complex samples.Source: §2.1 · pp. 3–4

2.2 Uncertainty Quantification for LLMs

  • Logit-based uncertainty methods are computationally efficient, but most ignore label priors and therefore struggle to distinguish low-frequency hesitation from high-frequency overconfidence.First-token distributions have been studied as an efficient uncertainty signal, but the paper identifies prior bias as an unresolved limitation.Source: §2.2 · p. 4

2.3 Retrieval-Augmented Generation and Dynamic Inference

  • Always-on RAG adds computational cost and latency and can introduce distracting context, while UCPOF gates retrieval and prompt correction with LSFU.Prior dynamic frameworks are described as having gating signals that may fail to account for label-prior bias.Source: §2.3 · pp. 4–5

2.4 Formal Definitions

  • The formal definitions specify ICL as conditional label prediction, RAG as generation with retrieved documents, and UCPOF as a score-threshold choice between static prompting and RAG.The decision rule uses the LSFU score S(x), threshold τ, and a statically optimized prompt C∗.Source: §2.4 · p. 5

3 Theoretical Motivation and Analysis

The theoretical motivation treats the first-token distribution as an early proxy for predictive confidence. The paper describes properties of its uncertainty function and frames prompt selection as a risk-aware trade-off between errors and computational cost.

3.1 First-Token Confidence Hypothesis

  • The First-Token Confidence Hypothesis treats concentration in the first-token distribution as an early proxy for overall predictive confidence and risk.The authors motivate this with the claim that uncertainty at the initial decision stage can be amplified during subsequent generation.Source: §3.1 · p. 6

3.2 Properties of the Proposed Uncertainty Metric

  • The proposed uncertainty function decreases as the first-token distribution concentrates, is bounded, responds continuously to small perturbations, and can connect to classical entropy- or margin-based measures.The stated properties are intended to support comparability, threshold decisions, and stable early uncertainty assessment.Source: §3.2 · p. 6

3.3 Risk-Aware Prompt Decision Interpretation

  • The risk-aware decision interpretation uses U(x) below a threshold for a basic prompt and above it to trigger enhanced strategies such as retrieval or verification.The authors describe this as a practical approximation of selective prediction that trades prediction risk against computational cost.Source: §3.3 · p. 7

4 Method

The method introduces LSFU as a prior-calibrated risk score for prompt optimization, then uses it to select static exemplars and gate conditional prompt correction. UCPOF prepares a static prompt and threshold offline, then uses them to decide whether online inference needs retrieval and correction.

4.1 Log-Scale Focal Uncertainty (LSFU)

  • LSFU combines top-K first-token entropy with a predicted-class prior to provide a log-scaled risk criterion for exemplar selection and dynamic correction.The focal-loss-inspired prior factor suppresses high-prior-class noise and emphasizes risk for low-prior classes; a log transform preserves discriminability, with ε for numerical stability.Source: §4.1 · pp. 7–8

4.2 Static Optimization: LSFU-based Gold Shot Selection

  • Gold-shot selection chooses correctly predicted training samples with the lowest LSFU values to build a static prompt with strong guidance and low noise.The selected exemplars are intended to provide clear decision anchors while avoiding the blurred guidance of boundary cases.Source: §4.2 · pp. 8–9

4.3 Uncertainty-Calibrated Prompt Optimization Framework (UCPOF)

  • UCPOF makes static-prompt inference the baseline and uses LSFU to trigger reflective-prompt correction only when the sample is insufficiently guided.The framework is organized around offline preparation and online inference.Source: §4.3 · pp. 9–10

4.3.1 Offline Preparation

  • Offline preparation builds the gold-shot prompt, sets a threshold to cover at least 90% of calibration-set errors, and constructs a vector knowledge base for retrieval.The threshold is intended to transfer to unseen test sets of the same type under the stated assumption of relatively stable within-domain uncertainty distributions.Source: §4.3.1 · pp. 10–11

4.3.2 Online Inference

  • Online inference predicts with the static prompt, returns that prediction when its LSFU is at or below T, and otherwise retrieves labeled references for reflective correction.The reflective prompt includes an initial-prediction notice, retrieved exemplars, and an explicit correction instruction before a second inference.Source: §4.3.2 · p. 11

5 Theoretical Analysis and Validation

LSFU is evaluated as a way to rank prediction risk and support selective prediction, while the gating-threshold analysis selects a 90% error-coverage target as a cost-benefit trade-off. The reported findings include separation between correct and incorrect prediction distributions and rising accuracy as high-LSFU samples are removed.

5.1 Effectiveness of the LSFU Metric

  • The prior modulation factor weights risk nonlinearly, suppressing pseudo-uncertainty for high-prior classes while amplifying risk for low-prior classes.The analysis describes the low-prior weighting as highlighting long-tail cases where entropy may appear normal despite insufficient evidence.Source: §5.1 · p. 12
  • Correct and incorrect predictions show separated LSFU density centers, and removing high-LSFU samples raises accuracy as coverage decreases.The density distributions overlap; the risk-coverage trend is reported for CASIE and ACE.Source: §5.1 · p. 13, Figure 3 · p. 13, Figure 4 · p. 14
  • At 10% coverage on CASIE, accuracy reaches 88.6%, while eliminating 20% of high-risk ACE samples raises accuracy from about 51% to 60.9%.The reported examples are presented as a safety boundary for the static prompt.Source: §5.1 · p. 14

5.2 Strategy for Selecting the Gating Threshold T

  • 90% error coverage is selected at the elbow of the error-coverage versus RAG-trigger-rate curve as a trade-off between efficiency and accuracy.The passage reports steep gains below 85% coverage and flattening above 95%, where capturing the final 5% of overconfidence errors requires an aggressive threshold.Source: §5.2 · p. 14, Figure 5 · p. 15

6 Experiments

Experiments assess static exemplar selection, dynamic prompt correction, and cost-effectiveness across six datasets spanning different domains and distributional characteristics. UCPOF improves average accuracy over the static baseline and reduces retrieval relative to Full RAG.

6.1 Experimental Setup

  • The setup covers six datasets and three models, using accuracy, AUC, and average inference cost to evaluate prompt optimization.Qwen2.5-7B-Instruct is the main model; Llama-3-8B-Instruct and Qwen2.5-14B-Instruct assess cross-model generalization and scalability.Source: §6 · p. 15, §6.1 · p. 15

6.2 Main Results

  • 6.03% average accuracy increase over the static Baseline; UCPOF improves accuracy across all datasets.The largest gain is +11.01% on ACE.Source: Table 2 · p. 16, §6.2 · p. 16

6.3.1 Gold Shot Selection Strategy

  • Lowest-LSFU exemplars yield the best average accuracy (73.81%), while removing Gold Shot lowers UCPOF average accuracy by 2.20% (78.09% →75.89%).Using the highest-LSFU, hardest samples as exemplars reduces accuracy on ACE, CB, and R8.Source: Table 3 · p. 17, §6.3.1 · p. 17, Table 4 · p. 17

6.3.2 Necessity of the Prior Probability

  • Adding Pprior improves AUC consistently across datasets and lowers ECE and Brier score, with the largest calibration improvement on CB.The AUC comparison is between LSFU with prior and ordinary First Token Entropy without prior.Source: §6.3.2 · p. 17, Table 5 · p. 18, Table 7 · p. 19

6.4 Efficiency Analysis

  • 5.75 points higher overall average accuracy than Full RAG (78.09 vs. 72.34), while reducing average retrieval trigger rate by 50.66% (49.34% vs. 100%).Full RAG falls to 48.21% on CB versus 78.57% for Baseline, while on AgNews Full RAG scores 91.42% versus 89.99% for UCPOF.Source: §6.2 · p. 16, §6.4 · p. 18
  • UCPOF lies near the upper-left of the efficiency–accuracy Pareto frontier and stops before the high-cost region where forced retrieval can reduce accuracy.Accuracy saturates at roughly a 60% trigger rate; the tuned threshold yields a ∼77% test trigger rate.Source: §6.4 · p. 18, Figure 6 · p. 19

7 Discussion and Conclusion

The paper presents LSFU and UCPOF for prompt optimization, balancing accuracy and inference efficiency through static prompt optimization and dynamic correction. It also reports limits in task scope, prompt construction, and threshold generalization.

7 Discussion and Conclusion

  • The paper proposes prior-aware Log-Scale Focal Uncertainty (LSFU) and UCPOF to address prompt sensitivity, distorted uncertainty measurement, and full-RAG noise.The conclusions describe LSFU as a basis for prompt optimization and UCPOF as a two-stage static-and-dynamic framework.Source: §7 · p. 19
  • 78.09% average accuracy (6.03% higher than Baseline) with lower inference cost than full RAG; UCPOF pairs static optimization with LSFU-gated dynamic correction.The efficiency comparison covers retrieval trigger rate and test accuracy.Source: Table 8 · p. 20, §7 · p. 20
  • 1.75% higher average accuracy than random selection came from using low-LSFU, high-confidence samples as few-shot exemplars.The paper describes these exemplars as providing decision anchors and reducing cognitive noise and hallucination-propagation risks.Source: §7 · p. 20
  • The framework is limited to classification, does not fully optimize exemplar arrangement or instruction details, and relies on domain-distribution stability for dynamic-threshold generalization.The paper identifies broader task adaptation and prompt-structure optimization as future work.Source: §7 · pp. 20–21

A Efficiency Analysis on Other Datasets

The appendix reports efficiency–accuracy Pareto curves and retrieval-trigger-rate versus accuracy curves for datasets other than ACE, using the protocol from Section 6.4.

A Efficiency Analysis on Other Datasets

  • The appendix extends the Section 6.4 efficiency analysis to datasets other than ACE, reporting Pareto and retrieval-trigger-rate versus accuracy curves.The passage specifies the protocol but does not provide curve outcomes.Source: A Efficiency Analysis on Other Datasets · p. 23

B Cross-Model Scalability

Cross-model experiments report UCPOF advantages over static baselines, Gold Shot, and full RAG, including on larger Qwen and lower-parameter Llama models.

B Cross-Model Scalability

  • Experiments tested whether LSFU depended on a specific architecture using Llama-3 and Qwen models of different sizes.Efficiency analyses are also shown for CASIE and AgNews.Source: B Cross-Model Scalability · p. 23, Figure 8 · p. 24, Figure 9 · p. 24
  • 2.17%–3.56% average accuracy improvement was reported for UCPOF over static baseline, static optimization (Gold Shot), and full RAG across cross-model scenarios.The authors state that LSFU's prior-calibration logic does not depend on a specific model pretraining paradigm.Source: B Cross-Model Scalability · p. 23
  • 80.38% average accuracy for Qwen2.5-14B exceeded full RAG at 80.10%.The passage says UCPOF maintained performance advantages as model parameters increased.Source: B Cross-Model Scalability · p. 23
  • 83.93% on the CB task for Llama-3-8B exceeded full RAG at 82.14%.The authors identify CB as semantically sensitive and describe the improvement as most significant on this low-parameter model.Source: B Cross-Model Scalability · p. 23

C Prior-Ablation Results on Other Models

The appendix reports prior-ablation results on additional model backbones.

C Prior-Ablation Results on Other Models

D Additional Experiments on Prompt Structure Validation

Supplementary experiments provide data across models and parameters using the “Task→Example→Input” structure to support Section 5.2 and strengthen prompt-optimization reliability.

D Additional Experiments on Prompt Structure Validation

E Validating the Top-K Parameter in LSFU

The top-K entropy term is stable across the tested K values: K = 50 captures the relevant high-probability mass, while larger K yields negligible benefit at additional cost.

E Validating the Top-K Parameter in LSFU

Loading 2603.18009v1…