GrACE: A Generative Approach to Better Confidence Elicitation and Efficient Test-Time Scaling in Large Language Models

1. Problem Statement

Large Language Models (LLMs) are increasingly deployed in high-stakes domains (e.g., healthcare, finance), necessitating reliable confidence estimation to ensure safety and trust. Current methods for eliciting confidence suffer from significant limitations:

• Post-generation Methods: Approaches like $P(True)$ or using auxiliary models require an additional evaluation stage after the LLM generates a response. This introduces high computational latency and degrades real-time interaction.
• Verbalized Confidence: Methods where the model outputs natural language phrases (e.g., "likely," "not sure") or numerical strings during generation are often poorly calibrated. LLMs tend to be overconfident in verbal expressions, and discrete vocabulary fails to capture nuanced confidence intervals.
• Test-Time Scaling (TTS) Inefficiency: Current TTS strategies (e.g., Self-Consistency) often rely on fixed sampling budgets or pre-sampling to determine stopping criteria, leading to inefficient resource usage. They lack a reliable, real-time metric to decide when to stop sampling or how to weight samples.

The core challenge is to develop a method that provides discriminative (distinguishes correct from incorrect answers) and calibrated (matches actual accuracy) confidence scores on-the-fly during generation, without auxiliary models or extra inference steps.

2. Methodology: GrACE Framework

The authors propose GrACE (Generative Approach to Confidence Elicitation), a framework that integrates confidence estimation directly into the generation process.

A. Generative Confidence Mechanism

Instead of generating text or using a separate model, GrACE introduces a special token, <CNF>, appended to the vocabulary.

• Process: The model is fine-tuned to generate the <CNF> token at the end of its response.
• Confidence Calculation: The confidence score ($c$) is derived in real-time by calculating the similarity between the model's last hidden state ($z_L$) and the embedding of the <CNF> token ($e_{<CNF>}$).
• Formula: $c = \text{sim}(z_L, e_{<CNF>})$. In practice, this is implemented using the softmax probability of the <CNF> token in the vocabulary projection, leveraging the model's native geometry.

B. Calibration Target Construction

To ensure the generated confidence is well-calibrated, GrACE does not use binary labels (correct/incorrect) directly, which leads to poor calibration. Instead, it uses a grouping strategy:

1. Probing: A linear probe is trained on a subset of data to predict an uncalibrated confidence score ($s$) based on the model's hidden states.
2. K-Fold Binning: The training data is partitioned into $k$ folds. The probe is trained on $k-1$ folds and used to score the remaining fold.
3. Grouping & Target Assignment: Data points are grouped into bins based on their uncalibrated scores ($s$). The empirical accuracy of each bin is calculated and assigned as the calibration target ($t$) for all samples in that bin.
4. Result: This creates a "sharp" distribution of targets that aligns predicted confidence with actual group accuracy.

C. Training Process

The model is fine-tuned using a combined loss function:

• Calibration Loss ($\mathcal{L}_C$): Mean Squared Error (MSE) between the generated confidence (similarity to <CNF>) and the calibration target ($t$).
• Supervised Fine-Tuning Loss ($\mathcal{L}_{SFT}$): Standard language modeling loss on the original response text to preserve the model's utility. The <CNF> token is excluded from this loss to prevent conflict.
• Efficiency: Only the LoRA (Low-Rank Adaptation) parameters and the embedding of <CNF> are fine-tuned, making the process parameter-efficient.

D. Application to Test-Time Scaling (TTS)

GrACE enables two new strategies for TTS:

1. GrACE-SC (Weighted Self-Consistency): Instead of a simple majority vote, final answers are selected by weighting samples based on their GrACE confidence scores.
2. GrACE-ES (Early Stopping): Sampling continues until a sample is generated with a confidence score exceeding a threshold $\tau$. If the budget is exhausted without meeting the threshold, the weighted vote is used. This drastically reduces the number of samples required.

3. Key Contributions

1. Novel Mechanism: Introduced a generative, on-the-fly confidence elicitation method using a special token and hidden-state similarity, eliminating the need for post-generation evaluation.
2. Calibration Strategy: Developed a binning-based target construction method that aligns confidence with empirical accuracy, addressing the overconfidence issues of verbalized methods.
3. Efficient TTS Integration: Proposed confidence-based aggregation and early-stopping strategies that improve decision accuracy while significantly reducing computational costs.
4. State-of-the-Art Performance: Demonstrated superior performance across multiple backbone models (Phi-3, Llama-2, Llama-3) and datasets (TriviaQA, SciQ, MathQA, ARC).

4. Experimental Results

• Open-Ended Generation: GrACE outperformed six baselines (including $P(True)$, Verbal, Apricot, and ActCab).
  • Calibration: Achieved the lowest Expected Calibration Error (ECE) and Brier scores. For example, on Llama-3.1-8B, GrACE reduced ECE by 2.6% compared to the next best method (Apricot).
  • Discrimination: Achieved the highest AUROC, indicating strong ability to distinguish correct from incorrect answers.
  • Generalization: GrACE showed robust performance in Out-of-Domain (OOD) settings, with minimal performance degradation compared to methods that rely heavily on semantic shortcuts.
• Test-Time Scaling:
  • Accuracy: GrACE-SC improved accuracy over standard Self-Consistency by 1.54% on average.
  • Efficiency: GrACE-ES significantly reduced the average number of samples needed. On the ARC_C dataset, 80.1% of queries required only a single sample to meet the confidence threshold, compared to a fixed budget of 8 for standard methods.

5. Significance

GrACE represents a paradigm shift in LLM reliability. By moving confidence estimation from a post-hoc or verbalized process to a native, continuous, and generative mechanism, it solves the trade-off between computational efficiency and calibration quality.

• Practical Deployment: It enables real-time, reliable confidence estimation without auxiliary models, making it feasible for high-stakes applications where latency and trust are critical.
• Resource Optimization: The ability to dynamically stop sampling based on reliable confidence scores offers a practical solution for optimizing inference costs in Test-Time Scaling, a key area of current LLM research.
• Scalability: The method is lightweight (fine-tuning only LoRA and one token embedding) and generalizes well across different model architectures and domains.