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

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.

## introduction large language models (llms) have achieved impressive results over various tasks and applications @xcite @xcite . nevertheless, even the most advanced llms are inevitably prone to making mistakes during text generation. their responses often contain hallucinations or non-factual claims @xcite , posing significant challenges for llm deployment in safetycritical domains. many studies have investigated methods for assessing the truthfulness of llm responses @xcite @xcite . however, many of the proposed techniques have limited practical applicability, as they often rely on external knowledge sources or require ensembling multiple large llms, leading to high computational costs that make them economically unfeasible for many use cases. one of the most promising approaches to addressing this challenge is uncertainty quantification (uq) @xcite @xcite @xcite . this research direction recognizes that we will never have complete information about model predictions due to the limited amount of training data and ambiguity of the tasks, and suggests general ways to estimate the reliability of predictions under different conditions. recently, a suite of uq methods specifically designed for text generation with llms has been developed @xcite @xcite @xcite . however, many of these methods are either ineffective or come with a substantial computational overhead, limiting their practicality for largescale or real-time applications. for text classification and regression tasks, researchers have identified several groups of techniques that maintain a balance between effectiveness and computational efficiency @xcite @xcite @xcite . one such class of approach is so-called densitybased uncertainty scores @xcite @xcite . these methods use embeddings of instances obtained from the top layers of a classification model to fit the density of the training distribution in the latent space. the likelihood of the input data under this estimated distribution is then used for confidence estimation. this has been demonstrated to achieve excellent results in out-ofdistribution detection tasks @xcite , token embeddings mahalanobis distance t tokens average over tokens pca n layers sequence probability input text decoder layer embeddings nx linear layer linear regressor uncertainty transformer output text figure 1 : an illustration of the proposed method. after each decoder layer, the embeddings of each generated token are extracted. subsequently, we compute the mahalanobis distance for each token and layer and then average over all tokens in the generated sequence. finally, we train a linear regression on the pca decomposition of the calculated features with the addition of sequence probability to predict the uncertainty of the generation. and proven to be useful for selective text classification @xcite . despite being computationally lightweight, these techniques often outperform more resource-intensive methods, such as deep ensembles @xcite and monte carlo dropout @xcite . unfortunately, the reported performance of density-based scores for text generation so far has been notably low @xcite . recent work has demonstrated that the internal states of llms carry a lot of information about their uncertainty @xcite @xcite @xcite . these techniques train a supplementary model on top of the activations of llm internal layers. however, they often rely on simplistic features and fail to incorporate more advanced, well-established density-based uq methods, limiting their ability to capture uncertainty. in this work, we address this gap by adapting density-based techniques for the uq of llms and propose a new supervised method based on densitybased features. specifically, we adapt one of the most robust methods for uq in the classification tasks, namely mahalanobis distance (md; @xcite ), and train a linear regression on top of the md scores from various layers of the llm. these features are supplemented with a probability of the generated sequence. figure 1 illustrates the scheme of the proposed supervised uq method. our extensive experimental evaluation demonstrates that the proposed method provides substantial improvement over the state of the art. our key contributions are as follows. ## related work many effective uq methods, such as deep ensembles @xcite and monte carlo (mc) dropout @xcite , require sampling multiple predictions from a model, which leads to substantial computational and memory overheads. a key challenge in uq is developing techniques that balance effectiveness with computational efficiency. among the most promising approaches in this regard are densitybased methods @xcite @xcite @xcite . these methods leverage latent representations of instances to model the training data distribution, then estimate how likely a new instance belongs to that distribution. @xcite propose to use mahalanobis distance (md) as a measure of uncertainty for out-of-distribution detection in computer vision tasks. podolskiy et al. for llms, @xcite and @xcite proposed uq methods that sample multiple predictions and leverage their diversity. in the context of black-box llms, where we have no access to the logits or embeddings of a model, @xcite propose the use of lexical dissimilarity of sampled texts as a measure of uncertainty. @xcite leverage a similarity matrix between responses for deeper analysis of the diversity of the sampled generations. some methods also combine sampling diversity measures with the probability of each generation @xcite @xcite @xcite . recently, it was demonstrated that md is an efficient approach for out-of-distribution detection in sequence-to-sequence models @xcite @xcite . however, for selective generation tasks, densitybased methods so far have substantially underperformed compared to trivial baselines @xcite . supervised methods are another research direction for uq of llms. @xcite demonstrate that the internal states of the model contain information about uncertainty, and propose to train a multi-layer perceptron over the hidden llm representation to predict the truthfulness of the model responses. @xcite enhance this idea by training a deep neural network with recurrent and convolutional layers. furthermore, this method uses embeddings from all layers and incorporates features based on the probability and the dynamics of the generated tokens through layers. in contrast to these methods, we employ a simple linear model, but focus on more accurate feature extraction from internal layers, using well-established density-based uq methods. ## background on density-based methods recently, mahalanobis distance (md) and robust density estimation (rde) were adapted @xcite to the text generation task by considering the marginal distribution of the training dataset. following the assumption of a gaussian distribution of training instance representations, the md method @xcite calculates a centroid of the training data µ and the empirical covariance matrix σ. for a given instance x, the uncertainty score is defined as the mahalanobis distance: rde @xcite operates within the reduced dimensionality of h l (x) via the kernel pca decomposition. to ensure the robustness of the covariance matrix, it uses the minimum covariance determinant estimate @xcite . finally, the uncertainty score is computed as the mahalanobis distance in the reduced dimensionality. @xcite proposed a modification of md -relative mahalanobis distance (rmd). it takes into account a background contrastive md score. the score aims to assess how close the test instance is to the in-domain training data compared to the background data. the uncertainty score based on rmd is given by the following equation: for the sequence-to-sequence tasks, it was proposed to use the last encoder and decoder layer for extracting hidden representation of the model @xcite . in contrast, recent research @xcite indicates that the middle layers of the model may be more suitable for decoder-only models. llm responses, each accompanied by an assessment of its correctness. the assessment can be based on ground truth answers (as in tasks like question-answering, machine translation, or summarization) or through alternative means, such as human annotation or another big llm. ## conclusion we have introduced a series of new supervised uq methods based on layer-wise features derived from the mahalanobis distance. we show that calculating md over token-level embeddings yields much better results than previous attempts that leverage sequence-level embeddings. training a linear regression on top of the layer-wise scores allows us to produce even better uncertainty scores and outperform the state-of-the-art supervised and unsupervised uq methods in selective classification across eleven datasets and in claim-level fact-checking. we also show that the proposed methods are computationally efficient and have the potential for generalization, which makes them useful in real-world llm-based applications. in future work, we aim to improve the generalization capabilities of the supervised uq methods on out-of-domain data by investigating new features and a more robust training pipeline.