Learned tokens improve large language model confidence estimates
Learned Readout Tokens for Extracting Large Language Model Confidence
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Large language models often give answers whose reliability we want to measure. The authors propose a new method that creates a special token to better understand how confident these models are in their answers. This new token helps assess correctness more accurately than older methods that just looked at existing parts of the model's output. They tested their approach on several datasets and models, showing consistent improvements in confidence prediction.
What this means in practice
- •For ai deployment teams: Improve trustworthiness checks for language model outputs by integrating learned readout tokens to more accurately flag unreliable responses.
- •For automated customer support developers: Enhance confidence scoring in chatbots to reduce errors impacting end users by using probes with learned tokens within the language model architecture.
Authors
Nima Seifi, Martin Ester
Abstract
Large language models (LLMs) are increasingly used within everyday applications where their responses have significant downstream impacts. For responsible deployment of such systems, we desire insight into how much we can trust their individual outputs. A popular approach for estimating LLM confidence is to train probes on the hidden states of the LLM's tokens to predict the correctness of the model's output. These probes typically read hidden states from tokens that are already in the model's vocabulary. In this work, we introduce a new approach to confidence probes that learn a new readout token and utilize its hidden states to train the probe. We show that our learned readout approach ScratchProbe outperforms its existing-token counterpart on a variety of datasets and LLMs. We also extend our method to a multi-token approach which we show provides further performance improvements. Finally, since the learned readout changes only the position a probe reads from, any existing probe can adopt it. Applied to a state-of-the-art probe, it improves AUROC on every backbone, in domain and out of domain.