Hidden Failures in Robustness: Why Supervised Uncertainty Quantification Needs Better Evaluation

2026-04-13

AI Summary

Existing probe-based supervised uncertainty quantification methods exhibit insufficient robustness under distribution shifts, with performance degrading notably in long-form text generation. This work systematically evaluates over 2,000 probes across diverse models, tasks, and out-of-distribution settings, employing multidimensional ablation studies to dissect the impact of representation layers, feature types, and token aggregation strategies. The analysis reveals that intermediate-layer representations combined with cross-token aggregation are critical for improving out-of-distribution generalization. Building on these insights, the authors propose a hybrid fallback strategy that substantially enhances probe reliability. This study provides empirical evidence and practical design guidelines for developing more robust uncertainty quantification probes in language generation.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Natural Language Processing: Safety and Robustness

Application Category

• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI

Abstract

Recent work has shown that the hidden states of large language models contain signals useful for uncertainty estimation and hallucination detection, motivating a growing interest in efficient probe-based approaches. Yet it remains unclear how robust existing methods are, and which probe designs provide uncertainty estimates that are reliable under distribution shift. We present a systematic study of supervised uncertainty probes across models, tasks, and OOD settings, training over 2,000 probes while varying the representation layer, feature type, and token aggregation strategy. Our evaluation highlights poor robustness in current methods, particularly in the case of long-form generations. We also find that probe robustness is driven less by architecture and more by the probe inputs. Middle-layer representations generalise more reliably than final-layer hidden states, and aggregating across response tokens is consistently more robust than relying on single-token features. These differences are often largely invisible in-distribution but become more important under distribution shift. Informed by our evaluation, we explore a simple hybrid back-off strategy for improving robustness, arguing that better evaluation is a prerequisite for building more robust probes.

Problem

• uncertainty quantification
• robustness
• distribution shift
• large language models
• probing

Innovation

• uncertainty quantification
• robustness evaluation
• distribution shift
• probe-based methods
• hidden states