Calibration-Aware Policy Optimization for Reasoning LLMs
Organizations: National Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institution of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Beijing National Research Center for Information Science and Technology, Tsinghua University
Abstract
Group Relative Policy Optimization (GRPO) enhances LLM reasoning but often induces overconfidence, where incorrect responses yield lower perplexity than correct ones, degrading relative calibration as described by the Area Under the Curve (AUC). Existing approaches either yield limited improvements in calibration or sacrifice gains in reasoning accuracy. We first prove that this degradation in GRPO-style algorithms stems from their uncertainty-agnostic advantage estimation, which inevitably misaligns optimization gradients with calibration. This leads to improved accuracy at the expense of degraded calibration. We then propose Calibration-Aware Policy Optimization (CAPO). It adopts a logistic AUC surrogate loss that is theoretically consistent and admits regret bound, enabling uncertainty-aware advantage estimation. By further incorporating a noise masking mechanism, CAPO achieves stable learning dynamics that jointly optimize calibration and accuracy. Experiments on multiple mathematical reasoning benchmarks show that CAPO-1.5B significantly improves calibration by up to 15% while achieving accuracy comparable to or better than GRPO, and further boosts accuracy on downstream inference-time scaling tasks by up to 5%. Moreover, when allowed to abstain under low-confidence conditions, CAPO achieves a Pareto-optimal precision-coverage trade-off, highlighting its practical value for hallucination mitigation.
Figures & tables
| Model | Ours | GRPO | GSPO | CoDaPO | CDE | SimKO |
|---|---|---|---|---|---|---|
| Qwen2.5-Math-1.5B | 25.33 | 20.33 | 20.00 | 21.67 | 16.67 | 11.67 |
| Qwen2.5-Math-7B | 38.33 | 33.33 | 32.21 | 31.66 | 31.66 | 23.33 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Configuration | Value |
| train_batch_size | 128 |
| ppo_mini_batch_size | 64 |
| max_prompt_length | 1050 |
| max_response_length | 3046 |
| ref-high | 2.5 |
| ref-low | 1.05 |
Language Model Calibration
Language model calibration studies whether models’ confidence estimates match their actual accuracy across tasks, prompts, and deployment settings.
See all 381 papers