Improving the Reliability of Large Language Models by Leveraging Uncertainty-Aware In-Context Learning

TL;DR

The paper tackles hallucination in large language models by proposing an uncertainty-aware in-context learning framework that uses a calibration dataset to fine-tune the model for self-correction based on token-logit-derived uncertainty. It introduces a two-stage process: first compute an uncertainty score from token probabilities using single-inference methods and map these to discrete scores, then perform uncertainty-aware correction via calibration fine-tuning and test-time correction. Key findings show that logit outputs partially reflect inherent uncertainty and that the model can autonomously recognize uncertainty to improve responses, leading to notable accuracy gains and more reliable outputs. This approach offers a practical, relatively lightweight path to reduce unreliable generations in real-world deployments without extensive retraining.

Abstract

In recent years, large-scale language models (LLMs) have gained attention for their impressive text generation capabilities. However, these models often face the challenge of "hallucination," which undermines their reliability. In this study, we introduce an uncertainty-aware in-context learning framework to empower the model to enhance or reject its output in response to uncertainty. Human-defined methods for estimating uncertainty typically assume that "uncertainty is lower when the model's response is correct compared to when it is incorrect." However, setting a precise threshold to distinguish correctness is challenging. Therefore, we introduce uncertainty information as an intermediary variable that implicitly influences the model's behavior. Our innovative uncertainty-aware in-context learning framework involves fine-tuning the LLM using a calibration dataset. Our aim is to improve the model's responses by filtering out answers with high uncertainty while considering the model's knowledge limitations. We evaluate the model's knowledge by examining multiple responses to the same question for the presence of a correct answer. When the model lacks relevant knowledge, the response should indicate that the question cannot be answered. Conversely, when the model has relevant knowledge, the response should provide the correct answer. Extensive experiments confirm the effectiveness of our framework, leading to two key findings. First, the logit output values of the LLM partly reflect inherent uncertainty. Second, our model autonomously recognizes uncertainty, resulting in improved responses.

Figures (3)

• Figure 1: An illustration of automatic correction of LLM responses based on uncertainty score. Two scenarios are considered: firstly, the model has the relevant knowledge (illustrated in Figure (a)), and secondly, the model lacks relevant knowledge (illustrated in Figure (b)). If the model has the relevant knowledge, the incorrect response should be modified to the correct answer . Conversely, if the model lacks the relevant knowledge, the response should be altered to indicate that it cannot answer the question.
• Figure 2: An illustration of our uncertainty-aware framework. It reveals a two-step process comprising uncertainty score calculation and uncertainty-aware correction. In the first step, known as uncertainty score calculation, we undertake two subtasks: uncertainty estimation and uncertainty score mapping. This sequence results in the generation of an uncertainty score for each response within the answer set. Subsequently, we incorporate these uncertainty scores into the LLM, facilitating adaptive self-correction. Through a supervised fine-tuning procedure, we empower the model to explore and expand its knowledge boundaries. This fine-tuning process enables the model to autonomously refine its responses by taking into account the calculated uncertainty scores.
• Figure 3: An illustration of the change in model behavior as a result of uncertainty-aware fine-tuning. "Our-model" refer to the "Ours-Vicuna(Log-sum)" mentioned before.