UCPO: Uncertainty-Aware Policy Optimization

TL;DR

This paper introduces UCPO, a novel reinforcement learning framework that enhances the reliability and calibration of large language models by addressing advantage bias and dynamically adjusting uncertainty rewards.

Contribution

UCPO proposes Ternary Advantage Decoupling and Dynamic Uncertainty Reward Adjustment to improve uncertainty handling in RL for LLMs, reducing bias and overconfidence.

Findings

UCPO outperforms existing methods in mathematical reasoning tasks.

It significantly improves model calibration and reliability.

The framework effectively balances reward signals in uncertain environments.

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.