Trace method improves confidence scores in language model answers
TRACE: Single-Pass Decoding-Trace Risk Localization for Generation Calibration
Computation and Language
Summary
Sometimes, AI language models give answers that sound right but have small mistakes like wrong numbers or facts. Previous ways to guess if an answer is reliable often miss these little errors because they look only at the overall answer. The authors created TRACE, which carefully watches how uncertain the AI is while it creates each word, spotting risky parts in the answer. This helps produce better confidence scores about the whole answer being correct, making the AI safer and more trustworthy.
What this means in practice
- •For ai product teams: Improve reliability indicators for AI-generated answers by detecting local risks during text creation without extra computations or external checks.
- •For enterprise chatbot developers: Enhance chatbot answer trustworthiness by better spotting when specific facts or numbers may be wrong, improving user confidence in automated support.
Authors
Yuebin Xu, Xuemei Peng, Junlan Chen, Zhiyi Chen, Zeyi Wen
Abstract
Reliable confidence estimation is essential for large language model deployment. However, answer-level calibration remains challenging because generation errors are often localized: a response may be fluent and high-probability overall while still failing at a critical number, entity, or factual claim. Existing estimators compress token probabilities, sequence likelihoods, entropy, or beam statistics into a global score, which can dilute such local risk signals. We propose TRACE, a single-pass, decoded-answer-preserving confidence estimator that treats decoding-time uncertainty as a trajectory through three steps: (i) recording token-level surprisal and predictive entropy during decoding, (ii) applying local risk operators to preserve uncertainty spikes, and (iii) converting localized trace risk into answer-level confidence. TRACE produces a label-free risk score, while TRACE+ calibrates trace-only features into probabilities using a held-out split, without extra generations or external verifiers. We evaluate four tasks against 19 calibration baselines, and TRACE+ reduces Brier from 0.149 to 0.137 and improves AUROC from 0.758 to 0.792 over the strongest likelihood baseline. Across seven LLMs, TRACE+ improves over the best non-TRACE baseline pool from 0.136 to 0.120 Brier and from 0.764 to 0.817 AUROC. Results show that localizing decoding-time risk provides a general approach to calibration.