UCPO: Uncertainty-Aware Policy Optimization

AI summary

Overview

• Research area: Reinforcement learning (RL) for aligning large language models (LLMs), with a specific focus on uncertainty-aware learning and calibration of model abstention behavior.
• Technical level: Advanced. The paper builds directly on the GRPO advantage estimator, presents formal derivations of failure modes, and introduces two new optimization mechanisms (TAD and DURA) whose behavior depends on group reward statistics.
• Scope: This paper diagnoses why fixed uncertainty rewards inside GRPO produce either overconfidence or reward hacking, and proposes the Uncertainty-Aware Policy Optimization (UCPO) framework to fix it via Ternary Advantage Decoupling and Dynamic Uncertainty Reward Adjustment.

What This Paper Is About

LLMs tend to make confident but wrong assertions when a question falls outside their knowledge, which limits trustworthiness in high-stakes settings. A natural fix is to reward an LLM for expressing uncertainty (for example, using a fixed intermediate reward between "wrong" and "right"), but the authors show that this naive approach — which they call GRPO-UC — creates an advantage bias that pushes models toward either overconfidence or degenerate over-refusal. The paper's goal is to redesign the reward and advantage computation so that "right", "wrong", and "uncertain" responses can coexist in a stable equilibrium.

Key Contributions

1. Diagnosis of advantage bias in ternary reward spaces. The authors formally analyze why inserting a fixed uncertainty reward into GRPO produces two failure modes: "Majoritarian Suppression" of uncertainty in high-performance regimes (where the uncertainty advantage turns negative) and "Reward Hacking" in low-performance regimes (where the uncertainty advantage dominates the gradient and drives an all-uncertainty collapse).

2. Ternary Advantage Decoupling (TAD). A method that partitions a rollout group into a deterministic set (right ∪ wrong) and an uncertainty set, then normalizes them through two independent channels so that semantic signals do not interfere. The uncertainty advantage is anchored to the right-sample advantage via a gain γ(q).

3. Dynamic Uncertainty Reward Adjustment (DURA). A dual-term formula that modulates γ(q) in real time using the observed ratios of right (P_r), wrong (P_w), and uncertain (P_u) rollouts, including an "Uncertainty Gain" term and an "Uncertainty Suppression" term (with a default weight w = 1).

4. Supporting mechanisms and empirical validation. Non-Ternary Filtering (NTF) discards samples missing either correct or incorrect rollouts, and Low-Resource Extensions (LRE) add batch-level smoothing and non-linear mapping for small rollout budgets. The authors validate UCPO on Qwen3-8B and Llama-3.1-8B-Instruct across math/text reasoning and general tasks, and release code at https://github.com/xzhouzeng/ucpo.

Main Findings

• UCPO achieves the highest average PAQ across all evaluated domains. On Math & Text Reasoning, UCPO reaches an average PAQ of 79.63% on Qwen3-8B and 28.45% on Llama-3.1-8B-Instruct; on General Tasks it reaches 79.68% and 58.58% respectively.

• The primary gain comes from converting erroneous definitive answers into explicit uncertainty. Figure 5 shows that the PAQ improvement is mainly driven by moving responses from the "hallucination" category to the "uncertainty" category, rather than by making the model reason better.

• Fixed-reward GRPO-UC is brittle across task difficulty. In Math & Text Reasoning on Llama-3.1-8B-Instruct, GRPO-UC with r_u = 0.5 reaches the highest baseline average PAQ (24.90) but collapses F1 to 9.01; with r_u = 0.8 the F1 falls to 4.25. In General Tasks on the same model, a low r_u = 0.2 yields the strongest F1 among GRPO-UC variants (50.96), while the higher settings trade F1 for PAQ (r_u = 0.5 gives PAQ 56.55 / F1 45.25; r_u = 0.8 gives PAQ 55.66 / F1 44.65).

• Standard GRPO never learns uncertainty. The paper reports that GRPO consistently maintains a near-zero uncertainty ratio across training because its binary Right/Wrong reward offers no incentive for expressing doubt.

• Training dynamics confirm the ternary imbalance. On high-accuracy tasks (Qwen3-8B, Math & Text Reasoning), GRPO-UC's uncertainty ratio fluctuates near 0%. On low-accuracy tasks (Llama-3.1-8B-Instruct, Math & Text Reasoning), GRPO-UC with r_u ≥ 0.5 triggers reward hacking and the uncertainty ratio surges to 100%.

• Removing DURA causes performance collapse. In the ablation (Table 3), the variant without DURA over-optimizes for uncertainty rewards, showing a reward-hacking surge to roughly 80% uncertainty in both domains.

• Removing TAD degrades calibration. Without TAD, PAQ drops (Math PAQ 22.56; General PAQ 51.17), and the corresponding "uncertainty" metric swings to 50.33 (Math) and 12.40 (General), indicating that deterministic gradients overshadow the calibration signal, especially in simpler tasks.

• NTF and LRE improve robustness. Excluding NTF induces training fluctuations and suboptimal convergence; the full configuration (TAD + DURA + NTF + LRE) yields the strongest F1 and PAQ on Llama-3.1-8B-Instruct (Math: PAQ 28.45, F1 22.65; General: PAQ 58.58, F1 43.10).

• Parameter sensitivity follows a clear trade-off. For the DURA hyperparameter w, the paper states γ(q) ∈ (−w, 1) with sufficient rollouts, and w = 1 gives γ(q) = 0 when P_r = P_w = P_u = 1/3. Smaller w favors PAQ while larger w improves F1 at the cost of PAQ. For group size, G = 8 gives the best average PAQ while G = 16 further improves average F1 over G = 4.

• Metric behavior caveats are reported. On Qwen3-8B, where mathematical reasoning is already heavily trained, improvements in Accuracy or F1 may not be uniform, and UCPO's gains are read mainly as improved uncertainty calibration (reflected in PAQ) rather than stronger reasoning. On multiple-choice general tasks, lower F1 can arise because uncertain answers may be correct by chance, so UCPO's conversion of unreliable answers into explicit uncertainty can reduce coverage while improving reliability.

Methodology in Plain English

The researchers start from GRPO, a popular RL method that groups multiple sampled responses to the same prompt and normalizes rewards within the group. They first add an intermediate uncertainty reward (a fixed value between wrong and right) and show mathematically — with toy proportions P_r, P_w, and P_u — that this creates a bias: the uncertainty advantage becomes negative when most responses are correct, and becomes dominant when correct responses are rare.

To fix this, they split each group of rollouts into two channels. Deterministic responses (right or wrong) are normalized only against each other. Uncertain responses are instead given an advantage equal to a scaled version of the right-answer advantage, scaled by a factor γ(q). This stops uncertainty from being drowned out by many correct answers.

Next, they make γ(q) adaptive. It has two parts: a "gain" term that boosts honesty when the model makes many errors relative to uncertainty, and a "suppression" term that penalizes unnecessary avoidance as the model gets more things right. Together these act as a feedback controller that keeps the rate of uncertainty expression in a healthy middle range during training.

They also add NTF to discard prompts where the deterministic channel is degenerate (all correct or all wrong, mirroring GRPO's existing zero-advantage handling) and LRE to stabilize γ(q) estimation when only a few rollouts are available.

Experiments use DAPO-Math-17k for math/text training, with evaluation on AIME24, AMC, MATH500, Minerva, and Olympiad Bench, and MMLU-Redux2 (1,000 instances held out for testing, remainder for training) plus GPQA-Diamond for general tasks. Responses are labeled as Accuracy, Hallucination, or Uncertainty, and scored with PAQ = Acc/(Acc+Hal) and an F1 score balancing truthfulness and informativeness. Training uses 8 A100 GPUs with group size G = 8 and 0.6 decoding temperature for evaluation, with metrics averaged over three independent responses per instance. Baselines are Baseline, Prompt-UC, GRPO (Right = 1, Wrong = 0), and GRPO-UC with r_u ∈ {0.2, 0.5, 0.8}.

Why This Matters

Impact on research. This paper reframes uncertainty alignment as a reward-design problem rather than a data-annotation problem, and provides a formal account of why previous fixed-reward approaches fail. The identification of "advantage bias" in ternary reward spaces is a reusable concept for any RL pipeline that introduces a third intermediate reward category. By showing that hyperparameter-free dynamics can match or exceed the best tuned fixed reward, UCPO challenges the assumption that every new reward signal needs its own hyperparameter sweep.

Real-world applications.

• High-stakes question answering (medical, legal, financial): a model that says "I don't know" rather than fabricating an answer reduces downstream risk from confident errors.
• Customer-facing assistants: calibrated abstention prevents fabricated policy details or nonexistent product features from being presented to users as fact.
• Scientific or engineering copilots: a reasoning assistant that flags uncertainty at its knowledge boundary helps users know when to verify externally.
• Autonomous agent pipelines: where one agent's hallucination becomes another agent's input, uncertain outputs can be routed to human review instead of propagating.

Industry relevance. The authors are affiliated with Ant Group and Zhejiang University, and UCPO is trained on only 8 A100 GPUs, which is modest by current RL fine-tuning standards. The lack of a separate value function (inherited from GRPO) and the elimination of exhaustive reward tuning are practical advantages for teams deploying RLHF-style pipelines without large compute budgets. The release of code at https://github.com/xzhouzeng/ucpo further lowers the barrier to adoption.

Future Directions

1. Ternary signal distribution effects. The conclusion explicitly identifies as an unresolved limitation that the ratios of the different rollout types (right, wrong, uncertain) may themselves influence uncertainty learning, and notes this was observed but not fully explored.

2. More effective balancing strategies. The authors state that future work will investigate the specific impact of ternary signal distributions on training dynamics and explore more effective balancing strategies.

3. Low-resource robustness. The LRE extensions (batch-level smoothing and non-linear mapping) are described only in the supplementary material; how far they scale down, and how well the γ(q) estimator behaves with very few rollouts, remains an open question.

4. Better guarantees and deployment safeguards. The impact statement acknowledges that UCPO does not guarantee detection of all unknown queries and that practical deployment should pair it with task-specific risk assessment, oversight, and external verification to manage both over-abstention and residual errors — implying follow-up work on detection coverage and abstention calibration is needed.

Target Audience

This paper is most valuable to RL researchers and engineers working on LLM post-training, especially those building RLHF or GRPO-based pipelines who need their models to abstain gracefully rather than hallucinate. It is also relevant to trustworthiness, calibration, and uncertainty-quantification researchers who want an RL-based alternative to prompting or supervised fine-tuning with abstention labels. Practitioners deploying LLMs in regulated or high-stakes domains (medicine, law, finance) will find the PAQ/F1 framing and the reported failure modes of fixed uncertainty rewards directly actionable. Readers without a background in policy gradient methods or advantage estimation should expect to spend time on the preliminary and appendix sections, since the paper's core argument is mathematical.

Authors’ abstract

The key to building trustworthy large language models (LLMs) lies in endowing them with inherent uncertainty expression capabilities, thereby mitigating overconfident errors in 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 adapts 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 of the model beyond their knowledge boundaries.