Do Not Design, Learn: A Trainable Scoring Function for Uncertainty Estimation in Generative LLMs

TL;DR

This work tackles the challenge of uncertainty estimation in generative LLMs by showing that existing probability-based scoring functions (e.g., LNS, MARS, TokenSAR, SE) have limitations in handling token interactions, biased probabilities, and multilingual applicability. It introduces Learnable Response Scoring (LARS), a trainable, RoBERTa-based scorer that consumes the prompt, generated tokens, and per-token probabilities encoded via a few-hot scheme, and is trained with calibrated supervision to produce well-calibrated UE scores. Across multiple QA and math-reasoning datasets, models, and languages, LARS yields substantial gains in AUROC and PRR over all baselines, and demonstrates robustness in OOD data and cross-domain settings, with data-efficient calibration (as few as 1k labeled examples). The approach offers a practical, scalable path to reliable UE in closed- and open-weight LLMs, enabling safer and more trustworthy AI-assisted reasoning in diverse real-world contexts.

Abstract

Uncertainty estimation (UE) of generative large language models (LLMs) is crucial for evaluating the reliability of generated sequences. A significant subset of UE methods utilize token probabilities to assess uncertainty, aggregating multiple token probabilities into a single UE score using a scoring function. Existing scoring functions for probability-based UE, such as length-normalized scoring and semantic contribution-based weighting, are designed to solve certain aspects of the problem but exhibit limitations, including the inability to handle biased probabilities and complex semantic dependencies between tokens. To address these issues, in this work, we propose Learnable Response Scoring (LARS) function, a novel scoring function that leverages supervised data to capture complex dependencies between tokens and probabilities, thereby producing more reliable and calibrated response scores in computing the uncertainty of LLM generations. Our comprehensive experiments across question-answering and arithmetical reasoning tasks with various datasets demonstrate that LARS significantly outperforms existing scoring functions, achieving improvements of up to 16\% AUROC score.

Figures (7)

• Figure 1: (Left) Answer generation using a generative LLM. (Mid Left) Overview of the proposed scoring function LARS. It utilizes the question, answer tokens, and token probabilities. Token probabilities are fed to LARS model as special probability tokens. (Mid Right) Illustration of few-hot represented embedding vectors of probability tokens. (Right) Overview of probability-based UE methods taking different sampled answer scores such as LNS malinin2021uncertainty, MARS bakman2024mars, or LARS (this work), and outputting a single UE value.
• Figure 2: Average accuracy and probability assignments of Llama2-7b-chat for specific entities in TriviaQA.
• Figure 3: Comparison of different probability association methods for LARS on TriviaQA (top) and NaturalQA (bottom).
• Figure 4: AUROC scores of LARS for different amount of questions in calibration data on TriviaQA. For each UE method and model, the best score across baseline scoring functions is provided as a reference.
• Figure 5: Comparison of different probability association methods for LARS on TriviaQA (top) and NaturalQA (bottom).