Uncertainty-aware Language Modeling for Selective Question Answering

Abstract: We present an automatic LLM conversion approach that produces uncertainty-aware LLMs capable of estimating uncertainty with every prediction. Our approach is model- and data-agnostic, is computationally-efficient, and does not rely on external models or systems. We evaluate converted models on the selective question answering setting -- to answer as many questions as possible while maintaining a given accuracy, forgoing providing predictions when necessary. As part of our results, we test BERT and Llama 2 model variants on the SQuAD extractive QA task and the TruthfulQA generative QA task. We show that using the uncertainty estimates provided by our approach to selectively answer questions leads to significantly higher accuracy over directly using model probabilities.

• The paper introduces a method to convert LLMs to uncertainty-aware models by incorporating both aleatoric and epistemic uncertainty estimates.
• It employs the Capsa framework to modify models like BERT on SQuAD and Llama 2 on TruthfulQA, ensuring computationally efficient uncertainty integration.
• The study shows that uncertainty-informed predictions significantly improve selective QA accuracy compared to traditional logit-based confidence measures.

Uncertainty-aware Language Modeling for Selective Question Answering

This essay examines the methodology and results of the paper "Uncertainty-aware Language Modeling for Selective Question Answering" (2311.15451). The objective is to explore the integration of uncertainty estimates into LLMs to enhance the accuracy and reliability of selective question answering tasks.

Introduction

The paper presents a mechanism to enable LLMs to produce uncertainty estimates alongside predictions, thereby allowing them to selectively answer questions based on confidence levels. The need for such a mechanism arises due to the inherent limitations of LLMs in handling ambiguous information, out-of-domain data, and inconsistent training inputs. Traditional approaches have relied on predicted softmax probabilities to gauge prediction confidence, but this paper argues for a more robust uncertainty quantification (UQ) approach that aligns better with model prediction accuracy.

Figure 1: Robust, uncertainty-aware language modeling. Our methodology converts LLMs — agnostic of architecture — into uncertainty-aware variants and applies to generative (i.e., next-token prediction, left) and extractive (i.e., sub-context answering, right) models.

Methodology

Selective Question Answering

The core strategy outlined in the study involves converting existing LLMs into uncertainty-aware models capable of providing both aleatoric and epistemic uncertainty estimates. Aleatoric uncertainty captures data-related noise, while epistemic uncertainty addresses uncertainties in the model's knowledge.

Models and Datasets

Two types of models and datasets are examined: extractive models, using BERT on the SQuAD dataset, and generative models, using Llama 2 on the TruthfulQA dataset. Both models are adapted to output uncertainty measures for their predictions, which are then employed to determine whether a model should answer or abstain from a given question.

Conversion Process

The conversion process utilizes the Capsa framework to integrate UQ metrics into existing LLMs, transforming them into models that can manage and report uncertainty alongside their outputs. This model modification process is computationally efficient and architecture agnostic, ensuring broad applicability.

Results

Performance on Test Datasets

The study’s findings highlight that uncertainty-aware models signify a marked improvement in both coverage and accuracy of selective QA tasks compared to baseline models using logit-based confidence measures. Specifically, the uncertainty estimates derived from these converted models provide a more reliable means of determining the likelihood of a correct prediction.

Figure 2: Selective answering accuracy by confidence level. Increasing values of logit probability do not correspond to increased question answering ability — despite being often misconceived as a measure of confidence. Our methods report a reliable measure of confidence — with increased confidence corresponding to increased accuracy.

The study observes that in generative and extractive QA scenarios, high-confidence predictions using traditional logit probabilities actually result in lower accuracy, demonstrating the efficacy of UQ methods over conventional confidence estimates.

Implications and Future Directions

This research advances the field of AI by integrating UQ into LLM predictions for more reliable question answering, especially under uncertainty. Such techniques are critical for models deployed in applications where incorrect predictions could be costly, such as legal, medical, or customer service domains.

The implications extend beyond QA tasks to any domain requiring robust, adaptable LLMs capable of self-assessing reliability. Future work could explore extending this framework to integrate with external knowledge bases dynamically or apply these methods to broader NLP tasks that necessitate robustness and reliability.

Conclusion

The paper effectively illustrates the need for uncertainty-aware language modeling in enhancing LLM performance on QA tasks. By adopting uncertainty measures, these models achieve higher accuracy and reliability, offering significant benefits for applications requiring precise, contextually appropriate responses. This shift towards uncertainty-informed predictions represents a crucial step in the evolution of AI systems, aiming to balance the trade-offs between coverage and accuracy most effectively.