💬 NLP

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

This paper introduces a novel supervised uncertainty quantification method for large language models that adapts Mahalanobis Distance to token-level embeddings across multiple layers, demonstrating superior accuracy and generalization over existing approaches for truthfulness elicitation in both sequence-level generation and fact-checking tasks.

Original authors: Artem Vazhentsev, Lyudmila Rvanova, Ivan Lazichny, Alexander Panchenko, Maxim Panov, Timothy Baldwin, Artem Shelmanov

Published 2026-04-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Artem Vazhentsev, Lyudmila Rvanova, Ivan Lazichny, Alexander Panchenko, Maxim Panov, Timothy Baldwin, Artem Shelmanov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, well-read robot (a Large Language Model, or LLM) that can write stories, answer questions, and summarize news. The problem is, this robot sometimes "hallucinates"—it confidently makes up facts or tells lies, just like a student who guesses the answer on a test but is actually wrong.

The authors of this paper wanted to build a "lie detector" for this robot. They wanted a way to know, before we trust the robot's answer, whether the robot is actually unsure or if it's just bluffing.

Here is how they did it, explained simply:

The Problem with Old Methods

Previously, scientists tried two main ways to catch the robot lying:

The "Ask Again" Method: Ask the robot the same question 10 times. If it gives 10 different answers, it's probably confused. Downside: This is slow and expensive, like asking a friend the same question 10 times just to see if they change their story.
The "Confidence Score" Method: Look at how confident the robot sounds. Downside: Robots are great at sounding confident even when they are completely wrong.

The New Idea: The "Token-Level Density" Detective

The authors realized that the robot's "brain" (its internal layers) holds clues about its uncertainty. They decided to use a mathematical tool called Mahalanobis Distance (MD).

The Analogy: The "Normal Crowd" vs. The "Outlier"
Imagine the robot has a mental library of "good, correct answers" it learned during training.

The Old Way (Sequence Level): When the robot finishes a whole sentence, the old method looked at the entire sentence as one big blob and asked, "Does this whole blob look like the 'good' answers?" The problem was that a sentence can have one weird word that ruins the whole thing, but the "blob" might still look okay.
The New Way (Token Level): The authors decided to look at every single word (token) the robot writes, one by one.
- They ask: "Does this specific word look like the words found in the robot's 'good' library?"
- If the robot writes a word that is very strange or far away from its library of correct words, that word gets a high "suspicion score."
- They do this for every word in the sentence, then average the scores to get a final "uncertainty score" for the whole answer.

How They Built It (The Recipe)

Gather the "Good" Examples: They took a bunch of questions and the robot's answers. They filtered these answers to keep only the ones that were definitely correct (like a teacher grading a test and keeping only the A+ papers).
Map the "Good" Words: They looked at the internal math (embeddings) of the words in those A+ papers to create a map of what "correct" looks like for every layer of the robot's brain.
The "Suspicion" Test: When the robot generates a new answer, the system checks every single word against that "correct" map.
- If a word is far away from the "correct" map, it gets a high distance score (High Uncertainty).
- If it's close, it gets a low score (Low Uncertainty).
The "Coach" (Linear Regression): Since looking at every single layer of the robot's brain is complicated, they trained a simple "coach" (a linear regression model). This coach looks at all the suspicion scores from the different layers and the robot's own confidence, then says, "Based on all these clues, how likely is this answer to be true?"

What They Found

They tested this new "Token-Level Detective" on 11 different tasks, from summarizing news articles to answering tricky science questions.

It's Faster: Unlike the "Ask Again" method, this doesn't require the robot to run multiple times. It's almost as fast as the robot just writing the answer normally.
It's Smarter: It caught lies much better than previous methods. In many tests, it was the best method available.
It Works on New Stuff: Even when they tested the robot on topics it hadn't seen during training (like switching from summarizing news to answering medical questions), the method still worked well.
The "Hybrid" Boost: They found that combining their new method with the robot's own confidence score (a "Hybrid" approach) made it even better at spotting lies.

The Bottom Line

The paper claims that by looking at every single word the robot writes and checking how "strange" it is compared to the robot's training, we can create a very fast, accurate, and cheap way to tell if the robot is telling the truth or making things up. It's like giving the robot a mirror so it can see when it's about to say something weird, before it even says it.

1. Problem Statement

Large Language Models (LLMs) are prone to generating hallucinations and non-factual claims, posing significant risks in safety-critical applications. While Uncertainty Quantification (UQ) is a primary approach to detecting these errors, existing methods face two major limitations:

Ineffectiveness of Density-Based Methods: Traditional density-based UQ techniques (e.g., Mahalanobis Distance) have been highly successful in text classification but have failed to perform well in text generation tasks when applied at the sequence level. Previous attempts yielded performance comparable to random selection.
Computational Cost: State-of-the-art UQ methods for LLMs often rely on consistency-based approaches (e.g., Semantic Entropy, sampling multiple generations) or deep ensembles, which introduce massive computational overhead (300–700% increase in inference time), making them impractical for real-time or large-scale applications.
Feature Limitations: Recent supervised methods that utilize internal LLM states often rely on simplistic features or fail to incorporate robust, established density-based metrics effectively.

2. Methodology

The authors propose a new supervised, token-level, density-based UQ framework that adapts the Mahalanobis Distance (MD) for generative models. The core innovation lies in shifting from sequence-level embeddings to token-level embeddings and aggregating them intelligently.

Key Components:

Token-Level Embedding Extraction:
- Instead of averaging embeddings to form a single sequence vector, the method extracts hidden state embeddings for every generated token from multiple layers of the LLM.
- Embedding Selection: To construct the reference distribution (centroid and covariance matrix), the method filters training data to include only high-quality (correct) responses. This ensures the density model represents "truthful" generations.
Layer-Wise Mahalanobis Distance (MD):
- For each layer $l$ , the method computes the Mahalanobis distance between a token's embedding and the centroid of the selected "correct" token embeddings from that layer.
- Relative MD (RMD): The method also implements Relative Mahalanobis Distance, which subtracts the distance to a background distribution (e.g., general C4 dataset) to better distinguish in-domain truthfulness from general noise.
- Aggregation: The token-level scores are averaged across the sequence to produce a sequence-level score (Average Token-level MD or ATMD / ATRMD).
Supervised Linear Regression:
- Since the best-performing layer varies by task, the method treats the layer-wise MD scores as features.
- A Linear Regression model is trained to predict the uncertainty score (targeted as the negative quality metric, e.g., $-ROUGE$ or $-Accuracy$) using these layer-wise scores.
- PCA: To handle multicollinearity among layer features, the top 10 Principal Components (PCA) are used as input to the regressor.
- Hybrid Score (HUQ): The final uncertainty score can be a hybrid of the supervised MD score and the sequence probability (Maximum Sequence Probability, MSP), combined via a learned weighting function.

3. Key Contributions

Token-Level Adaptation: The paper demonstrates that applying density-based methods at the token level (rather than sequence level) is crucial for success in LLM generation, overcoming the failure of previous sequence-level density approaches.
Efficient Supervised Framework: They introduce a computationally lightweight supervised method that combines layer-wise density scores with linear regression, avoiding the need for expensive sampling or ensembling.
Comprehensive Evaluation: The method was validated across 11 diverse datasets (covering summarization, QA with long/short answers, and multiple-choice QA) and two tasks: sequence-level selective generation and claim-level fact-checking.
Generalization: The method shows strong robustness in out-of-domain settings, maintaining performance even when trained on different task types (e.g., training on summarization/QA and testing on multiple-choice QA).

4. Experimental Results

The experiments utilized Llama 8b v3.1, Gemma 9b v2, and Mistral 7b models.

Selective Generation Performance:
- The proposed methods (SATRMD+MSP and HUQ-SATRMD) significantly outperformed all baselines, including state-of-the-art unsupervised methods (Semantic Entropy, SAR) and other supervised methods (SAPLMA, Factoscope).
- Metrics: Measured by Prediction Rejection Ratio (PRR). The proposed methods achieved the highest mean ranks across datasets. For example, on the MMLU dataset, SATRMD+MSP significantly outperformed the next best method.
- Layer Analysis: Contrary to classification tasks where the last layer is often best, the results showed that middle layers often provide the most informative embeddings for uncertainty in decoder-only LLMs.
Fact-Checking Performance:
- On claim-level fact-checking (Mistral 7b), HUQ-SATRMD achieved the best ROC-AUC (0.750), outperforming the baseline CCP by 3.4%.
Computational Efficiency:
- Overhead: The proposed methods introduce a minimal computational overhead of 5.3% – 7.6% compared to standard LLM inference.
- Comparison: This is drastically lower than sampling-based methods like Semantic Entropy or SAR, which incur 315% – 700% overhead.
Robustness:
- The method is robust to the size of the training dataset (performing well with as few as 200–500 instances).
- It is not highly sensitive to the correctness threshold used for filtering training embeddings.

5. Significance

This work bridges the gap between effective density-based UQ in classification and the challenging domain of LLM text generation. By shifting the granularity to the token level and utilizing a supervised linear combination of layer-wise features, the authors provide a solution that is:

More Accurate: Significantly outperforms current state-of-the-art UQ methods in detecting hallucinations.
More Efficient: Makes real-time uncertainty estimation feasible for production environments by avoiding expensive sampling.
Generalizable: Demonstrates the ability to transfer uncertainty estimation capabilities across different domains and task types, a critical requirement for deploying LLMs in diverse real-world applications.

The paper concludes that density-based methods, when properly adapted to the token-level architecture of LLMs, are a superior alternative to both trivial probability-based baselines and computationally expensive sampling methods.