Toward Better Generalisation in Uncertainty Estimators: Leveraging Data-Agnostic Features

TL;DR

This paper investigates combining data-agnostic features with hidden-state features to improve the out-of-domain generalisation of uncertainty estimators in LLMs, revealing mixed results and the importance of feature selection.

Contribution

It introduces a hybrid approach of data-agnostic and hidden-state features for uncertainty estimation, analyzing their impact on out-of-domain generalisation in LLMs.

Findings

Data-agnostic features often improve generalisation

In some cases, data-agnostic features degrade performance

Feature importance analysis shows underweighting of data-agnostic features

Abstract

Large Language Models (LLMs) often generate responses that are factually incorrect yet expressed with high confidence, which can pose serious risks for end users. To address this, it is essential for LLMs not only to produce answers but also to provide accurate estimates of their correctness. Uncertainty quantification methods have been introduced to assess the quality of LLM outputs, with factual accuracy being a key aspect of that quality. Among these methods, those that leverage hidden states to train probes have shown particular promise, as these internal representations encode information relevant to the factuality of responses, making this approach the focus of this paper. However, the probe trained on the hidden states of one dataset often struggles to generalise to another dataset of a different task or domain. To address this limitation, we explore combining data-agnostic features with hidden-state features and assess whether this hybrid feature set enhances out-of-domain performance. We further examine whether selecting only the most informative hidden-state features, thereby discarding task-specific noise, enables the data-agnostic features to contribute more effectively. The experiment results indicate that although introducing data-agnostic features generally enhances generalisation performance in most cases, in certain scenarios their inclusion degrades performance. A similar pattern emerges when retaining only the most important hidden-state features - adding data-agnostic features does not consistently further enhance performance compared to using the full set of hidden-state features. A closer analysis reveals that, in some specific cases, the trained probe underweights the data-agnostic features relative to the hidden-state features, which we believe is the main reason why the results are inconclusive.