UAlign: Leveraging Uncertainty Estimations for Factuality Alignment on Large Language Models

TL;DR

UAlign addresses the problem of factuality in LLMs by explicitly modeling knowledge boundaries through two uncertainty estimations, confidence and semantic entropy, and incorporating these estimates as input features to prompts during alignment. The framework constructs a knowledge-boundary dataset via multi-sample responses, trains uncertainty estimators and a reward model, and then applies PPO to improve factuality while encouraging refusals for unknowns. Empirical results show that UAlign enhances reliability and generalization on in-domain QA tasks and out-of-domain multilingual QA, outperforming prompt-, SFT-, RL-, and inference-based baselines. The approach offers a principled, uncertainty-aware path to safer, more truthful LLM deployment with modest additional compute, particularly when sampling is controlled and augmented by efficient inference strategies.

Abstract

Despite demonstrating impressive capabilities, Large Language Models (LLMs) still often struggle to accurately express the factual knowledge they possess, especially in cases where the LLMs' knowledge boundaries are ambiguous. To improve LLMs' factual expressions, we propose the UAlign framework, which leverages Uncertainty estimations to represent knowledge boundaries, and then explicitly incorporates these representations as input features into prompts for LLMs to Align with factual knowledge. First, we prepare the dataset on knowledge question-answering (QA) samples by calculating two uncertainty estimations, including confidence score and semantic entropy, to represent the knowledge boundaries for LLMs. Subsequently, using the prepared dataset, we train a reward model that incorporates uncertainty estimations and then employ the Proximal Policy Optimization (PPO) algorithm for factuality alignment on LLMs. Experimental results indicate that, by integrating uncertainty representations in LLM alignment, the proposed UAlign can significantly enhance the LLMs' capacities to confidently answer known questions and refuse unknown questions on both in-domain and out-of-domain tasks, showing reliability improvements and good generalizability over various prompt- and training-based baselines.

Figures (11)

• Figure 1: Examples of LLMs with (a) ambiguous and (b) explicit knowledge boundaries to answer questions.
• Figure 2: Illustration of UAlign dataset preparation process.
• Figure 3: Illustration of (a) SFT and (b) PPO alignment processes of UAlign framework. Note that for simplicity, we only present one estimation model in the figure but there are actually two.
• Figure 4: Illustration of the effects of different uses of uncertainty estimations under varying knowledge boundaries perceived by LLMs.
• Figure 5: Results of AUORC$\uparrow$ of several uncertainty estimation methods on TVQA using Llama-3 and Mistral.