Token-Level Density-Based Uncertainty Quantification Methods for Eliciting Truthfulness of Large Language Models

2025-02-20

AI Summary

This study addresses the challenge of uncertainty quantification (UQ) for text generated by large language models (LLMs). We propose a supervised, token-level uncertainty scoring method grounded in embedding density estimation. Our core innovation is the first adaptation of Mahalanobis distance to generative LLMs, leveraging multi-layer Transformer token embeddings to train a lightweight, interpretable linear regression model. Unlike conventional density-based UQ approaches—which fail under generative settings due to distributional shift—our method overcomes this limitation and unifies support for both sequence-level selective generation and statement-level factual verification. Evaluated across 11 diverse, cross-domain datasets, our approach consistently outperforms existing UQ baselines, achieving superior accuracy, minimal computational overhead, and strong generalization across architectures and domains.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: (Large) Language Models
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Uncertainty quantification (UQ) is a prominent approach for eliciting truthful answers from large language models (LLMs). To date, information-based and consistency-based UQ have been the dominant UQ methods for text generation via LLMs. Density-based methods, despite being very effective for UQ in text classification with encoder-based models, have not been very successful with generative LLMs. In this work, we adapt Mahalanobis Distance (MD) - a well-established UQ technique in classification tasks - for text generation and introduce a new supervised UQ method. Our method extracts token embeddings from multiple layers of LLMs, computes MD scores for each token, and uses linear regression trained on these features to provide robust uncertainty scores. Through extensive experiments on eleven datasets, we demonstrate that our approach substantially improves over existing UQ methods, providing accurate and computationally efficient uncertainty scores for both sequence-level selective generation and claim-level fact-checking tasks. Our method also exhibits strong generalization to out-of-domain data, making it suitable for a wide range of LLM-based applications.

Problem

• Quantify uncertainty in large language models
• Adapt Mahalanobis Distance for text generation
• Improve truthfulness in sequence-level selective generation

Innovation

• Token-level Mahalanobis Distance adaptation
• Multi-layer token embeddings extraction
• Linear regression for uncertainty scores