Robust Hallucination Detection in LLMs via Adaptive Token Selection

TL;DR

This work tackles hallucination detection in LLM outputs by introducing HaMI, a MIL-based framework that adaptively selects token-level cues most indicative of hallucinations and jointly optimizes their selection with a dedicated detector. It further enhances internal representations by incorporating predictive uncertainty, notably semantic entropy, to boost discriminability. Across four QA benchmarks and multiple model scales, HaMI (and HaMI*) consistently outperforms strong uncertainty- and representation-based baselines, with substantial gains on larger models. The results demonstrate robustness to generation length, strong cross-dataset generalization, and clear ablation-driven insights into token selection, uncertainty fusion, and layer-wise representation effects, suggesting broad applicability beyond QA tasks.

Abstract

Hallucinations in large language models (LLMs) pose significant safety concerns that impede their broader deployment. Recent research in hallucination detection has demonstrated that LLMs' internal representations contain truthfulness hints, which can be harnessed for detector training. However, the performance of these detectors is heavily dependent on the internal representations of predetermined tokens, fluctuating considerably when working on free-form generations with varying lengths and sparse distributions of hallucinated entities. To address this, we propose HaMI, a novel approach that enables robust detection of hallucinations through adaptive selection and learning of critical tokens that are most indicative of hallucinations. We achieve this robustness by an innovative formulation of the Hallucination detection task as Multiple Instance (HaMI) learning over token-level representations within a sequence, thereby facilitating a joint optimisation of token selection and hallucination detection on generation sequences of diverse forms. Comprehensive experimental results on four hallucination benchmarks show that HaMI significantly outperforms existing state-of-the-art approaches.

Figures (8)

• Figure 1: Tokens that contain the most sufficient information related to correctness may appear at various positions within the sequence.
• Figure 2: The framework of our proposed HaMI. The LLM is prompted to generate answer tokens accompanied by token representations $h_i$. The network receives sequences of token representations for both positive $\mathcal{B^+}$ and negative $\mathcal{B^-}$ bags as inputs at the training stage. The hallucination detector assigns a hallucination score to each token instance. We choose the top $k$ largest scores from both bags, subsequently maximising the discriminative margin between them by minimising a MIL loss as described in Eq. \ref{['eq: mil']}. Given the sequential nature of the generations, a constraint on the smoothness of hallucination scores of adjacent tokens is also added to HaMI.
• Figure 3: (a) AUROC results of cross-dataset generalisation on four datasets using LLaMA-3.1-8B. (b) AUROC w.r.t. dimensionality of the feature layer (left) and the number of network layers (right) based on LLaMA-3.1-8B.
• Figure 4: Adaptive token selection results showing tokens with the top-2 highest hallucination scores.
• Figure 5: (a) Performance of HaMI w.r.t. internal representations extracted from different layers. (b) The impact of augmentation strength $\lambda$ in Eq. \ref{['eq: enhancement']} on HaMI. All results are based on LLaMA-3.1-8B.