UCPO: Uncertainty-Aware Policy Optimization

TL;DR

UCPO tackles uncertainty handling in LLMs by revealing why fixed, binary reward schemes cause either overconfidence or avoidance, and by introducing a ternary, uncertainty-aware RL framework. The method combines Ternary Advantage Decoupling (TAD), which isolates deterministic and uncertain signals, with Dynamic Uncertainty Reward Adjustment (DURA), which adaptively tunes the uncertainty gain in real time. Across math reasoning and general tasks, UCPO achieves superior uncertainty calibration (higher PAQ) and reduced hallucinations, while maintaining learning stability and robustness under varying difficulties and resource constraints. The approach also demonstrates compatibility with diverse RL methods (e.g., DAPO) and offers practical benefits for safer, more reliable LLM deployment in real-world settings.

Abstract

The key to building trustworthy Large Language Models (LLMs) lies in endowing them with inherent uncertainty expression capabilities to mitigate the hallucinations that restrict their high-stakes applications. However, existing RL paradigms such as GRPO often suffer from Advantage Bias due to binary decision spaces and static uncertainty rewards, inducing either excessive conservatism or overconfidence. To tackle this challenge, this paper unveils the root causes of reward hacking and overconfidence in current RL paradigms incorporating uncertainty-based rewards, based on which we propose the UnCertainty-Aware Policy Optimization (UCPO) framework. UCPO employs Ternary Advantage Decoupling to separate and independently normalize deterministic and uncertain rollouts, thereby eliminating advantage bias. Furthermore, a Dynamic Uncertainty Reward Adjustment mechanism is introduced to calibrate uncertainty weights in real-time according to model evolution and instance difficulty. Experimental results in mathematical reasoning and general tasks demonstrate that UCPO effectively resolves the reward imbalance, significantly improving the reliability and calibration of the model beyond their knowledge boundaries.

Figures (13)

• Figure 1: Illustration of reward imbalance in uncertainty alignment: static rewards trigger overconfidence or avoidance degeneracy, whereas UCPO stabilizes the policy through adaptive calibration.
• Figure 2: The Ternary Imbalance Problem in GRPO-UC (a-b) contrasted with the balanced advantage distribution in UCPO (c-d). Each point in the ternary plots represents a specific combination of Right, Wrong, and Uncertain proportions within a group of $G$ outputs.
• Figure 3: Architecture of the UCPO Framework.
• Figure 4: Evolution of the uncertainty ratio over training steps, comparing baseline GRPO, the proposed UCPO, and GRPO-UC variants with different reward coefficients $r_u$.
• Figure 5: Aggregated distribution of Accuracy, Hallucination and Uncertainty across different alignment methods. Proportions are averaged over all datasets independently within the Math & Text Reasoning and General Tasks domains.